def SetupMPI(): try: from mpi4py import MPI global _has_mpi, COMM, RANK, SIZE _has_mpi = core.IsOperatorWithEngine("CreateCommonWorld", "MPI") COMM = MPI.COMM_WORLD RANK = COMM.Get_rank() SIZE = COMM.Get_size() except ImportError: _has_mpi = False
def test_is_operator_with_engine(self): self.assertTrue(core.IsOperatorWithEngine('Relu', 'DEFAULT')) self.assertFalse(core.IsOperatorWithEngine('Relu', 'NOEXIST'))
import hypothesis.strategies as st import numpy as np import unittest from caffe2.python import core, workspace, dyndep import caffe2.python.hypothesis_test_util as hu if workspace.has_gpu_support: dyndep.InitOpsLibrary("@/caffe2/caffe2/mpi:mpi_ops_gpu") else: dyndep.InitOpsLibrary("@/caffe2/caffe2/mpi:mpi_ops") try: from mpi4py import MPI _has_mpi = core.IsOperatorWithEngine("CreateCommonWorld", "MPI") COMM = MPI.COMM_WORLD RANK = COMM.Get_rank() SIZE = COMM.Get_size() except ImportError: _has_mpi = False @unittest.skipIf(not _has_mpi, "MPI is not available. Skipping.") class TestMPI(hu.HypothesisTestCase): @given(X=hu.tensor(), root=st.integers(min_value=0, max_value=SIZE - 1), device_option=st.sampled_from(hu.device_options), **hu.gcs) def test_broadcast(self, X, root, device_option, gc, dc): # Use mpi4py's broadcast to make sure that all nodes inherit the