fix(desktop): keep a verified CUDA torch, and repair a missing torchvision

desktop/src-tauri/src/main.rs

@@ -2664,7 +2664,8 @@ fn installed_version(site: &Path, package: &str) -> Option<String> {
 }
 
 /// 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<String> {
 /// 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<String> {
+fn torchvision_repair_spec(torch: &str, torchvision: Option<&str>) -> Option<String> {
     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]

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<BTreeSet<String>> {
     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<BTreeSet<String>> {
             }
         }
     }
-    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<usize> {
             "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<String> = fs::read_dir(&backup)?
+        .map(|entry| Ok(entry?.file_name().to_string_lossy().into_owned()))
+        .collect::<io::Result<_>>()?;
     // 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<String> = ["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]