diff --git a/docs/orchestrator.md b/docs/orchestrator.md index a189c1584..e5db197b5 100644 --- a/docs/orchestrator.md +++ b/docs/orchestrator.md @@ -51,11 +51,12 @@ public: // --- Intermediate-buffer allocation (runtime-owned lifetime) --- Tensor alloc(const std::vector &shape, DataType dtype); - // --- Internal lifecycle (invoked by Worker::run only) --- + // --- Internal lifecycle (invoked by Python Worker.submit/RunHandle) --- RunId begin_run(); void close_run_submission(RunId run_id); void fail_run_submission(RunId run_id, std::exception_ptr error); void wait_run(RunId run_id); + bool wait_run_for(RunId run_id, double timeout_seconds); bool run_done(RunId run_id) const; void release_run(RunId run_id); void scope_begin(); @@ -72,9 +73,32 @@ struct SubmitResult { TaskSlot task_slot; }; // internal only; not bound to Pyt `config`, since SUB has no per-call config. The run lifecycle and outer `scope_begin` / `scope_end` are invoked from Python -`Worker.run` through private bindings. They are not part of the user-facing -orch-fn API. `Worker.run` remains blocking: it closes submission, waits for the -matching run fence, performs run-owned cleanup, and returns `None`. +`Worker.submit` through private bindings. They are not part of the user-facing +orch-fn API. `Worker.submit` invokes the orchestration callback and closes DAG +submission synchronously, then returns a `RunHandle` before L3 device work has +necessarily completed. Graph-construction errors therefore remain synchronous; +device or endpoint errors are attached to the handle and raised by `wait()` or +`result()`. + +`RunHandle.done` polls the matching native run fence. `wait(timeout)` supports +bounded waits without cancelling or corrupting the run, and repeated waits +replay the same terminal result. The handle keeps its `Worker`, callback +arguments, configuration, and run-owned cleanup state alive until completion. +`Worker.close()` rejects new submissions and drains every accepted handle +before tearing down the worker tree. + +`Worker.run` remains source-compatible and blocking: + +```python +worker.run(orchestration, args, config) +# Equivalent to: +worker.submit(orchestration, args, config).wait() +``` + +The current L2 backend is synchronous, so L2 `submit()` executes the existing +blocking path and returns an already-completed handle. L3 asynchronous return +does not yet imply overlapping runs: a later submit waits for the prior run's +fence and cleanup before building the next DAG. Remote L3 submit adds two hidden pieces of metadata: final eligible worker-id sets and optional `RemoteTaskArgsSidecar` entries aligned by tensor index. diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index 4185031c0..63b0f7e43 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -367,6 +367,10 @@ inline void bind_worker(nb::module_ &m) { "_wait_run", &Orchestrator::wait_run, nb::arg("run_id"), nb::call_guard(), "Block until one run is terminal and raise only that run's error." ) + .def( + "_wait_run_for", &Orchestrator::wait_run_for, nb::arg("run_id"), nb::arg("timeout_seconds"), + nb::call_guard(), "Wait up to timeout_seconds for one run to become terminal." + ) .def("_run_done", &Orchestrator::run_done, nb::arg("run_id")) .def("_release_run", &Orchestrator::release_run, nb::arg("run_id")); diff --git a/python/simpler/orchestrator.py b/python/simpler/orchestrator.py index 29e3537df..789918ec2 100644 --- a/python/simpler/orchestrator.py +++ b/python/simpler/orchestrator.py @@ -6,7 +6,7 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- -"""Orchestrator — DAG builder exposed to the user's orch function during Worker.run(). +"""Orchestrator — DAG builder passed to a Worker submit/run callback. A thin Python facade over the C++ ``Orchestrator``. The Worker creates one Orchestrator handle at init, retrieves the C++ object via ``Worker.get_orchestrator()``, @@ -24,10 +24,12 @@ def my_orch(orch, args, cfg): sub_args.add_tensor(make_tensor_arg(output_tensor), TensorArgType.INPUT) orch.submit_sub(sub_handle, sub_args) - w.run(my_orch, my_args, my_config) + handle = w.submit(my_orch, my_args, my_config) + handle.wait() -Scope/drain lifecycle is managed by ``Worker.run()``; users never call those -directly. +Scope and submission-close lifecycle is managed by ``Worker.submit()``; +completion is managed by its ``RunHandle``. ``Worker.run()`` remains the +blocking ``submit(...).wait()`` compatibility entry point. """ from __future__ import annotations @@ -440,7 +442,7 @@ def create_l3_l2_queue(self, *, worker_id: int, depth: int, input_arena_bytes: i # deeper heap ring (``min(depth, MAX_RING_DEPTH-1)``) so their # memory reclaims independently of the outer scope. ``scope_end`` is # non-blocking — it releases scope refs and returns; call - # ``Worker.run``/``drain`` for a synchronous wait. + # the ``RunHandle`` returned by ``Worker.submit`` for a synchronous wait. # # Usage:: # @@ -543,7 +545,7 @@ def alloc(self, shape: Sequence[int], dtype: DataType) -> Tensor: return tensor # ------------------------------------------------------------------ - # Internal (called by Worker.run) + # Internal (called by Worker.submit) # ------------------------------------------------------------------ def _scope_begin(self) -> None: @@ -564,6 +566,9 @@ def _fail_run_submission(self, run_id: int, message: str = "") -> None: def _wait_run(self, run_id: int) -> None: self._o._wait_run(run_id) + def _wait_run_for(self, run_id: int, timeout_seconds: float) -> bool: + return bool(self._o._wait_run_for(run_id, timeout_seconds)) + def _run_done(self, run_id: int) -> bool: return bool(self._o._run_done(run_id)) diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 3be93c2ef..4280051eb 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -1876,6 +1876,125 @@ class _StartupCancelled(BaseException): cancellable (``close()`` fails fast while INITIALIZING).""" +class RunHandle: + """Completion handle returned by :meth:`Worker.submit`. + + A handle owns the run's Python keepalives and keeps its Worker alive until + the native completion fence has fired and run-owned resources are cleaned + up. Waiting is idempotent; every waiter observes the same terminal result. + """ + + def __init__(self, worker: Worker, run_id: int, keepalive: tuple[Any, ...]) -> None: + self._worker = worker + self._run_id: int | None = run_id + self._keepalive: tuple[Any, ...] | None = keepalive + self._cv = threading.Condition() + self._wait_in_progress = False + self._terminal = False + self._error: BaseException | None = None + + @classmethod + def _completed(cls, worker: Worker) -> RunHandle: + handle = cls.__new__(cls) + handle._worker = worker + handle._run_id = None + handle._keepalive = None + handle._cv = threading.Condition() + handle._wait_in_progress = False + handle._terminal = True + handle._error = None + return handle + + @staticmethod + def _deadline(timeout: float | None) -> float | None: + if timeout is None: + return None + value = float(timeout) + if value < 0 or not math.isfinite(value): + raise ValueError("RunHandle timeout must be a non-negative finite number of seconds") + return time.monotonic() + value + + @property + def done(self) -> bool: + """True once this run is terminal and its cleanup has been published. + + Reads False while another waiter is still publishing cleanup, so a run + whose native fence has already fired can report False until that waiter + finishes. + """ + with self._cv: + if self._terminal: + return True + # A waiter may have crossed the native fence and released the run + # identity while it is still publishing cleanup. Avoid querying a + # native id that can disappear in that interval. + if self._wait_in_progress: + return False + run_id = self._run_id + assert run_id is not None + return self._worker._run_handle_done(run_id) + + def wait(self, timeout: float | None = None) -> None: + """Wait for completion, raising ``TimeoutError`` or the run's error.""" + deadline = self._deadline(timeout) + with self._cv: + while not self._terminal and self._wait_in_progress: + remaining = None if deadline is None else deadline - time.monotonic() + if remaining is not None and remaining <= 0: + raise TimeoutError("RunHandle.wait() timed out") + self._cv.wait(timeout=remaining) + if self._terminal: + error = self._error + if error is not None: + raise error + return + self._wait_in_progress = True + run_id = self._run_id + + assert run_id is not None + native_error: BaseException | None = None + try: + remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) + completed = self._worker._wait_run_handle(run_id, remaining) + except Exception as exc: # native run failures are terminal + completed = True + native_error = exc + except BaseException: + with self._cv: + self._wait_in_progress = False + self._cv.notify_all() + raise + + if not completed: + with self._cv: + self._wait_in_progress = False + self._cv.notify_all() + raise TimeoutError("RunHandle.wait() timed out") + + error = self._worker._finalize_run_handle(self, run_id, native_error) + with self._cv: + self._error = error + self._run_id = None + self._keepalive = None + self._terminal = True + self._wait_in_progress = False + self._cv.notify_all() + if error is not None: + raise error + + def result(self, timeout: float | None = None) -> None: + """Alias for :meth:`wait`; successful runs have no return value.""" + self.wait(timeout) + + def _wait_for_serialization(self) -> None: + """Drain this run before another submission, without poisoning it.""" + try: + self.wait() + except Exception: + # The result remains cached for this handle's public wait/result. + pass + + def _forked_child_main(buf: memoryview, label: str, setup, serve, make_group_leader: bool = False) -> None: """Run a forked child to completion, always terminating via ``os._exit``. @@ -2036,6 +2155,13 @@ def __init__( self._last_rollback: dict[str, list[int]] | None = None self._hierarchical_start_mu = threading.Lock() self._hierarchical_start_cv = threading.Condition(self._hierarchical_start_mu) + # Asynchronous completion currently admits only one live device run. A + # new submit drains the prior accepted handle under this gate. + self._submit_mu = threading.Lock() + # Guarded by _hierarchical_start_cv. Handles are installed before their + # orchestration callback can enqueue work and retired after fence-owned + # cleanup, so close() can drain the exact accepted set. + self._accepted_run_handles: set[RunHandle] = set() # Absolute time.monotonic() deadline for the current startup epoch, set # once at init() and shared by every child group and recursive descendant # so the whole tree comes up within a single startup_timeout_s budget. @@ -5630,99 +5756,142 @@ def _find_host_buf_entry(self, addr: int, nbytes: int) -> _HostBufEntry | None: # run — uniform entry point # ------------------------------------------------------------------ - def run(self, callable, args=None, config=None) -> None: - """Execute one task (L2) or one DAG (L3+) synchronously. + def submit(self, callable, args=None, config=None) -> RunHandle: + """Submit one task (L2) or one DAG (L3+) and return its completion handle. Dispatch: - L2: ``callable`` is a ``CallableHandle`` returned by ``Worker.register(chip_callable)``. Routes to the private slot - carried by the handle. + carried by the handle. The current L2 backend remains blocking, so + the returned handle is already complete. - L3+: ``callable`` is a Python orch fn invoked with the - ``Orchestrator`` handle. + ``Orchestrator`` handle. Graph construction completes synchronously; + device completion is reported by the returned handle. ``args`` : TaskArgs (optional) ``config``: CallConfig (optional, default-constructed if None) - Returns ``None``. Per-stage run timing (host wall, on-NPU device wall + + Only one live device run is admitted: a later submission waits for the + previous handle's fence and cleanup before building its DAG. + """ + with self._operation_lease("submit"): + return self._submit_locked(callable, args, config) + + def run(self, callable, args=None, config=None) -> None: + """Execute one task or DAG synchronously as ``submit(...).wait()``. + + Per-stage run timing (host wall, on-NPU device wall + AICPU phase breakdown) is no longer returned — the platform emits it as ``[STRACE]`` log markers from each L2 ``simpler_run``, so the L3 dispatcher and its L2 children are observed uniformly. Parse the markers with ``simpler_setup.tools.strace_timing`` (see ``docs/dfx/host-trace.md``). """ - with self._operation_lease("run"): - self._run_locked(callable, args, config) + self.submit(callable, args=args, config=config).wait() - def _run_locked(self, callable, args, config) -> None: + def _submit_locked(self, callable, args, config) -> RunHandle: cfg = config if config is not None else CallConfig() if self.level == 2: assert self._chip_worker is not None state = self._resolve_handle(callable, expected_namespace="LOCAL_CHIP") self._chip_worker._run_slot(state.slot_id, args, cfg) - return None + return RunHandle._completed(self) + with self._submit_mu: + # Cleanup is Worker-global while only one live device run is + # admitted. Drain prior handles before a new callback can mutate + # those resources; errors remain attached only to their origin. + with self._hierarchical_start_cv: + prior_handles = tuple(self._accepted_run_handles) + for handle in prior_handles: + handle._wait_for_serialization() + return self._submit_l3_locked(callable, args, cfg) + + def _submit_l3_locked(self, callable, args, cfg: CallConfig) -> RunHandle: assert self._orch is not None assert self._worker is not None run_id = self._orch._begin_run() + handle = RunHandle(self, run_id, (callable, args, cfg)) + with self._hierarchical_start_cv: + self._accepted_run_handles.add(handle) + self._hierarchical_start_cv.notify_all() + scope_open = False - submission_failed = True - submission_error: BaseException | None = None try: self._orch._scope_begin() scope_open = True callable(self._orch, args, cfg) - submission_failed = False + scope_open = False + self._orch._scope_end() + self._orch._close_run_submission(run_id) except BaseException as e: - submission_error = e - raise - finally: try: if scope_open: + scope_open = False self._orch._scope_end() - except BaseException as e: - submission_failed = True - if submission_error is None: - submission_error = e - raise finally: try: - # Closing the run must not be able to strand `run_id` in the - # engine's registry: `_release_run` and the run-owned - # resource cleanup below both live under this `try`. - if submission_failed: - self._orch._fail_run_submission( - run_id, - "" if submission_error is None else _format_exc("orchestration", submission_error), - ) - else: - self._orch._close_run_submission(run_id) - try: - self._orch._wait_run(run_id) - except Exception as e: - self._poison_l3_l2_region_from_endpoint_error(e) - # Preserve a synchronous graph-construction exception - # already propagating from the orchestration callback. - if not submission_failed: - raise + self._orch._fail_run_submission(run_id, _format_exc("orchestration", e)) + except Exception: + # The orchestration cause is the useful one, so a failure + # while closing the run must not replace it; the handle's + # own wait still reports the run's terminal state. + pass finally: - try: - # Run-owned resources remain live until the per-run - # completion fence fires, even on failure. - self._release_active_remote_slot_refs() - self._flush_pending_remote_frees() - try: - self._cleanup_l3_l2_regions() - finally: - self._l3_l2_orch_comm_host_buffers.clear() - self._execute_pending_domain_releases() - if self._live_domains: - self._release_all_live_domains() - finally: - self._orch._release_run(run_id) - # L3+ returns None like every other worker level; per-L2-child timing - # is emitted as `[STRACE]` markers from each simpler_run. - return None + # Graph-construction failures remain synchronous, but any + # tasks already submitted still own their resources until + # the run fence fires. + handle._wait_for_serialization() + raise + return handle + + def _run_handle_done(self, run_id: int) -> bool: + assert self._orch is not None + return self._orch._run_done(run_id) + + def _wait_run_handle(self, run_id: int, timeout: float | None) -> bool: + assert self._orch is not None + if timeout is None: + self._orch._wait_run(run_id) + return True + return self._orch._wait_run_for(run_id, timeout) + + def _finalize_run_handle( + self, handle: RunHandle, run_id: int, native_error: BaseException | None + ) -> BaseException | None: + """Run fence-owned cleanup exactly once and return the cached result.""" + result = native_error + + def _step(fn) -> None: + nonlocal result + try: + fn() + except BaseException as exc: # noqa: BLE001 + if result is None: + result = exc + + if native_error is not None: + _step(lambda: self._poison_l3_l2_region_from_endpoint_error(native_error)) + _step(self._release_active_remote_slot_refs) + _step(self._flush_pending_remote_frees) + try: + _step(self._cleanup_l3_l2_regions) + finally: + self._l3_l2_orch_comm_host_buffers.clear() + _step(self._execute_pending_domain_releases) + if self._live_domains: + _step(self._release_all_live_domains) + orch = self._orch + if orch is None: + if result is None: + result = RuntimeError("RunHandle cleanup lost its native Orchestrator") + else: + _step(lambda: orch._release_run(run_id)) + with self._hierarchical_start_cv: + self._accepted_run_handles.discard(handle) + self._hierarchical_start_cv.notify_all() + return result @property def aicpu_dlopen_count(self) -> int: @@ -5824,11 +5993,13 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r attempt: _CloseAttempt | None = None result: BaseException | None = None teardown_tree = False + drain_complete = False + handles_to_drain: tuple[RunHandle, ...] = () try: with self._hierarchical_start_cv: if threading.get_ident() in self._lease_depth: raise RuntimeError( - "Worker.close(): cannot be called from within a run() / create_host_buffer() " + "Worker.close(): cannot be called from within a run() / submit() / create_host_buffer() " "operation on this thread" ) if self._lifecycle is _Lifecycle.INITIALIZING: @@ -5893,11 +6064,29 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r f"({_ROLLBACK_GRACEFUL_TIMEOUT_S}s); teardown deferred (worker stays CLOSED)" ) if result is None: - teardown_tree = self._has_live_resources() - # Latch terminal: once we commit to teardown no later close() - # re-drives it, whatever the outcome. - if teardown_tree: - self._teardown_attempted = True + handles_to_drain = tuple(self._accepted_run_handles) + drain_complete = True + if drain_complete: + # CLOSED prevents new admissions and active operations have + # drained, so this is the complete accepted set. Wait outside + # the lifecycle CV: handle retirement acquires the same lock. + for handle in handles_to_drain: + try: + handle.wait() + except BaseException as exc: # noqa: BLE001 + if result is None: + result = exc + with self._hierarchical_start_cv: + if self._accepted_run_handles: + # An asynchronous interruption left at least one fence + # undrained. Keep teardown retryable and the tree intact. + drain_complete = False + else: + teardown_tree = self._has_live_resources() + # Fence drain makes this close terminal even when the + # run error is the only remaining outcome and no native + # resource needs teardown. + self._teardown_attempted = teardown_tree or result is not None if teardown_tree: self._teardown_ready_tree() except BaseException as exc: # noqa: BLE001 diff --git a/src/common/hierarchical/orchestrator.cpp b/src/common/hierarchical/orchestrator.cpp index c462998fb..523156572 100644 --- a/src/common/hierarchical/orchestrator.cpp +++ b/src/common/hierarchical/orchestrator.cpp @@ -11,6 +11,7 @@ #include "orchestrator.h" +#include #include #include #include @@ -110,10 +111,10 @@ void Orchestrator::fail_run_submission(RunId run_id, std::exception_ptr error) { if (error) record_run_error(run_id, std::move(error)); { std::lock_guard runs_lk(runs_mu_); - if (building_run_id_ != run_id) { - throw std::logic_error("Orchestrator::fail_run_submission: run is not building"); - } - building_run_id_ = INVALID_RUN_ID; + // Failing a run closes it out from whatever state submission reached, + // so the fence always becomes reachable. Refusing a run that is no + // longer building would leave its waiter blocked forever. + if (building_run_id_ == run_id) building_run_id_ = INVALID_RUN_ID; } { std::lock_guard lk(run->completion_mu); @@ -139,6 +140,22 @@ void Orchestrator::wait_run(RunId run_id) { if (error) std::rethrow_exception(error); } +bool Orchestrator::wait_run_for(RunId run_id, double timeout_seconds) { + auto run = get_run(run_id); + std::exception_ptr error; + { + std::unique_lock lk(run->completion_mu); + if (!run->completion_cv.wait_for(lk, std::chrono::duration(timeout_seconds), [&run] { + return is_terminal(run->phase.load(std::memory_order_acquire)); + })) { + return false; + } + error = run->first_error; + } + if (error) std::rethrow_exception(error); + return true; +} + bool Orchestrator::run_done(RunId run_id) const { return is_terminal(get_run(run_id)->phase.load(std::memory_order_acquire)); } diff --git a/src/common/hierarchical/orchestrator.h b/src/common/hierarchical/orchestrator.h index e27d1ef9e..5273def13 100644 --- a/src/common/hierarchical/orchestrator.h +++ b/src/common/hierarchical/orchestrator.h @@ -121,6 +121,7 @@ class Orchestrator { void close_run_submission(RunId run_id); void fail_run_submission(RunId run_id, std::exception_ptr error = nullptr); void wait_run(RunId run_id); + bool wait_run_for(RunId run_id, double timeout_seconds); bool run_done(RunId run_id) const; bool run_failed(RunId run_id) const; void release_run(RunId run_id); diff --git a/tests/ut/cpp/hierarchical/test_orchestrator.cpp b/tests/ut/cpp/hierarchical/test_orchestrator.cpp index 428596ddb..df6ecb682 100644 --- a/tests/ut/cpp/hierarchical/test_orchestrator.cpp +++ b/tests/ut/cpp/hierarchical/test_orchestrator.cpp @@ -443,6 +443,15 @@ TEST_F(OrchestratorFixture, EmptyRunCompletesWhenSubmissionCloses) { orch.release_run(run_id); } +TEST_F(OrchestratorFixture, TimedWaitCanRetryAfterTimeout) { + EXPECT_FALSE(orch.wait_run_for(run_id, 0.0)); + EXPECT_FALSE(orch.run_done(run_id)); + + orch.close_run_submission(run_id); + EXPECT_TRUE(orch.wait_run_for(run_id, 0.0)); + orch.release_run(run_id); +} + TEST_F(OrchestratorFixture, OneTaskRunCompletesAfterConsumption) { auto result = orch.submit_next_level(C(80), single_tensor_args(0x8000, TensorArgType::OUTPUT), cfg, 0); EXPECT_EQ(S(result.task_slot).run_id, run_id); @@ -540,6 +549,25 @@ TEST_F(OrchestratorFixture, ConsumingASlotOwnedByAReleasedRunIsIgnored) { EXPECT_FALSE(orch.run_done(run_id)); } +// Failing a run is the recovery path: refusing a run whose submission already +// closed would leave the caller's fence wait blocked forever. +TEST_F(OrchestratorFixture, FailingAnAlreadyClosedRunStillResolvesTheFence) { + auto task = orch.submit_next_level(C(86), single_tensor_args(0x8600, TensorArgType::OUTPUT), cfg, 0); + orch.close_run_submission(run_id); + EXPECT_FALSE(orch.run_done(run_id)); + + EXPECT_NO_THROW( + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("late submission failure"))) + ); + + S(task.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + EXPECT_TRUE(orch.on_consumed(task.task_slot)); + EXPECT_TRUE(orch.run_done(run_id)); + EXPECT_TRUE(orch.run_failed(run_id)); + EXPECT_THROW(orch.wait_run(run_id), std::runtime_error); + orch.release_run(run_id); +} + TEST_F(OrchestratorFixture, FailedSubmissionCarriesItsMessageToTheFence) { orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("orchestration: ValueError: bad arg"))); EXPECT_TRUE(orch.run_failed(run_id)); diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index 5be779288..269eed068 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -13,9 +13,13 @@ """ import ctypes +import gc import struct import threading +import time +import weakref from multiprocessing.shared_memory import SharedMemory +from typing import cast import pytest import simpler.worker as worker_mod @@ -45,6 +49,7 @@ _CTRL_UNREGISTER, _IDLE, _OFF_STATE, + RunHandle, Worker, _buffer_field_addr, _mailbox_addr, @@ -1387,6 +1392,247 @@ def orch(o, args, cfg): counter_shm.unlink() +class TestRunHandle: + def test_submit_returns_before_completion_and_timeout_is_retryable(self): + state_shm = SharedMemory(create=True, size=8) + state_buf = state_shm.buf + assert state_buf is not None + _set_flag(state_buf, 0, 0) + _set_flag(state_buf, 4, 0) + try: + hw = Worker(level=3, num_sub_workers=1) + + def delayed(_args): + while _get_flag(state_buf, 0) == 0: + time.sleep(0.001) + _set_flag(state_buf, 4, 1) + + target = hw.register(delayed) + hw.init() + callback_done = False + + def orch(o, _args, _cfg): + nonlocal callback_done + o.submit_sub(target) + callback_done = True + + handle = hw.submit(orch) + assert isinstance(handle, RunHandle) + assert callback_done + assert not handle.done + with pytest.raises(TimeoutError, match="timed out"): + handle.wait(0.01) + assert not handle.done + + _set_flag(state_buf, 0, 1) + handle.wait(5.0) + assert handle.done + assert _get_flag(state_buf, 4) == 1 + handle.wait() + assert handle.result() is None + hw.close() + finally: + state_shm.close() + state_shm.unlink() + + def test_graph_error_is_synchronous(self): + hw = Worker(level=3, num_sub_workers=0) + hw.init() + try: + + def bad_graph(_orch, _args, _cfg): + raise ValueError("bad graph") + + with pytest.raises(ValueError, match="bad graph"): + hw.submit(bad_graph) + assert not hw._accepted_run_handles + finally: + hw.close() + + def test_failed_run_does_not_poison_next_submit(self): + counter_shm, counter_buf = _make_shared_counter() + try: + hw = Worker(level=3, num_sub_workers=1) + + def fail(_args): + raise RuntimeError("first run failed") + + failed_target = hw.register(fail) + good_target = hw.register(lambda _args: _increment_counter(counter_buf)) + hw.init() + failed = hw.submit(lambda o, _args, _cfg: o.submit_sub(failed_target)) + good = hw.submit(lambda o, _args, _cfg: o.submit_sub(good_target)) + + with pytest.raises(RuntimeError, match="first run failed"): + failed.wait() + good.wait() + assert _read_counter(counter_buf) == 1 + hw.close() + finally: + counter_shm.close() + counter_shm.unlink() + + def test_close_drains_accepted_handle_and_rejects_later_submit(self): + state_shm = SharedMemory(create=True, size=4) + state_buf = state_shm.buf + assert state_buf is not None + _set_flag(state_buf, 0, 0) + try: + hw = Worker(level=3, num_sub_workers=1) + + def delayed(_args): + while _get_flag(state_buf, 0) == 0: + time.sleep(0.001) + + target = hw.register(delayed) + hw.init() + handle = hw.submit(lambda o, _args, _cfg: o.submit_sub(target)) + releaser = threading.Thread(target=lambda: (time.sleep(0.1), _set_flag(state_buf, 0, 1))) + releaser.start() + try: + hw.close() + finally: + _set_flag(state_buf, 0, 1) + releaser.join(5.0) + assert handle.done + with pytest.raises(RuntimeError, match="requires an initialized"): + hw.submit(lambda *_args: None) + finally: + state_shm.close() + state_shm.unlink() + + def test_submit_close_race_accepts_and_drains_admitted_run(self): + callback_entered = threading.Event() + callback_release = threading.Event() + result: dict[str, object] = {} + hw = Worker(level=3, num_sub_workers=0) + hw.init() + + def orch(_o, _args, _cfg): + callback_entered.set() + assert callback_release.wait(5.0) + + submitter = threading.Thread(target=lambda: result.setdefault("handle", hw.submit(orch))) + submitter.start() + assert callback_entered.wait(3.0) + releaser = threading.Thread(target=lambda: (time.sleep(0.1), callback_release.set())) + releaser.start() + try: + hw.close() + finally: + callback_release.set() + submitter.join(5.0) + releaser.join(5.0) + handle = result["handle"] + assert isinstance(handle, RunHandle) + assert handle.done + + def test_orchestration_callable_is_kept_alive_until_completion(self): + class Keepalive: + pass + + state_shm = SharedMemory(create=True, size=4) + state_buf = state_shm.buf + assert state_buf is not None + _set_flag(state_buf, 0, 0) + token = Keepalive() + token_ref = weakref.ref(token) + try: + hw = Worker(level=3, num_sub_workers=1) + + def delayed(_args): + while _get_flag(state_buf, 0) == 0: + time.sleep(0.001) + + target = hw.register(delayed) + hw.init() + + def orch(o, _args, _cfg, held=token): + assert held is not None + o.submit_sub(target) + + handle = hw.submit(orch) + del orch, token + gc.collect() + assert token_ref() is not None + _set_flag(state_buf, 0, 1) + handle.wait(5.0) + gc.collect() + assert token_ref() is None + hw.close() + finally: + _set_flag(state_buf, 0, 1) + state_shm.close() + state_shm.unlink() + + def test_run_delegates_to_submit_and_wait(self, monkeypatch): + events: list[tuple] = [] + + class FakeHandle: + def wait(self): + events.append(("wait",)) + + def fake_submit(self, callable, args=None, config=None): + events.append(("submit", callable, args, config)) + return FakeHandle() + + monkeypatch.setattr(Worker, "submit", fake_submit) + worker = Worker(level=3, num_sub_workers=0) + + def callback(*_args): + return None + + worker.run(callback, args="args", config="config") + assert events == [("submit", callback, "args", "config"), ("wait",)] + + def test_serialization_drain_does_not_swallow_async_interrupt(self, monkeypatch): + handle = RunHandle._completed(Worker(level=3, num_sub_workers=0)) + + def interrupted_wait(_self, _timeout=None): + raise KeyboardInterrupt + + monkeypatch.setattr(RunHandle, "wait", interrupted_wait) + with pytest.raises(KeyboardInterrupt): + handle._wait_for_serialization() + + def test_done_query_cannot_race_native_run_release(self): + done_entered = threading.Event() + done_release = threading.Event() + wait_entered = threading.Event() + + class FakeWorker: + def _run_handle_done(self, run_id): + assert run_id == 1 + done_entered.set() + assert done_release.wait(5.0) + return True + + def _wait_run_handle(self, run_id, timeout): + assert run_id == 1 + wait_entered.set() + return True + + def _finalize_run_handle(self, handle, run_id, error): + assert run_id == 1 + return error + + handle = RunHandle(cast(Worker, FakeWorker.__new__(FakeWorker)), 1, ()) + observed: list[bool] = [] + done_thread = threading.Thread(target=lambda: observed.append(handle.done)) + wait_thread = threading.Thread(target=handle.wait) + done_thread.start() + assert done_entered.wait(3.0) + wait_thread.start() + try: + assert not wait_entered.wait(0.1) + finally: + done_release.set() + done_thread.join(5.0) + wait_thread.join(5.0) + assert observed == [True] + assert wait_entered.is_set() + + # --------------------------------------------------------------------------- # Test: multiple SUB workers execute in parallel # --------------------------------------------------------------------------- diff --git a/tests/ut/py/test_worker/test_startup_readiness.py b/tests/ut/py/test_worker/test_startup_readiness.py index abaa718c6..82be5ed56 100644 --- a/tests/ut/py/test_worker/test_startup_readiness.py +++ b/tests/ut/py/test_worker/test_startup_readiness.py @@ -39,7 +39,7 @@ import pytest from simpler.task_interface import CallConfig, TaskArgs -from simpler.worker import Worker +from simpler.worker import RunHandle, Worker # Hard wall budget for a single failure scenario — comfortably above the # injected startup_timeout_s values below, well under any real hang. @@ -695,6 +695,9 @@ def init(self, *_a, **_k): def _register_callable_at_slot(self, *_a, **_k): # pragma: no cover pass + def _run_slot(self, *_a, **_k): + pass + def finalize(self): pass @@ -758,6 +761,20 @@ def test_l2_init_then_close_does_not_hang(self, monkeypatch): # A second close is a clean no-op (does not re-block on the epoch cv). w.close() + def test_l2_submit_returns_completed_handle(self, monkeypatch): + from simpler.task_interface import ChipCallable # noqa: PLC0415 + + w = self._make_l2(monkeypatch) + callable = ChipCallable.build(signature=[], func_name="x", binary=b"\x00", children=[]) + callable_handle = w.register(callable) + with _hard_timeout(_TEST_WALL_BUDGET_S): + w.init() + run_handle = w.submit(callable_handle) + assert isinstance(run_handle, RunHandle) + assert run_handle.done + assert run_handle.wait() is None + w.close() + def test_l2_close_without_init_is_noop(self, monkeypatch): w = self._make_l2(monkeypatch) with _hard_timeout(_TEST_WALL_BUDGET_S): @@ -1047,8 +1064,9 @@ def test_close_drains_in_flight_operation(self, monkeypatch): def paused_run(self, *_a, **_k): entered.set() assert release.wait(10.0) + return RunHandle._completed(self) - monkeypatch.setattr(Worker, "_run_locked", paused_run) + monkeypatch.setattr(Worker, "_submit_locked", paused_run) w = self._make_l2(monkeypatch) # owner = this (main) thread with _hard_timeout(_TEST_WALL_BUDGET_S): w.init() @@ -1076,8 +1094,9 @@ def test_reentrant_close_from_operation_rejected(self, monkeypatch): def reentrant_run(self, *_a, **_k): result["close_err"] = _run_catch(self.close) + return RunHandle._completed(self) - monkeypatch.setattr(Worker, "_run_locked", reentrant_run) + monkeypatch.setattr(Worker, "_submit_locked", reentrant_run) w = self._make_l2(monkeypatch) # owner = main with _hard_timeout(_TEST_WALL_BUDGET_S): w.init() @@ -1101,8 +1120,9 @@ def test_close_timeout_defers_teardown_and_retry_completes(self, monkeypatch): def paused_run(self, *_a, **_k): entered.set() assert release.wait(10.0) + return RunHandle._completed(self) - monkeypatch.setattr(Worker, "_run_locked", paused_run) + monkeypatch.setattr(Worker, "_submit_locked", paused_run) w = self._make_l2(monkeypatch) # owner = main with _hard_timeout(_TEST_WALL_BUDGET_S): w.init() @@ -1136,8 +1156,9 @@ def test_close_retry_still_drains_before_teardown(self, monkeypatch): def paused_run(self, *_a, **_k): entered.set() assert release.wait(10.0) + return RunHandle._completed(self) - monkeypatch.setattr(Worker, "_run_locked", paused_run) + monkeypatch.setattr(Worker, "_submit_locked", paused_run) w = self._make_l2(monkeypatch) # owner = main with _hard_timeout(_TEST_WALL_BUDGET_S): w.init()