From 63349861800b9f5ba5f40fcbaa361a7df68516df Mon Sep 17 00:00:00 2001 From: Thales <> Date: Tue, 3 Nov 2026 05:43:00 +0000 Subject: [PATCH] fix(desktop): keep a verified CUDA torch, and repair a missing torchvision A torch swap commit deleted the backup tree with the manifest still in it, so a delete that stopped part way (a file held open) left a backup that the next launch restored over the verified CUDA install. The manifest is now removed first, and the next snapshot clears a manifest-less leftover before moving anything into it. owned_entries and restore kept RECORD and manifest names that exist() accepts but that are not entries: on Windows "C:" joins to a drive's current directory and "..." or ". " are trimmed to site-packages itself, which discard_installed or restore would then delete. Names are now matched against read_dir instead. The launch check reported a missing torchvision beside +cpu torch as healthy, so a pip run that stopped between uninstall and install left the karaoke split broken for good. An absent version.py now gets the CPU torchvision install; an unreadable one is still left alone. --- desktop/src-tauri/src/main.rs | 79 +++++++++++++++------ desktop/src-tauri/src/torchswap.rs | 109 +++++++++++++++++++++++++++-- 2 files changed, 160 insertions(+), 28 deletions(-) diff --git a/desktop/src-tauri/src/main.rs b/desktop/src-tauri/src/main.rs index e05d3ff9..76016b1b 100644 --- a/desktop/src-tauri/src/main.rs +++ b/desktop/src-tauri/src/main.rs @@ -2664,7 +2664,8 @@ fn installed_version(site: &Path, package: &str) -> Option { } /// The pip spec that repairs a CUDA torchvision sitting beside CPU torch, or -/// None when there is nothing to repair. +/// a missing torchvision beside it (a pip run that stopped between uninstall +/// and install). `None` when there is nothing to repair. /// /// That pairing is what the 0.20.0 Linux packages shipped: CPU torch, and the /// cu124 torchvision PyPI serves by default, whose compiled ops need CUDA @@ -2678,12 +2679,17 @@ fn installed_version(site: &Path, package: &str) -> Option { /// the CPU package. A torch line with no known torchvision is left alone /// rather than guessed at. #[cfg(any(not(target_os = "macos"), test))] -fn torchvision_repair_spec(torch: &str, torchvision: &str) -> Option { +fn torchvision_repair_spec(torch: &str, torchvision: Option<&str>) -> Option { let (torch_public, torch_tag) = torch.split_once('+')?; - let (_, vision_tag) = torchvision.split_once('+')?; - if torch_tag != "cpu" || !vision_tag.starts_with("cu") { + if torch_tag != "cpu" { return None; } + if let Some(torchvision) = torchvision { + let (_, vision_tag) = torchvision.split_once('+')?; + if !vision_tag.starts_with("cu") { + return None; + } + } torchvision_for_torch(torch_public).map(|vision| format!("torchvision=={vision}+cpu")) } @@ -3376,20 +3382,27 @@ fn repair_torchvision_if_needed(state: &BackendState, app: &tauri::AppHandle) -> if torchswap::has_backup(&site) { return "ok"; } - let (Some(torch), Some(vision)) = ( - installed_version(&site, "torch"), - installed_version(&site, "torchvision"), - ) else { + let Some(torch) = installed_version(&site, "torch") else { return "ok"; }; - let Some(spec) = torchvision_repair_spec(&torch, &vision) else { + let vision = installed_version(&site, "torchvision"); + // Only an absent version.py means torchvision is gone. One that is there + // but unreadable is not worth a download on every launch. + if vision.is_none() && site.join("torchvision").join("version.py").exists() { + return "ok"; + } + let Some(spec) = torchvision_repair_spec(&torch, vision.as_deref()) else { return "ok"; }; log_setup_step("torchvision-repair"); + let found = match &vision { + Some(vision) => format!("torchvision {vision} is a CUDA build"), + None => "torchvision is missing".to_string(), + }; append_to_setup_log( &data_dir, - &format!("torchvision {vision} is a CUDA build beside torch {torch}; installing {spec}"), + &format!("{found} beside torch {torch}; installing {spec}"), ); let label = "torchvision repair"; // Puts a line on the setup screen before pip has said anything. @@ -10011,19 +10024,33 @@ configuration: --cc=clang\n"; use super::torchvision_repair_spec as spec; // The 0.20.0 Linux packages, CPU and NVIDIA alike (#778). assert_eq!( - spec("2.6.0+cpu", "0.21.0+cu124").as_deref(), + spec("2.6.0+cpu", Some("0.21.0+cu124")).as_deref(), Some("torchvision==0.21.0+cpu") ); // Healthy: a fixed package, or the reporter's hand repair. - assert_eq!(spec("2.6.0+cpu", "0.21.0+cpu"), None); + assert_eq!(spec("2.6.0+cpu", Some("0.21.0+cpu")), None); // An NVIDIA install after the CUDA install: a working pair. - assert_eq!(spec("2.6.0+cu124", "0.21.0+cu124"), None); - assert_eq!(spec("2.8.0+cu128", "0.23.0+cu128"), None); + assert_eq!(spec("2.6.0+cu124", Some("0.21.0+cu124")), None); + assert_eq!(spec("2.8.0+cu128", Some("0.23.0+cu128")), None); // No local tag at all, as some PyPI builds report: nothing to judge. - assert_eq!(spec("2.6.0", "0.21.0"), None); - assert_eq!(spec("2.6.0+cpu", "0.21.0"), None); + assert_eq!(spec("2.6.0", Some("0.21.0")), None); + assert_eq!(spec("2.6.0+cpu", Some("0.21.0")), None); // A torch line with no known torchvision is not guessed at. - assert_eq!(spec("2.7.1+cpu", "0.22.1+cu126"), None); + assert_eq!(spec("2.7.1+cpu", Some("0.22.1+cu126")), None); + } + + /// A pip run stopped after it uninstalled torchvision and before it put + /// the new one in: every package ships torchvision, so it is put back. + #[test] + fn a_missing_torchvision_beside_cpu_torch_is_repaired() { + use super::torchvision_repair_spec as spec; + assert_eq!( + spec("2.6.0+cpu", None).as_deref(), + Some("torchvision==0.21.0+cpu") + ); + // A CUDA torch is the CUDA install's to manage, not this repair's. + assert_eq!(spec("2.6.0+cu124", None), None); + assert_eq!(spec("2.7.1+cpu", None), None); } #[test] @@ -10041,15 +10068,25 @@ configuration: --cc=clang\n"; write("torch", "2.6.0+cpu"); write("torchvision", "0.21.0+cu124"); let torch = super::installed_version(site.path(), "torch").unwrap(); - let vision = super::installed_version(site.path(), "torchvision").unwrap(); + let vision = super::installed_version(site.path(), "torchvision"); assert_eq!( - super::torchvision_repair_spec(&torch, &vision).as_deref(), + super::torchvision_repair_spec(&torch, vision.as_deref()).as_deref(), Some("torchvision==0.21.0+cpu") ); // After the repair, the next launch has nothing to do. write("torchvision", "0.21.0+cpu"); - let vision = super::installed_version(site.path(), "torchvision").unwrap(); - assert_eq!(super::torchvision_repair_spec(&torch, &vision), None); + let vision = super::installed_version(site.path(), "torchvision"); + assert_eq!( + super::torchvision_repair_spec(&torch, vision.as_deref()), + None + ); + // Uninstalled and never reinstalled. + fs::remove_dir_all(site.path().join("torchvision")).unwrap(); + let vision = super::installed_version(site.path(), "torchvision"); + assert_eq!( + super::torchvision_repair_spec(&torch, vision.as_deref()).as_deref(), + Some("torchvision==0.21.0+cpu") + ); } #[test] diff --git a/desktop/src-tauri/src/torchswap.rs b/desktop/src-tauri/src/torchswap.rs index 522ca625..41c97ae7 100644 --- a/desktop/src-tauri/src/torchswap.rs +++ b/desktop/src-tauri/src/torchswap.rs @@ -55,9 +55,11 @@ fn is_torch_dist_info(name: &str) -> bool { /// their RECORD. Only names that exist are returned. pub fn owned_entries(site: &Path) -> io::Result> { let mut owned = BTreeSet::new(); + let mut listed = BTreeSet::new(); for entry in fs::read_dir(site)? { let entry = entry?; let name = entry.file_name().to_string_lossy().into_owned(); + listed.insert(name.clone()); if !is_torch_dist_info(&name) { continue; } @@ -83,7 +85,10 @@ pub fn owned_entries(site: &Path) -> io::Result> { } } } - owned.retain(|name| name != BACKUP_DIR && site.join(name).exists()); + // Matched against the listing, not checked with exists(): a damaged + // RECORD can name "C:", which Windows joins to that drive's current + // directory, or "..." and ". ", which it trims to site itself. + owned.retain(|name| name != BACKUP_DIR && listed.contains(name)); Ok(owned) } @@ -114,6 +119,12 @@ pub fn snapshot(site: &Path) -> io::Result { "no torch packages found to back up", )); } + // A backup with no manifest is a commit that could not finish deleting + // it. Nothing in it is to be restored, and a stale entry left in it would + // block the move below. + if backup.exists() { + fs::remove_dir_all(&backup)?; + } fs::create_dir_all(&backup)?; // The manifest goes first and reaches the disk before anything moves, so // a crash part way through leaves a record of what to put back. @@ -151,16 +162,16 @@ pub fn restore(site: &Path) -> io::Result<()> { .filter(|n| !n.is_empty() && *n != "." && *n != ".." && !n.contains(['/', '\\'])) .map(str::to_string) .collect(); + // From the listing for the same reason as in owned_entries. + let saved: BTreeSet = fs::read_dir(&backup)? + .map(|entry| Ok(entry?.file_name().to_string_lossy().into_owned())) + .collect::>()?; // Moved names come back, replacing whatever is there now. A name the // manifest lists but the backup lacks was never moved, so it is still the // original and stays. - for name in &names { - let saved = backup.join(name); - if !saved.exists() { - continue; - } + for name in names.intersection(&saved) { remove_entry(&site.join(name))?; - fs::rename(&saved, site.join(name))?; + fs::rename(backup.join(name), site.join(name))?; } // Then anything torch-owned that the original did not have: the CUDA // build's own dist-info, its extra top-level packages. @@ -175,6 +186,10 @@ pub fn restore(site: &Path) -> io::Result<()> { /// The swap worked: the backup is no longer needed. pub fn commit(site: &Path) -> io::Result<()> { let backup = backup_root(site); + // The manifest goes first, so a delete that stops part way (a file held + // open) leaves no half-emptied backup for the next launch to restore over + // the working install. + remove_entry(&backup.join(MANIFEST))?; if backup.exists() { fs::remove_dir_all(backup)?; } @@ -380,6 +395,86 @@ mod tests { assert!(site.join("torch").join("lib").join("c10_cuda.dll").exists()); } + /// A file in the backup held open, as an antivirus scan can, stops the + /// commit part way. What is left must not count as a backup, or the next + /// launch restores half a CPU torch over the verified CUDA one. The open + /// file sits in functorch, which NTFS lists before manifest.txt. + #[cfg(windows)] + #[test] + fn a_commit_that_stops_part_way_leaves_no_backup_to_restore() { + use std::os::windows::fs::OpenOptionsExt; + let tmp = tempfile::tempdir().unwrap(); + let site = fake_site(tmp.path(), "2.6.0+cpu", None); + snapshot(&site).unwrap(); + install_cuda(&site); + let held = fs::OpenOptions::new() + .read(true) + .share_mode(0) + .open(backup_root(&site).join("functorch").join("__init__.py")) + .unwrap(); + + assert!(commit(&site).is_err()); + assert!(!has_backup(&site)); + restore(&site).unwrap(); + assert!(site.join("torch").join("lib").join("c10_cuda.dll").exists()); + + // The leftover is cleared by the next snapshot, not moved into. + drop(held); + discard_installed(&site).unwrap(); + let site = fake_site(tmp.path(), "2.6.0+cpu", None); + snapshot(&site).unwrap(); + restore(&site).unwrap(); + assert_eq!( + fs::read_to_string(site.join("functorch").join("__init__.py")).unwrap(), + "2.6.0+cpu" + ); + assert!(!backup_root(&site).exists()); + } + + /// A damaged RECORD naming a drive must not make the drive's current + /// directory look like a torch entry, which discard_installed would + /// delete. Only reads, so the unfixed code is safe to run against it. + #[test] + fn a_record_line_naming_no_entry_in_site_is_not_owned() { + let tmp = tempfile::tempdir().unwrap(); + let site = fake_site(tmp.path(), "2.6.0+cpu", None); + let record = site.join("torch-2.6.0+cpu.dist-info").join("RECORD"); + let mut lines = fs::read_to_string(&record).unwrap(); + // Windows trims trailing dots and spaces, so the last three are site. + lines.push_str("C:\\x,,\nC:,,\n...\\x,,\n. /x,,\n.. \\x,,\n"); + fs::write(&record, lines).unwrap(); + let owned = owned_entries(&site).unwrap(); + let expected: BTreeSet = ["functorch", "torch", "torch-2.6.0+cpu.dist-info"] + .map(String::from) + .into(); + assert_eq!(owned, expected); + } + + /// The same for a damaged manifest, which an older build could write + /// from such a RECORD. "..." joins to the backup dir and to site itself; + /// acting on it would delete site-packages. This test's site is a temp + /// dir, so the unfixed code is safe to run against it. + #[test] + fn a_manifest_line_that_names_no_entry_is_ignored() { + let tmp = tempfile::tempdir().unwrap(); + let site = fake_site(tmp.path(), "2.6.0+cpu", None); + snapshot(&site).unwrap(); + install_cuda(&site); + let manifest = backup_root(&site).join(MANIFEST); + let mut lines = fs::read_to_string(&manifest).unwrap(); + lines.push_str("...\n. \n"); + fs::write(&manifest, lines).unwrap(); + + restore(&site).unwrap(); + + assert_eq!( + fs::read_to_string(site.join("torch").join("lib").join("c10.dll")).unwrap(), + "c10 2.6.0+cpu" + ); + assert!(site.join("numpy").join("__init__.py").exists()); + assert!(!has_backup(&site)); + } + /// A crash part way through the snapshot: some entries moved, some not. /// Restoring puts back what moved and keeps what never did. #[test]