rusty_stack/installers/components/
aiter.rs

1//! AITER installer — ports `scripts/install_aiter.sh`.
2//!
3//! Constructs correct git clone URL/branch + pip install commands for
4//! AMD Iterative Tensor Runtime (AITER). AITER depends on PyTorch and ROCm.
5//!
6//! # Validation Assertions
7//!
8//! - **VAL-INSTALL-016**: AITER correct git clone and pip install
9//! - **VAL-INSTALL-046**: AITER declares dependency on PyTorch and ROCm
10
11use crate::installers::common::RocmEnv;
12use std::path::PathBuf;
13
14// ===========================================================================
15// Types
16// ===========================================================================
17
18/// Installation method.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum InstallMethod {
21    /// Install globally.
22    Global,
23    /// Install in a virtual environment.
24    Venv,
25    /// Try global, fallback to venv.
26    Auto,
27}
28
29impl std::fmt::Display for InstallMethod {
30    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31        match self {
32            InstallMethod::Global => write!(f, "global"),
33            InstallMethod::Venv => write!(f, "venv"),
34            InstallMethod::Auto => write!(f, "auto"),
35        }
36    }
37}
38
39/// Configuration for the AITER installer.
40#[derive(Debug, Clone)]
41pub struct AiterConfig {
42    /// ROCm version string (e.g., "7.2.0").
43    pub rocm_version: String,
44    /// GPU architecture (e.g., "gfx1100").
45    pub gpu_arch: String,
46    /// Python binary to use.
47    pub python_bin: String,
48    /// Installation method.
49    pub method: InstallMethod,
50    /// Whether to run in dry-run mode.
51    pub dry_run: bool,
52    /// Whether strict ROCm mode is enabled.
53    pub strict_rocm: bool,
54    /// Custom AITER repo URL (overrides default).
55    pub repo_url: Option<String>,
56}
57
58impl Default for AiterConfig {
59    fn default() -> Self {
60        Self {
61            rocm_version: "7.2.0".to_string(),
62            gpu_arch: "gfx1100".to_string(),
63            python_bin: "python3".to_string(),
64            method: InstallMethod::Auto,
65            dry_run: false,
66            strict_rocm: true,
67            repo_url: None,
68        }
69    }
70}
71
72/// A constructed shell command.
73#[derive(Debug, Clone, PartialEq, Eq)]
74pub struct ShellCommand {
75    /// The program to run.
76    pub program: String,
77    /// Arguments to pass.
78    pub args: Vec<String>,
79    /// Environment variables to set.
80    pub env: Vec<(String, String)>,
81    /// Working directory for the command.
82    pub working_dir: Option<PathBuf>,
83}
84
85impl ShellCommand {
86    /// Format as a shell command string.
87    pub fn to_command_string(&self) -> String {
88        let env_prefix = self
89            .env
90            .iter()
91            .map(|(k, v)| format!("{k}={v}"))
92            .collect::<Vec<_>>()
93            .join(" ");
94        let cmd = if self.args.is_empty() {
95            self.program.clone()
96        } else {
97            format!("{} {}", self.program, self.args.join(" "))
98        };
99        if env_prefix.is_empty() {
100            cmd
101        } else {
102            format!("{env_prefix} {cmd}")
103        }
104    }
105}
106
107/// The AITER installer.
108pub struct AiterInstaller {
109    config: AiterConfig,
110}
111
112impl AiterInstaller {
113    /// Create a new AITER installer with the given config.
114    pub fn new(config: AiterConfig) -> Self {
115        Self { config }
116    }
117
118    /// Create with default config.
119    pub fn with_defaults() -> Self {
120        Self::new(AiterConfig::default())
121    }
122
123    // -----------------------------------------------------------------------
124    // Dependencies (VAL-INSTALL-046)
125    // -----------------------------------------------------------------------
126
127    /// AITER depends on PyTorch and ROCm.
128    pub fn dependencies(&self) -> &[&str] {
129        &["pytorch", "rocm"]
130    }
131
132    /// Validate that all dependencies are satisfied.
133    pub fn validate_dependencies(&self, installed_components: &[&str]) -> anyhow::Result<()> {
134        for dep in self.dependencies() {
135            if !installed_components.contains(dep) {
136                anyhow::bail!("AITER requires '{}' to be installed first", dep);
137            }
138        }
139        Ok(())
140    }
141
142    // -----------------------------------------------------------------------
143    // Git clone URL (VAL-INSTALL-016)
144    // -----------------------------------------------------------------------
145
146    /// Get the git clone URL for AITER.
147    pub fn git_clone_url(&self) -> &str {
148        self.config
149            .repo_url
150            .as_deref()
151            .unwrap_or("https://github.com/ROCm/aiter.git")
152    }
153
154    /// Construct the git clone command.
155    pub fn build_git_clone_command(&self, target_dir: &str) -> ShellCommand {
156        ShellCommand {
157            program: "git".to_string(),
158            args: vec![
159                "clone".to_string(),
160                "--recursive".to_string(),
161                self.git_clone_url().to_string(),
162                target_dir.to_string(),
163            ],
164            env: vec![],
165            working_dir: None,
166        }
167    }
168
169    // -----------------------------------------------------------------------
170    // Command construction (VAL-INSTALL-016)
171    // -----------------------------------------------------------------------
172
173    /// Construct the pip install command for AITER build dependencies.
174    ///
175    /// The original script installs: packaging, pybind11, pandas, einops,
176    /// psutil, numpy, setuptools, wheel, typing-extensions, cmake, ninja.
177    pub fn build_deps_install_command(&self) -> ShellCommand {
178        let use_break = self.config.method == InstallMethod::Global
179            || self.config.method == InstallMethod::Auto;
180
181        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
182        if use_break {
183            args.push("--break-system-packages".to_string());
184        }
185        args.extend(["--no-cache-dir".to_string(), "--upgrade".to_string()]);
186
187        let deps = [
188            "packaging",
189            "pybind11",
190            "pandas",
191            "einops",
192            "psutil",
193            "numpy",
194            "setuptools",
195            "wheel",
196            "typing-extensions",
197            "cmake",
198            "ninja",
199        ];
200        for dep in deps {
201            args.push(dep.to_string());
202        }
203
204        ShellCommand {
205            program: self.config.python_bin.clone(),
206            args,
207            env: vec![],
208            working_dir: None,
209        }
210    }
211
212    /// Construct the pip install command for AITER from source.
213    ///
214    /// `pip install --no-cache-dir --no-deps .` — `--no-deps` is MANDATORY for
215    /// the hard-prime No-CUDA tenet: AITER's metadata declares torch/triton-family
216    /// deps, and without `--no-deps` pip would resolve/upgrade torch (a CUDA
217    /// build) from PyPI, overriding the ROCm torch. Runtime/build deps are
218    /// installed explicitly via `build_deps_install_command` (filtered). The
219    /// `--break-system-packages` flag is added for global installs.
220    pub fn build_pip_install_command(&self, src_dir: &str) -> ShellCommand {
221        let use_break = self.config.method == InstallMethod::Global
222            || self.config.method == InstallMethod::Auto;
223
224        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
225        if use_break {
226            args.push("--break-system-packages".to_string());
227        }
228        args.extend([
229            "--no-cache-dir".to_string(),
230            "--no-deps".to_string(),
231            ".".to_string(),
232        ]);
233
234        let primary_arch = self.resolve_primary_arch();
235        let rocm_path = "/opt/rocm".to_string();
236
237        ShellCommand {
238            program: self.config.python_bin.clone(),
239            args,
240            env: vec![
241                ("GPU_ARCH".to_string(), primary_arch.clone()),
242                ("PYTORCH_ROCM_ARCH".to_string(), primary_arch),
243                ("GPU_ARCHS".to_string(), self.config.gpu_arch.clone()),
244                ("ROCM_PATH".to_string(), rocm_path),
245                (
246                    "HSA_OVERRIDE_GFX_VERSION".to_string(),
247                    self.config
248                        .gpu_arch
249                        .split(';')
250                        .next()
251                        .map(|a| {
252                            let raw = a.trim().strip_prefix("gfx").unwrap_or("1100");
253                            match raw.len() {
254                                4 => format!("{}.{}.{}", &raw[0..2], &raw[2..3], &raw[3..4]),
255                                3 => format!("{}.{}.{}", &raw[0..1], &raw[1..2], &raw[2..3]),
256                                _ => "11.0.0".to_string(),
257                            }
258                        })
259                        .unwrap_or_else(|| "11.0.0".to_string()),
260                ),
261            ],
262            working_dir: Some(PathBuf::from(src_dir)),
263        }
264    }
265    pub fn build_rocm_env(&self, rocm_env: &RocmEnv) -> Vec<(String, String)> {
266        let rocm_path = rocm_env
267            .path()
268            .map(|p| p.to_string_lossy().to_string())
269            .unwrap_or_else(|| "/opt/rocm".to_string());
270
271        let primary_arch = self.resolve_primary_arch();
272
273        vec![
274            ("GPU_ARCH".to_string(), primary_arch.clone()),
275            ("PYTORCH_ROCM_ARCH".to_string(), primary_arch.clone()),
276            ("GPU_ARCHS".to_string(), primary_arch),
277            ("ROCM_PATH".to_string(), rocm_path),
278            ("HSA_OVERRIDE_GFX_VERSION".to_string(), "11.0.0".to_string()),
279        ]
280    }
281
282    /// Construct the AITER JIT directory environment.
283    pub fn aiter_jit_env(&self) -> Vec<(String, String)> {
284        let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
285        let jit_dir = format!("{home}/.mlstack/aiter/jit");
286        vec![("AITER_JIT_DIR".to_string(), jit_dir)]
287    }
288
289    /// Resolve the primary GPU architecture (first item before semicolon).
290    pub fn resolve_primary_arch(&self) -> String {
291        self.config
292            .gpu_arch
293            .split(';')
294            .next()
295            .unwrap_or("gfx1100")
296            .trim()
297            .to_string()
298    }
299}
300
301// ===========================================================================
302// Tests
303// ===========================================================================
304
305#[cfg(test)]
306mod tests {
307    use super::*;
308    use std::path::PathBuf;
309
310    // --- VAL-INSTALL-016: AITER correct git clone and pip install ---
311
312    #[test]
313    fn test_git_clone_url() {
314        let installer = AiterInstaller::with_defaults();
315        assert_eq!(
316            installer.git_clone_url(),
317            "https://github.com/ROCm/aiter.git"
318        );
319    }
320
321    #[test]
322    fn test_git_clone_url_custom() {
323        let installer = AiterInstaller::new(AiterConfig {
324            repo_url: Some("https://example.com/aiter.git".to_string()),
325            ..Default::default()
326        });
327        assert_eq!(installer.git_clone_url(), "https://example.com/aiter.git");
328    }
329
330    #[test]
331    fn test_git_clone_command() {
332        let installer = AiterInstaller::with_defaults();
333        let cmd = installer.build_git_clone_command("/tmp/aiter");
334        assert_eq!(cmd.program, "git");
335        assert!(cmd.args.contains(&"clone".to_string()));
336        assert!(cmd.args.contains(&"--recursive".to_string()));
337        assert!(cmd.args.iter().any(|a| a.contains("ROCm/aiter")));
338        assert!(cmd.args.contains(&"/tmp/aiter".to_string()));
339    }
340
341    #[test]
342    fn test_deps_install_command() {
343        let installer = AiterInstaller::new(AiterConfig {
344            method: InstallMethod::Global,
345            ..Default::default()
346        });
347        let cmd = installer.build_deps_install_command();
348        assert_eq!(cmd.program, "python3");
349        assert!(cmd.args.contains(&"packaging".to_string()));
350        assert!(cmd.args.contains(&"pybind11".to_string()));
351        assert!(cmd.args.contains(&"pandas".to_string()));
352        assert!(cmd.args.contains(&"einops".to_string()));
353        assert!(cmd.args.contains(&"cmake".to_string()));
354        assert!(cmd.args.contains(&"ninja".to_string()));
355        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
356    }
357
358    #[test]
359    fn test_deps_install_command_venv() {
360        let installer = AiterInstaller::new(AiterConfig {
361            method: InstallMethod::Venv,
362            ..Default::default()
363        });
364        let cmd = installer.build_deps_install_command();
365        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
366    }
367
368    #[test]
369    fn test_pip_install_command() {
370        let installer = AiterInstaller::new(AiterConfig {
371            method: InstallMethod::Global,
372            ..Default::default()
373        });
374        let cmd = installer.build_pip_install_command("/tmp/aiter");
375        assert_eq!(cmd.program, "python3");
376        assert!(cmd.args.contains(&"--no-cache-dir".to_string()));
377        assert!(cmd.args.contains(&".".to_string()));
378        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
379        // Must NOT use --no-build-isolation (build deps are needed)
380        assert!(!cmd.args.contains(&"--no-build-isolation".to_string()));
381        // Must set ROCm environment variables
382        assert!(cmd.env.iter().any(|(k, _)| k == "GPU_ARCH"));
383        assert!(cmd.env.iter().any(|(k, _)| k == "PYTORCH_ROCM_ARCH"));
384        assert!(cmd.env.iter().any(|(k, _)| k == "ROCM_PATH"));
385        assert!(cmd.env.iter().any(|(k, _)| k == "HSA_OVERRIDE_GFX_VERSION"));
386    }
387
388    #[test]
389    fn test_build_rocm_env() {
390        let installer = AiterInstaller::new(AiterConfig {
391            gpu_arch: "gfx1100".to_string(),
392            ..Default::default()
393        });
394        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
395        let env = installer.build_rocm_env(&rocm_env);
396        assert!(env.iter().any(|(k, v)| k == "GPU_ARCH" && v == "gfx1100"));
397        assert!(env
398            .iter()
399            .any(|(k, v)| k == "PYTORCH_ROCM_ARCH" && v == "gfx1100"));
400        assert!(env
401            .iter()
402            .any(|(k, v)| k == "HSA_OVERRIDE_GFX_VERSION" && v == "11.0.0"));
403    }
404
405    #[test]
406    fn test_aiter_jit_env() {
407        let installer = AiterInstaller::with_defaults();
408        let env = installer.aiter_jit_env();
409        assert!(env.iter().any(|(k, _)| k == "AITER_JIT_DIR"));
410        assert!(env.iter().any(|(_, v)| v.contains(".mlstack/aiter/jit")));
411    }
412
413    #[test]
414    fn test_resolve_primary_arch() {
415        let installer = AiterInstaller::new(AiterConfig {
416            gpu_arch: "gfx1100;gfx1101;gfx1102".to_string(),
417            ..Default::default()
418        });
419        assert_eq!(installer.resolve_primary_arch(), "gfx1100");
420    }
421
422    #[test]
423    fn test_resolve_primary_arch_single() {
424        let installer = AiterInstaller::new(AiterConfig {
425            gpu_arch: "gfx1100".to_string(),
426            ..Default::default()
427        });
428        assert_eq!(installer.resolve_primary_arch(), "gfx1100");
429    }
430
431    // --- VAL-INSTALL-046: AITER declares dependency on PyTorch and ROCm ---
432
433    #[test]
434    fn test_dependencies() {
435        let installer = AiterInstaller::with_defaults();
436        assert!(installer.dependencies().contains(&"pytorch"));
437        assert!(installer.dependencies().contains(&"rocm"));
438    }
439
440    #[test]
441    fn test_validate_dependencies_success() {
442        let installer = AiterInstaller::with_defaults();
443        assert!(installer
444            .validate_dependencies(&["pytorch", "rocm"])
445            .is_ok());
446    }
447
448    #[test]
449    fn test_validate_dependencies_missing_pytorch() {
450        let installer = AiterInstaller::with_defaults();
451        let result = installer.validate_dependencies(&["rocm"]);
452        assert!(result.is_err());
453        assert!(result.unwrap_err().to_string().contains("pytorch"));
454    }
455
456    #[test]
457    fn test_validate_dependencies_missing_rocm() {
458        let installer = AiterInstaller::with_defaults();
459        let result = installer.validate_dependencies(&["pytorch"]);
460        assert!(result.is_err());
461        assert!(result.unwrap_err().to_string().contains("rocm"));
462    }
463}
rusty_stack/installers/components/aiter.rs

rusty_stack/installers/components/
aiter.rs
