Release the GPU after every stage

songtrace serve is long-lived but only wants the GPU in bursts, and until now
it kept everything it ever loaded. align cached wav2vec2 in a module dict that
was never cleared, torchcrepe caches its model on its own infer() function,
and even once the last reference to a model goes, PyTorch keeps the freed
blocks in its allocator rather than returning them to the driver. On a shared
box that is the whole GPU held against one idle process.

Stages now hand it back as soon as they are done with it. A stage that caches
a model registers how to drop it with @on_release, so models.py needs to know
nothing about the pipeline, and run_pipeline calls release() in a finally
after each stage — not once per job, because transcribe sits between separate
and align waiting on the Whisper server, which can be minutes of a GPU nothing
is using.

release() answers the bytes it handed back and the stage logs a line when that
is not zero, which is the only part of this visible from outside. It will not
initialise CUDA just to clean up after a run that never touched it.

The cost is reloading each model once per job rather than once per process,
seconds against stages that run for minutes. The CUDA context itself stays
until the process exits; releasing that would mean a subprocess per stage.

README.md

@@ -87,6 +87,30 @@ Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/)
 installed on the host; `docker-compose.yml` requests one GPU via the
 `nvidia` device driver.
 
+## Sharing the GPU
+
+`songtrace serve` is long-lived but only wants the GPU in bursts, so it hands
+it back after every stage rather than holding it for the life of the process.
+Three stages load a model — `separate` (Demucs), `align` (wav2vec2) and
+`pitch` (torchcrepe) — and each one's weights are dropped as soon as that
+stage finishes, followed by `torch.cuda.empty_cache()` to return the
+allocator's freed blocks to the driver instead of keeping them for reuse. This
+matters most between `separate` and `align`: `transcribe` sits between them
+waiting on your Whisper server, which can be minutes of a GPU nothing is
+using.
+
+Stages that gave memory back say so in the job log:
+
+```
+[separate] released 3704 MiB of GPU memory
+```
+
+The cost is reloading each model once per job instead of once per process,
+which is seconds against stages that run for minutes. What this does not
+release is the CUDA context itself, a few hundred MB that belongs to the
+process until it exits; freeing that would mean running each stage in a
+subprocess. If you need every last MB back, stop the container between runs.
+
 ## Bare-metal
 
 ```sh

src/songtrace/models.py

@@ -0,0 +1,47 @@
+"""Handing the GPU back between stages.
+
+The GPU this runs on is shared, and `songtrace serve` outlives every job it
+runs, so a stage that caches its weights holds VRAM until the process exits —
+including while a later stage is waiting on the Whisper server, which can be
+minutes of a GPU nothing is using. Two things have to be undone: the caches
+themselves, and PyTorch's allocator, which keeps freed blocks for reuse rather
+than returning them to the driver.
+
+A stage that caches a model registers how to drop it with @on_release, so this
+module needs to know nothing about the pipeline. Registration happens at import
+time, which run_pipeline has already done by the time it calls release().
+"""
+import gc
+import sys
+from typing import Callable
+
+_callbacks: list[Callable[[], None]] = []
+
+
+def on_release(fn: Callable[[], None]) -> Callable[[], None]:
+    """Register a callback that drops a cached model."""
+    _callbacks.append(fn)
+    return fn
+
+
+def release() -> int:
+    """Drop every cached model and return its VRAM to the driver.
+
+    Answers the bytes handed back, which is the only part of this a caller can
+    show someone. Safe whenever: a stage that never loaded anything has nothing
+    to drop, and the allocator is left alone unless this process actually
+    initialised CUDA.
+    """
+    for callback in _callbacks:
+        callback()
+    # weights reach the allocator's free list only once the last reference
+    # goes, and a module graph is cyclic enough to outlive refcounting
+    gc.collect()
+    torch = sys.modules.get("torch")
+    if torch is None or not torch.cuda.is_available():
+        return 0
+    if not torch.cuda.is_initialized():
+        return 0
+    before = torch.cuda.memory_reserved()
+    torch.cuda.empty_cache()
+    return max(before - torch.cuda.memory_reserved(), 0)

src/songtrace/pipeline.py

@@ -3,6 +3,7 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Callable
 
+from . import models
 from .config import Config
 
 class PipelineError(Exception):
@@ -74,6 +75,14 @@ def run_pipeline(inp: SongInput, cfg: Config, log=print,
             log(f"[{stage.name}] cached, skipping")
             continue
         log(f"[{stage.name}] running")
-        stage.run(ctx)
+        # per stage rather than per job: each model is used by one stage, so
+        # holding it past that only keeps VRAM from whatever else wants it
+        try:
+            stage.run(ctx)
+        finally:
+            # logged only when it actually gave something back, so the line is
+            # evidence rather than noise on a box where this matters
+            if freed := models.release():
+                log(f"[{stage.name}] released {freed / 2**20:.0f} MiB of GPU memory")
     return PipelineResult(output_dir=cfg.output_dir / song_slug(inp),
                           warnings=read_warnings(work))

src/songtrace/stages/align.py

@@ -2,11 +2,14 @@ import json
 from pathlib import Path
 
 from ..lyrics import alignment_token, normalize_lyrics
+from ..models import on_release
 from ..pipeline import Ctx, PipelineError, Stage
 from .transcribe import classify_lyrics_match, load_transcript, lyrics_similarity, vocal_sections
 
 SIMILARITY_THRESHOLD = 0.4
 
+# a stage aligns section by section, so the model is cached across the calls
+# within one run and dropped at the end of it
 _model_cache = {}
 
 
@@ -17,6 +20,11 @@ def _get_model(device: str):
     return _model_cache[device]
 
 
+@on_release
+def unload_model() -> None:
+    _model_cache.clear()
+
+
 def force_align_section(vocals_path: Path, start: float, end: float,
                         tokens: list[str], device: str) -> list[tuple[float, float, float]]:
     import torch

src/songtrace/stages/pitch.py

@@ -1,10 +1,23 @@
+import sys
 from pathlib import Path
 
 import numpy as np
 
+from ..models import on_release
 from ..pipeline import Ctx, Stage
 
 
+@on_release
+def unload_model() -> None:
+    """torchcrepe caches its model as attributes on its own infer function."""
+    torchcrepe = sys.modules.get("torchcrepe")
+    if torchcrepe is None:
+        return
+    for attr in ("model", "capacity"):
+        if hasattr(torchcrepe.infer, attr):
+            delattr(torchcrepe.infer, attr)
+
+
 def track_pitch(vocals_path: Path, model: str, device: str):
     import torchaudio
     import torchcrepe

tests/test_models.py

@@ -0,0 +1,129 @@
+import sys
+import types
+
+import pytest
+
+from songtrace import models
+from songtrace.stages import align, pitch
+
+
+@pytest.fixture
+def isolated_callbacks(monkeypatch):
+    """release() with only the callbacks a test registers."""
+    monkeypatch.setattr(models, "_callbacks", [])
+    return models._callbacks
+
+
+def fake_torch(available=True, initialised=True, emptied=None, reserved=None):
+    """A torch whose allocator reports `reserved` bytes, then 0 once emptied."""
+    state = {"reserved": 4 * 2 ** 20 if reserved is None else reserved}
+
+    def empty_cache():
+        emptied.append(True)
+        state["reserved"] = 0
+
+    mod = types.ModuleType("torch")
+    mod.cuda = types.SimpleNamespace(
+        is_available=lambda: available,
+        is_initialized=lambda: initialised,
+        memory_reserved=lambda: state["reserved"],
+        empty_cache=empty_cache)
+    return mod
+
+
+def test_release_without_torch_or_callbacks_does_nothing(monkeypatch,
+                                                         isolated_callbacks):
+    monkeypatch.delitem(sys.modules, "torch", raising=False)
+    assert models.release() == 0
+
+
+def test_release_runs_every_registered_callback(isolated_callbacks):
+    called = []
+    models.on_release(lambda: called.append("a"))
+    models.on_release(lambda: called.append("b"))
+    models.release()
+    assert called == ["a", "b"]
+
+
+def test_on_release_returns_the_function_so_it_works_as_a_decorator():
+    def drop():
+        pass
+    assert models.on_release(drop) is drop
+    models._callbacks.remove(drop)
+
+
+def test_release_empties_the_allocator_when_cuda_is_in_use(monkeypatch,
+                                                           isolated_callbacks):
+    emptied = []
+    monkeypatch.setitem(sys.modules, "torch", fake_torch(emptied=emptied))
+    assert models.release() == 4 * 2 ** 20
+    assert emptied == [True]
+
+
+def test_release_leaves_the_allocator_alone_without_a_gpu(monkeypatch,
+                                                          isolated_callbacks):
+    emptied = []
+    monkeypatch.setitem(sys.modules, "torch",
+                        fake_torch(available=False, emptied=emptied))
+    assert models.release() == 0
+    assert emptied == []
+
+
+def test_release_does_not_initialise_cuda_just_to_clean_up(monkeypatch,
+                                                           isolated_callbacks):
+    """A run that never touched the GPU must not create a context on the way out."""
+    emptied = []
+    monkeypatch.setitem(sys.modules, "torch",
+                        fake_torch(initialised=False, emptied=emptied))
+    assert models.release() == 0
+    assert emptied == []
+
+
+def test_align_drops_its_cached_model():
+    align._model_cache["cuda"] = object()
+    align.unload_model()
+    assert align._model_cache == {}
+
+
+def test_align_registers_its_unload():
+    assert align.unload_model in models._callbacks
+
+
+def test_pitch_drops_the_model_torchcrepe_cached(monkeypatch):
+    torchcrepe = types.ModuleType("torchcrepe")
+
+    def infer():
+        pass
+
+    infer.model = object()
+    infer.capacity = "full"
+    torchcrepe.infer = infer
+    monkeypatch.setitem(sys.modules, "torchcrepe", torchcrepe)
+    pitch.unload_model()
+    assert not hasattr(infer, "model")
+    assert not hasattr(infer, "capacity")
+
+
+def test_pitch_unload_is_fine_when_torchcrepe_was_never_imported(monkeypatch):
+    monkeypatch.delitem(sys.modules, "torchcrepe", raising=False)
+    pitch.unload_model()
+
+
+def test_pitch_unload_is_fine_when_no_model_was_cached(monkeypatch):
+    torchcrepe = types.ModuleType("torchcrepe")
+    torchcrepe.infer = lambda: None
+    monkeypatch.setitem(sys.modules, "torchcrepe", torchcrepe)
+    pitch.unload_model()
+
+
+def test_pitch_registers_its_unload():
+    assert pitch.unload_model in models._callbacks
+
+
+def test_release_reports_nothing_freed_when_the_allocator_was_already_empty(
+        monkeypatch, isolated_callbacks):
+    emptied = []
+    monkeypatch.setitem(sys.modules, "torch",
+                        fake_torch(emptied=emptied, reserved=0))
+    assert models.release() == 0
+    assert emptied == [True]

tests/test_pipeline.py

@@ -53,3 +53,49 @@ def test_read_warnings_is_empty_without_a_notes_file(tmp_path):
 def test_read_warnings_reads_the_notes_file(tmp_path):
     (tmp_path / "notes_sec.json").write_text(json.dumps({"warnings": ["w1", "w2"]}))
     assert read_warnings(tmp_path) == ["w1", "w2"]
+
+def test_gpu_is_released_after_every_stage(tmp_path, monkeypatch):
+    """Not just at the end: a later stage may wait minutes on the Whisper server."""
+    from songtrace import models
+    cfg, inp, calls = make_cfg(tmp_path), make_inp(), []
+    monkeypatch.setattr(models, "release", lambda: calls.append("release"))
+    stages = [fake_stage("s1", calls), fake_stage("s2", calls)]
+    run_pipeline(inp, cfg, log=lambda m: None, stages=stages)
+    assert calls == ["s1", "release", "s2", "release"]
+
+def test_gpu_is_released_when_a_stage_fails(tmp_path, monkeypatch):
+    import pytest
+    from songtrace import models
+    from songtrace.pipeline import PipelineError
+    cfg, inp, calls = make_cfg(tmp_path), make_inp(), []
+    monkeypatch.setattr(models, "release", lambda: calls.append("release"))
+    def boom(ctx):
+        raise PipelineError("boom")
+    stages = [Stage(name="s1", outputs=lambda ctx: [ctx.work / "s1.done"], run=boom)]
+    with pytest.raises(PipelineError):
+        run_pipeline(inp, cfg, log=lambda m: None, stages=stages)
+    assert calls == ["release"]
+
+def test_a_cached_stage_does_not_pay_for_a_release(tmp_path, monkeypatch):
+    from songtrace import models
+    cfg, inp, calls = make_cfg(tmp_path), make_inp(), []
+    stages = [fake_stage("s1", calls)]
+    run_pipeline(inp, cfg, log=lambda m: None, stages=stages)
+    calls.clear()
+    monkeypatch.setattr(models, "release", lambda: calls.append("release"))
+    run_pipeline(inp, cfg, log=lambda m: None, stages=stages)
+    assert calls == []
+
+def test_a_release_that_freed_memory_is_reported_in_the_log(tmp_path, monkeypatch):
+    from songtrace import models
+    cfg, inp, lines = make_cfg(tmp_path), make_inp(), []
+    monkeypatch.setattr(models, "release", lambda: 512 * 2 ** 20)
+    run_pipeline(inp, cfg, log=lines.append, stages=[fake_stage("s1", [])])
+    assert "[s1] released 512 MiB of GPU memory" in lines
+
+def test_a_release_that_freed_nothing_stays_out_of_the_log(tmp_path, monkeypatch):
+    from songtrace import models
+    cfg, inp, lines = make_cfg(tmp_path), make_inp(), []
+    monkeypatch.setattr(models, "release", lambda: 0)
+    run_pipeline(inp, cfg, log=lines.append, stages=[fake_stage("s1", [])])
+    assert not any("released" in line for line in lines)