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("wandb".to_string());
125
126        ShellCommand {
127            program: self.config.python_bin.clone(),
128            args,
129            env: vec![],
130        }
131    }
132}
133
134// ===========================================================================
135// Tests
136// ===========================================================================
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    // --- VAL-INSTALL-021: W&B correct pip command ---
143
144    #[test]
145    fn test_install_command_global() {
146        let installer = WandbInstaller::new(WandbConfig {
147            method: InstallMethod::Global,
148            ..Default::default()
149        });
150        let cmd = installer.build_install_command();
151        assert_eq!(cmd.program, "python3");
152        assert!(cmd.args.contains(&"-m".to_string()));
153        assert!(cmd.args.contains(&"pip".to_string()));
154        assert!(cmd.args.contains(&"install".to_string()));
155        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
156        assert!(cmd.args.contains(&"wandb".to_string()));
157    }
158
159    #[test]
160    fn test_install_command_venv() {
161        let installer = WandbInstaller::new(WandbConfig {
162            method: InstallMethod::Venv,
163            ..Default::default()
164        });
165        let cmd = installer.build_install_command();
166        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
167        assert!(cmd.args.contains(&"wandb".to_string()));
168    }
169
170    #[test]
171    fn test_install_command_auto() {
172        let installer = WandbInstaller::new(WandbConfig {
173            method: InstallMethod::Auto,
174            ..Default::default()
175        });
176        let cmd = installer.build_install_command();
177        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
178        assert!(cmd.args.contains(&"wandb".to_string()));
179    }
180
181    #[test]
182    fn test_command_string_format() {
183        let cmd = ShellCommand {
184            program: "python3".to_string(),
185            args: vec![
186                "-m".to_string(),
187                "pip".to_string(),
188                "install".to_string(),
189                "wandb".to_string(),
190            ],
191            env: vec![],
192        };
193        assert_eq!(cmd.to_command_string(), "python3 -m pip install wandb");
194    }
195
196    #[test]
197    fn test_command_string_with_env() {
198        let cmd = ShellCommand {
199            program: "python3".to_string(),
200            args: vec![
201                "-m".to_string(),
202                "pip".to_string(),
203                "install".to_string(),
204                "wandb".to_string(),
205            ],
206            env: vec![("ROCM_PATH".to_string(), "/opt/rocm".to_string())],
207        };
208        let s = cmd.to_command_string();
209        assert!(s.contains("ROCM_PATH=/opt/rocm"));
210        assert!(s.contains("python3 -m pip install wandb"));
211    }
212}
rusty_stack/installers/components/wandb.rs

rusty_stack/installers/components/
wandb.rs
