fix(separate): tear the worker down on any failure, and honour cancel before the CPU retry

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})"

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"]