From 5aec234acadec614aa5efc2bf70da61e518e662a Mon Sep 17 00:00:00 2001
From: Hariharan Devarajan <mani.hariharan@gmail.com>
Date: Tue, 13 Aug 2024 14:07:28 -0700
Subject: [PATCH] only intialize and finalize on DLIOMPI (#214)

---
 .gitignore                      |  1 +
 dlio_benchmark/main.py          | 21 ++++++++++-----------
 dlio_benchmark/utils/utility.py |  6 +++---
 3 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/.gitignore b/.gitignore
index ede9fd70..ae91d41f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -157,3 +157,4 @@ dmypy.json
 .DS_Store
 /.idea/
 /venv-quartz/
+venv-openmpi
\ No newline at end of file
diff --git a/dlio_benchmark/main.py b/dlio_benchmark/main.py
index 2246edbf..84388ccd 100644
--- a/dlio_benchmark/main.py
+++ b/dlio_benchmark/main.py
@@ -16,8 +16,6 @@
 """
 import os
 import math
-from mpi4py import MPI
-comm = MPI.COMM_WORLD
 import logging
 from time import time, sleep
 import json
@@ -50,8 +48,7 @@
 from dlio_benchmark.utils.utility import Profile, PerfTrace
 
 dlp = Profile(MODULE_DLIO_BENCHMARK)
-# To make sure the output folder is the same in all the nodes. We have to do this. 
-comm.Barrier()
+# To make sure the output folder is the same in all the nodes. We have to do this.
 import hydra
 
 class DLIOBenchmark(object):
@@ -383,19 +380,21 @@ def finalize(self):
         self.comm.barrier()
         self.args.finalize_dftracer(self.dftracer)
 
-
 @hydra.main(version_base=None, config_path="configs", config_name="config")
-def main(cfg: DictConfig) -> None:
-
-    """
-    The main method to start the benchmark runtime.
-    """
-    DLIOMPI.get_instance().initialize()
+def run_benchmark(cfg: DictConfig):    
     benchmark = DLIOBenchmark(cfg['workload'])
     os.environ["DARSHAN_DISABLE"] = "1"
     benchmark.initialize()
     benchmark.run()
     benchmark.finalize()
+
+
+def main() -> None:
+    """
+    The main method to start the benchmark runtime.
+    """
+    DLIOMPI.get_instance().initialize()
+    run_benchmark()
     DLIOMPI.get_instance().finalize()
 
 if __name__ == '__main__':
diff --git a/dlio_benchmark/utils/utility.py b/dlio_benchmark/utils/utility.py
index 8872f2ec..6bbde3e8 100644
--- a/dlio_benchmark/utils/utility.py
+++ b/dlio_benchmark/utils/utility.py
@@ -65,10 +65,8 @@ def iter(self, a):
 
 # MPI cannot be initialized automatically, or read_thread spawn/forkserver
 # child processes will abort trying to open a non-existant PMI_fd file.
-import mpi4py
-mpi4py.rc.initialize = False
-from mpi4py import MPI
 
+import mpi4py
 p = psutil.Process()
 
 
@@ -112,6 +110,7 @@ def classname(cls):
         return cls.__qualname__
 
     def initialize(self):
+        from mpi4py import MPI
         if self.mpi_state == MPIState.UNINITIALIZED:
             # MPI may have already been initialized by dlio_benchmark_test.py
             if not MPI.Is_initialized():
@@ -181,6 +180,7 @@ def nnodes(self):
         else:
             return self.mpi_size//self.mpi_ppn
     def finalize(self):
+        from mpi4py import MPI
         if self.mpi_state == MPIState.MPI_INITIALIZED and MPI.Is_initialized():
             MPI.Finalize()