diff --git a/app/core/model_cache.py b/app/core/model_cache.py index 0de34568..14bca045 100644 --- a/app/core/model_cache.py +++ b/app/core/model_cache.py @@ -28,6 +28,7 @@ import logging import shutil +import traceback from collections.abc import Callable, Iterable from pathlib import Path from typing import TypeVar @@ -91,6 +92,11 @@ def load_or_heal(load: Callable[[], T], stale: Callable[[], Iterable[Path]]) -> fresh = [p for p in paths if p not in _healed] if not fresh: raise + # audio-separator reads its indexes with json.load(open(...)) and + # never closes the file, so the one that failed is still open in the + # traceback's frames, and Windows will not delete an open file. + # Clearing the frames' locals closes it; the traceback itself stays. + traceback.clear_frames(first.__traceback__) removed = [p for p in fresh if _remove(p)] _healed.update(fresh) if not removed: @@ -125,10 +131,16 @@ def vocal_split_artifacts(models_dir: Path, model_file: str) -> list[Path]: they are the names audio-separator actually writes: a single "model_data.json" is not one of them, so listing that instead heals the checkpoint and leaves the index that produced the error in place. + + download_checks.json, the list of models audio-separator offers, is read + before any of them, and is written in place like the rest: a setup ended + mid-fetch (the wizard's Skip, #405) leaves it half written, and json.load + fails on it every time after. """ root = models_dir / "audio-separator" return [ root / model_file, root / "vr_model_data.json", root / "mdx_model_data.json", + root / "download_checks.json", ] diff --git a/app/pipeline/warmup.py b/app/pipeline/warmup.py index 0a033382..dc34d406 100644 --- a/app/pipeline/warmup.py +++ b/app/pipeline/warmup.py @@ -8,10 +8,23 @@ steps; Docker has no equivalent step and keeps the pre-existing lazy-download-on-first-use behavior (see docs/models.md). -One line per model is printed to stdout so the caller can report per-model -status: "WARMUP_OK " on success, "WARMUP_FAILED " on -failure. A model failing to download here is never fatal to setup -- exactly -like the lazy path it replaces, a missing model degrades only that one +Lines on stdout tell the caller where it is (#405): "WARMUP_START +" before each model, then "WARMUP_OK " on success or +"WARMUP_FAILED " on failure, followed by "WARMUP_OFFLINE " +when the download host could not be reached or stopped answering. When the +disk cannot hold the models, the only line is "WARMUP_NO_SPACE +" and nothing is downloaded. + +The steps run inside bounded_downloads (app/core/model_download.py): a +torch.hub download (the separation and beat models) gets the separation +model's read timeout and tries, so a dead connection ends it in about a +minute and a half, and every other socket that sets no timeout of its own +(Whisper's) gets the read timeout. The Hugging Face hub and audio-separator +set their own, audio-separator 300 s per read, which the shell's stall limit +outlasts. + +A model failing to download here is never fatal to setup -- exactly like the +lazy path it replaces, a missing model degrades only that one feature (e.g. no beat grid, or the karaoke split unavailable until it can download later) rather than blocking the app from starting. @@ -31,18 +44,38 @@ from __future__ import annotations +import http.client import os +import shutil +import socket import sys from app.core.config import ( BEAT_MODEL_CHECKPOINT, DEMUCS_MODEL, + MODEL_DOWNLOAD_ATTEMPTS, MODELS_DIR, SECTION_MODEL, + TIMEOUT_MODEL_DOWNLOAD_READ, TRANSCRIBE_MODEL_CPU, TRANSCRIBE_MODEL_GPU, VOCAL_SPLIT_MODEL, ) +from app.core.model_download import bounded_downloads + +# What each model adds to the data folder, in bytes: measured on a finished +# install (models/ and cache/huggingface/, 2026-11-03), except Whisper's +# "small", which is its published size. Only the space check before anything +# downloads uses these, so a model renamed through the environment is just not +# counted. +_DOWNLOAD_BYTES = { + "demucs": 54_996_327, + "beat_this": 81_058_141, + "sections": 11_204_803, + "vocal_split": 52_834_731, + "turbo": 1_617_941_637, + "small": 483_617_219, +} def _warm_demucs() -> None: @@ -90,24 +123,30 @@ def _warm_sections() -> None: load_pretrained_model(model_name=SECTION_MODEL, device="cpu") -def _warm_whisper() -> None: - """The Whisper model the lyrics stage would use here, and only when that - stage would run: under "auto" only a CUDA install transcribes, and the - GPU model is 1.6 GB. Last in the list for the same reason, so a slow - connection costs the other models nothing. A download cut short by the - setup timeout heals itself: Whisper checks the file's SHA-256 on every - load and fetches it again when it does not match.""" +def _whisper_model() -> str | None: + """The Whisper model the lyrics stage would use here, or None when that + stage would not run: under "auto" only a CUDA install transcribes.""" from app.core.settings import get_demucs_device, transcribe_lyrics_enabled device = get_demucs_device() if not transcribe_lyrics_enabled(device): + return None + return TRANSCRIBE_MODEL_GPU if device == "cuda" else TRANSCRIBE_MODEL_CPU + + +def _warm_whisper() -> None: + """The GPU model is 1.6 GB, so this is last in the list: a slow + connection costs the other models nothing. A download cut short by the + setup timeout heals itself: Whisper checks the file's SHA-256 on every + load and fetches it again when it does not match.""" + name = _whisper_model() + if name is None: return import whisper from app.pipeline.transcribe import whisper_models_dir - name = TRANSCRIBE_MODEL_GPU if device == "cuda" else TRANSCRIBE_MODEL_CPU whisper.load_model(name, device="cpu", download_root=str(whisper_models_dir())) @@ -137,17 +176,85 @@ def _remember_vocal_split(error: Exception | None) -> None: print(f"WARMUP_NOTE vocal_split status not recorded: {e}", flush=True) +def _free_bytes() -> int: + """Free space on the disk MODELS_DIR is on, which may not exist yet.""" + folder = MODELS_DIR + while not folder.exists() and folder != folder.parent: + folder = folder.parent + return shutil.disk_usage(folder).free + + +def _bytes_needed() -> int: + """About what the downloads still need: every model's size, less what + MODELS_DIR already holds, so an install that has them all needs nothing. + + It errs towards trying. Whatever else sits in MODELS_DIR counts as + present, and the sections model's few MB in the Hugging Face cache are + always counted as needed rather than walking a cache the user may share + with other tools. A check that passes wrongly costs a download that fails + on the full disk with its own error, which the wizard reports; one that + fails wrongly would skip every model for nothing. A torch.hub ".partial" + is a download cut short, not a model.""" + total = 0 + for name, _ in _STEPS: + key = _whisper_model() if name == "whisper" else name + total += _DOWNLOAD_BYTES.get(key or "", 0) + present = sum( + path.stat().st_size + for path in (MODELS_DIR.rglob("*") if MODELS_DIR.is_dir() else ()) + if path.is_file() and not path.is_symlink() and path.suffix != ".partial" + ) + return max(0, total - present) + + +def _enough_space() -> bool: + try: + needed, free = _bytes_needed(), _free_bytes() + except Exception as e: + # Only a courtesy: a full disk still fails each download with its own + # error. Not a reason to download nothing. + print(f"WARMUP_NOTE space not checked: {e}", flush=True) + return True + if needed > free: + print(f"WARMUP_NO_SPACE {needed} {free}", flush=True) + return False + return True + + +def _unreachable(error: BaseException) -> bool: + """Whether a download failed because its host could not be reached or + stopped answering: no network, no DNS, a refused or dropped connection, a + read that timed out. The whole chain is read, since every library wraps + these in errors of its own (requests, httpx, the hub, ModelDownloadError).""" + seen: set[int] = set() + current: BaseException | None = error + while current is not None and id(current) not in seen: + if isinstance( + current, (socket.gaierror, TimeoutError, ConnectionError, http.client.IncompleteRead) + ): + return True + seen.add(id(current)) + current = current.__cause__ or current.__context__ + return False + + def main() -> int: - for name, fn in _STEPS: - try: - fn() - print(f"WARMUP_OK {name}", flush=True) - if name == "vocal_split": - _remember_vocal_split(None) - except Exception as e: - print(f"WARMUP_FAILED {name} {e}", flush=True) - if name == "vocal_split": - _remember_vocal_split(e) + if not _enough_space(): + return 0 + with bounded_downloads(TIMEOUT_MODEL_DOWNLOAD_READ, MODEL_DOWNLOAD_ATTEMPTS): + for number, (name, fn) in enumerate(_STEPS, start=1): + print(f"WARMUP_START {name} {number} {len(_STEPS)}", flush=True) + try: + fn() + print(f"WARMUP_OK {name}", flush=True) + if name == "vocal_split": + _remember_vocal_split(None) + except Exception as e: + print(f"WARMUP_FAILED {name} {e}", flush=True) + if _unreachable(e): + print(f"WARMUP_OFFLINE {name}", flush=True) + if name == "vocal_split": + _remember_vocal_split(e) return 0 # best-effort throughout -- see module docstring diff --git a/desktop/src-tauri/src/backendjob.rs b/desktop/src-tauri/src/backendjob.rs index b4db32f8..6d89092f 100644 --- a/desktop/src-tauri/src/backendjob.rs +++ b/desktop/src-tauri/src/backendjob.rs @@ -54,7 +54,7 @@ mod windows_impl { }; use windows_sys::Win32::System::JobObjects::{ AssignProcessToJobObject, CreateJobObjectW, JobObjectExtendedLimitInformation, - SetInformationJobObject, JOBOBJECT_EXTENDED_LIMIT_INFORMATION, + SetInformationJobObject, TerminateJobObject, JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, }; use windows_sys::Win32::System::Threading::{ @@ -67,8 +67,21 @@ mod windows_impl { // SAFETY: a job handle is a kernel handle, usable from any thread; it is // only closed once, in Drop. unsafe impl Send for KillOnCloseJob {} + // SAFETY: shared use only terminates the job, which Windows allows from + // any thread on the same handle at once. + unsafe impl Sync for KillOnCloseJob {} impl KillOnCloseJob { + /// End every process in the job now, while StemDeck keeps running: + /// the setup wizard's Skip on the model step (#405). + pub fn terminate(&self) -> io::Result<()> { + // SAFETY: the handle is owned by `self` and open until Drop. + if unsafe { TerminateJobObject(self.0, 1) } == 0 { + return Err(io::Error::last_os_error()); + } + Ok(()) + } + fn new() -> io::Result { // SAFETY: plain Win32 calls; the handle is owned by the returned // value, and `info` is a zeroed plain-data struct of the size given. diff --git a/desktop/src-tauri/src/main.rs b/desktop/src-tauri/src/main.rs index 02b3b5e6..ccd6d1cb 100644 --- a/desktop/src-tauri/src/main.rs +++ b/desktop/src-tauri/src/main.rs @@ -160,14 +160,46 @@ struct BackendHandles { _job: Option, } +struct SetupChild { + label: String, + #[cfg(unix)] + pid: u32, + /// The job it runs in (run_setup_child): ending the job ends whatever the + /// child started as well. + #[cfg(windows)] + job: Option>, +} + +impl SetupChild { + fn stop(&self) -> Result<(), String> { + // SAFETY: the pid was stored right after spawn and is cleared once the + // child has been waited on; SIGTERM is best-effort. + #[cfg(unix)] + unsafe { + libc::kill(self.pid as libc::pid_t, libc::SIGTERM) + }; + #[cfg(windows)] + match &self.job { + Some(job) => job + .terminate() + .map_err(|e| format!("could not stop the {}: {e}", self.label))?, + // spawn_in_job could not make a job and said so; there is no + // handle left here to end the child with. + None => return Err(format!("could not stop the {}: it has no job", self.label)), + } + Ok(()) + } +} + #[derive(Default)] struct BackendStateInner { handles: Option, /// True while start_backend is executing; prevents concurrent starts (#145). starting: bool, - /// PID of an in-progress setup-time subprocess (pip install, or the model - /// warmup download, #275); killed by stop_backend on window close (#140). - setup_child_pid: Option, + /// The setup-time subprocess in progress (pip install, or the model warmup + /// download, #275), stopped by stop_backend on window close (#140) and by + /// the wizard's Skip (#405). + setup_child: Option, /// Save destinations the user has picked but not yet downloaded to (#338). /// /// The export is two commands so the UI can tell "choosing a folder" apart @@ -569,6 +601,7 @@ fn main() { ensure_torch_device, repair_torchvision, warmup_models, + skip_model_warmup, start_backend, build_target, open_url, @@ -1748,15 +1781,89 @@ fn ensure_external_assets() -> Result { }) } -#[derive(Serialize)] +#[derive(Serialize, Default, Debug, PartialEq)] #[serde(rename_all = "camelCase")] struct ModelWarmupStatus { demucs_ready: bool, beat_this_ready: bool, sections_ready: bool, vocal_split_ready: bool, + whisper_ready: bool, + /// The models that failed because their download host could not be + /// reached or stopped answering (WARMUP_OFFLINE). + offline: Vec, + /// Set when the disk could not hold the models, so none was downloaded. + no_space: Option, +} + +#[derive(Serialize, Debug, PartialEq)] +struct NoSpace { + needed: u64, + free: u64, +} + +impl ModelWarmupStatus { + /// The status `app/pipeline/warmup.py` reported, from its stdout. + fn from_stdout(stdout: &str) -> Self { + let mut status = Self::default(); + for line in stdout.lines() { + match line.split(' ').collect::>().as_slice() { + ["WARMUP_OK", "demucs"] => status.demucs_ready = true, + ["WARMUP_OK", "beat_this"] => status.beat_this_ready = true, + ["WARMUP_OK", "sections"] => status.sections_ready = true, + ["WARMUP_OK", "vocal_split"] => status.vocal_split_ready = true, + ["WARMUP_OK", "whisper"] => status.whisper_ready = true, + ["WARMUP_OFFLINE", model] => status.offline.push(model.to_string()), + ["WARMUP_NO_SPACE", needed, free] => { + if let (Ok(needed), Ok(free)) = (needed.parse(), free.parse()) { + status.no_space = Some(NoSpace { needed, free }); + } + } + _ => {} + } + } + status + } +} + +/// One `WARMUP_START ` line, sent to the setup screen as +/// it arrives so it can say which model it is on (#405). +#[derive(Clone, Serialize, Debug, PartialEq)] +struct WarmupProgress { + model: String, + number: u32, + count: u32, } +fn warmup_progress(line: &str) -> Option { + let mut parts = line.strip_prefix("WARMUP_START ")?.split(' '); + let model = parts.next()?.to_string(); + let number = parts.next()?.parse().ok()?; + let count = parts.next()?.parse().ok()?; + Some(WarmupProgress { + model, + number, + count, + }) +} + +const MODEL_WARMUP_LABEL: &str = "model warmup"; + +/// How long the model warmup may say nothing on either stream before it is +/// taken for hung (#405). It is not the download bound: the downloads in it +/// give up by themselves and say so, and tqdm writes to stderr while bytes +/// arrive. This catches a warmup stuck past those bounds. The longest silence +/// any of them allows is audio-separator's 300 s read timeout: the real +/// warmup behind a proxy that never answers was quiet on both streams for +/// 300 s there and at most 77 s anywhere else. A minute over that, so a dead +/// connection still gets its own message, naming the model. +const WARMUP_STALL_TIMEOUT: Duration = Duration::from_secs(6 * 60); + +/// Backstop for a warmup that keeps talking and never finishes. The largest +/// model is 1.6 GB: about 35 minutes at 0.77 MB/s, the slowest connection a +/// report has measured (#502). +const WARMUP_HARD_CAP: Duration = Duration::from_secs(2 * 60 * 60); + /// Eagerly downloads/caches the ML models StemDeck uses (Demucs, beat-this, /// automatic song sections, and the on-demand lead/backing vocal-split karaoke model, #275) via /// `app/pipeline/warmup.py`, so a user's first real job doesn't pay for any @@ -1765,7 +1872,19 @@ struct ModelWarmupStatus { /// still proceeds, and that one feature falls back to its existing /// lazy-download-on-first-use behavior, same as before this step existed. #[tauri::command] -fn warmup_models(state: tauri::State) -> Result { +async fn warmup_models(app: tauri::AppHandle) -> Result { + // Off the main thread. A command that is not async runs on it, so while + // the models downloaded the window could not even be closed (#405), and + // the wizard's Skip could not have reached the shell. + tauri::async_runtime::spawn_blocking(move || run_model_warmup(&app.state(), &app)) + .await + .map_err(|e| e.to_string())? +} + +fn run_model_warmup( + state: &BackendState, + app: &tauri::AppHandle, +) -> Result { log_setup_step("model-warmup"); let root = app_root()?; let data_dir = local_data_dir()?; @@ -1791,27 +1910,32 @@ fn warmup_models(state: tauri::State) -> Result status.demucs_ready = true, - "WARMUP_OK beat_this" => status.beat_this_ready = true, - "WARMUP_OK sections" => status.sections_ready = true, - "WARMUP_OK vocal_split" => status.vocal_split_ready = true, - _ if line.starts_with("WARMUP_FAILED") => { - append_to_setup_log(&data_dir, &format!("model warmup: {line}")); - } - _ => {} + if ["WARMUP_FAILED", "WARMUP_NO_SPACE", "WARMUP_NOTE"] + .iter() + .any(|kind| line.starts_with(kind)) + { + append_to_setup_log(&data_dir, &format!("model warmup: {line}")); } } if !output.status.success() { @@ -1824,32 +1948,73 @@ fn warmup_models(state: tauri::State) -> Result Result, ) -> Result { - // Held to the end of this function: dropping the job ends the child. - let (child, _job) = - backendjob::spawn_in_job(command).map_err(|e| format!("failed to start {label}: {e}"))?; - if let Ok(mut inner) = state.inner.lock() { - inner.setup_child_pid = Some(child.id()); + // One at a time: there is one slot to stop a child through, and two + // downloads into the same cache would only race each other. The check and + // the spawn share the lock so two callers cannot both pass it. + let mut inner = state.inner.lock().map_err(|e| e.to_string())?; + if let Some(running) = &inner.setup_child { + return Err(format!( + "{label} not started: the {} is still running", + running.label + )); } + // Held to the end of this function: dropping the last handle to the job + // ends the child. + let (child, job) = + backendjob::spawn_in_job(command).map_err(|e| format!("failed to start {label}: {e}"))?; + let _job = job.map(Arc::new); + inner.setup_child = Some(SetupChild { + label: label.to_string(), + #[cfg(unix)] + pid: child.id(), + #[cfg(windows)] + job: _job.clone(), + }); + drop(inner); let output = wait(child); if let Ok(mut inner) = state.inner.lock() { - inner.setup_child_pid = None; + inner.setup_child = None; } output } +/// The wizard's Skip on the model step (#405). Stops the model warmup and +/// nothing else, so a stray call can never cut a pip install short. Each +/// model then downloads the first time a feature needs it, as it would have +/// if this step did not exist. +#[tauri::command] +fn skip_model_warmup(state: tauri::State) -> Result<(), String> { + if stop_setup_child(&state, MODEL_WARMUP_LABEL)? { + if let Ok(data_dir) = local_data_dir() { + append_to_setup_log(&data_dir, "model warmup: skipped by the user"); + } + } + Ok(()) +} + +/// Stops the running setup child if it is the one named `label`, saying +/// whether it did. Nothing to stop is not an error: it finished first. +fn stop_setup_child(state: &BackendState, label: &str) -> Result { + let inner = state.inner.lock().map_err(|e| e.to_string())?; + match &inner.setup_child { + Some(child) if child.label == label => child.stop().map(|()| true), + _ => Ok(false), + } +} + /// Spawns the Python/uvicorn backend, waits for it to become healthy, and returns its URL. #[tauri::command] fn start_backend( @@ -3230,7 +3395,7 @@ const PIP_STALL_TIMEOUT: Duration = Duration::from_secs(5 * 60); /// Backstop for a pip that never goes quiet but never finishes either. Nothing /// legitimate approaches this: the reporter whose install prompted the change /// took about 17 minutes for the whole 3 GB pass on a slow link. The window -/// close handler kills the child through `setup_child_pid` long before here. +/// close handler kills the child through `setup_child` long before here. #[cfg(not(target_os = "macos"))] const PIP_HARD_CAP: Duration = Duration::from_secs(4 * 60 * 60); @@ -4139,18 +4304,18 @@ async fn start_audio_drag( } fn stop_backend(state: &BackendState) { - let (handles, _setup_child_pid) = match state.inner.lock() { - Ok(mut guard) => (guard.handles.take(), guard.setup_child_pid.take()), + let (handles, setup_child) = match state.inner.lock() { + Ok(mut guard) => (guard.handles.take(), guard.setup_child.take()), Err(_) => return, }; - // Kill any in-progress setup-time subprocess (pip install, model warmup) - // so it doesn't corrupt the venv/cache if the window is closed mid-setup (#140). - // On Windows its job (run_setup_child) ends it when StemDeck exits. - #[cfg(unix)] - if let Some(pid) = _setup_child_pid { - // SAFETY: pid was stored immediately after spawn; we send SIGTERM best-effort. - unsafe { libc::kill(pid as libc::pid_t, libc::SIGTERM) }; + // Stop any in-progress setup-time subprocess (pip install, model warmup) + // so it doesn't corrupt the venv/cache if the window is closed mid-setup + // (#140), on Windows too (#405). + if let Some(child) = setup_child { + if let Err(e) = child.stop() { + eprintln!("[stemdeck] {e}"); + } } let Some(mut handles) = handles else { return }; @@ -7417,7 +7582,7 @@ fn update_setup_config( /// not block: emitting a Tauri event or appending a line to a log is fine, /// anything slower would stall the very pipe this exists to drain. /// -/// stderr is still read to EOF rather than streamed. It carries the failure +/// stderr is collected whole rather than streamed. It carries the failure /// text used to classify an error, which is only wanted once, at the end. fn child_output_streaming( mut child: Child, @@ -7429,11 +7594,13 @@ fn child_output_streaming( where F: FnMut(&str) + Send + 'static, { - // Every line stamps this, and the wait loop reads it. Only stdout counts as - // activity: pip writes its progress there, and stderr is read to EOF in one - // go rather than line by line, so it has nothing per-line to stamp with. + // Output on either stream stamps this, and the wait loop reads it. pip + // writes its progress to stdout; the model warmup's downloads write theirs + // (tqdm) to stderr, and its stdout is quiet for as long as one model takes, + // which on a slow connection is far longer than the stall budget (#405). let last_line = Arc::new(Mutex::new(Instant::now())); let stamp = Arc::clone(&last_line); + let stderr_stamp = Arc::clone(&last_line); let stdout_reader = child.stdout.take().map(|pipe| { thread::spawn(move || { let mut collected = Vec::new(); @@ -7451,9 +7618,22 @@ where }); let stderr_reader = child.stderr.take().map(|mut pipe| { thread::spawn(move || { - let mut buf = Vec::new(); - let _ = pipe.read_to_end(&mut buf); - buf + let mut collected = Vec::new(); + let mut chunk = [0u8; 8192]; + loop { + match pipe.read(&mut chunk) { + Ok(0) => break, + Ok(n) => { + if let Ok(mut at) = stderr_stamp.lock() { + *at = Instant::now(); + } + collected.extend_from_slice(&chunk[..n]); + } + Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {} + Err(_) => break, + } + } + collected }) }); let deadline = Deadline::Stall { @@ -8036,8 +8216,18 @@ mod tests { let state = super::BackendState::default(); let output = super::run_setup_child(&mut command, &state, "setup child", |mut child| { - let pid = state.inner.lock().unwrap().setup_child_pid; - assert_eq!(pid, Some(child.id()), "the close handler can find it"); + let recorded = state + .inner + .lock() + .unwrap() + .setup_child + .as_ref() + .map(|c| (c.label.clone(), c.job.is_some())); + assert_eq!( + recorded, + Some(("setup child".to_string(), true)), + "the close handler can find it" + ); let deadline = std::time::Instant::now() + Duration::from_secs(10); while !dir.path().join("started").exists() && std::time::Instant::now() < deadline { std::thread::sleep(Duration::from_millis(50)); @@ -8052,7 +8242,7 @@ mod tests { .expect("ran"); assert!(output.status.success()); - assert_eq!(state.inner.lock().unwrap().setup_child_pid, None); + assert!(state.inner.lock().unwrap().setup_child.is_none()); assert!( dir.path().join("started").exists(), "the setup child's own child never ran" @@ -8064,6 +8254,243 @@ mod tests { ); } + /// A setup child that starts a child of its own, which writes `started` in + /// `dir` at once and would write `late` about 3 seconds on, and then runs + /// for about 20 seconds itself. `started` coming from the grandchild means + /// it exists by the time a test sees the file. + fn long_setup_child(dir: &Path) -> Command { + #[cfg(windows)] + let command = { + use std::os::windows::process::CommandExt; + let mut command = Command::new("cmd"); + command.args(["/d", "/c"]).raw_arg( + "start /b cmd /d /c \"type nul > started & ping -n 4 127.0.0.1 >nul & type nul > late\" & ping -n 21 127.0.0.1 >nul", + ); + command + }; + #[cfg(unix)] + let command = { + let mut command = Command::new("sh"); + command.args(["-c", "touch started; (sleep 3; touch late) & sleep 20"]); + command + }; + let mut command = command; + command + .current_dir(dir) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()); + command + } + + /// Runs `long_setup_child` as a setup child labelled `label` on a thread + /// of its own, returning once it has started. The handle yields how long + /// it ran. + fn start_long_setup_child( + state: &std::sync::Arc, + dir: &Path, + label: &'static str, + ) -> std::thread::JoinHandle { + let mut command = long_setup_child(dir); + let state = std::sync::Arc::clone(state); + let runner = std::thread::spawn(move || { + let started = std::time::Instant::now(); + super::run_setup_child(&mut command, &state, label, |mut child| { + let status = child.wait().map_err(|e| e.to_string())?; + Ok(std::process::Output { + status, + stdout: Vec::new(), + stderr: Vec::new(), + }) + }) + .expect("ran"); + started.elapsed() + }); + let deadline = std::time::Instant::now() + Duration::from_secs(10); + while !dir.join("started").exists() && std::time::Instant::now() < deadline { + std::thread::sleep(Duration::from_millis(50)); + } + assert!( + dir.join("started").exists(), + "the setup child never started" + ); + runner + } + + #[test] + fn closing_the_window_mid_setup_ends_the_setup_child() { + // #405: on Windows, closing the window during the model download left + // the warmup running (stop_backend only signalled it on unix). It has + // to end at once, and so does what it started. + let dir = make_tmp(); + let state = std::sync::Arc::new(super::BackendState::default()); + let runner = start_long_setup_child(&state, dir.path(), "model warmup"); + + super::stop_backend(&state); + + let ran = runner.join().expect("joined"); + assert!( + ran < Duration::from_secs(10), + "the setup child ran on for {ran:?}" + ); + #[cfg(windows)] + { + std::thread::sleep(Duration::from_secs(5)); + assert!( + !dir.path().join("late").exists(), + "a process the setup child started outlived the close" + ); + } + } + + #[test] + fn skip_stops_the_model_warmup_and_nothing_else() { + let dir = make_tmp(); + let state = std::sync::Arc::new(super::BackendState::default()); + let runner = start_long_setup_child(&state, dir.path(), super::MODEL_WARMUP_LABEL); + assert!(super::stop_setup_child(&state, super::MODEL_WARMUP_LABEL).expect("stopped")); + let ran = runner.join().expect("joined"); + assert!( + ran < Duration::from_secs(10), + "the warmup ran on for {ran:?}" + ); + + // A pip install is never cut short by Skip: it runs to its end. + let dir = make_tmp(); + let runner = start_long_setup_child(&state, dir.path(), "CUDA torch install"); + assert!(!super::stop_setup_child(&state, super::MODEL_WARMUP_LABEL).expect("no error")); + let ran = runner.join().expect("joined"); + assert!( + ran > Duration::from_secs(15), + "the pip install was stopped at {ran:?}" + ); + + // Nothing running is not an error: the warmup finished first. + assert!(!super::stop_setup_child(&state, super::MODEL_WARMUP_LABEL).expect("no error")); + } + + #[test] + fn a_second_setup_child_is_refused_while_one_runs() { + // Two at once would share the one slot that stops a child (the second + // to start overwrites it, the first to end clears it), and download + // into the same cache side by side. + let dir = make_tmp(); + let state = std::sync::Arc::new(super::BackendState::default()); + let runner = start_long_setup_child(&state, dir.path(), "CUDA torch install"); + + let second = make_tmp(); + let refused = super::run_setup_child( + &mut long_setup_child(second.path()), + &state, + super::MODEL_WARMUP_LABEL, + |_child| unreachable!("the second child must not be started"), + ); + + let err = refused.expect_err("a second setup child was started"); + assert!(err.contains("CUDA torch install is still running"), "{err}"); + super::stop_backend(&state); + runner.join().expect("joined"); + } + + #[test] + fn output_on_stderr_alone_keeps_a_streamed_child_alive() { + // #405: the model warmup's downloads report progress on stderr (tqdm), + // and its stdout is silent for as long as one model takes. Under a + // stall budget that counted stdout only, a healthy slow download was + // killed. Eight seconds of stderr ticks against a three second budget. + #[cfg(windows)] + let mut command = { + use std::os::windows::process::CommandExt; + let mut command = Command::new("cmd"); + command + .args(["/d", "/c"]) + .raw_arg("for /l %i in (1,1,8) do @(echo tick 1>&2 & ping -n 2 127.0.0.1 >nul)"); + command + }; + #[cfg(unix)] + let mut command = { + let mut command = Command::new("sh"); + command.args([ + "-c", + "for i in 1 2 3 4 5 6 7 8; do echo tick >&2; sleep 1; done", + ]); + command + }; + command + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + let child = command.spawn().expect("spawn"); + + let output = super::child_output_streaming( + child, + Duration::from_secs(3), + Duration::from_secs(120), + "stderr only", + |_line| {}, + ) + .expect("a child that keeps writing to stderr must be left to finish"); + + assert!(output.status.success()); + assert_eq!( + String::from_utf8_lossy(&output.stderr) + .lines() + .filter(|line| line.trim() == "tick") + .count(), + 8, + "every stderr line must still be collected" + ); + } + + #[test] + fn the_warmup_stall_budget_outlasts_one_download_try() { + // audio-separator waits up to 300 s for a read before it gives up by + // itself, the longest silence of any download in the warmup. The stall + // budget must outlast that, or it would end the whole warmup where the + // model's own failure, and the models after it, were still to come. + assert!(super::WARMUP_STALL_TIMEOUT >= Duration::from_secs(300 + 30)); + assert!(super::WARMUP_STALL_TIMEOUT <= Duration::from_secs(15 * 60)); + assert!(super::WARMUP_HARD_CAP >= Duration::from_secs(60 * 60)); + } + + #[test] + fn warmup_progress_reads_the_start_lines() { + assert_eq!( + super::warmup_progress("WARMUP_START beat_this 2 5"), + Some(super::WarmupProgress { + model: "beat_this".to_string(), + number: 2, + count: 5, + }) + ); + assert_eq!(super::warmup_progress("WARMUP_OK beat_this"), None); + assert_eq!(super::warmup_progress("WARMUP_START beat_this"), None); + } + + #[test] + fn the_warmup_status_says_what_arrived_and_why_not() { + let status = super::ModelWarmupStatus::from_stdout( + "WARMUP_START demucs 1 5\nWARMUP_OK demucs\n\ + WARMUP_START beat_this 2 5\n\ + WARMUP_FAILED beat_this model download failed: x (DNS lookup failed)\n\ + WARMUP_OFFLINE beat_this\nWARMUP_START whisper 5 5\nWARMUP_OK whisper\n", + ); + assert!(status.demucs_ready && status.whisper_ready); + assert!(!status.beat_this_ready && !status.sections_ready); + assert_eq!(status.offline, vec!["beat_this".to_string()]); + assert_eq!(status.no_space, None); + + let full = super::ModelWarmupStatus::from_stdout("WARMUP_NO_SPACE 200000000 1000\n"); + assert_eq!( + full.no_space, + Some(super::NoSpace { + needed: 200_000_000, + free: 1000 + }) + ); + assert!(!full.demucs_ready && full.offline.is_empty()); + } + #[test] #[cfg(not(target_os = "macos"))] fn the_pip_stall_budget_outlasts_a_quiet_wheel_install() { diff --git a/desktop/ui/index.html b/desktop/ui/index.html index bb124441..55f91904 100644 --- a/desktop/ui/index.html +++ b/desktop/ui/index.html @@ -28,9 +28,12 @@

StemDeck

+
+ +
diff --git a/desktop/ui/setup.css b/desktop/ui/setup.css index cc3cc7f1..5200f89a 100644 --- a/desktop/ui/setup.css +++ b/desktop/ui/setup.css @@ -177,6 +177,15 @@ h1 { 100% { transform: translateX(400%); } } +.notice { + margin: 16px 0 0; + padding: 12px; + border: 1px solid rgba(216, 168, 74, 0.3); + border-radius: 8px; + color: #f1cf82; + background: rgba(216, 168, 74, 0.08); +} + .details { max-height: 180px; overflow: auto; @@ -204,6 +213,11 @@ button { cursor: pointer; } +button:disabled { + opacity: 0.5; + cursor: default; +} + .hidden { display: none; } diff --git a/desktop/ui/setup.js b/desktop/ui/setup.js index e0f06976..270c85f9 100644 --- a/desktop/ui/setup.js +++ b/desktop/ui/setup.js @@ -11,6 +11,9 @@ function startBackend() { const statusEl = document.getElementById("status"); const detailsEl = document.getElementById("details"); const retryBtn = document.getElementById("retry"); +const skipModelsBtn = document.getElementById("skip-models"); +const continueBtn = document.getElementById("continue"); +const noticeEl = document.getElementById("notice"); const steps = [...document.querySelectorAll("[data-step]")]; function setStep(name, state) { @@ -129,6 +132,143 @@ async function startPipProgress(stopTimedMessages) { }; } +// The names app/pipeline/warmup.py gives its models, as this screen says them. +const MODEL_LABELS = { + demucs: "stem separation", + beat_this: "beat detection", + sections: "song sections", + vocal_split: "karaoke split", + whisper: "lyrics transcription", +}; + +function modelsLater(separationMissing) { + return ( + "StemDeck works without them: each one downloads the first time a feature needs it" + + (separationMissing ? ", so the first separation downloads the separation model before it starts." : ".") + ); +} + +/** + * What to tell someone whose models did not all arrive (#405), or null when + * they did. Until this, a missing model only reached the console, and the + * first sign was a feature that quietly did less. + */ +function modelWarmupNotice(status, skipped) { + if (status?.noSpace) { + const gb = (bytes) => (bytes / 1e9).toFixed(1); + return ( + `Not enough free disk space for the AI models: about ${gb(status.noSpace.needed)} GB ` + + `needed, ${gb(status.noSpace.free)} GB free, so none was downloaded. StemDeck works ` + + "without them: each one downloads the first time a feature needs it and there is " + + "room, so the first separation downloads the separation model before it starts." + ); + } + // camelCase, because ModelWarmupStatus is #[serde(rename_all = + // "camelCase")]. Reading the Rust field names instead gives undefineds, + // `missing` is empty every time, and this says nothing. + const missing = [ + ["demucs", status?.demucsReady], + ["beat_this", status?.beatThisReady], + ["sections", status?.sectionsReady], + ["vocal_split", status?.vocalSplitReady], + ["whisper", status?.whisperReady], + ] + .filter(([, ready]) => ready === false) + .map(([name]) => name); + if (!missing.length) return null; + const said = (names) => names.map((name) => MODEL_LABELS[name]).join(", "); + const later = modelsLater(missing.includes("demucs")); + if (skipped) return `AI model download skipped, so these were not downloaded: ${said(missing)}. ${later}`; + const offline = missing.filter((name) => status.offline?.includes(name)); + const failed = missing.filter((name) => !offline.includes(name)); + const causes = []; + if (offline.length) { + causes.push( + "StemDeck could not reach the model download servers (no internet connection, or the " + + `connection stopped answering), so these AI models were not downloaded: ${said(offline)}.` + ); + } + if (failed.length) causes.push(`These AI models could not be downloaded: ${said(failed)}.`); + return `${causes.join(" ")} ${later} logs/setup.log in the StemDeck data folder has the details.`; +} + +/** Shows `text` until the Continue button is pressed. */ +function showNotice(text) { + noticeEl.textContent = text; + noticeEl.classList.remove("hidden"); + continueBtn.classList.remove("hidden"); + continueBtn.focus(); + return new Promise((resolve) => { + continueBtn.addEventListener( + "click", + () => { + continueBtn.classList.add("hidden"); + resolve(); + }, + { once: true } + ); + }); +} + +/** + * The model step: downloads the AI models, saying which one it is on, and + * offers Skip. Returns what to tell the user afterwards, or null. + * + * Best-effort, per model: a model that does not download here falls back to + * its lazy download on first use, so nothing here fails setup. + */ +async function warmupModels() { + let current = "Starting the AI model download..."; + let since = Date.now(); + const show = () => setStatus(`${current} Elapsed: ${formatElapsed(since)}.`); + let skipped = false; + show(); + const timer = window.setInterval(show, 1000); + let unlisten = null; + try { + unlisten = await window.__TAURI__.event.listen("model-warmup-progress", (event) => { + if (skipped) return; + const { model, number, count } = event.payload ?? {}; + // Offered from the first model on: only once the warmup has said + // something is it sure to be running, and so stoppable. + skipModelsBtn.classList.remove("hidden"); + current = `Getting AI model ${number} of ${count}: ${MODEL_LABELS[model] ?? model}...`; + since = Date.now(); + show(); + }); + + skipModelsBtn.disabled = false; + skipModelsBtn.onclick = async () => { + skipModelsBtn.disabled = true; + skipped = true; + current = "Skipping the AI model download..."; + show(); + try { + await invoke("skip_model_warmup"); + } catch (err) { + skipped = false; + skipModelsBtn.disabled = false; + current = `The download could not be stopped (${err}), so it goes on...`; + show(); + } + }; + + try { + return modelWarmupNotice(await invoke("warmup_models"), skipped); + } catch (err) { + // The warmup could not run at all (no Python), or stopped answering. + const what = skipped ? "AI model download skipped." : `The AI model download stopped: ${err}.`; + return `${what} ${modelsLater(true)}`; + } + } finally { + window.clearInterval(timer); + unlisten?.(); + skipModelsBtn.classList.add("hidden"); + skipModelsBtn.onclick = null; + setStatus("AI model step finished."); + } +} + function isMac() { return /mac/i.test(navigator.userAgentData?.platform ?? navigator.platform ?? ""); } @@ -329,9 +469,26 @@ function ffmpegFailureHint(message) { ); } +// Retry can be on screen while setup goes on (a GPU warning does not stop it), +// and a second run beside the first would share the model step's buttons. +let setupRunning = false; + async function runSetup() { + if (setupRunning) return; + setupRunning = true; + try { + await setup(); + } finally { + setupRunning = false; + } +} + +async function setup() { detailsEl.classList.add("hidden"); retryBtn.classList.add("hidden"); + noticeEl.classList.add("hidden"); + continueBtn.classList.add("hidden"); + skipModelsBtn.classList.add("hidden"); for (const step of steps) step.classList.remove("active", "done", "error"); try { @@ -524,59 +681,13 @@ async function runSetup() { // model through torchvision. await repairTorchvision(); + let modelNotice = null; await runStep("model", async () => { - const stopProgress = startProgressStatus([ - { - afterSeconds: 0, - text: "Downloading AI models... this can take a few minutes on first run.", - }, - { - afterSeconds: 60, - text: "Still downloading AI models... slow networks can delay this.", - }, - { - afterSeconds: 300, - text: "Still working... large models can take a while on a slow connection.", - }, - ]); - try { - // Best-effort, per model: any model that doesn't download here just - // falls back to its existing lazy-download-on-first-use behavior -- - // never fails setup over this (warmup_models itself never throws for - // an individual model failure; this catch is only for the whole - // subprocess failing to run at all, e.g. a missing Python). - // - // The per-model result was thrown away here, so a user whose beat - // model never arrived saw a clean setup and then a permanently worse - // beat grid, with nothing anywhere connecting the two (#502). Naming - // it at least puts it in the log that gets attached to bug reports. - // - // Deliberately not setStatus: the next step overwrites the status line - // immediately and then navigates away to the backend, so anything - // written here is invisible. Telling the user properly belongs in the - // main UI, where the affected feature actually lives, and is its own - // piece of work. - // camelCase, because ModelWarmupStatus is #[serde(rename_all = - // "camelCase")]. Reading the Rust field names instead gives four - // undefineds, `missing` is empty every time, and this says nothing. - const status = await invoke("warmup_models"); - const missing = [ - [status?.demucsReady, "stem separation"], - [status?.beatThisReady, "beat detection"], - [status?.sectionsReady, "song sections"], - [status?.vocalSplitReady, "karaoke split"], - ] - .filter(([ready]) => ready === false) - .map(([, label]) => label); - if (missing.length) { - console.warn("models not downloaded during setup:", missing.join(", ")); - } - } catch (err) { - console.warn("model warmup failed (will download lazily on first use):", err); - } finally { - stopProgress(); - } + modelNotice = await warmupModels(); }); + // Before the backend step, which overwrites the status line and then + // leaves this page: anything said only in passing here was never seen. + if (modelNotice) await showNotice(modelNotice); await runStep("backend", async () => { setStatus(gpuSummary ? `${gpuSummary} - starting backend...` : "Starting StemDeck backend..."); diff --git a/tests/e2e/setup-wizard.spec.mjs b/tests/e2e/setup-wizard.spec.mjs new file mode 100644 index 00000000..67099e6d --- /dev/null +++ b/tests/e2e/setup-wizard.spec.mjs @@ -0,0 +1,169 @@ +// The desktop setup wizard's model step (#405): it says which model it is on, +// offers Skip, and says what did not arrive before it leaves the page. +// +// desktop/ui is the page the desktop app opens before any backend exists, so +// it is served here from disk under a made-up origin, with a stand-in for the +// Tauri bridge that the test drives. No Rust runs: the shell's side (stopping +// the warmup, the stall deadline) is covered by its cargo tests. +import { test, expect } from "@playwright/test"; +import { readFile } from "node:fs/promises"; +import path from "node:path"; +import { fileURLToPath } from "node:url"; + +const UI = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "../../desktop/ui"); +const ORIGIN = "http://setup.test"; + +const ALL_READY = { + demucsReady: true, + beatThisReady: true, + sectionsReady: true, + vocalSplitReady: true, + whisperReady: true, + offline: [], + noSpace: null, +}; + +async function openWizard(page) { + await page.route(`${ORIGIN}/**`, async (route) => { + const name = new URL(route.request().url()).pathname.slice(1); + if (name === "started.html") { + return route.fulfill({ contentType: "text/html", body: "

backend

" }); + } + const types = { html: "text/html", js: "text/javascript", css: "text/css" }; + const body = await readFile(path.join(UI, path.basename(name))); + return route.fulfill({ contentType: types[path.extname(name).slice(1)], body }); + }); + await page.addInitScript((origin) => { + const listeners = {}; + let finishWarmup = null; + const answers = { + probe_runtime: { pythonReady: true, ffmpegReady: true, torchDevice: null, pythonPath: "python" }, + runtime_pack_status: { manifest: { version: "1" }, installedVersion: "1" }, + ensure_torch_device: { gpuDetected: false, torchDevice: "cpu", reason: "no-gpu" }, + repair_torchvision: { outcome: "ok" }, + start_backend: { url: `${origin}/started.html` }, + }; + window.__setup = { + calls: [], + emit(name, payload) { + for (const fn of listeners[name] ?? []) fn({ payload }); + }, + finishWarmup(status) { + finishWarmup(status); + }, + }; + window.__TAURI__ = { + core: { + invoke(cmd) { + window.__setup.calls.push(cmd); + if (cmd === "warmup_models") return new Promise((resolve) => (finishWarmup = resolve)); + if (cmd === "skip_model_warmup") { + // What the shell does: the warmup is stopped during model 2, so + // it returns with only the first one ready. + finishWarmup({ demucsReady: true, beatThisReady: false, sectionsReady: false, + vocalSplitReady: false, whisperReady: false, offline: [], noSpace: null }); + } + return Promise.resolve(answers[cmd] ?? null); + }, + }, + event: { + async listen(name, fn) { + (listeners[name] ??= []).push(fn); + return () => { + listeners[name] = listeners[name].filter((other) => other !== fn); + }; + }, + }, + }; + }, ORIGIN); + await page.goto(`${ORIGIN}/index.html`); + await expect(page.locator('[data-step="model"]')).toHaveClass(/active/); + await expect.poll(() => page.evaluate(() => window.__setup.calls)).toContain("warmup_models"); +} + +const finish = (page, status) => page.evaluate((s) => window.__setup.finishWarmup(s), status); + +test("the model step names the model it is on, and Skip says what the first use costs", async ({ page }) => { + await openWizard(page); + const skip = page.locator("#skip-models"); + // Not before the warmup has said anything: until then it may not be running yet. + await expect(skip).toBeHidden(); + + await page.evaluate(() => + window.__setup.emit("model-warmup-progress", { model: "beat_this", number: 2, count: 5 }) + ); + await expect(page.locator("#status")).toContainText("Getting AI model 2 of 5: beat detection..."); + await expect(skip).toBeVisible(); + + await skip.click(); + + const notice = page.locator("#notice"); + await expect(notice).toContainText( + "AI model download skipped, so these were not downloaded: beat detection, song sections, " + + "karaoke split, lyrics transcription. StemDeck works without them: each one downloads " + + "the first time a feature needs it." + ); + // The separation model had arrived before Skip, so nothing about the first separation. + await expect(notice).not.toContainText("first separation"); + await expect(page.locator("#status")).toHaveText("AI model step finished."); + await expect(skip).toBeHidden(); + expect(await page.evaluate(() => window.__setup.calls)).toContain("skip_model_warmup"); + // Read before the backend step leaves the page. + expect(await page.evaluate(() => window.__setup.calls)).not.toContain("start_backend"); + + await page.locator("#continue").click(); + await expect(page).toHaveURL(`${ORIGIN}/started.html`); +}); + +test("offline: the notice names what did not arrive and why", async ({ page }) => { + await openWizard(page); + await finish(page, { + ...ALL_READY, + beatThisReady: false, + sectionsReady: false, + offline: ["beat_this", "sections"], + }); + + const notice = page.locator("#notice"); + await expect(notice).toContainText( + "could not reach the model download servers (no internet connection, or the connection " + + "stopped answering), so these AI models were not downloaded: beat detection, song sections." + ); + await expect(notice).not.toContainText("These AI models could not be downloaded"); + // The separation model arrived, so nothing about the first separation. + await expect(notice).not.toContainText("first separation"); + await page.locator("#continue").click(); + await expect(page).toHaveURL(`${ORIGIN}/started.html`); +}); + +test("a model that failed for another reason is not blamed on the connection", async ({ page }) => { + await openWizard(page); + await finish(page, { ...ALL_READY, demucsReady: false, vocalSplitReady: false, offline: ["demucs"] }); + + const notice = page.locator("#notice"); + await expect(notice).toContainText("so these AI models were not downloaded: stem separation."); + await expect(notice).toContainText("These AI models could not be downloaded: karaoke split."); + await expect(notice).toContainText("the first separation downloads the separation model before it starts"); +}); + +test("a full disk: the notice says how much room the models need", async ({ page }) => { + await openWizard(page); + await finish(page, { + ...ALL_READY, + demucsReady: false, + beatThisReady: false, + sectionsReady: false, + vocalSplitReady: false, + whisperReady: false, + noSpace: { needed: 1_800_000_000, free: 400_000_000 }, + }); + + await expect(page.locator("#notice")).toContainText("about 1.8 GB needed, 0.4 GB free, so none was downloaded"); +}); + +test("everything arrived: setup goes straight on", async ({ page }) => { + await openWizard(page); + await finish(page, ALL_READY); + + await expect(page).toHaveURL(`${ORIGIN}/started.html`); +}); diff --git a/tests/test_model_cache.py b/tests/test_model_cache.py index bd58bd1c..4c526f19 100644 --- a/tests/test_model_cache.py +++ b/tests/test_model_cache.py @@ -149,3 +149,59 @@ def load(): # 2 for the first call (try, heal, retry), then 1 each for the two after it. assert len(attempts) == 4 assert artifact.exists(), "later attempts must not keep deleting it" + + +def test_a_half_written_vocal_split_index_is_fetched_again(tmp_path, monkeypatch): + """audio-separator writes its model list (download_checks.json) in place, so + a setup ended mid-fetch leaves it half written. Its own loader then fails + on that file every time, before any model is looked at.""" + import json + import logging + + from audio_separator.separator import Separator + from audio_separator.separator import separator as separator_module + + from app.core.model_cache import vocal_split_artifacts + + root = tmp_path / "audio-separator" + root.mkdir() + (root / "download_checks.json").write_text('{"vr_download_list": {') + index = { + key: {} + for key in ( + "demucs_download_list", + "vr_download_list", + "mdx_download_list", + "mdx_download_vip_list", + "mdx23c_download_list", + "mdx23c_download_vip_list", + "roformer_download_list", + ) + } + body = json.dumps(index).encode() + + class Response: + status_code = 200 + headers = {"content-length": str(len(body))} + + def iter_content(self, chunk_size): + yield body + + fetched = [] + + def get(url, **kwargs): + fetched.append(url) + return Response() + + monkeypatch.setattr(separator_module.requests, "get", get) + separator = Separator.__new__(Separator) + separator.model_file_dir = str(root) + separator.logger = logging.getLogger("test.audio_separator") + + load_or_heal( + separator.list_supported_model_files, + lambda: vocal_split_artifacts(tmp_path, "UVR_MDXNET_KARA_2.onnx"), + ) + + assert [url.rsplit("/", 1)[1] for url in fetched] == ["download_checks.json"] + assert json.loads((root / "download_checks.json").read_text()) == index diff --git a/tests/test_pipeline_warmup.py b/tests/test_pipeline_warmup.py index b5855150..c5d75675 100644 --- a/tests/test_pipeline_warmup.py +++ b/tests/test_pipeline_warmup.py @@ -1,8 +1,108 @@ +import hashlib import os +import socket +import subprocess +import sys +import threading +import time +import urllib.request +from pathlib import Path from unittest.mock import patch +import pytest +import torch.hub + from app.pipeline import warmup +ROOT = Path(__file__).resolve().parents[1] +# The real one, before the fixture below replaces it for every test. +_FREE_BYTES = warmup._free_bytes + + +@pytest.fixture(autouse=True) +def _model_folders(monkeypatch, tmp_path): + """The space check reads the models folder and the disk: the folder is + this test's own, with room to spare unless a test says not. And no + request to the local test servers goes through a proxy this machine may + have set, which would also make it a real network request.""" + models = tmp_path / "models" + monkeypatch.setattr(warmup, "MODELS_DIR", models) + monkeypatch.setattr(warmup, "_free_bytes", lambda: 1 << 50) + for name in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"): + monkeypatch.delenv(name, raising=False) + monkeypatch.delenv(name.lower(), raising=False) + monkeypatch.setenv("NO_PROXY", "127.0.0.1,localhost") + return models + + +class _Server: + """A local HTTP server on a throwaway port. `respond(conn)` answers each + request; `sent` counts the body bytes it has written.""" + + def __init__(self, respond): + self._respond = respond + self.sent = 0 + self.released = threading.Event() + self._sock = socket.create_server(("127.0.0.1", 0)) + self.port = self._sock.getsockname()[1] + threading.Thread(target=self._serve, daemon=True).start() + + def url(self, name: str) -> str: + return f"http://127.0.0.1:{self.port}/models/{name}" + + def _serve(self) -> None: + while True: + try: + conn, _ = self._sock.accept() + except OSError: + return + threading.Thread(target=self._handle, args=(conn,), daemon=True).start() + + def _handle(self, conn: socket.socket) -> None: + request = b"" + while b"\r\n\r\n" not in request: + chunk = conn.recv(4096) + if not chunk: + return + request += chunk + try: + self._respond(self, conn) + except OSError: + pass + finally: + conn.close() + + def close(self) -> None: + self.released.set() + self._sock.close() + + +@pytest.fixture() +def serve(): + servers: list[_Server] = [] + + def start(respond) -> _Server: + server = _Server(respond) + servers.append(server) + return server + + yield start + for server in servers: + server.close() + + +def _send_slowly(body: bytes, chunk: int, pause: float): + def respond(server: _Server, conn: socket.socket) -> None: + conn.sendall(f"HTTP/1.1 200 OK\r\nContent-Length: {len(body)}\r\n\r\n".encode()) + for start in range(0, len(body), chunk): + if server.released.is_set(): + return + conn.sendall(body[start : start + chunk]) + server.sent += min(chunk, len(body) - start) + time.sleep(pause) + + return respond + def test_section_warmup_loads_cpu_model(): with patch("allin1_infer.models.load_pretrained_model") as load: @@ -16,7 +116,7 @@ def test_warmup_continues_after_individual_failure(monkeypatch, capsys): def fail(): calls.append("fail") - raise RuntimeError("offline") + raise RuntimeError("bad checkpoint") def succeed(): calls.append("succeed") @@ -25,19 +125,235 @@ def succeed(): assert warmup.main() == 0 assert calls == ["fail", "succeed"] + # Each model is announced before it runs, so the setup screen can say + # which one it is on (#405). A failure that is not the network's is not + # reported as one. + assert capsys.readouterr().out.splitlines() == [ + "WARMUP_START sections 1 2", + "WARMUP_FAILED sections bad checkpoint", + "WARMUP_START demucs 2 2", + "WARMUP_OK demucs", + ] + + +def test_no_network_is_told_apart_from_other_failures(monkeypatch, capsys): + """Offline: no name resolves. The step's error is whatever its library + wraps that in, here urllib's URLError.""" + + def no_dns(*args, **kwargs): + raise socket.gaierror(11001, "getaddrinfo failed") + + monkeypatch.setattr(socket, "getaddrinfo", no_dns) + monkeypatch.setattr( + warmup, + "_STEPS", + (("beat_this", lambda: urllib.request.urlopen("http://models.example/beat.ckpt")),), + ) + + warmup.main() + + lines = capsys.readouterr().out.splitlines() + assert lines[0] == "WARMUP_START beat_this 1 1" + assert lines[1].startswith("WARMUP_FAILED beat_this ") + assert lines[2:] == ["WARMUP_OFFLINE beat_this"] + + +def test_a_host_that_answers_with_an_error_is_not_offline(serve, monkeypatch, capsys): + def not_found(server, conn): + conn.sendall(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n") + + server = serve(not_found) + monkeypatch.setattr( + warmup, "_STEPS", (("demucs", lambda: urllib.request.urlopen(server.url("x.th"))),) + ) + + warmup.main() + + lines = capsys.readouterr().out.splitlines() + assert lines[1].startswith("WARMUP_FAILED demucs HTTP Error 404") + assert len(lines) == 2 + + +def test_a_stalled_model_host_ends_the_model_instead_of_the_step( + serve, monkeypatch, tmp_path, capsys +): + """The host takes the connection and sends nothing. The model fails the way + the separation model's download does (#125), with its file and folder + named, and the next model still gets its turn, instead of this one holding + the whole step until the shell's deadline.""" + server = serve(lambda server, conn: server.released.wait(20)) + dst = tmp_path / "checkpoints" / "model-deadbeef.th" + dst.parent.mkdir() + monkeypatch.setattr(warmup, "TIMEOUT_MODEL_DOWNLOAD_READ", 1) + monkeypatch.setattr(warmup, "MODEL_DOWNLOAD_ATTEMPTS", 1) + monkeypatch.setattr( + warmup, + "_STEPS", + ( + ( + "demucs", + lambda: torch.hub.download_url_to_file( + server.url(dst.name), str(dst), progress=False + ), + ), + ("beat_this", lambda: None), + ), + ) + + started = time.monotonic() + warmup.main() + elapsed = time.monotonic() - started + + lines = capsys.readouterr().out.splitlines() + assert elapsed < 10, f"the stalled download held the step for {elapsed:.0f} s" + assert lines[1].startswith("WARMUP_FAILED demucs model download failed: ") + assert "(no data for 1 s)" in lines[1] + assert lines[2:] == [ + "WARMUP_OFFLINE demucs", + "WARMUP_START beat_this 2 2", + "WARMUP_OK beat_this", + ] + assert not dst.exists() + + +def test_not_enough_space_downloads_nothing_and_says_how_much(monkeypatch, capsys): + calls = [] + monkeypatch.setattr(warmup, "_STEPS", (("demucs", lambda: calls.append("demucs")),)) + monkeypatch.setattr(warmup, "_free_bytes", lambda: 1_000) + + assert warmup.main() == 0 + + assert calls == [] + needed = warmup._DOWNLOAD_BYTES["demucs"] + assert capsys.readouterr().out.splitlines() == [f"WARMUP_NO_SPACE {needed} 1000"] + + +def test_models_already_on_disk_need_no_space(monkeypatch, capsys, _model_folders): + """A second setup run (a new version) on a nearly full disk still checks + the models it has instead of refusing for want of room it does not need.""" + monkeypatch.setattr(warmup, "_DOWNLOAD_BYTES", {"demucs": 4_000}) + monkeypatch.setattr(warmup, "_STEPS", (("demucs", lambda: None),)) + monkeypatch.setattr(warmup, "_free_bytes", lambda: 1_000) + checkpoints = _model_folders / "torch" / "hub" / "checkpoints" + checkpoints.mkdir(parents=True) + (checkpoints / "demucs.th").write_bytes(b"\0" * 4_000) + + warmup.main() + + assert capsys.readouterr().out.splitlines() == ["WARMUP_START demucs 1 1", "WARMUP_OK demucs"] + + +def test_a_download_cut_short_is_not_counted_as_a_model(monkeypatch, capsys, _model_folders): + monkeypatch.setattr(warmup, "_DOWNLOAD_BYTES", {"demucs": 4_000}) + monkeypatch.setattr(warmup, "_STEPS", (("demucs", lambda: None),)) + monkeypatch.setattr(warmup, "_free_bytes", lambda: 1_000) + checkpoints = _model_folders / "torch" / "hub" / "checkpoints" + checkpoints.mkdir(parents=True) + (checkpoints / "demucs.th.0123abcd.partial").write_bytes(b"\0" * 4_000) + + warmup.main() + + assert capsys.readouterr().out.splitlines() == ["WARMUP_NO_SPACE 4000 1000"] + + +def test_a_space_check_that_cannot_run_does_not_stop_the_downloads(monkeypatch, capsys): + def broken(): + raise OSError("settings unreadable") + + monkeypatch.setattr(warmup, "_bytes_needed", broken) + monkeypatch.setattr(warmup, "_STEPS", (("demucs", lambda: None),)) + + warmup.main() + assert capsys.readouterr().out.splitlines() == [ - "WARMUP_FAILED sections offline", + "WARMUP_NOTE space not checked: settings unreadable", + "WARMUP_START demucs 1 1", "WARMUP_OK demucs", ] +def test_free_space_is_read_before_the_models_folder_exists(tmp_path): + folder = tmp_path / "data" / "models" + with patch.object(warmup, "MODELS_DIR", folder): + assert _FREE_BYTES() > 0 + assert not folder.exists() + + +def test_whisper_counts_only_when_lyrics_would_be_transcribed(monkeypatch): + monkeypatch.setattr(warmup, "_STEPS", (("whisper", lambda: None),)) + monkeypatch.setattr(warmup, "_whisper_model", lambda: None) + assert warmup._bytes_needed() == 0 + + monkeypatch.setattr(warmup, "_whisper_model", lambda: "turbo") + assert warmup._bytes_needed() == warmup._DOWNLOAD_BYTES["turbo"] + + +_KILL_CHILD = """ +import sys, torch.hub +from app.pipeline import warmup +url, dst = sys.argv[1], sys.argv[2] +warmup._bytes_needed = lambda: 0 +warmup._STEPS = (("demucs", lambda: torch.hub.download_url_to_file(url, dst, progress=False)),) +sys.exit(warmup.main()) +""" + + +def test_a_download_stopped_midway_leaves_no_file_that_looks_complete(serve, tmp_path): + """The wizard's Skip, and closing the window, end the warmup mid-file + with a hard kill. What decides whether that is safe is where the + half file sits: torch.hub, which fetches the separation and beat models, + writes to a temporary name and moves the file into place only once it is + whole, so the next run, or the first Split, downloads it again instead of + loading half a checkpoint.""" + body = os.urandom(8 << 20) + server = serve(_send_slowly(body, 64 << 10, 0.05)) + checkpoints = tmp_path / "models" / "torch" / "hub" / "checkpoints" + checkpoints.mkdir(parents=True) + dst = checkpoints / "htdemucs-0123abcd.th" + env = { + **os.environ, + "STEMDECK_DATA_DIR": str(tmp_path), + "TORCH_HOME": str(tmp_path / "models" / "torch"), + "PYTHONUNBUFFERED": "1", + } + child = subprocess.Popen( + [sys.executable, "-c", _KILL_CHILD, server.url(dst.name), str(dst)], + cwd=ROOT, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + ) + try: + deadline = time.monotonic() + 60 + while server.sent < (1 << 20) and time.monotonic() < deadline: + assert child.poll() is None, "the download ended before it could be stopped" + time.sleep(0.05) + assert server.sent >= (1 << 20), "the download never started" + finally: + child.kill() + child.wait(timeout=30) + assert child.stdout is not None + child.stdout.close() + + assert not dst.exists(), "a half download sits where a whole checkpoint would" + leftovers = [p.name for p in checkpoints.iterdir()] + assert all(name.endswith(".partial") for name in leftovers), leftovers + + # And the next try is a fresh, whole download. + server.released.set() + whole = serve(_send_slowly(body, 1 << 20, 0)) + torch.hub.download_url_to_file( + whole.url(dst.name), str(dst), hashlib.sha256(body).hexdigest()[:8], progress=False + ) + assert dst.read_bytes() == body + + def test_section_warmup_disables_hugging_face_symlinks(monkeypatch): """Unelevated Windows cannot create the cache symlinks the hub prefers. Stubbed rather than patched through the real package so the guarantee is checked even where the optional inference dependency is not installed. """ - import sys import types monkeypatch.delenv("HF_HUB_DISABLE_SYMLINKS", raising=False) diff --git a/tests/test_vocal_split_status.py b/tests/test_vocal_split_status.py index 7ffb016a..f87c2df8 100644 --- a/tests/test_vocal_split_status.py +++ b/tests/test_vocal_split_status.py @@ -165,6 +165,8 @@ def broken(): raise ImportError("No module named 'torchvision'") monkeypatch.setattr(warmup, "_STEPS", (("vocal_split", broken),)) + # The space check reads this machine's disk and model folders; not here. + monkeypatch.setattr(warmup, "_bytes_needed", lambda: 0) assert warmup.main() == 0 assert "WARMUP_FAILED vocal_split" in capsys.readouterr().out @@ -177,6 +179,8 @@ def test_a_good_warmup_clears_an_old_failure(monkeypatch): vss.record(False, DEPENDENCY_MISSING, source="warmup") monkeypatch.setattr(warmup, "_STEPS", (("vocal_split", lambda: None),)) + # The space check reads this machine's disk and model folders; not here. + monkeypatch.setattr(warmup, "_bytes_needed", lambda: 0) warmup.main() vss.reset_state()