diff --git a/app/main.py b/app/main.py index a274d40..1bd4b08 100644 --- a/app/main.py +++ b/app/main.py @@ -176,6 +176,12 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: _background_tasks.add(wt) wt.add_done_callback(_background_tasks.discard) yield + # Tear down the persistent demucs worker (#309) so a clean shutdown never + # leaves it as an orphaned process -- it has no parent-death watchdog of + # its own, unlike the desktop backend itself. + from app.pipeline.separate import _kill_worker + + _kill_worker() # Phones hitting the self-hosted server URL get the mobile UI; everything diff --git a/app/pipeline/demucs_worker.py b/app/pipeline/demucs_worker.py new file mode 100644 index 0000000..cbd1799 --- /dev/null +++ b/app/pipeline/demucs_worker.py @@ -0,0 +1,110 @@ +"""Persistent demucs worker (#309). + +Run as its own process: `python -m app.pipeline.demucs_worker `. +Loads the model once, then serves jobs one at a time over stdin/stderr -- +eliminating the torch import + model load + CUDA kernel warmup cost that +dominated repeated per-job subprocess spawns on GPU (measured 35-42% of the +separate stage on an RTX 3080; see #288/#309). The parent (separate.py) +keeps this process alive across consecutive successful jobs on the same +device and only tears it down on cancel or a genuine failure. + +Protocol: + - Parent writes one JSON line to stdin per job: + {"source": "", "job_dir": "", "shifts": 1} + - Normal demucs progress (tqdm, \\r-delimited "NN%" lines) streams to + stderr exactly as it would from the demucs CLI -- apply_model(progress= + True) is the same call the CLI itself makes, unchanged from what + separate.py's reader already parsed from a one-shot subprocess. + - On completion the worker writes one more stderr line and, only on + failure, exits: + "@@DONE@@" -- job ok, worker keeps serving + "@@ERROR@@" -- job failed, worker exits(1) + A job failure always exits the worker rather than trying to keep + serving: after an exception mid-inference (OOM-adjacent or not), GPU + memory / CUDA context state for future jobs isn't something we can + vouch for, so the parent respawns fresh rather than risk reusing a + worker in an unknown state. This matches today's behavior, where any + failure already meant "process is dead, next attempt spawns fresh" -- + the reuse win only applies to the happy path. + - EOF on stdin (parent closed the pipe) ends the worker's loop cleanly. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +from app.core.config import DEMUCS_MODEL + + +def _run_one_job(model, device: str, req: dict) -> None: + from demucs.apply import apply_model + from demucs.audio import save_audio + from demucs.separate import load_track + + source = Path(req["source"]) + job_dir = Path(req["job_dir"]) + shifts = int(req.get("shifts", 1)) + + # Identical to demucs.separate.main()'s per-track body (same functions, + # same default split/overlap/segment/clip/bit-depth) -- we're not + # reimplementing the audio pipeline, just calling it repeatedly on an + # already-loaded model instead of once per fresh process. + wav = load_track(source, model.audio_channels, model.samplerate) + ref = wav.mean(0) + wav = wav - ref.mean() + wav = wav / ref.std() + sources = apply_model( + model, + wav[None], + device=device, + shifts=shifts, + split=True, + overlap=0.25, + progress=True, + num_workers=0, + segment=None, + )[0] + sources = sources * ref.std() + sources = sources + ref.mean() + + out_dir = job_dir / DEMUCS_MODEL / source.stem + out_dir.mkdir(parents=True, exist_ok=True) + for stem_tensor, name in zip(sources, model.sources, strict=True): + save_audio( + stem_tensor, + str(out_dir / f"{name}.wav"), + samplerate=model.samplerate, + clip="rescale", + bits_per_sample=16, + as_float=False, + ) + + +def main() -> None: + device = sys.argv[1] if len(sys.argv) > 1 else "cpu" + + from demucs.pretrained import get_model + + model = get_model(DEMUCS_MODEL) + model.eval() + model.cpu() + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + req = json.loads(line) + _run_one_job(model, device, req) + except Exception as e: + sys.stderr.write(f"@@ERROR@@{json.dumps(str(e))}\n") + sys.stderr.flush() + sys.exit(1) + sys.stderr.write("@@DONE@@\n") + sys.stderr.flush() + + +if __name__ == "__main__": + main() diff --git a/app/pipeline/separate.py b/app/pipeline/separate.py index b9566f2..ecdf630 100644 --- a/app/pipeline/separate.py +++ b/app/pipeline/separate.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import logging import os import re @@ -23,34 +24,41 @@ # GPU processing can be silent for minutes; 30 min covers legitimate pauses # while still catching genuine hangs (GPU deadlock, OOM stall, etc.). +# Persistent worker (#309): only one job ever runs at a time (_pipeline_lock +# in runner.py), so there is exactly one worker to track, not a pool. Reused +# across consecutive successful jobs on the same device; torn down on cancel, +# a device change, or any job failure -- see demucs_worker.py's docstring for +# why a failure always kills the worker rather than trying to keep serving. +_worker: dict[str, object] = {} -def _demucs_cmd(device: str, source: Path, job_dir: Path) -> list[str]: - """Build the demucs CLI invocation. Module-level seam so tests can swap - in a stub executable without touching the process-management machinery.""" - cmd = [ - sys.executable, - "-m", - "demucs", - "-n", - DEMUCS_MODEL, - "-d", - device, - ] - # "best" quality (Settings -> General): demucs re-runs separation on a - # randomly time-shifted copy of the input and averages the two passes -- - # measurably cleaner stems, ~2x the separation time. Applies on any - # device; read fresh per job like get_demucs_device() below. - if get_separation_quality() == "best": - cmd += ["--shifts", "2"] - cmd += ["-o", str(job_dir), str(source)] - return cmd +def _spawn_worker_cmd(device: str) -> list[str]: + """Build the persistent-worker invocation. Module-level seam so tests can + swap in a stub executable without touching the process-management + machinery (mirrors the old _demucs_cmd seam).""" + return [sys.executable, "-m", "app.pipeline.demucs_worker", device] -def _run_demucs(job: Job, source: Path, job_dir: Path, device: str) -> tuple[int, list[str]]: - """One demucs attempt on `device`: spawn, stream progress, watchdog stalls. - Returns (returncode, stderr_tail). Raises JobCancelled when the exit was - caused by POST /cancel. The retry policy lives in separate().""" +def _kill_worker() -> None: + proc = _worker.pop("proc", None) + _worker.pop("device", None) + if proc is not None and proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + + +def _get_worker(device: str) -> subprocess.Popen: + """Return a live worker bound to `device`, reusing the current one if it + already matches and is still alive, spawning fresh otherwise (first call, + a device change, or the previous worker died/was torn down).""" + proc = _worker.get("proc") + if proc is not None and _worker.get("device") == device and proc.poll() is None: + return proc + _kill_worker() + env = os.environ.copy() try: import certifi @@ -60,32 +68,60 @@ def _run_demucs(job: Job, source: Path, job_dir: Path, device: str) -> tuple[int except ModuleNotFoundError: pass - spawn_at = time.monotonic() proc = subprocess.Popen( - _demucs_cmd(device, source, job_dir), + _spawn_worker_cmd(device), + stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True, - bufsize=0, + bufsize=1, env=env, ) - if proc.stderr is None: - raise RuntimeError("demucs subprocess has no stderr pipe") + _worker["proc"] = proc + _worker["device"] = device + return proc + + +def _run_demucs(job: Job, source: Path, job_dir: Path, device: str) -> tuple[int, list[str]]: + """One demucs job dispatched to the persistent worker for `device`: + reuse-or-spawn, stream progress, watchdog stalls. + + Returns (returncode, stderr_tail). Raises JobCancelled when the exit was + caused by POST /cancel. The retry policy lives in separate().""" + spawn_at = time.monotonic() + proc = _get_worker(device) + if proc.stdin is None or proc.stderr is None: + raise RuntimeError("demucs worker has no stdin/stderr pipe") set_proc(job.id, proc) - # Time from spawn to demucs's first progress line -- process/model-load - # startup cost, as opposed to actual separation work (#288). Measurement - # only: subprocess isolation (kill-on-cancel, crash containment) is a - # design feature we keep regardless of what this turns out to be. + + shifts = 2 if get_separation_quality() == "best" else 1 + req = json.dumps({"source": str(source), "job_dir": str(job_dir), "shifts": shifts}) + "\n" + try: + proc.stdin.write(req) + proc.stdin.flush() + except (BrokenPipeError, OSError): + # The worker died between _get_worker() and here -- e.g. a cancel on + # the previous job raced this dispatch. Treat as an ordinary failure; + # separate()'s retry policy handles it exactly like a nonzero exit. + set_proc(job.id, None) + _kill_worker() + return 1, ["demucs worker is not accepting input (died before dispatch)"] + + # Time from dispatch to the first progress line -- near-zero for a reused + # warm worker, the full process/model-load cost for a freshly spawned one + # (#288/#309). Subprocess isolation (kill-on-cancel, crash containment) + # is a design feature we keep regardless of which case this run hits. startup_recorded = False # tqdm uses \r to redraw -- read char-by-char and split on \r or \n. - # Keep the last few non-progress lines so we can surface them if demucs - # exits non-zero (otherwise the only signal would be a bare exit code). + # Keep the last few non-progress lines so we can surface them if the job + # fails (otherwise the only signal would be a bare exit code). buf = "" tail: list[str] = [] last_output: list[float] = [time.monotonic()] - # Event set by the reader loop when the process exits normally so the - # watchdog can wake up immediately instead of waiting out its 30 s sleep. + job_ok: bool | None = None # None while streaming; True/False once decided + # Event set by the reader loop when the job finishes so the watchdog can + # wake up immediately instead of waiting out its 30 s sleep. _done_evt = threading.Event() def _watchdog() -> None: @@ -94,7 +130,7 @@ def _watchdog() -> None: return if time.monotonic() - last_output[0] > TIMEOUT_DEMUCS_STALL: logger.warning( - "demucs stalled for %ss with no output, terminating job %s", + "demucs worker stalled for %ss with no output, terminating job %s", TIMEOUT_DEMUCS_STALL, job.id, ) @@ -107,6 +143,9 @@ def _watchdog() -> None: while True: ch = proc.stderr.read(1) if not ch: + # EOF -- the worker process itself exited (crash, or it + # already wrote @@ERROR@@ and is shutting down). + job_ok = False break last_output[0] = time.monotonic() if ch in ("\r", "\n"): @@ -114,6 +153,18 @@ def _watchdog() -> None: buf = "" if not line: continue + if line == "@@DONE@@": + job_ok = True + break + if line.startswith("@@ERROR@@"): + msg = line[len("@@ERROR@@") :] + try: + msg = json.loads(msg) + except json.JSONDecodeError: + pass + tail.append(str(msg)) + job_ok = False + break m = _PCT_RE.search(line) if m: if not startup_recorded: @@ -131,19 +182,24 @@ def _watchdog() -> None: tail.pop(0) else: buf += ch - - proc.wait() finally: _done_evt.set() set_proc(job.id, None) wt.join(timeout=2) - # POST /cancel calls proc.terminate() directly, which causes the read loop - # above to hit EOF and proc.wait() to return a nonzero status. Translate - # that into JobCancelled before the generic "demucs failed" path. + # 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() + + # POST /cancel calls proc.terminate() directly, which causes the read + # loop above to hit EOF. Translate that into JobCancelled before the + # generic "demucs failed" path. if job.cancel_requested: raise JobCancelled() - return proc.returncode, tail + return (0, tail) if job_ok else (1, tail) def separate(job: Job, source: Path, job_dir: Path) -> Path: diff --git a/tests/test_separate_fallback.py b/tests/test_separate_fallback.py index 826809e..3200d70 100644 --- a/tests/test_separate_fallback.py +++ b/tests/test_separate_fallback.py @@ -1,8 +1,10 @@ -"""Tests for the GPU->CPU separation fallback (#276). +"""Tests for the persistent demucs worker (#309) and the GPU->CPU separation +fallback (#276). -The demucs invocation is swapped for stub Python one-liners via the -_demucs_cmd seam, so the real process machinery (Popen, stderr streaming, -watchdog, cancel translation) runs end-to-end without demucs or a GPU. +The worker invocation is swapped for stub Python scripts via the +_spawn_worker_cmd seam, so the real process machinery (Popen, stdin +dispatch, stderr streaming, watchdog, cancel translation, worker reuse) runs +end-to-end without demucs or a GPU. """ from __future__ import annotations @@ -16,31 +18,43 @@ from app.pipeline import separate as sep_mod from app.pipeline.errors import SeparationError +# A persistent worker stub: reads one JSON job request per line for as long +# as stdin stays open, always succeeding (writes a stem WAV where the real +# worker would, then "100%" + "@@DONE@@" to stderr) and keeps serving. +_SUCCESS_WORKER = """ +import sys, json, os +for line in sys.stdin: + req = json.loads(line) + d = os.path.join(req["job_dir"], "htdemucs_6s", "source") + os.makedirs(d, exist_ok=True) + open(os.path.join(d, "vocals.wav"), "wb").write(b"RIFF") + sys.stderr.write("100%\\n@@DONE@@\\n") + sys.stderr.flush() +""" + +# A worker stub that fails its first (only) dispatched job with a CUDA-OOM +# -shaped message, then exits -- matching demucs_worker.py's real behavior of +# never trying to keep serving after a failure. +_FAILING_WORKER = """ +import sys, json +sys.stdin.readline() +sys.stderr.write("@@ERROR@@" + json.dumps("CUDA out of memory. Tried 2 GiB") + "\\n") +sys.stderr.flush() +sys.exit(1) +""" -def _stub_cmds(fail_devices: set[str], calls: list[str]): - """A _demucs_cmd replacement: fails with a CUDA-OOM message on the given - devices, succeeds (writing a stem WAV where demucs would) elsewhere.""" - def fake_cmd(device: str, source: Path, job_dir: Path) -> list[str]: +def _stub_spawns(fail_devices: set[str], calls: list[str]): + """A _spawn_worker_cmd replacement. `calls` records one entry per SPAWNED + worker process (not per dispatched job) -- reuse across jobs on the same + device means fewer calls than jobs, which the reuse tests assert on.""" + + def fake_spawn(device: str) -> list[str]: calls.append(device) - if device in fail_devices: - code = ( - "import sys;" - " sys.stderr.write('RuntimeError: CUDA out of memory. Tried 2 GiB\\n');" - " sys.exit(1)" - ) - else: - out_dir = job_dir / sep_mod.DEMUCS_MODEL / source.stem - code = ( - "import os, sys;" - f" d = {str(out_dir)!r};" - " os.makedirs(d, exist_ok=True);" - " open(os.path.join(d, 'vocals.wav'), 'wb').write(b'RIFF');" - " sys.stderr.write('100%|separated\\n')" - ) + code = _FAILING_WORKER if device in fail_devices else _SUCCESS_WORKER return [sys.executable, "-c", code] - return fake_cmd + return fake_spawn @pytest.fixture() @@ -50,12 +64,21 @@ def job(tmp_path: Path): return j +@pytest.fixture(autouse=True) +def _reset_worker(): + """Each test starts and ends with no lingering worker reference -- a + prior test's stub process must never leak into the next test.""" + sep_mod._worker.clear() + yield + sep_mod._kill_worker() + + def test_gpu_failure_falls_back_to_cpu(job, tmp_path, monkeypatch, caplog): import logging calls: list[str] = [] monkeypatch.setattr(sep_mod, "get_demucs_device", lambda: "cuda") - monkeypatch.setattr(sep_mod, "_demucs_cmd", _stub_cmds({"cuda"}, calls)) + monkeypatch.setattr(sep_mod, "_spawn_worker_cmd", _stub_spawns({"cuda"}, calls)) with caplog.at_level(logging.WARNING, logger="stemdeck.pipeline"): stems_root = sep_mod.separate(job, tmp_path / "source.wav", tmp_path) @@ -70,26 +93,62 @@ def test_gpu_failure_falls_back_to_cpu(job, tmp_path, monkeypatch, caplog): assert "CUDA out of memory" in warning -def test_demucs_cmd_omits_shifts_at_standard_quality(monkeypatch, tmp_path): +def test_dispatch_omits_extra_shifts_at_standard_quality(monkeypatch, tmp_path): monkeypatch.setattr(sep_mod, "get_separation_quality", lambda: "standard") - cmd = sep_mod._demucs_cmd("cpu", tmp_path / "source.wav", tmp_path) - assert "--shifts" not in cmd + # Drive a real job and have the stub echo the dispatched request's + # "shifts" value back via a marker file -- simplest way to inspect what + # separate() actually sent without patching json.dumps at the call site. + echo_worker = """ +import sys, json, os +for line in sys.stdin: + req = json.loads(line) + d = os.path.join(req["job_dir"], "htdemucs_6s", "source") + os.makedirs(d, exist_ok=True) + open(os.path.join(d, "vocals.wav"), "wb").write(b"RIFF") + open(os.path.join(req["job_dir"], "shifts.txt"), "w").write(str(req["shifts"])) + sys.stderr.write("100%\\n@@DONE@@\\n") + sys.stderr.flush() +""" + monkeypatch.setattr(sep_mod, "get_demucs_device", lambda: "cpu") + monkeypatch.setattr( + sep_mod, "_spawn_worker_cmd", lambda device: [sys.executable, "-c", echo_worker] + ) + + sep_mod.separate(Job(id="abcdefabc277"), tmp_path / "source.wav", tmp_path) + + assert (tmp_path / "shifts.txt").read_text() == "1" + + +def test_dispatch_includes_shifts_2_at_best_quality(monkeypatch, tmp_path): + echo_worker = """ +import sys, json, os +for line in sys.stdin: + req = json.loads(line) + d = os.path.join(req["job_dir"], "htdemucs_6s", "source") + os.makedirs(d, exist_ok=True) + open(os.path.join(d, "vocals.wav"), "wb").write(b"RIFF") + open(os.path.join(req["job_dir"], "shifts.txt"), "w").write(str(req["shifts"])) + sys.stderr.write("100%\\n@@DONE@@\\n") + sys.stderr.flush() +""" + monkeypatch.setattr(sep_mod, "get_separation_quality", lambda: "best") + monkeypatch.setattr(sep_mod, "get_demucs_device", lambda: "cpu") + monkeypatch.setattr( + sep_mod, "_spawn_worker_cmd", lambda device: [sys.executable, "-c", echo_worker] + ) + sep_mod.separate(Job(id="abcdefabc278"), tmp_path / "source.wav", tmp_path) -def test_demucs_cmd_includes_shifts_2_at_best_quality(monkeypatch, tmp_path): - monkeypatch.setattr(sep_mod, "get_separation_quality", lambda: "best") - cmd = sep_mod._demucs_cmd("cpu", tmp_path / "source.wav", tmp_path) - i = cmd.index("--shifts") - assert cmd[i + 1] == "2" + assert (tmp_path / "shifts.txt").read_text() == "2" def test_records_startup_timing_on_first_progress_line(job, tmp_path, monkeypatch): - """#288: measurement only -- time from Popen to the first progress line - demucs emits, so the real subprocess/model-load startup cost can be - quantified before deciding whether it's worth a design change.""" + """#288/#309: time from dispatch to the first progress line demucs + emits -- the full spawn/model-load cost for a fresh worker, near-zero + for a reused warm one.""" calls: list[str] = [] monkeypatch.setattr(sep_mod, "get_demucs_device", lambda: "cpu") - monkeypatch.setattr(sep_mod, "_demucs_cmd", _stub_cmds(set(), calls)) + monkeypatch.setattr(sep_mod, "_spawn_worker_cmd", _stub_spawns(set(), calls)) sep_mod.separate(job, tmp_path / "source.wav", tmp_path) @@ -101,7 +160,7 @@ def test_records_startup_timing_on_first_progress_line(job, tmp_path, monkeypatc def test_gpu_success_needs_no_fallback(job, tmp_path, monkeypatch): calls: list[str] = [] monkeypatch.setattr(sep_mod, "get_demucs_device", lambda: "cuda") - monkeypatch.setattr(sep_mod, "_demucs_cmd", _stub_cmds(set(), calls)) + monkeypatch.setattr(sep_mod, "_spawn_worker_cmd", _stub_spawns(set(), calls)) stems_root = sep_mod.separate(job, tmp_path / "source.wav", tmp_path) @@ -114,7 +173,7 @@ def test_gpu_success_needs_no_fallback(job, tmp_path, monkeypatch): def test_cpu_failure_does_not_retry(job, tmp_path, monkeypatch): calls: list[str] = [] monkeypatch.setattr(sep_mod, "get_demucs_device", lambda: "cpu") - monkeypatch.setattr(sep_mod, "_demucs_cmd", _stub_cmds({"cpu"}, calls)) + monkeypatch.setattr(sep_mod, "_spawn_worker_cmd", _stub_spawns({"cpu"}, calls)) with pytest.raises(SeparationError) as exc_info: sep_mod.separate(job, tmp_path / "source.wav", tmp_path) @@ -127,7 +186,7 @@ def test_cpu_failure_does_not_retry(job, tmp_path, monkeypatch): def test_both_attempts_failing_raises_with_both_tails(job, tmp_path, monkeypatch): calls: list[str] = [] monkeypatch.setattr(sep_mod, "get_demucs_device", lambda: "mps") - monkeypatch.setattr(sep_mod, "_demucs_cmd", _stub_cmds({"mps", "cpu"}, calls)) + monkeypatch.setattr(sep_mod, "_spawn_worker_cmd", _stub_spawns({"mps", "cpu"}, calls)) with pytest.raises(SeparationError) as exc_info: sep_mod.separate(job, tmp_path / "source.wav", tmp_path) @@ -144,7 +203,7 @@ def test_both_attempts_failing_raises_with_both_tails(job, tmp_path, monkeypatch def test_cancel_during_gpu_attempt_skips_fallback(job, tmp_path, monkeypatch): calls: list[str] = [] monkeypatch.setattr(sep_mod, "get_demucs_device", lambda: "cuda") - monkeypatch.setattr(sep_mod, "_demucs_cmd", _stub_cmds({"cuda"}, calls)) + monkeypatch.setattr(sep_mod, "_spawn_worker_cmd", _stub_spawns({"cuda"}, calls)) job.cancel_requested = True # POST /cancel arrived before/mid attempt with pytest.raises(JobCancelled): @@ -157,31 +216,91 @@ def test_partial_gpu_output_cleared_before_retry(job, tmp_path, monkeypatch): """A failed GPU attempt's partial stems must not leak into the CPU run.""" calls: list[str] = [] marker = tmp_path / sep_mod.DEMUCS_MODEL / "partial-garbage.wav" + marker_repr = str(marker).replace("\\", "\\\\") - def fake_cmd(device: str, source: Path, job_dir: Path) -> list[str]: + def fake_spawn(device: str) -> list[str]: calls.append(device) if device == "cuda": - # Simulate demucs dying after writing partial output. + # Simulate the worker dying after writing partial output. code = ( - "import os, sys;" - f" os.makedirs(os.path.dirname({str(marker)!r}), exist_ok=True);" - f" open({str(marker)!r}, 'wb').write(b'junk');" - " sys.stderr.write('RuntimeError: CUDA error\\n');" - " sys.exit(1)" + "import os, sys, json\n" + "sys.stdin.readline()\n" + f"os.makedirs(os.path.dirname('{marker_repr}'), exist_ok=True)\n" + f"open('{marker_repr}', 'wb').write(b'junk')\n" + "sys.stderr.write('@@ERROR@@' + json.dumps('CUDA error') + chr(10))\n" + "sys.stderr.flush()\n" + "sys.exit(1)\n" ) - return [sys.executable, "-c", code] - out_dir = tmp_path / sep_mod.DEMUCS_MODEL / source.stem - code = ( - "import os;" - f" os.makedirs({str(out_dir)!r}, exist_ok=True);" - f" open(os.path.join({str(out_dir)!r}, 'vocals.wav'), 'wb').write(b'RIFF')" - ) + else: + code = _SUCCESS_WORKER return [sys.executable, "-c", code] monkeypatch.setattr(sep_mod, "get_demucs_device", lambda: "cuda") - monkeypatch.setattr(sep_mod, "_demucs_cmd", fake_cmd) + monkeypatch.setattr(sep_mod, "_spawn_worker_cmd", fake_spawn) sep_mod.separate(job, tmp_path / "source.wav", tmp_path) assert calls == ["cuda", "cpu"] assert not marker.exists(), "partial GPU output must be cleared before the CPU retry" + + +# ─── #309: worker reuse ──────────────────────────────────────────────────── + + +def test_worker_reused_across_consecutive_jobs_on_same_device(tmp_path, monkeypatch): + calls: list[str] = [] + monkeypatch.setattr(sep_mod, "get_demucs_device", lambda: "cpu") + monkeypatch.setattr(sep_mod, "_spawn_worker_cmd", _stub_spawns(set(), calls)) + + for i in range(3): + job = Job(id=f"abcdefabc30{i}") + (tmp_path / "source.wav").write_bytes(b"RIFF") + sep_mod.separate(job, tmp_path / "source.wav", tmp_path) + + assert calls == ["cpu"] # one spawn serving all three jobs + + +def test_worker_respawned_on_device_change(tmp_path, monkeypatch): + calls: list[str] = [] + monkeypatch.setattr(sep_mod, "_spawn_worker_cmd", _stub_spawns(set(), calls)) + + devices = iter(["cpu", "cuda"]) + monkeypatch.setattr(sep_mod, "get_demucs_device", lambda: next(devices)) + + sep_mod.separate(Job(id="abcdefabc310"), tmp_path / "source.wav", tmp_path) + sep_mod.separate(Job(id="abcdefabc311"), tmp_path / "source.wav", tmp_path) + + assert calls == ["cpu", "cuda"] + + +def test_worker_not_reused_after_failure(tmp_path, monkeypatch): + """A failed job's worker is never handed the next job -- GPU/CUDA state + afterward isn't something we can vouch for (see demucs_worker.py).""" + calls: list[str] = [] + monkeypatch.setattr(sep_mod, "get_demucs_device", lambda: "cpu") + monkeypatch.setattr(sep_mod, "_spawn_worker_cmd", _stub_spawns({"cpu"}, calls)) + + with pytest.raises(SeparationError): + sep_mod.separate(Job(id="abcdefabc320"), tmp_path / "source.wav", tmp_path) + # The next job on the same device spawns a fresh worker rather than + # reusing the one that just failed. + monkeypatch.setattr(sep_mod, "_spawn_worker_cmd", _stub_spawns(set(), calls)) + sep_mod.separate(Job(id="abcdefabc321"), tmp_path / "source.wav", tmp_path) + + assert calls == ["cpu", "cpu"] # two spawns: the failure, then the retry + + +def test_cancel_kills_worker_next_job_spawns_fresh(tmp_path, monkeypatch): + calls: list[str] = [] + monkeypatch.setattr(sep_mod, "get_demucs_device", lambda: "cpu") + monkeypatch.setattr(sep_mod, "_spawn_worker_cmd", _stub_spawns({"cpu"}, calls)) + cancelled_job = Job(id="abcdefabc330") + cancelled_job.cancel_requested = True + + with pytest.raises(JobCancelled): + sep_mod.separate(cancelled_job, tmp_path / "source.wav", tmp_path) + + monkeypatch.setattr(sep_mod, "_spawn_worker_cmd", _stub_spawns(set(), calls)) + sep_mod.separate(Job(id="abcdefabc331"), tmp_path / "source.wav", tmp_path) + + assert calls == ["cpu", "cpu"] # two spawns: cancelled, then fresh