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/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()