branderson
f187bfd6c0
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.
130 lines
4.0 KiB
Python
130 lines
4.0 KiB
Python
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]
|