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

rusty_stack/installers/components/
pytorch.rs
