From 615a7261e6b0f77b51264de1529164fdc124d490 Mon Sep 17 00:00:00 2001 From: cjkindel Date: Tue, 15 Sep 2026 11:50:01 -0700 Subject: [PATCH 1/8] Claim a worker spawn before it suspends `spawn_worker` checked `_managed_worker_processes` for a duplicate, then wrote its entry 68 lines later, with the static-URL await and the fork in between. Two spawns for one library both got past the check and both forked; only one could be recorded, leaving the other's process untracked and holding that library's dependencies in memory until its own heartbeat lapsed. Two `StartWorkerRequest`s for one library name overlap whenever `_library_file_path_to_info` carries duplicate entries for it (a FAILURE entry alongside the LOADED copy), or when `_start_workers` runs twice before the first worker registers. The key is now claimed in the same step as the check, and released in a `finally` so a failed fork does not leave the library permanently unspawnable. --- .../retained_mode/managers/worker_manager.py | 147 ++++++++++-------- tests/unit/app/test_app_worker.py | 47 ++++++ 2 files changed, 128 insertions(+), 66 deletions(-) diff --git a/src/griptape_nodes/retained_mode/managers/worker_manager.py b/src/griptape_nodes/retained_mode/managers/worker_manager.py index 3410987325..036887af1a 100644 --- a/src/griptape_nodes/retained_mode/managers/worker_manager.py +++ b/src/griptape_nodes/retained_mode/managers/worker_manager.py @@ -122,6 +122,11 @@ def __init__( # Subprocesses spawned by this orchestrator (library_name → process) self._managed_worker_processes: dict[str, asyncio.subprocess.Process] = {} + # Worker keys whose spawn has been claimed but has not yet reached the registry above. + # The registry entry cannot serve as the claim: it is only written once the subprocess + # exists, and the work in between suspends. + self._spawns_in_flight: set[str] = set() + # The event loop that spawned the worker subprocesses. asyncio.subprocess.Process # binds its exit Future to its creating loop, so proc.wait() is only legal on this # loop. Eviction can run on a different loop (the websocket-tasks loop), which is @@ -368,75 +373,85 @@ async def spawn_worker(self, args: list[str], worker_key: str) -> None: worker_key is an opaque identifier used to track the process and prevent duplicate spawns. Callers are responsible for constructing the args list. """ - if worker_key in self._managed_worker_processes: + if worker_key in self._managed_worker_processes or worker_key in self._spawns_in_flight: logger.error("Worker for key '%s' already spawned; refusing duplicate spawn.", worker_key) return - # Spawn with the orchestrator's PRE-project environ so the worker boots with the - # same clean env baseline a fresh engine would have. Inheriting the live os.environ - # would bake the orchestrator's current-project env vars into the worker's restore - # baseline, leaving the worker unable to unset them on a later project switch. - base_environ = self.engine.project_manager.get_pre_project_environ() - worker_environ = {**base_environ, "GTN_ENGINE_ID": str(uuid.uuid4())} - # Stamp the spawning orchestrator's id so the worker can report it in its discovery - # heartbeat (orchestrator_engine_id), letting clients identify and nest worker engines. - # The orchestrator always has an id by the time it spawns a worker; guard the None - # case anyway so a subprocess env value is never None. - orchestrator_engine_id = self.engine.engine_identity_manager.active_engine_id - if orchestrator_engine_id is not None: - worker_environ["GTN_ORCHESTRATOR_ENGINE_ID"] = orchestrator_engine_id - # Worker stdout is a pipe when the orchestrator is hosted by a GUI app (e.g. the - # desktop app); unbuffered output keeps worker log lines from stalling in Python's - # block buffer and from being lost on a crash. - worker_environ["PYTHONUNBUFFERED"] = "1" - - # PYTHONPATH precedes site-packages, making this library-first with the engine's own - # environment as the fallback. It must be the environment rather than a later sys.path - # splice: sys.modules never reconsiders a module this process has already imported. - execution_site_packages = self.engine.library_manager.execution_site_packages(worker_key) - if execution_site_packages is not None: - # Prepended, not assigned: a launcher-set PYTHONPATH (embedding hosts, source checkouts) - # is part of the environment the engine itself booted with, and dropping it only in - # exec-deps workers would lose those modules in exactly one process kind. - inherited_pythonpath = worker_environ.get("PYTHONPATH") - worker_environ["PYTHONPATH"] = ( - execution_site_packages + os.pathsep + inherited_pythonpath - if inherited_pythonpath - else execution_site_packages - ) - logger.debug( - "Worker for library '%s' will resolve imports from %s first", - worker_key, - execution_site_packages, - ) + # Claimed here, in the same step as the check above, so no await separates them. Two + # spawns for one key would otherwise both get past a registry-only guard and both fork, + # and only one can be recorded -- leaving the other's process untracked, holding its + # library's dependencies in memory until its own heartbeat lapses. + self._spawns_in_flight.add(worker_key) + try: + # Spawn with the orchestrator's PRE-project environ so the worker boots with the + # same clean env baseline a fresh engine would have. Inheriting the live os.environ + # would bake the orchestrator's current-project env vars into the worker's restore + # baseline, leaving the worker unable to unset them on a later project switch. + base_environ = self.engine.project_manager.get_pre_project_environ() + worker_environ = {**base_environ, "GTN_ENGINE_ID": str(uuid.uuid4())} + # Stamp the spawning orchestrator's id so the worker can report it in its discovery + # heartbeat (orchestrator_engine_id), letting clients identify and nest worker engines. + # The orchestrator always has an id by the time it spawns a worker; guard the None + # case anyway so a subprocess env value is never None. + orchestrator_engine_id = self.engine.engine_identity_manager.active_engine_id + if orchestrator_engine_id is not None: + worker_environ["GTN_ORCHESTRATOR_ENGINE_ID"] = orchestrator_engine_id + # Worker stdout is a pipe when the orchestrator is hosted by a GUI app (e.g. the + # desktop app); unbuffered output keeps worker log lines from stalling in Python's + # block buffer and from being lost on a crash. + worker_environ["PYTHONUNBUFFERED"] = "1" + + # PYTHONPATH precedes site-packages, making this library-first with the engine's own + # environment as the fallback. It must be the environment rather than a later sys.path + # splice: sys.modules never reconsiders a module this process has already imported. + execution_site_packages = self.engine.library_manager.execution_site_packages(worker_key) + if execution_site_packages is not None: + # Prepended, not assigned: a launcher-set PYTHONPATH (embedding hosts, source checkouts) + # is part of the environment the engine itself booted with, and dropping it only in + # exec-deps workers would lose those modules in exactly one process kind. + inherited_pythonpath = worker_environ.get("PYTHONPATH") + worker_environ["PYTHONPATH"] = ( + execution_site_packages + os.pathsep + inherited_pythonpath + if inherited_pythonpath + else execution_site_packages + ) + logger.debug( + "Worker for library '%s' will resolve imports from %s first", + worker_key, + execution_site_packages, + ) - # No workspace variable here: GTN_CONFIG_ outranks the runtime project override, so a worker - # handed one could never follow its orchestrator onto a project's workspace again. The - # workspace arrives with the project, adopted from the registration reply. - - # Hand the worker the URL of the static server the orchestrator is ALREADY serving - # this workspace on. Without it the worker starts its own server, wins an arbitrary - # OS-assigned port, and hands back asset URLs on that port -- which die when the - # worker is evicted and are already dead by the time a saved workflow is reopened. - # Both processes share the workspace on disk, so the orchestrator's long-lived - # server is the right place to serve anything a worker writes. - static_base_url = await self._orchestrator_static_server_base_url() - if static_base_url is not None: - worker_environ[ORCHESTRATOR_STATIC_SERVER_BASE_URL_ENV] = static_base_url - # Hand the orchestrator's own stdout/stderr to the worker explicitly so worker log - # lines land in the same stream as orchestrator logs. Implicit inheritance is - # POSIX-only: on Windows, redirected std handles (e.g. the desktop app's pipes) are - # not passed to a child unless subprocess sends them via STARTF_USESTDHANDLES, so - # the worker would log to an invisible console instead. - proc = await asyncio.create_subprocess_exec( - *args, - env=worker_environ, - stdout=sys.stdout, - stderr=sys.stderr, - ) - # Record the loop that owns this subprocess so termination can hop back to it. - # All spawns run on the engine event-queue loop, so this is idempotent. - self._spawn_loop = asyncio.get_running_loop() - self._managed_worker_processes[worker_key] = proc + # No workspace variable here: GTN_CONFIG_ outranks the runtime project override, so a worker + # handed one could never follow its orchestrator onto a project's workspace again. The + # workspace arrives with the project, adopted from the registration reply. + + # Hand the worker the URL of the static server the orchestrator is ALREADY serving + # this workspace on. Without it the worker starts its own server, wins an arbitrary + # OS-assigned port, and hands back asset URLs on that port -- which die when the + # worker is evicted and are already dead by the time a saved workflow is reopened. + # Both processes share the workspace on disk, so the orchestrator's long-lived + # server is the right place to serve anything a worker writes. + static_base_url = await self._orchestrator_static_server_base_url() + if static_base_url is not None: + worker_environ[ORCHESTRATOR_STATIC_SERVER_BASE_URL_ENV] = static_base_url + # Hand the orchestrator's own stdout/stderr to the worker explicitly so worker log + # lines land in the same stream as orchestrator logs. Implicit inheritance is + # POSIX-only: on Windows, redirected std handles (e.g. the desktop app's pipes) are + # not passed to a child unless subprocess sends them via STARTF_USESTDHANDLES, so + # the worker would log to an invisible console instead. + proc = await asyncio.create_subprocess_exec( + *args, + env=worker_environ, + stdout=sys.stdout, + stderr=sys.stderr, + ) + # Record the loop that owns this subprocess so termination can hop back to it. + # All spawns run on the engine event-queue loop, so this is idempotent. + self._spawn_loop = asyncio.get_running_loop() + self._managed_worker_processes[worker_key] = proc + finally: + # Released even when the fork raises: the claim outliving a failed spawn would + # silently refuse every later attempt for this library. + self._spawns_in_flight.discard(worker_key) logger.info("Spawned worker for key '%s' (pid %s)", worker_key, proc.pid) async def reset_workers(self) -> None: diff --git a/tests/unit/app/test_app_worker.py b/tests/unit/app/test_app_worker.py index 0d20a9d15a..23eee936f4 100644 --- a/tests/unit/app/test_app_worker.py +++ b/tests/unit/app/test_app_worker.py @@ -690,6 +690,53 @@ async def test_duplicate_spawn_is_noop(self, worker_manager: WorkerManager) -> N mock_exec.assert_not_called() + @pytest.mark.asyncio + async def test_concurrent_spawns_for_one_key_fork_once(self, worker_manager: WorkerManager) -> None: + """Two spawns racing for one library must produce one subprocess. + + The registry entry is written only once the process exists, and the work in between + suspends, so checking the registry alone lets the second caller through. The loser's + process would then be untracked, holding its library's dependencies until its own + heartbeat lapsed. + """ + worker_manager.engine.library_manager.execution_site_packages.return_value = None # type: ignore[union-attr] + + async def _suspend_then_answer() -> None: + # Yields inside the window between the duplicate check and the registry write. + await asyncio.sleep(0) + + with ( + patch.object(worker_manager, "_orchestrator_static_server_base_url", _suspend_then_answer), + patch("asyncio.create_subprocess_exec", return_value=_managed_proc_mock()) as mock_exec, + ): + await asyncio.gather( + worker_manager.spawn_worker(["/usr/bin/gtn", "engine"], "My Library"), + worker_manager.spawn_worker(["/usr/bin/gtn", "engine"], "My Library"), + ) + + mock_exec.assert_called_once() + assert list(worker_manager._managed_worker_processes) == ["My Library"] + + @pytest.mark.asyncio + async def test_a_failed_fork_does_not_keep_the_key_claimed(self, worker_manager: WorkerManager) -> None: + """A spawn that raises must leave the key spawnable. + + The claim outliving a failed fork would silently refuse every later attempt for that + library, which reads as a worker that never starts and never says why. + """ + worker_manager.engine.library_manager.execution_site_packages.return_value = None # type: ignore[union-attr] + + with patch("asyncio.create_subprocess_exec", side_effect=OSError("no interpreter")): + with pytest.raises(OSError, match="no interpreter"): + await worker_manager.spawn_worker(["/usr/bin/gtn", "engine"], "My Library") + + assert "My Library" not in worker_manager._spawns_in_flight + + with patch("asyncio.create_subprocess_exec", return_value=_managed_proc_mock()) as mock_exec: + await worker_manager.spawn_worker(["/usr/bin/gtn", "engine"], "My Library") + + mock_exec.assert_called_once() + @pytest.mark.asyncio async def test_spawns_subprocess_with_provided_args(self, worker_manager: WorkerManager) -> None: mock_proc = MagicMock() From 42eefcd1ccace1f0f3e70656172ddece5c25230f Mon Sep 17 00:00:00 2001 From: cjkindel Date: Tue, 15 Sep 2026 11:50:55 -0700 Subject: [PATCH 2/8] Read a spawn task's cancellation before its exception `_log_spawn_error` asked `task.exception()` first, which raises `CancelledError` on a cancelled task. Raising from a done-callback surfaces as loop-level "Exception in callback" noise and skips the `note_worker_unavailable` below it, so nothing records that no worker is coming. Cancellation reaches this callback at loop teardown, where nothing is waiting on the worker, so the cost today is the noise rather than a stalled run. Note on A1 of the tracking issue, which covers the same function: the claim that the spawn task can be garbage collected while parked does not hold. `Event.wait()` appends a future to the event's `_waiters`, and the awaiting task installs `Task.__wakeup` as that future's done-callback, so the task stays reachable from the manager for as long as it is parked. No change made there. --- .../retained_mode/managers/worker_manager.py | 5 +++ tests/unit/app/test_app_worker.py | 34 +++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/src/griptape_nodes/retained_mode/managers/worker_manager.py b/src/griptape_nodes/retained_mode/managers/worker_manager.py index 036887af1a..23beb08e26 100644 --- a/src/griptape_nodes/retained_mode/managers/worker_manager.py +++ b/src/griptape_nodes/retained_mode/managers/worker_manager.py @@ -809,6 +809,11 @@ def _log_spawn_error(self, task: asyncio.Task, library_name: str) -> None: caller cannot tell that a bad interpreter or an OSError stopped the worker ever existing. Refusals that return rather than raise are invisible here and record themselves. """ + # Asked before `task.exception()`, which raises on a cancelled task. From a done-callback + # that surfaces as loop-level "Exception in callback" noise and skips the refusal below. + # Cancellation reaches here at loop teardown, where no run is waiting on a worker. + if task.cancelled(): + return exc = task.exception() if exc is None: return diff --git a/tests/unit/app/test_app_worker.py b/tests/unit/app/test_app_worker.py index 23eee936f4..c8182c78ad 100644 --- a/tests/unit/app/test_app_worker.py +++ b/tests/unit/app/test_app_worker.py @@ -1050,6 +1050,40 @@ async def test_returns_success_immediately(self, worker_manager: WorkerManager) assert isinstance(result, worker_events.StartWorkerResultSuccess) +class TestLogSpawnError: + @pytest.mark.asyncio + async def test_a_cancelled_spawn_does_not_raise_from_the_callback(self, worker_manager: WorkerManager) -> None: + """A cancelled spawn task must not make its own done-callback raise. + + `task.exception()` raises on a cancelled task, and a done-callback that raises becomes + loop-level "Exception in callback" noise with the refusal below it skipped. + """ + + async def _never() -> None: + await asyncio.sleep(3600) + + task = asyncio.create_task(_never()) + task.cancel() + await asyncio.gather(task, return_exceptions=True) + assert task.cancelled() + + worker_manager._log_spawn_error(task, "My Library") + + @pytest.mark.asyncio + async def test_a_failed_spawn_still_records_a_refusal(self, worker_manager: WorkerManager) -> None: + async def _raise() -> None: + msg = "no interpreter" + raise OSError(msg) + + task = asyncio.create_task(_raise()) + await asyncio.gather(task, return_exceptions=True) + + with patch.object(worker_manager, "note_worker_unavailable") as mock_refuse: + worker_manager._log_spawn_error(task, "My Library") + + mock_refuse.assert_called_once() + + class TestSpawnWhenSessionReady: @pytest.mark.asyncio async def test_skips_wait_when_session_already_active(self, worker_manager: WorkerManager) -> None: From e8f8e7da874a2060e978570bf3fb991e67a5be19 Mon Sep 17 00:00:00 2001 From: cjkindel Date: Tue, 15 Sep 2026 11:55:01 -0700 Subject: [PATCH 3/8] Stop tracking a request whose run was cancelled `route_to_worker` registers a request then awaits it. Every other exit from that await removes the entry: a response pops it in `_try_match`, eviction in `cancel_requests_by_tag`. Flow cancellation re-raised without removing anything, so the entry stayed in `_pending_requests` for the life of the process -- one per cancelled node execution. Each leaked entry keeps its worker's tag, so a later `cancel_requests_by_tag` walks and re-settles long-dead requests, and `pending_count`, the only visibility into that map, reports a number that never comes down. Nothing fails loudly; it degrades over a long editing session. `_cancel_request` did exactly the removal needed but was async, and awaiting inside an except block while a cancellation is being delivered adds a suspension point in the one place it is least wanted. Its body never awaited, so it becomes the public synchronous `discard_request` and its five internal callers drop their `await`. --- .../api_client/request_client.py | 47 ++++++++------- .../retained_mode/managers/worker_manager.py | 38 ++++++++---- tests/unit/app/test_app_worker.py | 60 +++++++++++++++++++ 3 files changed, 111 insertions(+), 34 deletions(-) diff --git a/src/griptape_nodes/api_client/request_client.py b/src/griptape_nodes/api_client/request_client.py index 72abf3fddb..dac3ec543e 100644 --- a/src/griptape_nodes/api_client/request_client.py +++ b/src/griptape_nodes/api_client/request_client.py @@ -74,7 +74,7 @@ def __init__( # Map of request_id -> pending request where tag identifies the originating worker/caller self._pending_requests: dict[str, _PendingRequest] = {} # threading.Lock, not asyncio.Lock: this guards state reached from more than one loop -- - # _try_match runs on the transport loop while _track_request and _cancel_request run on the + # _try_match runs on the transport loop while _track_request and discard_request run on the # loop that issued the request. An asyncio.Lock binds to the loop that first awaits it and # only checks on the CONTENDED path, so a cross-loop acquire raises intermittently while the # uncontended fast path excludes nothing at all. Every section it guards is synchronous @@ -167,12 +167,12 @@ async def request( except TimeoutError: logger.error("Request %s timed out", request_id) - await self._cancel_request(request_id) + self.discard_request(request_id) raise except Exception as e: logger.error("Request %s failed: %s", request_id, e) - await self._cancel_request(request_id) + self.discard_request(request_id) raise else: logger.debug("Request %s completed successfully", request_id) @@ -230,12 +230,12 @@ async def request_to_orchestrator( except TimeoutError: logger.error("Forwarded request %s timed out", request_id) - await self._cancel_request(request_id) + self.discard_request(request_id) raise except Exception as e: logger.error("Forwarded request %s failed: %s", request_id, e) - await self._cancel_request(request_id) + self.discard_request(request_id) raise else: logger.debug("Forwarded request %s completed", request_id) @@ -321,7 +321,7 @@ async def request_batch( except (TimeoutError, Exception) as e: logger.error("Batch request failed: %s", e) for request_id in request_ids: - await self._cancel_request(request_id) + self.discard_request(request_id) raise else: logger.debug("Batch of %d requests completed", len(inner_events)) @@ -375,6 +375,25 @@ async def fail_requests_by_tag(self, tag: str, error: Exception) -> None: RequestClient._settle(lambda entry=entry: entry.future.set_exception(error)) logger.debug("Failed request %s (tag=%s): %s", rid, tag, error) + def discard_request(self, request_id: str) -> None: + """Stop tracking a request and cancel its future. + + Synchronous so a caller cancelled mid-await can call it from its own except block without + introducing a suspension point while the cancellation is being delivered. + + Args: + request_id: Request identifier + """ + with self._lock: + entry = self._pending_requests.pop(request_id, None) + + if entry is None: + logger.debug("Request already completed or unknown: %s", request_id) + return + + RequestClient._settle(entry.future.cancel) + logger.debug("Cancelled request: %s", request_id) + async def _track_request( self, request_id: str, @@ -473,22 +492,6 @@ async def _reject_request(self, request_id: str, error: Exception) -> None: with self._lock: self._reject_request_unlocked(request_id, error) - async def _cancel_request(self, request_id: str) -> None: - """Cancel a pending request and clean up its tracking. - - Args: - request_id: Request identifier - """ - with self._lock: - entry = self._pending_requests.pop(request_id, None) - - if entry is None: - logger.debug("Request already completed or unknown: %s", request_id) - return - - RequestClient._settle(entry.future.cancel) - logger.debug("Cancelled request: %s", request_id) - @property def pending_count(self) -> int: """Get number of currently pending requests. diff --git a/src/griptape_nodes/retained_mode/managers/worker_manager.py b/src/griptape_nodes/retained_mode/managers/worker_manager.py index 23beb08e26..bf954e4d4e 100644 --- a/src/griptape_nodes/retained_mode/managers/worker_manager.py +++ b/src/griptape_nodes/retained_mode/managers/worker_manager.py @@ -525,21 +525,35 @@ async def route_to_worker( request_id, tag=worker_engine_id, resolve_failures_as_payload=True ) - await self.forward_event_to_worker( - event_request.model_copy(update={"request_id": request_id}), - worker_engine_id=worker_engine_id, - worker_request_topic=worker_request_topic, - ) - # No wall-clock timeout here: long-running AI workloads (diffusion, multi-pass refinement) - # routinely exceed any sensible default. Worker liveness is enforced by the heartbeat loop, - # which evicts a silent worker and fails its in-flight requests with WorkerGoneError, so a - # dead worker still surfaces to the caller without a per-request ceiling. + # The publish is inside the `try` because it is an await like any other: a cancellation + # delivered while it is parked on send_message, or a raise out of the transport, would + # otherwise leave the entry tracked with nobody left to settle it. + # + # No wall-clock timeout on the response: long-running AI workloads (diffusion, multi-pass + # refinement) routinely exceed any sensible default. Worker liveness is enforced by the + # heartbeat loop, which evicts a silent worker and fails its in-flight requests with + # WorkerGoneError, so a dead worker still surfaces to the caller without a per-request + # ceiling. # # The future is settled by whichever loop the transport runs on, which is not this one. # wrap_future adapts it for this loop and installs the threadsafe wakeup. A CancelledError - # out of here means only one thing -- the caller was cancelled -- because a worker going - # away raises WorkerGoneError instead. - return await asyncio.wrap_future(future) + # out of the response await means only one thing -- the caller was cancelled -- because a + # worker going away raises WorkerGoneError instead. + try: + await self.forward_event_to_worker( + event_request.model_copy(update={"request_id": request_id}), + worker_engine_id=worker_engine_id, + worker_request_topic=worker_request_topic, + ) + return await asyncio.wrap_future(future) + except BaseException: + # BaseException, not Exception: cancellation is the common case and it is not an + # Exception. The other ways out remove the request themselves -- a response pops it in + # _try_match, eviction in fail_requests_by_tag -- so this is the one exit that has to + # clean up after itself. Without it the entry outlives the run, one per cancelled node + # execution, each still carrying its worker's tag for fail_requests_by_tag to walk. + self._tx.request_client.discard_request(request_id) + raise async def _orchestrator_static_server_base_url(self) -> str | None: """The base URL this engine serves the workspace on, awaited until initialization decides it. diff --git a/tests/unit/app/test_app_worker.py b/tests/unit/app/test_app_worker.py index c8182c78ad..12d585b1ce 100644 --- a/tests/unit/app/test_app_worker.py +++ b/tests/unit/app/test_app_worker.py @@ -64,6 +64,11 @@ async def fail_requests_by_tag(self, tag: str, error: Exception) -> None: if not entry.future.done(): entry.future.set_exception(error) + def discard_request(self, request_id: str) -> None: + entry = self._pending_requests.pop(request_id, None) + if entry is not None and not entry.future.done(): + entry.future.cancel() + @pytest.fixture def worker_manager() -> WorkerManager: @@ -1234,6 +1239,61 @@ async def test_flow_cancellation_still_cancels(self, worker_manager: WorkerManag with pytest.raises(asyncio.CancelledError): await task + @pytest.mark.asyncio + async def test_flow_cancellation_stops_tracking_the_request(self, worker_manager: WorkerManager) -> None: + """A cancelled run must not leave its request behind in the pending map. + + Nothing else pops it on this path, so the entry would outlive the run -- one per cancelled + node execution for the life of the process -- and each keeps its worker's tag, so a later + fail_requests_by_tag walks and re-settles long-dead requests. + """ + assert isinstance(worker_manager._tx.request_client, _FakeRequestClient) + fake_rc = worker_manager._tx.request_client + event_request = EventRequest(request=ExecuteNodeRequest(node_name="MyNode", parameter_values={})) + + task = asyncio.create_task(worker_manager.route_to_worker(event_request, _ENGINE, _WORKER_REQUEST_TOPIC)) + # Long enough to be parked on the response, which is the await the cancellation has to + # unwind from for the entry to be the caller's to remove. + await asyncio.sleep(0.01) + assert len(fake_rc._pending_requests) == 1 + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert fake_rc._pending_requests == {} + + @pytest.mark.asyncio + async def test_cancellation_during_the_publish_stops_tracking_the_request( + self, worker_manager: WorkerManager + ) -> None: + """The publish is an await too, so the entry has to be released if it unwinds there. + + Registering before publishing is deliberate -- the response can arrive the instant the + request lands -- which leaves a window where the entry exists and the publish has not + returned. A cancellation delivered in that window leaks exactly the entry the response-side + guard exists to release. + """ + assert isinstance(worker_manager._tx.request_client, _FakeRequestClient) + fake_rc = worker_manager._tx.request_client + event_request = EventRequest(request=ExecuteNodeRequest(node_name="MyNode", parameter_values={})) + publishing = asyncio.Event() + + async def _park_in_the_publish(*_args: object, **_kwargs: object) -> None: + publishing.set() + await asyncio.sleep(3600) + + with patch.object(worker_manager, "forward_event_to_worker", new=_park_in_the_publish): + task = asyncio.create_task(worker_manager.route_to_worker(event_request, _ENGINE, _WORKER_REQUEST_TOPIC)) + await publishing.wait() + assert len(fake_rc._pending_requests) == 1 + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert fake_rc._pending_requests == {} + class TestGetTopicsToSubscribe: def test_orchestrator_includes_base_request_topic(self, worker_manager: WorkerManager) -> None: From 778657a49e288066d369d98303654e67be13fc5d Mon Sep 17 00:00:00 2001 From: cjkindel Date: Tue, 15 Sep 2026 11:55:25 -0700 Subject: [PATCH 4/8] Read the pending request map under its lock `pending_count` and `pending_request_ids` read `_pending_requests` while holding nothing, and it is mutated from more than one loop: `_try_match` on the transport loop, `track_request` and `cancel_requests_by_tag` from whichever loop issued the request. `len()` only risks a stale answer, which is fine for a diagnostic. Building the id list iterates, so a concurrent insert or pop raises "dictionary changed size during iteration" -- a diagnostic that can raise is worse than one that answers a moment out of date. Both are diagnostics with no production caller today, so nothing holds the lock across a call into them. --- src/griptape_nodes/api_client/request_client.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/griptape_nodes/api_client/request_client.py b/src/griptape_nodes/api_client/request_client.py index dac3ec543e..43878d7e31 100644 --- a/src/griptape_nodes/api_client/request_client.py +++ b/src/griptape_nodes/api_client/request_client.py @@ -499,7 +499,8 @@ def pending_count(self) -> int: Returns: Count of pending requests """ - return len(self._pending_requests) + with self._lock: + return len(self._pending_requests) @property def pending_request_ids(self) -> list[str]: @@ -508,7 +509,11 @@ def pending_request_ids(self) -> list[str]: Returns: List of request_id strings """ - return list(self._pending_requests.keys()) + # Building the list iterates the map, so an insert or pop from another loop raises + # "dictionary changed size during iteration". A diagnostic that can raise is worse than + # one that answers a moment out of date. + with self._lock: + return list(self._pending_requests.keys()) async def _try_match(self, message: dict[str, Any]) -> bool: """Attempt to match an incoming message to a pending request. From 73d8b245792004e37f8f67031bb648e174c08a22 Mon Sep 17 00:00:00 2001 From: cjkindel Date: Tue, 15 Sep 2026 11:57:29 -0700 Subject: [PATCH 5/8] Cover the spawn side of the static-URL handover `_orchestrator_static_server_base_url` picks between a direct read and a thread hop, and between two warnings when no URL arrives. Nothing exercised either choice: inverting the settled check left the suite green, and so did swapping the two warnings. Both choices matter to whoever has to act on them. Taking the hop when the answer is already in parks an uncancellable default-executor thread that teardown joins, and blaming the settle timeout for a resolution that raised points an operator at slow startup instead of the bind failure that actually happened. Five tests: the direct read, the hop with its timeout, each warning, and a cloud backend staying quiet because there its URLs outlive the worker anyway. Each fails if the branch it covers is inverted. --- tests/unit/app/test_app_worker.py | 103 +++++++++++++++++++++++++++++- 1 file changed, 102 insertions(+), 1 deletion(-) diff --git a/tests/unit/app/test_app_worker.py b/tests/unit/app/test_app_worker.py index 12d585b1ce..fb2974d34b 100644 --- a/tests/unit/app/test_app_worker.py +++ b/tests/unit/app/test_app_worker.py @@ -10,6 +10,7 @@ import asyncio import concurrent.futures import json +import logging import os import sys import threading @@ -21,6 +22,7 @@ import pytest from griptape_nodes.api_client.request_client import _PendingRequest +from griptape_nodes.drivers.storage.local_storage_driver import LocalStorageDriver from griptape_nodes.retained_mode.events import worker_events from griptape_nodes.retained_mode.events.app_events import CurrentProjectChanged from griptape_nodes.retained_mode.events.base_events import EventRequest @@ -29,7 +31,11 @@ ExecuteNodeResultSuccess, ) from griptape_nodes.retained_mode.managers.project_manager import SYSTEM_DEFAULTS_KEY -from griptape_nodes.retained_mode.managers.worker_manager import WorkerManager, WorkerRegistration +from griptape_nodes.retained_mode.managers.worker_manager import ( + _STATIC_URL_SETTLE_TIMEOUT_S, + WorkerManager, + WorkerRegistration, +) from griptape_nodes.utils.version_utils import engine_version _SESSION = "sess-abc" @@ -812,6 +818,101 @@ async def test_spawn_env_omits_orchestrator_id_when_unknown(self, worker_manager assert "GTN_ORCHESTRATOR_ENGINE_ID" not in env +class TestOrchestratorStaticServerBaseUrl: + """The spawn side of the static-URL handover. + + Resolution and spawn are both listeners on AppInitializationComplete and fan out as unordered + concurrent tasks, so the URL is awaited rather than sampled. Which wait runs, and which of the + two warnings a missing URL earns, are what an operator reads when a worker's asset URLs come + back dead, so both choices are pinned here. + """ + + @pytest.mark.asyncio + async def test_a_settled_url_is_read_without_a_thread_hop(self, worker_manager: WorkerManager) -> None: + """The decision is normally already in, and the hop is the path with a cost. + + A blocking wait handed to a thread cannot be cancelled, so taking it when nothing needs it + parks a default-executor thread that teardown then joins. + """ + static_files_manager = worker_manager.engine.static_files_manager # type: ignore[union-attr] + static_files_manager.static_server_base_url_settled = True + static_files_manager.wait_for_static_server_base_url.return_value = "http://orchestrator:4242" + + with patch("asyncio.to_thread", new=AsyncMock()) as mock_to_thread: + result = await worker_manager._orchestrator_static_server_base_url() + + assert result == "http://orchestrator:4242" + mock_to_thread.assert_not_called() + static_files_manager.wait_for_static_server_base_url.assert_called_once_with(0) + + @pytest.mark.asyncio + async def test_an_undecided_url_is_waited_for_off_the_loop(self, worker_manager: WorkerManager) -> None: + """The blocking wait must not run on the event loop, which is serving everything else.""" + static_files_manager = worker_manager.engine.static_files_manager # type: ignore[union-attr] + static_files_manager.static_server_base_url_settled = False + static_files_manager.wait_for_static_server_base_url.return_value = "http://orchestrator:4242" + + result = await worker_manager._orchestrator_static_server_base_url() + + assert result == "http://orchestrator:4242" + static_files_manager.wait_for_static_server_base_url.assert_called_once_with(_STATIC_URL_SETTLE_TIMEOUT_S) + + @pytest.mark.asyncio + async def test_a_settled_absence_blames_resolution_rather_than_the_wait( + self, worker_manager: WorkerManager, caplog: pytest.LogCaptureFixture + ) -> None: + """Initialization deciding there is no server settles in microseconds. + + Blaming the settle timeout for it points an operator at slow startup when the real lead is + an earlier resolution failure, which under local storage is the only way to reach here. + """ + static_files_manager = worker_manager.engine.static_files_manager # type: ignore[union-attr] + static_files_manager.static_server_base_url_settled = True + static_files_manager.wait_for_static_server_base_url.return_value = None + static_files_manager.storage_driver = MagicMock(spec=LocalStorageDriver) + + with caplog.at_level(logging.WARNING): + result = await worker_manager._orchestrator_static_server_base_url() + + assert result is None + assert "Check for an earlier failure resolving the static server" in caplog.text + + @pytest.mark.asyncio + async def test_a_url_that_never_arrives_blames_the_wait( + self, worker_manager: WorkerManager, caplog: pytest.LogCaptureFixture + ) -> None: + static_files_manager = worker_manager.engine.static_files_manager # type: ignore[union-attr] + static_files_manager.static_server_base_url_settled = False + static_files_manager.wait_for_static_server_base_url.return_value = None + static_files_manager.storage_driver = MagicMock(spec=LocalStorageDriver) + + with caplog.at_level(logging.WARNING): + result = await worker_manager._orchestrator_static_server_base_url() + + assert result is None + assert "never decided" in caplog.text + + @pytest.mark.asyncio + async def test_a_cloud_backend_without_a_url_says_nothing( + self, worker_manager: WorkerManager, caplog: pytest.LogCaptureFixture + ) -> None: + """On a cloud backend a worker's URLs come from the same bucket and outlive it. + + There is nothing to warn about, and warning anyway trains people to ignore the case where + the URLs really do die with the worker. + """ + static_files_manager = worker_manager.engine.static_files_manager # type: ignore[union-attr] + static_files_manager.static_server_base_url_settled = True + static_files_manager.wait_for_static_server_base_url.return_value = None + static_files_manager.storage_driver = MagicMock() + + with caplog.at_level(logging.WARNING): + result = await worker_manager._orchestrator_static_server_base_url() + + assert result is None + assert caplog.records == [] + + class TestResetWorkers: @pytest.mark.asyncio async def test_terminates_all_processes(self, worker_manager: WorkerManager) -> None: From 6b2ab209726500c0438b8fe12c8acc17e5f27557 Mon Sep 17 00:00:00 2001 From: cjkindel Date: Tue, 15 Sep 2026 12:10:46 -0700 Subject: [PATCH 6/8] Type the mocked static-files manager in the worker tests The static-URL tests set `static_server_base_url_settled`, which is a property on the real manager, so pyright rejected the assignment even though the object is a MagicMock. Casting once per test says what the object actually is instead of suppressing each line. --- tests/unit/app/test_app_worker.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/tests/unit/app/test_app_worker.py b/tests/unit/app/test_app_worker.py index fb2974d34b..81a5b1edb3 100644 --- a/tests/unit/app/test_app_worker.py +++ b/tests/unit/app/test_app_worker.py @@ -737,9 +737,11 @@ async def test_a_failed_fork_does_not_keep_the_key_claimed(self, worker_manager: """ worker_manager.engine.library_manager.execution_site_packages.return_value = None # type: ignore[union-attr] - with patch("asyncio.create_subprocess_exec", side_effect=OSError("no interpreter")): - with pytest.raises(OSError, match="no interpreter"): - await worker_manager.spawn_worker(["/usr/bin/gtn", "engine"], "My Library") + with ( + patch("asyncio.create_subprocess_exec", side_effect=OSError("no interpreter")), + pytest.raises(OSError, match="no interpreter"), + ): + await worker_manager.spawn_worker(["/usr/bin/gtn", "engine"], "My Library") assert "My Library" not in worker_manager._spawns_in_flight @@ -834,7 +836,7 @@ async def test_a_settled_url_is_read_without_a_thread_hop(self, worker_manager: A blocking wait handed to a thread cannot be cancelled, so taking it when nothing needs it parks a default-executor thread that teardown then joins. """ - static_files_manager = worker_manager.engine.static_files_manager # type: ignore[union-attr] + static_files_manager = cast("MagicMock", worker_manager.engine.static_files_manager) static_files_manager.static_server_base_url_settled = True static_files_manager.wait_for_static_server_base_url.return_value = "http://orchestrator:4242" @@ -848,7 +850,7 @@ async def test_a_settled_url_is_read_without_a_thread_hop(self, worker_manager: @pytest.mark.asyncio async def test_an_undecided_url_is_waited_for_off_the_loop(self, worker_manager: WorkerManager) -> None: """The blocking wait must not run on the event loop, which is serving everything else.""" - static_files_manager = worker_manager.engine.static_files_manager # type: ignore[union-attr] + static_files_manager = cast("MagicMock", worker_manager.engine.static_files_manager) static_files_manager.static_server_base_url_settled = False static_files_manager.wait_for_static_server_base_url.return_value = "http://orchestrator:4242" @@ -866,7 +868,7 @@ async def test_a_settled_absence_blames_resolution_rather_than_the_wait( Blaming the settle timeout for it points an operator at slow startup when the real lead is an earlier resolution failure, which under local storage is the only way to reach here. """ - static_files_manager = worker_manager.engine.static_files_manager # type: ignore[union-attr] + static_files_manager = cast("MagicMock", worker_manager.engine.static_files_manager) static_files_manager.static_server_base_url_settled = True static_files_manager.wait_for_static_server_base_url.return_value = None static_files_manager.storage_driver = MagicMock(spec=LocalStorageDriver) @@ -881,7 +883,7 @@ async def test_a_settled_absence_blames_resolution_rather_than_the_wait( async def test_a_url_that_never_arrives_blames_the_wait( self, worker_manager: WorkerManager, caplog: pytest.LogCaptureFixture ) -> None: - static_files_manager = worker_manager.engine.static_files_manager # type: ignore[union-attr] + static_files_manager = cast("MagicMock", worker_manager.engine.static_files_manager) static_files_manager.static_server_base_url_settled = False static_files_manager.wait_for_static_server_base_url.return_value = None static_files_manager.storage_driver = MagicMock(spec=LocalStorageDriver) @@ -901,7 +903,7 @@ async def test_a_cloud_backend_without_a_url_says_nothing( There is nothing to warn about, and warning anyway trains people to ignore the case where the URLs really do die with the worker. """ - static_files_manager = worker_manager.engine.static_files_manager # type: ignore[union-attr] + static_files_manager = cast("MagicMock", worker_manager.engine.static_files_manager) static_files_manager.static_server_base_url_settled = True static_files_manager.wait_for_static_server_base_url.return_value = None static_files_manager.storage_driver = MagicMock() From fb204ae630de99346add0edac222df13397542d9 Mon Sep 17 00:00:00 2001 From: cjkindel Date: Thu, 17 Sep 2026 11:01:49 -0700 Subject: [PATCH 7/8] Apply review findings to the spawn and request-leak fixes A regression in the spawn claim: `reset_workers` cleared the process registry but not the new claim set, so a spawn suspended across a library reload kept its claim and the reload's own spawn for that library was refused. The refusal records nothing (a worker is normally on its way when a key is claimed) so the next run would wait out the whole startup grace and then blame the library load. Cleared with the registry, with a test that fails without it. `note_worker_unavailable` is deliberately still NOT called from the duplicate guard: it records a reason and settles the library, which would falsely fail a library that already has, or is about to have, a worker. Comments that were wrong or that answered a reviewer are gone rather than reworded again: a claim that `list(dict.keys())` can raise "dictionary changed size during iteration" (it cannot, the whole build is one C loop under the GIL) and a `discard_request` docstring citing a suspension point that never existed. `test_an_undecided_url_is_waited_for_off_the_loop` asserted its own name only loosely; it now records the executing thread and fails if the blocking wait runs on the event loop. --- .../api_client/request_client.py | 6 ---- .../retained_mode/managers/worker_manager.py | 3 ++ tests/unit/app/test_app_worker.py | 28 +++++++++++++++++++ 3 files changed, 31 insertions(+), 6 deletions(-) diff --git a/src/griptape_nodes/api_client/request_client.py b/src/griptape_nodes/api_client/request_client.py index 43878d7e31..d74dff038c 100644 --- a/src/griptape_nodes/api_client/request_client.py +++ b/src/griptape_nodes/api_client/request_client.py @@ -378,9 +378,6 @@ async def fail_requests_by_tag(self, tag: str, error: Exception) -> None: def discard_request(self, request_id: str) -> None: """Stop tracking a request and cancel its future. - Synchronous so a caller cancelled mid-await can call it from its own except block without - introducing a suspension point while the cancellation is being delivered. - Args: request_id: Request identifier """ @@ -509,9 +506,6 @@ def pending_request_ids(self) -> list[str]: Returns: List of request_id strings """ - # Building the list iterates the map, so an insert or pop from another loop raises - # "dictionary changed size during iteration". A diagnostic that can raise is worse than - # one that answers a moment out of date. with self._lock: return list(self._pending_requests.keys()) diff --git a/src/griptape_nodes/retained_mode/managers/worker_manager.py b/src/griptape_nodes/retained_mode/managers/worker_manager.py index bf954e4d4e..f8bfd6c22e 100644 --- a/src/griptape_nodes/retained_mode/managers/worker_manager.py +++ b/src/griptape_nodes/retained_mode/managers/worker_manager.py @@ -498,6 +498,9 @@ async def reset_workers(self) -> None: except Exception as e: logger.debug("Failed to unsubscribe from '%s' during reset: %s", response_topic, e) self._managed_worker_processes.clear() + # Cleared with the registry, not left behind: a spawn still inside its awaits holds this + # library's claim, and a claim surviving the reset refuses the reload's own spawn for it. + self._spawns_in_flight.clear() self._workers.clear() self._worker_last_seen.clear() diff --git a/tests/unit/app/test_app_worker.py b/tests/unit/app/test_app_worker.py index 81a5b1edb3..c67a9aa84e 100644 --- a/tests/unit/app/test_app_worker.py +++ b/tests/unit/app/test_app_worker.py @@ -853,10 +853,21 @@ async def test_an_undecided_url_is_waited_for_off_the_loop(self, worker_manager: static_files_manager = cast("MagicMock", worker_manager.engine.static_files_manager) static_files_manager.static_server_base_url_settled = False static_files_manager.wait_for_static_server_base_url.return_value = "http://orchestrator:4242" + waiting_thread: list[str] = [] + + def _record_thread(timeout_s: float) -> str: # noqa: ARG001 + waiting_thread.append(threading.current_thread().name) + return "http://orchestrator:4242" + + static_files_manager.wait_for_static_server_base_url.side_effect = _record_thread result = await worker_manager._orchestrator_static_server_base_url() assert result == "http://orchestrator:4242" + # The wait blocks whichever thread runs it, so running it here would stall every other + # spawn and every request this loop is serving for the whole settle timeout. + assert len(waiting_thread) == 1 + assert waiting_thread[0] != threading.current_thread().name static_files_manager.wait_for_static_server_base_url.assert_called_once_with(_STATIC_URL_SETTLE_TIMEOUT_S) @pytest.mark.asyncio @@ -916,6 +927,23 @@ async def test_a_cloud_backend_without_a_url_says_nothing( class TestResetWorkers: + @pytest.mark.asyncio + async def test_a_claim_does_not_outlive_the_reset(self, worker_manager: WorkerManager) -> None: + """A reload resets and then spawns again, so a surviving claim would refuse its own spawn. + + The refusal records nothing -- a worker is normally on its way when a key is claimed -- so + the next run would wait out the whole startup grace and then blame the library load. + """ + worker_manager.engine.library_manager.execution_site_packages.return_value = None # type: ignore[union-attr] + worker_manager._spawns_in_flight.add("My Library") + + await worker_manager.reset_workers() + + with patch("asyncio.create_subprocess_exec", return_value=_managed_proc_mock()) as mock_exec: + await worker_manager.spawn_worker(["/usr/bin/gtn", "engine"], "My Library") + + mock_exec.assert_called_once() + @pytest.mark.asyncio async def test_terminates_all_processes(self, worker_manager: WorkerManager) -> None: proc_a, proc_b = _managed_proc_mock(), _managed_proc_mock() From 7047794abadaab9cf215e3ca1f2d7a1bf7694591 Mon Sep 17 00:00:00 2001 From: cjkindel Date: Wed, 16 Sep 2026 11:29:25 -0700 Subject: [PATCH 8/8] Release a spawn claim only to the attempt that took it `reset_workers` drops the claims so a reload can spawn its libraries again, which leaves a spawn suspended between its claim and the registry write resuming to find the key claimed by the reload's spawn. Releasing by name alone freed that one, so the library was admitted again while a spawn was genuinely in flight -- and a second fork leaves one of the two processes untracked, which is the leak the claim exists to prevent. Claims are now per attempt: the map holds a token identifying the holder, and the release is conditional on the stored token still being the one this attempt took. The test drives the interleaving: park a spawn, reset, let the reload claim the key, then release the stale spawn with its fork failing so no registry entry shadows a wrongly-freed claim. It fails against an unconditional release. The comment on the claim also drops its comparison with the guard this replaced and states the constraint instead: the registry entry is written only once the subprocess exists, and the work in between suspends. --- .../retained_mode/managers/worker_manager.py | 28 +++++++----- tests/unit/app/test_app_worker.py | 44 ++++++++++++++++++- 2 files changed, 59 insertions(+), 13 deletions(-) diff --git a/src/griptape_nodes/retained_mode/managers/worker_manager.py b/src/griptape_nodes/retained_mode/managers/worker_manager.py index f8bfd6c22e..3711a33393 100644 --- a/src/griptape_nodes/retained_mode/managers/worker_manager.py +++ b/src/griptape_nodes/retained_mode/managers/worker_manager.py @@ -122,10 +122,10 @@ def __init__( # Subprocesses spawned by this orchestrator (library_name → process) self._managed_worker_processes: dict[str, asyncio.subprocess.Process] = {} - # Worker keys whose spawn has been claimed but has not yet reached the registry above. - # The registry entry cannot serve as the claim: it is only written once the subprocess - # exists, and the work in between suspends. - self._spawns_in_flight: set[str] = set() + # Worker keys whose spawn has been claimed but has not yet reached the registry above, + # each mapped to a token identifying the attempt that holds it. Keyed by attempt rather + # than by name alone so a spawn can only ever release its own claim. + self._spawns_in_flight: dict[str, object] = {} # The event loop that spawned the worker subprocesses. asyncio.subprocess.Process # binds its exit Future to its creating loop, so proc.wait() is only legal on this @@ -376,11 +376,12 @@ async def spawn_worker(self, args: list[str], worker_key: str) -> None: if worker_key in self._managed_worker_processes or worker_key in self._spawns_in_flight: logger.error("Worker for key '%s' already spawned; refusing duplicate spawn.", worker_key) return - # Claimed here, in the same step as the check above, so no await separates them. Two - # spawns for one key would otherwise both get past a registry-only guard and both fork, - # and only one can be recorded -- leaving the other's process untracked, holding its - # library's dependencies in memory until its own heartbeat lapses. - self._spawns_in_flight.add(worker_key) + # Claimed in the same step as the check above, so no await separates them. The registry + # entry cannot serve as this guard: it is written only once the subprocess exists, and the + # work in between suspends. A second fork for one library leaves one of the two processes + # untracked, holding that library's dependencies until its own heartbeat lapses. + claim = object() + self._spawns_in_flight[worker_key] = claim try: # Spawn with the orchestrator's PRE-project environ so the worker boots with the # same clean env baseline a fresh engine would have. Inheriting the live os.environ @@ -449,9 +450,12 @@ async def spawn_worker(self, args: list[str], worker_key: str) -> None: self._spawn_loop = asyncio.get_running_loop() self._managed_worker_processes[worker_key] = proc finally: - # Released even when the fork raises: the claim outliving a failed spawn would - # silently refuse every later attempt for this library. - self._spawns_in_flight.discard(worker_key) + # Released even when the fork raises, or the claim would silently refuse every later + # attempt for this library -- but only while this attempt still holds it. A reset drops + # the claims so a reload can spawn again, so a spawn suspended across one resumes to + # find the key belonging to the reload's spawn, and freeing that admits a third fork. + if self._spawns_in_flight.get(worker_key) is claim: + del self._spawns_in_flight[worker_key] logger.info("Spawned worker for key '%s' (pid %s)", worker_key, proc.pid) async def reset_workers(self) -> None: diff --git a/tests/unit/app/test_app_worker.py b/tests/unit/app/test_app_worker.py index c67a9aa84e..6db8fd92b0 100644 --- a/tests/unit/app/test_app_worker.py +++ b/tests/unit/app/test_app_worker.py @@ -750,6 +750,48 @@ async def test_a_failed_fork_does_not_keep_the_key_claimed(self, worker_manager: mock_exec.assert_called_once() + @pytest.mark.asyncio + async def test_a_spawn_outlived_by_a_reset_does_not_free_the_next_claim( + self, worker_manager: WorkerManager + ) -> None: + """A spawn only releases a claim it still holds. + + A reset drops the claims so a reload can spawn again, which leaves a spawn suspended across + it resuming to find the key claimed by the reload's spawn. Releasing by name alone frees + that one, and the library is admitted for a third fork while a spawn is genuinely in flight. + """ + worker_manager.engine.library_manager.execution_site_packages.return_value = None # type: ignore[union-attr] + released = asyncio.Event() + + async def _park_until_released() -> None: + await released.wait() + + # The stale spawn, parked mid-flight between its claim and the registry write. Its fork + # fails, so no registry entry is left behind to shadow a wrongly-freed claim. + with ( + patch.object(worker_manager, "_orchestrator_static_server_base_url", _park_until_released), + patch("asyncio.create_subprocess_exec", side_effect=OSError("stale spawn died")), + ): + stale = asyncio.create_task(worker_manager.spawn_worker(["/usr/bin/gtn", "engine"], "My Library")) + await asyncio.sleep(0.01) + assert "My Library" in worker_manager._spawns_in_flight + + await worker_manager.reset_workers() + reload_claim = object() + worker_manager._spawns_in_flight["My Library"] = reload_claim + + released.set() + with pytest.raises(OSError, match="stale spawn died"): + await stale + + # The stale spawn has finished and must have left the reload's claim standing. + assert worker_manager._spawns_in_flight.get("My Library") is reload_claim + + with patch("asyncio.create_subprocess_exec") as mock_exec: + await worker_manager.spawn_worker(["/usr/bin/gtn", "engine"], "My Library") + + mock_exec.assert_not_called() + @pytest.mark.asyncio async def test_spawns_subprocess_with_provided_args(self, worker_manager: WorkerManager) -> None: mock_proc = MagicMock() @@ -935,7 +977,7 @@ async def test_a_claim_does_not_outlive_the_reset(self, worker_manager: WorkerMa the next run would wait out the whole startup grace and then blame the library load. """ worker_manager.engine.library_manager.execution_site_packages.return_value = None # type: ignore[union-attr] - worker_manager._spawns_in_flight.add("My Library") + worker_manager._spawns_in_flight["My Library"] = object() await worker_manager.reset_workers()