rusty_stack/installers/components/

pytorch.rs

//! PyTorch installer — ports `scripts/install_pytorch_rocm.sh`.
//!
//! Constructs correct pip install commands with the right index URL for
//! ROCm version + Python version wheel selection.
//!
//! # Validation Assertion
//!
//! - **VAL-INSTALL-003**: PyTorch installer correct wheel selection

use std::fmt;

// ===========================================================================
// Types
// ===========================================================================

/// PyTorch release channel.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TorchChannel {
    /// Stable releases only.
    Stable,
    /// Latest (stable fallback to nightly).
    Latest,
    /// Nightly builds.
    Nightly,
}

impl fmt::Display for TorchChannel {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            TorchChannel::Stable => write!(f, "stable"),
            TorchChannel::Latest => write!(f, "latest"),
            TorchChannel::Nightly => write!(f, "nightly"),
        }
    }
}

impl TorchChannel {
    /// Parse from a choice number (1-3).
    pub fn from_choice(choice: u8) -> Option<Self> {
        match choice {
            1 => Some(TorchChannel::Stable),
            2 => Some(TorchChannel::Latest),
            3 => Some(TorchChannel::Nightly),
            _ => None,
        }
    }
}

/// AMD's PyTorch wheel index for ROCm 10.x (TheRock) releases.
pub const AMD_WHL_NEXT_INDEX: &str = "https://stable.repo.amd.com/rocm/whl-next/";

/// Installation method.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum InstallMethod {
    /// Install globally.
    Global,
    /// Install in a virtual environment.
    Venv,
    /// Try global, fallback to venv.
    Auto,
}

impl fmt::Display for InstallMethod {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            InstallMethod::Global => write!(f, "global"),
            InstallMethod::Venv => write!(f, "venv"),
            InstallMethod::Auto => write!(f, "auto"),
        }
    }
}

/// Configuration for the PyTorch installer.
#[derive(Debug, Clone)]
pub struct PyTorchConfig {
    /// PyTorch release channel.
    pub channel: TorchChannel,
    /// Installation method.
    pub method: InstallMethod,
    /// Python binary to use.
    pub python_bin: String,
    /// Whether to force reinstall.
    pub force_reinstall: bool,
    /// Whether to run in dry-run mode.
    pub dry_run: bool,
}

impl Default for PyTorchConfig {
    fn default() -> Self {
        Self {
            channel: TorchChannel::Latest,
            method: InstallMethod::Auto,
            python_bin: "python3".to_string(),
            force_reinstall: false,
            dry_run: false,
        }
    }
}

/// A constructed pip install command.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PipCommand {
    /// The program to run.
    pub program: String,
    /// Arguments to pass.
    pub args: Vec<String>,
}

impl PipCommand {
    /// Format as a shell command string.
    pub fn to_command_string(&self) -> String {
        if self.args.is_empty() {
            self.program.clone()
        } else {
            format!("{} {}", self.program, self.args.join(" "))
        }
    }
}

/// The PyTorch installer.
pub struct PyTorchInstaller {
    config: PyTorchConfig,
}

impl PyTorchInstaller {
    /// Create a new PyTorch installer with the given config.
    pub fn new(config: PyTorchConfig) -> Self {
        Self { config }
    }

    /// Create with default config.
    pub fn with_defaults() -> Self {
        Self::new(PyTorchConfig::default())
    }

    // -----------------------------------------------------------------------
    // Index URL construction (VAL-INSTALL-003)
    // -----------------------------------------------------------------------

    /// Get the PyTorch index URL for a given ROCm major.minor version.
    ///
    /// Matches the logic in the original script:
    /// - ROCm 7.x: tries Radeon manylinux first, then PyTorch nightly
    /// - ROCm 6.4+: uses PyTorch nightly builds
    /// - ROCm 6.3: uses stable builds
    /// - ROCm 6.0-6.2: uses stable builds for 6.2
    /// - ROCm 5.x: uses stable builds for 5.7
    /// - Fallback: ROCm 6.3 stable
    pub fn index_url_for_rocm(&self, rocm_mm: &str) -> String {
        let mm: f32 = rocm_mm.parse().unwrap_or(7.2);
        if mm >= 10.0 {
            // ROCm 10.x (TheRock) PyTorch wheels live on AMD's own PEP 503 index,
            // tagged `+rocm10.x.y`, with per-GPU `[device-*]` extras.
            AMD_WHL_NEXT_INDEX.to_string()
        } else if mm >= 7.0 {
            // PyTorch's OWN PEP 503 simple index for ROCm 7.x. NOT AMD's Radeon
            // manylinux listing (repo.radeon.com/rocm/manylinux/rocm-rel-X.Y/) —
            // that is a flat --find-links directory with no per-package /torch/
            // subpage, so `pip --index-url` against it fails with "No matching
            // distribution found for torch". PyTorch's index serves proper
            // manylinux wheels (torch-X.Y+rocm7.2) + torchvision/torchaudio.
            format!("https://download.pytorch.org/whl/rocm{rocm_mm}")
        } else if mm >= 6.4 {
            format!("https://download.pytorch.org/whl/nightly/rocm{rocm_mm}")
        } else if mm >= 6.3 {
            format!("https://download.pytorch.org/whl/rocm{rocm_mm}")
        } else if mm >= 6.0 {
            "https://download.pytorch.org/whl/rocm6.2".to_string()
        } else if mm >= 5.0 {
            "https://download.pytorch.org/whl/rocm5.7".to_string()
        } else {
            "https://download.pytorch.org/whl/rocm6.3".to_string()
        }
    }

    /// Get the PyTorch nightly index URL for a given ROCm version.
    pub fn nightly_index_url(&self, rocm_mm: &str) -> String {
        format!("https://download.pytorch.org/whl/nightly/rocm{rocm_mm}")
    }

    /// Get the Radeon manylinux index URL.
    pub fn radeon_index_url(&self, rocm_mm: &str) -> String {
        format!("https://repo.radeon.com/rocm/manylinux/rocm-rel-{rocm_mm}/")
    }

    /// Get the fallback index URL (ROCm 7.0 stable).
    pub fn fallback_index_url(&self) -> String {
        "https://download.pytorch.org/whl/rocm7.0".to_string()
    }

    // -----------------------------------------------------------------------
    // Command construction
    // -----------------------------------------------------------------------

    /// Construct pip install command for PyTorch with ROCm support.
    ///
    /// This matches the original script's pip install logic:
    /// - Uses `uv pip install` if uv is available
    /// - Uses `python3 -m pip install` as fallback
    /// - Adds `--index-url` for the correct ROCm wheel index
    /// - Adds `--break-system-packages` for global installs
    /// - Installs torch, torchvision, torchaudio
    ///
    /// `target_version` (the planner/manifest target) pins torch to an exact
    /// version + `--upgrade` for update-driven installs, so pip actually advances
    /// the installed wheel instead of no-op'ing on an already-satisfied one. A
    /// `None`/empty target (fresh TUI install) leaves torch unpinned → latest.
    /// PEP 440: `torch==2.13.0` matches the ROCm wheel `2.13.0+rocm7.2`, so the
    /// pin resolves on the PyTorch ROCm index.
    pub fn build_install_command(
        &self,
        rocm_mm: &str,
        use_uv: bool,
        target_version: Option<&str>,
    ) -> PipCommand {
        let index_url = self.index_url_for_rocm(rocm_mm);
        let is_global = self.config.method == InstallMethod::Global
            || self.config.method == InstallMethod::Auto;

        // For global installs, always use python3 -m pip with --break-system-packages
        // instead of `uv pip install --system`, which fails on uv-managed Python
        // installations with "externally managed" errors.
        let effective_use_uv = use_uv && !is_global;

        let mut args = Vec::new();

        if effective_use_uv {
            args.push("pip".to_string());
            args.push("install".to_string());
        } else {
            args.push("-m".to_string());
            args.push("pip".to_string());
            args.push("install".to_string());
            if is_global {
                args.push("--break-system-packages".to_string());
            }
        }

        args.push("--no-cache".to_string());
        // Pin to the exact target + upgrade for update-driven installs. Without
        // this pip sees the installed wheel as already satisfying an unpinned
        // `torch` and exits 0 with the OLD version still active (the false-success
        // root cause). torchvision/torchaudio stay unpinned so pip's resolver
        // picks versions compatible with the pinned torch.
        let pin = target_version.map(str::trim).filter(|v| !v.is_empty());
        if pin.is_some() {
            args.push("--upgrade".to_string());
        }
        args.push("--index-url".to_string());
        args.push(index_url);
        // ROCm 10.x wheels need the `[device-all]` extra to pull the ROCm runtime.
        let torch_name = if rocm_mm.parse::<f32>().map(|m| m >= 10.0).unwrap_or(false) {
            "torch[device-all]"
        } else {
            "torch"
        };
        match pin {
            Some(v) => args.push(format!("{torch_name}=={v}")),
            None => args.push(torch_name.to_string()),
        }
        args.push("torchvision".to_string());
        args.push("torchaudio".to_string());

        let program = if effective_use_uv {
            "uv".to_string()
        } else {
            self.config.python_bin.clone()
        };

        PipCommand { program, args }
    }

    /// Construct the fallback command when primary install fails.
    /// Tries the ROCm 7.0 stable index.
    pub fn build_fallback_command(&self, use_uv: bool) -> PipCommand {
        let is_global = self.config.method == InstallMethod::Global
            || self.config.method == InstallMethod::Auto;

        // For global installs, always use python3 -m pip with --break-system-packages
        // instead of `uv pip install --system`, which fails on uv-managed Python.
        let effective_use_uv = use_uv && !is_global;

        let mut args = Vec::new();
        if effective_use_uv {
            args.push("pip".to_string());
            args.push("install".to_string());
        } else {
            args.push("-m".to_string());
            args.push("pip".to_string());
            args.push("install".to_string());
            if is_global {
                args.push("--break-system-packages".to_string());
            }
        }

        args.push("--no-cache".to_string());
        args.push("--index-url".to_string());
        args.push(self.fallback_index_url());
        args.push("torch".to_string());
        args.push("torchvision".to_string());
        args.push("torchaudio".to_string());

        let program = if effective_use_uv {
            "uv".to_string()
        } else {
            self.config.python_bin.clone()
        };

        PipCommand { program, args }
    }

    /// Construct the common ML dependencies install command.
    ///
    /// The original script installs `torchsde` and `sentencepiece` after PyTorch.
    pub fn build_common_deps_command(&self, use_uv: bool) -> PipCommand {
        let is_global = self.config.method == InstallMethod::Global
            || self.config.method == InstallMethod::Auto;

        // For global installs, always use python3 -m pip with --break-system-packages
        // instead of `uv pip install --system`, which fails on uv-managed Python.
        let effective_use_uv = use_uv && !is_global;

        let mut args = Vec::new();
        if effective_use_uv {
            args.push("pip".to_string());
            args.push("install".to_string());
        } else {
            args.push("-m".to_string());
            args.push("pip".to_string());
            args.push("install".to_string());
            args.push("--break-system-packages".to_string());
        }
        args.push("--no-cache-dir".to_string());
        args.push("--upgrade".to_string());
        args.push("torchsde".to_string());
        args.push("sentencepiece".to_string());

        let program = if effective_use_uv {
            "uv".to_string()
        } else {
            self.config.python_bin.clone()
        };
        PipCommand { program, args }
    }
}

// ===========================================================================
// Tests
// ===========================================================================

#[cfg(test)]
mod tests {
    use super::*;

    // --- VAL-INSTALL-003: PyTorch installer correct wheel selection ---

    #[test]
    fn test_index_url_rocm_72() {
        let installer = PyTorchInstaller::with_defaults();
        let url = installer.index_url_for_rocm("7.2");
        assert_eq!(url, "https://download.pytorch.org/whl/rocm7.2");
    }

    #[test]
    fn test_index_url_rocm_101_uses_amd_whl_next() {
        let installer = PyTorchInstaller::with_defaults();
        assert_eq!(installer.index_url_for_rocm("10.1"), AMD_WHL_NEXT_INDEX);
        let cmd = installer.build_install_command("10.1", false, Some("2.14.0"));
        assert!(cmd.args.contains(&AMD_WHL_NEXT_INDEX.to_string()));
        assert!(cmd.args.contains(&"torch[device-all]==2.14.0".to_string()));
        let old = installer.build_install_command("7.2", false, None);
        assert!(old.args.contains(&"torch".to_string()));
    }

    #[test]
    fn test_index_url_rocm_70() {
        let installer = PyTorchInstaller::with_defaults();
        let url = installer.index_url_for_rocm("7.0");
        assert_eq!(url, "https://download.pytorch.org/whl/rocm7.0");
    }

    #[test]
    fn test_index_url_rocm_64() {
        let installer = PyTorchInstaller::with_defaults();
        let url = installer.index_url_for_rocm("6.4");
        assert_eq!(url, "https://download.pytorch.org/whl/nightly/rocm6.4");
    }

    #[test]
    fn test_index_url_rocm_63() {
        let installer = PyTorchInstaller::with_defaults();
        let url = installer.index_url_for_rocm("6.3");
        assert_eq!(url, "https://download.pytorch.org/whl/rocm6.3");
    }

    #[test]
    fn test_index_url_rocm_62() {
        let installer = PyTorchInstaller::with_defaults();
        let url = installer.index_url_for_rocm("6.2");
        assert_eq!(url, "https://download.pytorch.org/whl/rocm6.2");
    }

    #[test]
    fn test_index_url_rocm_57() {
        let installer = PyTorchInstaller::with_defaults();
        let url = installer.index_url_for_rocm("5.7");
        assert_eq!(url, "https://download.pytorch.org/whl/rocm5.7");
    }

    #[test]
    fn test_index_url_fallback() {
        let installer = PyTorchInstaller::with_defaults();
        let url = installer.index_url_for_rocm("4.0");
        assert_eq!(url, "https://download.pytorch.org/whl/rocm6.3");
    }

    #[test]
    fn test_nightly_index_url() {
        let installer = PyTorchInstaller::with_defaults();
        let url = installer.nightly_index_url("7.2");
        assert_eq!(url, "https://download.pytorch.org/whl/nightly/rocm7.2");
    }

    #[test]
    fn test_radeon_index_url() {
        let installer = PyTorchInstaller::with_defaults();
        let url = installer.radeon_index_url("7.2");
        assert_eq!(url, "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/");
    }

    #[test]
    fn test_build_install_command_uv_global_uses_pip() {
        // Global + uv: should use python3 -m pip (not uv --system) to avoid
        // "externally managed" errors on uv-managed Python installations.
        let installer = PyTorchInstaller::new(PyTorchConfig {
            method: InstallMethod::Global,
            ..Default::default()
        });
        let cmd = installer.build_install_command("7.2", true, None);
        assert_eq!(cmd.program, "python3");
        assert!(cmd.args.contains(&"-m".to_string()));
        assert!(cmd.args.contains(&"pip".to_string()));
        assert!(cmd.args.contains(&"install".to_string()));
        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
        assert!(cmd.args.contains(&"--index-url".to_string()));
        assert!(cmd.args.iter().any(|a| a.contains("rocm7.2")));
        assert!(cmd.args.contains(&"torch".to_string()));
        // No target → unpinned, no --upgrade.
        assert!(!cmd.args.contains(&"--upgrade".to_string()));
        assert!(!cmd.args.iter().any(|a| a.starts_with("torch==")));
    }

    #[test]
    fn test_build_install_command_pip() {
        let installer = PyTorchInstaller::new(PyTorchConfig {
            method: InstallMethod::Global,
            python_bin: "python3".to_string(),
            ..Default::default()
        });
        let cmd = installer.build_install_command("6.4", false, None);
        assert_eq!(cmd.program, "python3");
        assert!(cmd.args.contains(&"-m".to_string()));
        assert!(cmd.args.contains(&"pip".to_string()));
        assert!(cmd.args.contains(&"--break-system-packages".to_string()));
        assert!(cmd.args.iter().any(|a| a.contains("rocm6.4")));
    }

    #[test]
    fn test_build_install_command_venv_no_break_system() {
        let installer = PyTorchInstaller::new(PyTorchConfig {
            method: InstallMethod::Venv,
            ..Default::default()
        });
        let cmd = installer.build_install_command("7.2", true, None);
        assert!(!cmd.args.contains(&"--break-system-packages".to_string()));
        assert!(!cmd.args.contains(&"--system".to_string()));
    }

    #[test]
    fn test_build_install_command_pins_target_for_update() {
        // An update-driven install pins torch to the manifest target + adds
        // --upgrade, so pip advances the wheel instead of no-op'ing on the
        // already-satisfied old one. torchvision/torchaudio stay unpinned.
        let installer = PyTorchInstaller::new(PyTorchConfig {
            method: InstallMethod::Global,
            ..Default::default()
        });
        let cmd = installer.build_install_command("7.2", false, Some("2.13.0"));
        assert!(cmd.args.contains(&"--upgrade".to_string()));
        assert!(cmd.args.contains(&"torch==2.13.0".to_string()));
        assert!(cmd.args.contains(&"torchvision".to_string()));
        assert!(cmd.args.contains(&"torchaudio".to_string()));
        // An empty/whitespace target is treated as "no target".
        let cmd_empty = installer.build_install_command("7.2", false, Some("   "));
        assert!(!cmd_empty.args.contains(&"--upgrade".to_string()));
        assert!(cmd_empty.args.contains(&"torch".to_string()));
    }

    #[test]
    fn test_build_fallback_command() {
        let installer = PyTorchInstaller::new(PyTorchConfig {
            method: InstallMethod::Global,
            ..Default::default()
        });
        let cmd = installer.build_fallback_command(true);
        assert!(cmd.args.iter().any(|a| a.contains("rocm7.0")));
        assert!(cmd.args.contains(&"torch".to_string()));
    }

    #[test]
    fn test_build_common_deps_command() {
        let installer = PyTorchInstaller::with_defaults();
        let cmd = installer.build_common_deps_command(true);
        assert!(cmd.args.contains(&"torchsde".to_string()));
        assert!(cmd.args.contains(&"sentencepiece".to_string()));
    }

    #[test]
    fn test_channel_from_choice() {
        assert_eq!(TorchChannel::from_choice(1), Some(TorchChannel::Stable));
        assert_eq!(TorchChannel::from_choice(2), Some(TorchChannel::Latest));
        assert_eq!(TorchChannel::from_choice(3), Some(TorchChannel::Nightly));
        assert_eq!(TorchChannel::from_choice(0), None);
    }
}