flodl_cli/libtorch/
detect.rs

1//! libtorch installation detection and .arch metadata parsing.
2
3use std::fs;
4use std::path::Path;
5
6use crate::util::system::GpuInfo;
7
8// ---------------------------------------------------------------------------
9// Types
10// ---------------------------------------------------------------------------
11
12/// Metadata about an installed libtorch variant (from `.arch` file).
13#[derive(Debug)]
14pub struct LibtorchInfo {
15    /// Relative path from project root (e.g. "precompiled/cu128", "builds/sm61-sm120").
16    pub path: String,
17    pub torch_version: Option<String>,
18    pub cuda_version: Option<String>,
19    pub archs: Option<String>,
20    pub source: Option<String>,
21}
22
23// ---------------------------------------------------------------------------
24// Detection
25// ---------------------------------------------------------------------------
26
27/// Read the active libtorch variant from `<root>/libtorch/.active` and
28/// parse its `.arch` metadata.
29///
30/// On heterogeneous rigs (multiple hosts sharing the same checkout via
31/// NFS / virtiofs / S3-FUSE) a single `.active` file can't represent
32/// "PT 2.10 on the Blackwell host AND PT 2.7 on the Pascal host" at
33/// the same time. The `FDL_LIBTORCH_CASE` env var selects an
34/// alternative pointer file `libtorch/.active.<case>`; the file's
35/// content is read identically to `.active`.
36///
37/// Setting `FDL_LIBTORCH_CASE=<case>` with no corresponding pointer
38/// file is a hard misconfiguration: this function logs the missing
39/// file to stderr and returns `None` (callers surface this as "libtorch
40/// not configured") rather than silently falling back to `.active`,
41/// which would otherwise mask the user's explicit selection.
42pub fn read_active(root: &Path) -> Option<LibtorchInfo> {
43    let lt_dir = root.join("libtorch");
44    let pointer = match std::env::var("FDL_LIBTORCH_CASE") {
45        Ok(case) if !case.trim().is_empty() => {
46            let case = case.trim();
47            let case_file = lt_dir.join(format!(".active.{case}"));
48            if !case_file.exists() {
49                eprintln!(
50                    "fdl: FDL_LIBTORCH_CASE={case} but `{}` does not exist. \
51                     Create it with `fdl libtorch use <variant> --as {case}` \
52                     or hand-write the variant path (e.g. \
53                     `precompiled/cu128`).",
54                    case_file.display(),
55                );
56                return None;
57            }
58            case_file
59        }
60        _ => lt_dir.join(".active"),
61    };
62    read_active_from(&pointer, &lt_dir)
63}
64
65/// Read any `.active*` pointer file and resolve its content (a
66/// relative path like `precompiled/cu128` or `builds/sm61-sm120`)
67/// against `libtorch_root` to produce a [`LibtorchInfo`].
68///
69/// Used by [`read_active`] (default `.active`), by callers that have
70/// the pointer file path directly (e.g. cluster.yml's per-host
71/// `arch:` naming a case pointer, resolving to
72/// `…/libtorch/.active.<case>`), and by tests
73/// that need to validate a pointer-file shape without setting the
74/// `FDL_LIBTORCH_CASE` env var.
75pub fn read_active_from(pointer: &Path, libtorch_root: &Path) -> Option<LibtorchInfo> {
76    let active = fs::read_to_string(pointer).ok()?;
77    let path = active.trim().to_string();
78    if path.is_empty() {
79        return None;
80    }
81    let arch_dir = libtorch_root.join(&path);
82    Some(libtorch_info_from_dir(path, &arch_dir))
83}
84
85/// Build a [`LibtorchInfo`] for a variant directory: `path` is the string
86/// recorded in the info (a relative variant path like `precompiled/cu128`
87/// or an absolute directory), `arch_dir` is the directory whose `.arch`
88/// file supplies the metadata. The four metadata fields stay `None` when
89/// the `.arch` file is absent or unreadable. One home for the parse that
90/// `read_active_from`, `run::resolve_libtorch_at`, and probe's
91/// `check_libtorch*` all used to copy inline.
92pub(crate) fn libtorch_info_from_dir(path: String, arch_dir: &Path) -> LibtorchInfo {
93    let mut info = LibtorchInfo {
94        path,
95        torch_version: None,
96        cuda_version: None,
97        archs: None,
98        source: None,
99    };
100    if let Ok(content) = fs::read_to_string(arch_dir.join(".arch")) {
101        parse_arch_into(&content, &mut info);
102    }
103    info
104}
105
106/// Fill a [`LibtorchInfo`]'s metadata fields from `.arch` file content
107/// (`torch=` / `cuda=` / `archs=` / `source=` lines; unknown lines ignored).
108fn parse_arch_into(content: &str, info: &mut LibtorchInfo) {
109    for line in content.lines() {
110        if let Some(val) = line.strip_prefix("torch=") {
111            info.torch_version = Some(val.to_string());
112        } else if let Some(val) = line.strip_prefix("cuda=") {
113            info.cuda_version = Some(val.to_string());
114        } else if let Some(val) = line.strip_prefix("archs=") {
115            info.archs = Some(val.to_string());
116        } else if let Some(val) = line.strip_prefix("source=") {
117            info.source = Some(val.to_string());
118        }
119    }
120}
121
122/// Record per-GPU arch coverage for a resolved variant and push a loud,
123/// actionable issue for every GPU the libtorch build does not cover — or a
124/// single issue when the variant carries no `.arch` metadata. Returns
125/// `(gpu_index, covered)` pairs in GPU order. One home for the coverage
126/// loop probe's three `check_libtorch*` paths used to copy inline (with
127/// drifted wording).
128pub(crate) fn arch_coverage(
129    info: &LibtorchInfo,
130    gpus: &[GpuInfo],
131    issues: &mut Vec<String>,
132) -> Vec<(u8, bool)> {
133    let mut archs_match = Vec::new();
134    if let Some(archs) = &info.archs {
135        for g in gpus {
136            let ok = arch_compatible(g, archs);
137            archs_match.push((g.index, ok));
138            if !ok {
139                issues.push(format!(
140                    "GPU {} ({}, {}) not covered by libtorch archs `{}`. \
141                     Rebuild libtorch with this arch or activate a \
142                     compatible variant.",
143                    g.index,
144                    g.short_name(),
145                    g.sm_version(),
146                    archs
147                ));
148            }
149        }
150    } else {
151        issues.push(
152            "libtorch is present but `.arch` metadata is missing — cannot \
153             verify GPU compatibility. Place an `.arch` file in the variant \
154             directory (cuda=, torch=, archs=, source=)."
155                .into(),
156        );
157    }
158    archs_match
159}
160
161/// List all installed libtorch variants under `<root>/libtorch/`.
162///
163/// Scans `precompiled/` and `builds/` subdirectories.
164pub fn list_variants(root: &Path) -> Vec<String> {
165    let mut variants = Vec::new();
166    let lt_dir = root.join("libtorch");
167
168    for subdir in ["precompiled", "builds"] {
169        let dir = lt_dir.join(subdir);
170        if let Ok(entries) = fs::read_dir(&dir) {
171            for entry in entries.flatten() {
172                if entry.path().join("lib").is_dir() {
173                    if let Some(name) = entry.file_name().to_str() {
174                        variants.push(format!("{}/{}", subdir, name));
175                    }
176                }
177            }
178        }
179    }
180
181    variants.sort();
182    variants
183}
184
185/// Check whether a GPU's compute capability is covered by the libtorch
186/// variant's compiled architectures (from the .arch file).
187pub fn arch_compatible(gpu: &GpuInfo, archs: &str) -> bool {
188    let exact = format!("{}.{}", gpu.sm_major, gpu.sm_minor);
189    archs.contains(&exact) || archs.contains(&format!("{}", gpu.sm_major))
190}
191
192/// Check whether a libtorch variant directory looks valid (has lib/).
193pub fn is_valid_variant(root: &Path, variant: &str) -> bool {
194    root.join(format!("libtorch/{}/lib", variant)).is_dir()
195}
196
197/// Set the active libtorch variant by writing `<root>/libtorch/.active`.
198pub fn set_active(root: &Path, variant: &str) -> Result<(), String> {
199    let lt_dir = root.join("libtorch");
200    fs::create_dir_all(&lt_dir)
201        .map_err(|e| format!("cannot create libtorch/: {}", e))?;
202    fs::write(lt_dir.join(".active"), format!("{}\n", variant))
203        .map_err(|e| format!("cannot write libtorch/.active: {}", e))
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209    use crate::util::test_env::env_lock;
210    use std::path::PathBuf;
211    use std::sync::atomic::{AtomicU64, Ordering};
212    use std::time::{SystemTime, UNIX_EPOCH};
213
214    // Per-process counter so concurrent test binaries don't collide
215    // on the unique-suffix algorithm.
216    static SCRATCH_SEQ: AtomicU64 = AtomicU64::new(0);
217
218    /// Hand-rolled scratch dir under the system temp dir + RAII
219    /// cleanup. flodl-cli keeps external deps minimal (the serde
220    /// ecosystem only — no utility crates like `tempfile`) so we
221    /// cannot pull in `tempfile`.
222    struct Scratch(PathBuf);
223    impl Scratch {
224        fn new() -> Self {
225            let nanos = SystemTime::now()
226                .duration_since(UNIX_EPOCH)
227                .map(|d| d.as_nanos())
228                .unwrap_or(0);
229            let seq = SCRATCH_SEQ.fetch_add(1, Ordering::Relaxed);
230            let dir = std::env::temp_dir()
231                .join(format!("fdl-libtorch-resolver-{}-{}", nanos, seq));
232            fs::create_dir_all(&dir).expect("create scratch");
233            Self(dir)
234        }
235        fn path(&self) -> &std::path::Path { &self.0 }
236    }
237    impl Drop for Scratch {
238        fn drop(&mut self) {
239            let _ = fs::remove_dir_all(&self.0);
240        }
241    }
242
243    /// Build a fake project root with two synthetic libtorch variants
244    /// (`precompiled/v1` and `builds/v2`), each with a `lib/` dir and
245    /// `.arch` metadata. The names are deliberately abstract — the
246    /// resolver doesn't care about variant naming, just the
247    /// `<kind>/<name>` shape it reads from the pointer file.
248    fn make_root() -> Scratch {
249        let s = Scratch::new();
250        let v1 = s.path().join("libtorch/precompiled/v1");
251        fs::create_dir_all(v1.join("lib")).unwrap();
252        fs::write(
253            v1.join(".arch"),
254            "torch=1.0\ncuda=1.0\narchs=0.0\nsource=precompiled\n",
255        ).unwrap();
256        let v2 = s.path().join("libtorch/builds/v2");
257        fs::create_dir_all(v2.join("lib")).unwrap();
258        fs::write(
259            v2.join(".arch"),
260            "torch=2.0\ncuda=2.0\narchs=1.0\nsource=build\n",
261        ).unwrap();
262        s
263    }
264
265    #[test]
266    fn read_active_default_pointer() {
267        let _guard = env_lock();
268        // SAFETY: serialized via env_lock().
269        unsafe { std::env::remove_var("FDL_LIBTORCH_CASE"); }
270        let root = make_root();
271        fs::write(
272            root.path().join("libtorch/.active"),
273            "precompiled/v1\n",
274        ).unwrap();
275        let info = read_active(root.path()).expect("read_active");
276        assert_eq!(info.path, "precompiled/v1");
277        assert_eq!(info.torch_version.as_deref(), Some("1.0"));
278    }
279
280    #[test]
281    fn fdl_libtorch_case_selects_alternate_pointer() {
282        let _guard = env_lock();
283        let root = make_root();
284        fs::write(
285            root.path().join("libtorch/.active"),
286            "builds/v2\n",
287        ).unwrap();
288        fs::write(
289            root.path().join("libtorch/.active.alt"),
290            "precompiled/v1\n",
291        ).unwrap();
292        // SAFETY: serialized via env_lock().
293        unsafe { std::env::set_var("FDL_LIBTORCH_CASE", "alt"); }
294        let info = read_active(root.path()).expect("read_active");
295        // SAFETY: serialized via env_lock().
296        unsafe { std::env::remove_var("FDL_LIBTORCH_CASE"); }
297        assert_eq!(info.path, "precompiled/v1");
298        assert_eq!(info.torch_version.as_deref(), Some("1.0"));
299    }
300
301    #[test]
302    fn fdl_libtorch_case_missing_file_returns_none_loudly() {
303        let _guard = env_lock();
304        let root = make_root();
305        fs::write(
306            root.path().join("libtorch/.active"),
307            "builds/v2\n",
308        ).unwrap();
309        // No `.active.nonexistent` file.
310        // SAFETY: serialized via env_lock().
311        unsafe { std::env::set_var("FDL_LIBTORCH_CASE", "nonexistent"); }
312        let info = read_active(root.path());
313        // SAFETY: serialized via env_lock().
314        unsafe { std::env::remove_var("FDL_LIBTORCH_CASE"); }
315        assert!(info.is_none(),
316            "explicit case with missing file must not silently fall back to .active");
317    }
318
319    #[test]
320    fn read_active_from_resolves_pointer_directly() {
321        let _guard = env_lock();
322        let root = make_root();
323        let pointer = root.path().join("libtorch/.active.alt");
324        fs::write(&pointer, "builds/v2\n").unwrap();
325        let info = read_active_from(
326            &pointer, &root.path().join("libtorch"),
327        ).expect("read_active_from");
328        assert_eq!(info.path, "builds/v2");
329        assert_eq!(info.archs.as_deref(), Some("1.0"));
330    }
331}
flodl_cli/libtorch/detect.rs

flodl_cli/libtorch/
detect.rs
