rusty_stack/installers/components/
rocm_smi.rs

1//! ROCm SMI installer — ports `scripts/install_rocm_smi.sh`.
2//!
3//! Constructs correct package manager commands per distro for installing
4//! ROCm System Management Interface. Also supports pip-based installation
5//! from the cloned repository.
6//!
7//! # Validation Assertion
8//!
9//! - **VAL-INSTALL-018**: ROCm SMI correct package command
10
11use crate::installers::common::DistroFacade;
12use crate::platform::detection::DistroFamily;
13use std::path::PathBuf;
14
15// ===========================================================================
16// Types
17// ===========================================================================
18
19/// A constructed shell command.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub struct ShellCommand {
22    /// The program to run.
23    pub program: String,
24    /// Arguments to pass.
25    pub args: Vec<String>,
26    /// Environment variables to set.
27    pub env: Vec<(String, String)>,
28    /// Working directory for the command.
29    pub working_dir: Option<PathBuf>,
30}
31
32impl ShellCommand {
33    /// Format as a shell command string.
34    pub fn to_command_string(&self) -> String {
35        let env_prefix = self
36            .env
37            .iter()
38            .map(|(k, v)| format!("{k}={v}"))
39            .collect::<Vec<_>>()
40            .join(" ");
41        let cmd = if self.args.is_empty() {
42            self.program.clone()
43        } else {
44            format!("{} {}", self.program, self.args.join(" "))
45        };
46        if env_prefix.is_empty() {
47            cmd
48        } else {
49            format!("{env_prefix} {cmd}")
50        }
51    }
52}
53
54/// Configuration for the ROCm SMI installer.
55#[derive(Debug, Clone)]
56pub struct RocmSmiConfig {
57    /// Python binary to use.
58    pub python_bin: String,
59    /// Whether to run in dry-run mode.
60    pub dry_run: bool,
61    /// Whether to force reinstall.
62    pub force: bool,
63}
64
65impl Default for RocmSmiConfig {
66    fn default() -> Self {
67        Self {
68            python_bin: "python3".to_string(),
69            dry_run: false,
70            force: false,
71        }
72    }
73}
74
75/// The ROCm SMI installer.
76pub struct RocmSmiInstaller {
77    config: RocmSmiConfig,
78}
79
80impl RocmSmiInstaller {
81    /// Create a new ROCm SMI installer with the given config.
82    pub fn new(config: RocmSmiConfig) -> Self {
83        Self { config }
84    }
85
86    /// Create with default config.
87    pub fn with_defaults() -> Self {
88        Self::new(RocmSmiConfig::default())
89    }
90
91    // -----------------------------------------------------------------------
92    // Git clone
93    // -----------------------------------------------------------------------
94
95    /// Get the ROCm SMI library git clone URL.
96    pub fn git_clone_url(&self) -> &'static str {
97        "https://github.com/RadeonOpenCompute/rocm_smi_lib.git"
98    }
99
100    /// Construct the git clone command.
101    pub fn build_git_clone_command(&self, target_dir: &str) -> ShellCommand {
102        ShellCommand {
103            program: "git".to_string(),
104            args: vec![
105                "clone".to_string(),
106                self.git_clone_url().to_string(),
107                target_dir.to_string(),
108            ],
109            env: vec![],
110            working_dir: None,
111        }
112    }
113
114    // -----------------------------------------------------------------------
115    // Package manager commands (VAL-INSTALL-018)
116    // -----------------------------------------------------------------------
117
118    /// Construct the package install command for the given distro.
119    ///
120    /// Maps to the correct package manager commands per distro:
121    /// - Debian/Ubuntu: `sudo apt install -y rocm-smi`
122    /// - Fedora/RHEL: `sudo dnf install -y rocm-smi`
123    /// - Arch: `sudo pacman -S --needed --noconfirm rocm-smi`
124    /// - SUSE: `sudo zypper install -y rocm-smi`
125    pub fn build_package_install_command(&self, distro: &DistroFacade) -> ShellCommand {
126        let (program, args) = match distro.family() {
127            DistroFamily::Debian => (
128                "sudo".to_string(),
129                vec![
130                    "apt-get".to_string(),
131                    "install".to_string(),
132                    "-y".to_string(),
133                    "rocm-smi".to_string(),
134                ],
135            ),
136            DistroFamily::Rhel => (
137                "sudo".to_string(),
138                vec![
139                    "dnf".to_string(),
140                    "install".to_string(),
141                    "-y".to_string(),
142                    "rocm-smi".to_string(),
143                ],
144            ),
145            DistroFamily::Arch => (
146                "sudo".to_string(),
147                vec![
148                    "pacman".to_string(),
149                    "-S".to_string(),
150                    "--needed".to_string(),
151                    "--noconfirm".to_string(),
152                    "rocm-smi-lib".to_string(),
153                ],
154            ),
155            DistroFamily::Suse => (
156                "sudo".to_string(),
157                vec![
158                    "zypper".to_string(),
159                    "install".to_string(),
160                    "-y".to_string(),
161                    "rocm-smi".to_string(),
162                ],
163            ),
164            _ => (
165                "sudo".to_string(),
166                vec![
167                    "apt-get".to_string(),
168                    "install".to_string(),
169                    "-y".to_string(),
170                    "rocm-smi".to_string(),
171                ],
172            ),
173        };
174
175        ShellCommand {
176            program,
177            args,
178            env: vec![],
179            working_dir: None,
180        }
181    }
182
183    /// Construct the package update command for the given distro.
184    pub fn build_package_update_command(&self, distro: &DistroFacade) -> ShellCommand {
185        let (program, args) = match distro.family() {
186            DistroFamily::Debian => (
187                "sudo".to_string(),
188                vec![
189                    "apt-get".to_string(),
190                    "update".to_string(),
191                    "-y".to_string(),
192                ],
193            ),
194            DistroFamily::Rhel => (
195                "sudo".to_string(),
196                vec!["dnf".to_string(), "makecache".to_string()],
197            ),
198            DistroFamily::Arch => (
199                "sudo".to_string(),
200                vec!["pacman".to_string(), "-Sy".to_string()],
201            ),
202            DistroFamily::Suse => (
203                "sudo".to_string(),
204                vec!["zypper".to_string(), "refresh".to_string()],
205            ),
206            _ => (
207                "sudo".to_string(),
208                vec![
209                    "apt-get".to_string(),
210                    "update".to_string(),
211                    "-y".to_string(),
212                ],
213            ),
214        };
215
216        ShellCommand {
217            program,
218            args,
219            env: vec![],
220            working_dir: None,
221        }
222    }
223
224    /// Construct the pip install command for rocm-smi from source.
225    pub fn build_pip_install_command(&self, src_dir: &str) -> ShellCommand {
226        ShellCommand {
227            program: self.config.python_bin.clone(),
228            args: vec![
229                "-m".to_string(),
230                "pip".to_string(),
231                "install".to_string(),
232                "--break-system-packages".to_string(),
233                ".".to_string(),
234            ],
235            env: vec![],
236            working_dir: Some(PathBuf::from(src_dir)),
237        }
238    }
239
240    /// Construct the ROCm environment setup commands.
241    pub fn build_rocm_env_setup(&self) -> Vec<(String, String)> {
242        vec![
243            ("ROCM_PATH".to_string(), "/opt/rocm".to_string()),
244            (
245                "PATH".to_string(),
246                "/opt/rocm/bin:".to_string() + &std::env::var("PATH").unwrap_or_default(),
247            ),
248            (
249                "LD_LIBRARY_PATH".to_string(),
250                "/opt/rocm/lib:".to_string()
251                    + &std::env::var("LD_LIBRARY_PATH").unwrap_or_default(),
252            ),
253            ("HSA_OVERRIDE_GFX_VERSION".to_string(), "11.0.0".to_string()),
254            ("PYTORCH_ROCM_ARCH".to_string(), "gfx1100".to_string()),
255        ]
256    }
257}
258
259// ===========================================================================
260// Tests
261// ===========================================================================
262
263#[cfg(test)]
264mod tests {
265    use super::*;
266    use crate::platform::detection::{DistroInfo, PackageManager};
267
268    // --- VAL-INSTALL-018: ROCm SMI correct package command ---
269
270    #[test]
271    fn test_git_clone_url() {
272        let installer = RocmSmiInstaller::with_defaults();
273        assert_eq!(
274            installer.git_clone_url(),
275            "https://github.com/RadeonOpenCompute/rocm_smi_lib.git"
276        );
277    }
278
279    #[test]
280    fn test_git_clone_command() {
281        let installer = RocmSmiInstaller::with_defaults();
282        let cmd = installer.build_git_clone_command("/tmp/rocm_smi");
283        assert_eq!(cmd.program, "git");
284        assert!(cmd.args.contains(&"clone".to_string()));
285        assert!(cmd
286            .args
287            .iter()
288            .any(|a| a.contains("RadeonOpenCompute/rocm_smi_lib")));
289    }
290
291    #[test]
292    fn test_package_install_debian() {
293        let installer = RocmSmiInstaller::with_defaults();
294        let distro = DistroFacade::from_info(DistroInfo {
295            id: "ubuntu".to_string(),
296            family: DistroFamily::Debian,
297            pkg_manager: PackageManager::Apt,
298            ..Default::default()
299        });
300        let cmd = installer.build_package_install_command(&distro);
301        assert_eq!(cmd.program, "sudo");
302        assert!(cmd.args.contains(&"apt-get".to_string()));
303        assert!(cmd.args.contains(&"install".to_string()));
304        assert!(cmd.args.contains(&"-y".to_string()));
305        assert!(cmd.args.contains(&"rocm-smi".to_string()));
306    }
307
308    #[test]
309    fn test_package_install_rhel() {
310        let installer = RocmSmiInstaller::with_defaults();
311        let distro = DistroFacade::from_info(DistroInfo {
312            id: "fedora".to_string(),
313            family: DistroFamily::Rhel,
314            pkg_manager: PackageManager::Dnf,
315            ..Default::default()
316        });
317        let cmd = installer.build_package_install_command(&distro);
318        assert_eq!(cmd.program, "sudo");
319        assert!(cmd.args.contains(&"dnf".to_string()));
320        assert!(cmd.args.contains(&"install".to_string()));
321        assert!(cmd.args.contains(&"rocm-smi".to_string()));
322    }
323
324    #[test]
325    fn test_package_install_arch() {
326        let installer = RocmSmiInstaller::with_defaults();
327        let distro = DistroFacade::from_info(DistroInfo {
328            id: "arch".to_string(),
329            family: DistroFamily::Arch,
330            pkg_manager: PackageManager::Pacman,
331            ..Default::default()
332        });
333        let cmd = installer.build_package_install_command(&distro);
334        assert_eq!(cmd.program, "sudo");
335        assert!(cmd.args.contains(&"pacman".to_string()));
336        assert!(cmd.args.contains(&"-S".to_string()));
337        assert!(cmd.args.contains(&"--needed".to_string()));
338        assert!(cmd.args.contains(&"--noconfirm".to_string()));
339        assert!(cmd.args.contains(&"rocm-smi-lib".to_string()));
340    }
341
342    #[test]
343    fn test_package_install_suse() {
344        let installer = RocmSmiInstaller::with_defaults();
345        let distro = DistroFacade::from_info(DistroInfo {
346            id: "opensuse".to_string(),
347            family: DistroFamily::Suse,
348            pkg_manager: PackageManager::Zypper,
349            ..Default::default()
350        });
351        let cmd = installer.build_package_install_command(&distro);
352        assert_eq!(cmd.program, "sudo");
353        assert!(cmd.args.contains(&"zypper".to_string()));
354        assert!(cmd.args.contains(&"install".to_string()));
355        assert!(cmd.args.contains(&"rocm-smi".to_string()));
356    }
357
358    #[test]
359    fn test_package_update_debian() {
360        let installer = RocmSmiInstaller::with_defaults();
361        let distro = DistroFacade::from_info(DistroInfo {
362            id: "ubuntu".to_string(),
363            family: DistroFamily::Debian,
364            pkg_manager: PackageManager::Apt,
365            ..Default::default()
366        });
367        let cmd = installer.build_package_update_command(&distro);
368        assert!(cmd.args.contains(&"apt-get".to_string()));
369        assert!(cmd.args.contains(&"update".to_string()));
370    }
371
372    #[test]
373    fn test_pip_install_command() {
374        let installer = RocmSmiInstaller::with_defaults();
375        let cmd = installer.build_pip_install_command("/tmp/rocm_smi_lib");
376        assert_eq!(cmd.program, "python3");
377        assert!(cmd.args.contains(&"-m".to_string()));
378        assert!(cmd.args.contains(&"pip".to_string()));
379        assert!(cmd.args.contains(&"install".to_string()));
380        assert!(cmd.args.contains(&".".to_string()));
381    }
382
383    #[test]
384    fn test_rocm_env_setup() {
385        let installer = RocmSmiInstaller::with_defaults();
386        let env = installer.build_rocm_env_setup();
387        assert!(env
388            .iter()
389            .any(|(k, v)| k == "ROCM_PATH" && v == "/opt/rocm"));
390        assert!(env
391            .iter()
392            .any(|(k, v)| k == "HSA_OVERRIDE_GFX_VERSION" && v == "11.0.0"));
393    }
394}
rusty_stack/installers/components/rocm_smi.rs

rusty_stack/installers/components/
rocm_smi.rs
