Update mc_multicpu_test.py

Author: fancaiyu

MCintegration/mc_multicpu_test.py

@@ -2,18 +2,10 @@
 import torch.distributed as dist
 import torch.multiprocessing as mp
 import os
-import pytest
 from integrators import MonteCarlo, MarkovChainMonteCarlo
 
-@pytest.fixture
-def rank():
-    return 0
 
-@pytest.fixture
-def world_size():
-    return 8
-
-def test_init_process(rank, world_size, fn, backend="gloo"):
+def init_process(rank, world_size, fn, backend="gloo"):
     # Set MASTER_ADDR and MASTER_PORT appropriately
     # Assuming environment variables are set by the cluster's job scheduler
     master_addr = os.getenv("MASTER_ADDR", "localhost")
@@ -26,7 +18,7 @@ def test_init_process(rank, world_size, fn, backend="gloo"):
     fn(rank, world_size)
 
 
-def test_run_mcmc(rank, world_size):
+def run_mcmc(rank, world_size):
     # Instantiate the MarkovChainMonteCarlo class
     bounds = [(-1, 1), (-1, 1)]
     n_eval = 8000000
@@ -60,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(test_init_process, args=(world_size, test_run_mcmc), nprocs=world_size, join=True)
+    mp.spawn(init_process, args=(world_size, run_mcmc), nprocs=world_size, join=True)