diff --git a/.gitignore b/.gitignore index c6bd3c0dd..b585a50da 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,8 @@ x.log *.csv *.DS_Store +IBCDFO + .vscode build/ diff --git a/docs/data_structures/libE_specs/libE_specs_general.rst b/docs/data_structures/libE_specs/libE_specs_general.rst index 567451e46..15f0676f2 100644 --- a/docs/data_structures/libE_specs/libE_specs_general.rst +++ b/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. diff --git a/libensemble/gen_classes/queue_gen.py b/libensemble/gen_classes/queue_gen.py new file mode 100644 index 000000000..538753c6b --- /dev/null +++ b/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() diff --git a/libensemble/manager.py b/libensemble/manager.py index 2f0e6df25..dc42975fe 100644 --- a/libensemble/manager.py +++ b/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 diff --git a/libensemble/specs.py b/libensemble/specs.py index 821fcb7c9..d87030829 100644 --- a/libensemble/specs.py +++ b/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``""" diff --git a/libensemble/tests/regression_tests/test_queue_gen.py b/libensemble/tests/regression_tests/test_queue_gen.py new file mode 100644 index 000000000..8df0dfe35 --- /dev/null +++ b/libensemble/tests/regression_tests/test_queue_gen.py @@ -0,0 +1,66 @@ +"""Test for QueueGenerator / QueueService + +Uses QueueService to spin up libE in a thread with a trivial doubler sim, +submits N work items, drains results, shuts down, joins. + + +Run with: + python test_queue_gen.py +""" + +# Do not change these lines - they are parsed by run-tests.sh +# TESTSUITE_COMMS: local +# TESTSUITE_NPROCS: 4 + +import math +import time + +import numpy as np +from gest_api.vocs import VOCS + +from libensemble.gen_classes.queue_gen import QueueService +from libensemble.specs import ExitCriteria, LibeSpecs, SimSpecs + +NWORKERS = 4 +N_SUBMITS = 10 + + +def doubler_sim(InputArray, _, sim_specs): + """Trivial sim: returns 2*x. Sleeps a bit to mimic real work.""" + time.sleep(0.5) + out = np.zeros(1, dtype=sim_specs["out"]) + out["y"] = 2.0 * InputArray["x"][0] + return out + + +def main(): + vocs = VOCS(variables={"x": [-100.0, 100.0]}, objectives={"y": "MINIMIZE"}) + sim_specs = SimSpecs(sim_f=doubler_sim, inputs=["x"], outputs=[("y", float)]) + libE_specs = LibeSpecs(nworkers=NWORKERS, service_mode=True) + exit_criteria = ExitCriteria(sim_max=N_SUBMITS) + + service = QueueService(vocs, sim_specs, libE_specs, exit_criteria) + service.start() + print(f"libE thread started ({NWORKERS} workers)") + + # Submit work (1, 2, 3....), then signal shutdown. + for i in range(N_SUBMITS): + service.submit({"x": float(i)}) + service.shutdown() + print(f"submitted {N_SUBMITS} items + shutdown sentinel") + + # Block-drain until all results collected (or timeout) + results = service.collect_results(N_SUBMITS, timeout=60) + print(f"\ncollected {len(results)}/{N_SUBMITS}") + + # Verify y == 2*x + ok = all(math.isclose(r["y"], 2 * r["x"]) for r in results) + print("PASS" if ok else "FAIL") + + # Wait for libE to wind down + service.join(timeout=10) + print("done") + + +if __name__ == "__main__": + main() diff --git a/libensemble/tests/regression_tests/test_queue_gen_idle_timeout.py b/libensemble/tests/regression_tests/test_queue_gen_idle_timeout.py new file mode 100644 index 000000000..583609acc --- /dev/null +++ b/libensemble/tests/regression_tests/test_queue_gen_idle_timeout.py @@ -0,0 +1,59 @@ +"""Test for service_mode_idle_timeout. + +Producer stays silent after a single submit. With service_mode_idle_timeout +set, libE should exit by itself after the timeout instead of hanging. + +Run with: + python test_queue_gen_idle_timeout.py +""" + +# Do not change these lines - they are parsed by run-tests.sh +# TESTSUITE_COMMS: local +# TESTSUITE_NPROCS: 2 + +import time + +import numpy as np +from gest_api.vocs import VOCS + +from libensemble.gen_classes.queue_gen import QueueService +from libensemble.specs import ExitCriteria, LibeSpecs, SimSpecs + +NWORKERS = 2 +IDLE_TIMEOUT = 3.0 # seconds + + +def doubler_sim(InputArray, _, sim_specs): + out = np.zeros(1, dtype=sim_specs["out"]) + out["y"] = 2.0 * InputArray["x"][0] + return out + + +def main(): + vocs = VOCS(variables={"x": [-100.0, 100.0]}, objectives={"y": "MINIMIZE"}) + sim_specs = SimSpecs(sim_f=doubler_sim, inputs=["x"], outputs=[("y", float)]) + libE_specs = LibeSpecs( + nworkers=NWORKERS, + service_mode=True, + service_mode_idle_timeout=IDLE_TIMEOUT, + ) + exit_criteria = ExitCriteria(sim_max=1000) # large so idle_timeout fires first + + service = QueueService(vocs, sim_specs, libE_specs, exit_criteria) + service.start() + print(f"libE thread started ({NWORKERS} workers), idle_timeout={IDLE_TIMEOUT}s") + + service.submit({"x": 1.0}) + print("submitted 1 item, now going silent") + + t0 = time.time() + service.join(timeout=IDLE_TIMEOUT + 10) + elapsed = time.time() - t0 + print(f"libE thread exited after {elapsed:.2f}s (expected ~{IDLE_TIMEOUT}s)") + + assert not service.is_alive(), "libE didn't exit after idle_timeout" + print("PASS") + + +if __name__ == "__main__": + main() diff --git a/libensemble/tests/regression_tests/test_queue_gen_streaming.py b/libensemble/tests/regression_tests/test_queue_gen_streaming.py new file mode 100644 index 000000000..b7fd8e6cb --- /dev/null +++ b/libensemble/tests/regression_tests/test_queue_gen_streaming.py @@ -0,0 +1,73 @@ +"""Test for QueueGenerator / QueueService — streaming with replacement. + +Submit NWORKERS items so each worker has work, then as each result comes +back submit a replacement until TOTAL items have been submitted. Drain +the remaining results and shutdown. + +Run with: + python test_queue_gen_streaming.py +""" + +# Do not change these lines - they are parsed by run-tests.sh +# TESTSUITE_COMMS: local +# TESTSUITE_NPROCS: 4 + +import math +import time + +import numpy as np +from gest_api.vocs import VOCS + +from libensemble.gen_classes.queue_gen import QueueService +from libensemble.specs import ExitCriteria, LibeSpecs, SimSpecs + +NWORKERS = 4 +TOTAL = 16 # total items to submit (>> NWORKERS so the rolling window matters) + + +def doubler_sim(InputArray, _, sim_specs): + """Trivial sim: returns 2*x. Sleeps a bit to mimic real work.""" + time.sleep(0.5) + out = np.zeros(1, dtype=sim_specs["out"]) + out["y"] = 2.0 * InputArray["x"][0] + return out + + +def main(): + vocs = VOCS(variables={"x": [-100.0, 100.0]}, objectives={"y": "MINIMIZE"}) + sim_specs = SimSpecs(sim_f=doubler_sim, inputs=["x"], outputs=[("y", float)]) + libE_specs = LibeSpecs(nworkers=NWORKERS, service_mode=True, + service_mode_idle_timeout=30) + exit_criteria = ExitCriteria(sim_max=TOTAL) + + service = QueueService(vocs, sim_specs, libE_specs, exit_criteria) + service.start() + print(f"libE thread started ({NWORKERS} workers)") + + # Submit one item per worker + for i in range(NWORKERS): + service.submit({"x": float(i)}) + submitted = NWORKERS + print(f"submitted initial {NWORKERS} items") + + # As each result arrives, submit a replacement until TOTAL submitted + results = [] + for r in service.stream_results(TOTAL, timeout=60): + results.append(r) + print(f" got ({len(results)}/{TOTAL}): {r}") + if submitted < TOTAL: + service.submit({"x": float(submitted)}) + submitted += 1 + + service.shutdown() + print(f"\nsubmitted {submitted}, collected {len(results)}/{TOTAL}") + + ok = all(math.isclose(r["y"], 2 * r["x"]) for r in results) + print("PASS" if ok else "FAIL") + + service.join(timeout=10) + print("done") + + +if __name__ == "__main__": + main() diff --git a/libensemble/tests/unit_tests/test_manager_main.py b/libensemble/tests/unit_tests/test_manager_main.py index 4c8a14be3..45b871938 100644 --- a/libensemble/tests/unit_tests/test_manager_main.py +++ b/libensemble/tests/unit_tests/test_manager_main.py @@ -1,11 +1,14 @@ import time +from types import SimpleNamespace +from unittest.mock import Mock import numpy as np import numpy.lib.recfunctions +import pytest import libensemble.manager as man import libensemble.tests.unit_tests.setup as setup -from libensemble.message_numbers import PERSIS_STOP +from libensemble.message_numbers import EVAL_SIM_TAG, PERSIS_STOP libE_specs = {"comms": "local", "gen_on_worker": True} @@ -64,10 +67,6 @@ def test_term_test_3(): def test_send_work_order_persis_stop_empty_rows(): - # Regression test: a PERSIS_STOP work order with zero H_rows (e.g. final_gen_send - # with nothing left to forward to a persistent gen) must still send the meta-message. - # Otherwise, a persistent worker blocked in PersistentSupport.recv() waiting on this - # message hangs forever, since nothing would ever be sent to it. hist, sim_specs, gen_specs, exit_criteria, al = setup.hist_setup2() final_libE_specs = dict(libE_specs, final_gen_send=True) fake_wcomm = FakeComm() @@ -81,12 +80,56 @@ def test_send_work_order_persis_stop_empty_rows(): } mgr._send_work_order(work, 1) - assert len(fake_wcomm.sent) == 1, "Expected the PERSIS_STOP meta-message to be sent even with zero H_rows" + assert len(fake_wcomm.sent) == 1 tag, sent_work = fake_wcomm.sent[0] assert tag == PERSIS_STOP assert len(sent_work["libE_info"]["H_rows"]) == 0 +def test_send_work_order_with_partial_cache_hit(): + mgr = SimpleNamespace( + hist=SimpleNamespace( + H=np.zeros(2, dtype=[("x", float)]), + cache_set=True, + ), + use_cache=True, + resources=None, + cache_hit=False, + wcomms=[Mock()], + ) + mgr._update_state_from_cache = Mock(return_value=[0]) + work = { + "tag": EVAL_SIM_TAG, + "H_fields": ["x"], + "libE_info": {"H_rows": np.array([0, 1])}, + } + + man.Manager._send_work_order(mgr, work, 0) + + np.testing.assert_array_equal(work["libE_info"]["H_rows"], [0, 1]) + sent_work = mgr.wcomms[0].send.call_args_list[0].args[1] + assert sent_work["libE_info"]["H_rows"] == [1] + + +def test_run_rejects_idle_allocator_without_service_mode(): + mgr = SimpleNamespace( + libE_specs={}, + exit_criteria={}, + hist=SimpleNamespace(trim_H=Mock(return_value=np.zeros(0)), save_cache=Mock()), + W=np.zeros(1, dtype=man.Manager.worker_dtype), + wcomms=[], + _service_idle_start=None, + term_test=Mock(side_effect=[False, False]), + _kill_cancelled_sims=Mock(), + _receive_from_workers_or_cache=Mock(return_value={}), + _alloc_work=Mock(return_value=({}, {}, 0)), + _final_receive_and_kill=Mock(return_value=({}, 0, 0)), + ) + + with pytest.raises(man.LoggedException, match="alloc_f did not return any work"): + man.Manager.run(mgr, {}) + + if __name__ == "__main__": test_term_test_1() test_term_test_2() diff --git a/libensemble/tests/unit_tests/test_queue_gen.py b/libensemble/tests/unit_tests/test_queue_gen.py new file mode 100644 index 000000000..38647648c --- /dev/null +++ b/libensemble/tests/unit_tests/test_queue_gen.py @@ -0,0 +1,155 @@ +import time +from queue import Queue +from unittest.mock import MagicMock + +import pytest +from gest_api.vocs import VOCS + +from libensemble.gen_classes.queue_gen import QueueGenerator, QueueService + + +def _make_gen(): + vocs = VOCS(variables={"x": [-1, 1]}) + return QueueGenerator(vocs, input_queue=Queue(), output_queue=Queue()) + + +def test_getstate_strips_queues(): + # Simulates what happens when gen_specs (including the generator) is + # pickled to spawn a sim-only worker under comms="local". + gen = _make_gen() + state = gen.__getstate__() + assert state["input_queue"] is None + assert state["output_queue"] is None + # The live object itself is untouched. + assert gen.input_queue is not None + assert gen.output_queue is not None + + +def test_suggest_raises_clear_error_without_queues(): + # Simulates a QueueGenerator that ended up running somewhere other than + # the manager (e.g. gen_on_worker=True, bypassing QueueService's guard) + # and therefore only has the stripped-down pickled state. + gen = _make_gen() + gen.input_queue = None + gen.output_queue = None + with pytest.raises(RuntimeError, match="must run on the manager"): + gen.suggest(1) + + +def test_ingest_raises_clear_error_without_queues(): + gen = _make_gen() + gen.input_queue = None + gen.output_queue = None + with pytest.raises(RuntimeError, match="must run on the manager"): + gen.ingest([{"x": 0.0}]) + + +def test_queue_service_rejects_gen_on_worker(): + from libensemble.specs import LibeSpecs + + vocs = VOCS(variables={"x": [-1, 1]}) + libE_specs = LibeSpecs(nworkers=2, gen_on_worker=True) + with pytest.raises(AssertionError, match="gen_on_worker"): + QueueService(vocs, sim_specs={}, libE_specs=libE_specs, exit_criteria={"sim_max": 1}) + + +def test_suggest_returns_empty_after_shutdown_seen(): + gen = _make_gen() + gen.poll_timeout = 0.05 + gen._shutdown_seen = True + assert gen.suggest(1) == [] + assert gen.suggest(5) == [] + + +def test_suggest_shutdown_sentinel_sets_flag(): + gen = _make_gen() + gen.poll_timeout = 0.1 + sentinel = QueueGenerator.shutdown_sentinel() + gen.input_queue.put(sentinel) + assert gen.suggest(1) == [] + assert gen._shutdown_seen is True + assert gen.suggest(1) == [] + + +def test_finalize_forwards_results(): + gen = _make_gen() + gen.finalize([{"x": 1.0, "y": 2.0}]) + assert gen.output_queue.get_nowait() == {"x": 1.0, "y": 2.0} + gen.finalize(None) + gen.finalize([]) + assert gen.output_queue.empty() + + +def test_shutdown_sentinel_identity(): + s1 = QueueGenerator.shutdown_sentinel() + s2 = QueueGenerator.shutdown_sentinel() + assert s1 is s2 + + +def test_get_completed_drains(): + vocs = VOCS(variables={"x": [-1, 1]}) + from libensemble.specs import ExitCriteria, LibeSpecs + + service = QueueService(vocs, sim_specs={}, libE_specs=LibeSpecs(nworkers=1), exit_criteria=ExitCriteria(sim_max=1)) + assert service.get_completed() == [] + service.output_queue.put({"x": 1.0}) + service.output_queue.put({"x": 2.0}) + out = service.get_completed() + assert out == [{"x": 1.0}, {"x": 2.0}] + assert service.get_completed() == [] + + +def test_collect_results_timeout_empty(): + vocs = VOCS(variables={"x": [-1, 1]}) + from libensemble.specs import ExitCriteria, LibeSpecs + + service = QueueService(vocs, sim_specs={}, libE_specs=LibeSpecs(nworkers=1), exit_criteria=ExitCriteria(sim_max=1)) + start = time.time() + out = service.collect_results(1, timeout=1.0) + elapsed = time.time() - start + assert out == [] + assert 0.9 <= elapsed < 2.0 + + service.output_queue.put({"x": 9.0}) + out = service.collect_results(1, timeout=1.0) + assert out == [{"x": 9.0}] + + +def test_stream_results_timeout_and_yield(): + vocs = VOCS(variables={"x": [-1, 1]}) + from libensemble.specs import ExitCriteria, LibeSpecs + + service = QueueService(vocs, sim_specs={}, libE_specs=LibeSpecs(nworkers=1), exit_criteria=ExitCriteria(sim_max=1)) + start = time.time() + out = list(service.stream_results(n=1, timeout=1.0)) + elapsed = time.time() - start + assert out == [] + assert 0.9 <= elapsed < 2.0 + + service.output_queue.put({"x": 1.0}) + service.output_queue.put({"x": 2.0}) + out = list(service.stream_results(n=2, timeout=2.0)) + assert out == [{"x": 1.0}, {"x": 2.0}] + + +def test_shutdown_wait_joins(): + vocs = VOCS(variables={"x": [-1, 1]}) + from libensemble.specs import ExitCriteria, LibeSpecs + + service = QueueService(vocs, sim_specs={}, libE_specs=LibeSpecs(nworkers=1), exit_criteria=ExitCriteria(sim_max=1)) + mock_thread = MagicMock() + service._thread = mock_thread + service.shutdown(wait=True) + assert service.input_queue.get_nowait() is QueueGenerator.shutdown_sentinel() + mock_thread.join.assert_called_once() + + mock_thread.reset_mock() + service.shutdown(wait=False) + mock_thread.join.assert_not_called() + + +if __name__ == "__main__": + test_getstate_strips_queues() + test_suggest_raises_clear_error_without_queues() + test_ingest_raises_clear_error_without_queues() + test_queue_service_rejects_gen_on_worker()