rusty_stack/installers/components/
triton.rs

1//! Triton installer — ports `scripts/install_triton_multi.sh`.
2//!
3//! Constructs correct pip install with version constraints and ROCm flags.
4//! Triton is built from source (ROCm/triton fork) with GPU arch and
5//! TRITON_ROCM=1 flags.
6//!
7//! # Validation Assertion
8//!
9//! - **VAL-INSTALL-004**: Triton installer correct pip command
10
11use std::fmt;
12use std::path::PathBuf;
13
14// ===========================================================================
15// Types
16// ===========================================================================
17
18/// Triton branch selection.
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
20pub enum TritonBranch {
21    /// Stable release branch (release/internal/3.6.x).
22    Stable,
23    /// MLIR-optimized branch (release/internal/3.6.x).
24    Mlir,
25    /// Preview/main performance branch.
26    Preview,
27}
28
29impl fmt::Display for TritonBranch {
30    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31        match self {
32            TritonBranch::Stable => write!(f, "release/internal/3.6.x"),
33            TritonBranch::Mlir => write!(f, "release/internal/3.6.x"),
34            TritonBranch::Preview => write!(f, "main_perf"),
35        }
36    }
37}
38
39/// Configuration for the Triton installer.
40#[derive(Debug, Clone)]
41pub struct TritonConfig {
42    /// ROCm channel (affects branch selection).
43    pub rocm_channel: String,
44    /// GPU architecture (e.g., "gfx1100").
45    pub gpu_arch: String,
46    /// Python binary to use.
47    pub python_bin: String,
48    /// Whether to run in dry-run mode.
49    pub dry_run: bool,
50}
51
52impl Default for TritonConfig {
53    fn default() -> Self {
54        Self {
55            rocm_channel: "latest".to_string(),
56            gpu_arch: "gfx1100".to_string(),
57            python_bin: "python3".to_string(),
58            dry_run: false,
59        }
60    }
61}
62
63/// A constructed shell command.
64#[derive(Debug, Clone, PartialEq, Eq)]
65pub struct ShellCommand {
66    /// The program to run.
67    pub program: String,
68    /// Arguments to pass.
69    pub args: Vec<String>,
70    /// Environment variables to set.
71    pub env: Vec<(String, String)>,
72    /// Working directory for the command.
73    pub working_dir: Option<PathBuf>,
74}
75
76impl ShellCommand {
77    /// Format as a shell command string.
78    pub fn to_command_string(&self) -> String {
79        let env_prefix = self
80            .env
81            .iter()
82            .map(|(k, v)| format!("{k}={v}"))
83            .collect::<Vec<_>>()
84            .join(" ");
85        let cmd = if self.args.is_empty() {
86            self.program.clone()
87        } else {
88            format!("{} {}", self.program, self.args.join(" "))
89        };
90        if env_prefix.is_empty() {
91            cmd
92        } else {
93            format!("{env_prefix} {cmd}")
94        }
95    }
96}
97
98/// The Triton installer.
99pub struct TritonInstaller {
100    config: TritonConfig,
101}
102
103impl TritonInstaller {
104    /// Create a new Triton installer with the given config.
105    pub fn new(config: TritonConfig) -> Self {
106        Self { config }
107    }
108
109    /// Create with default config.
110    pub fn with_defaults() -> Self {
111        Self::new(TritonConfig::default())
112    }
113
114    // -----------------------------------------------------------------------
115    // Branch selection
116    // -----------------------------------------------------------------------
117
118    /// Determine the Triton branch to use based on channel and GPU arch.
119    ///
120    /// Matches the original script logic:
121    /// - Preview channel or gfx12xx GPUs → preview branch
122    /// - Otherwise → MLIR branch with stable fallback
123    pub fn select_branch(&self) -> TritonBranch {
124        if self.config.rocm_channel == "preview" || self.config.gpu_arch.starts_with("gfx12") {
125            TritonBranch::Preview
126        } else {
127            TritonBranch::Mlir
128        }
129    }
130
131    /// Get the git clone URL.
132    pub fn git_clone_url(&self) -> &'static str {
133        "https://github.com/ROCm/triton.git"
134    }
135
136    // -----------------------------------------------------------------------
137    // Command construction (VAL-INSTALL-004)
138    // -----------------------------------------------------------------------
139
140    /// Construct the git clone command.
141    pub fn build_git_clone_command(&self, target_dir: &str) -> ShellCommand {
142        ShellCommand {
143            program: "git".to_string(),
144            args: vec![
145                "clone".to_string(),
146                self.git_clone_url().to_string(),
147                target_dir.to_string(),
148            ],
149            env: vec![],
150            working_dir: None,
151        }
152    }
153
154    /// Construct the git checkout command for the selected branch.
155    pub fn build_git_checkout_command(&self) -> ShellCommand {
156        let branch = self.select_branch();
157        ShellCommand {
158            program: "git".to_string(),
159            args: vec!["checkout".to_string(), branch.to_string()],
160            env: vec![],
161            working_dir: None,
162        }
163    }
164
165    /// Construct the build prerequisites pip install command.
166    ///
167    /// Installs pybind11, cmake, ninja, and other build tools.
168    pub fn build_prerequisites_command(&self) -> ShellCommand {
169        let use_break = self.supports_break_system_packages();
170        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
171        if use_break {
172            args.push("--break-system-packages".to_string());
173        }
174        args.extend([
175            "--upgrade".to_string(),
176            "--no-cache-dir".to_string(),
177            "pip".to_string(),
178            "setuptools".to_string(),
179            "wheel".to_string(),
180            "packaging".to_string(),
181            "pybind11".to_string(),
182            "cmake".to_string(),
183            "ninja".to_string(),
184        ]);
185
186        ShellCommand {
187            program: self.config.python_bin.clone(),
188            args,
189            env: vec![],
190            working_dir: None,
191        }
192    }
193
194    /// Construct the Triton pip install from source command.
195    ///
196    /// This builds Triton from the cloned source with ROCm flags:
197    /// - GPU_ARCHS=gfx1100
198    /// - TRITON_ROCM=1
199    /// - MAX_JOBS=<nproc-1>
200    pub fn build_pip_install_command(&self, src_dir: &str) -> ShellCommand {
201        let use_break = self.supports_break_system_packages();
202        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
203        if use_break {
204            args.push("--break-system-packages".to_string());
205        }
206        args.extend([
207            "--no-build-isolation".to_string(),
208            "--no-deps".to_string(),
209            ".".to_string(),
210        ]);
211
212        let nproc = std::thread::available_parallelism()
213            .map(|n| n.get().saturating_sub(1))
214            .unwrap_or(4);
215
216        let env = vec![
217            ("GPU_ARCHS".to_string(), self.config.gpu_arch.clone()),
218            ("TRITON_ROCM".to_string(), "1".to_string()),
219            ("MAX_JOBS".to_string(), nproc.to_string()),
220        ];
221
222        ShellCommand {
223            program: self.config.python_bin.clone(),
224            args,
225            env,
226            working_dir: Some(PathBuf::from(src_dir)),
227        }
228    }
229
230    /// Check whether pip supports --break-system-packages.
231    fn supports_break_system_packages(&self) -> bool {
232        // Python 3.11+ has this flag; assume true on modern systems
233        true
234    }
235}
236
237// ===========================================================================
238// Tests
239// ===========================================================================
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244
245    // --- VAL-INSTALL-004: Triton installer correct pip command ---
246
247    #[test]
248    fn test_git_clone_url() {
249        let installer = TritonInstaller::with_defaults();
250        assert_eq!(
251            installer.git_clone_url(),
252            "https://github.com/ROCm/triton.git"
253        );
254    }
255
256    #[test]
257    fn test_git_clone_command() {
258        let installer = TritonInstaller::with_defaults();
259        let cmd = installer.build_git_clone_command("/tmp/triton-rocm/triton");
260        assert_eq!(cmd.program, "git");
261        assert!(cmd.args.contains(&"clone".to_string()));
262        assert!(cmd.args.iter().any(|a| a.contains("ROCm/triton")));
263    }
264
265    #[test]
266    fn test_git_checkout_mlir() {
267        let installer = TritonInstaller::new(TritonConfig {
268            rocm_channel: "latest".to_string(),
269            gpu_arch: "gfx1100".to_string(),
270            ..Default::default()
271        });
272        let cmd = installer.build_git_checkout_command();
273        assert_eq!(cmd.program, "git");
274        assert!(cmd.args.contains(&"checkout".to_string()));
275        // MLIR branch for non-preview, non-gfx12
276        assert!(cmd.args.iter().any(|a| a == "release/internal/3.6.x"));
277    }
278
279    #[test]
280    fn test_git_checkout_preview() {
281        let installer = TritonInstaller::new(TritonConfig {
282            rocm_channel: "preview".to_string(),
283            gpu_arch: "gfx1100".to_string(),
284            ..Default::default()
285        });
286        let cmd = installer.build_git_checkout_command();
287        assert!(cmd.args.iter().any(|a| a == "main_perf"));
288    }
289
290    #[test]
291    fn test_git_checkout_gfx12() {
292        let installer = TritonInstaller::new(TritonConfig {
293            rocm_channel: "latest".to_string(),
294            gpu_arch: "gfx1200".to_string(),
295            ..Default::default()
296        });
297        assert_eq!(installer.select_branch(), TritonBranch::Preview);
298    }
299
300    #[test]
301    fn test_prerequisites_command() {
302        let installer = TritonInstaller::with_defaults();
303        let cmd = installer.build_prerequisites_command();
304        assert_eq!(cmd.program, "python3");
305        assert!(cmd.args.contains(&"pybind11".to_string()));
306        assert!(cmd.args.contains(&"cmake".to_string()));
307        assert!(cmd.args.contains(&"ninja".to_string()));
308    }
309
310    #[test]
311    fn test_pip_install_command() {
312        let installer = TritonInstaller::new(TritonConfig {
313            gpu_arch: "gfx1100".to_string(),
314            ..Default::default()
315        });
316        let cmd = installer.build_pip_install_command("/tmp/triton-rocm/triton/python");
317        assert_eq!(cmd.program, "python3");
318        assert!(cmd.args.contains(&"--no-build-isolation".to_string()));
319        assert!(cmd.args.contains(&"--no-deps".to_string()));
320        assert!(cmd.args.contains(&".".to_string()));
321
322        // Verify ROCm environment variables
323        assert!(cmd
324            .env
325            .iter()
326            .any(|(k, v)| k == "GPU_ARCHS" && v == "gfx1100"));
327        assert!(cmd.env.iter().any(|(k, v)| k == "TRITON_ROCM" && v == "1"));
328        assert!(cmd.env.iter().any(|(k, _)| k == "MAX_JOBS"));
329    }
330
331    #[test]
332    fn test_branch_selection_latest() {
333        let installer = TritonInstaller::new(TritonConfig {
334            rocm_channel: "latest".to_string(),
335            gpu_arch: "gfx1100".to_string(),
336            ..Default::default()
337        });
338        assert_eq!(installer.select_branch(), TritonBranch::Mlir);
339    }
340
341    #[test]
342    fn test_branch_selection_stable() {
343        let installer = TritonInstaller::new(TritonConfig {
344            rocm_channel: "stable".to_string(),
345            gpu_arch: "gfx1100".to_string(),
346            ..Default::default()
347        });
348        assert_eq!(installer.select_branch(), TritonBranch::Mlir);
349    }
350
351    #[test]
352    fn test_branch_selection_preview() {
353        let installer = TritonInstaller::new(TritonConfig {
354            rocm_channel: "preview".to_string(),
355            gpu_arch: "gfx1100".to_string(),
356            ..Default::default()
357        });
358        assert_eq!(installer.select_branch(), TritonBranch::Preview);
359    }
360
361    #[test]
362    fn test_command_string_format() {
363        let cmd = ShellCommand {
364            program: "git".to_string(),
365            args: vec![
366                "clone".to_string(),
367                "https://github.com/ROCm/triton.git".to_string(),
368            ],
369            env: vec![],
370            working_dir: None,
371        };
372        assert_eq!(
373            cmd.to_command_string(),
374            "git clone https://github.com/ROCm/triton.git"
375        );
376    }
377
378    #[test]
379    fn test_command_string_with_env() {
380        let cmd = ShellCommand {
381            program: "pip".to_string(),
382            args: vec!["install".to_string()],
383            env: vec![
384                ("TRITON_ROCM".to_string(), "1".to_string()),
385                ("GPU_ARCHS".to_string(), "gfx1100".to_string()),
386            ],
387            working_dir: None,
388        };
389        let s = cmd.to_command_string();
390        assert!(s.contains("TRITON_ROCM=1"));
391        assert!(s.contains("GPU_ARCHS=gfx1100"));
392    }
393}
rusty_stack/installers/components/triton.rs

rusty_stack/installers/components/
triton.rs
