Repository: choderalab/openmmtools

Issue:
Context cache failure to access `set*`
Unfortunately, I don't have a minimal example yet, but I did want to note that a few days ago, as I was trying to use the `ContextCache` with a `celery` worker, I ran into a strange bug. Using the `AlchemicalNonequilibriumLangevinIntegrator`, whenever it tried to recover a context and integrator, I got an exception that there was no attribute `set*` (including `setTemperature`, etc.).

 After some debugging, I noticed that the `AlchemicalNonequilibriumLangevinIntegrator` was becoming a regular `CustomIntegrator` somehow, but I'm not quite sure how or why.

This isn't super high priority, since I have a workaround for now, but since the `ContextCache` API is really awesome and convenient, it would be great to figure out what's going on.

Discussion:

Patch:
diff --git a/openmmtools/cache.py b/openmmtools/cache.py
index f6d3ede..87ec8d6 100644
--- a/openmmtools/cache.py
+++ b/openmmtools/cache.py
@@ -17,7 +17,7 @@ Provide cache classes to handle creation of OpenMM Context objects.
 import copy
 import collections
 
-from simtk import openmm
+from simtk import openmm, unit
 
 from openmmtools import integrators
 
@@ -114,12 +114,26 @@ class LRUCache(object):
     @time_to_live.setter
     def time_to_live(self, new_time_to_live):
         # Update entries only if we are changing the ttl.
-        ttl_diff = new_time_to_live - self._ttl
-        if ttl_diff == 0:
+        if new_time_to_live == self._ttl:
             return
+
+        # Update expiration of cache entries.
         for entry in self._data.values():
-            entry.expiration += ttl_diff
-        self._remove_expired()
+            # If there was no time to live before, just let entries
+            # expire in new_time_to_live accesses
+            if self._ttl is None:
+                entry.expiration = self._n_access + new_time_to_live
+            # If we don't want expiration anymore, delete the field.
+            # This way we save memory in case there are a lot of entries.
+            elif new_time_to_live is None:
+                del entry.expiration
+            # Otherwise just add/subtract the difference.
+            else:
+                entry.expiration += new_time_to_live - self._ttl
+
+        # Purge cache only if there is a time to live.
+        if new_time_to_live is not None:
+            self._remove_expired()
         self._ttl = new_time_to_live
 
     def empty(self):
@@ -366,11 +380,18 @@ class ContextCache(object):
             thermodynamic_state_id = self._generate_state_id(thermodynamic_state)
             matching_context_ids = [context_id for context_id in self._lru
                                     if context_id[0] == thermodynamic_state_id]
-            if len(matching_context_ids) > 0:
-                context = self._lru[matching_context_ids[0]]  # Return first found.
+            if len(matching_context_ids) == 0:
+                # We have to create a new Context.
+                integrator = self._get_default_integrator(thermodynamic_state.temperature)
+            elif len(matching_context_ids) == 1:
+                # Only one match.
+                context = self._lru[matching_context_ids[0]]
             else:
-                # We have to create a new Context. Use a likely-to-be-used Integrator.
-                integrator = integrators.GeodesicBAOABIntegrator(temperature=thermodynamic_state.temperature)
+                # Multiple matches, prefer non-default Integrator.
+                for context_id in matching_context_ids:
+                    if context_id[1] != self._default_integrator_id():
+                        context = self._lru[context_id]
+                        break
 
         if context is None:
             # Determine the Context id matching the pair state-integrator.
@@ -447,6 +468,7 @@ class ContextCache(object):
 
         """
         standard_integrator = copy.deepcopy(integrator)
+        integrators.RestorableIntegrator.restore_interface(standard_integrator)
         for attribute, std_value in cls._COMPATIBLE_INTEGRATOR_ATTRIBUTES.items():
             try:
                 getattr(standard_integrator, 'set' + attribute)(std_value)
@@ -459,7 +481,13 @@ class ContextCache(object):
         """Return a unique key for the ThermodynamicState."""
         # We take advantage of the cached _standard_system_hash property
         # to generate a compatible hash for the thermodynamic state.
-        return str(thermodynamic_state._standard_system_hash)
+        return thermodynamic_state._standard_system_hash
+
+    @classmethod
+    def _generate_integrator_id(cls, integrator):
+        """Return a unique key for the given Integrator."""
+        standard_integrator = cls._standardize_integrator(integrator)
+        return openmm.XmlSerializer.serialize(standard_integrator).__hash__()
 
     @classmethod
     def _generate_context_id(cls, thermodynamic_state, integrator):
@@ -472,10 +500,25 @@ class ContextCache(object):
 
         """
         state_id = cls._generate_state_id(thermodynamic_state)
-        standard_integrator = cls._standardize_integrator(integrator)
-        integrator_id = openmm.XmlSerializer.serialize(standard_integrator).__hash__()
+        integrator_id = cls._generate_integrator_id(integrator)
         return state_id, integrator_id
 
+    @staticmethod
+    def _get_default_integrator(temperature):
+        """Return a new instance of the default integrator."""
+        # Use a likely-to-be-used Integrator.
+        return integrators.GeodesicBAOABIntegrator(temperature=temperature)
+
+    @classmethod
+    def _default_integrator_id(cls):
+        """Return the unique key of the default integrator."""
+        if cls._cached_default_integrator_id is None:
+            default_integrator = cls._get_default_integrator(300*unit.kelvin)
+            default_integrator_id = cls._generate_integrator_id(default_integrator)
+            cls._cached_default_integrator_id = default_integrator_id
+        return cls._cached_default_integrator_id
+    _cached_default_integrator_id = None
+
 
 # =============================================================================
 # DUMMY CONTEXT CACHE
@@ -523,7 +566,7 @@ class DummyContextCache(object):
 # GLOBAL CONTEXT CACHE
 # =============================================================================
 
-global_context_cache = ContextCache(capacity=3, time_to_live=50)
+global_context_cache = ContextCache(capacity=None, time_to_live=None)
 
 
 # =============================================================================
diff --git a/openmmtools/integrators.py b/openmmtools/integrators.py
index ecdb460..ce87c38 100644
--- a/openmmtools/integrators.py
+++ b/openmmtools/integrators.py
@@ -38,6 +38,7 @@ this program.  If not, see <http://www.gnu.org/licenses/>.
 import numpy as np
 import logging
 import re
+import zlib
 
 import simtk.unit
 
@@ -213,9 +214,18 @@ class RestorableIntegrator(mm.CustomIntegrator,PrettyPrintableIntegrator):
 
     @staticmethod
     def _compute_class_hash(integrator_class):
-        """Return a numeric hash for the integrator class."""
-        # We need to convert to float because some digits may be lost in the conversion
-        return float(hash(integrator_class.__name__))
+        """Return a numeric hash for the integrator class.
+
+        The hash will become part of the Integrator serialization,
+        so it is important for it consistent across processes in case
+        the integrator is sent to a remote worker. The hash() built-in
+        function is seeded by the PYTHONHASHSEED environmental variable,
+        so we can't use it here.
+
+        We also need to convert to float because some digits may be
+        lost in the conversion.
+        """
+        return float(zlib.adler32(integrator_class.__name__.encode()))
 
 
 class ThermostatedIntegrator(RestorableIntegrator):


Test patch:
diff --git a/openmmtools/tests/test_cache.py b/openmmtools/tests/test_cache.py
index de9b008..ef8a8ef 100644
--- a/openmmtools/tests/test_cache.py
+++ b/openmmtools/tests/test_cache.py
@@ -124,14 +124,32 @@ def test_lru_cache_capacity_property():
 
 def test_lru_cache_time_to_live_property():
     """Decreasing the time to live updates the expiration of elements."""
-    cache = LRUCache(time_to_live=50)
-    for i in range(4):
-        cache[str(i)] = i
+    def add_4_elements(_cache):
+        for i in range(4):
+            _cache[str(i)] = i
+
+    cache = LRUCache(time_to_live=None)
+    add_4_elements(cache)
     assert len(cache) == 4
+
+    # Setting time to live to 1 cause all 4 entries to expire on next access.
     cache.time_to_live = 1
-    assert len(cache) == 1
+    assert len(cache) == 4
     assert cache.time_to_live == 1
-    assert '3' in cache
+    cache['4'] = 4
+    assert len(cache) == 1
+    assert '4' in cache
+
+    # Increase time_to_live.
+    cache.time_to_live = 2
+    add_4_elements(cache)
+    assert len(cache) == 2
+    assert '2' in cache and '3' in cache
+
+    # Setting it back to None makes it limitless.
+    cache.time_to_live = None
+    add_4_elements(cache)
+    assert len(cache) == 4
 
 
 # =============================================================================
@@ -226,19 +244,25 @@ class TestContextCache(object):
         cache = ContextCache()
         state1, state2 = self.incompatible_states[:2]
 
-        # First we create a Context in state1.
-        cache.get_context(state1, copy.deepcopy(self.verlet_2fs))
+        # If there are no previous Contexts, a default integrator is used to create a new one.
+        context, default_integrator = cache.get_context(state1)
         assert len(cache) == 1
 
-        # When we don't specify the integrator, it first looks for cached Contexts.
+        # Now we create another Context in state1 with a different integrator.
+        assert type(self.verlet_2fs) is not type(default_integrator)  # test precondition
+        cache.get_context(state1, copy.deepcopy(self.verlet_2fs))
+        assert len(cache) == 2
+
+        # When we don't specify the integrator, it first looks
+        # for cached Contexts, and non-default contexts are preferred.
         context, integrator = cache.get_context(state1)
-        assert len(cache) == 1
+        assert len(cache) == 2
         assert state1.is_context_compatible(context)
-        assert isinstance(integrator, openmm.VerletIntegrator)
+        assert isinstance(integrator, type(self.verlet_2fs)), type(integrator)
 
         # With an incompatible state, a new Context is created.
         cache.get_context(state2)
-        assert len(cache) == 2
+        assert len(cache) == 3
 
     def test_cache_capacity_ttl(self):
         """Check that the cache capacity and time_to_live work as expected."""