various zrw / gen-worker indexing fixes. defensive programming around worker indexes and zrws. mypy types. vibe-coded new test cases for test_mpi_runners_zrw_subnode_uneven

Author: jlnav

libensemble/libE.py

@@ -385,7 +385,7 @@ def libE_mpi(sim_specs, gen_specs, exit_criteria, persis_info, alloc_specs, libE
         # Run manager or worker code, depending
         if is_manager:
             if resources is not None:
-                resources.set_resource_manager(nworkers)
+                resources.set_resource_manager(nworkers + 1)
             return libE_mpi_manager(
                 mpi_comm, sim_specs, gen_specs, exit_criteria, persis_info, alloc_specs, libE_specs, H0
             )

libensemble/resources/resources.py

@@ -67,8 +67,8 @@ def __init__(self, libE_specs: dict, platform_info: dict = {}, top_level_dir: st
         """Initiate a new resources object"""
         self.top_level_dir = top_level_dir or os.getcwd()
         self.glob_resources = GlobalResources(libE_specs=libE_specs, platform_info=platform_info, top_level_dir="")
-        self.resource_manager = None  # For Manager
-        self.worker_resources = None  # For Workers
+        self.resource_manager: ResourceManager | None = None  # For Manager
+        self.worker_resources: WorkerResources | None = None  # For Workers
 
     def set_worker_resources(self, num_workers: int, workerid: int) -> None:
         """Initiate the worker resources component of resources"""
@@ -167,8 +167,6 @@ def __init__(self, libE_specs: dict, platform_info: dict = {}, top_level_dir: st
         self.top_level_dir = top_level_dir
         self.dedicated_mode = libE_specs.get("dedicated_mode", False)
         self.zero_resource_workers = libE_specs.get("zero_resource_workers", [])
-        if 0 not in self.zero_resource_workers:
-            self.zero_resource_workers.append(0)
         self.num_resource_sets = libE_specs.get("num_resource_sets", None)
         self.enforce_worker_core_bounds = libE_specs.get("enforce_worker_core_bounds", False)
         self.gpus_per_group = libE_specs.get("gpus_per_group")

libensemble/resources/worker_resources.py

@@ -105,7 +105,7 @@ def free_rsets(self, worker=None):
     def get_index_list(num_workers: int, num_rsets: int, zero_resource_list: list[int | Any]) -> list[int | None]:
         """Map WorkerID to index into a nodelist"""
         index = 0
-        index_list = []
+        index_list: list[int | None] = []
         for i in range(0, num_workers):
             if i in zero_resource_list:
                 index_list.append(None)
@@ -116,6 +116,11 @@ def get_index_list(num_workers: int, num_rsets: int, zero_resource_list: list[in
                 else:
                     index_list.append(index)
                 index += 1
+
+        for i in zero_resource_list:
+            if i >= num_workers:
+                logger.warning(f"Worker index {i} from zero_resource_workers is out of range (0-{num_workers - 1})")
+
         return index_list
 
 
@@ -364,7 +369,7 @@ def get_local_nodelist(
         local_nodelist = list(OrderedDict.fromkeys(team_list))  # Maintain order of nodes
         logger.debug(f"Worker's local_nodelist is {local_nodelist}")
 
-        slots = {}
+        slots: dict[str, list[int]] = {}
         for node in local_nodelist:
             slots[node] = []
 

libensemble/tests/functionality_tests/test_mpi_gpu_settings.py

@@ -66,7 +66,7 @@
 # Main block is necessary only when using local comms with spawn start method (default on macOS and Windows).
 if __name__ == "__main__":
     nworkers, is_manager, libE_specs, _ = parse_args()
-    libE_specs["num_resource_sets"] = nworkers - 1  # Persistent gen does not need resources
+    libE_specs["num_resource_sets"] = nworkers  # Persistent gen does not need resources
     libE_specs["use_workflow_dir"] = True  # Only a place for Open MPI machinefiles
 
     if libE_specs["comms"] == "tcp":
@@ -88,8 +88,8 @@
         "persis_in": ["f", "x", "sim_id"],
         "out": [("priority", float), ("resource_sets", int), ("x", float, n)],
         "user": {
-            "initial_batch_size": nworkers - 1,
-            "max_resource_sets": nworkers - 1,  # Any sim created can req. 1 worker up to all.
+            "initial_batch_size": nworkers,
+            "max_resource_sets": nworkers,  # Any sim created can req. 1 worker up to all.
             "lb": np.array([-3, -2]),
             "ub": np.array([3, 2]),
         },

libensemble/tests/functionality_tests/test_mpi_gpu_settings_env.py

@@ -43,7 +43,7 @@
 # Main block is necessary only when using local comms with spawn start method (default on macOS and Windows).
 if __name__ == "__main__":
     nworkers, is_manager, libE_specs, _ = parse_args()
-    libE_specs["num_resource_sets"] = nworkers - 1  # Persistent gen does not need resources
+    libE_specs["num_resource_sets"] = nworkers  # Persistent gen does not need resources
     libE_specs["use_workflow_dir"] = True  # Only a place for Open MPI machinefiles
 
     # Optional for organization of output scripts
@@ -70,8 +70,8 @@
         "persis_in": ["f", "x", "sim_id"],
         "out": [("priority", float), ("resource_sets", int), ("x", float, n)],
         "user": {
-            "initial_batch_size": nworkers - 1,
-            "max_resource_sets": nworkers - 1,  # Any sim created can req. 1 worker up to all.
+            "initial_batch_size": nworkers,
+            "max_resource_sets": nworkers,  # Any sim created can req. 1 worker up to all.
             "lb": np.array([-3, -2]),
             "ub": np.array([3, 2]),
         },

libensemble/tests/functionality_tests/test_mpi_runners_zrw_subnode_uneven.py

@@ -139,26 +139,39 @@
         exp_srun.append(srun_p1 + str(nodename) + srun_p2 + str(ntasks) + srun_p3 + str(ntasks) + srun_p4)
 
     test_list = test_list_base
-    exp_list = exp_srun
+    exp_list_dynamic = exp_srun.copy()
+    if nworkers == 5:
+        # Iteration 0 (Dynamic): Workers 1, 2 -> node-2; 3, 4, 5 -> node-1
+        n1 = srun_p1 + "node-1" + srun_p2 + "5" + srun_p3 + "5" + srun_p4
+        n2 = srun_p1 + "node-2" + srun_p2 + "8" + srun_p3 + "8" + srun_p4
+        exp_list_dynamic = [n2, n2, n1, n1, n1]
+        # Iteration 1 (Static): Worker 1 is gen. Workers 2, 3 -> node-1; 4, 5 -> node-2
+        exp_list_static = [n1, n1, n2, n2]
+
     sim_specs["user"] = {
         "tests": test_list,
-        "expect": exp_list,
         "persis_gens": n_gens,
     }
 
     iterations = 2
     for prob_id in range(iterations):
         if prob_id == 0:
             # Uses dynamic scheduler - will find node 2 slots first (as fewer)
-            libE_specs["num_resource_sets"] = nworkers - 1  # Any worker can be the gen
-            sim_specs["user"]["offset_for_scheduler"] = True  # Changes expected values
+            libE_specs["gen_on_worker"] = False
+            libE_specs["num_resource_sets"] = nworkers
+            sim_specs["user"]["expect"] = exp_list_dynamic
+            sim_specs["user"]["offset_for_scheduler"] = False
+            sim_specs["user"]["persis_gens"] = 0
             persis_info = add_unique_random_streams({}, nworkers + 1)
 
         else:
             # Uses static scheduler - will find node 1 slots first
+            libE_specs["gen_on_worker"] = True
+            sim_specs["user"]["expect"] = exp_list_static
             del libE_specs["num_resource_sets"]
             libE_specs["zero_resource_workers"] = [1]  # Gen must be worker 1
             sim_specs["user"]["offset_for_scheduler"] = False
+            sim_specs["user"]["persis_gens"] = 1
             persis_info = add_unique_random_streams({}, nworkers + 1)
 
         # Perform the run

libensemble/tests/functionality_tests/test_persistent_uniform_gen_decides_stop.py

@@ -36,6 +36,8 @@
         n = 2
         init_batch_size = nworkers - ngens
 
+        libE_specs["gen_on_worker"] = True
+
         if ngens >= nworkers:
             sys.exit("The number of generators must be less than the number of workers -- aborting...")