add multicpu test

.github/workflows/CI.yml

@@ -23,8 +23,8 @@ jobs:
         os:
           - ubuntu-latest
         arch:
-          - x64
           - x86
+          - x64
 
     steps:
       - uses: actions/checkout@v2
@@ -43,6 +43,11 @@ jobs:
 
       - name: Run tests
         run: |
+          export RANK=0
+          export LOCAL_RANK=0
+          export WORLD_SIZE=1
+          export MASTER_ADDR=localhost
+          export MASTER_PORT=12345
           export PYTHONPATH=MCintegration
           pytest --cov --cov-report=xml --ignore=examples
 

examples/mc_multicpu_test.py → MCintegration/mc_multicpu_test.py

@@ -2,8 +2,7 @@
 import torch.distributed as dist
 import torch.multiprocessing as mp
 import os
-import MCintegration
-from MCintegration import MonteCarlo, MarkovChainMonteCarlo
+from integrators import MonteCarlo, MarkovChainMonteCarlo
 
 
 def init_process(rank, world_size, fn, backend="gloo"):
@@ -53,7 +52,6 @@ def two_integrands(x, f):
         # Only rank 0 prints the result
         print("MarkovChainMonteCarlo Result:", mcmc_result)
 
-
-if __name__ == "__main__":
+def test_mcmc():
     world_size = 8  # Number of processes to launch
     mp.spawn(init_process, args=(world_size, run_mcmc), nprocs=world_size, join=True)

examples/mc_multigpu_test.py

@@ -10,8 +10,8 @@
 )
 
 
-backend = "nccl"
-# backend = "gloo"
+# backend = "nccl"
+backend = "gloo"
 # set_seed(42)
 setup(backend=backend)
 device = get_device()