disable jax if PYMBAR_DISABLE_JAX is set#569
Conversation
| def test_env_values(self, monkeypatch, env_value, expected): | ||
| """Test various environment variable values.""" | ||
| monkeypatch.setenv("PYMBAR_DISABLE_JAX", env_value) | ||
| assert _setup_jax_accleration() is expected |
There was a problem hiding this comment.
Do we want to specify an error message in case the assertion fails? Something like "received X but expected Y"?
There was a problem hiding this comment.
pytest will do that for us
ijpulidos
left a comment
ijpulidos
left a comment
There was a problem hiding this comment.
This is looking great, I think this works for our purposes. Just some minor changes/comments.
In the long term I wonder if we want to make the MBAR object itself specify the backend via an argument.
Co-authored-by: Iván Pulido <2949729+ijpulidos@users.noreply.github.com>
Co-authored-by: Iván Pulido <2949729+ijpulidos@users.noreply.github.com>
…x-at-runtime-option' into feat/add-disable-jax-at-runtime-option
|
#509 is how we will fix this long term, but this should work for now. |
| ("yes", True), | ||
| ("YES", True), | ||
| ("YeS", True), | ||
| ("0", True), |
There was a problem hiding this comment.
I would expect 0=False and 1=True, no?
There was a problem hiding this comment.
duh! I was thinking about 0 being the "good" exit code and that was "true" to be BUT as soon as you mention that, it is the opposite of the convention
Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com>
ijpulidos
left a comment
ijpulidos
left a comment
There was a problem hiding this comment.
Looking great. One minor comment, but otherwise LGTM.
|
|
||
| However, many packages released on conda-forge specify `pymbar` as a dependency, making it impossible to install the non-jax accelerated of PyMbar. | ||
| In `pymbar >= 4.2.0` JAX acceleration can be controlled with the environmental variable `PYMBAR_DISABLE_JAX`. | ||
| Setting `PYMBAR_DISABLE_JAX` to a truth-y value (`TRUE`, `YES`, `0`) will disable JAX acceleration even if the `jax` package is installed. |
There was a problem hiding this comment.
I guess we mean here 1 instead of 0 as a truth-y value?
…x-at-runtime-option' into feat/add-disable-jax-at-runtime-option
3 participants
This will let a user disable jax at runtime. A larger refactor is needed to test this since at import time the
PYMBAR_DISABLE_JAXvariable is checked, so even if monkey patched the module has already been imported. This also means that once pymbar has been imported, settingPYMBAR_DISABLE_JAXwill do nothing.TODO: add docs