Rebase of #1737 (libE service mode) onto develop + comms=local pickling fix

.gitignore

@@ -18,6 +18,8 @@ x.log
 *.csv
 *.DS_Store
 
+IBCDFO
+
 .vscode
 build/
 

docs/data_structures/libE_specs/libE_specs_general.rst

@@ -39,6 +39,20 @@ General
     List of workers that should run only generators. All other workers will run
     only simulator functions.
 
+**service_mode** [bool] = ``False``:
+    If ``True``, the manager tolerates idle workers + an alloc returning no work
+    (instead of asserting). Used when libEnsemble is driven by an external
+    producer (e.g. a queue-backed Generator fed by an MCP server) and may
+    legitimately have nothing to dispatch for periods of time. The manager
+    sleeps briefly between checks instead of panicking. Termination is the
+    caller's responsibility (e.g. via ``exit_criteria`` or external stop).
+
+**service_mode_idle_timeout** [float]:
+    In ``service_mode``, exit after this many seconds with no active workers
+    and no work dispatched. Default ``None`` means run forever waiting for
+    new submissions. Useful so an external producer doesn't leave libE running
+    after it goes silent.
+
 **cache_long_sims** [bool] = ``False``:
     Cache simulation results with runtimes >1s to disk. Subsequent runs with the same
     configuration, callable code and state, VOCS, and H0 will access this cache.

libensemble/gen_classes/queue_gen.py

@@ -0,0 +1,250 @@
+"""Queue-backed generator: work comes from an external producer (e.g. an MCP
+tool, an LLM agent, a REST API) via an input queue. Completed results go back
+out via an output queue.
+
+This lets libE run in 'service mode' — driven by an external loop instead of
+its own gen function.
+
+Usage:
+    from queue import Queue
+    from gest_api.vocs import VOCS
+    from libensemble.gen_classes.queue_gen import QueueGenerator
+
+    in_q = Queue()
+    out_q = Queue()
+    gen = QueueGenerator(VOCS(variables={"x": [-1, 1]}),
+                         input_queue=in_q, output_queue=out_q)
+    in_q.put({"x": 0.5})  # external producer feeds work
+"""
+
+import threading
+import time
+from queue import Empty, Queue
+from typing import Any, List, Optional
+
+from gest_api import Generator
+from gest_api.vocs import VOCS
+
+_SHUTDOWN = object()  # sentinel: external producer signals "no more work"
+
+
+class QueueGenerator(Generator):
+    """Generator that pulls work from an external Queue and pushes results back.
+
+    suggest(n):
+      - Blocks up to ``poll_timeout`` waiting for the FIRST item (so the libE
+        manager doesn't spin hot when nothing is queued).
+      - Then drains up to ``n - 1`` more items non-blockingly.
+      - Returns [] on timeout (libE will call again).
+      - Returns [] if a shutdown sentinel is seen.
+
+    ingest(results):
+      - Forwards each result dict to the output queue verbatim.
+    """
+
+    def __init__(
+        self,
+        vocs: VOCS,
+        *,
+        input_queue: Queue,
+        output_queue: Queue,
+        poll_timeout: float = 1.0,
+    ):
+        self.vocs = vocs
+        self.input_queue = input_queue
+        self.output_queue = output_queue
+        self.poll_timeout = poll_timeout
+        self._shutdown_seen = False
+        super().__init__(vocs)
+
+    def _validate_vocs(self, vocs: VOCS) -> None:
+        assert len(self.vocs.variable_names), "VOCS must contain variables."
+
+    def __getstate__(self):
+        # This generator shares plain queue.Queue objects directly with an
+        # external producer living in *this* process, so it must keep running
+        # in-process (QueueService enforces this via a gen_on_worker guard,
+        # but QueueGenerator can also be wired up directly, bypassing that
+        # check). Under libE_specs["comms"] = "local" (multiprocessing),
+        # gen_specs - including this generator instance - is still pickled to
+        # spawn every *sim*-only worker process, even though they never touch
+        # it. Drop the live Queue objects (which hold unpicklable
+        # threading.Lock state) so that pickling doesn't crash; those sim
+        # workers never use them. If gen_on_worker=True sends this generator
+        # to a worker that's actually expected to call suggest()/ingest(),
+        # the guards below raise a clear error instead of an AttributeError.
+        state = self.__dict__.copy()
+        state["input_queue"] = None
+        state["output_queue"] = None
+        return state
+
+    def _check_queues(self) -> None:
+        if self.input_queue is None or self.output_queue is None:
+            raise RuntimeError(
+                "QueueGenerator has no queues in this process. This generator must run "
+                "on the manager (do not set libE_specs['gen_on_worker'] = True), since its "
+                "queues are plain in-memory objects shared with the external producer and "
+                "are not available in a separate worker process."
+            )
+
+    def suggest(self, num_points: Optional[int]) -> List[dict]:
+        self._check_queues()
+        if self._shutdown_seen:
+            return []
+        n = num_points or 1
+        items: List[dict] = []
+        try:
+            first = self.input_queue.get(timeout=self.poll_timeout)
+        except Empty:
+            return []
+        if first is _SHUTDOWN:
+            self._shutdown_seen = True
+            return []
+        items.append(first)
+        for _ in range(n - 1):
+            try:
+                nxt = self.input_queue.get_nowait()
+            except Empty:
+                break
+            if nxt is _SHUTDOWN:
+                self._shutdown_seen = True
+                break
+            items.append(nxt)
+        return items
+
+    def ingest(self, results: List[dict]) -> None:
+        self._check_queues()
+        for r in results:
+            self.output_queue.put(r)
+
+    def finalize(self, results: List[dict] = None, *args: Any, **kwargs: Any):
+        if results:
+            self.ingest(results)
+        return None
+
+    @staticmethod
+    def shutdown_sentinel():
+        """External producer puts this on the input queue to signal end."""
+        return _SHUTDOWN
+
+
+class QueueService:
+    """Service-mode wrapper: spawns libE in a thread with a QueueGenerator and
+    exposes submit/get_completed/shutdown to an external producer.
+
+    Hides the queue/thread/generator plumbing every producer would otherwise
+    repeat. The producer just creates a service and feeds it work:
+
+        service = QueueService(vocs, sim_specs, libE_specs, exit_criteria)
+        service.start()
+        service.submit({"x": 1.0})
+        for r in service.get_completed():
+            ...
+        service.shutdown()
+    """
+
+    def __init__(
+        self,
+        vocs: VOCS,
+        sim_specs,
+        libE_specs,
+        exit_criteria,
+        *,
+        persis_in: Optional[List[str]] = None,
+        batch_size: int = 0,
+        poll_timeout: float = 1.0,
+    ):
+        from libensemble import Ensemble
+        from libensemble.specs import GenSpecs
+
+        # QueueGenerator shares plain queue.Queue objects directly with the
+        # external producer's process - it must run on the manager (the
+        # default), not a spawned worker, which wouldn't have access to the
+        # same queue objects (and whose copy would be stripped of them, see
+        # QueueGenerator.__getstate__/_check_queues).
+        assert not getattr(
+            libE_specs, "gen_on_worker", False
+        ), "QueueGenerator must run on the manager; do not set libE_specs['gen_on_worker'] = True."
+
+        # final_gen_send must be True for QueueGenerator — otherwise the last
+        # batch of completed sims is never ingested and never reaches the
+        # output_queue, so consumers silently miss results.
+        libE_specs.final_gen_send = True
+
+        self.input_queue: Queue = Queue()
+        self.output_queue: Queue = Queue()
+        self._gen = QueueGenerator(
+            vocs,
+            input_queue=self.input_queue,
+            output_queue=self.output_queue,
+            poll_timeout=poll_timeout,
+        )
+        gen_specs = GenSpecs(
+            generator=self._gen,
+            vocs=vocs,
+            persis_in=persis_in or [],
+            batch_size=batch_size,
+        )
+        self._ensemble = Ensemble(sim_specs, gen_specs, exit_criteria, libE_specs)
+        self._thread: Optional[threading.Thread] = None
+
+    def start(self) -> None:
+        """Spawn the libE thread."""
+        self._thread = threading.Thread(target=self._ensemble.run, daemon=True)
+        self._thread.start()
+
+    def submit(self, item: dict) -> None:
+        """Submit one work item."""
+        self.input_queue.put(item)
+
+    def get_completed(self) -> List[dict]:
+        """Drain all completed results (non-blocking)."""
+        out = []
+        while True:
+            try:
+                out.append(self.output_queue.get_nowait())
+            except Empty:
+                break
+        return out
+
+    def collect_results(self, n: int, timeout: float = 60.0) -> List[dict]:
+        """Block-drain until ``n`` results collected or ``timeout`` elapses.
+        Returns whatever was collected (may be < n on timeout)."""
+        results: List[dict] = []
+        deadline = time.time() + timeout
+        while len(results) < n and time.time() < deadline:
+            try:
+                results.append(self.output_queue.get(timeout=1))
+            except Empty:
+                pass
+        return results
+
+    def stream_results(self, n: Optional[int] = None, timeout: float = 60.0):
+        """Yield results as they arrive. Stops after ``n`` yielded or
+        ``timeout`` seconds elapse with no new result. ``n=None`` streams
+        indefinitely until timeout."""
+        deadline = time.time() + timeout
+        yielded = 0
+        while (n is None or yielded < n) and time.time() < deadline:
+            try:
+                r = self.output_queue.get(timeout=1)
+            except Empty:
+                continue
+            deadline = time.time() + timeout  # reset on activity
+            yielded += 1
+            yield r
+
+    def shutdown(self, wait: bool = False) -> None:
+        """Signal libE to stop accepting new work and drain in-flight.
+        If ``wait=True``, block until the libE thread exits."""
+        self.input_queue.put(_SHUTDOWN)
+        if wait:
+            self.join()
+
+    def join(self, timeout: Optional[float] = None) -> None:
+        """Wait for libE thread to exit."""
+        if self._thread is not None:
+            self._thread.join(timeout=timeout)
+
+    def is_alive(self) -> bool:
+        return self._thread is not None and self._thread.is_alive()

libensemble/manager.py

@@ -222,6 +222,7 @@ def __init__(
         self.gen_specs = gen_specs
         self.exit_criteria = exit_criteria
         self.elapsed = lambda: timer.elapsed
+        self._service_idle_start: float | None = None  # service_mode idle tracker
         self.wcomms = wcomms
         self.WorkerExc = False
         self.persis_pending: list[int] = []
@@ -1105,9 +1106,27 @@ def run(self, persis_info: dict) -> tuple[dict, int, int]:
                     self._check_work_order(Work[w], w)
                     self._send_work_order(Work[w], w)
                     self._update_state_on_alloc(Work[w], w)
-                assert self.term_test() or any(
-                    self.W["active"] != 0
-                ), "alloc_f did not return any work, although all workers are idle."
+                if not (self.term_test() or any(self.W["active"] != 0)):
+                    if self.libE_specs.get("service_mode"):
+                        # Service mode: external producer may have nothing
+                        # to dispatch right now. Throttle and re-poll instead
+                        # of asserting (which would treat this as deadlock).
+                        idle_timeout = self.libE_specs.get("service_mode_idle_timeout")
+                        if idle_timeout is not None:
+                            if self._service_idle_start is None:
+                                self._service_idle_start = time.time()
+                            elif time.time() - self._service_idle_start > idle_timeout:
+                                logger.info(
+                                    f"service_mode idle for {idle_timeout}s with no work, exiting"
+                                )
+                                break
+                        time.sleep(0.1)
+                        continue
+                    raise AssertionError(
+                        "alloc_f did not return any work, although all workers are idle."
+                    )
+                # work was dispatched OR workers active OR terminating — reset idle clock
+                self._service_idle_start = None
         except WorkerException as e:
             report_worker_exc(e)
             raise LoggedException(e.args[0], e.args[1]) from None

libensemble/specs.py

@@ -476,6 +476,24 @@ class LibeSpecs(BaseModel):
     By default, the generator runs on the manager process as a thread (Worker 0).
     """
 
+    service_mode: bool = False
+    """
+    If True, the manager tolerates idle workers + an alloc returning no work
+    (instead of asserting). Used when libEnsemble is driven by an external
+    producer (e.g. a queue-backed Generator fed by an MCP server) and may
+    legitimately have nothing to dispatch for periods of time. The manager
+    sleeps briefly between checks instead of panicking. Termination is the
+    caller's responsibility (e.g. via ``exit_criteria`` or external stop).
+    """
+
+    service_mode_idle_timeout: float | None = None
+    """
+    In ``service_mode``, exit after this many seconds with no active workers
+    and no work dispatched. ``None`` (default) means run forever waiting for
+    new submissions. Useful so an MCP server doesn't leave libE running after
+    the agent goes silent.
+    """
+
     mpi_comm: object | None = None
     """ libEnsemble MPI communicator. Default: ``MPI.COMM_WORLD``"""