diff --git a/backend/fluksio/flow/runs.py b/backend/fluksio/flow/runs.py index 39fee35..3afeb2e 100644 --- a/backend/fluksio/flow/runs.py +++ b/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( diff --git a/backend/fluksio/flow/workers.py b/backend/fluksio/flow/workers.py index 06180bf..4abf517 100644 --- a/backend/fluksio/flow/workers.py +++ b/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: diff --git a/backend/tests/flow/test_workers.py b/backend/tests/flow/test_workers.py index a2be684..335c708 100644 --- a/backend/tests/flow/test_workers.py +++ b/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()