rusty_stack/installers/components/
deepspeed.rs

1//! DeepSpeed installer — ports `scripts/install_deepspeed.sh`.
2//!
3//! Constructs correct pip install with ROCm build flags.
4//! DeepSpeed depends on PyTorch (validated via preflight check).
5//!
6//! # Validation Assertion
7//!
8//! - **VAL-INSTALL-006**: DeepSpeed installer correct pip command
9
10use crate::installers::common::RocmEnv;
11use std::fmt;
12
13// ===========================================================================
14// Types
15// ===========================================================================
16
17/// Installation method.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum InstallMethod {
20    /// Install globally.
21    Global,
22    /// Install in a virtual environment.
23    Venv,
24    /// Try global, fallback to venv.
25    Auto,
26}
27
28impl fmt::Display for InstallMethod {
29    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
30        match self {
31            InstallMethod::Global => write!(f, "global"),
32            InstallMethod::Venv => write!(f, "venv"),
33            InstallMethod::Auto => write!(f, "auto"),
34        }
35    }
36}
37
38/// Configuration for the DeepSpeed installer.
39#[derive(Debug, Clone)]
40pub struct DeepSpeedConfig {
41    /// Installation method.
42    pub method: InstallMethod,
43    /// Python binary to use.
44    pub python_bin: String,
45    /// Whether to force reinstall.
46    pub force_reinstall: bool,
47    /// Whether to run in dry-run mode.
48    pub dry_run: bool,
49}
50
51impl Default for DeepSpeedConfig {
52    fn default() -> Self {
53        Self {
54            method: InstallMethod::Auto,
55            python_bin: "python3".to_string(),
56            force_reinstall: false,
57            dry_run: false,
58        }
59    }
60}
61
62/// A constructed pip command.
63#[derive(Debug, Clone, PartialEq, Eq)]
64pub struct PipCommand {
65    /// The program to run.
66    pub program: String,
67    /// Arguments to pass.
68    pub args: Vec<String>,
69    /// Environment variables to set.
70    pub env: Vec<(String, String)>,
71}
72
73impl PipCommand {
74    /// Format as a shell command string.
75    pub fn to_command_string(&self) -> String {
76        let env_prefix = self
77            .env
78            .iter()
79            .map(|(k, v)| format!("{k}={v}"))
80            .collect::<Vec<_>>()
81            .join(" ");
82        let cmd = if self.args.is_empty() {
83            self.program.clone()
84        } else {
85            format!("{} {}", self.program, self.args.join(" "))
86        };
87        if env_prefix.is_empty() {
88            cmd
89        } else {
90            format!("{env_prefix} {cmd}")
91        }
92    }
93}
94
95/// The DeepSpeed installer.
96pub struct DeepSpeedInstaller {
97    config: DeepSpeedConfig,
98}
99
100impl DeepSpeedInstaller {
101    /// Create a new DeepSpeed installer with the given config.
102    pub fn new(config: DeepSpeedConfig) -> Self {
103        Self { config }
104    }
105
106    /// Create with default config.
107    pub fn with_defaults() -> Self {
108        Self::new(DeepSpeedConfig::default())
109    }
110
111    // -----------------------------------------------------------------------
112    // ROCm build environment (VAL-INSTALL-006)
113    // -----------------------------------------------------------------------
114
115    /// Get the ROCm environment variables needed for DeepSpeed build.
116    ///
117    /// DeepSpeed auto-detects ROCm when ROCM_HOME is set.
118    /// The original script sets:
119    /// - ROCM_HOME=/opt/rocm
120    /// - HIP_PATH=/opt/rocm
121    /// - HSA_OVERRIDE_GFX_VERSION
122    /// - PYTORCH_ROCM_ARCH
123    pub fn rocm_build_env(&self, rocm_env: &RocmEnv) -> Vec<(String, String)> {
124        let rocm_path = rocm_env
125            .path()
126            .map(|p| p.to_string_lossy().to_string())
127            .unwrap_or_else(|| "/opt/rocm".to_string());
128
129        vec![
130            ("ROCM_HOME".to_string(), rocm_path.clone()),
131            ("HIP_PATH".to_string(), rocm_path),
132            ("HSA_OVERRIDE_GFX_VERSION".to_string(), "11.0.0".to_string()),
133            ("PYTORCH_ROCM_ARCH".to_string(), "gfx1100".to_string()),
134        ]
135    }
136
137    // -----------------------------------------------------------------------
138    // Command construction
139    // -----------------------------------------------------------------------
140
141    /// Construct the pip install command for DeepSpeed dependencies.
142    ///
143    /// Installs packaging, ninja, pydantic, jsonschema, einops.
144    pub fn build_deps_install_command(&self) -> PipCommand {
145        let is_global = self.config.method == InstallMethod::Global
146            || self.config.method == InstallMethod::Auto;
147
148        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
149        if is_global {
150            args.push("--break-system-packages".to_string());
151        }
152        args.extend([
153            "packaging".to_string(),
154            "ninja".to_string(),
155            "pydantic".to_string(),
156            "jsonschema".to_string(),
157            "einops".to_string(),
158        ]);
159
160        PipCommand {
161            program: self.config.python_bin.clone(),
162            args,
163            env: vec![],
164        }
165    }
166
167    /// Construct the pip install command for DeepSpeed.
168    ///
169    /// The original script uses `pip install deepspeed einops` with
170    /// ROCM_HOME and HIP_PATH set. It does NOT set DS_ACCELERATOR=rocm
171    /// as that caused ValueError in newer versions.
172    pub fn build_install_command(&self) -> PipCommand {
173        let is_global = self.config.method == InstallMethod::Global
174            || self.config.method == InstallMethod::Auto;
175
176        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
177        if is_global {
178            args.push("--break-system-packages".to_string());
179        }
180        // --no-deps: DeepSpeed's metadata depends on `torch`; without this pip
181        // resolves/upgrade torch from PyPI (a CUDA build), overriding the ROCm
182        // torch Rusty installed — violating "never override a core component".
183        // Runtime deps are installed explicitly via build_deps_install_command.
184        // (Stage 3: single-source deps / no-override.)
185        args.push("--no-deps".to_string());
186        if self.config.force_reinstall {
187            args.push("--force-reinstall".to_string());
188        }
189        args.push("deepspeed".to_string());
190        args.push("einops".to_string());
191
192        PipCommand {
193            program: self.config.python_bin.clone(),
194            args,
195            env: vec![], // ROCm env set separately
196        }
197    }
198
199    /// Construct the retry install command with --no-deps.
200    ///
201    /// Used when the primary install fails.
202    pub fn build_no_deps_command(&self) -> PipCommand {
203        let mut args = vec![
204            "-m".to_string(),
205            "pip".to_string(),
206            "install".to_string(),
207            "--break-system-packages".to_string(),
208            "--no-deps".to_string(),
209        ];
210        args.push("deepspeed".to_string());
211
212        PipCommand {
213            program: self.config.python_bin.clone(),
214            args,
215            env: vec![],
216        }
217    }
218
219    /// Construct the force-reinstall retry command.
220    ///
221    /// Uses `--no-deps` so a force-reinstall never pulls/overrides the ROCm
222    /// torch (Stage 3: never override a core component).
223    pub fn build_force_reinstall_command(&self) -> PipCommand {
224        let mut args = vec![
225            "-m".to_string(),
226            "pip".to_string(),
227            "install".to_string(),
228            "--break-system-packages".to_string(),
229            "--force-reinstall".to_string(),
230            "--no-deps".to_string(),
231        ];
232        args.push("deepspeed".to_string());
233
234        PipCommand {
235            program: self.config.python_bin.clone(),
236            args,
237            env: vec![],
238        }
239    }
240
241    /// Get the list of required dependencies to verify after install.
242    pub fn required_dependencies(&self) -> &[&str] {
243        &[
244            "packaging",
245            "ninja",
246            "pydantic",
247            "jsonschema",
248            "einops",
249            "deepspeed",
250        ]
251    }
252}
253
254// ===========================================================================
255// Tests
256// ===========================================================================
257
258#[cfg(test)]
259mod tests {
260    use super::*;
261    use std::path::PathBuf;
262
263    // --- VAL-INSTALL-006: DeepSpeed installer correct pip command ---
264
265    #[test]
266    fn test_build_install_command_global() {
267        let installer = DeepSpeedInstaller::new(DeepSpeedConfig {
268            method: InstallMethod::Global,
269            ..Default::default()
270        });
271        let cmd = installer.build_install_command();
272        assert_eq!(cmd.program, "python3");
273        assert!(cmd.args.contains(&"-m".to_string()));
274        assert!(cmd.args.contains(&"pip".to_string()));
275        assert!(cmd.args.contains(&"install".to_string()));
276        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
277        assert!(cmd.args.contains(&"deepspeed".to_string()));
278        assert!(cmd.args.contains(&"einops".to_string()));
279    }
280
281    #[test]
282    fn test_build_install_command_venv() {
283        let installer = DeepSpeedInstaller::new(DeepSpeedConfig {
284            method: InstallMethod::Venv,
285            ..Default::default()
286        });
287        let cmd = installer.build_install_command();
288        // Venv should NOT have --break-system-packages
289        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
290        assert!(cmd.args.contains(&"deepspeed".to_string()));
291    }
292
293    #[test]
294    fn test_build_install_command_force_reinstall() {
295        let installer = DeepSpeedInstaller::new(DeepSpeedConfig {
296            force_reinstall: true,
297            ..Default::default()
298        });
299        let cmd = installer.build_install_command();
300        assert!(cmd.args.contains(&"--force-reinstall".to_string()));
301    }
302
303    #[test]
304    fn test_build_deps_command() {
305        let installer = DeepSpeedInstaller::with_defaults();
306        let cmd = installer.build_deps_install_command();
307        assert!(cmd.args.contains(&"packaging".to_string()));
308        assert!(cmd.args.contains(&"ninja".to_string()));
309        assert!(cmd.args.contains(&"pydantic".to_string()));
310        assert!(cmd.args.contains(&"jsonschema".to_string()));
311        assert!(cmd.args.contains(&"einops".to_string()));
312    }
313
314    #[test]
315    fn test_build_no_deps_command() {
316        let installer = DeepSpeedInstaller::with_defaults();
317        let cmd = installer.build_no_deps_command();
318        assert!(cmd.args.contains(&"--no-deps".to_string()));
319        assert!(cmd.args.contains(&"deepspeed".to_string()));
320    }
321
322    #[test]
323    fn test_build_force_reinstall_command() {
324        let installer = DeepSpeedInstaller::with_defaults();
325        let cmd = installer.build_force_reinstall_command();
326        assert!(cmd.args.contains(&"--force-reinstall".to_string()));
327        assert!(cmd.args.contains(&"deepspeed".to_string()));
328    }
329
330    #[test]
331    fn test_rocm_build_env() {
332        let installer = DeepSpeedInstaller::with_defaults();
333        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
334        let env = installer.rocm_build_env(&rocm_env);
335        assert!(env
336            .iter()
337            .any(|(k, v)| k == "ROCM_HOME" && v == "/opt/rocm"));
338        assert!(env.iter().any(|(k, v)| k == "HIP_PATH" && v == "/opt/rocm"));
339        assert!(env
340            .iter()
341            .any(|(k, v)| k == "HSA_OVERRIDE_GFX_VERSION" && v == "11.0.0"));
342        assert!(env
343            .iter()
344            .any(|(k, v)| k == "PYTORCH_ROCM_ARCH" && v == "gfx1100"));
345    }
346
347    #[test]
348    fn test_rocm_build_env_no_rocm() {
349        let installer = DeepSpeedInstaller::with_defaults();
350        let rocm_env = RocmEnv::none();
351        let env = installer.rocm_build_env(&rocm_env);
352        assert!(env
353            .iter()
354            .any(|(k, v)| k == "ROCM_HOME" && v == "/opt/rocm"));
355    }
356
357    #[test]
358    fn test_required_dependencies() {
359        let installer = DeepSpeedInstaller::with_defaults();
360        let deps = installer.required_dependencies();
361        assert!(deps.contains(&"deepspeed"));
362        assert!(deps.contains(&"einops"));
363        assert!(deps.contains(&"ninja"));
364    }
365
366    #[test]
367    fn test_command_string_format() {
368        let cmd = PipCommand {
369            program: "python3".to_string(),
370            args: vec![
371                "-m".to_string(),
372                "pip".to_string(),
373                "install".to_string(),
374                "deepspeed".to_string(),
375            ],
376            env: vec![("ROCM_HOME".to_string(), "/opt/rocm".to_string())],
377        };
378        let s = cmd.to_command_string();
379        assert!(s.contains("ROCM_HOME=/opt/rocm"));
380        assert!(s.contains("python3 -m pip install deepspeed"));
381    }
382}
rusty_stack/installers/components/deepspeed.rs

rusty_stack/installers/components/
deepspeed.rs
