Clean up global context cache use

[zhang-ivy]

I recently was using this code snippet to run one of my energy validation tests:

    from perses.tests.utils import validate_endstate_energies_point
   
    outdir = "/data/chodera/zhangi/perses_benchmark/repex/33/0/"

    for endstate in [0, 1]:
        with open(os.path.join(outdir, f"0_apo.pickle"), "rb") as f:
            htf = pickle.load(f)
            validate_endstate_energies_point(htf, endstate=endstate, minimize=True)

However, I'm getting this error:

Input In [2], in run_rest_factory_test(name)
     45 with open(os.path.join(outdir, f"0_{phase}.pickle"), "rb") as f:
     46     htf = pickle.load(f)
---> 48 validate_endstate_energies_point(htf, endstate=endstate, minimize=True)

File ~/miniconda3/envs/openmm-dev/lib/python3.9/site-packages/perses-0.9.5+94.g0fb13173.dirty-py3.9.egg/perses/tests/utils.py:950, in validate_endstate_energies_point(htf, endstate, minimize)
    948     hybrid_positions = sampler_state.positions
    949 context_hybrid.setPositions(hybrid_positions)
--> 950 components_hybrid = compute_potential_components(context_hybrid, beta=beta)
    952 # Get energy components of original system
    953 thermostate_other = states.ThermodynamicState(system=system, temperature=temperature)

File ~/miniconda3/envs/openmm-dev/lib/python3.9/site-packages/perses-0.9.5+94.g0fb13173.dirty-py3.9.egg/perses/tests/utils.py:307, in compute_potential_components(context, beta, platform)
    304 # Make a deep copy of the system.
    305 import copy
--> 307 from perses.dispersed.utils import configure_platform
    308 platform = configure_platform(platform.getName(), fallback_platform_name='Reference', precision='double')
    310 system = context.getSystem()

File ~/miniconda3/envs/openmm-dev/lib/python3.9/site-packages/perses-0.9.5+94.g0fb13173.dirty-py3.9.egg/perses/dispersed/utils.py:90, in <module>
     87     return platform
     89 #########
---> 90 cache.global_context_cache.platform = configure_platform(utils.get_fastest_platform().getName())
     91 #########
     92 
     93 #smc functions
     94 def compute_survival_rate(sMC_particle_ancestries):

File ~/miniconda3/envs/openmm-dev/lib/python3.9/site-packages/openmmtools-0.21.1+14.g025fc04-py3.9.egg/openmmtools/cache.py:335, in ContextCache.platform(self, new_platform)
    332 @platform.setter
    333 def platform(self, new_platform):
    334     if len(self._lru) > 0:
--> 335         raise RuntimeError('Cannot change platform of a non-empty ContextCache')
    336     if new_platform is None:
    337         self._platform_properties = None

RuntimeError: Cannot change platform of a non-empty ContextCache

This is because this line tries to set the platform of the global context cache, but was already set when this line was called, because feptasks.py has import openmmtools.cache as cache, which creates a global context cache here.

If I comment out this line, the problem goes away.

I think we need to go through the all the files in this repo and make sure that global context caches are being used properly.

[zhang-ivy]

related: #963

[ijpulidos]

@zhang-ivy Thanks for reporting this. I agree we need to monitor and clean things up to match the current context cache behavior from openmmtools. Can you share the pickled file so I can test and check this?

I believe the issue is not because the global context cache is created but instead, as the error says, it is because it's trying to change its platform after being used (non-empty), but I'd like to confirm this is indeed what's happening.

[zhang-ivy]

0_apo.pickle.zip

[zhang-ivy]

I believe the issue is not because the global context cache is created but instead, as the error says, it is because it's trying to change its platform after being used (non-empty), but I'd like to confirm this is indeed what's happening.

The error about changing the platform arises because either 1) the global context cache is being created somewhere else first (where it should not be created) or 2) the global context cache platform is attempting to be changed (where it should not be changed). We will have to discuss this further to identify which one is the problem, as it depends on what we want the desired behavior to be.

[ijpulidos]

We should remove the manual global context cache configuration in perses. It should happen automatically in openmmtools. CC choderalab/openmmtools#558

[ijpulidos]

@zhang-ivy I'm getting the following error while trying to load your serialized objects

Warning: importing 'simtk.openmm' is deprecated.  Import 'openmm' instead.
INFO:numexpr.utils:NumExpr defaulting to 8 threads.
Traceback (most recent call last):
  File "/home/user/workdir/debugging/perses-968/test.py", line 10, in <module>
    htf = pickle.load(f)
  File "/home/user/miniconda3/envs/perses-dev/lib/python3.9/site-packages/openmm/openmm.py", line 1916, in __setstate__
    system = XmlSerializer.deserialize(serializationString)
  File "/home/user/miniconda3/envs/perses-dev/lib/python3.9/site-packages/openmm/openmm.py", line 6928, in deserialize
    return XmlSerializer._deserializeForce(inputString)
  File "/home/user/miniconda3/envs/perses-dev/lib/python3.9/site-packages/openmm/openmm.py", line 6850, in _deserializeForce
    return _openmm.XmlSerializer__deserializeForce(inputString)
openmm.OpenMMException: Unsupported version number

Which I presume is because you are probably using a dev version of openmm or similar. Can you share the openmm version you are using? Just to be able to reproduce and test the solution to this issue. Thanks.

[zhang-ivy]

Yes its the nightly dev version. Here is the commit hash: bbd751bf1a8db98c65d856c1159bf79c8bec2e15

You can also use this test which builds the htf for you: https://github.com/choderalab/perses/blob/rest-over-protocol/perses/tests/test_relative.py#L1025 in the rest-over-protocol branch. I haven't tried reproducing the error using the test though, so I'm not sure what the behavior will be.

[zhang-ivy]

@ijpulidos : I just realized that the openmm version I used here is not from one of the nightly dev builds. I built one of Peter's branches from source (before he merged it into master).. in any case, I already sent you the environment yaml, but here it is again:
openmm-dev.yml.zip

[zhang-ivy]

related to: choderalab/openmmtools#558