diff --git a/studio/src-tauri/src/install.rs b/studio/src-tauri/src/install.rs index 74e560a921..bdefca22a1 100644 --- a/studio/src-tauri/src/install.rs +++ b/studio/src-tauri/src/install.rs @@ -54,12 +54,29 @@ pub(crate) fn managed_install_in_progress() -> bool { .unwrap_or(false) } +/// Whether this process created the marker rather than finding one an earlier +/// interrupted install left. Clearing someone else's drops the only signal that +/// its venv is half-written. Process-wide because the marker is one file. +static MARKER_CREATED_HERE: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + fn create_install_in_progress_marker() -> Result<(), String> { let path = install_in_progress_marker_path()?; - create_install_in_progress_marker_at(&path) + let created = create_install_in_progress_marker_at(&path)?; + MARKER_CREATED_HERE.store(created, std::sync::atomic::Ordering::Relaxed); + Ok(()) } -fn create_install_in_progress_marker_at(path: &Path) -> Result<(), String> { +/// Clear only a marker this attempt created. Used where the venv was never +/// touched: a failed spawn, and the elevation exits. +fn clear_own_install_marker() { + if MARKER_CREATED_HERE.load(std::sync::atomic::Ordering::Relaxed) { + clear_install_marker_best_effort(); + } +} + +/// Ok(true) when this call created the file, Ok(false) when one was there. +fn create_install_in_progress_marker_at(path: &Path) -> Result { let parent = path .parent() .ok_or_else(|| format!("Invalid install marker path: {}", path.display()))?; @@ -71,8 +88,8 @@ fn create_install_in_progress_marker_at(path: &Path) -> Result<(), String> { .create_new(true) .open(path) { - Ok(_) => Ok(()), - Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => Ok(()), + Ok(_) => Ok(true), + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => Ok(false), Err(error) => Err(format!("Failed to create {}: {}", path.display(), error)), } } @@ -90,6 +107,7 @@ fn clear_install_marker_best_effort() { if let Err(msg) = clear_install_in_progress_marker() { warn!("[install] {}", msg); } + MARKER_CREATED_HERE.store(false, std::sync::atomic::Ordering::Relaxed); } fn clear_install_in_progress_marker_at(path: &Path) -> Result<(), String> { @@ -516,10 +534,11 @@ 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. Unless another installer - // owns the marker, where clearing it drops its recovery signal. + // Nothing ran, so nothing is half-installed. Unless the marker is + // not ours: another installer owns it, or an earlier interrupted + // one left it, and either way its signal is not ours to drop. if msg != INSTALL_ALREADY_RUNNING { - clear_install_marker_best_effort(); + clear_own_install_marker(); } diagnostics::finish_attempt( &diagnostics, @@ -648,7 +667,7 @@ pub fn record_pending_elevation_canceled( return false; }; // The resumed run the elevation exit left the marker for is not happening. - clear_install_marker_best_effort(); + clear_own_install_marker(); diagnostics::finish_attempt( diagnostics, &attempt, @@ -922,7 +941,7 @@ fn finish_elevation_failure( message: String, ) { // Terminal, like a cancelled prompt: the resumed run is not happening. - clear_install_marker_best_effort(); + clear_own_install_marker(); if let Some(attempt) = attempt { diagnostics::finish_attempt( diagnostics, @@ -971,10 +990,11 @@ mod tests { )); let marker = directory.join(INSTALL_IN_PROGRESS_MARKER); - create_install_in_progress_marker_at(&marker).unwrap(); + assert!(create_install_in_progress_marker_at(&marker).unwrap()); assert!(marker.is_file()); - create_install_in_progress_marker_at(&marker).unwrap(); + // A second attempt finds the first one's marker and does not own it. + assert!(!create_install_in_progress_marker_at(&marker).unwrap()); assert!(marker.is_file()); clear_install_in_progress_marker_at(&marker).unwrap(); diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index a933235d6a..f06da08d90 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import csv import importlib.util import hashlib import hmac @@ -301,18 +302,25 @@ def _recorded_files_missing(installed) -> bool: """Whether RECORD lists files the venv no longer has. An interrupted replace can recreate a package directory and stop part-way - through filling it, so the directory existing proves nothing. 6656 files - across studio.txt's closure cost 158ms, against a probe that imports the - backend, and __pycache__ is skipped because deleting it is not damage. + through filling it, so the directory existing proves nothing. RECORD is read + directly because Distribution.files drops entries that no longer exist, + which is exactly the set this looks for. 6656 files across studio.txt's + closure cost 158ms, and __pycache__ is skipped: deleting it is not damage. """ - for recorded in installed.files or (): - parts = PurePosixPath(str(recorded)).parts + # egg-info has no RECORD, and its file lists say nothing about completeness. + record = installed.read_text("RECORD") + if record is None: + return False + for row in csv.reader(record.splitlines()): + if not row or not row[0]: + continue + parts = PurePosixPath(row[0]).parts # .dist-info is the metadata itself; .data and ../ land outside the tree. if parts[0] == ".." or parts[0].endswith((".dist-info", ".data")): continue if "__pycache__" in parts or parts[-1].endswith(".pyc"): continue - if not installed.locate_file(recorded).exists(): + if not installed.locate_file(row[0]).exists(): return True return False diff --git a/unsloth_cli/tests/test_studio_runtime_readiness.py b/unsloth_cli/tests/test_studio_runtime_readiness.py index 24a687a4c2..15fff78b11 100644 --- a/unsloth_cli/tests/test_studio_runtime_readiness.py +++ b/unsloth_cli/tests/test_studio_runtime_readiness.py @@ -93,7 +93,9 @@ def test_desktop_runtime_check_accepts_a_prerelease_over_a_floor(monkeypatch, ca monkeypatch.setattr( importlib.import_module("importlib.metadata"), "distribution", - lambda _name: SimpleNamespace(version = "2.0.0b1", files = [], requires = None), + lambda _name: SimpleNamespace( + version = "2.0.0b1", files = [], requires = None, read_text = lambda _n: None, + ), ) studio.desktop_runtime_check(_json_output = True) @@ -112,7 +114,9 @@ def test_desktop_runtime_check_rejects_version_mismatch(monkeypatch, capsys, tmp monkeypatch.setattr( importlib.import_module("importlib.metadata"), "distribution", - lambda _name: SimpleNamespace(version = "1.0", files = [], requires = None), + lambda _name: SimpleNamespace( + version = "1.0", files = [], requires = None, read_text = lambda _n: None, + ), ) with pytest.raises(typer.Exit): @@ -137,7 +141,9 @@ def test_desktop_runtime_check_rejects_metadata_without_an_unpacked_package( monkeypatch.setattr( importlib.import_module("importlib.metadata"), "distribution", - lambda _name: SimpleNamespace(version = "0.140.5", files = None, requires = None), + lambda _name: SimpleNamespace( + version = "0.140.5", files = None, requires = None, read_text = lambda _n: None, + ), ) with pytest.raises(typer.Exit): @@ -156,7 +162,9 @@ def _fake_distributions(monkeypatch, installed): version, requires = installed[name] except KeyError: raise metadata.PackageNotFoundError(name) from None - return SimpleNamespace(version = version, files = [], requires = requires) + return SimpleNamespace( + version = version, files = [], requires = requires, read_text = lambda _n: None, + ) monkeypatch.setattr(metadata, "distribution", _distribution) @@ -256,8 +264,12 @@ def test_a_root_pin_is_checked_even_when_a_dependency_names_it_first(monkeypatch def test_a_record_without_its_files_is_reported_missing(monkeypatch, capsys, tmp_path): """An interrupted replace can recreate the package directory and stop - part-way through filling it, so the directory existing proves nothing.""" + part-way through filling it, so the directory existing proves nothing. + + Built as a real dist-info rather than a stub: Distribution.files drops + entries that no longer exist, which is the whole set this looks for.""" studio = importlib.import_module("unsloth_cli.commands.studio") + metadata = importlib.import_module("importlib.metadata") backend = tmp_path / "backend" requirements = backend / "requirements" requirements.mkdir(parents = True) @@ -266,23 +278,20 @@ def test_a_record_without_its_files_is_reported_missing(monkeypatch, capsys, tmp 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/processors.py", - "structlog/__pycache__/__init__.cpython-311.pyc", - "structlog-25.1.0.dist-info/RECORD", - ], - locate_file = lambda name: site_packages / name, + dist_info = site_packages / "structlog-25.1.0.dist-info" + dist_info.mkdir(parents = True) + (dist_info / "METADATA").write_text( + "Metadata-Version: 2.1\nName: structlog\nVersion: 25.1.0\n", encoding = "utf-8", ) - monkeypatch.setattr( - importlib.import_module("importlib.metadata"), - "distribution", - lambda _name: installed, + (dist_info / "RECORD").write_text( + "structlog/__init__.py,,\n" + "structlog/processors.py,,\n" + "structlog/__pycache__/__init__.cpython-313.pyc,,\n" + "structlog-25.1.0.dist-info/RECORD,,\n", + encoding = "utf-8", ) + installed = next(iter(metadata.distributions(path = [str(site_packages)]))) + monkeypatch.setattr(metadata, "distribution", lambda _name: installed) with pytest.raises(typer.Exit): studio.desktop_runtime_check(_json_output = True) @@ -295,6 +304,7 @@ def test_a_record_without_its_files_is_reported_missing(monkeypatch, capsys, tmp studio.desktop_runtime_check(_json_output = True) assert json.loads(capsys.readouterr().out)["module"] == "structlog" + # Complete, and the never-written .pyc must not count as damage. (site_packages / "structlog" / "processors.py").touch() studio.desktop_runtime_check(_json_output = True) assert json.loads(capsys.readouterr().out) == {"runtime_ready": True}