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

rusty_stack/installers/components/
aiter.rs
