feat(pipeline): persistent demucs worker

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

app/pipeline/demucs_worker.py

@@ -0,0 +1,110 @@
+"""Persistent demucs worker (#309).
+
+Run as its own process: `python -m app.pipeline.demucs_worker <device>`.
+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": "<path>", "job_dir": "<path>", "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@@<json-encoded message>"   -- 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()

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,13 +143,28 @@ 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"):
                 line = buf.strip()
                 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: