pytorch/caffe2/python/context_test.py
Alisson Gusatti Azzolini b7783a1976 Make ContextManager thread-safe
Summary: ContextManager was thread local. This caused issues because the context registration needs to be global. What needs to be thread local is the current context.

Reviewed By: jhcross

Differential Revision: D4556050

fbshipit-source-id: 5de1c0d9fd0a778c4cb1eadef01f9a1ab488f603
2017-02-13 19:45:35 -08:00

36 lines
951 B
Python

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
from caffe2.python import context, test_util
from threading import Thread
@context.define_context()
class MyContext(object):
pass
class TestContext(test_util.TestCase):
def use_my_context(self):
try:
for _ in range(100):
with MyContext() as a:
for _ in range(100):
self.assertTrue(MyContext.current() == a)
except Exception as e:
self._exceptions.append(e)
def testMultiThreaded(self):
threads = []
self._exceptions = []
for _ in range(8):
thread = Thread(target=self.use_my_context)
thread.start()
threads.append(thread)
for t in threads:
t.join()
for e in self._exceptions:
raise e