feat(pipeline): persistent demucs worker - #310
Merged
Merged
Conversation
Replaces the fresh-subprocess-per-job model with a warm worker process that loads the demucs model once and serves jobs one at a time over a stdin/stderr protocol, reusing the same process across consecutive successful jobs on the same device instead of paying spawn + import + model-load + CUDA warmup on every single job. Measured on an RTX 3080 (see #288's data): startup was 35-42% of the separate stage for a fresh worker. With reuse, a warm second job drops separate_startup from ~5s to ~0.6s and total job time from ~13.5s to ~6.7s -- roughly half, for every job after the first on a given device. app/pipeline/demucs_worker.py: the worker script (run via `python -m app.pipeline.demucs_worker <device>`). Calls the exact same demucs library functions the CLI itself calls (load_track, apply_model, save_audio, same default split/overlap/segment/clip/bit-depth) -- not a reimplementation of the audio pipeline, just the same calls made repeatedly on an already-loaded model instead of once per fresh process. Verified bit-for-bit identical output against the old subprocess-CLI path on a real track (with shifts=0, since demucs's own apply_model applies a random time-shift internally whenever shifts>=1, independent of this change -- both paths share that variance equally). app/pipeline/separate.py: _run_demucs now reuses-or-spawns a worker via _get_worker(device) instead of always spawning; dispatches one JSON line per job and reads progress from stderr exactly as before (same tqdm-driven "NN%" lines, same watchdog-stall detection). A worker is torn down -- never reused for the next job -- after a cancel or any job failure: GPU/CUDA state afterward isn't something we can vouch for, so only the happy path keeps the process warm. A device change (Settings, or the GPU->CPU fallback within one job) always gets a fresh worker. app/main.py: kill the worker on clean app shutdown so it's never left as an orphaned process. Closes #309
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Closes #309, a follow-up to #288's real-hardware measurement (35-42% of the GPU
separate()stage was subprocess-spawn + model-load overhead on an RTX 3080).app/pipeline/demucs_worker.py(new): the worker script, run aspython -m app.pipeline.demucs_worker <device>. Loads the model once, then serves jobs one at a time over a stdin (dispatch) / stderr (progress + completion) protocol. Calls the exact same demucs library functions the CLI itself calls internally (load_track,apply_model,save_audio, same default split/overlap/segment/clip-mode/bit-depth) -- this is not a reimplementation of the audio pipeline, just the same calls made repeatedly on an already-loaded model instead of once per fresh process.app/pipeline/separate.py:_run_demucsnow reuses-or-spawns a worker via_get_worker(device)instead of always spawning a fresh subprocess. Dispatches one JSON line per job, reads progress from stderr exactly as before (same tqdm-driven"NN%"lines, same stall watchdog). A worker is torn down, never reused, after a cancel or any job failure -- GPU/CUDA state afterward isn't something we can vouch for, so only the happy path stays warm. A device change (Settings, or the GPU->CPU fallback mid-job) always gets a fresh worker.separate()'s retry/fallback policy is untouched.app/main.py: kills the worker on clean app shutdown so it's never left as an orphaned process (it has no parent-death watchdog of its own).Correctness validation (real hardware, not simulated)
Bit-for-bit identical output: ran the same real track through the old
python -m demucsCLI subprocess and the new worker, both withshifts=0(demucs's ownapply_modelapplies a random time-shift internally whenevershifts>=1by design -- both paths share that variance equally, soshifts=0isolates genuine correctness from that). All 6 stems came back byte-identical.End-to-end reuse + correctness on a real RTX 3080: 2 consecutive real jobs through
separate()--separate_startupSame worker process reused, startup dropped from 5.3s to 0.6s, total job time nearly halved. Job 2's stem output verified sane (correct shapes, RMS values consistent with the same track's known presence data).
Test plan
tests/test_separate_fallback.pyrewritten for the new_spawn_worker_cmdseam (persistent stub scripts instead of one-shot stubs): all prior GPU->CPU fallback / cancel / partial-output-cleanup coverage preserved, plus new tests for worker reuse across consecutive jobs, respawn on device change, no-reuse-after-failure, and cancel-kills-worker.205 passed, 12 skippedruff check/ruff format --checkclean