From b20818e2497c858ed4e531dae901402662c4d9d3 Mon Sep 17 00:00:00 2001 From: Thales Pereira <31625914+thcp@users.noreply.github.com> Date: Mon, 31 Aug 2026 21:39:42 +0100 Subject: [PATCH] fix(separate): tear the worker down on any failure, and honour cancel before the CPU retry The persistent worker teardown sat after the finally block, not inside 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 -- propagated without it, so _worker still held the process and the next job's _get_worker() saw a matching device and a live poll() and reused a worker whose CUDA state followed an exception. ml-pipeline.md is explicit that any non-success must tear it down. A second path missed it entirely: the pipe check raises before the try, so a worker that came back without stdin/stderr stayed cached and would be handed to every subsequent job. Found while writing the test for the first one. separate() also had no cancel check between its two attempts. The rmtree of a multi-GB partial result takes seconds and nothing is registered for cancel during it, so a cancel landing there was invisible and the full CPU pass ran to completion -- 10+ minutes -- before JobCancelled was finally raised. The UI showed "Cancelling" throughout. _kill_worker now reaps after kill(). Without communicate() a worker wedged in an uninterruptible CUDA call becomes a zombie whose pipes close only when the Popen refcount happens to drop; vocal_split.py already pairs the two. An entry-point cancel check was tried and removed. It broke test_cancel_kills_worker_next_job_spawns_fresh, which verifies that a cancelled job tears the worker down so the next one spawns fresh -- a check that never spawns has no worker to tear down, and that test encodes the rule this change is meant to protect. The queue worker already gates dispatch on cancellation, so the fallback check covers the gap that was actually reported. Verified: reverting each half fails its test. Refs #514 --- app/pipeline/separate.py | 33 +++++-- tests/test_separate_worker_lifecycle.py | 116 ++++++++++++++++++++++++ 2 files changed, 142 insertions(+), 7 deletions(-) create mode 100644 tests/test_separate_worker_lifecycle.py 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"]