rusty_stack/installers/components/
amdgpu_drivers.rs

1//! AMD GPU drivers installer — ports `scripts/install_amdgpu_drivers.sh`.
2//!
3//! Constructs correct package commands per distro family for AMDGPU driver
4//! installation.
5//!
6//! # Validation Assertions
7//!
8//! - **VAL-INSTALL-029**: AMD GPU drivers correct package command
9
10use crate::installers::common::DistroFacade;
11use crate::platform::detection::DistroFamily;
12
13// ===========================================================================
14// Types
15// ===========================================================================
16
17/// A constructed shell command.
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct ShellCommand {
20    /// The program to run.
21    pub program: String,
22    /// Arguments to pass.
23    pub args: Vec<String>,
24    /// Environment variables to set.
25    pub env: Vec<(String, String)>,
26}
27
28impl ShellCommand {
29    /// Format as a shell command string.
30    pub fn to_command_string(&self) -> String {
31        let env_prefix = self
32            .env
33            .iter()
34            .map(|(k, v)| format!("{k}={v}"))
35            .collect::<Vec<_>>()
36            .join(" ");
37        let cmd = if self.args.is_empty() {
38            self.program.clone()
39        } else {
40            format!("{} {}", self.program, self.args.join(" "))
41        };
42        if env_prefix.is_empty() {
43            cmd
44        } else {
45            format!("{env_prefix} {cmd}")
46        }
47    }
48}
49
50/// Configuration for the AMDGPU drivers installer.
51#[derive(Debug, Clone, Default)]
52pub struct AmdgpuConfig {
53    /// Whether to run in dry-run mode.
54    pub dry_run: bool,
55    /// Whether to force reinstallation.
56    pub force: bool,
57}
58
59/// The AMDGPU drivers installer.
60pub struct AmdgpuInstaller {
61    #[allow(dead_code)]
62    config: AmdgpuConfig,
63}
64
65impl AmdgpuInstaller {
66    /// Create a new AMDGPU drivers installer with the given config.
67    pub fn new(config: AmdgpuConfig) -> Self {
68        Self { config }
69    }
70
71    /// Create with default config.
72    pub fn with_defaults() -> Self {
73        Self::new(AmdgpuConfig::default())
74    }
75
76    // -----------------------------------------------------------------------
77    // Package commands per distro (VAL-INSTALL-029)
78    // -----------------------------------------------------------------------
79
80    /// Get the required system packages for the given distro.
81    pub fn required_packages(&self, distro: &DistroFacade) -> Vec<&str> {
82        match distro.family() {
83            DistroFamily::Debian => vec![
84                "amdgpu-dkms",
85                "linux-headers-$(uname -r)",
86                "linux-modules-extra-$(uname -r)",
87            ],
88            DistroFamily::Rhel => vec!["amdgpu-dkms", "kernel-devel"],
89            DistroFamily::Arch => vec!["amdgpu"],
90            DistroFamily::Suse => vec!["amdgpu-dkms", "kernel-default-devel"],
91            _ => vec!["amdgpu-dkms"],
92        }
93    }
94
95    /// Get the ROCm packages for the given distro.
96    pub fn rocm_packages(&self, _distro: &DistroFacade) -> Vec<&str> {
97        vec!["rocm-dev", "rocm-libs", "rocprofiler"]
98    }
99
100    /// Construct the package install command for the given distro.
101    ///
102    /// The original script uses:
103    /// - apt: `sudo apt install -y <packages>`
104    /// - dnf: `sudo dnf install -y <packages>`
105    /// - pacman: `sudo pacman -S --noconfirm <packages>`
106    /// - zypper: `sudo zypper install -y <packages>`
107    pub fn build_package_install_command(
108        &self,
109        distro: &DistroFacade,
110        packages: &[&str],
111    ) -> ShellCommand {
112        let (program, mut base_args) = match distro.family() {
113            DistroFamily::Debian => (
114                "sudo".to_string(),
115                vec![
116                    "apt-get".to_string(),
117                    "install".to_string(),
118                    "-y".to_string(),
119                ],
120            ),
121            DistroFamily::Rhel => (
122                "sudo".to_string(),
123                vec!["dnf".to_string(), "install".to_string(), "-y".to_string()],
124            ),
125            DistroFamily::Arch => (
126                "sudo".to_string(),
127                vec![
128                    "pacman".to_string(),
129                    "-S".to_string(),
130                    "--noconfirm".to_string(),
131                ],
132            ),
133            DistroFamily::Suse => (
134                "sudo".to_string(),
135                vec![
136                    "zypper".to_string(),
137                    "install".to_string(),
138                    "-y".to_string(),
139                ],
140            ),
141            _ => (
142                "sudo".to_string(),
143                vec![
144                    "apt-get".to_string(),
145                    "install".to_string(),
146                    "-y".to_string(),
147                ],
148            ),
149        };
150
151        for pkg in packages {
152            base_args.push(pkg.to_string());
153        }
154
155        ShellCommand {
156            program,
157            args: base_args,
158            env: vec![],
159        }
160    }
161
162    /// Construct the apt-specific download and install command for amdgpu-install.
163    ///
164    /// The original script (Ubuntu-specific):
165    /// 1. `wget https://repo.radeon.com/amdgpu-install/7.2.4/ubuntu/$CODENAME/amdgpu-install_7.2.4.70204-1_all.deb`
166    /// 2. `sudo apt install -y ./amdgpu-install_7.2.4.70204-1_all.deb`
167    /// 3. `sudo apt update`
168    /// 4. `sudo apt install -y linux-headers-$(uname -r) linux-modules-extra-$(uname -r)`
169    /// 5. `sudo apt install -y amdgpu-dkms`
170    pub fn build_amdgpu_install_deb_commands(&self, codename: &str) -> Vec<ShellCommand> {
171        let deb_url = format!(
172            "https://repo.radeon.com/amdgpu-install/7.2.4/ubuntu/{}/amdgpu-install_7.2.4.70204-1_all.deb",
173            codename
174        );
175        let deb_file = "amdgpu-install_7.2.4.70204-1_all.deb";
176
177        vec![
178            // Download
179            ShellCommand {
180                program: "wget".to_string(),
181                args: vec!["-q".to_string(), deb_url],
182                env: vec![],
183            },
184            // Install deb
185            ShellCommand {
186                program: "sudo".to_string(),
187                args: vec![
188                    "apt".to_string(),
189                    "install".to_string(),
190                    "-y".to_string(),
191                    format!("./{deb_file}"),
192                ],
193                env: vec![],
194            },
195            // Update
196            ShellCommand {
197                program: "sudo".to_string(),
198                args: vec!["apt".to_string(), "update".to_string()],
199                env: vec![],
200            },
201            // Linux headers
202            ShellCommand {
203                program: "sudo".to_string(),
204                args: vec![
205                    "apt".to_string(),
206                    "install".to_string(),
207                    "-y".to_string(),
208                    "linux-headers-$(uname -r)".to_string(),
209                    "linux-modules-extra-$(uname -r)".to_string(),
210                ],
211                env: vec![],
212            },
213            // AMDGPU DKMS
214            ShellCommand {
215                program: "sudo".to_string(),
216                args: vec![
217                    "apt".to_string(),
218                    "install".to_string(),
219                    "-y".to_string(),
220                    "amdgpu-dkms".to_string(),
221                ],
222                env: vec![],
223            },
224        ]
225    }
226
227    /// Construct environment variable setup commands.
228    pub fn build_env_setup_commands(&self, gpu_arch: &str) -> Vec<ShellCommand> {
229        vec![
230            ShellCommand {
231                program: "export".to_string(),
232                args: vec!["HSA_OVERRIDE_GFX_VERSION=11.0.0".to_string()],
233                env: vec![],
234            },
235            ShellCommand {
236                program: "export".to_string(),
237                args: vec![format!("PYTORCH_ROCM_ARCH={gpu_arch}")],
238                env: vec![],
239            },
240            ShellCommand {
241                program: "export".to_string(),
242                args: vec!["ROCM_PATH=/opt/rocm".to_string()],
243                env: vec![],
244            },
245        ]
246    }
247
248    /// Construct the verification command.
249    pub fn build_verify_command(&self) -> ShellCommand {
250        ShellCommand {
251            program: "lsmod".to_string(),
252            args: vec![],
253            env: vec![],
254        }
255    }
256}
257
258// ===========================================================================
259// Tests
260// ===========================================================================
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265    use crate::platform::detection::{DistroFamily, DistroInfo, PackageManager};
266
267    /// Helper to create a DistroFacade for testing.
268    fn make_distro(family: DistroFamily) -> DistroFacade {
269        DistroFacade::from_info(DistroInfo {
270            id: match family {
271                DistroFamily::Debian => "ubuntu",
272                DistroFamily::Rhel => "fedora",
273                DistroFamily::Arch => "arch",
274                DistroFamily::Suse => "opensuse",
275                DistroFamily::Unknown => "unknown",
276            }
277            .to_string(),
278            name: "Test Distro".to_string(),
279            version: "1.0".to_string(),
280            codename: "test".to_string(),
281            id_like: String::new(),
282            family,
283            pkg_manager: match family {
284                DistroFamily::Debian => PackageManager::Apt,
285                DistroFamily::Rhel => PackageManager::Dnf,
286                DistroFamily::Arch => PackageManager::Pacman,
287                DistroFamily::Suse => PackageManager::Zypper,
288                DistroFamily::Unknown => PackageManager::Apt,
289            },
290        })
291    }
292
293    // -----------------------------------------------------------------------
294    // VAL-INSTALL-029: AMD GPU drivers correct package command
295    // -----------------------------------------------------------------------
296
297    #[test]
298    fn test_package_install_command_debian() {
299        let installer = AmdgpuInstaller::with_defaults();
300        let distro = make_distro(DistroFamily::Debian);
301        let cmd = installer.build_package_install_command(&distro, &["amdgpu-dkms"]);
302
303        assert_eq!(cmd.program, "sudo");
304        assert!(cmd.args.contains(&"apt-get".to_string()));
305        assert!(cmd.args.contains(&"install".to_string()));
306        assert!(cmd.args.contains(&"-y".to_string()));
307        assert!(cmd.args.contains(&"amdgpu-dkms".to_string()));
308    }
309
310    #[test]
311    fn test_package_install_command_rhel() {
312        let installer = AmdgpuInstaller::with_defaults();
313        let distro = make_distro(DistroFamily::Rhel);
314        let cmd = installer.build_package_install_command(&distro, &["amdgpu-dkms"]);
315
316        assert_eq!(cmd.program, "sudo");
317        assert!(cmd.args.contains(&"dnf".to_string()));
318        assert!(cmd.args.contains(&"install".to_string()));
319        assert!(cmd.args.contains(&"-y".to_string()));
320    }
321
322    #[test]
323    fn test_package_install_command_arch() {
324        let installer = AmdgpuInstaller::with_defaults();
325        let distro = make_distro(DistroFamily::Arch);
326        let cmd = installer.build_package_install_command(&distro, &["amdgpu"]);
327
328        assert_eq!(cmd.program, "sudo");
329        assert!(cmd.args.contains(&"pacman".to_string()));
330        assert!(cmd.args.contains(&"-S".to_string()));
331        assert!(cmd.args.contains(&"--noconfirm".to_string()));
332        assert!(cmd.args.contains(&"amdgpu".to_string()));
333    }
334
335    #[test]
336    fn test_package_install_command_suse() {
337        let installer = AmdgpuInstaller::with_defaults();
338        let distro = make_distro(DistroFamily::Suse);
339        let cmd = installer.build_package_install_command(&distro, &["amdgpu-dkms"]);
340
341        assert_eq!(cmd.program, "sudo");
342        assert!(cmd.args.contains(&"zypper".to_string()));
343        assert!(cmd.args.contains(&"install".to_string()));
344        assert!(cmd.args.contains(&"-y".to_string()));
345    }
346
347    #[test]
348    fn test_required_packages_debian() {
349        let installer = AmdgpuInstaller::with_defaults();
350        let distro = make_distro(DistroFamily::Debian);
351        let pkgs = installer.required_packages(&distro);
352        assert!(pkgs.contains(&"amdgpu-dkms"));
353    }
354
355    #[test]
356    fn test_required_packages_arch() {
357        let installer = AmdgpuInstaller::with_defaults();
358        let distro = make_distro(DistroFamily::Arch);
359        let pkgs = installer.required_packages(&distro);
360        assert!(pkgs.contains(&"amdgpu"));
361    }
362
363    #[test]
364    fn test_amdgpu_install_deb_commands() {
365        let installer = AmdgpuInstaller::with_defaults();
366        let cmds = installer.build_amdgpu_install_deb_commands("noble");
367        assert_eq!(cmds.len(), 5);
368
369        // First command should be wget
370        assert_eq!(cmds[0].program, "wget");
371        assert!(cmds[0].args.iter().any(|a| a.contains("repo.radeon.com")));
372        assert!(cmds[0].args.iter().any(|a| a.contains("noble")));
373
374        // Second command should install the deb
375        assert!(cmds[1].args.iter().any(|a| a.contains("amdgpu-install")));
376
377        // Third should be apt update
378        assert!(cmds[2].args.contains(&"update".to_string()));
379
380        // Fourth should install linux headers
381        assert!(cmds[3].args.iter().any(|a| a.contains("linux-headers")));
382
383        // Fifth should install amdgpu-dkms
384        assert!(cmds[4].args.contains(&"amdgpu-dkms".to_string()));
385    }
386
387    #[test]
388    fn test_env_setup_commands() {
389        let installer = AmdgpuInstaller::with_defaults();
390        let cmds = installer.build_env_setup_commands("gfx1100");
391        assert_eq!(cmds.len(), 3);
392
393        let cmd_strings: Vec<String> = cmds.iter().map(|c| c.to_command_string()).collect();
394        assert!(cmd_strings
395            .iter()
396            .any(|s| s.contains("HSA_OVERRIDE_GFX_VERSION=11.0.0")));
397        assert!(cmd_strings
398            .iter()
399            .any(|s| s.contains("PYTORCH_ROCM_ARCH=gfx1100")));
400        assert!(cmd_strings
401            .iter()
402            .any(|s| s.contains("ROCM_PATH=/opt/rocm")));
403    }
404
405    #[test]
406    fn test_verify_command() {
407        let installer = AmdgpuInstaller::with_defaults();
408        let cmd = installer.build_verify_command();
409        assert_eq!(cmd.program, "lsmod");
410    }
411}
rusty_stack/installers/components/amdgpu_drivers.rs

rusty_stack/installers/components/
amdgpu_drivers.rs
