From 7b22a8584f44fd7ea4f9e2486efc0726b07b4f03 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 27 Jul 2026 11:44:37 +0000 Subject: [PATCH] Fix the readiness probe deadlock and marker lifecycle for PR #7490 The runtime probe waits on child exit with stdout piped and undrained, so a CLI printing more than the pipe buffer holds blocks mid-write and the probe times out. It imports the whole backend, which is exactly what emits stray stdout, so a healthy install gets reported broken and sent to repair after a 10s stall on every launch. Drain concurrently instead, here and in the desktop-capabilities probe next to it. The install marker was cleared only on success, so a failed update, or a cancelled elevation prompt, left it behind and pinned a working install into repair it may not be able to finish offline. Clear it on every terminal outcome and treat marker bookkeeping as best effort rather than failing the install over it. Skip pip flag lines in studio.txt: an unparseable line reported the install broken, and repair reinstalls the same file, so that never cleared. Allow prereleases against a floor for the same reason. --- studio/src-tauri/src/install.rs | 44 ++++++++----------- studio/src-tauri/src/preflight.rs | 39 ++++++++++++++++ studio/src-tauri/src/preflight/managed.rs | 29 +++++++++--- unsloth_cli/commands/studio.py | 17 +++++-- .../tests/test_studio_runtime_readiness.py | 38 ++++++++++++++++ 5 files changed, 131 insertions(+), 36 deletions(-) diff --git a/studio/src-tauri/src/install.rs b/studio/src-tauri/src/install.rs index 72c610ffcc..27996d30a9 100644 --- a/studio/src-tauri/src/install.rs +++ b/studio/src-tauri/src/install.rs @@ -82,6 +82,17 @@ fn clear_install_in_progress_marker() -> Result<(), String> { clear_install_in_progress_marker_at(&path) } +/// The marker means "a run started and never reported an outcome", so every +/// terminal outcome clears it, failures included. The runtime probe is the +/// backstop for a broken venv, whereas a marker left behind by a failed update +/// pins a working install into a repair it may not be able to finish offline. +/// Never fatal: losing the marker costs a fast path, not correctness. +fn clear_install_marker_best_effort() { + if let Err(msg) = clear_install_in_progress_marker() { + warn!("[install] {}", msg); + } +} + fn clear_install_in_progress_marker_at(path: &Path) -> Result<(), String> { match std::fs::remove_file(path) { Ok(()) => Ok(()), @@ -496,18 +507,7 @@ fn run_install_with_event_mode( ); if let Err(msg) = create_install_in_progress_marker() { - diagnostics::finish_attempt( - &diagnostics, - &attempt, - None, - false, - Some(format!("create_install_marker: {msg}")), - ); - clear_current_attempt(&state); - if event_mode.emit_terminal_events() { - emit_failed(&app, &msg); - } - return Err(msg); + warn!("[install] {}", msg); } let (stdout, stderr) = match spawn_script(&script, &args, &state) { @@ -542,20 +542,7 @@ fn run_install_with_event_mode( match result { Ok((status, _)) if status.success() => { - if let Err(msg) = clear_install_in_progress_marker() { - diagnostics::finish_attempt( - &diagnostics, - &attempt, - Some(status.to_string()), - false, - Some(format!("clear_install_marker: {msg}")), - ); - clear_current_attempt(&state); - if event_mode.emit_terminal_events() { - emit_failed(&app, &msg); - } - return Err(msg); - } + clear_install_marker_best_effort(); diagnostics::finish_attempt( &diagnostics, &attempt, @@ -589,6 +576,7 @@ fn run_install_with_event_mode( let _ = app.emit(event_mode.needs_elevation_event(), &packages); Err("NEEDS_ELEVATION".to_string()) } else { + clear_install_marker_best_effort(); let msg = format!("Installer exited with code {}", code); diagnostics::finish_attempt( &diagnostics, @@ -611,6 +599,7 @@ fn run_install_with_event_mode( Err(msg) } Err(msg) => { + clear_install_marker_best_effort(); diagnostics::finish_attempt(&diagnostics, &attempt, None, false, Some(msg.clone())); clear_current_attempt(&state); if event_mode.emit_terminal_events() { @@ -650,6 +639,9 @@ pub fn record_pending_elevation_canceled( let Some(attempt) = attempt else { return false; }; + // The elevation exit left the marker in place for the resumed run that is + // now not happening. + clear_install_marker_best_effort(); diagnostics::finish_attempt( diagnostics, &attempt, diff --git a/studio/src-tauri/src/preflight.rs b/studio/src-tauri/src/preflight.rs index 1e129e85b0..e7abaa6715 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -636,6 +636,45 @@ exit 1 } } + #[cfg(unix)] + #[tokio::test] + async fn managed_runtime_probe_survives_more_stdout_than_the_pipe_holds() { + // The probe imports the whole backend, so import-time output can exceed + // the 64 KiB pipe buffer. Waiting on exit before draining wedges the + // child mid-write and reports a healthy install as broken. + let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await; + let _cache_home = ManagedCapabilityCacheHome::new("runtime-noisy"); + remove_managed_capability_cache(); + + let fake = fake_cli( + "runtime-noisy", + r#"#!/bin/sh +if [ "$1" = "studio" ] && [ "$2" = "desktop-runtime-check" ] && [ "$3" = "--json" ]; then + awk 'BEGIN { while (i++ < 4000) print "import chatter on stdout" }' + printf '{"runtime_ready":true}\n' + exit 0 +fi +if [ "$1" = "studio" ] && [ "$2" = "desktop-capabilities" ] && [ "$3" = "--json" ]; then + printf '{"desktop_protocol_version":1,"desktop_manageability_version":1,"supports_api_only":true,"supports_provision_desktop_auth":true,"supports_desktop_backend_ownership":true,"version":"2026.5.3"}' + exit 0 +fi +exit 1 +"#, + ); + + let started = std::time::Instant::now(); + let probe = probe_managed_bin(fake.bin.clone()).await; + assert!( + matches!(probe, ManagedProbe::Ready { .. }), + "noisy but healthy CLI must probe Ready, got {probe:?}" + ); + assert!( + started.elapsed() < std::time::Duration::from_secs(9), + "probe hit the 10s timeout instead of draining stdout" + ); + remove_managed_capability_cache(); + } + #[cfg(unix)] #[tokio::test] async fn managed_runtime_probe_runs_before_capability_cache() { diff --git a/studio/src-tauri/src/preflight/managed.rs b/studio/src-tauri/src/preflight/managed.rs index 57b8b53c81..a45131584a 100644 --- a/studio/src-tauri/src/preflight/managed.rs +++ b/studio/src-tauri/src/preflight/managed.rs @@ -306,11 +306,21 @@ async fn probe_cli_runtime(bin: &Path) -> Result<(), String> { return Err("studio_runtime_probe_failed".to_string()); }; + // Drain while waiting: the backend import this probe runs can print more than + // the pipe buffer holds, and waiting first deadlocks on it, timing out a + // healthy install into repair. + let reader = tokio::spawn(async move { + let mut buffer = Vec::new(); + let _ = stdout.read_to_end(&mut buffer).await; + buffer + }); + let status = match tokio::time::timeout(Duration::from_secs(10), child.wait()).await { Ok(Ok(status)) => status, _ => { let _ = child.kill().await; let _ = child.wait().await; + reader.abort(); info!( "Managed runtime probe timed out in {}ms", started.elapsed().as_millis() @@ -319,10 +329,9 @@ async fn probe_cli_runtime(bin: &Path) -> Result<(), String> { } }; - let mut output = Vec::new(); - if stdout.read_to_end(&mut output).await.is_err() { + let Ok(output) = reader.await else { return Err("studio_runtime_probe_failed".to_string()); - } + }; let payload = String::from_utf8_lossy(&output) .lines() .rev() @@ -389,11 +398,19 @@ async fn probe_cli_capability(bin: &Path) -> Option { return None; }; + // Same reason as the runtime probe above. + let reader = tokio::spawn(async move { + let mut buffer = Vec::new(); + let _ = stdout.read_to_end(&mut buffer).await; + buffer + }); + match tokio::time::timeout(Duration::from_secs(10), child.wait()).await { Ok(Ok(status)) if status.success() => {} Err(_) => { let _ = child.kill().await; let _ = child.wait().await; + reader.abort(); info!( "Managed desktop-capabilities probe timed out in {}ms", started.elapsed().as_millis() @@ -401,6 +418,7 @@ async fn probe_cli_capability(bin: &Path) -> Option { return None; } _ => { + reader.abort(); info!( "Managed desktop-capabilities probe exited unsuccessfully in {}ms", started.elapsed().as_millis() @@ -409,10 +427,9 @@ async fn probe_cli_capability(bin: &Path) -> Option { } } - let mut output = Vec::new(); - if stdout.read_to_end(&mut output).await.is_err() { + let Ok(output) = reader.await else { return None; - } + }; let capability = serde_json::from_slice::(&output).ok(); info!( diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index d6ccdf1a83..cf5215e69a 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -272,21 +272,30 @@ def _load_run_module(): def _missing_studio_requirement(run_mod): from importlib.metadata import PackageNotFoundError, distribution - from packaging.requirements import Requirement + from packaging.requirements import InvalidRequirement, Requirement requirements = Path(run_mod.__file__).with_name("requirements") / "studio.txt" for line in requirements.read_text(encoding = "utf-8").splitlines(): line = line.partition("#")[0].strip() - if not line: + # pip flags (-r, --extra-index-url) are not requirements. Skipping them + # matters: an unparseable line would report the install broken, and + # repair reinstalls the same file, so the failure would never clear. + if not line or line.startswith("-"): + continue + try: + requirement = Requirement(line) + except InvalidRequirement: continue - requirement = Requirement(line) if requirement.marker and not requirement.marker.evaluate(): continue try: installed = distribution(requirement.name) except PackageNotFoundError: return requirement.name - if requirement.specifier and not requirement.specifier.contains(installed.version): + # prereleases=True: a prerelease satisfying a floor is not a broken install. + if requirement.specifier and not requirement.specifier.contains( + installed.version, prereleases = True + ): return requirement.name return None diff --git a/unsloth_cli/tests/test_studio_runtime_readiness.py b/unsloth_cli/tests/test_studio_runtime_readiness.py index 23c3ad80b2..9eace51d60 100644 --- a/unsloth_cli/tests/test_studio_runtime_readiness.py +++ b/unsloth_cli/tests/test_studio_runtime_readiness.py @@ -64,6 +64,44 @@ def test_desktop_runtime_check_catches_later_startup_dependency(monkeypatch, cap assert payload["module"] == "definitely-missing-studio-package" +def test_desktop_runtime_check_ignores_pip_flag_lines(monkeypatch, capsys, tmp_path): + """An unparseable line must not fail the check: repair reinstalls the same + file, so the install would be declared broken forever.""" + studio = importlib.import_module("unsloth_cli.commands.studio") + backend = tmp_path / "backend" + requirements = backend / "requirements" + requirements.mkdir(parents = True) + (requirements / "studio.txt").write_text( + "--extra-index-url https://example.invalid/simple\n-r base.txt\n", + encoding = "utf-8", + ) + run_mod = SimpleNamespace(__file__ = str(backend / "run.py")) + monkeypatch.setattr(studio, "_load_run_module", lambda: run_mod) + + studio.desktop_runtime_check(_json_output = True) + + assert json.loads(capsys.readouterr().out) == {"runtime_ready": True} + + +def test_desktop_runtime_check_accepts_a_prerelease_over_a_floor(monkeypatch, capsys, tmp_path): + studio = importlib.import_module("unsloth_cli.commands.studio") + backend = tmp_path / "backend" + requirements = backend / "requirements" + requirements.mkdir(parents = True) + (requirements / "studio.txt").write_text("example-package>=1.0\n", encoding = "utf-8") + run_mod = SimpleNamespace(__file__ = str(backend / "run.py")) + monkeypatch.setattr(studio, "_load_run_module", lambda: run_mod) + monkeypatch.setattr( + importlib.import_module("importlib.metadata"), + "distribution", + lambda _name: SimpleNamespace(version = "2.0.0b1"), + ) + + studio.desktop_runtime_check(_json_output = True) + + assert json.loads(capsys.readouterr().out) == {"runtime_ready": True} + + def test_desktop_runtime_check_rejects_version_mismatch(monkeypatch, capsys, tmp_path): studio = importlib.import_module("unsloth_cli.commands.studio") backend = tmp_path / "backend"