rusty_stack/installers/components/
wandb.rs

1//! W&B (Weights & Biases) installer — ports `scripts/install_wandb.sh`.
2//!
3//! Constructs correct pip install command for the wandb package.
4//!
5//! # Validation Assertion
6//!
7//! - **VAL-INSTALL-021**: W&B correct pip command
8
9// ===========================================================================
10// Types
11// ===========================================================================
12
13/// Installation method.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum InstallMethod {
16    /// Install globally.
17    Global,
18    /// Install in a virtual environment.
19    Venv,
20    /// Try global, fallback to venv.
21    Auto,
22}
23
24impl std::fmt::Display for InstallMethod {
25    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        match self {
27            InstallMethod::Global => write!(f, "global"),
28            InstallMethod::Venv => write!(f, "venv"),
29            InstallMethod::Auto => write!(f, "auto"),
30        }
31    }
32}
33
34/// Configuration for the W&B installer.
35#[derive(Debug, Clone)]
36pub struct WandbConfig {
37    /// Python binary to use.
38    pub python_bin: String,
39    /// Installation method.
40    pub method: InstallMethod,
41    /// Whether to run in dry-run mode.
42    pub dry_run: bool,
43    /// Whether strict ROCm mode is enabled.
44    pub strict_rocm: bool,
45}
46
47impl Default for WandbConfig {
48    fn default() -> Self {
49        Self {
50            python_bin: "python3".to_string(),
51            method: InstallMethod::Auto,
52            dry_run: false,
53            strict_rocm: true,
54        }
55    }
56}
57
58/// A constructed shell command.
59#[derive(Debug, Clone, PartialEq, Eq)]
60pub struct ShellCommand {
61    /// The program to run.
62    pub program: String,
63    /// Arguments to pass.
64    pub args: Vec<String>,
65    /// Environment variables to set.
66    pub env: Vec<(String, String)>,
67}
68
69impl ShellCommand {
70    /// Format as a shell command string.
71    pub fn to_command_string(&self) -> String {
72        let env_prefix = self
73            .env
74            .iter()
75            .map(|(k, v)| format!("{k}={v}"))
76            .collect::<Vec<_>>()
77            .join(" ");
78        let cmd = if self.args.is_empty() {
79            self.program.clone()
80        } else {
81            format!("{} {}", self.program, self.args.join(" "))
82        };
83        if env_prefix.is_empty() {
84            cmd
85        } else {
86            format!("{env_prefix} {cmd}")
87        }
88    }
89}
90
91/// The W&B installer.
92pub struct WandbInstaller {
93    config: WandbConfig,
94}
95
96impl WandbInstaller {
97    /// Create a new W&B installer with the given config.
98    pub fn new(config: WandbConfig) -> Self {
99        Self { config }
100    }
101
102    /// Create with default config.
103    pub fn with_defaults() -> Self {
104        Self::new(WandbConfig::default())
105    }
106
107    // -----------------------------------------------------------------------
108    // Command construction (VAL-INSTALL-021)
109    // -----------------------------------------------------------------------
110
111    /// Construct the pip install command for wandb.
112    ///
113    /// The original script installs:
114    /// `pip install wandb`
115    /// With optional `--break-system-packages` for global installs.
116    pub fn build_install_command(&self) -> ShellCommand {
117        let use_break = self.config.method == InstallMethod::Global
118            || self.config.method == InstallMethod::Auto;
119
120        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
121        if use_break {
122            args.push("--break-system-packages".to_string());
123        }
124        args.push("--upgrade".to_string());
125        args.push("wandb".to_string());
126
127        ShellCommand {
128            program: self.config.python_bin.clone(),
129            args,
130            env: vec![],
131        }
132    }
133}
134
135// ===========================================================================
136// Tests
137// ===========================================================================
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    // --- VAL-INSTALL-021: W&B correct pip command ---
144
145    #[test]
146    fn test_install_command_global() {
147        let installer = WandbInstaller::new(WandbConfig {
148            method: InstallMethod::Global,
149            ..Default::default()
150        });
151        let cmd = installer.build_install_command();
152        assert_eq!(cmd.program, "python3");
153        assert!(cmd.args.contains(&"-m".to_string()));
154        assert!(cmd.args.contains(&"pip".to_string()));
155        assert!(cmd.args.contains(&"install".to_string()));
156        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
157        assert!(cmd.args.contains(&"wandb".to_string()));
158    }
159
160    #[test]
161    fn test_install_command_venv() {
162        let installer = WandbInstaller::new(WandbConfig {
163            method: InstallMethod::Venv,
164            ..Default::default()
165        });
166        let cmd = installer.build_install_command();
167        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
168        assert!(cmd.args.contains(&"wandb".to_string()));
169    }
170
171    #[test]
172    fn test_install_command_auto() {
173        let installer = WandbInstaller::new(WandbConfig {
174            method: InstallMethod::Auto,
175            ..Default::default()
176        });
177        let cmd = installer.build_install_command();
178        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
179        assert!(cmd.args.contains(&"wandb".to_string()));
180    }
181
182    #[test]
183    fn test_command_string_format() {
184        let cmd = ShellCommand {
185            program: "python3".to_string(),
186            args: vec![
187                "-m".to_string(),
188                "pip".to_string(),
189                "install".to_string(),
190                "wandb".to_string(),
191            ],
192            env: vec![],
193        };
194        assert_eq!(cmd.to_command_string(), "python3 -m pip install wandb");
195    }
196
197    #[test]
198    fn test_command_string_with_env() {
199        let cmd = ShellCommand {
200            program: "python3".to_string(),
201            args: vec![
202                "-m".to_string(),
203                "pip".to_string(),
204                "install".to_string(),
205                "wandb".to_string(),
206            ],
207            env: vec![("ROCM_PATH".to_string(), "/opt/rocm".to_string())],
208        };
209        let s = cmd.to_command_string();
210        assert!(s.contains("ROCM_PATH=/opt/rocm"));
211        assert!(s.contains("python3 -m pip install wandb"));
212    }
213}
rusty_stack/installers/components/wandb.rs

rusty_stack/installers/components/
wandb.rs
