diff --git a/studio/src-tauri/src/desktop_backend_owner.rs b/studio/src-tauri/src/desktop_backend_owner.rs index 5e0c0425e6..a0400eea5b 100644 --- a/studio/src-tauri/src/desktop_backend_owner.rs +++ b/studio/src-tauri/src/desktop_backend_owner.rs @@ -268,6 +268,9 @@ impl BackendOwnerState { /// Only while the file is still ours. Cleanup after an exit can take seconds, /// and a start in that window already wrote the next metadata to this path. pub(crate) fn remove(self) { + // Held across the check and the unlink: a start that lands between the + // two would otherwise have its file deleted by the backend it replaced. + let _guard = lock_metadata(); if let Ok(Some(on_disk)) = read_metadata(&self.path) { if on_disk.token_sha256 != self.metadata.token_sha256 { warn!( @@ -316,7 +319,16 @@ impl BackendOwnerState { } } +/// Serializes activation against cleanup. Both run in this process, and the +/// check in remove() only holds if no write can land inside it. +static METADATA_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + +fn lock_metadata() -> std::sync::MutexGuard<'static, ()> { + METADATA_LOCK.lock().unwrap_or_else(|e| e.into_inner()) +} + fn write_metadata(path: &Path, metadata: &DesktopBackendMetadata) -> Result<(), String> { + let _guard = lock_metadata(); let parent = path .parent() .ok_or_else(|| "desktop owner metadata path has no parent".to_string())?; diff --git a/studio/src-tauri/src/preflight.rs b/studio/src-tauri/src/preflight.rs index b9046eaa26..fe9e4e6fcd 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -712,6 +712,8 @@ sleep 120 async fn managed_runtime_probe_survives_more_stdout_than_the_pipe_holds() { // The probe imports the whole backend, so output can exceed the 64 KiB // pipe buffer, and waiting before draining wedges the child mid-write. + // This CLI writes ~1 MiB before the payload, so the drain both keeps up + // and keeps only the tail, which still has to contain the payload. let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await; let _cache_home = ManagedCapabilityCacheHome::new("runtime-noisy"); remove_managed_capability_cache(); @@ -720,7 +722,7 @@ sleep 120 "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" }' + awk 'BEGIN { while (i++ < 40000) print "import chatter on stdout" }' printf '{"runtime_ready":true}\n' exit 0 fi diff --git a/studio/src-tauri/src/preflight/managed.rs b/studio/src-tauri/src/preflight/managed.rs index ba38523688..a3a67189b5 100644 --- a/studio/src-tauri/src/preflight/managed.rs +++ b/studio/src-tauri/src/preflight/managed.rs @@ -303,6 +303,30 @@ async fn run_cli_probe(bin: &Path, args: &[&str]) -> bool { } } +/// Room for the payload and what preceded it. Bounded because the binary being +/// probed is a possibly damaged install, and keeping every byte lets an import +/// stuck printing grow this buffer for the whole timeout. +const PROBE_TAIL_BYTES: usize = 64 * 1024; + +/// Drained while the wait runs: the backend import can print more than the pipe +/// holds, and waiting first deadlocks, timing out a healthy install into repair. +fn drain_probe_tail(mut stdout: tokio::process::ChildStdout) -> tokio::task::JoinHandle> { + tokio::spawn(async move { + let mut tail = Vec::new(); + let mut chunk = [0u8; 8192]; + while let Ok(read) = stdout.read(&mut chunk).await { + if read == 0 { + break; + } + tail.extend_from_slice(&chunk[..read]); + if tail.len() > PROBE_TAIL_BYTES { + tail.drain(..tail.len() - PROBE_TAIL_BYTES); + } + } + tail + }) +} + /// The probe could not answer: it would not run, or 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"; @@ -353,17 +377,11 @@ async fn probe_cli_runtime(bin: &Path) -> Result<(), String> { ); return Err(RUNTIME_PROBE_FAILED.to_string()); }; - let Some(mut stdout) = child.stdout.take() else { + let Some(stdout) = child.stdout.take() else { return Err(RUNTIME_PROBE_FAILED.to_string()); }; - // Drain while waiting: the backend import can print more than the pipe holds, - // and waiting first deadlocks, 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 reader = drain_probe_tail(stdout); let status = match tokio::time::timeout(Duration::from_secs(10), child.wait()).await { Ok(Ok(status)) => status, @@ -445,16 +463,11 @@ async fn probe_cli_capability(bin: &Path) -> Option { ); return None; }; - let Some(mut stdout) = child.stdout.take() else { + let Some(stdout) = child.stdout.take() else { 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 - }); + let reader = drain_probe_tail(stdout); match tokio::time::timeout(Duration::from_secs(10), child.wait()).await { Ok(Ok(status)) if status.success() => {} diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 43c6634865..285689b894 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -18,7 +18,7 @@ import types import urllib.error import urllib.request from datetime import datetime, timezone -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import List, Optional import typer @@ -297,6 +297,23 @@ def _studio_requirement_roots(run_mod): return roots +def _recorded_packages_missing(installed) -> bool: + """Whether RECORD lists top-level packages the venv no longer has. + + An interrupted replace removes the package and can leave its RECORD, so + reading one is not proof the modules are there. Stats the top-level names + only: a handful per distribution, and the failure is never partial. + """ + packages = set() + for recorded in installed.files or (): + top = PurePosixPath(str(recorded)).parts[0] + # .dist-info is the metadata itself; .data and ../ land outside the tree. + if top == ".." or top.endswith((".dist-info", ".data")): + continue + packages.add(top) + return any(not installed.locate_file(name).exists() for name in packages) + + def _missing_studio_requirement(run_mod): """First requirement studio.txt needs that this venv cannot supply. @@ -318,7 +335,7 @@ def _missing_studio_requirement(run_mod): return requirement.name # Metadata outlives the package: wheels store METADATA first, so a killed # unpack leaves a version behind. RECORD is last, so its absence marks it. - if installed.files is None: + if installed.files is None or _recorded_packages_missing(installed): return requirement.name # Only studio.txt pins are enforced: install_python_stack uses --no-deps, # so transitive bounds read unsatisfied in venvs that work. A prerelease diff --git a/unsloth_cli/tests/test_studio_runtime_readiness.py b/unsloth_cli/tests/test_studio_runtime_readiness.py index 454638258b..ed14a1e8f7 100644 --- a/unsloth_cli/tests/test_studio_runtime_readiness.py +++ b/unsloth_cli/tests/test_studio_runtime_readiness.py @@ -252,3 +252,35 @@ def test_a_root_pin_is_checked_even_when_a_dependency_names_it_first(monkeypatch studio.desktop_runtime_check(_json_output = True) assert json.loads(capsys.readouterr().out)["module"] == "huggingface-hub" + + +def test_a_record_without_its_package_is_reported_missing(monkeypatch, capsys, tmp_path): + """An interrupted replace can delete the package and leave its RECORD, so a + readable file list is not proof the modules are there.""" + studio = importlib.import_module("unsloth_cli.commands.studio") + backend = tmp_path / "backend" + requirements = backend / "requirements" + requirements.mkdir(parents = True) + (requirements / "studio.txt").write_text("structlog\n", encoding = "utf-8") + run_mod = SimpleNamespace(__file__ = str(backend / "run.py")) + monkeypatch.setattr(studio, "_load_run_module", lambda: run_mod) + + site_packages = tmp_path / "site-packages" + (site_packages / "structlog-25.1.0.dist-info").mkdir(parents = True) + installed = SimpleNamespace( + version = "25.1.0", + requires = None, + files = ["structlog/__init__.py", "structlog-25.1.0.dist-info/RECORD"], + locate_file = lambda name: site_packages / name, + ) + monkeypatch.setattr( + importlib.import_module("importlib.metadata"), "distribution", lambda _name: installed, + ) + + with pytest.raises(typer.Exit): + studio.desktop_runtime_check(_json_output = True) + assert json.loads(capsys.readouterr().out)["module"] == "structlog" + + (site_packages / "structlog").mkdir() + studio.desktop_runtime_check(_json_output = True) + assert json.loads(capsys.readouterr().out) == {"runtime_ready": True}