Parallel performance profiler

src/Drivers/hiopbbpy/BODriverEX.py

@@ -72,11 +72,12 @@ def cons_jac_vec(x):
 
 
 if __name__ == "__main__":
+  do_profiling = True
   for prob_type in prob_type_l:
     print()
     # ----- evaluator
-    obj_evaluator = MPIEvaluator()
-    opt_evaluator = MPIEvaluator(function_mode=False)
+    obj_evaluator = MPIEvaluator(profiling=do_profiling, task_name="MPI_OBJ_EVAL")
+    opt_evaluator = MPIEvaluator(function_mode=False, profiling=do_profiling, task_name="MPI_OPT_EVAL")
     if prob_type == "LpNorm":
       problem = LpNormProblem(nx, xlimits)
     else:

src/Drivers/hiopbbpy/EvaluationManagerCI.py

@@ -9,59 +9,89 @@
 import argparse
 import time
 import sys
+import os
+import socket
+import threading
 from hiopbbpy.utils import EvaluationManager, is_running_with_mpi
 from concurrent.futures import ThreadPoolExecutor
 
-def _fn_for_test(x, sleep_time=0.1):
-  if is_running_with_mpi():
-    from mpi4py import MPI
-    comm = MPI.COMM_WORLD
-    rank = comm.Get_rank()
-    print(f"Rank {rank}: Processing {x}")
-  time.sleep(sleep_time)  # Simulate some work
-  return x * x
+def _fn_for_test(x, sleep_time=0.1, slow_first=False, driver_rank=0):
+    hostname = socket.gethostname()
+    pid = os.getpid()
+
+    if slow_first and x == 0:
+        actual_sleep = 3 * sleep_time
+    else:
+        actual_sleep = sleep_time
 
+    print(
+        f"rank={driver_rank} pid={pid} host={hostname}: processing x={x}",
+        flush=True,
+    )
 
+    time.sleep(actual_sleep)
+    return x * x
 
 if __name__ == "__main__":
   # Arguments for command line
   parser = argparse.ArgumentParser(
     description="Execute n function calls with t duration.",
     epilog="To properly run the example with mpi4py, use: env MPI4PY_FUTURES_MAX_WORKERS=<N> mpiexec -n 1 python evaluation_manager.py",
   )
-  parser.add_argument("-n", type=int, default=100, help="Number of tasks to execute")
+  parser.add_argument("-n", type=int, default=20, help="Number of tasks to execute")
   parser.add_argument(
       "-t", "--sleep_time", type=float, default=1, help="Sleep time for each task"
   )
+  parser.add_argument(
+      "-p",
+      "--profile",
+      action="store_true",
+      help="Enable profiling",
+  )
+  parser.add_argument(
+    "-s",
+    "--slow_first",
+    action="store_true",
+    help="Make the first task slower (3x sleep time)",
+  )
   args = parser.parse_args()
 
+  # Choose executor type
   if is_running_with_mpi():
+    from mpi4py import MPI
+    driver_rank = MPI.COMM_WORLD.Get_rank()
     executor_type = "mpi"
   else:
+    driver_rank = 0
     executor_type = "cpu"
 
   # Set up logging
   logging.basicConfig(level=logging.INFO)
 
   # Create manager
   cpu_executor = ThreadPoolExecutor()
-  manager = EvaluationManager(cpu_executor=cpu_executor)
+  manager = EvaluationManager(
+      cpu_executor=cpu_executor,
+      profiling=args.profile,
+      task_name="CI_TASK"
+  )
 
-  # Submit tasks to the manager
+  # Submit tasks
   t0 = time.perf_counter()
   manager.submit_tasks(
       _fn_for_test,
       [i for i in range(args.n)],
       execute_at=executor_type,
       sleep_time=args.sleep_time,
+      slow_first=args.slow_first,
+      driver_rank=driver_rank,
   )
 
   # Do some other work while tasks are running
-  print("Doing other work", end="", flush=True)
   for i in range(5):
-    print(".", end="", flush=True)
+    print("Doing other work (Master)", flush=True)
     time.sleep(args.sleep_time)
-  print(" Done.")
+  print("Doing other work (Master) --- Done.")
 
   # Wait for all tasks to complete
   print("Waiting for tasks to complete...")
@@ -76,5 +106,4 @@ def _fn_for_test(x, sleep_time=0.1):
 
   # Clean up
   del manager
-  retval = 0
-  sys.exit(retval)
+  sys.exit(0)

src/hiopbbpy/surrogate_modeling/krg.py

@@ -16,7 +16,7 @@ def __init__(self, theta, xlimits, ndim, corr="pow_exp", noise0=None, random_sta
     super().__init__(ndim, xlimits)
     if random_state is None:
       random_state = 42
-    design_space = DesignSpace(xlimits, random_state=random_state)
+    design_space = DesignSpace(xlimits, seed=random_state)
     if noise0 is None:
       self.surrogatesmt = KRG(design_space=design_space,
                        print_global=False,