Merge pull request 'Release the GPU after every stage' (#11) from gpu-unload into main

Reviewed-on: #11
Reviewed-by: branderson <brad@branderson.io>

README.md

@@ -108,6 +108,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)