diff --git a/src/griptape_nodes/api_client/request_client.py b/src/griptape_nodes/api_client/request_client.py index 72abf3fddb..d74dff038c 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,22 @@ 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. + + 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 +489,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. @@ -496,7 +496,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]: @@ -505,7 +506,8 @@ def pending_request_ids(self) -> list[str]: Returns: List of request_id strings """ - return list(self._pending_requests.keys()) + 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. diff --git a/src/griptape_nodes/retained_mode/managers/worker_manager.py b/src/griptape_nodes/retained_mode/managers/worker_manager.py index 3410987325..3711a33393 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, + # 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 # loop. Eviction can run on a different loop (the websocket-tasks loop), which is @@ -368,75 +373,89 @@ 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 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 + # 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, 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: @@ -483,6 +502,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() @@ -510,21 +532,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. @@ -794,6 +830,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 0d20a9d15a..6db8fd92b0 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" @@ -64,6 +70,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: @@ -690,6 +701,97 @@ 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")), + 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_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() @@ -760,7 +862,130 @@ 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 = 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" + + 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 = 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 + 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 = 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) + + 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 = 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) + + 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 = 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() + + 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_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["My Library"] = object() + + 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() @@ -1003,6 +1228,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: @@ -1153,6 +1412,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: