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    /// Construct the ROCm environment setup for wandb.
134    pub fn build_rocm_env(&self) -> Vec<(String, String)> {
135        vec![
136            ("ROCM_PATH".to_string(), "/opt/rocm".to_string()),
137            (
138                "PATH".to_string(),
139                "/opt/rocm/bin:".to_string() + &std::env::var("PATH").unwrap_or_default(),
140            ),
141            (
142                "LD_LIBRARY_PATH".to_string(),
143                "/opt/rocm/lib:".to_string()
144                    + &std::env::var("LD_LIBRARY_PATH").unwrap_or_default(),
145            ),
146            ("HSA_OVERRIDE_GFX_VERSION".to_string(), "11.0.0".to_string()),
147            ("PYTORCH_ROCM_ARCH".to_string(), "gfx1100".to_string()),
148        ]
149    }
150}
151
152// ===========================================================================
153// Tests
154// ===========================================================================
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159
160    // --- VAL-INSTALL-021: W&B correct pip command ---
161
162    #[test]
163    fn test_install_command_global() {
164        let installer = WandbInstaller::new(WandbConfig {
165            method: InstallMethod::Global,
166            ..Default::default()
167        });
168        let cmd = installer.build_install_command();
169        assert_eq!(cmd.program, "python3");
170        assert!(cmd.args.contains(&"-m".to_string()));
171        assert!(cmd.args.contains(&"pip".to_string()));
172        assert!(cmd.args.contains(&"install".to_string()));
173        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
174        assert!(cmd.args.contains(&"wandb".to_string()));
175    }
176
177    #[test]
178    fn test_install_command_venv() {
179        let installer = WandbInstaller::new(WandbConfig {
180            method: InstallMethod::Venv,
181            ..Default::default()
182        });
183        let cmd = installer.build_install_command();
184        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
185        assert!(cmd.args.contains(&"wandb".to_string()));
186    }
187
188    #[test]
189    fn test_install_command_auto() {
190        let installer = WandbInstaller::new(WandbConfig {
191            method: InstallMethod::Auto,
192            ..Default::default()
193        });
194        let cmd = installer.build_install_command();
195        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
196        assert!(cmd.args.contains(&"wandb".to_string()));
197    }
198
199    #[test]
200    fn test_command_string_format() {
201        let cmd = ShellCommand {
202            program: "python3".to_string(),
203            args: vec![
204                "-m".to_string(),
205                "pip".to_string(),
206                "install".to_string(),
207                "wandb".to_string(),
208            ],
209            env: vec![],
210        };
211        assert_eq!(cmd.to_command_string(), "python3 -m pip install wandb");
212    }
213
214    #[test]
215    fn test_command_string_with_env() {
216        let cmd = ShellCommand {
217            program: "python3".to_string(),
218            args: vec![
219                "-m".to_string(),
220                "pip".to_string(),
221                "install".to_string(),
222                "wandb".to_string(),
223            ],
224            env: vec![("ROCM_PATH".to_string(), "/opt/rocm".to_string())],
225        };
226        let s = cmd.to_command_string();
227        assert!(s.contains("ROCM_PATH=/opt/rocm"));
228        assert!(s.contains("python3 -m pip install wandb"));
229    }
230
231    #[test]
232    fn test_rocm_env() {
233        let installer = WandbInstaller::with_defaults();
234        let env = installer.build_rocm_env();
235        assert!(env
236            .iter()
237            .any(|(k, v)| k == "ROCM_PATH" && v == "/opt/rocm"));
238        assert!(env
239            .iter()
240            .any(|(k, v)| k == "HSA_OVERRIDE_GFX_VERSION" && v == "11.0.0"));
241    }
242}
rusty_stack/installers/components/wandb.rs

rusty_stack/installers/components/
wandb.rs
