rusty_stack/installers/components/
pytorch.rs

1//! PyTorch installer — ports `scripts/install_pytorch_rocm.sh`.
2//!
3//! Constructs correct pip install commands with the right index URL for
4//! ROCm version + Python version wheel selection.
5//!
6//! # Validation Assertion
7//!
8//! - **VAL-INSTALL-003**: PyTorch installer correct wheel selection
9
10use crate::installers::common::RocmEnv;
11use std::fmt;
12
13// ===========================================================================
14// Types
15// ===========================================================================
16
17/// PyTorch release channel.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
19pub enum TorchChannel {
20    /// Stable releases only.
21    Stable,
22    /// Latest (stable fallback to nightly).
23    Latest,
24    /// Nightly builds.
25    Nightly,
26}
27
28impl fmt::Display for TorchChannel {
29    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
30        match self {
31            TorchChannel::Stable => write!(f, "stable"),
32            TorchChannel::Latest => write!(f, "latest"),
33            TorchChannel::Nightly => write!(f, "nightly"),
34        }
35    }
36}
37
38impl TorchChannel {
39    /// Parse from a choice number (1-3).
40    pub fn from_choice(choice: u8) -> Option<Self> {
41        match choice {
42            1 => Some(TorchChannel::Stable),
43            2 => Some(TorchChannel::Latest),
44            3 => Some(TorchChannel::Nightly),
45            _ => None,
46        }
47    }
48}
49
50/// Installation method.
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52pub enum InstallMethod {
53    /// Install globally.
54    Global,
55    /// Install in a virtual environment.
56    Venv,
57    /// Try global, fallback to venv.
58    Auto,
59}
60
61impl fmt::Display for InstallMethod {
62    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
63        match self {
64            InstallMethod::Global => write!(f, "global"),
65            InstallMethod::Venv => write!(f, "venv"),
66            InstallMethod::Auto => write!(f, "auto"),
67        }
68    }
69}
70
71/// Configuration for the PyTorch installer.
72#[derive(Debug, Clone)]
73pub struct PyTorchConfig {
74    /// PyTorch release channel.
75    pub channel: TorchChannel,
76    /// Installation method.
77    pub method: InstallMethod,
78    /// Python binary to use.
79    pub python_bin: String,
80    /// Whether to force reinstall.
81    pub force_reinstall: bool,
82    /// Whether to run in dry-run mode.
83    pub dry_run: bool,
84}
85
86impl Default for PyTorchConfig {
87    fn default() -> Self {
88        Self {
89            channel: TorchChannel::Latest,
90            method: InstallMethod::Auto,
91            python_bin: "python3".to_string(),
92            force_reinstall: false,
93            dry_run: false,
94        }
95    }
96}
97
98/// A constructed pip install command.
99#[derive(Debug, Clone, PartialEq, Eq)]
100pub struct PipCommand {
101    /// The program to run.
102    pub program: String,
103    /// Arguments to pass.
104    pub args: Vec<String>,
105}
106
107impl PipCommand {
108    /// Format as a shell command string.
109    pub fn to_command_string(&self) -> String {
110        if self.args.is_empty() {
111            self.program.clone()
112        } else {
113            format!("{} {}", self.program, self.args.join(" "))
114        }
115    }
116}
117
118/// The PyTorch installer.
119pub struct PyTorchInstaller {
120    config: PyTorchConfig,
121}
122
123impl PyTorchInstaller {
124    /// Create a new PyTorch installer with the given config.
125    pub fn new(config: PyTorchConfig) -> Self {
126        Self { config }
127    }
128
129    /// Create with default config.
130    pub fn with_defaults() -> Self {
131        Self::new(PyTorchConfig::default())
132    }
133
134    // -----------------------------------------------------------------------
135    // Index URL construction (VAL-INSTALL-003)
136    // -----------------------------------------------------------------------
137
138    /// Get the PyTorch index URL for a given ROCm major.minor version.
139    ///
140    /// Matches the logic in the original script:
141    /// - ROCm 7.x: tries Radeon manylinux first, then PyTorch nightly
142    /// - ROCm 6.4+: uses PyTorch nightly builds
143    /// - ROCm 6.3: uses stable builds
144    /// - ROCm 6.0-6.2: uses stable builds for 6.2
145    /// - ROCm 5.x: uses stable builds for 5.7
146    /// - Fallback: ROCm 6.3 stable
147    pub fn index_url_for_rocm(&self, rocm_mm: &str) -> String {
148        let mm: f32 = rocm_mm.parse().unwrap_or(7.2);
149        if mm >= 7.0 {
150            // For ROCm 7.x, use the Radeon manylinux index
151            format!("https://repo.radeon.com/rocm/manylinux/rocm-rel-{rocm_mm}/")
152        } else if mm >= 6.4 {
153            format!("https://download.pytorch.org/whl/nightly/rocm{rocm_mm}")
154        } else if mm >= 6.3 {
155            format!("https://download.pytorch.org/whl/rocm{rocm_mm}")
156        } else if mm >= 6.0 {
157            "https://download.pytorch.org/whl/rocm6.2".to_string()
158        } else if mm >= 5.0 {
159            "https://download.pytorch.org/whl/rocm5.7".to_string()
160        } else {
161            "https://download.pytorch.org/whl/rocm6.3".to_string()
162        }
163    }
164
165    /// Get the PyTorch nightly index URL for a given ROCm version.
166    pub fn nightly_index_url(&self, rocm_mm: &str) -> String {
167        format!("https://download.pytorch.org/whl/nightly/rocm{rocm_mm}")
168    }
169
170    /// Get the Radeon manylinux index URL.
171    pub fn radeon_index_url(&self, rocm_mm: &str) -> String {
172        format!("https://repo.radeon.com/rocm/manylinux/rocm-rel-{rocm_mm}/")
173    }
174
175    /// Get the fallback index URL (ROCm 7.0 stable).
176    pub fn fallback_index_url(&self) -> String {
177        "https://download.pytorch.org/whl/rocm7.0".to_string()
178    }
179
180    // -----------------------------------------------------------------------
181    // Command construction
182    // -----------------------------------------------------------------------
183
184    /// Construct pip install command for PyTorch with ROCm support.
185    ///
186    /// This matches the original script's pip install logic:
187    /// - Uses `uv pip install` if uv is available
188    /// - Uses `python3 -m pip install` as fallback
189    /// - Adds `--index-url` for the correct ROCm wheel index
190    /// - Adds `--break-system-packages` for global installs
191    /// - Installs torch, torchvision, torchaudio
192    pub fn build_install_command(&self, rocm_mm: &str, use_uv: bool) -> PipCommand {
193        let index_url = self.index_url_for_rocm(rocm_mm);
194        let is_global = self.config.method == InstallMethod::Global
195            || self.config.method == InstallMethod::Auto;
196
197        // For global installs, always use python3 -m pip with --break-system-packages
198        // instead of `uv pip install --system`, which fails on uv-managed Python
199        // installations with "externally managed" errors.
200        let effective_use_uv = use_uv && !is_global;
201
202        let mut args = Vec::new();
203
204        if effective_use_uv {
205            args.push("pip".to_string());
206            args.push("install".to_string());
207        } else {
208            args.push("-m".to_string());
209            args.push("pip".to_string());
210            args.push("install".to_string());
211            if is_global {
212                args.push("--break-system-packages".to_string());
213            }
214        }
215
216        args.push("--no-cache".to_string());
217        args.push("--index-url".to_string());
218        args.push(index_url);
219        args.push("torch".to_string());
220        args.push("torchvision".to_string());
221        args.push("torchaudio".to_string());
222
223        let program = if effective_use_uv {
224            "uv".to_string()
225        } else {
226            self.config.python_bin.clone()
227        };
228
229        PipCommand { program, args }
230    }
231
232    /// Construct the fallback command when primary install fails.
233    /// Tries the ROCm 7.0 stable index.
234    pub fn build_fallback_command(&self, use_uv: bool) -> PipCommand {
235        let is_global = self.config.method == InstallMethod::Global
236            || self.config.method == InstallMethod::Auto;
237
238        // For global installs, always use python3 -m pip with --break-system-packages
239        // instead of `uv pip install --system`, which fails on uv-managed Python.
240        let effective_use_uv = use_uv && !is_global;
241
242        let mut args = Vec::new();
243        if effective_use_uv {
244            args.push("pip".to_string());
245            args.push("install".to_string());
246        } else {
247            args.push("-m".to_string());
248            args.push("pip".to_string());
249            args.push("install".to_string());
250            if is_global {
251                args.push("--break-system-packages".to_string());
252            }
253        }
254
255        args.push("--no-cache".to_string());
256        args.push("--index-url".to_string());
257        args.push(self.fallback_index_url());
258        args.push("torch".to_string());
259        args.push("torchvision".to_string());
260        args.push("torchaudio".to_string());
261
262        let program = if effective_use_uv {
263            "uv".to_string()
264        } else {
265            self.config.python_bin.clone()
266        };
267
268        PipCommand { program, args }
269    }
270
271    /// Construct the common ML dependencies install command.
272    ///
273    /// The original script installs `torchsde` and `sentencepiece` after PyTorch.
274    pub fn build_common_deps_command(&self, use_uv: bool) -> PipCommand {
275        let is_global = self.config.method == InstallMethod::Global
276            || self.config.method == InstallMethod::Auto;
277
278        // For global installs, always use python3 -m pip with --break-system-packages
279        // instead of `uv pip install --system`, which fails on uv-managed Python.
280        let effective_use_uv = use_uv && !is_global;
281
282        let mut args = Vec::new();
283        if effective_use_uv {
284            args.push("pip".to_string());
285            args.push("install".to_string());
286        } else {
287            args.push("-m".to_string());
288            args.push("pip".to_string());
289            args.push("install".to_string());
290            args.push("--break-system-packages".to_string());
291        }
292        args.push("--no-cache-dir".to_string());
293        args.push("--upgrade".to_string());
294        args.push("torchsde".to_string());
295        args.push("sentencepiece".to_string());
296
297        let program = if effective_use_uv {
298            "uv".to_string()
299        } else {
300            self.config.python_bin.clone()
301        };
302        PipCommand { program, args }
303    }
304
305    /// Construct the ROCm environment variable exports.
306    ///
307    /// Returns a list of (VAR, VALUE) pairs matching the original script.
308    pub fn rocm_env_exports(&self, rocm_env: &RocmEnv) -> Vec<(String, String)> {
309        let mut exports = vec![
310            ("HSA_OVERRIDE_GFX_VERSION".to_string(), "11.0.0".to_string()),
311            ("PYTORCH_ROCM_ARCH".to_string(), "gfx1100".to_string()),
312            (
313                "ROCM_PATH".to_string(),
314                rocm_env
315                    .path()
316                    .map(|p| p.to_string_lossy().to_string())
317                    .unwrap_or_else(|| "/opt/rocm".to_string()),
318            ),
319        ];
320
321        // HSA_TOOLS_LIB - check for rocprofiler library
322        // Try the ROCm 7.x layout first (lib/rocprofiler-sdk/), then the old layout (lib/)
323        let rocm_lib = rocm_env.path().and_then(|p| {
324            let new_layout = p.join("lib/rocprofiler-sdk/librocprofiler-sdk-tool.so");
325            let old_layout = p.join("lib/librocprofiler-sdk-tool.so");
326            if new_layout.exists() {
327                Some(new_layout)
328            } else if old_layout.exists() {
329                Some(old_layout)
330            } else {
331                None
332            }
333        });
334        if let Some(lib) = rocm_lib {
335            exports.push((
336                "HSA_TOOLS_LIB".to_string(),
337                lib.to_string_lossy().to_string(),
338            ));
339        } else {
340            exports.push(("HSA_TOOLS_LIB".to_string(), "0".to_string()));
341        }
342
343        exports
344    }
345}
346
347// ===========================================================================
348// Tests
349// ===========================================================================
350
351#[cfg(test)]
352mod tests {
353    use super::*;
354
355    // --- VAL-INSTALL-003: PyTorch installer correct wheel selection ---
356
357    #[test]
358    fn test_index_url_rocm_72() {
359        let installer = PyTorchInstaller::with_defaults();
360        let url = installer.index_url_for_rocm("7.2");
361        assert_eq!(url, "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/");
362    }
363
364    #[test]
365    fn test_index_url_rocm_70() {
366        let installer = PyTorchInstaller::with_defaults();
367        let url = installer.index_url_for_rocm("7.0");
368        assert_eq!(url, "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0/");
369    }
370
371    #[test]
372    fn test_index_url_rocm_64() {
373        let installer = PyTorchInstaller::with_defaults();
374        let url = installer.index_url_for_rocm("6.4");
375        assert_eq!(url, "https://download.pytorch.org/whl/nightly/rocm6.4");
376    }
377
378    #[test]
379    fn test_index_url_rocm_63() {
380        let installer = PyTorchInstaller::with_defaults();
381        let url = installer.index_url_for_rocm("6.3");
382        assert_eq!(url, "https://download.pytorch.org/whl/rocm6.3");
383    }
384
385    #[test]
386    fn test_index_url_rocm_62() {
387        let installer = PyTorchInstaller::with_defaults();
388        let url = installer.index_url_for_rocm("6.2");
389        assert_eq!(url, "https://download.pytorch.org/whl/rocm6.2");
390    }
391
392    #[test]
393    fn test_index_url_rocm_57() {
394        let installer = PyTorchInstaller::with_defaults();
395        let url = installer.index_url_for_rocm("5.7");
396        assert_eq!(url, "https://download.pytorch.org/whl/rocm5.7");
397    }
398
399    #[test]
400    fn test_index_url_fallback() {
401        let installer = PyTorchInstaller::with_defaults();
402        let url = installer.index_url_for_rocm("4.0");
403        assert_eq!(url, "https://download.pytorch.org/whl/rocm6.3");
404    }
405
406    #[test]
407    fn test_nightly_index_url() {
408        let installer = PyTorchInstaller::with_defaults();
409        let url = installer.nightly_index_url("7.2");
410        assert_eq!(url, "https://download.pytorch.org/whl/nightly/rocm7.2");
411    }
412
413    #[test]
414    fn test_radeon_index_url() {
415        let installer = PyTorchInstaller::with_defaults();
416        let url = installer.radeon_index_url("7.2");
417        assert_eq!(url, "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/");
418    }
419
420    #[test]
421    fn test_build_install_command_uv_global_uses_pip() {
422        // Global + uv: should use python3 -m pip (not uv --system) to avoid
423        // "externally managed" errors on uv-managed Python installations.
424        let installer = PyTorchInstaller::new(PyTorchConfig {
425            method: InstallMethod::Global,
426            ..Default::default()
427        });
428        let cmd = installer.build_install_command("7.2", true);
429        assert_eq!(cmd.program, "python3");
430        assert!(cmd.args.contains(&"-m".to_string()));
431        assert!(cmd.args.contains(&"pip".to_string()));
432        assert!(cmd.args.contains(&"install".to_string()));
433        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
434        assert!(cmd.args.contains(&"--index-url".to_string()));
435        assert!(cmd.args.iter().any(|a| a.contains("rocm-rel-7.2")));
436        assert!(cmd.args.contains(&"torch".to_string()));
437    }
438
439    #[test]
440    fn test_build_install_command_pip() {
441        let installer = PyTorchInstaller::new(PyTorchConfig {
442            method: InstallMethod::Global,
443            python_bin: "python3".to_string(),
444            ..Default::default()
445        });
446        let cmd = installer.build_install_command("6.4", false);
447        assert_eq!(cmd.program, "python3");
448        assert!(cmd.args.contains(&"-m".to_string()));
449        assert!(cmd.args.contains(&"pip".to_string()));
450        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
451        assert!(cmd.args.iter().any(|a| a.contains("rocm6.4")));
452    }
453
454    #[test]
455    fn test_build_install_command_venv_no_break_system() {
456        let installer = PyTorchInstaller::new(PyTorchConfig {
457            method: InstallMethod::Venv,
458            ..Default::default()
459        });
460        let cmd = installer.build_install_command("7.2", true);
461        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
462        assert!(!cmd.args.contains(&"--system".to_string()));
463    }
464
465    #[test]
466    fn test_build_fallback_command() {
467        let installer = PyTorchInstaller::new(PyTorchConfig {
468            method: InstallMethod::Global,
469            ..Default::default()
470        });
471        let cmd = installer.build_fallback_command(true);
472        assert!(cmd.args.iter().any(|a| a.contains("rocm7.0")));
473        assert!(cmd.args.contains(&"torch".to_string()));
474    }
475
476    #[test]
477    fn test_build_common_deps_command() {
478        let installer = PyTorchInstaller::with_defaults();
479        let cmd = installer.build_common_deps_command(true);
480        assert!(cmd.args.contains(&"torchsde".to_string()));
481        assert!(cmd.args.contains(&"sentencepiece".to_string()));
482    }
483
484    #[test]
485    fn test_rocm_env_exports() {
486        let installer = PyTorchInstaller::with_defaults();
487        let rocm_env = RocmEnv::from_known(
488            Some(std::path::PathBuf::from("/opt/rocm")),
489            "7.2.0".to_string(),
490        );
491        let exports = installer.rocm_env_exports(&rocm_env);
492        assert!(exports
493            .iter()
494            .any(|(k, v)| k == "HSA_OVERRIDE_GFX_VERSION" && v == "11.0.0"));
495        assert!(exports
496            .iter()
497            .any(|(k, v)| k == "PYTORCH_ROCM_ARCH" && v == "gfx1100"));
498        assert!(exports
499            .iter()
500            .any(|(k, v)| k == "ROCM_PATH" && v == "/opt/rocm"));
501        assert!(exports.iter().any(|(k, _)| k == "HSA_TOOLS_LIB"));
502    }
503
504    #[test]
505    fn test_channel_from_choice() {
506        assert_eq!(TorchChannel::from_choice(1), Some(TorchChannel::Stable));
507        assert_eq!(TorchChannel::from_choice(2), Some(TorchChannel::Latest));
508        assert_eq!(TorchChannel::from_choice(3), Some(TorchChannel::Nightly));
509        assert_eq!(TorchChannel::from_choice(0), None);
510    }
511}
rusty_stack/installers/components/pytorch.rs

rusty_stack/installers/components/
pytorch.rs
