Remove MPI

Author: andreicuceu

lya_2pt/scripts/run_mpi.py

@@ -2,7 +2,6 @@
 
 import argparse
 import time
-from mpi4py import MPI
 from configparser import ConfigParser
 
 from lya_2pt.interface import Interface
@@ -16,42 +15,39 @@ def main():
     parser.add_argument('-i', '--config', type=str, default=None,
                         help=('Path to config file'))
 
+    parser.add_argument('-r', '--rank', type=int, default=0, help=('Rank of the job'))
+    parser.add_argument('-s', '--size', type=int, default=1, help=('Number of jobs'))
+
     args = parser.parse_args()
 
+    if args.rank >= args.size:
+        raise MPIError(f"Rank {args.rank} is greater than the number of MPI processes {args.size}")
+
     config = ConfigParser()
     config.read(args.config)
 
-    # Initilize MPI objects
-    mpi_comm = MPI.COMM_WORLD
-    mpi_rank = mpi_comm.Get_rank()
-    mpi_size = mpi_comm.Get_size()
-
     lya2pt = Interface(config)
 
-    if len(lya2pt.files) < mpi_size:
+    if len(lya2pt.files) < args.size:
         raise MPIError(f"Less files than MPI processes. "
                        f"Found {len(lya2pt.files)} healpix files and running "
-                       f"{mpi_size} MPI processes. This is wasteful. "
+                       f"{args.size} MPI processes. This is wasteful. "
                        "Please lower the numper of MPI processes.")
 
-    if mpi_rank == 0:
-        total_t1 = time.time()
-        print('Reading tracers...', flush=True)
-    lya2pt.read_tracers(mpi_rank=mpi_rank)
+    total_t1 = time.time()
+
+    print(f'Rank {args.rank}: Reading tracers...', flush=True)
+    lya2pt.read_tracers(mpi_rank=args.rank)
 
-    if mpi_rank == 0:
-        print('Starting computation...', flush=True)
-    lya2pt.run(mpi_size=mpi_size, mpi_rank=mpi_rank)
+    print(f'Rank {args.rank}: Starting computation...', flush=True)
+    lya2pt.run(mpi_size=args.size, mpi_rank=args.rank)
 
-    if mpi_rank == 0:
-        print('Writing results...', flush=True)
-    lya2pt.write_results(mpi_rank=mpi_rank)
+    print(f'Rank {args.rank}: Writing results...', flush=True)
+    lya2pt.write_results(mpi_rank=args.rank)
 
-    mpi_comm.Barrier()
-    if mpi_rank == 0:
-        total_t2 = time.time()
-        print(f'Total time: {(total_t2-total_t1):.3f} sec', flush=True)
-        print('Done', flush=True)
+    total_t2 = time.time()
+    print(f'Rank {args.rank}: Total time: {(total_t2-total_t1):.3f} sec', flush=True)
+    print(f'Rank {args.rank}: Done', flush=True)
 
 
 if __name__ == '__main__':