Skip to content

Release the GIL for the duration of LambdaLever::setLambda. - #481

Merged
lohedges merged 2 commits into
develfrom
feature_scoped_gil_release
Sep 9, 2026
Merged

Release the GIL for the duration of LambdaLever::setLambda.#481
lohedges merged 2 commits into
develfrom
feature_scoped_gil_release

Conversation

@lohedges

@lohedges lohedges commented Sep 8, 2026

Copy link
Copy Markdown
Contributor

Profiling a replica exchange run with py-spy record --gil showed that setLambda accounted for 82% of all GIL-held time, serialising the worker threads against one another during the lambda switches that follow each dynamics block. The obvious fix, adding bp::release_gil_policy() to the wrapper, does nothing, because release_gil_policy is compiled out codebase-wide by SIRE_DISABLE_GIL_POLICY, and it cannot simply be re-enabled: boost/python/detail/caller.hpp runs the result converter between precall and postcall, so with the policy live every wrapped function that returns a converted value builds a Python object with the GIL released (enabling it segfaults during import sire). This adds Helpers/scoped_gil_release.hpp, an RAII guard that releases the GIL inside the wrapped call so that argument and result conversion still hold it, and uses it for setLambda via a thin hand-written wrapper. setLambda touches no Python: its OpenMM::Context argument is converted before the release, and the only QM interaction is a plain setter. A comment on SIRE_DISABLE_GIL_POLICY now records why it stays.

Measured on a 12-window, 4-worker repex run: GIL-held time fell from 28.6% to 6.4% of wall, for +6.4% throughput (mean cycle (This was a small system, so the gain is likely smaller on a full protein-ligand box, where MD dominates a larger share of the cycle.)

The same approach can be applied to other functions as profiling identifies them. The recipe is to run py-spy record --gil on a representative threaded workload, look for functions holding the GIL for a large share of the time, and wrap those with SireHelpers::ScopedGILRelease. Each candidate needs checking by hand first: the function must not re-enter Python, either directly or via a callback, and any from-Python argument conversion must complete before the GIL is released (boost constructs the arg_from_python converters ahead of the call, so this holds for ordinary arguments). This is deliberately case-by-case rather than blanket, i.e. the codebase-wide route is closed for the reasons above, and a wrapper is only worth adding where profiling shows it pays.

Debugged and implemented with help of Claude Code.

  • I confirm that I have merged the latest version of devel into this branch before issuing this pull request (e.g. by running git pull origin devel): [y]
  • I confirm that I have added a test for any new functionality in this pull request: [n]
  • I confirm that I have added documentation (e.g. a new tutorial page or detailed guide) for any new functionality in this pull request: [n]
  • I confirm that I have added a changelog entry to the changelog (we will add a link to this PR as part of the review): [y]
  • I confirm that I have permission to release this code under the GPL3 license: [y]

@chryswoods chryswoods left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍

@lohedges

lohedges commented Sep 9, 2026

Copy link
Copy Markdown
Contributor Author

I've run some checks on takakura and the results are nice. This first plot shows how the production time scales vs GROMACS using MPS on a single GPU.

oversubscription_benchmark_gil_release

The next shows what happens when you don't use MPS (or oversubscribe), but simulate an increasing number of windows on all 3 GPUs. (Note that I've deliberately used a replica count that is a multiple of three to divide work equally between the GPUs.)

oversubscription_benchmark_multigpu_no_mps

The upshot is that the benefit scales with how much concurrent work is processed, i.e. both the number of workers and the number of windows they cycle through, so you get a benefit from MPS on a single GPU, or from dividing the work over multiple GPUs, with a double benefit from both.

@lohedges
lohedges merged commit 535fab1 into devel Sep 9, 2026
9 of 10 checks passed
@lohedges
lohedges deleted the feature_scoped_gil_release branch September 9, 2026 11:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants