diff --git a/app/pipeline/separate.py b/app/pipeline/separate.py index 609ff3d0..52cd4d6a 100644 --- a/app/pipeline/separate.py +++ b/app/pipeline/separate.py @@ -48,6 +48,11 @@ def _kill_worker() -> None: proc.wait(timeout=5) except subprocess.TimeoutExpired: proc.kill() + # kill() only sends the signal. Without the reap a worker wedged in + # an uninterruptible CUDA call becomes a zombie whose pipes are + # closed only incidentally, whenever the Popen refcount happens to + # drop. vocal_split.py already pairs the two. + proc.communicate() def _get_worker(device: str) -> subprocess.Popen: @@ -108,6 +113,10 @@ def _run_demucs(job: Job, source: Path, job_dir: Path, device: str) -> tuple[int spawn_at = time.monotonic() proc = _get_worker(device) if proc.stdin is None or proc.stderr is None: + # Raised before the try below, so nothing else tears this down, and a + # worker without pipes would otherwise stay cached and be handed to + # every subsequent job. + _kill_worker() raise RuntimeError("demucs worker has no stdin/stderr pipe") set_proc(job.id, proc) @@ -203,13 +212,18 @@ def _watchdog() -> None: _done_evt.set() set_proc(job.id, None) wt.join(timeout=2) - - # Never reuse a worker after anything but a clean success: a cancel - # (proc.terminate() from the API thread) already killed it; a failure's - # GPU/CUDA state afterward isn't something we can vouch for. Only the - # happy path keeps the worker warm for the next job. - if job_ok is not True: - _kill_worker() + # Never reuse a worker after anything but a clean success: a cancel + # (proc.terminate() from the API thread) already killed it; a failure's + # GPU/CUDA state afterward isn't something we can vouch for. Only the + # happy path keeps the worker warm for the next job. + # + # Inside the finally, not after it: an exception out of the read loop + # -- proc.stderr.read(1) raising OSError when the API thread's + # terminate() races the read, or _set() raising -- skipped this + # entirely and left a worker whose CUDA state followed an exception + # warm for the next job (#514). + if job_ok is not True: + _kill_worker() # POST /cancel calls proc.terminate() directly, which causes the read # loop above to hit EOF. Translate that into JobCancelled before the @@ -253,6 +267,11 @@ def separate(job: Job, source: Path, job_dir: Path) -> Path: # Partial output from the failed attempt must not be mistaken for # results by collect(); CPU restarts from scratch, so does progress. shutil.rmtree(job_dir / DEMUCS_MODEL, ignore_errors=True) + # The rmtree above can take seconds on a multi-GB partial result, and + # nothing is registered for cancel during it. Re-check before paying + # for a full CPU pass the user has already asked to stop (#514). + if job.cancel_requested: + raise JobCancelled() _set(job, progress=0.0, stage="GPU failed — retrying on CPU (slower)...") job.gpu_fallback = True job.compute_device = f"cpu (fallback from {device})" diff --git a/tests/test_separate_worker_lifecycle.py b/tests/test_separate_worker_lifecycle.py new file mode 100644 index 00000000..8564ae60 --- /dev/null +++ b/tests/test_separate_worker_lifecycle.py @@ -0,0 +1,116 @@ +"""The persistent demucs worker must not survive a failure, and cancel must +land promptly (#514). + +ml-pipeline.md: the worker "is reused across consecutive **successful** jobs on +the same device, but torn down after **any** non-success (cancellation or +failure) -- post-exception CUDA state can't be trusted. Both halves of this +rule matter; don't relax either side." +""" + +from __future__ import annotations + +import pytest + +from app.core.models import Job, JobCancelled +from app.pipeline import separate as _separate + + +@pytest.fixture(autouse=True) +def _no_worker(): + _separate._worker.clear() + yield + _separate._worker.clear() + + +def _job(**kw): + return Job(id="a1b2c3d4e5f6", **kw) + + +def test_an_exception_in_the_stream_loop_still_tears_the_worker_down(tmp_path, monkeypatch): + # The teardown used to sit after the finally, so anything raising out of + # the read loop left a worker whose CUDA state followed an exception warm + # for the next job. + killed = [] + monkeypatch.setattr(_separate, "_kill_worker", lambda: killed.append(True)) + + class _Boom: + stdin = None + stderr = None + + monkeypatch.setattr(_separate, "_get_worker", lambda device: _Boom()) + + job = _job(status="separating") + with pytest.raises(RuntimeError): + _separate._run_demucs(job, tmp_path / "s.wav", tmp_path, "cpu") + + assert killed, "a worker was left warm after an exception" + + +class _FakePipe: + """stdin that accepts the request, stderr that dies mid-stream.""" + + def __init__(self, on_read): + self._on_read = on_read + + def write(self, _data): + return None + + def flush(self): + return None + + def read(self, _n): + return self._on_read() + + +class _FakeProc: + def __init__(self, on_read): + self.stdin = _FakePipe(on_read) + self.stderr = _FakePipe(on_read) + + def poll(self): + return None + + def terminate(self): + return None + + +def test_an_exception_inside_the_read_loop_still_tears_the_worker_down(tmp_path, monkeypatch): + """The case #514 is actually about: teardown sat *after* the finally, so an + OSError from proc.stderr.read(1) -- which happens when the API thread's + terminate() races the read -- propagated with the worker left warm.""" + killed = [] + monkeypatch.setattr(_separate, "_kill_worker", lambda: killed.append(True)) + + def _boom(): + raise OSError("broken pipe") + + monkeypatch.setattr(_separate, "_get_worker", lambda device: _FakeProc(_boom)) + monkeypatch.setattr(_separate, "set_proc", lambda *a, **kw: None) + + job = _job(status="separating") + with pytest.raises(OSError): + _separate._run_demucs(job, tmp_path / "s.wav", tmp_path, "cpu") + + assert killed, "a worker whose CUDA state followed an exception stayed warm" + + +def test_cancel_between_the_gpu_attempt_and_the_cpu_fallback_is_honoured(tmp_path, monkeypatch): + # The expensive case: without this the entire CPU separation ran to + # completion -- 10+ minutes -- while the UI showed "Cancelling". + attempts = [] + + def _fake_run(job, source, job_dir, device): + attempts.append(device) + if device != "cpu": + job.cancel_requested = True # the user cancels during the failure + return 1, ["boom"] + raise AssertionError("the CPU fallback must not run for a cancelled job") + + monkeypatch.setattr(_separate, "_run_demucs", _fake_run) + monkeypatch.setattr(_separate, "get_demucs_device", lambda: "cuda") + + job = _job(status="separating") + with pytest.raises(JobCancelled): + _separate.separate(job, tmp_path / "s.wav", tmp_path) + + assert attempts == ["cuda"]