flodl-cli 0.5.0

libtorch manager and GPU diagnostic tool for Rust deep learning
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105

//! libtorch installation detection and .arch metadata parsing.

use std::fs;
use std::path::Path;

use crate::util::system::GpuInfo;

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

/// Metadata about an installed libtorch variant (from `.arch` file).
pub struct LibtorchInfo {
    /// Relative path from project root (e.g. "precompiled/cu128", "builds/sm61-sm120").
    pub path: String,
    pub torch_version: Option<String>,
    pub cuda_version: Option<String>,
    pub archs: Option<String>,
    pub source: Option<String>,
}

// ---------------------------------------------------------------------------
// Detection
// ---------------------------------------------------------------------------

/// Read the active libtorch variant from `<root>/libtorch/.active` and parse
/// its `.arch` metadata.
pub fn read_active(root: &Path) -> Option<LibtorchInfo> {
    let active_path = root.join("libtorch/.active");
    let active = fs::read_to_string(active_path).ok()?;
    let path = active.trim().to_string();
    if path.is_empty() {
        return None;
    }

    let arch_path = root.join(format!("libtorch/{}/.arch", path));
    let mut info = LibtorchInfo {
        path,
        torch_version: None,
        cuda_version: None,
        archs: None,
        source: None,
    };

    if let Ok(content) = fs::read_to_string(arch_path) {
        for line in content.lines() {
            if let Some(val) = line.strip_prefix("torch=") {
                info.torch_version = Some(val.to_string());
            } else if let Some(val) = line.strip_prefix("cuda=") {
                info.cuda_version = Some(val.to_string());
            } else if let Some(val) = line.strip_prefix("archs=") {
                info.archs = Some(val.to_string());
            } else if let Some(val) = line.strip_prefix("source=") {
                info.source = Some(val.to_string());
            }
        }
    }

    Some(info)
}

/// List all installed libtorch variants under `<root>/libtorch/`.
///
/// Scans `precompiled/` and `builds/` subdirectories.
pub fn list_variants(root: &Path) -> Vec<String> {
    let mut variants = Vec::new();
    let lt_dir = root.join("libtorch");

    for subdir in ["precompiled", "builds"] {
        let dir = lt_dir.join(subdir);
        if let Ok(entries) = fs::read_dir(&dir) {
            for entry in entries.flatten() {
                if entry.path().join("lib").is_dir() {
                    if let Some(name) = entry.file_name().to_str() {
                        variants.push(format!("{}/{}", subdir, name));
                    }
                }
            }
        }
    }

    variants.sort();
    variants
}

/// Check whether a GPU's compute capability is covered by the libtorch
/// variant's compiled architectures (from the .arch file).
pub fn arch_compatible(gpu: &GpuInfo, archs: &str) -> bool {
    let exact = format!("{}.{}", gpu.sm_major, gpu.sm_minor);
    archs.contains(&exact) || archs.contains(&format!("{}", gpu.sm_major))
}

/// Check whether a libtorch variant directory looks valid (has lib/).
pub fn is_valid_variant(root: &Path, variant: &str) -> bool {
    root.join(format!("libtorch/{}/lib", variant)).is_dir()
}

/// Set the active libtorch variant by writing `<root>/libtorch/.active`.
pub fn set_active(root: &Path, variant: &str) -> Result<(), String> {
    let lt_dir = root.join("libtorch");
    fs::create_dir_all(&lt_dir)
        .map_err(|e| format!("cannot create libtorch/: {}", e))?;
    fs::write(lt_dir.join(".active"), format!("{}\n", variant))
        .map_err(|e| format!("cannot write libtorch/.active: {}", e))
}