BUG placed context stack inside thread-local data space (#1555)

* placed context stack inside thread-local data space

* fixed thread-local context manager, and added a regression test

* fixed docstring of #1552 regression test

Author: burnpanck

pymc3/model.py

@@ -1,3 +1,5 @@
+import threading
+
 import numpy as np
 import theano
 import theano.tensor as tt
@@ -95,6 +97,7 @@ class Context(object):
     """Functionality for objects that put themselves in a context using
     the `with` statement.
     """
+    contexts = threading.local()
 
     def __enter__(self):
         type(self).get_contexts().append(self)
@@ -105,10 +108,11 @@ def __exit__(self, typ, value, traceback):
 
     @classmethod
     def get_contexts(cls):
-        if not hasattr(cls, "contexts"):
-            cls.contexts = []
-
-        return cls.contexts
+        # no race-condition here, cls.contexts is a thread-local object
+        # be sure not to override contexts in a subclass however!
+        if not hasattr(cls.contexts, 'stack'):
+            cls.contexts.stack = []
+        return cls.contexts.stack
 
     @classmethod
     def get_context(cls):

pymc3/tests/test_modelcontext.py

@@ -0,0 +1,48 @@
+import threading
+import unittest
+
+from pymc3 import Model, Normal
+
+
+class TestModelContext(unittest.TestCase):
+    def test_thread_safety(self):
+        """ Regression test for issue #1552: Thread safety of model context manager
+
+        This test creates two threads that attempt to construct two
+        unrelated models at the same time.
+        For repeatable testing, the two threads are syncronised such
+        that thread A enters the context manager first, then B,
+        then A attempts to declare a variable while B is still in the context manager.
+        """
+        aInCtxt,bInCtxt,aDone = [threading.Event() for k in range(3)]
+        modelA = Model()
+        modelB = Model()
+        def make_model_a():
+            with modelA:
+                aInCtxt.set()
+                bInCtxt.wait()
+                a = Normal('a',0,1)
+            aDone.set()
+        def make_model_b():
+            aInCtxt.wait()
+            with modelB:
+                bInCtxt.set()
+                aDone.wait()
+                b = Normal('b', 0, 1)
+        threadA = threading.Thread(target=make_model_a)
+        threadB = threading.Thread(target=make_model_b)
+        threadA.start()
+        threadB.start()
+        threadA.join()
+        threadB.join()
+        # now let's see which model got which variable
+        # previous to #1555, the variables would be swapped:
+        # - B enters it's model context after A, but before a is declared -> a goes into B
+        # - A leaves it's model context before B attempts to declare b. A's context manager
+        #   takes B from the stack, such that b ends up in model A
+        self.assertEqual(
+            (
+                list(modelA.named_vars),
+                list(modelB.named_vars),
+            ), (['a'],['b'])
+        )