diff --git a/studio/src-tauri/src/desktop_backend_owner.rs b/studio/src-tauri/src/desktop_backend_owner.rs index c7d0a7b309..95b606b111 100644 --- a/studio/src-tauri/src/desktop_backend_owner.rs +++ b/studio/src-tauri/src/desktop_backend_owner.rs @@ -265,7 +265,21 @@ impl BackendOwnerState { self.write() } + /// Only while the file is still ours. Cleanup after a backend exits can run + /// for seconds terminating descendants, and a start in that window has + /// already written the next backend's metadata to this same path. pub(crate) fn remove(self) { + if let Ok(Some(on_disk)) = read_metadata(&self.path) { + if on_disk.token_sha256 != self.metadata.token_sha256 { + warn!( + "Leaving desktop backend owner metadata at {}: a newer backend owns it", + self.path.display() + ); + return; + } + } + // Unreadable counts as ours: a file nothing can parse would otherwise + // outlive every backend and fail each later ownership check. remove_metadata_file(&self.path); } @@ -978,6 +992,34 @@ mod tests { dir.join("desktop_backend.json") } + #[test] + fn a_replaced_owner_file_survives_the_previous_backend_cleanup() { + // Cleanup after an exit can take seconds; a start in that window has + // already written the next backend's metadata to the same path. + let path = temp_metadata_path("replaced-owner"); + let exited = BackendOwnerState::from_metadata(path.clone(), metadata(1, Some(8888))); + let mut current = metadata(2, Some(8899)); + current.token = "next-backend-token".to_string(); + current.token_sha256 = token_sha256("next-backend-token"); + write_metadata(&path, ¤t).unwrap(); + + exited.remove(); + + let left = read_metadata(&path).unwrap().unwrap(); + assert_eq!(left.token_sha256, token_sha256("next-backend-token")); + } + + #[test] + fn an_owner_still_holding_the_file_removes_it() { + let path = temp_metadata_path("own-owner"); + let owner = BackendOwnerState::from_metadata(path.clone(), metadata(1, Some(8888))); + write_metadata(&path, &metadata(1, Some(8888))).unwrap(); + + owner.remove(); + + assert!(read_metadata(&path).unwrap().is_none()); + } + fn closed_port() -> u16 { let listener = std::net::TcpListener::bind(("127.0.0.1", 0)).unwrap(); let port = listener.local_addr().unwrap().port(); diff --git a/studio/src-tauri/src/install.rs b/studio/src-tauri/src/install.rs index 3d4f8a3e9d..fc149aba0f 100644 --- a/studio/src-tauri/src/install.rs +++ b/studio/src-tauri/src/install.rs @@ -193,6 +193,10 @@ fn emit_complete(app: &AppHandle) { // ── Spawn ── +/// A start rejected because one is already running. Distinguished from a real +/// spawn failure so the loser of the race leaves the winner's marker alone. +const INSTALL_ALREADY_RUNNING: &str = "Installation is already running."; + /// Spawns the install script in a process group. /// Returns (stdout, stderr) handles for streaming. /// The GroupChild is stored in state so stop_install() can kill the entire tree. @@ -209,7 +213,7 @@ fn spawn_script( > { let mut install = state.lock().map_err(|e| e.to_string())?; if install.child.is_some() { - return Err("Installation is already running.".to_string()); + return Err(INSTALL_ALREADY_RUNNING.to_string()); } install.intentional_stop = false; install.needed_packages.clear(); @@ -514,8 +518,12 @@ fn run_install_with_event_mode( let (stdout, stderr) = match spawn_script(&script, &args, &state) { Ok(handles) => handles, Err(msg) => { - // Nothing ran, so nothing is half-installed. - clear_install_marker_best_effort(); + // Nothing ran, so nothing is half-installed. Unless another + // installer is already inside the venv, in which case the marker + // is its own and clearing it here throws away its recovery signal. + if msg != INSTALL_ALREADY_RUNNING { + 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 599f5c0fa2..3032bb4e63 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -410,6 +410,7 @@ mod tests { for reason in [ "studio_runtime_missing_dependency", "studio_runtime_import_failed", + "studio_runtime_check_unsupported", "install_incomplete", "desktop_backend_version_too_old", ] { @@ -750,10 +751,12 @@ exit 1 #[cfg(unix)] #[tokio::test] - async fn managed_cli_predating_the_runtime_check_is_not_forced_into_repair() { + async fn a_cli_predating_the_runtime_check_is_stale_for_its_own_reason() { // Every published CLI satisfies MIN_DESKTOP_BACKEND_VERSION but has no // `desktop-runtime-check`, so click exits 2 with an empty stdout. That - // must not repair an install that still launches Studio. + // is the CLI the interrupted installs shipped with, and -h passes + // without touching the backend, so launchable is not ready. Update it + // first, and keep the reason apart from a CLI that cannot answer. let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await; let _cache_home = ManagedCapabilityCacheHome::new("runtime-check-absent"); remove_managed_capability_cache(); @@ -775,11 +778,12 @@ exit 1 ); let probe = probe_managed_bin(old.bin.clone()).await; assert!( - matches!(probe, ManagedProbe::Ready { .. }), - "CLI without the runtime-check subcommand must stay Ready, got {probe:?}" + matches!(&probe, ManagedProbe::Stale { reason, .. } + if reason == "studio_runtime_check_unsupported"), + "CLI without the runtime-check subcommand must be stale for it, got {probe:?}" ); - // The fallback must not rescue a CLI that cannot launch at all. + // And a CLI that cannot launch at all keeps the plain probe failure. remove_managed_capability_cache(); let broken = fake_cli( "runtime-check-unlaunchable", diff --git a/studio/src-tauri/src/preflight/managed.rs b/studio/src-tauri/src/preflight/managed.rs index 08cca786ce..4a6e552fcd 100644 --- a/studio/src-tauri/src/preflight/managed.rs +++ b/studio/src-tauri/src/preflight/managed.rs @@ -303,10 +303,16 @@ async fn run_cli_probe(bin: &Path, args: &[&str]) -> bool { } } +/// The probe could not answer: the binary would not run, or it exited without +/// a payload. Distinct from a CLI that is simply too old to know the command. +const RUNTIME_PROBE_FAILED: &str = "studio_runtime_probe_failed"; /// A hung binary, kept out of the fallback below: an older CLI rejects the /// unknown command at once, so only a broken one reaches the timeout, and /// running two more probes on it would treble the wait before repair. const RUNTIME_PROBE_TIMEOUT: &str = "studio_runtime_probe_timeout"; +/// A CLI predating the runtime probe. Repairable, and the update installs one +/// that can answer, so the next preflight checks the venv for real. +const RUNTIME_CHECK_UNSUPPORTED: &str = "studio_runtime_check_unsupported"; /// True when the CLI is older than `desktop-runtime-check` yet still launches. /// Such a CLI exits with a usage error and no JSON, which is indistinguishable @@ -347,10 +353,10 @@ async fn probe_cli_runtime(bin: &Path) -> Result<(), String> { "Managed runtime probe failed to spawn in {}ms", started.elapsed().as_millis() ); - return Err("studio_runtime_probe_failed".to_string()); + return Err(RUNTIME_PROBE_FAILED.to_string()); }; let Some(mut stdout) = child.stdout.take() else { - return Err("studio_runtime_probe_failed".to_string()); + return Err(RUNTIME_PROBE_FAILED.to_string()); }; // Drain while waiting: the backend import this probe runs can print more than @@ -377,7 +383,7 @@ async fn probe_cli_runtime(bin: &Path) -> Result<(), String> { }; let Ok(output) = reader.await else { - return Err("studio_runtime_probe_failed".to_string()); + return Err(RUNTIME_PROBE_FAILED.to_string()); }; let payload = String::from_utf8_lossy(&output) .lines() @@ -396,11 +402,11 @@ async fn probe_cli_runtime(bin: &Path) -> Result<(), String> { } Some("backend_import_failed") => "studio_runtime_import_failed", Some("backend_startup_failed") => super::STUDIO_RUNTIME_STARTUP_FAILED, - _ => "studio_runtime_probe_failed", + _ => RUNTIME_PROBE_FAILED, }; Err(reason.to_string()) } - None => Err("studio_runtime_probe_failed".to_string()), + None => Err(RUNTIME_PROBE_FAILED.to_string()), }; info!( "Managed runtime probe finished ok={} in {}ms", @@ -520,21 +526,21 @@ pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe { // fingerprint only proves protocol compatibility, not that Studio's backend // imports are complete after an interrupted dependency transaction. if let Err(reason) = probe_cli_runtime(&bin).await { - // Every released CLI predates this subcommand, so a missing-command exit - // must not strand an install the launch probe still accepts. - if reason != "studio_runtime_probe_failed" || !predates_runtime_check(&bin).await { - info!( - "Managed preflight: runtime unusable for {:?} reason={} in {}ms", - bin, - reason, - started.elapsed().as_millis() - ); - return ManagedProbe::Stale { bin, reason }; - } + // A CLI too old for this subcommand is the one the interrupted installs + // shipped with, and -h passes without touching the backend, so accepting + // it here leaves those users on the same crash. Update it, then ask. + let reason = if reason == RUNTIME_PROBE_FAILED && predates_runtime_check(&bin).await { + RUNTIME_CHECK_UNSUPPORTED.to_string() + } else { + reason + }; info!( - "Managed preflight: cli predates the runtime probe for {:?}; using the launch probe", - bin + "Managed preflight: runtime unusable for {:?} reason={} in {}ms", + bin, + reason, + started.elapsed().as_millis() ); + return ManagedProbe::Stale { bin, reason }; } if let Some(fingerprint) = managed_bin_fingerprint(&bin) { diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 121dbb9649..6baea99f74 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -309,14 +309,11 @@ def _missing_studio_requirement(run_mod): from importlib.metadata import PackageNotFoundError, distribution from packaging.requirements import InvalidRequirement, Requirement - pending = [(root, True) for root in reversed(_studio_requirement_roots(run_mod))] + pending = [] seen = set() - while pending: - requirement, is_root = pending.pop() - key = _canonical_distribution_name(requirement.name) - if key in seen: - continue - seen.add(key) + + def visit(requirement, is_root): + """The name to report, or None; queues the requirement's dependencies.""" try: installed = distribution(requirement.name) except PackageNotFoundError: @@ -345,7 +342,27 @@ def _missing_studio_requirement(run_mod): # No extra is requested, so extras-only dependencies do not apply. if parsed.marker and not parsed.marker.evaluate({"extra": ""}): continue - pending.append((parsed, False)) + pending.append(parsed) + return None + + # Every root is checked before the walk starts. Reached as a dependency + # first, a root would mark itself seen and never meet its own pin: datasets + # asks for huggingface-hub>=0.25,<2 and studio.txt pins it to ==0.36.2. + roots = _studio_requirement_roots(run_mod) + seen.update(_canonical_distribution_name(root.name) for root in roots) + for requirement in roots: + missing = visit(requirement, True) + if missing: + return missing + while pending: + requirement = pending.pop() + key = _canonical_distribution_name(requirement.name) + if key in seen: + continue + seen.add(key) + missing = visit(requirement, False) + if missing: + return missing return None diff --git a/unsloth_cli/tests/test_studio_runtime_readiness.py b/unsloth_cli/tests/test_studio_runtime_readiness.py index bc843a5bcb..e215a669f3 100644 --- a/unsloth_cli/tests/test_studio_runtime_readiness.py +++ b/unsloth_cli/tests/test_studio_runtime_readiness.py @@ -230,3 +230,31 @@ def test_desktop_runtime_check_reports_a_rejected_setting_instead_of_exiting(mon assert payload["runtime_ready"] is False assert payload["reason"] == "backend_startup_failed" assert "UNSLOTH_CPU_THREADS" in payload["error"] + + +def test_a_root_pin_is_checked_even_when_a_dependency_names_it_first( + monkeypatch, capsys, tmp_path +): + """datasets asks for huggingface-hub>=0.25,<2 and studio.txt pins ==0.36.2. + Reached as a dependency first, the pin would never get to decide.""" + studio = importlib.import_module("unsloth_cli.commands.studio") + backend = tmp_path / "backend" + requirements = backend / "requirements" + requirements.mkdir(parents = True) + (requirements / "studio.txt").write_text( + "datasets==4.3.0\nhuggingface-hub==0.36.2\n", encoding = "utf-8", + ) + run_mod = SimpleNamespace(__file__ = str(backend / "run.py")) + monkeypatch.setattr(studio, "_load_run_module", lambda: run_mod) + _fake_distributions( + monkeypatch, + { + "datasets": ("4.3.0", ["huggingface-hub>=0.25,<2"]), + "huggingface-hub": ("1.25.1", None), + }, + ) + + with pytest.raises(typer.Exit): + studio.desktop_runtime_check(_json_output = True) + + assert json.loads(capsys.readouterr().out)["module"] == "huggingface-hub"