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        if self.config.force_reinstall {
181            args.push("--force-reinstall".to_string());
182        }
183        args.push("deepspeed".to_string());
184        args.push("einops".to_string());
185
186        PipCommand {
187            program: self.config.python_bin.clone(),
188            args,
189            env: vec![], // ROCm env set separately
190        }
191    }
192
193    /// Construct the retry install command with --no-deps.
194    ///
195    /// Used when the primary install fails.
196    pub fn build_no_deps_command(&self) -> PipCommand {
197        let mut args = vec![
198            "-m".to_string(),
199            "pip".to_string(),
200            "install".to_string(),
201            "--break-system-packages".to_string(),
202            "--no-deps".to_string(),
203        ];
204        args.push("deepspeed".to_string());
205
206        PipCommand {
207            program: self.config.python_bin.clone(),
208            args,
209            env: vec![],
210        }
211    }
212
213    /// Construct the force-reinstall retry command.
214    pub fn build_force_reinstall_command(&self) -> PipCommand {
215        let mut args = vec![
216            "-m".to_string(),
217            "pip".to_string(),
218            "install".to_string(),
219            "--break-system-packages".to_string(),
220            "--force-reinstall".to_string(),
221        ];
222        args.push("deepspeed".to_string());
223
224        PipCommand {
225            program: self.config.python_bin.clone(),
226            args,
227            env: vec![],
228        }
229    }
230
231    /// Get the list of required dependencies to verify after install.
232    pub fn required_dependencies(&self) -> &[&str] {
233        &[
234            "packaging",
235            "ninja",
236            "pydantic",
237            "jsonschema",
238            "einops",
239            "deepspeed",
240        ]
241    }
242}
243
244// ===========================================================================
245// Tests
246// ===========================================================================
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251    use std::path::PathBuf;
252
253    // --- VAL-INSTALL-006: DeepSpeed installer correct pip command ---
254
255    #[test]
256    fn test_build_install_command_global() {
257        let installer = DeepSpeedInstaller::new(DeepSpeedConfig {
258            method: InstallMethod::Global,
259            ..Default::default()
260        });
261        let cmd = installer.build_install_command();
262        assert_eq!(cmd.program, "python3");
263        assert!(cmd.args.contains(&"-m".to_string()));
264        assert!(cmd.args.contains(&"pip".to_string()));
265        assert!(cmd.args.contains(&"install".to_string()));
266        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
267        assert!(cmd.args.contains(&"deepspeed".to_string()));
268        assert!(cmd.args.contains(&"einops".to_string()));
269    }
270
271    #[test]
272    fn test_build_install_command_venv() {
273        let installer = DeepSpeedInstaller::new(DeepSpeedConfig {
274            method: InstallMethod::Venv,
275            ..Default::default()
276        });
277        let cmd = installer.build_install_command();
278        // Venv should NOT have --break-system-packages
279        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
280        assert!(cmd.args.contains(&"deepspeed".to_string()));
281    }
282
283    #[test]
284    fn test_build_install_command_force_reinstall() {
285        let installer = DeepSpeedInstaller::new(DeepSpeedConfig {
286            force_reinstall: true,
287            ..Default::default()
288        });
289        let cmd = installer.build_install_command();
290        assert!(cmd.args.contains(&"--force-reinstall".to_string()));
291    }
292
293    #[test]
294    fn test_build_deps_command() {
295        let installer = DeepSpeedInstaller::with_defaults();
296        let cmd = installer.build_deps_install_command();
297        assert!(cmd.args.contains(&"packaging".to_string()));
298        assert!(cmd.args.contains(&"ninja".to_string()));
299        assert!(cmd.args.contains(&"pydantic".to_string()));
300        assert!(cmd.args.contains(&"jsonschema".to_string()));
301        assert!(cmd.args.contains(&"einops".to_string()));
302    }
303
304    #[test]
305    fn test_build_no_deps_command() {
306        let installer = DeepSpeedInstaller::with_defaults();
307        let cmd = installer.build_no_deps_command();
308        assert!(cmd.args.contains(&"--no-deps".to_string()));
309        assert!(cmd.args.contains(&"deepspeed".to_string()));
310    }
311
312    #[test]
313    fn test_build_force_reinstall_command() {
314        let installer = DeepSpeedInstaller::with_defaults();
315        let cmd = installer.build_force_reinstall_command();
316        assert!(cmd.args.contains(&"--force-reinstall".to_string()));
317        assert!(cmd.args.contains(&"deepspeed".to_string()));
318    }
319
320    #[test]
321    fn test_rocm_build_env() {
322        let installer = DeepSpeedInstaller::with_defaults();
323        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
324        let env = installer.rocm_build_env(&rocm_env);
325        assert!(env
326            .iter()
327            .any(|(k, v)| k == "ROCM_HOME" && v == "/opt/rocm"));
328        assert!(env.iter().any(|(k, v)| k == "HIP_PATH" && v == "/opt/rocm"));
329        assert!(env
330            .iter()
331            .any(|(k, v)| k == "HSA_OVERRIDE_GFX_VERSION" && v == "11.0.0"));
332        assert!(env
333            .iter()
334            .any(|(k, v)| k == "PYTORCH_ROCM_ARCH" && v == "gfx1100"));
335    }
336
337    #[test]
338    fn test_rocm_build_env_no_rocm() {
339        let installer = DeepSpeedInstaller::with_defaults();
340        let rocm_env = RocmEnv::none();
341        let env = installer.rocm_build_env(&rocm_env);
342        assert!(env
343            .iter()
344            .any(|(k, v)| k == "ROCM_HOME" && v == "/opt/rocm"));
345    }
346
347    #[test]
348    fn test_required_dependencies() {
349        let installer = DeepSpeedInstaller::with_defaults();
350        let deps = installer.required_dependencies();
351        assert!(deps.contains(&"deepspeed"));
352        assert!(deps.contains(&"einops"));
353        assert!(deps.contains(&"ninja"));
354    }
355
356    #[test]
357    fn test_command_string_format() {
358        let cmd = PipCommand {
359            program: "python3".to_string(),
360            args: vec![
361                "-m".to_string(),
362                "pip".to_string(),
363                "install".to_string(),
364                "deepspeed".to_string(),
365            ],
366            env: vec![("ROCM_HOME".to_string(), "/opt/rocm".to_string())],
367        };
368        let s = cmd.to_command_string();
369        assert!(s.contains("ROCM_HOME=/opt/rocm"));
370        assert!(s.contains("python3 -m pip install deepspeed"));
371    }
372}
rusty_stack/installers/components/deepspeed.rs

rusty_stack/installers/components/
deepspeed.rs
