diff --git a/src/griptape_nodes/retained_mode/managers/library/workers.py b/src/griptape_nodes/retained_mode/managers/library/workers.py index 419ca5febe..7082a7e783 100644 --- a/src/griptape_nodes/retained_mode/managers/library/workers.py +++ b/src/griptape_nodes/retained_mode/managers/library/workers.py @@ -259,13 +259,27 @@ async def maybe_start_workers_for_existing_session(self) -> None: worker_manager.set_session_ready() await self._start_workers() - async def _start_workers(self) -> None: + async def start_worker_for_library(self, library_name: str) -> None: + """Start the worker for one library, if it executes in one, leaving every other library alone. + + Without an active session nothing starts here: `on_session_started` starts it with the rest. + """ + if self.engine.library_manager.is_worker or not self.engine.get_session_id(): + return + await self._start_workers(only_library_name=library_name) + + async def _start_workers(self, only_library_name: str | None = None) -> None: """Issue StartWorkerRequest for every library whose nodes execute in a worker. Asks WorkerManager to spawn a subprocess. Used on session start (both initial and subsequent) so that worker creation is always tied to an active session. + + Args: + only_library_name: Consider only this library, for a caller that stopped one worker. """ for library_info in self.engine.library_manager._library_file_path_to_info.values(): + if only_library_name is not None and library_info.library_name != only_library_name: + continue # `enabled` matters: executes_in_worker is set before the lifecycle is overwritten with # DISABLED, so without this a library the user turned off still got an idle worker # subprocess. diff --git a/src/griptape_nodes/retained_mode/managers/worker_manager.py b/src/griptape_nodes/retained_mode/managers/worker_manager.py index a85e02b667..b748511ec1 100644 --- a/src/griptape_nodes/retained_mode/managers/worker_manager.py +++ b/src/griptape_nodes/retained_mode/managers/worker_manager.py @@ -145,8 +145,9 @@ def __init__( # 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] = {} + # than by name alone so a spawn can only ever release its own claim. The claim is set once + # the attempt ends, so a caller that must not race the spawn can wait for it. + self._spawns_in_flight: dict[str, asyncio.Event] = {} # 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 @@ -452,7 +453,7 @@ async def spawn_worker(self, args: list[str], worker_key: str) -> None: # 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() + claim = asyncio.Event() self._spawns_in_flight[worker_key] = claim try: # Spawn with the orchestrator's PRE-project environ so the worker boots with the @@ -531,6 +532,7 @@ async def spawn_worker(self, args: list[str], worker_key: str) -> None: # 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] + claim.set() logger.debug("Spawned worker for key '%s' (pid %s)", worker_key, proc.pid) async def reset_workers(self) -> None: @@ -546,41 +548,74 @@ async def reset_workers(self) -> None: len(self._managed_worker_processes), list(self._managed_worker_processes.keys()), ) - await asyncio.gather( - *( - self._terminate_via_spawn_loop(library_name, proc) - for library_name, proc in list(self._managed_worker_processes.items()) - ) - ) + await self._stop_workers(library_names=None, action="reload") + # 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() + + async def stop_worker_for_library(self, library_name: str) -> None: + """Terminate the worker serving one library and forget it, leaving every other worker running. + + A spawn already under way is waited out first and its process stopped with the rest, so once + this returns no worker for the library is running or about to start. + """ + claim = self._spawns_in_flight.get(library_name) + if claim is not None: + await claim.wait() + await self._stop_workers(library_names={library_name}, action="reset") + self.note_worker_unavailable(library_name, "its worker process was stopped while the library is being reset.") + + async def _stop_workers(self, *, library_names: set[str] | None, action: str) -> None: + """Terminate the named libraries' workers (every worker when None) and forget them. + + Anything awaiting one of them is failed with the reason, and its response topic is + unsubscribed. + """ + processes = [ + (library_name, proc) + for library_name, proc in self._managed_worker_processes.items() + if library_names is None or library_name in library_names + ] + await asyncio.gather(*(self._terminate_via_spawn_loop(library_name, proc) for library_name, proc in processes)) + # Read after the termination, not before: a worker can register while its process is being + # stopped, and an id missed here stays routable to a dead process. + worker_ids = [ + wid + for wid, registration in self._workers.items() + if library_names is None or registration.worker_key in library_names + ] # Settle anything still awaiting one of these workers, BEFORE the registry is cleared. # Clearing it is what makes this the last chance: the heartbeat loop can only evict ids it # can still see, so after this nothing reaches these requests at all. route_to_worker has no # wall-clock ceiling, so a node dispatched into a worker this call terminates would await a # future that never settles. if self._transport is not None: - for wid in list(self._workers): + for wid in worker_ids: registration = self._workers[wid] await self._tx.request_client.fail_requests_by_tag( wid, worker_events.WorkerGoneError( - f"worker '{wid}' was shut down to reload library '{registration.worker_key}'" + f"worker '{wid}' was shut down to {action} library '{registration.worker_key}'" if registration.worker_key - else f"worker '{wid}' was shut down to reload libraries" + else f"worker '{wid}' was shut down to {action} libraries" ), ) session_id = self.engine.get_session_id() if session_id and self._transport is not None: - for wid in list(self._workers): + for wid in worker_ids: response_topic = f"sessions/{session_id}/workers/{wid}/response" try: await self._tx.unsubscribe_from_topic(response_topic) 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() + logger.debug("Failed to unsubscribe from '%s' while stopping workers: %s", response_topic, e) + if library_names is None: + self._managed_worker_processes.clear() + self._workers.clear() + return + for library_name, _proc in processes: + self._managed_worker_processes.pop(library_name, None) + for wid in worker_ids: + self._workers.pop(wid, None) async def route_to_worker( self, diff --git a/tests/unit/app/test_app_worker.py b/tests/unit/app/test_app_worker.py index 4258778633..46232ed71e 100644 --- a/tests/unit/app/test_app_worker.py +++ b/tests/unit/app/test_app_worker.py @@ -969,7 +969,7 @@ async def _park_until_released() -> None: assert "My Library" in worker_manager._spawns_in_flight await worker_manager.reset_workers() - reload_claim = object() + reload_claim = asyncio.Event() worker_manager._spawns_in_flight["My Library"] = reload_claim released.set() @@ -1169,7 +1169,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.environment.execution_site_packages.return_value = None # type: ignore[union-attr] - worker_manager._spawns_in_flight["My Library"] = object() + worker_manager._spawns_in_flight["My Library"] = asyncio.Event() await worker_manager.reset_workers() @@ -1251,6 +1251,124 @@ async def test_unsubscribes_response_topic_for_each_registered_worker(self, work assert f"sessions/{_SESSION}/workers/eng-2/response" in unsubscribed +class TestStopWorkerForLibrary: + """Stopping the worker of one library, as a library reset does before deleting its environment.""" + + @pytest.mark.asyncio + async def test_terminates_only_the_process_of_that_library(self, worker_manager: WorkerManager) -> None: + proc_a, proc_b = _managed_proc_mock(), _managed_proc_mock() + worker_manager._managed_worker_processes["Lib A"] = proc_a + worker_manager._managed_worker_processes["Lib B"] = proc_b + + await worker_manager.stop_worker_for_library("Lib A") + + proc_a.terminate.assert_called_once() + proc_b.terminate.assert_not_called() + assert list(worker_manager._managed_worker_processes) == ["Lib B"] + + @pytest.mark.asyncio + async def test_forgets_only_the_registration_of_that_library(self, worker_manager: WorkerManager) -> None: + worker_manager._workers["eng-a"] = WorkerRegistration( + request_topic=f"sessions/{_SESSION}/workers/eng-a/request", worker_key="Lib A" + ) + worker_manager._workers["eng-b"] = WorkerRegistration( + request_topic=f"sessions/{_SESSION}/workers/eng-b/request", worker_key="Lib B" + ) + + await worker_manager.stop_worker_for_library("Lib A") + + assert list(worker_manager._workers) == ["eng-b"] + unsubscribed = [call.args[0] for call in worker_manager._tx.unsubscribe_from_topic.call_args_list] # type: ignore[union-attr] + assert unsubscribed == [f"sessions/{_SESSION}/workers/eng-a/response"] + + @pytest.mark.asyncio + async def test_waits_out_a_spawn_in_flight_and_stops_what_it_started(self, worker_manager: WorkerManager) -> None: + """A spawn parked between its claim and the registry write must not leave a live worker behind. + + Without the wait, the stop would find nothing to terminate, and the spawn would then resume and + register a worker importing from the environment the caller is about to delete. + """ + worker_manager.engine.library_manager.environment.execution_site_packages.return_value = None # type: ignore[union-attr] + released = asyncio.Event() + + async def park_until_released() -> None: + await released.wait() + + proc = _managed_proc_mock() + with ( + patch.object(worker_manager, "_orchestrator_static_server_base_url", park_until_released), + patch("asyncio.create_subprocess_exec", return_value=proc), + ): + spawn = asyncio.create_task(worker_manager.spawn_worker(["/usr/bin/gtn", "engine"], "Lib A")) + await asyncio.sleep(0.01) + stop = asyncio.create_task(worker_manager.stop_worker_for_library("Lib A")) + await asyncio.sleep(0.01) + assert not stop.done() + + released.set() + await spawn + await stop + + proc.terminate.assert_called_once() + assert "Lib A" not in worker_manager._managed_worker_processes + assert "Lib A" not in worker_manager._spawns_in_flight + + @pytest.mark.asyncio + async def test_settles_requests_awaiting_that_worker(self, worker_manager: WorkerManager) -> None: + """A node running in the worker fails with the reason instead of waiting on a process that is gone.""" + assert isinstance(worker_manager._tx.request_client, _FakeRequestClient) + worker_manager._workers = {_ENGINE: WorkerRegistration(request_topic=_WORKER_REQUEST_TOPIC, worker_key="Lib")} + event_request = EventRequest(request=ExecuteNodeRequest(node_name="MyNode", parameter_values={})) + + async def stop_mid_flight() -> None: + await asyncio.sleep(0) + await worker_manager.stop_worker_for_library("Lib") + + asyncio.create_task(stop_mid_flight()) # noqa: RUF006 + + with pytest.raises(worker_events.WorkerGoneError, match="shut down to reset library 'Lib'"): + await worker_manager.route_to_worker(event_request, _ENGINE, _WORKER_REQUEST_TOPIC) + + @pytest.mark.asyncio + async def test_forgets_a_worker_that_registers_while_its_process_is_stopping( + self, worker_manager: WorkerManager + ) -> None: + """A booting worker can register during the termination await; it must not stay routable.""" + worker_manager._managed_worker_processes["Lib A"] = _managed_proc_mock() + + async def register_during_termination(_library_name: str, _proc: object) -> None: + worker_manager._workers["eng-late"] = WorkerRegistration( + request_topic=f"sessions/{_SESSION}/workers/eng-late/request", worker_key="Lib A" + ) + + with patch.object(worker_manager, "_terminate_via_spawn_loop", side_effect=register_during_termination): + await worker_manager.stop_worker_for_library("Lib A") + + assert worker_manager.get_worker_for_key("Lib A") is None + + @pytest.mark.asyncio + async def test_releases_anything_waiting_for_that_library_to_become_executable( + self, worker_manager: WorkerManager + ) -> None: + """A node waiting on the readiness gate learns why instead of waiting out the load timeout.""" + worker_manager.expect_worker("Lib A") + + await worker_manager.stop_worker_for_library("Lib A") + + assert worker_manager.has_settled("Lib A") + assert worker_manager.worker_unavailable_reason("Lib A") is not None + + @pytest.mark.asyncio + async def test_is_a_no_op_for_a_library_without_a_worker(self, worker_manager: WorkerManager) -> None: + proc = _managed_proc_mock() + worker_manager._managed_worker_processes["Lib B"] = proc + + await worker_manager.stop_worker_for_library("Lib A") + + proc.terminate.assert_not_called() + assert list(worker_manager._managed_worker_processes) == ["Lib B"] + + class TestTerminateViaSpawnLoop: """Cross-loop worker termination. diff --git a/tests/unit/retained_mode/managers/test_library_worker_config.py b/tests/unit/retained_mode/managers/test_library_worker_config.py index 1f4000bdcd..713801cd6e 100644 --- a/tests/unit/retained_mode/managers/test_library_worker_config.py +++ b/tests/unit/retained_mode/managers/test_library_worker_config.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from dataclasses import replace from typing import Any, cast from unittest.mock import AsyncMock, MagicMock @@ -503,6 +504,30 @@ async def test_a_library_with_met_requirements_spawns(self) -> None: # Nothing standing in the way, so a stale account of a previous attempt is cleared. assert manager._library_file_path_to_info["/some/path.json"].execution_unavailable_reason is None + @pytest.mark.asyncio + async def test_starting_the_worker_of_one_library_leaves_the_others_alone(self) -> None: + """A reset stops one library's worker and starts only that one again.""" + manager = self._manager(unmet=False) + other = replace(manager._library_file_path_to_info["/some/path.json"], library_name="Other Lib") + manager._library_file_path_to_info["/other/path.json"] = other + + await manager.workers._start_workers(only_library_name="Lib") + + ahandle_request = cast("MagicMock", manager._engine).ahandle_request + ahandle_request.assert_awaited_once() + assert ahandle_request.await_args.args[0].library_name == "Lib" + + @pytest.mark.asyncio + async def test_starting_one_worker_without_a_session_does_nothing(self) -> None: + """Without a session the spawn is left to on_session_started, like every other worker.""" + manager = self._manager(unmet=False) + engine = cast("MagicMock", manager._engine) + engine.get_session_id.return_value = None + + await manager.workers.start_worker_for_library("Lib") + + engine.ahandle_request.assert_not_awaited() + class TestLibraryDependencyResolution: """A dependency declaration names a REPO; the registry is keyed by library NAME.