diff --git a/app/api/jobs.py b/app/api/jobs.py index f2be8f41..4f16765f 100644 --- a/app/api/jobs.py +++ b/app/api/jobs.py @@ -326,6 +326,16 @@ def cancel_job(job_id: str) -> dict: if job is None: raise HTTPException(status_code=404, detail="job not found") if job.status in ("done", "error", "cancelled"): + # A vocal split only ever runs on a done job, so this early return made + # it uncancellable by construction: the flag was never even set, while + # the split held _pipeline_lock and stalled the whole import queue for + # its full duration (#519). Terminating the worker is enough -- the + # split's own error path marks it failed and releases the lock. + if job.vocal_split == "running": + job.cancel_requested = True + proc = registry_get_proc(job_id) + if proc is not None and proc.poll() is None: + proc.terminate() return job.to_state() job.cancel_requested = True diff --git a/app/core/process.py b/app/core/process.py index 609372ba..965c4e21 100644 --- a/app/core/process.py +++ b/app/core/process.py @@ -38,3 +38,53 @@ def process_exists(pid: int) -> bool: kernel32.CloseHandle(handle) return True return ctypes.get_last_error() != ERROR_INVALID_PARAMETER + + +_PARENT_POLL_SECONDS = 1.0 + + +def _watch_parent(parent_pid: int) -> None: + """Exit as soon as the process that spawned us is gone. + + A worker's stdin EOF only covers a parent that exits between jobs. + Mid-inference the worker reads nothing, and the parent may have been killed + in a way that ran no cleanup at all (SIGKILL, Force Quit, Task Manager, a + crash). Without this the worker keeps running -- holding a GPU, in the + demucs and vocal-split cases -- with nobody left to collect the result. + + os._exit rather than sys.exit: this runs on a daemon thread, and raising + SystemExit there would not interrupt inference running in C code. Nothing + here needs flushing. + """ + import sys + import time + + while True: + if not process_exists(parent_pid): + sys.stderr.write("@@ERROR@@parent process exited\n") + sys.stderr.flush() + os._exit(1) + time.sleep(_PARENT_POLL_SECONDS) + + +def arm_parent_watchdog() -> None: + """Start the parent-death watchdog if the parent asked for one. + + Shared by every long-running worker. It lived in demucs_worker.py, which is + why vocal_split_worker and section_worker never had it: a Force-Quit during + a vocal split orphaned an onnxruntime process holding the GPU, and a + section pass outlived the parent whose TIMEOUT_SECTIONS was its only bound + (#519). + """ + import threading + + raw = os.environ.get("STEMDECK_PARENT_PID", "").strip() + if not raw: + return + try: + parent_pid = int(raw) + except ValueError: + return + if parent_pid <= 0 or parent_pid == os.getpid(): + return + threading.Thread(target=_watch_parent, args=(parent_pid,), daemon=True).start() diff --git a/app/pipeline/demucs_worker.py b/app/pipeline/demucs_worker.py index 482cc0ae..dc679ca8 100644 --- a/app/pipeline/demucs_worker.py +++ b/app/pipeline/demucs_worker.py @@ -35,14 +35,11 @@ from __future__ import annotations import json -import os import sys -import threading -import time from pathlib import Path from app.core.config import DEMUCS_MODEL -from app.core.process import process_exists +from app.core.process import arm_parent_watchdog def _run_one_job(model, device: str, req: dict) -> None: @@ -89,47 +86,9 @@ def _run_one_job(model, device: str, req: dict) -> None: ) -_PARENT_POLL_SECONDS = 1.0 - - -def _watch_parent(parent_pid: int) -> None: - """Exit as soon as the process that spawned us is gone. - - The stdin EOF in the loop below only covers a parent that exits between - jobs. Mid-separation the worker is inside torch and reads nothing, and the - parent may have been killed in a way that ran no cleanup at all (SIGKILL, - Force Quit, Task Manager, a crash). Without this, the worker would keep a - GPU busy with nobody left to collect the result. - - os._exit rather than sys.exit: this runs on a daemon thread, and raising - SystemExit there would not interrupt inference running in C code. Nothing - here needs flushing -- a half-written model directory is cleared before the - job is retried. - """ - while True: - if not process_exists(parent_pid): - sys.stderr.write("@@ERROR@@parent process exited\n") - sys.stderr.flush() - os._exit(1) - time.sleep(_PARENT_POLL_SECONDS) - - -def _arm_parent_watchdog() -> None: - raw = os.environ.get("STEMDECK_PARENT_PID", "").strip() - if not raw: - return - try: - parent_pid = int(raw) - except ValueError: - return - if parent_pid <= 0 or parent_pid == os.getpid(): - return - threading.Thread(target=_watch_parent, args=(parent_pid,), daemon=True).start() - - def main() -> None: device = sys.argv[1] if len(sys.argv) > 1 else "cpu" - _arm_parent_watchdog() + arm_parent_watchdog() from demucs.pretrained import get_model diff --git a/app/pipeline/runner.py b/app/pipeline/runner.py index 07603ab4..f4b3c6a0 100644 --- a/app/pipeline/runner.py +++ b/app/pipeline/runner.py @@ -14,6 +14,7 @@ from app.core.models import Job, JobCancelled, _set from app.core.redact import redact from app.core.registry import persist as persist_registry +from app.core.registry import set_proc from app.pipeline.analyze import analyze from app.pipeline.beatgrid import compute_beat_grid from app.pipeline.collect import ( @@ -49,6 +50,30 @@ def _check_cancel(job: Job) -> None: raise JobCancelled() +def _run_registered_ffmpeg(job: Job, cmd: list[str], timeout: int) -> tuple[int, bytes]: + """Run ffmpeg with the process registered, so cancel can reach it. + + subprocess.run() cannot be interrupted: POST /cancel sets the flag, but + nothing looks at it until the call returns, so a cancel during a large + upload's transcode was a no-op for up to TIMEOUT_FFMPEG per call -- twice + over on the .mp4 path, which runs both this and the video extract (#519). + + Mirrors collect._run_ffmpeg, which registers for exactly this reason. + """ + proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE) + set_proc(job.id, proc) + try: + try: + _, stderr = proc.communicate(timeout=timeout) + except subprocess.TimeoutExpired: + proc.kill() + proc.communicate() + raise + return proc.returncode, stderr or b"" + finally: + set_proc(job.id, None) + + def _extract_video_track(job: Job, source: Path, job_dir: Path) -> None: """For an .mp4 upload, preserve a silent video-only track at video.mp4 so the studio can later mux it with a custom stem mix @@ -76,7 +101,7 @@ def _extract_video_track(job: Job, source: Path, job_dir: Path) -> None: str(dest), ] try: - result = subprocess.run(cmd, capture_output=True, timeout=TIMEOUT_FFMPEG) + returncode, _ = _run_registered_ffmpeg(job, cmd, TIMEOUT_FFMPEG) except (OSError, subprocess.SubprocessError) as e: # ffmpeg missing or timed out. Distinct from an .mp4 that simply has no # video stream, and the only one of the two worth surfacing (#436). @@ -84,7 +109,7 @@ def _extract_video_track(job: Job, source: Path, job_dir: Path) -> None: job.video_status = "failed" logger.warning("video extract failed for job %s: %s", job.id, e) return - if result.returncode != 0 or not dest.is_file() or dest.stat().st_size == 0: + if returncode != 0 or not dest.is_file() or dest.stat().st_size == 0: dest.unlink(missing_ok=True) job.video_status = "unavailable" logger.info("no video track preserved for job %s (source has no video stream?)", job.id) @@ -127,10 +152,10 @@ def _prepare_local_source(job: Job, source: Path, job_dir: Path) -> Path: "-y", str(dest), ] - result = subprocess.run(cmd, capture_output=True, timeout=TIMEOUT_FFMPEG) - if result.returncode != 0: + returncode, stderr = _run_registered_ffmpeg(job, cmd, TIMEOUT_FFMPEG) + if returncode != 0: raise RuntimeError( - "ffmpeg transcode failed: " + result.stderr.decode("utf-8", errors="replace").strip() + "ffmpeg transcode failed: " + stderr.decode("utf-8", errors="replace").strip() ) source.unlink(missing_ok=True) return dest diff --git a/app/pipeline/section_worker.py b/app/pipeline/section_worker.py index 0d336317..7998404d 100644 --- a/app/pipeline/section_worker.py +++ b/app/pipeline/section_worker.py @@ -10,6 +10,7 @@ import threading from pathlib import Path +from app.core.process import arm_parent_watchdog from app.pipeline.section_refine import refine_segments _HEARTBEAT_SECONDS = 10 @@ -66,6 +67,9 @@ def _load_beat_grid(path: Path | None) -> object | None: def main(argv: list[str] | None = None) -> int: + # Without this a CPU inference pass outlives the parent whose + # TIMEOUT_SECTIONS was its only bound (#519). + arm_parent_watchdog() args = _parser().parse_args(argv) for path in (args.stems_dir / f"{name}.wav" for name in ("bass", "drums", "other", "vocals")): if not path.is_file(): diff --git a/app/pipeline/sections.py b/app/pipeline/sections.py index 7e509b40..0a994f16 100644 --- a/app/pipeline/sections.py +++ b/app/pipeline/sections.py @@ -214,6 +214,10 @@ def _terminate(proc: subprocess.Popen) -> None: def _run_registered_process(job: Job, cmd: list[str]) -> tuple[int, list[str], list[str]]: """Run a child with cancellation, total timeout, and output-stall detection.""" env = os.environ.copy() + # The worker arms a watchdog on this and hard-exits when we disappear, so a + # kill that runs no cleanup (SIGKILL, Force Quit, Task Manager, a crash) + # cannot leave it running with nobody to collect the result (#519). + env["STEMDECK_PARENT_PID"] = str(os.getpid()) env["PYTHONIOENCODING"] = "utf-8:replace" proc = subprocess.Popen( cmd, diff --git a/app/pipeline/vocal_split.py b/app/pipeline/vocal_split.py index 65cb224e..a3d1ebf8 100644 --- a/app/pipeline/vocal_split.py +++ b/app/pipeline/vocal_split.py @@ -61,6 +61,10 @@ def split_vocals(job: Job, stems_dir: Path) -> list[str]: # the mismatch simply moves rather than being fixed. Demucs and audio- # separator both emit progress bars and can echo track metadata, neither of # which is guaranteed to be cp1252-safe. + # The worker arms a watchdog on this and hard-exits when we disappear, so a + # kill that runs no cleanup (SIGKILL, Force Quit, Task Manager, a crash) + # cannot leave it running with nobody to collect the result (#519). + env["STEMDECK_PARENT_PID"] = str(os.getpid()) env["PYTHONIOENCODING"] = "utf-8:replace" try: import certifi diff --git a/app/pipeline/vocal_split_worker.py b/app/pipeline/vocal_split_worker.py index 116b9eb1..fe305c6b 100644 --- a/app/pipeline/vocal_split_worker.py +++ b/app/pipeline/vocal_split_worker.py @@ -24,6 +24,7 @@ import sys from app.core.config import VOCAL_SPLIT_MODEL +from app.core.process import arm_parent_watchdog def _run(device: str, vocals_path: str, out_dir: str) -> None: @@ -57,6 +58,10 @@ def _run(device: str, vocals_path: str, out_dir: str) -> None: def main() -> None: + # A Force-Quit of the app otherwise orphans this process holding the GPU: + # onnxruntime reads nothing from stdin mid-inference, so EOF never arrives + # and nothing else bounds it (#519). + arm_parent_watchdog() if len(sys.argv) < 4: sys.stderr.write("@@ERROR@@usage: vocal_split_worker \n") sys.stderr.flush() diff --git a/tests/test_cancellation_reach.py b/tests/test_cancellation_reach.py new file mode 100644 index 00000000..8f0a79b8 --- /dev/null +++ b/tests/test_cancellation_reach.py @@ -0,0 +1,170 @@ +"""Cancellation has to reach the processes it is meant to stop (#519). + +python-fastapi.md: "Always register subprocess with set_proc(job_id, proc) +immediately after Popen(); deregister in finally." Several stages used +subprocess.run(), which cannot be interrupted -- the flag is set but nothing +looks at it until the call returns. +""" + +from __future__ import annotations + +import os + +import pytest + +from app.core import process as _process +from app.core.models import Job +from app.core.registry import register as registry_register +from app.pipeline import runner as _runner + + +def _job(**kw): + return Job(id="a1b2c3d4e5f6", **kw) + + +def test_ffmpeg_is_registered_so_cancel_can_reach_it(tmp_path, monkeypatch): + seen = {} + + def _capture(job_id, proc): + # Registered while it runs, cleared after: both halves matter. + seen.setdefault("during", proc if proc is not None else seen.get("during")) + seen["last"] = proc + + monkeypatch.setattr(_runner, "set_proc", _capture) + + job = _job() + rc, _ = _runner._run_registered_ffmpeg(job, ["sh", "-c", "exit 0"], 30) + + assert rc == 0 + assert seen["during"] is not None, "cancel could not have reached this process" + assert seen["last"] is None, "the registration must be cleared in a finally" + + +def test_a_failing_command_still_deregisters(tmp_path, monkeypatch): + cleared = [] + monkeypatch.setattr(_runner, "set_proc", lambda job_id, proc: cleared.append(proc)) + + job = _job() + rc, stderr = _runner._run_registered_ffmpeg(job, ["sh", "-c", "echo boom >&2; exit 3"], 30) + + assert rc == 3 + assert b"boom" in stderr, "stderr must still be captured for the error message" + assert cleared[-1] is None + + +# ─── parent-death watchdog ─── + + +def test_the_watchdog_arms_when_the_parent_asks(monkeypatch): + started = [] + monkeypatch.setenv("STEMDECK_PARENT_PID", "999999") + monkeypatch.setattr( + "threading.Thread", + lambda *a, **kw: type("T", (), {"start": lambda self: started.append(True)})(), + ) + + _process.arm_parent_watchdog() + + assert started + + +@pytest.mark.parametrize("value", ["", "not-a-number", "0", "-1"]) +def test_the_watchdog_stays_off_without_a_usable_parent_pid(monkeypatch, value): + started = [] + monkeypatch.setenv("STEMDECK_PARENT_PID", value) + monkeypatch.setattr( + "threading.Thread", + lambda *a, **kw: type("T", (), {"start": lambda self: started.append(True)})(), + ) + + _process.arm_parent_watchdog() + + assert not started + + +def test_the_watchdog_never_targets_our_own_pid(monkeypatch): + # Would hard-exit the worker the moment it started. + started = [] + monkeypatch.setenv("STEMDECK_PARENT_PID", str(os.getpid())) + monkeypatch.setattr( + "threading.Thread", + lambda *a, **kw: type("T", (), {"start": lambda self: started.append(True)})(), + ) + + _process.arm_parent_watchdog() + + assert not started + + +def test_every_worker_spawn_exports_the_parent_pid(): + # demucs_worker had this; the other two did not, so a Force-Quit orphaned + # an onnxruntime process holding the GPU. + import pathlib + + for path in ( + "app/pipeline/separate.py", + "app/pipeline/vocal_split.py", + "app/pipeline/sections.py", + ): + src = pathlib.Path(path).read_text() + assert "STEMDECK_PARENT_PID" in src, f"{path} spawns a worker without the watchdog" + + +def test_every_worker_arms_the_watchdog(): + import pathlib + + for path in ( + "app/pipeline/demucs_worker.py", + "app/pipeline/vocal_split_worker.py", + "app/pipeline/section_worker.py", + ): + src = pathlib.Path(path).read_text() + assert "arm_parent_watchdog()" in src, f"{path} never arms the watchdog" + + +# ─── a running vocal split must be cancellable ─── + + +class _LiveProc: + def __init__(self): + self.terminated = False + + def poll(self): + return None + + def terminate(self): + self.terminated = True + + +def test_cancelling_a_running_vocal_split_terminates_it(monkeypatch): + # cancel_job returns early for a done job -- and a vocal split only ever + # runs on a done job, so it was uncancellable by construction while holding + # _pipeline_lock and stalling the import queue. + from app.api import jobs as _jobs_api + + job = _job(status="done", title="Song") + job.vocal_split = "running" + registry_register(job) + + proc = _LiveProc() + monkeypatch.setattr(_jobs_api, "registry_get_proc", lambda job_id: proc) + + _jobs_api.cancel_job(job.id) + + assert job.cancel_requested is True + assert proc.terminated, "the split ran to completion with the cancel button doing nothing" + + +def test_cancelling_a_plain_done_job_still_does_nothing(monkeypatch): + from app.api import jobs as _jobs_api + + job = _job(status="done", title="Song") + registry_register(job) + + proc = _LiveProc() + monkeypatch.setattr(_jobs_api, "registry_get_proc", lambda job_id: proc) + + _jobs_api.cancel_job(job.id) + + assert not proc.terminated + assert job.cancel_requested is False diff --git a/tests/test_video_status.py b/tests/test_video_status.py index 7a6b3a80..d89a71e5 100644 --- a/tests/test_video_status.py +++ b/tests/test_video_status.py @@ -97,14 +97,16 @@ def _run_local_extract(tmp_path: Path, returncode: int, raises=None) -> Job: source = job_dir / "in.mp4" source.write_bytes(b"x") - def fake_run(cmd, **kwargs): + # Stubs _run_registered_ffmpeg rather than subprocess.run: the extract goes + # through Popen + set_proc now, so cancel can reach it (#519). + def fake_run(job_arg, cmd, timeout): if raises is not None: raise raises if returncode == 0: Path(cmd[-1]).write_bytes(b"fake mp4 payload") - return subprocess.CompletedProcess(cmd, returncode, b"", b"") + return returncode, b"" - with patch.object(runner_mod.subprocess, "run", fake_run): + with patch.object(runner_mod, "_run_registered_ffmpeg", fake_run): runner_mod._extract_video_track(job, source, job_dir) return job diff --git a/tests/test_worker_parent_watchdog.py b/tests/test_worker_parent_watchdog.py index d7901336..583f8970 100644 --- a/tests/test_worker_parent_watchdog.py +++ b/tests/test_worker_parent_watchdog.py @@ -45,8 +45,8 @@ def test_worker_exits_when_its_parent_disappears(tmp_path): script = textwrap.dedent( """ import sys, time - from app.pipeline.demucs_worker import _arm_parent_watchdog - _arm_parent_watchdog() + from app.core.process import arm_parent_watchdog + arm_parent_watchdog() # Busy the way a separation is busy: never reading stdin, so only the # watchdog can end this process. while True: @@ -83,18 +83,20 @@ def test_worker_exits_when_its_parent_disappears(tmp_path): def test_worker_ignores_an_unset_or_bogus_parent_pid(monkeypatch): """A worker run by hand (no STEMDECK_PARENT_PID) must not arm the watchdog and shoot itself.""" - from app.pipeline import demucs_worker + import threading as _threading + + from app.core import process as _process_mod started: list[object] = [] monkeypatch.setattr( - demucs_worker.threading, + _threading, "Thread", lambda *a, **k: started.append((a, k)) or _NoopThread(), ) for value in ("", " ", "not-a-number", "0", "-5", str(os.getpid())): monkeypatch.setenv("STEMDECK_PARENT_PID", value) - demucs_worker._arm_parent_watchdog() + _process_mod.arm_parent_watchdog() assert started == [], "watchdog armed on a pid it should have ignored"