#!/usr/bin/env python3 from hypothesis import given, settings import hypothesis.strategies as st from multiprocessing import Process, Queue import numpy as np import os import pickle import tempfile import shutil from caffe2.python import core, workspace, dyndep import caffe2.python.hypothesis_test_util as hu from gloo.python import IoError dyndep.InitOpsLibrary("@/caffe2/caffe2/distributed:file_store_handler_ops") dyndep.InitOpsLibrary("@/caffe2/caffe2/distributed:redis_store_handler_ops") dyndep.InitOpsLibrary("@/caffe2/caffe2/distributed:store_ops") dyndep.InitOpsLibrary("@/caffe2/caffe2/contrib/gloo:gloo_ops") dyndep.InitOpsLibrary("@/caffe2/caffe2/contrib/gloo:gloo_ops_gpu") op_engine = 'GLOO' class TemporaryDirectory: def __enter__(self): self.tmpdir = tempfile.mkdtemp() return self.tmpdir def __exit__(self, type, value, traceback): shutil.rmtree(self.tmpdir) class TestCase(hu.HypothesisTestCase): test_counter = 0 sync_counter = 0 def run_test_locally(self, fn, device_option=None, **kwargs): # Queue for assertion errors on subprocesses queue = Queue() # Capture any exception thrown by the subprocess def run_fn(*args, **kwargs): try: with core.DeviceScope(device_option): fn(*args, **kwargs) workspace.ResetWorkspace() queue.put(True) except Exception as ex: queue.put(ex) # Start N processes in the background procs = [] for i in range(kwargs['comm_size']): kwargs['comm_rank'] = i proc = Process( target=run_fn, kwargs=kwargs) proc.start() procs.append(proc) # Test complete, join background processes while len(procs) > 0: proc = procs.pop(0) while proc.is_alive(): proc.join(10) # Raise exception if we find any. Otherwise each worker # should put a True into the queue # Note that the following is executed ALSO after # the last process was joined, so if ANY exception # was raised, it will be re-raised here. self.assertFalse(queue.empty(), "Job failed without a result") o = queue.get() if isinstance(o, Exception): raise o else: self.assertTrue(o) def run_test_distributed(self, fn, device_option=None, **kwargs): comm_rank = os.getenv('COMM_RANK') self.assertIsNotNone(comm_rank) comm_size = os.getenv('COMM_SIZE') self.assertIsNotNone(comm_size) kwargs['comm_rank'] = int(comm_rank) kwargs['comm_size'] = int(comm_size) with core.DeviceScope(device_option): fn(**kwargs) workspace.ResetWorkspace() def create_common_world(self, comm_rank, comm_size, tmpdir=None, existing_cw=None): store_handler = "store_handler" # If REDIS_HOST is set, use RedisStoreHandler for rendezvous. if existing_cw is None: redis_host = os.getenv("REDIS_HOST") redis_port = int(os.getenv("REDIS_PORT", 6379)) if redis_host is not None: workspace.RunOperatorOnce( core.CreateOperator( "RedisStoreHandlerCreate", [], [store_handler], prefix=str(TestCase.test_counter) + "/", host=redis_host, port=redis_port)) else: workspace.RunOperatorOnce( core.CreateOperator( "FileStoreHandlerCreate", [], [store_handler], path=tmpdir)) common_world = "common_world" else: common_world = str(existing_cw) + ".forked" if existing_cw is not None: workspace.RunOperatorOnce( core.CreateOperator( "CloneCommonWorld", [existing_cw], [common_world], sync=True, engine=op_engine)) else: workspace.RunOperatorOnce( core.CreateOperator( "CreateCommonWorld", [store_handler], [common_world], size=comm_size, rank=comm_rank, sync=True, engine=op_engine)) return (store_handler, common_world) def synchronize(self, store_handler, value, comm_rank=None): TestCase.sync_counter += 1 blob = "sync_{}".format(TestCase.sync_counter) if comm_rank == 0: workspace.FeedBlob(blob, pickle.dumps(value)) workspace.RunOperatorOnce( core.CreateOperator( "StoreSet", [store_handler, blob], [])) else: workspace.RunOperatorOnce( core.CreateOperator( "StoreGet", [store_handler], [blob])) return pickle.loads(workspace.FetchBlob(blob)) def _test_broadcast(self, comm_rank=None, comm_size=None, blob_size=None, num_blobs=None, tmpdir=None, use_float16=False, ): store_handler, common_world = self.create_common_world( comm_rank=comm_rank, comm_size=comm_size, tmpdir=tmpdir) blob_size = self.synchronize( store_handler, blob_size, comm_rank=comm_rank) num_blobs = self.synchronize( store_handler, num_blobs, comm_rank=comm_rank) for i in range(comm_size): blobs = [] for j in range(num_blobs): blob = "blob_{}".format(j) offset = (comm_rank * num_blobs) + j value = np.full(blob_size, offset, np.float16 if use_float16 else np.float32) workspace.FeedBlob(blob, value) blobs.append(blob) net = core.Net("broadcast") net.Broadcast( [common_world] + blobs, blobs, root=i, engine=op_engine) workspace.CreateNet(net) workspace.RunNet(net.Name()) for j in range(num_blobs): np.testing.assert_array_equal( workspace.FetchBlob(blobs[j]), i * num_blobs) # Run the net a few more times to check the operator # works not just the first time it's called for _tmp in range(4): workspace.RunNet(net.Name()) @given(comm_size=st.integers(min_value=2, max_value=8), blob_size=st.integers(min_value=int(1e3), max_value=int(1e6)), num_blobs=st.integers(min_value=1, max_value=4), device_option=st.sampled_from([hu.cpu_do]), use_float16=st.booleans()) @settings(deadline=10000) def test_broadcast(self, comm_size, blob_size, num_blobs, device_option, use_float16): TestCase.test_counter += 1 if os.getenv('COMM_RANK') is not None: self.run_test_distributed( self._test_broadcast, blob_size=blob_size, num_blobs=num_blobs, use_float16=use_float16, device_option=device_option) else: with TemporaryDirectory() as tmpdir: self.run_test_locally( self._test_broadcast, comm_size=comm_size, blob_size=blob_size, num_blobs=num_blobs, device_option=device_option, tmpdir=tmpdir, use_float16=use_float16) def _test_allreduce(self, comm_rank=None, comm_size=None, blob_size=None, num_blobs=None, tmpdir=None, use_float16=False ): store_handler, common_world = self.create_common_world( comm_rank=comm_rank, comm_size=comm_size, tmpdir=tmpdir) blob_size = self.synchronize( store_handler, blob_size, comm_rank=comm_rank) num_blobs = self.synchronize( store_handler, num_blobs, comm_rank=comm_rank) blobs = [] for i in range(num_blobs): blob = "blob_{}".format(i) value = np.full(blob_size, (comm_rank * num_blobs) + i, np.float16 if use_float16 else np.float32) workspace.FeedBlob(blob, value) blobs.append(blob) net = core.Net("allreduce") net.Allreduce( [common_world] + blobs, blobs, engine=op_engine) workspace.CreateNet(net) workspace.RunNet(net.Name()) for i in range(num_blobs): np.testing.assert_array_equal( workspace.FetchBlob(blobs[i]), (num_blobs * comm_size) * (num_blobs * comm_size - 1) / 2) # Run the net a few more times to check the operator # works not just the first time it's called for _tmp in range(4): workspace.RunNet(net.Name()) def _test_allreduce_multicw(self, comm_rank=None, comm_size=None, tmpdir=None ): _store_handler, common_world = self.create_common_world( comm_rank=comm_rank, comm_size=comm_size, tmpdir=tmpdir) _, common_world2 = self.create_common_world( comm_rank=comm_rank, comm_size=comm_size, tmpdir=tmpdir, existing_cw=common_world) blob_size = int(1e4) num_blobs = 4 for cw in [common_world, common_world2]: blobs = [] for i in range(num_blobs): blob = "blob_{}".format(i) value = np.full(blob_size, (comm_rank * num_blobs) + i, np.float32) workspace.FeedBlob(blob, value) blobs.append(blob) net = core.Net("allreduce_multicw") net.Allreduce( [cw] + blobs, blobs, engine=op_engine) workspace.RunNetOnce(net) for i in range(num_blobs): np.testing.assert_array_equal( workspace.FetchBlob(blobs[i]), (num_blobs * comm_size) * (num_blobs * comm_size - 1) / 2) @given(comm_size=st.integers(min_value=2, max_value=8), blob_size=st.integers(min_value=int(1e3), max_value=int(1e6)), num_blobs=st.integers(min_value=1, max_value=4), device_option=st.sampled_from([hu.cpu_do]), use_float16=st.booleans()) @settings(deadline=10000) def test_allreduce(self, comm_size, blob_size, num_blobs, device_option, use_float16): TestCase.test_counter += 1 if os.getenv('COMM_RANK') is not None: self.run_test_distributed( self._test_allreduce, blob_size=blob_size, num_blobs=num_blobs, use_float16=use_float16, device_option=device_option) else: with TemporaryDirectory() as tmpdir: self.run_test_locally( self._test_allreduce, comm_size=comm_size, blob_size=blob_size, num_blobs=num_blobs, device_option=device_option, tmpdir=tmpdir, use_float16=use_float16) def _test_reduce_scatter(self, comm_rank=None, comm_size=None, blob_size=None, num_blobs=None, tmpdir=None, use_float16=False ): store_handler, common_world = self.create_common_world( comm_rank=comm_rank, comm_size=comm_size, tmpdir=tmpdir) blob_size = self.synchronize( store_handler, blob_size, comm_rank=comm_rank) num_blobs = self.synchronize( store_handler, num_blobs, comm_rank=comm_rank) blobs = [] for i in range(num_blobs): blob = "blob_{}".format(i) value = np.full(blob_size, (comm_rank * num_blobs) + i, np.float16 if use_float16 else np.float32) workspace.FeedBlob(blob, value) blobs.append(blob) # Specify distribution among ranks i.e. number of elements # scattered/distributed to each process. recv_counts = np.zeros(comm_size, dtype=np.int32) remaining = blob_size chunk_size = (blob_size + comm_size - 1) / comm_size for i in range(comm_size): recv_counts[i] = min(chunk_size, remaining) remaining = remaining - chunk_size if remaining > chunk_size else 0 recv_counts_blob = "recvCounts" workspace.FeedBlob(recv_counts_blob, recv_counts) blobs.append(recv_counts_blob) net = core.Net("reduce_scatter") net.ReduceScatter( [common_world] + blobs, blobs, engine=op_engine) workspace.CreateNet(net) workspace.RunNet(net.Name()) for i in range(num_blobs): np.testing.assert_array_equal( np.resize(workspace.FetchBlob(blobs[i]), recv_counts[comm_rank]), (num_blobs * comm_size) * (num_blobs * comm_size - 1) / 2) # Run the net a few more times to check the operator # works not just the first time it's called for _tmp in range(4): workspace.RunNet(net.Name()) @given(comm_size=st.integers(min_value=2, max_value=8), blob_size=st.integers(min_value=int(1e3), max_value=int(1e6)), num_blobs=st.integers(min_value=1, max_value=4), device_option=st.sampled_from([hu.cpu_do]), use_float16=st.booleans()) @settings(deadline=10000) def test_reduce_scatter(self, comm_size, blob_size, num_blobs, device_option, use_float16): TestCase.test_counter += 1 if os.getenv('COMM_RANK') is not None: self.run_test_distributed( self._test_reduce_scatter, blob_size=blob_size, num_blobs=num_blobs, use_float16=use_float16, device_option=device_option) else: with TemporaryDirectory() as tmpdir: self.run_test_locally( self._test_reduce_scatter, comm_size=comm_size, blob_size=blob_size, num_blobs=num_blobs, device_option=device_option, tmpdir=tmpdir, use_float16=use_float16) def _test_allgather(self, comm_rank=None, comm_size=None, blob_size=None, num_blobs=None, tmpdir=None, use_float16=False ): store_handler, common_world = self.create_common_world( comm_rank=comm_rank, comm_size=comm_size, tmpdir=tmpdir) blob_size = self.synchronize( store_handler, blob_size, comm_rank=comm_rank) num_blobs = self.synchronize( store_handler, num_blobs, comm_rank=comm_rank) blobs = [] for i in range(num_blobs): blob = "blob_{}".format(i) value = np.full(blob_size, (comm_rank * num_blobs) + i, np.float16 if use_float16 else np.float32) workspace.FeedBlob(blob, value) blobs.append(blob) net = core.Net("allgather") net.Allgather( [common_world] + blobs, ["Gathered"], engine=op_engine) workspace.CreateNet(net) workspace.RunNet(net.Name()) # create expected output expected_output = np.array([]) for i in range(comm_size): for j in range(num_blobs): value = np.full(blob_size, (i * num_blobs) + j, np.float16 if use_float16 else np.float32) expected_output = np.concatenate((expected_output, value)) np.testing.assert_array_equal( workspace.FetchBlob("Gathered"), expected_output) # Run the net a few more times to check the operator # works not just the first time it's called for _tmp in range(4): workspace.RunNet(net.Name()) @given(comm_size=st.integers(min_value=2, max_value=8), blob_size=st.integers(min_value=int(1e3), max_value=int(1e6)), num_blobs=st.integers(min_value=1, max_value=4), device_option=st.sampled_from([hu.cpu_do]), use_float16=st.booleans()) @settings(max_examples=10, deadline=None) def test_allgather(self, comm_size, blob_size, num_blobs, device_option, use_float16): TestCase.test_counter += 1 if os.getenv('COMM_RANK') is not None: self.run_test_distributed( self._test_allgather, blob_size=blob_size, num_blobs=num_blobs, use_float16=use_float16, device_option=device_option) else: with TemporaryDirectory() as tmpdir: self.run_test_locally( self._test_allgather, comm_size=comm_size, blob_size=blob_size, num_blobs=num_blobs, device_option=device_option, tmpdir=tmpdir, use_float16=use_float16) @given(device_option=st.sampled_from([hu.cpu_do])) @settings(deadline=10000) def test_forked_cw(self, device_option): TestCase.test_counter += 1 if os.getenv('COMM_RANK') is not None: self.run_test_distributed( self._test_allreduce_multicw, device_option=device_option) else: # Note: this test exercises the path where we fork a common world. # We therefore don't need a comm size larger than 2. It used to be # run with comm_size=8, which causes flaky results in a stress run. # The flakiness was caused by too many listening sockets being # created by Gloo context initialization (8 processes times # 7 sockets times 20-way concurrency, plus TIME_WAIT). with TemporaryDirectory() as tmpdir: self.run_test_locally( self._test_allreduce_multicw, comm_size=2, device_option=device_option, tmpdir=tmpdir) def _test_barrier( self, comm_rank=None, comm_size=None, tmpdir=None, ): store_handler, common_world = self.create_common_world( comm_rank=comm_rank, comm_size=comm_size, tmpdir=tmpdir ) net = core.Net("barrier") net.Barrier( [common_world], [], engine=op_engine) workspace.CreateNet(net) workspace.RunNet(net.Name()) # Run the net a few more times to check the operator # works not just the first time it's called for _tmp in range(4): workspace.RunNet(net.Name()) @given(comm_size=st.integers(min_value=2, max_value=8), device_option=st.sampled_from([hu.cpu_do])) @settings(deadline=10000) def test_barrier(self, comm_size, device_option): TestCase.test_counter += 1 if os.getenv('COMM_RANK') is not None: self.run_test_distributed( self._test_barrier, device_option=device_option) else: with TemporaryDirectory() as tmpdir: self.run_test_locally( self._test_barrier, comm_size=comm_size, device_option=device_option, tmpdir=tmpdir) def _test_close_connection( self, comm_rank=None, comm_size=None, tmpdir=None, ): ''' One node calls close connection, others wait it on barrier. Test will check that all will exit eventually. ''' # Caffe's for closers only: # https://www.youtube.com/watch?v=QMFwFgG9NE8 closer = comm_rank == comm_size // 2, store_handler, common_world = self.create_common_world( comm_rank=comm_rank, comm_size=comm_size, tmpdir=tmpdir ) net = core.Net("barrier_or_close") if not closer: net.Barrier( [common_world], [], engine=op_engine) else: net.DestroyCommonWorld( [common_world], [common_world], engine=op_engine) # Sleep a bit to ensure others start the barrier import time time.sleep(0.1) workspace.CreateNet(net) workspace.RunNet(net.Name()) @given(comm_size=st.integers(min_value=2, max_value=8), device_option=st.sampled_from([hu.cpu_do])) @settings(deadline=10000) def test_close_connection(self, comm_size, device_option): import time start_time = time.time() TestCase.test_counter += 1 if os.getenv('COMM_RANK') is not None: self.run_test_distributed( self._test_close_connection, device_option=device_option) else: with TemporaryDirectory() as tmpdir: self.run_test_locally( self._test_close_connection, comm_size=comm_size, device_option=device_option, tmpdir=tmpdir) # Check that test finishes quickly because connections get closed. # This assert used to check that the end to end runtime was less # than 2 seconds, but this may not always be the case if there # is significant overhead in starting processes. Ideally, this # assert is replaced by one that doesn't depend on time but rather # checks the success/failure status of the barrier that is run. self.assertLess(time.time() - start_time, 20.0) def _test_io_error( self, comm_rank=None, comm_size=None, tmpdir=None, ): ''' Only one node will participate in allreduce, resulting in an IoError ''' store_handler, common_world = self.create_common_world( comm_rank=comm_rank, comm_size=comm_size, tmpdir=tmpdir) if comm_rank == 0: blob_size = 1000 num_blobs = 1 blobs = [] for i in range(num_blobs): blob = "blob_{}".format(i) value = np.full( blob_size, (comm_rank * num_blobs) + i, np.float32 ) workspace.FeedBlob(blob, value) blobs.append(blob) net = core.Net("allreduce") net.Allreduce( [common_world] + blobs, blobs, engine=op_engine) workspace.CreateNet(net) workspace.RunNet(net.Name()) @given(comm_size=st.integers(min_value=2, max_value=8), device_option=st.sampled_from([hu.cpu_do])) @settings(deadline=10000) def test_io_error(self, comm_size, device_option): TestCase.test_counter += 1 with self.assertRaises(IoError): if os.getenv('COMM_RANK') is not None: self.run_test_distributed( self._test_io_error, device_option=device_option) else: with TemporaryDirectory() as tmpdir: self.run_test_locally( self._test_io_error, comm_size=comm_size, device_option=device_option, tmpdir=tmpdir) if __name__ == "__main__": import unittest unittest.main()