vibe coded approach lets see how this works

Author: jlnav

libensemble/libE.py

@@ -155,7 +155,7 @@ def libE(
     exit_criteria: ExitCriteria,
     persis_info: dict = {},
     alloc_specs: AllocSpecs = AllocSpecs(),
-    libE_specs: LibeSpecs = {},
+    libE_specs: LibeSpecs | dict = {},
     H0=None,
 ) -> (np.ndarray, dict, int):
     """
@@ -242,11 +242,16 @@ def libE(
     ]
     exit_criteria = specs_dump(ensemble.exit_criteria, by_alias=True, exclude_none=True)
 
-    # Restore objects that don't survive serialization via model_dump
-    if hasattr(ensemble.gen_specs, "generator") and ensemble.gen_specs.generator is not None:
-        gen_specs["generator"] = ensemble.gen_specs.generator
-    if hasattr(ensemble.gen_specs, "vocs") and ensemble.gen_specs.vocs is not None:
-        gen_specs["vocs"] = ensemble.gen_specs.vocs
+    if hasattr(ensemble.sim_specs, "simulator") and ensemble.sim_specs.simulator is not None:
+        sim_specs["simulator"] = ensemble.sim_specs.simulator
+    if hasattr(ensemble.sim_specs, "vocs") and ensemble.sim_specs.vocs is not None:
+        sim_specs["vocs"] = ensemble.sim_specs.vocs
+
+    if ensemble.gen_specs is not None:
+        if hasattr(ensemble.gen_specs, "generator") and ensemble.gen_specs.generator is not None:
+            gen_specs["generator"] = ensemble.gen_specs.generator
+        if hasattr(ensemble.gen_specs, "vocs") and ensemble.gen_specs.vocs is not None:
+            gen_specs["vocs"] = ensemble.gen_specs.vocs
 
     # Extract platform info from settings or environment
     platform_info = get_platform(libE_specs)
@@ -358,7 +363,7 @@ def libE_mpi(sim_specs, gen_specs, exit_criteria, persis_info, alloc_specs, libE
         logger.manager_warning("*WARNING* libEnsemble detected a NULL communicator")
         return [], persis_info, 3  # Process not in mpi_comm
 
-    assert libE_specs["mpi_comm"].Get_size() > 1, "Manager only - must be at least one worker (2 MPI tasks)"
+    assert libE_specs["mpi_comm"].Get_size() >= 1, "Manager only - must be at least one MPI task"
 
     with DupComm(libE_specs["mpi_comm"]) as mpi_comm:
         is_manager = mpi_comm.Get_rank() == 0
@@ -368,7 +373,6 @@ def libE_mpi(sim_specs, gen_specs, exit_criteria, persis_info, alloc_specs, libE
             local_host = socket.gethostname()
             libE_nodes = list(set(mpi_comm.allgather(local_host)))
             resources.add_comm_info(libE_nodes=libE_nodes)
-            nworkers = mpi_comm.Get_size() - 1
 
         exctr = Executor.executor
         if exctr is not None:
@@ -379,7 +383,8 @@ 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)
+                n_resource_workers = mpi_comm.Get_size()
+                resources.set_resource_manager(n_resource_workers)
             return libE_mpi_manager(
                 mpi_comm, sim_specs, gen_specs, exit_criteria, persis_info, alloc_specs, libE_specs, H0
             )
@@ -497,7 +502,8 @@ def libE_local(sim_specs, gen_specs, exit_criteria, persis_info, alloc_specs, li
 
     # Set manager resources after the forkpoint.
     if resources is not None:
-        resources.set_resource_manager(libE_specs["nworkers"])
+        n_resource_workers = libE_specs["nworkers"] + (not libE_specs.get("gen_on_worker", False))
+        resources.set_resource_manager(n_resource_workers)
 
     if not libE_specs["disable_log_files"]:
         exit_logger = manager_logging_config(specs=libE_specs)

libensemble/manager.py

@@ -232,7 +232,7 @@ def __init__(
             (1, "stop_val", self.term_test_stop_val),
         ]
 
-        gen_on_manager = self.libE_specs.get("gen_on_manager", False)
+        gen_on_manager = not self.libE_specs.get("gen_on_worker", False)
 
         self.W = np.zeros(len(self.wcomms) + gen_on_manager, dtype=Manager.worker_dtype)
         if gen_on_manager:
@@ -662,7 +662,7 @@ def _get_alloc_libE_info(self) -> dict:
             "use_resource_sets": self.use_resource_sets,
             "gen_num_procs": self.gen_num_procs,
             "gen_num_gpus": self.gen_num_gpus,
-            "gen_on_manager": self.libE_specs.get("gen_on_manager", False),
+            "gen_on_worker": self.libE_specs.get("gen_on_worker", False),
         }
 
     def _alloc_work(self, H: npt.NDArray, persis_info: dict) -> dict:

libensemble/resources/resources.py

@@ -166,7 +166,7 @@ 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", [])
+        self.zero_resource_workers = libE_specs.get("zero_resource_workers", [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,8 +105,8 @@ 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 = []
-        for i in range(1, num_workers + 1):
+        index_list: list[int | None] = []
+        for i in range(0, num_workers):
             if i in zero_resource_list:
                 index_list.append(None)
             else:
@@ -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/specs.py

@@ -344,9 +344,9 @@ class LibeSpecs(BaseModel):
     nworkers: int | None = 0
     """ Number of worker processes in ``"local"``, ``"threads"``, or ``"tcp"``."""
 
-    gen_on_manager: bool | None = False
-    """ Instructs Manager process to run generator functions.
-    This generator function can access/modify user objects by reference.
+    gen_on_worker: bool = False
+    """ Instructs libEnsemble to run generator functions on a worker rank.
+    By default, the generator runs on the manager process as a thread (Worker 0).
     """
 
     mpi_comm: object | None = None
@@ -635,10 +635,11 @@ class LibeSpecs(BaseModel):
     libEnsemble processes (manager and workers) are running.
     """
 
-    zero_resource_workers: list[int] | None = []
+    zero_resource_workers: list[int] | None = [0]
     """
     list of workers that require no resources. For when a fixed mapping of workers
     to resources is required. Otherwise, use ``num_resource_sets``.
+    By default, Worker 0 (manager thread) is a zero-resource worker.
     For use with supported allocation functions.
     """
 

libensemble/tests/functionality_tests/test_mpi_gpu_settings.py

@@ -66,7 +66,15 @@
 # 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
+
+    # If gen_on_worker is False (default), then all nworkers are available for sims.
+    # Worker 0 is the generator (and it is a zero_resource_worker by default).
+    if not libE_specs.get("gen_on_worker", False):
+        nsim_workers = nworkers
+    else:
+        nsim_workers = nworkers - 1
+
+    libE_specs["num_resource_sets"] = nsim_workers
     libE_specs["use_workflow_dir"] = True  # Only a place for Open MPI machinefiles
 
     if libE_specs["comms"] == "tcp":
@@ -88,8 +96,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": nsim_workers,
+            "max_resource_sets": nsim_workers,  # 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,15 @@
 # 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
+
+    # If gen_on_worker is False (default), then all nworkers are available for sims.
+    # Worker 0 is the generator (and it is a zero_resource_worker by default).
+    if not libE_specs.get("gen_on_worker", False):
+        nsim_workers = nworkers
+    else:
+        nsim_workers = nworkers - 1
+
+    libE_specs["num_resource_sets"] = nsim_workers
     libE_specs["use_workflow_dir"] = True  # Only a place for Open MPI machinefiles
 
     # Optional for organization of output scripts
@@ -70,8 +78,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": nsim_workers,
+            "max_resource_sets": nsim_workers,  # 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_persistent_uniform_gen_decides_stop.py

@@ -33,11 +33,18 @@
     nworkers, is_manager, libE_specs, _ = parse_args()
 
     for ngens in range(1, 3):
+        # If gen_on_worker is False (default), the first gen is on the manager (Worker 0).
+        # Subsequent gens (if ngens > 1) move to worker ranks.
+        if not libE_specs.get("gen_on_worker", False):
+            nsim_workers = nworkers - (ngens - 1)
+        else:
+            nsim_workers = nworkers - ngens
+
         n = 2
-        init_batch_size = nworkers - ngens
+        init_batch_size = nsim_workers
 
-        if ngens >= nworkers:
-            sys.exit("The number of generators must be less than the number of workers -- aborting...")
+        if nsim_workers <= 0:
+            sys.exit("The number of generators must be less than the available workers -- aborting...")
 
         sim_specs = {
             "sim_f": sim_f,

libensemble/tests/unit_tests/test_manager_main.py

@@ -6,7 +6,7 @@
 import libensemble.manager as man
 import libensemble.tests.unit_tests.setup as setup
 
-libE_specs = {"comms": "local"}
+libE_specs = {"comms": "local", "gen_on_worker": True}
 
 
 def test_term_test_1():

libensemble/tests/unit_tests/test_resources.py

@@ -694,23 +694,23 @@ def test_map_workerid_to_index():
 
     zero_resource_list = []
     index_list = ResourceManager.get_index_list(num_workers, num_rsets, zero_resource_list)
-    for workerID in range(1, num_workers + 1):
-        index = index_list[workerID - 1]
-        assert index == workerID - 1, "index incorrect. Received: " + str(index)
+    for workerID in range(0, num_workers):
+        index = index_list[workerID]
+        assert index == workerID, "index incorrect. Received: " + str(index)
 
-    zero_resource_list = [1]
+    zero_resource_list = [0]
     index_list = ResourceManager.get_index_list(num_workers, num_rsets, zero_resource_list)
-    for workerID in range(2, num_workers + 1):
-        index = index_list[workerID - 1]
-        assert index == workerID - 2, "index incorrect. Received: " + str(index)
+    for workerID in range(1, num_workers):
+        index = index_list[workerID]
+        assert index == workerID - 1, "index incorrect. Received: " + str(index)
 
-    zero_resource_list = [1, 2]
+    zero_resource_list = [0, 1]
     index_list = ResourceManager.get_index_list(num_workers, num_rsets, zero_resource_list)
-    for workerID in range(3, num_workers + 1):
-        index = index_list[workerID - 1]
-        assert index == workerID - 3, "index incorrect. Received: " + str(index)
+    for workerID in range(2, num_workers):
+        index = index_list[workerID]
+        assert index == workerID - 2, "index incorrect. Received: " + str(index)
 
-    zero_resource_list = [1, 3]
+    zero_resource_list = [0, 2]
     index_list = ResourceManager.get_index_list(num_workers, num_rsets, zero_resource_list)
 
     workerID = 2