#[cfg(unix)]
use crate::platform::linux;
use std::path::PathBuf;
#[derive(Debug, Clone)]
pub struct RocmEnv {
path: Option<PathBuf>,
version: String,
}
impl RocmEnv {
pub fn detect() -> Self {
#[cfg(unix)]
{
let path = linux::detect_rocm_path();
let version = linux::get_rocm_version();
Self { path, version }
}
#[cfg(not(unix))]
{
Self {
path: None,
version: String::new(),
}
}
}
pub fn from_known(path: Option<PathBuf>, version: String) -> Self {
Self { path, version }
}
pub fn none() -> Self {
Self {
path: None,
version: String::new(),
}
}
pub fn path(&self) -> Option<&PathBuf> {
self.path.as_ref()
}
pub fn version(&self) -> &str {
&self.version
}
pub fn is_detected(&self) -> bool {
self.path.is_some()
}
pub fn bin_path(&self) -> Option<PathBuf> {
self.path
.as_ref()
.map(|p| p.join("bin"))
.filter(|p| p.is_dir())
}
pub fn lib_path(&self) -> Option<PathBuf> {
self.path.as_ref().and_then(|p| {
let lib = p.join("lib");
if lib.is_dir() {
return Some(lib);
}
let lib64 = p.join("lib64");
if lib64.is_dir() {
return Some(lib64);
}
None
})
}
pub fn include_path(&self) -> Option<PathBuf> {
self.path
.as_ref()
.map(|p| p.join("include"))
.filter(|p| p.is_dir())
}
pub fn tool_exists(&self, tool_name: &str) -> bool {
if let Some(bin) = self.bin_path() {
let tool_path = bin.join(tool_name);
if tool_path.exists() {
return true;
}
}
which_exists(tool_name)
}
pub fn tool_path(&self, tool_name: &str) -> Option<PathBuf> {
if let Some(bin) = self.bin_path() {
let tool_path = bin.join(tool_name);
if tool_path.exists() {
return Some(tool_path);
}
}
if let Ok(output) = std::process::Command::new("which").arg(tool_name).output() {
if output.status.success() {
let path_str = String::from_utf8_lossy(&output.stdout).trim().to_string();
if !path_str.is_empty() {
return Some(PathBuf::from(path_str));
}
}
}
None
}
pub fn version_major_minor(&self) -> String {
extract_major_minor(&self.version)
}
pub fn series_candidates(&self) -> Vec<String> {
let detected = self.version_major_minor();
let mut candidates = Vec::new();
if !detected.is_empty() {
candidates.push(detected);
}
for fallback in &[
"7.2", "7.1", "7.0", "6.4", "6.3", "6.2", "6.1", "6.0", "5.7",
] {
let fb = fallback.to_string();
if !candidates.contains(&fb) {
candidates.push(fb);
}
}
candidates
}
pub fn torch_index_url(&self, channel: &str, series: &str) -> String {
match channel {
"nightly" => format!("https://download.pytorch.org/whl/nightly/rocm{series}"),
_ => format!("https://download.pytorch.org/whl/rocm{series}"),
}
}
pub fn torch_radeon_index_url(&self, series: &str) -> String {
format!("https://repo.radeon.com/rocm/manylinux/rocm-rel-{series}/")
}
}
impl Default for RocmEnv {
fn default() -> Self {
Self::detect()
}
}
fn extract_major_minor(version: &str) -> String {
let version = version.trim();
if version.is_empty() {
return String::new();
}
let chars: Vec<char> = version.chars().collect();
let mut result = String::new();
let mut dot_count = 0;
for &ch in &chars {
if ch == '.' {
dot_count += 1;
if dot_count > 1 {
break;
}
result.push(ch);
} else if ch.is_ascii_digit() {
result.push(ch);
} else {
break;
}
}
if result.contains('.') && result.len() >= 3 {
result
} else {
String::new()
}
}
fn which_exists(cmd: &str) -> bool {
std::process::Command::new("which")
.arg(cmd)
.output()
.map(|o| o.status.success())
.unwrap_or(false)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_extract_major_minor_standard() {
assert_eq!(extract_major_minor("7.2.0"), "7.2");
assert_eq!(extract_major_minor("6.4.3"), "6.4");
}
#[test]
fn test_extract_major_minor_two_part() {
assert_eq!(extract_major_minor("7.2"), "7.2");
}
#[test]
fn test_extract_major_minor_empty() {
assert_eq!(extract_major_minor(""), "");
}
#[test]
fn test_extract_major_minor_invalid() {
assert_eq!(extract_major_minor("abc"), "");
}
#[test]
fn test_rocm_env_detect_no_panic() {
let env = RocmEnv::detect();
let _ = env.version();
}
#[test]
fn test_rocm_env_from_known() {
let env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
assert!(env.is_detected());
assert_eq!(env.version(), "7.2.0");
assert_eq!(env.path(), Some(&PathBuf::from("/opt/rocm")));
}
#[test]
fn test_rocm_env_none() {
let env = RocmEnv::none();
assert!(!env.is_detected());
assert!(env.version().is_empty());
assert!(env.path().is_none());
}
#[test]
fn test_rocm_env_version_major_minor() {
let env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
assert_eq!(env.version_major_minor(), "7.2");
}
#[test]
fn test_rocm_env_version_major_minor_empty() {
let env = RocmEnv::none();
assert!(env.version_major_minor().is_empty());
}
#[test]
fn test_rocm_env_series_candidates() {
let env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
let candidates = env.series_candidates();
assert!(!candidates.is_empty());
assert_eq!(candidates[0], "7.2");
assert!(candidates.contains(&"7.1".to_string()));
assert!(candidates.contains(&"6.4".to_string()));
}
#[test]
fn test_rocm_env_series_candidates_no_duplicates() {
let env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
let candidates = env.series_candidates();
let mut seen = std::collections::HashSet::new();
for c in &candidates {
assert!(seen.insert(c.clone()), "Duplicate series candidate: {c}");
}
}
#[test]
fn test_rocm_env_torch_index_url() {
let env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
assert_eq!(
env.torch_index_url("stable", "7.2"),
"https://download.pytorch.org/whl/rocm7.2"
);
assert_eq!(
env.torch_index_url("nightly", "7.2"),
"https://download.pytorch.org/whl/nightly/rocm7.2"
);
}
#[test]
fn test_rocm_env_torch_radeon_index_url() {
let env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
assert_eq!(
env.torch_radeon_index_url("7.2"),
"https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/"
);
}
#[test]
fn test_rocm_env_default() {
let env = RocmEnv::default();
let _ = env.version();
}
#[test]
fn test_rocm_env_bin_path_with_mock() {
let env = RocmEnv::from_known(Some(PathBuf::from("/usr")), "1.0.0".to_string());
assert!(env.bin_path().is_some());
}
#[test]
fn test_rocm_env_lib_path_with_mock() {
let env = RocmEnv::from_known(Some(PathBuf::from("/usr")), "1.0.0".to_string());
assert!(env.lib_path().is_some());
}
#[test]
fn test_rocm_env_include_path_with_mock() {
let env = RocmEnv::from_known(Some(PathBuf::from("/usr")), "1.0.0".to_string());
assert!(env.include_path().is_some());
}
#[test]
fn test_rocm_env_bin_path_none() {
let env = RocmEnv::none();
assert!(env.bin_path().is_none());
}
#[test]
fn test_rocm_env_tool_exists_bash() {
let env = RocmEnv::from_known(Some(PathBuf::from("/usr")), "1.0.0".to_string());
assert!(env.tool_exists("bash"));
}
#[test]
fn test_rocm_env_tool_exists_nonexistent() {
let env = RocmEnv::none();
assert!(!env.tool_exists("nonexistent_tool_xyz_12345"));
}
#[test]
fn test_rocm_env_tool_path_bash() {
let env = RocmEnv::from_known(Some(PathBuf::from("/usr")), "1.0.0".to_string());
assert!(env.tool_path("bash").is_some());
}
#[test]
fn test_rocm_env_tool_path_nonexistent() {
let env = RocmEnv::none();
assert!(env.tool_path("nonexistent_tool_xyz_12345").is_none());
}
#[test]
fn test_rocm_env_delegates_to_platform_linux() {
let env = RocmEnv::detect();
let platform_path = linux::detect_rocm_path();
let platform_version = linux::get_rocm_version();
assert_eq!(env.path(), platform_path.as_ref());
assert_eq!(env.version(), platform_version);
}
#[test]
fn test_rocm_version_no_panic_on_malformed() {
let env = RocmEnv::from_known(None, "not-a-version".to_string());
let mm = env.version_major_minor();
assert!(mm.is_empty()); }
#[test]
fn test_rocm_version_empty_on_not_installed() {
let env = RocmEnv::none();
assert!(env.version().is_empty());
assert!(env.version_major_minor().is_empty());
}
}