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    /// The proven working approach uses `pip install --no-cache-dir .`
215    /// (NOT `--no-build-isolation --no-deps` which skips build dependencies).
216    /// The `--break-system-packages` flag is added for global installs.
217    pub fn build_pip_install_command(&self, src_dir: &str) -> ShellCommand {
218        let use_break = self.config.method == InstallMethod::Global
219            || self.config.method == InstallMethod::Auto;
220
221        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
222        if use_break {
223            args.push("--break-system-packages".to_string());
224        }
225        args.extend(["--no-cache-dir".to_string(), ".".to_string()]);
226
227        let primary_arch = self.resolve_primary_arch();
228        let rocm_path = "/opt/rocm".to_string();
229
230        ShellCommand {
231            program: self.config.python_bin.clone(),
232            args,
233            env: vec![
234                ("GPU_ARCH".to_string(), primary_arch.clone()),
235                ("PYTORCH_ROCM_ARCH".to_string(), primary_arch),
236                ("GPU_ARCHS".to_string(), self.config.gpu_arch.clone()),
237                ("ROCM_PATH".to_string(), rocm_path),
238                (
239                    "HSA_OVERRIDE_GFX_VERSION".to_string(),
240                    self.config
241                        .gpu_arch
242                        .split(';')
243                        .next()
244                        .map(|a| {
245                            let raw = a.trim().strip_prefix("gfx").unwrap_or("1100");
246                            match raw.len() {
247                                4 => format!("{}.{}.{}", &raw[0..2], &raw[2..3], &raw[3..4]),
248                                3 => format!("{}.{}.{}", &raw[0..1], &raw[1..2], &raw[2..3]),
249                                _ => "11.0.0".to_string(),
250                            }
251                        })
252                        .unwrap_or_else(|| "11.0.0".to_string()),
253                ),
254            ],
255            working_dir: Some(PathBuf::from(src_dir)),
256        }
257    }
258    pub fn build_rocm_env(&self, rocm_env: &RocmEnv) -> Vec<(String, String)> {
259        let rocm_path = rocm_env
260            .path()
261            .map(|p| p.to_string_lossy().to_string())
262            .unwrap_or_else(|| "/opt/rocm".to_string());
263
264        let primary_arch = self.resolve_primary_arch();
265
266        vec![
267            ("GPU_ARCH".to_string(), primary_arch.clone()),
268            ("PYTORCH_ROCM_ARCH".to_string(), primary_arch.clone()),
269            ("GPU_ARCHS".to_string(), primary_arch),
270            ("ROCM_PATH".to_string(), rocm_path),
271            ("HSA_OVERRIDE_GFX_VERSION".to_string(), "11.0.0".to_string()),
272        ]
273    }
274
275    /// Construct the AITER JIT directory environment.
276    pub fn aiter_jit_env(&self) -> Vec<(String, String)> {
277        let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
278        let jit_dir = format!("{home}/.mlstack/aiter/jit");
279        vec![("AITER_JIT_DIR".to_string(), jit_dir)]
280    }
281
282    /// Resolve the primary GPU architecture (first item before semicolon).
283    pub fn resolve_primary_arch(&self) -> String {
284        self.config
285            .gpu_arch
286            .split(';')
287            .next()
288            .unwrap_or("gfx1100")
289            .trim()
290            .to_string()
291    }
292}
293
294// ===========================================================================
295// Tests
296// ===========================================================================
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301    use std::path::PathBuf;
302
303    // --- VAL-INSTALL-016: AITER correct git clone and pip install ---
304
305    #[test]
306    fn test_git_clone_url() {
307        let installer = AiterInstaller::with_defaults();
308        assert_eq!(
309            installer.git_clone_url(),
310            "https://github.com/ROCm/aiter.git"
311        );
312    }
313
314    #[test]
315    fn test_git_clone_url_custom() {
316        let installer = AiterInstaller::new(AiterConfig {
317            repo_url: Some("https://example.com/aiter.git".to_string()),
318            ..Default::default()
319        });
320        assert_eq!(installer.git_clone_url(), "https://example.com/aiter.git");
321    }
322
323    #[test]
324    fn test_git_clone_command() {
325        let installer = AiterInstaller::with_defaults();
326        let cmd = installer.build_git_clone_command("/tmp/aiter");
327        assert_eq!(cmd.program, "git");
328        assert!(cmd.args.contains(&"clone".to_string()));
329        assert!(cmd.args.contains(&"--recursive".to_string()));
330        assert!(cmd.args.iter().any(|a| a.contains("ROCm/aiter")));
331        assert!(cmd.args.contains(&"/tmp/aiter".to_string()));
332    }
333
334    #[test]
335    fn test_deps_install_command() {
336        let installer = AiterInstaller::new(AiterConfig {
337            method: InstallMethod::Global,
338            ..Default::default()
339        });
340        let cmd = installer.build_deps_install_command();
341        assert_eq!(cmd.program, "python3");
342        assert!(cmd.args.contains(&"packaging".to_string()));
343        assert!(cmd.args.contains(&"pybind11".to_string()));
344        assert!(cmd.args.contains(&"pandas".to_string()));
345        assert!(cmd.args.contains(&"einops".to_string()));
346        assert!(cmd.args.contains(&"cmake".to_string()));
347        assert!(cmd.args.contains(&"ninja".to_string()));
348        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
349    }
350
351    #[test]
352    fn test_deps_install_command_venv() {
353        let installer = AiterInstaller::new(AiterConfig {
354            method: InstallMethod::Venv,
355            ..Default::default()
356        });
357        let cmd = installer.build_deps_install_command();
358        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
359    }
360
361    #[test]
362    fn test_pip_install_command() {
363        let installer = AiterInstaller::new(AiterConfig {
364            method: InstallMethod::Global,
365            ..Default::default()
366        });
367        let cmd = installer.build_pip_install_command("/tmp/aiter");
368        assert_eq!(cmd.program, "python3");
369        assert!(cmd.args.contains(&"--no-cache-dir".to_string()));
370        assert!(cmd.args.contains(&".".to_string()));
371        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
372        // Must NOT use --no-build-isolation (build deps are needed)
373        assert!(!cmd.args.contains(&"--no-build-isolation".to_string()));
374        // Must set ROCm environment variables
375        assert!(cmd.env.iter().any(|(k, _)| k == "GPU_ARCH"));
376        assert!(cmd.env.iter().any(|(k, _)| k == "PYTORCH_ROCM_ARCH"));
377        assert!(cmd.env.iter().any(|(k, _)| k == "ROCM_PATH"));
378        assert!(cmd.env.iter().any(|(k, _)| k == "HSA_OVERRIDE_GFX_VERSION"));
379    }
380
381    #[test]
382    fn test_build_rocm_env() {
383        let installer = AiterInstaller::new(AiterConfig {
384            gpu_arch: "gfx1100".to_string(),
385            ..Default::default()
386        });
387        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
388        let env = installer.build_rocm_env(&rocm_env);
389        assert!(env.iter().any(|(k, v)| k == "GPU_ARCH" && v == "gfx1100"));
390        assert!(env
391            .iter()
392            .any(|(k, v)| k == "PYTORCH_ROCM_ARCH" && v == "gfx1100"));
393        assert!(env
394            .iter()
395            .any(|(k, v)| k == "HSA_OVERRIDE_GFX_VERSION" && v == "11.0.0"));
396    }
397
398    #[test]
399    fn test_aiter_jit_env() {
400        let installer = AiterInstaller::with_defaults();
401        let env = installer.aiter_jit_env();
402        assert!(env.iter().any(|(k, _)| k == "AITER_JIT_DIR"));
403        assert!(env.iter().any(|(_, v)| v.contains(".mlstack/aiter/jit")));
404    }
405
406    #[test]
407    fn test_resolve_primary_arch() {
408        let installer = AiterInstaller::new(AiterConfig {
409            gpu_arch: "gfx1100;gfx1101;gfx1102".to_string(),
410            ..Default::default()
411        });
412        assert_eq!(installer.resolve_primary_arch(), "gfx1100");
413    }
414
415    #[test]
416    fn test_resolve_primary_arch_single() {
417        let installer = AiterInstaller::new(AiterConfig {
418            gpu_arch: "gfx1100".to_string(),
419            ..Default::default()
420        });
421        assert_eq!(installer.resolve_primary_arch(), "gfx1100");
422    }
423
424    // --- VAL-INSTALL-046: AITER declares dependency on PyTorch and ROCm ---
425
426    #[test]
427    fn test_dependencies() {
428        let installer = AiterInstaller::with_defaults();
429        assert!(installer.dependencies().contains(&"pytorch"));
430        assert!(installer.dependencies().contains(&"rocm"));
431    }
432
433    #[test]
434    fn test_validate_dependencies_success() {
435        let installer = AiterInstaller::with_defaults();
436        assert!(installer
437            .validate_dependencies(&["pytorch", "rocm"])
438            .is_ok());
439    }
440
441    #[test]
442    fn test_validate_dependencies_missing_pytorch() {
443        let installer = AiterInstaller::with_defaults();
444        let result = installer.validate_dependencies(&["rocm"]);
445        assert!(result.is_err());
446        assert!(result.unwrap_err().to_string().contains("pytorch"));
447    }
448
449    #[test]
450    fn test_validate_dependencies_missing_rocm() {
451        let installer = AiterInstaller::with_defaults();
452        let result = installer.validate_dependencies(&["pytorch"]);
453        assert!(result.is_err());
454        assert!(result.unwrap_err().to_string().contains("rocm"));
455    }
456}
rusty_stack/installers/components/aiter.rs

rusty_stack/installers/components/
aiter.rs
