Clean up global context cache use

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]

0_apo.pickle.zip

[zhang-ivy]

related to: https://github.com/choderalab/openmmtools/issues/558