migrate many alloc_f["user"] options to gen_specs. remove corresponding alloc_f import from many tests. mypy and formatting adjustments. first approach at describing options in specs.py

Author: jlnav

libensemble/alloc_funcs/give_sim_work_first.py

@@ -14,7 +14,7 @@ def give_sim_work_first(
     alloc_specs: dict,
     persis_info: dict,
     libE_info: dict,
-) -> tuple[dict]:
+) -> tuple[dict, dict]:
     """
     Decide what should be given to workers. This allocation function gives any
     available simulation work first, and only when all simulations are
@@ -41,7 +41,6 @@ def give_sim_work_first(
     """
 
     user = alloc_specs.get("user", {})
-
     if "cancel_sims_time" in user:
         # Cancel simulations that are taking too long
         rows = np.where(np.logical_and.reduce((H["sim_started"], ~H["sim_ended"], ~H["cancel_requested"])))[0]
@@ -53,8 +52,11 @@ def give_sim_work_first(
     if libE_info["sim_max_given"] or not libE_info["any_idle_workers"]:
         return {}, persis_info
 
-    # Initialize alloc_specs["user"] as user.
-    batch_give = user.get("give_all_with_same_priority", False)
+    # Initialize options - check gen_specs first
+    batch_give = gen_specs.get("give_all_with_same_priority", user.get("give_all_with_same_priority", False))
+    num_active_gens = gen_specs.get("num_active_gens", user.get("num_active_gens", 1))
+    batch_mode = gen_specs.get("batch_mode", user.get("batch_mode", False))
+
     gen_in = gen_specs.get("in", [])
 
     manage_resources = libE_info["use_resource_sets"]
@@ -77,11 +79,11 @@ def give_sim_work_first(
     else:
         for wid in support.avail_worker_ids(gen_workers=True):
             # Allow at most num_active_gens active generator instances
-            if gen_count >= user.get("num_active_gens", gen_count + 1):
+            if gen_count >= num_active_gens:
                 break
 
             # Do not start gen instances in batch mode if workers still working
-            if user.get("batch_mode") and not support.all_sim_ended(H):
+            if batch_mode and not support.all_sim_ended(H):
                 break
 
             # Give gen work

libensemble/alloc_funcs/start_only_persistent.py

@@ -25,7 +25,7 @@ def only_persistent_gens(W, H, sim_specs, gen_specs, alloc_specs, persis_info, l
 
     To be provided in calling script: E.g., ``alloc_specs["user"]["async_return"] = True``
 
-    init_sample_size: int, optional
+    initial_batch_size: int, optional
         Initial sample size - always return in batch. Default: 0
 
     num_active_gens: int, optional
@@ -59,16 +59,18 @@ def only_persistent_gens(W, H, sim_specs, gen_specs, alloc_specs, persis_info, l
     user = alloc_specs.get("user", {})
     manage_resources = libE_info["use_resource_sets"]
 
-    active_recv_gen = user.get("active_recv_gen", False)  # Persistent gen can handle irregular communications
-    init_sample_size = user.get("init_sample_size", 0)  # Always batch return until this many evals complete
-    batch_give = user.get("give_all_with_same_priority", False)
+    active_recv_gen = gen_specs.get("active_recv_gen", user.get("active_recv_gen", False))
+    initial_batch_size = gen_specs.get("initial_batch_size", user.get("initial_batch_size", 0))
+    batch_give = gen_specs.get("give_all_with_same_priority", user.get("give_all_with_same_priority", False))
 
     support = AllocSupport(W, manage_resources, persis_info, libE_info)
     gen_count = support.count_persis_gens()
     Work = {}
 
     # Asynchronous return to generator
-    async_return = user.get("async_return", False) and sum(H["sim_ended"]) >= init_sample_size
+    async_return = (
+        gen_specs.get("async_return", user.get("async_return", False)) and sum(H["sim_ended"]) >= initial_batch_size
+    )
 
     if gen_count < persis_info.get("num_gens_started", 0):
         # When a persistent worker is done, trigger a shutdown (returning exit condition of 1)
@@ -94,7 +96,7 @@ def only_persistent_gens(W, H, sim_specs, gen_specs, alloc_specs, persis_info, l
     # Now the give_sim_work_first part
     points_to_evaluate = ~H["sim_started"] & ~H["cancel_requested"]
     avail_workers = support.avail_worker_ids(persistent=False, zero_resource_workers=False, gen_workers=False)
-    if user.get("alt_type"):
+    if gen_specs.get("alt_type", user.get("alt_type")):
         avail_workers = list(
             set(support.avail_worker_ids(persistent=False, zero_resource_workers=False))
             | set(support.avail_worker_ids(persistent=EVAL_SIM_TAG, zero_resource_workers=False))
@@ -106,7 +108,7 @@ def only_persistent_gens(W, H, sim_specs, gen_specs, alloc_specs, persis_info, l
         sim_ids_to_send = support.points_by_priority(H, points_avail=points_to_evaluate, batch=batch_give)
 
         try:
-            if user.get("alt_type"):
+            if gen_specs.get("alt_type", user.get("alt_type")):
                 Work[wid] = support.sim_work(
                     wid, H, sim_specs["in"], sim_ids_to_send, persis_info.get(wid), persistent=True
                 )
@@ -122,7 +124,7 @@ def only_persistent_gens(W, H, sim_specs, gen_specs, alloc_specs, persis_info, l
         avail_workers = support.avail_worker_ids(persistent=False, zero_resource_workers=True, gen_workers=True)
 
         for wid in avail_workers:
-            if gen_count < user.get("num_active_gens", 1):
+            if gen_count < gen_specs.get("num_active_gens", user.get("num_active_gens", 1)):
                 # Finally, start a persistent generator as there is nothing else to do.
                 try:
                     Work[wid] = support.gen_work(

libensemble/gen_funcs/persistent_sampling.py

@@ -16,9 +16,9 @@
 ]
 
 
-def _get_user_params(user_specs):
+def _get_user_params(user_specs, gen_specs):
     """Extract user params"""
-    b = user_specs["initial_batch_size"]
+    b = gen_specs.get("initial_batch_size") or user_specs.get("initial_batch_size") or gen_specs.get("init_sample_size")
     ub = user_specs["ub"]
     lb = user_specs["lb"]
     n = len(lb)  # dimension
@@ -44,7 +44,7 @@ def persistent_uniform(_, persis_info, gen_specs, libE_info):
         `test_persistent_uniform_sampling_async.py <https://github.com/Libensemble/libensemble/blob/develop/libensemble/tests/functionality_tests/test_persistent_uniform_sampling_async.py>`_
     """  # noqa
 
-    b, n, lb, ub = _get_user_params(gen_specs["user"])
+    b, n, lb, ub = _get_user_params(gen_specs.get("user", {}), gen_specs)
     ps = PersistentSupport(libE_info, EVAL_GEN_TAG)
 
     # Send batches until manager sends stop tag
@@ -73,7 +73,7 @@ def persistent_uniform_final_update(_, persis_info, gen_specs, libE_info):
         `test_persistent_uniform_sampling_running_mean.py <https://github.com/Libensemble/libensemble/blob/develop/libensemble/tests/functionality_tests/test_persistent_uniform_sampling_running_mean.py>`_
     """  # noqa
 
-    b, n, lb, ub = _get_user_params(gen_specs["user"])
+    b, n, lb, ub = _get_user_params(gen_specs.get("user", {}), gen_specs)
     ps = PersistentSupport(libE_info, EVAL_GEN_TAG)
 
     def generate_corners(x, y):
@@ -145,7 +145,7 @@ def persistent_request_shutdown(_, persis_info, gen_specs, libE_info):
     .. seealso::
         `test_persistent_uniform_gen_decides_stop.py <https://github.com/Libensemble/libensemble/blob/develop/libensemble/tests/functionality_tests/test_persistent_uniform_gen_decides_stop.py>`_
     """  # noqa
-    b, n, lb, ub = _get_user_params(gen_specs["user"])
+    b, n, lb, ub = _get_user_params(gen_specs.get("user", {}), gen_specs)
     shutdown_limit = gen_specs["user"]["shutdown_limit"]
     f_count = 0
     ps = PersistentSupport(libE_info, EVAL_GEN_TAG)
@@ -173,7 +173,7 @@ def uniform_nonblocking(_, persis_info, gen_specs, libE_info):
     .. seealso::
         `test_persistent_uniform_sampling.py <https://github.com/Libensemble/libensemble/blob/develop/libensemble/tests/functionality_tests/test_persistent_uniform_sampling.py>`_
     """  # noqa
-    b, n, lb, ub = _get_user_params(gen_specs["user"])
+    b, n, lb, ub = _get_user_params(gen_specs.get("user", {}), gen_specs)
     ps = PersistentSupport(libE_info, EVAL_GEN_TAG)
 
     # Send batches until manager sends stop tag
@@ -223,7 +223,11 @@ def batched_history_matching(_, persis_info, gen_specs, libE_info):
     lb = gen_specs["user"]["lb"]
 
     n = len(lb)
-    b = gen_specs["user"]["initial_batch_size"]
+    b = (
+        gen_specs.get("initial_batch_size")
+        or gen_specs["user"].get("initial_batch_size")
+        or gen_specs.get("init_sample_size")
+    )
     q = gen_specs["user"]["num_best_vals"]
     ps = PersistentSupport(libE_info, EVAL_GEN_TAG)
 
@@ -250,7 +254,11 @@ def persistent_uniform_with_cancellations(_, persis_info, gen_specs, libE_info):
     ub = gen_specs["user"]["ub"]
     lb = gen_specs["user"]["lb"]
     n = len(lb)
-    b = gen_specs["user"]["initial_batch_size"]
+    b = (
+        gen_specs.get("initial_batch_size")
+        or gen_specs["user"].get("initial_batch_size")
+        or gen_specs.get("init_sample_size")
+    )
 
     # Start cancelling points from half initial batch onward
     cancel_from = b // 2  # Should get at least this many points back

libensemble/specs.py

@@ -200,11 +200,14 @@ class GenSpecs(BaseModel):
 
     initial_batch_size: int = 0
     """
-    Number of initial points to request that the generator create. If zero, falls back to ``batch_size``.
-    If both options are zero, defaults to the number of workers.
+    Initial sample size.
+    For standardized generators, this is the number of initial points to request that the
+    generator create. If zero, falls back to ``batch_size``.
+    For persistent generators, this is the number of points evaluated before switching
+    from batch return to asynchronous return (if ``async_return`` is True).
 
-    Note: Certain generators included with libEnsemble decide
-    batch sizes via ``gen_specs["user"]`` or other methods.
+    Note: Certain generators included with libEnsemble decide batch sizes via
+    ``gen_specs["user"]`` or other methods.
     """
 
     batch_size: int = 0
@@ -233,6 +236,53 @@ class GenSpecs(BaseModel):
     they will be automatically derived from VOCS.
     """
 
+    # Only used if using the only_persistent_gens allocation function (default)
+    num_active_gens: int = 1
+    """
+    Maximum number of persistent generators to start. Default: 1.
+    Only used if using the ``only_persistent_gens`` allocation function (the default).
+    """
+
+    async_return: bool = False
+    """
+    Return results to gen as they come in (after sample). Default: False (batch return).
+    Only used if using the ``only_persistent_gens`` allocation function (the default).
+    """
+
+    active_recv_gen: bool = False
+    """
+    Create gen in active receive mode. If True, the manager does not need to wait
+    for a return from the generator before sending further returned points.
+    Default: False. Only used if using the ``only_persistent_gens`` allocation function (the default).
+    """
+
+    give_all_with_same_priority: bool = False
+    """
+    If True, then all points with the same priority value are given as a batch to the sim.
+    Default: False. Only used if using the ``only_persistent_gens`` allocation function (the default).
+    """
+
+    alt_type: bool = False
+    """
+    If True, then the specialized allocator behavior for some persistent gens is used.
+    Only used if using the ``only_persistent_gens`` allocation function (the default).
+    """
+
+    batch_mode: bool = False
+    """
+    If True, then the generator will not be started if there are still simulations
+    running. Only used if using the ``give_sim_work_first`` allocation function.
+    """
+
+    @model_validator(mode="before")
+    def set_gen_specs_fields_from_user(cls, values):
+        """Set fields from user dict for backward compatibility."""
+        # init_sample_size is now initial_batch_size
+        if "init_sample_size" in values and "initial_batch_size" not in values:
+            values["initial_batch_size"] = values.pop("init_sample_size")
+
+        return values
+
     @model_validator(mode="after")
     def set_fields_from_vocs(self):
         """Set persis_in and outputs from VOCS if vocs is provided and fields are not set."""
@@ -290,10 +340,14 @@ class AllocSpecs(BaseModel):
         of ``give_sim_work_first``.
     """
 
-    user: dict | None = {"num_active_gens": 1}
+    user: dict | None = {}
     """
     A user-data dictionary to place bounds, constants, settings, or other parameters
     for customizing the allocation function.
+
+    .. note::
+        As of libEnsemble v2.0, generator-specific allocation options (e.g., ``async_return``,
+        ``num_active_gens``) have been moved to :class:`GenSpecs<libensemble.specs.GenSpecs>`.
     """
 
     outputs: list[tuple] = Field([], alias="out")

libensemble/tests/functionality_tests/test_GPU_gen_resources.py

@@ -32,7 +32,6 @@
 
 import numpy as np
 
-from libensemble.alloc_funcs.start_only_persistent import only_persistent_gens as alloc_f
 from libensemble.executors.mpi_executor import MPIExecutor
 from libensemble.gen_funcs.persistent_sampling_var_resources import uniform_sample_with_sim_gen_resources as gen_f
 
@@ -79,23 +78,17 @@
         "gen_f": gen_f,
         "persis_in": ["f", "x", "sim_id"],
         "out": [("num_procs", int), ("num_gpus", int), ("x", float, n)],
+        "initial_batch_size": nworkers - 1,
+        "give_all_with_same_priority": False,
+        "async_return": False,
         "user": {
-            "initial_batch_size": nworkers - 1,
             "max_procs": nworkers - 1,  # Any sim created can req. 1 worker up to all.
             "lb": np.array([-3, -2]),
             "ub": np.array([3, 2]),
             "dry_run": dry_run,
         },
     }
 
-    alloc_specs = {
-        "alloc_f": alloc_f,
-        "user": {
-            "give_all_with_same_priority": False,
-            "async_return": False,  # False batch returns
-        },
-    }
-
     exit_criteria = {"sim_max": 20}
     libE_specs["resource_info"] = {"cores_on_node": (nworkers * 2, nworkers * 4), "gpus_on_node": nworkers}
 
@@ -126,8 +119,6 @@
                 gen_specs["user"]["max_procs"] = max(nworkers - 2, 1)
 
             # Perform the run
-            H, persis_info, flag = libE(
-                sim_specs, gen_specs, exit_criteria, persis_info, libE_specs=libE_specs, alloc_specs=alloc_specs
-            )
+            H, persis_info, flag = libE(sim_specs, gen_specs, exit_criteria, persis_info, libE_specs=libE_specs)
 
 # All asserts are in gen and sim funcs

libensemble/tests/functionality_tests/test_asktell_sampling.py

@@ -19,7 +19,6 @@
 import libensemble.sim_funcs.six_hump_camel as six_hump_camel
 
 # Import libEnsemble items for this test
-from libensemble.alloc_funcs.start_only_persistent import only_persistent_gens as alloc_f
 from libensemble.gen_classes.sampling import UniformSample
 from libensemble.libE import libE
 from libensemble.sim_funcs.executor_hworld import executor_hworld as sim_f_exec
@@ -59,7 +58,6 @@ def sim_f(In):
 
     vocs = VOCS(variables=variables, objectives=objectives)
 
-    alloc_specs = {"alloc_f": alloc_f}
     exit_criteria = {"gen_max": 201}
     persis_info = add_unique_random_streams({}, nworkers + 1, seed=1234)
 
@@ -89,9 +87,7 @@ def sim_f(In):
             }
 
         gen_specs["generator"] = generator
-        H, persis_info, flag = libE(
-            sim_specs, gen_specs, exit_criteria, persis_info, alloc_specs, libE_specs=libE_specs
-        )
+        H, persis_info, flag = libE(sim_specs, gen_specs, exit_criteria, persis_info, libE_specs=libE_specs)
 
         if is_manager:
             print(H[["sim_id", "x", "f"]][:10])

libensemble/tests/functionality_tests/test_asktell_sampling_external_gen.py

@@ -20,12 +20,12 @@
 
 from libensemble import Ensemble
 
-# Import libEnsemble items for this test
-from libensemble.alloc_funcs.start_only_persistent import only_persistent_gens as alloc_f
-
 # from libensemble.gen_classes.external.sampling import UniformSampleArray
 from libensemble.gen_classes.external.sampling import UniformSample
-from libensemble.specs import AllocSpecs, ExitCriteria, GenSpecs, LibeSpecs, SimSpecs
+from libensemble.specs import ExitCriteria, GenSpecs, LibeSpecs, SimSpecs
+
+# Import libEnsemble items for this test
+
 
 # from gest_api.vocs import ContinuousVariable
 
@@ -75,15 +75,13 @@ def sim_f_scalar(In):
             vocs=vocs,
         )
 
-        alloc_specs = AllocSpecs(alloc_f=alloc_f)
         exit_criteria = ExitCriteria(gen_max=201)
 
         ensemble = Ensemble(
             parse_args=True,
             sim_specs=sim_specs,
             gen_specs=gen_specs,
             exit_criteria=exit_criteria,
-            alloc_specs=alloc_specs,
             libE_specs=libE_specs,
         )