feat(workers): stop and start the worker of a single library

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.

src/griptape_nodes/retained_mode/managers/worker_manager.py

@@ -143,8 +143,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
@@ -450,7 +451,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
@@ -529,6 +530,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:
@@ -544,41 +546,71 @@ 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")
+
+    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
+        ]
+        worker_ids = [
+            wid
+            for wid, registration in self._workers.items()
+            if library_names is None or registration.worker_key in library_names
+        ]
+        await asyncio.gather(*(self._terminate_via_spawn_loop(library_name, proc) for library_name, proc in processes))
         # 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,

tests/unit/app/test_app_worker.py

@@ -948,7 +948,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()
@@ -1148,7 +1148,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()
 
@@ -1230,6 +1230,95 @@ 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_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.
 

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,19 @@ 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"
+
 
 class TestLibraryDependencyResolution:
     """A dependency declaration names a REPO; the registry is keyed by library NAME.