diff --git a/README.md b/README.md index f7c448e..7ef9404 100644 --- a/README.md +++ b/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 diff --git a/src/songtrace/models.py b/src/songtrace/models.py new file mode 100644 index 0000000..9ce5d30 --- /dev/null +++ b/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) diff --git a/src/songtrace/pipeline.py b/src/songtrace/pipeline.py index 0553388..53b856d 100644 --- a/src/songtrace/pipeline.py +++ b/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)) diff --git a/src/songtrace/stages/align.py b/src/songtrace/stages/align.py index 59b5662..6c44467 100644 --- a/src/songtrace/stages/align.py +++ b/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 diff --git a/src/songtrace/stages/pitch.py b/src/songtrace/stages/pitch.py index 5d639f6..b131cdd 100644 --- a/src/songtrace/stages/pitch.py +++ b/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 diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000..ccf5365 --- /dev/null +++ b/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] diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 19d85d0..e0053aa 100644 --- a/tests/test_pipeline.py +++ b/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)