| import os |
| import sys |
| import subprocess |
| import cloudpickle |
| import base64 |
| import pytest |
| from functools import wraps |
|
|
| try: |
| from mpi4py import MPI |
| except ImportError: |
| MPI = None |
|
|
| def with_mpi(nproc=2, timeout=30, skip_if_no_mpi=True): |
| def outer_thunk(fn): |
| @wraps(fn) |
| def thunk(*args, **kwargs): |
| serialized_fn = base64.b64encode(cloudpickle.dumps(lambda: fn(*args, **kwargs))) |
| subprocess.check_call([ |
| 'mpiexec','-n', str(nproc), |
| sys.executable, |
| '-m', 'baselines.common.tests.test_with_mpi', |
| serialized_fn |
| ], env=os.environ, timeout=timeout) |
|
|
| if skip_if_no_mpi: |
| return pytest.mark.skipif(MPI is None, reason="MPI not present")(thunk) |
| else: |
| return thunk |
|
|
| return outer_thunk |
|
|
|
|
| if __name__ == '__main__': |
| if len(sys.argv) > 1: |
| fn = cloudpickle.loads(base64.b64decode(sys.argv[1])) |
| assert callable(fn) |
| fn() |
|
|
