Retire the GPU workers when a run that held a card finishes

A worker that has run a jax node keeps holding the GPU after the run: XLA
preallocates most of the VRAM at import and never releases it, so the next
process OOMs on preallocation while a warm idle worker sits on the card.
Pools are kept warm on purpose — a library reads its environment at import,
so a warm worker cannot be re-told — but the end of a run is a point where
the memory should go back, and the environments carrying a GPU assignment
are exactly the pools that ran on one. Idle ones go now, busy ones when
they return.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019Hra4ndWMCLU5F3KjUuVAc

backend/fluksio/flow/runs.py

@@ -1190,6 +1190,7 @@ class RunService:
             self._finish(run_id, status, reason, result, duration)
             run.status = status
             self._publish(run, "run_finished")
+            self._release_cards(run)
             # Its values were only ever this run's; nothing reads them once it
             # has a result. On Redis the namespace would expire anyway.
             if state is not None and status != "error":
@@ -1198,6 +1199,29 @@ class RunService:
                 except Exception:
                     logger.warning("Could not clear state of run %s", run_id)
 
+    def _release_cards(self, run: Run) -> None:
+        """Hand a GPU run's device memory back when the run is over.
+
+        The accountant frees the card the moment the node returns, but the
+        worker that ran on it is kept warm and a library that preallocated
+        most of the VRAM never gives it up — so the next process to want the
+        card found it taken by one sitting idle.
+        """
+        # ponytail: retires every CUDA pool rather than the ones this run used,
+        # which needs no bookkeeping — a concurrent GPU run's busy worker only
+        # dies when it returns, which is when its own memory should go back
+        # anyway. What it costs is the warm worker of a *live* flow's GPU node.
+        # Track the pools per run if that ever matters.
+        if not (run.needs or {}).get("gpus"):
+            return
+        pool = getattr(self.controller, "workers", None)
+        if pool is None:
+            return
+        try:
+            pool.retire_gpu_children()
+        except Exception:
+            logger.warning("Could not retire the GPU workers of run %s", run.id)
+
     def _record_node(self, run_id: str, outcome: NodeOutcome) -> None:
         outputs = _cacheable(outcome)
         row = RunNode(

backend/fluksio/flow/workers.py

@@ -321,6 +321,24 @@ class PythonWorkerPool:
         for child in children:
             child.respawn_all()
 
+    def retire_gpu_children(self) -> None:
+        """Retire the pools holding a card, so the VRAM goes back.
+
+        A library like JAX takes most of the device when it imports and never
+        releases it, so a warm worker that has run one such node is a held
+        card — and warm is the point of a pool, so nothing retires it. At the
+        end of a run there is something to key on: the environments carrying a
+        GPU assignment are exactly the pools that ran on one.
+        """
+        with self._lock:
+            children = [
+                child
+                for key, child in self._children.items()
+                if any(name == "CUDA_VISIBLE_DEVICES" for name, _ in key)
+            ]
+        for child in children:
+            child.respawn_all()
+
     def _drain(self) -> list[_Worker | None]:
         slots = []
         while True:

backend/tests/flow/test_workers.py

@@ -685,6 +685,19 @@ def test_retiring_workers_reaches_the_children(pool):
     assert child._generation > before
 
 
+def test_retiring_the_cards_leaves_the_other_pools_warm(pool):
+    """A library that preallocated the card only gives it back by dying."""
+    card = pool.for_env({"CUDA_VISIBLE_DEVICES": "0"})
+    threads = pool.for_env({"OMP_NUM_THREADS": "2"})
+    before = (card._generation, threads._generation)
+
+    pool.retire_gpu_children()
+
+    assert card._generation > before[0]
+    # Nothing to hand back, so nothing pays a cold start for it.
+    assert threads._generation == before[1]
+
+
 def test_cancelling_reaches_a_node_running_in_a_child(pool):
     child = pool.for_env({"OMP_NUM_THREADS": "2"})
     started = threading.Event()