diff --git a/Cargo.lock b/Cargo.lock index 5a49645671..1aa7df0fdf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5598,9 +5598,12 @@ dependencies = [ "rattler_conda_types", "regex", "serde", + "serde_json", "temp-env", + "tempfile", "thiserror 2.0.19", "tracing", + "windows-sys 0.61.2", "winver", ] diff --git a/crates/rattler-bin/src/commands/create.rs b/crates/rattler-bin/src/commands/create.rs index 12255c04bb..11e631421e 100644 --- a/crates/rattler-bin/src/commands/create.rs +++ b/crates/rattler-bin/src/commands/create.rs @@ -246,6 +246,7 @@ pub async fn create(opt: Opt, offline: bool) -> miette::Result<()> { VirtualPackages::detect_for_platform( install_platform, &VirtualPackageOverrides::from_env(), + rattler::default_cache_dir().ok().as_deref(), ) .map(|vpkgs| vpkgs.into_generic_virtual_packages().collect::>()) .into_diagnostic() diff --git a/crates/rattler-bin/src/commands/exec.rs b/crates/rattler-bin/src/commands/exec.rs index 7b9b31c1fc..8e1a7f0b61 100644 --- a/crates/rattler-bin/src/commands/exec.rs +++ b/crates/rattler-bin/src/commands/exec.rs @@ -262,13 +262,15 @@ async fn create_exec_prefix(options: CreateExecPrefixOptions<'_>) -> miette::Res tracing::debug!("loaded {} records from repodata", total_records); // Determine virtual packages of the current platform - let virtual_packages: Vec = - VirtualPackage::detect(&VirtualPackageOverrides::from_env()) - .into_diagnostic() - .context("failed to determine virtual packages")? - .into_iter() - .map(GenericVirtualPackage::from) - .collect(); + let virtual_packages: Vec = VirtualPackage::detect( + &VirtualPackageOverrides::from_env(), + rattler::default_cache_dir().ok().as_deref(), + ) + .into_diagnostic() + .context("failed to determine virtual packages")? + .into_iter() + .map(GenericVirtualPackage::from) + .collect(); let solver_task = SolverTask { specs: specs.to_vec(), diff --git a/crates/rattler-bin/src/commands/prefix.rs b/crates/rattler-bin/src/commands/prefix.rs index bf9e4f6b57..e9f6ff1621 100644 --- a/crates/rattler-bin/src/commands/prefix.rs +++ b/crates/rattler-bin/src/commands/prefix.rs @@ -233,6 +233,7 @@ fn validate_virtual_package_dependencies( let virtual_packages = rattler_virtual_packages::VirtualPackages::detect_for_platform( platform, &rattler_virtual_packages::VirtualPackageOverrides::from_env(), + rattler::default_cache_dir().ok().as_deref(), ) .into_diagnostic() .with_context(|| format!("failed to determine virtual packages for {platform}"))? diff --git a/crates/rattler-bin/src/commands/solve.rs b/crates/rattler-bin/src/commands/solve.rs index c0535d8a01..f76f0899d7 100644 --- a/crates/rattler-bin/src/commands/solve.rs +++ b/crates/rattler-bin/src/commands/solve.rs @@ -186,9 +186,13 @@ pub async fn solve(opt: Opt, offline: bool) -> miette::Result<()> { if let Some(virtual_packages) = &opt.virtual_package { parse_virtual_packages(virtual_packages) } else { - VirtualPackages::detect_for_platform(opt.platform, &VirtualPackageOverrides::from_env()) - .map(|vpkgs| vpkgs.into_generic_virtual_packages().collect::>()) - .into_diagnostic() + VirtualPackages::detect_for_platform( + opt.platform, + &VirtualPackageOverrides::from_env(), + rattler::default_cache_dir().ok().as_deref(), + ) + .map(|vpkgs| vpkgs.into_generic_virtual_packages().collect::>()) + .into_diagnostic() } })?; diff --git a/crates/rattler-bin/src/commands/virtual_packages.rs b/crates/rattler-bin/src/commands/virtual_packages.rs index a1b6fe19a3..903b57daf4 100644 --- a/crates/rattler-bin/src/commands/virtual_packages.rs +++ b/crates/rattler-bin/src/commands/virtual_packages.rs @@ -7,11 +7,37 @@ use rattler_virtual_packages::VirtualPackageOverrides; pub struct Opt {} pub fn virtual_packages(_opt: Opt) -> miette::Result<()> { - let virtual_packages = - rattler_virtual_packages::VirtualPackage::detect(&VirtualPackageOverrides::from_env()) - .into_diagnostic()?; - for package in virtual_packages { - println!("{}", GenericVirtualPackage::from(package.clone())); + let cache_dir = rattler::default_cache_dir().ok(); + tracing::debug!( + cache_dir = %cache_dir + .as_ref() + .map_or_else(|| "".to_string(), |path| path.display().to_string()), + "detecting virtual packages" + ); + + let virtual_packages = rattler_virtual_packages::VirtualPackage::detect( + &VirtualPackageOverrides::from_env(), + cache_dir.as_deref(), + ) + .into_diagnostic()?; + + let generic_virtual_packages = virtual_packages + .into_iter() + .map(GenericVirtualPackage::from) + .collect::>(); + let package_strings = generic_virtual_packages + .iter() + .map(ToString::to_string) + .collect::>(); + + tracing::debug!( + count = package_strings.len(), + packages = ?package_strings, + "detected virtual packages" + ); + + for package in generic_virtual_packages { + println!("{package}"); } Ok(()) } diff --git a/crates/rattler_virtual_packages/Cargo.toml b/crates/rattler_virtual_packages/Cargo.toml index d99621a964..ad4e1c3386 100644 --- a/crates/rattler_virtual_packages/Cargo.toml +++ b/crates/rattler_virtual_packages/Cargo.toml @@ -17,6 +17,8 @@ once_cell = { workspace = true } rattler_conda_types = { workspace = true, default-features = false } regex = { workspace = true } serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +tempfile = { workspace = true } thiserror = { workspace = true } tracing = { workspace = true } archspec = { workspace = true } @@ -26,6 +28,11 @@ plist = { workspace = true } [target.'cfg(target_os="windows")'.dependencies] winver = { workspace = true } +windows-sys = { workspace = true, features = [ + "Win32_System_SystemInformation", + "Win32_System_Registry", + "Win32_Devices_DeviceAndDriverInstallation", +] } [dev-dependencies] temp-env = { workspace = true } diff --git a/crates/rattler_virtual_packages/src/cuda.rs b/crates/rattler_virtual_packages/src/cuda.rs index 37eb693a6c..4b74329916 100644 --- a/crates/rattler_virtual_packages/src/cuda.rs +++ b/crates/rattler_virtual_packages/src/cuda.rs @@ -7,8 +7,9 @@ //! The CUDA driver version represents the maximum CUDA version supported by the installed //! NVIDIA drivers. This is detected via: //! -//! * CUDA driver library (libcuda): Standard method -//! * nvidia-smi command: Fallback on musl systems where dynamic library loading is not supported +//! * NVIDIA Management Library (NVML): Standard method +//! * CUDA driver library (libcuda) and the nvidia-smi command: Fallbacks for systems without +//! NVML, and for musl systems where dynamic library loading is not supported //! //! ## CUDA Compute Capability (`__cuda_arch`) //! @@ -18,13 +19,92 @@ use libloading::{Library, Symbol}; use once_cell::sync::OnceCell; use rattler_conda_types::Version; +use serde::{Deserialize, Serialize}; use std::process::Command; +use std::sync::atomic::{AtomicBool, Ordering}; use std::{ - mem::MaybeUninit, - os::raw::{c_int, c_uint, c_ulong}, + os::raw::{c_int, c_uint, c_void}, + path::Path, + ptr, str::FromStr, }; +mod cache; + +const NVML_SUCCESS: c_int = 0; +const NVML_ERROR_UNINITIALIZED: c_int = 1; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum NvmlCudaVersionError { + MissingSymbol, + Nvml(c_int), + InvalidVersion, +} + +impl NvmlCudaVersionError { + fn should_retry_after_init(self) -> bool { + matches!(self, Self::Nvml(NVML_ERROR_UNINITIALIZED)) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(super) enum CudaDetectionMethod { + NvmlNoInit, + NvmlInitialized, + Libcuda, + NvidiaSmi, +} + +impl CudaDetectionMethod { + fn as_str(self) -> &'static str { + match self { + Self::NvmlNoInit => "nvml_no_init", + Self::NvmlInitialized => "nvml_initialized", + Self::Libcuda => "libcuda", + Self::NvidiaSmi => "nvidia_smi", + } + } +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub(super) struct CudaInfoSources { + pub version: Option, + pub arch: Option, +} + +impl CudaInfoSources { + fn version_str(self) -> &'static str { + self.version + .map_or("", CudaDetectionMethod::as_str) + } + + fn arch_str(self) -> &'static str { + self.arch.map_or("", CudaDetectionMethod::as_str) + } +} + +struct DetectedCudaInfo { + info: CudaInfo, + sources: CudaInfoSources, +} + +/// Converts a CUDA driver version integer (as reported by NVML/libcuda) into a [`Version`]. +/// +/// The integer is encoded as `major * 1000 + minor * 10` (e.g. `12040` for CUDA 12.4). Because the +/// FFI out-parameters are zero-initialized, an implausible value (such as `0` from an out-param the +/// driver never wrote, a negative value, or a nonsensically large one) can slip through even on a +/// `SUCCESS` return. Only values whose CUDA major version lies in `1..=99` are accepted; anything +/// else is rejected so that garbage never propagates (or gets cached). +fn parse_cuda_driver_version(version: c_int) -> Option { + // CUDA major version 1..=99, i.e. the encoded integer must be within [1000, 99990]. + if !(1000..=99_990).contains(&version) { + tracing::trace!(version, "rejecting implausible CUDA driver version integer"); + return None; + } + Version::from_str(&format!("{}.{}", version / 1000, (version % 1000) / 10)).ok() +} + /// Validates that a string is in the format "major.minor" where both parts are digits. /// /// Returns `true` if the format is valid for CUDA compute capability. @@ -79,6 +159,17 @@ pub struct CudaInfo { pub arch_info: Option, } +fn display_cuda_version(version: Option<&Version>) -> String { + version.map_or_else(|| "".to_string(), ToString::to_string) +} + +fn display_cuda_arch(arch_info: Option<&CudaArchInfo>) -> String { + arch_info.map_or_else( + || "".to_string(), + |arch| format!("{}.{}", arch.major, arch.minor), + ) +} + /// Returns comprehensive CUDA information from the current platform. /// /// This function returns both the CUDA driver version and compute capability information @@ -87,17 +178,104 @@ pub struct CudaInfo { /// /// This is more efficient than calling [`cuda_version`] and [`cuda_arch`] separately /// because the CUDA library is loaded only once. -pub fn cuda_info() -> &'static CudaInfo { - static DETECTED_CUDA_INFO: OnceCell = OnceCell::new(); - DETECTED_CUDA_INFO.get_or_init(detect_cuda_info) +/// +/// Detection runs at most once per process; the in-memory result is reused afterwards. The on-disk +/// cache, however, is synced lazily: the first call that is given a `cache_dir` reads from and/or +/// writes to it. A call that passes `None` (e.g. `EnvOverride::detect_from_host`) does not disable +/// the disk cache for the rest of the process — a later call passing `Some(cache_dir)` still +/// persists the already-detected result. Pass `None` from every call to fully disable the disk +/// cache. +pub fn cuda_info(cache_dir: Option<&Path>) -> &'static CudaInfo { + static DETECTED_CUDA_INFO: OnceCell = OnceCell::new(); + // Whether the in-memory result has been synced with the on-disk cache (read from it or written + // to it). This lets a later call with a `cache_dir` persist a result first detected without one. + static PERSISTED: AtomicBool = AtomicBool::new(false); + cuda_info_impl( + &cache::CacheEnv::current(), + &DETECTED_CUDA_INFO, + &PERSISTED, + cache_dir, + ) +} + +/// Core of [`cuda_info`], generic over the state so it can be unit-tested with local state instead +/// of the process-global statics. +fn cuda_info_impl<'a>( + env: &cache::CacheEnv, + state: &'a OnceCell, + persisted: &AtomicBool, + cache_dir: Option<&Path>, +) -> &'a CudaInfo { + if let Some(detected) = state.get() { + tracing::trace!(info = ?detected.info, "using process-cached CUDA info"); + maybe_persist(env, detected, persisted, cache_dir); + return &detected.info; + } + + let detected = state.get_or_init(|| { + // Initializing the driver to detect the GPU can be slow, so the result is cached on disk + // and reused until the cache is invalidated (reboot, driver change, GPU change, TTL). + if let Some(cache_dir) = cache_dir { + tracing::trace!(cache_dir = %cache_dir.display(), "checking CUDA info cache"); + if let Some(cached) = cache::read_with_env(env, cache_dir) { + tracing::debug!( + version = %display_cuda_version(cached.info.version.as_ref()), + arch = %display_cuda_arch(cached.info.arch_info.as_ref()), + version_source = cached.sources.version_str(), + arch_source = cached.sources.arch_str(), + "using disk-cached CUDA info" + ); + // We are now in sync with disk; no need to write it back. + persisted.store(true, Ordering::Relaxed); + return cached; + } + } else { + tracing::trace!("CUDA info disk cache disabled"); + } + + tracing::trace!("detecting CUDA info from host"); + let detected = detect_cuda_info(); + tracing::debug!( + version = %display_cuda_version(detected.info.version.as_ref()), + arch = %display_cuda_arch(detected.info.arch_info.as_ref()), + version_source = detected.sources.version_str(), + arch_source = detected.sources.arch_str(), + "detected CUDA info from host" + ); + detected + }); + + // Persist a freshly detected result if this (or a later) call supplied a cache directory. + maybe_persist(env, detected, persisted, cache_dir); + &detected.info +} + +/// Writes the detected result to disk once, if a cache directory is available and it has not been +/// synced with disk yet. A benign race may write twice, which is safe because the write replaces +/// the file atomically. +fn maybe_persist( + env: &cache::CacheEnv, + detected: &DetectedCudaInfo, + persisted: &AtomicBool, + cache_dir: Option<&Path>, +) { + let Some(cache_dir) = cache_dir else { + return; + }; + if persisted.load(Ordering::Relaxed) { + return; + } + cache::write_with_env(env, cache_dir, &detected.info, detected.sources); + // The flag means "we have synced with disk"; set it regardless of write's best-effort outcome. + persisted.store(true, Ordering::Relaxed); } /// Returns the maximum CUDA version available on the current platform. /// /// This corresponds to the `__cuda` virtual package. The result is cached, -/// so subsequent calls are very fast. -pub fn cuda_version() -> Option { - cuda_info().version.clone() +/// so subsequent calls are very fast. See [`cuda_info`] for the `cache_dir` semantics. +pub fn cuda_version(cache_dir: Option<&Path>) -> Option { + cuda_info(cache_dir).version.clone() } /// Returns CUDA compute capability information from the current platform. @@ -109,206 +287,393 @@ pub fn cuda_version() -> Option { /// * No CUDA drivers are installed /// * No CUDA devices are detected /// * Device enumeration fails -/// * The system is using musl libc (dynamic library loading not supported) /// -/// The result is cached, so subsequent calls are very fast. -pub fn cuda_arch() -> Option { - cuda_info().arch_info.clone() +/// The result is cached, so subsequent calls are very fast. See [`cuda_info`] for the `cache_dir` +/// semantics. +pub fn cuda_arch(cache_dir: Option<&Path>) -> Option { + cuda_info(cache_dir).arch_info.clone() } /// Detects comprehensive CUDA information from the current system. /// /// This function performs unified detection of both CUDA driver version and compute -/// capability by loading the CUDA library once and querying all necessary information. +/// capability by loading NVML once and querying all necessary information. /// /// The detection process: -/// 1. Attempts to load the CUDA driver library (`libcuda`) -/// 2. Initializes the CUDA driver API -/// 3. Queries the driver version (for `__cuda` virtual package) -/// 4. Enumerates all CUDA devices and queries their compute capabilities -/// 5. Returns the minimum compute capability across all devices (for `__cuda_arch` virtual package) +/// 1. Attempts to load NVML (`libnvidia-ml`/`nvml.dll`) +/// 2. Queries the driver version (for `__cuda` virtual package); this does not require init +/// 3. Initializes NVML and enumerates all CUDA devices to query their compute capabilities +/// 4. Returns the minimum compute capability across all devices (for `__cuda_arch` virtual package) /// -/// On musl systems, only the version is detected via `nvidia-smi` since dynamic library +/// On musl systems, both are detected via the `nvidia-smi` command since dynamic library /// loading is not supported. -fn detect_cuda_info() -> CudaInfo { - if cfg!(target_env = "musl") { +fn detect_cuda_info() -> DetectedCudaInfo { + let mut detected = if cfg!(target_env = "musl") { + tracing::trace!("detecting CUDA info via nvidia-smi because musl cannot load NVML"); // Dynamically loading a library is not supported on musl so we have to fall-back to using - // the nvidia-smi command. Architecture detection requires library loading, so it's - // unavailable on musl. - CudaInfo { - version: detect_cuda_version_via_nvidia_smi(), - arch_info: None, + // the nvidia-smi command. + let version = detect_cuda_version_via_nvidia_smi(); + // Only query the compute capability when the driver version was found. On a GPU-less musl + // system the version query already failed, so running the arch query too would just spawn + // another doomed process and could produce an inconsistent `{version: None, arch: Some}`. + let arch_info = version + .as_ref() + .and_then(|_| detect_cuda_arch_via_nvidia_smi()); + DetectedCudaInfo { + sources: CudaInfoSources { + version: version.is_some().then_some(CudaDetectionMethod::NvidiaSmi), + arch: arch_info + .is_some() + .then_some(CudaDetectionMethod::NvidiaSmi), + }, + info: CudaInfo { version, arch_info }, } } else { - // Try to detect via libcuda which allows us to get both version and architecture info - detect_cuda_info_via_libcuda() + tracing::trace!("detecting CUDA info via NVML"); + // Prefer NVML because it is not affected by `CUDA_VISIBLE_DEVICES`, but fall back to the + // older probes so systems that expose libcuda (or nvidia-smi) without NVML still report + // `__cuda`. + let mut detected = detect_cuda_info_via_nvml(); + + if detected.info.version.is_none() { + tracing::debug!( + "NVML did not detect a CUDA driver version; trying libcuda/nvidia-smi fallbacks" + ); + if let Some((version, source)) = detect_cuda_version_fallbacks() { + detected.info.version = Some(version); + detected.sources.version = Some(source); + } + } + + if detected.info.version.is_some() && detected.info.arch_info.is_none() { + tracing::debug!( + "NVML did not detect CUDA compute capability; trying nvidia-smi fallback" + ); + detected.info.arch_info = detect_cuda_arch_via_nvidia_smi(); + if detected.info.arch_info.is_some() { + detected.sources.arch = Some(CudaDetectionMethod::NvidiaSmi); + } + } + + // Last resort for platforms that ship libcuda but neither NVML nor nvidia-smi (e.g. + // Jetson/Tegra). libcuda comes last because its device enumeration is affected by + // `CUDA_VISIBLE_DEVICES`. + if detected.info.version.is_some() && detected.info.arch_info.is_none() { + tracing::debug!( + "nvidia-smi did not detect CUDA compute capability; trying libcuda fallback" + ); + detected.info.arch_info = detect_cuda_arch_via_libcuda(); + if detected.info.arch_info.is_some() { + detected.sources.arch = Some(CudaDetectionMethod::Libcuda); + } + } + + detected + }; + + // Normalization for all paths: `__cuda_arch` is meaningless without `__cuda`. If no driver + // version was detected, drop any compute capability so callers can never observe + // arch-without-version. + if detected.info.version.is_none() + && (detected.info.arch_info.is_some() || detected.sources.arch.is_some()) + { + tracing::debug!( + "dropping CUDA compute capability because no CUDA driver version was detected" + ); + detected.info.arch_info = None; + detected.sources.arch = None; } + + detected } -/// Detects CUDA version and architecture information via the CUDA driver library. +/// Detects CUDA version and architecture information via the NVIDIA Management Library. /// -/// This function loads `libcuda` and uses the CUDA Driver API to query both the driver -/// version and device compute capabilities. This is more efficient than separate detection -/// because the library is loaded only once. +/// The library is loaded once and used to query both the driver version and device compute +/// capabilities. NVML is preferred over libcuda because it is not affected by +/// `CUDA_VISIBLE_DEVICES`. /// /// Returns a `CudaInfo` struct where: /// * `version` is `None` if the driver version cannot be determined /// * `arch_info` is `None` if no devices are present or device queries fail -fn detect_cuda_info_via_libcuda() -> CudaInfo { - // Try to open the CUDA library - let cuda_library = match cuda_library_paths() - .iter() - .find_map(|path| unsafe { Library::new(*path).ok() }) - { - Some(lib) => lib, - None => { - return CudaInfo { +fn detect_cuda_info_via_nvml() -> DetectedCudaInfo { + let mut library = None; + for path in nvml_library_paths() { + match unsafe { Library::new(*path) } { + Ok(loaded) => { + tracing::trace!(library_path = *path, "loaded NVML library"); + library = Some(loaded); + break; + } + Err(err) => { + tracing::trace!(library_path = *path, error = %err, "failed to load NVML library"); + } + } + } + + let Some(library) = library else { + tracing::debug!("could not load NVML library from any known path"); + return DetectedCudaInfo { + info: CudaInfo { version: None, arch_info: None, - }; - } + }, + sources: CudaInfoSources::default(), + }; }; - // Get entry points from the library - let cu_init: Symbol<'_, unsafe extern "C" fn(c_uint) -> c_ulong> = - match unsafe { cuda_library.get(b"cuInit\0") } { - Ok(init) => init, - Err(_) => { - return CudaInfo { - version: None, - arch_info: None, - }; - } - }; + // Attempt the cheap no-init query. Some drivers answer `nvmlSystemGetCudaDriverVersion` before + // `nvmlInit`, but it is not officially supported, so this result is only used as a fall back for + // when `nvmlInit` itself fails below. + let no_init_version = match cuda_version_from_nvml_library(&library) { + Ok(version) => { + tracing::trace!(%version, "detected CUDA driver version via NVML without init"); + Some(version) + } + Err(err) => { + tracing::trace!( + ?err, + "CUDA driver version query via NVML without init failed" + ); + None + } + }; - // Initialize the CUDA library - if unsafe { cu_init(0) } != 0 { - return CudaInfo { - version: None, - arch_info: None, - }; - } + // Compute capability requires enumerating devices, which needs NVML to be initialized. Since + // NVML is being initialized anyway, query the driver version while initialized too and prefer + // that officially supported result over the no-init query. + let (initialized_version, arch_info) = + detect_cuda_initialized_info_via_nvml(&library, true, true); - // Detect the driver version (can succeed even without devices) - let version = detect_cuda_version_from_library(&cuda_library); + // Prefer the version obtained from initialized NVML whenever it is available; only fall back to + // the no-init value when `nvmlInit` (and thus the initialized query) did not produce one. + let (version, version_source) = if let Some(version) = initialized_version { + (Some(version), Some(CudaDetectionMethod::NvmlInitialized)) + } else if let Some(version) = no_init_version { + (Some(version), Some(CudaDetectionMethod::NvmlNoInit)) + } else { + (None, None) + }; - // Detect architecture info (requires devices to be present) - let arch_info = detect_cuda_arch_from_library(&cuda_library); + let arch_source = arch_info + .as_ref() + .map(|_| CudaDetectionMethod::NvmlInitialized); - CudaInfo { version, arch_info } + DetectedCudaInfo { + info: CudaInfo { version, arch_info }, + sources: CudaInfoSources { + version: version_source, + arch: arch_source, + }, + } } -/// Detects CUDA driver version from an already-loaded CUDA library. +/// Queries the CUDA driver version from an already-loaded NVML library. /// -/// This function queries the CUDA driver version using `cuDriverGetVersion`. -/// The version can be detected even if no GPU devices are present on the system. -fn detect_cuda_version_from_library(cuda_library: &Library) -> Option { - let cu_driver_get_version: Symbol<'_, unsafe extern "C" fn(*mut c_int) -> c_ulong> = - unsafe { cuda_library.get(b"cuDriverGetVersion\0") }.ok()?; +/// Some drivers allow `nvmlSystemGetCudaDriverVersion` before `nvmlInit`, but others return +/// `NVML_ERROR_UNINITIALIZED`; callers may retry after initializing NVML for that error. +fn cuda_version_from_nvml_library(library: &Library) -> Result { + // Find the `nvmlSystemGetCudaDriverVersion_v2` function. If that function cannot be found, fall + // back to the `nvmlSystemGetCudaDriverVersion` function instead. + let nvml_system_get_cuda_driver_version: Symbol<'_, unsafe extern "C" fn(*mut c_int) -> c_int> = + unsafe { + library + .get(b"nvmlSystemGetCudaDriverVersion_v2\0") + .or_else(|_| library.get(b"nvmlSystemGetCudaDriverVersion\0")) + } + .map_err(|_err| NvmlCudaVersionError::MissingSymbol)?; - // Get the version from the library - let mut version_int = MaybeUninit::uninit(); - if unsafe { cu_driver_get_version(version_int.as_mut_ptr()) != 0 } { - return None; + // Zero-initialize the out-parameter so that a driver returning `NVML_SUCCESS` without actually + // writing it yields a deterministic `0`, which `parse_cuda_driver_version` rejects, rather than + // undefined behavior from reading uninitialized memory. + let mut cuda_driver_version: c_int = 0; + let result = unsafe { nvml_system_get_cuda_driver_version(&mut cuda_driver_version) }; + if result != NVML_SUCCESS { + return Err(NvmlCudaVersionError::Nvml(result)); } - let version = unsafe { version_int.assume_init() }; - // Convert the version integer to a version string - Version::from_str(&format!("{}.{}", version / 1000, (version % 1000) / 10)).ok() + parse_cuda_driver_version(cuda_driver_version).ok_or(NvmlCudaVersionError::InvalidVersion) } -/// Detects CUDA compute capability from an already-loaded CUDA library. +/// Queries information that requires initialized NVML. /// -/// This function enumerates all CUDA devices and queries their compute capabilities, -/// returning the **minimum** compute capability found across all devices along with -/// the name of the device that has this minimum capability. -/// -/// Returns `None` if: -/// * No CUDA devices are detected (`cuDeviceGetCount` returns 0) -/// * Device enumeration fails -/// * Any of the required CUDA Driver API functions cannot be loaded -fn detect_cuda_arch_from_library(cuda_library: &Library) -> Option { - // CUDA device attribute constants for querying compute capability - const CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR: c_int = 75; - const CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR: c_int = 76; - - // Get required function pointers from the library - let cu_device_get_count: Symbol<'_, unsafe extern "C" fn(*mut c_int) -> c_ulong> = - unsafe { cuda_library.get(b"cuDeviceGetCount\0") }.ok()?; +/// Returns `(version, arch_info)`. Each field is only queried when the corresponding `query_*` +/// argument is true. The initialized version query is a compatibility fallback for drivers that +/// reject `nvmlSystemGetCudaDriverVersion` before `nvmlInit`. +fn detect_cuda_initialized_info_via_nvml( + library: &Library, + query_version: bool, + query_arch: bool, +) -> (Option, Option) { + // NVML device handle (`nvmlDevice_t`) is an opaque pointer. + type NvmlDevice = *mut c_void; - let cu_device_get: Symbol<'_, unsafe extern "C" fn(*mut c_int, c_int) -> c_ulong> = - unsafe { cuda_library.get(b"cuDeviceGet\0") }.ok()?; + let Some(nvml_init): Option c_int>> = (unsafe { + library + .get(b"nvmlInit_v2\0") + .or_else(|_| library.get(b"nvmlInit\0")) + .ok() + }) else { + tracing::debug!("missing nvmlInit symbol"); + return (None, None); + }; - let cu_device_get_attribute: Symbol< - '_, - unsafe extern "C" fn(*mut c_int, c_int, c_int) -> c_ulong, - > = unsafe { cuda_library.get(b"cuDeviceGetAttribute\0") }.ok()?; + let Some(nvml_shutdown): Option c_int>> = + (unsafe { library.get(b"nvmlShutdown\0").ok() }) + else { + tracing::debug!("missing nvmlShutdown symbol"); + return (None, None); + }; - // Get the number of CUDA devices - let mut device_count = MaybeUninit::uninit(); - if unsafe { cu_device_get_count(device_count.as_mut_ptr()) } != 0 { - return None; + tracing::trace!(query_version, query_arch, "initializing NVML"); + let init_result = unsafe { nvml_init() }; + if init_result != NVML_SUCCESS { + tracing::debug!(return_code = init_result, "nvmlInit failed"); + return (None, None); } - let device_count = unsafe { device_count.assume_init() }; - // No devices found - if device_count == 0 { - return None; - } + let version = if query_version { + match cuda_version_from_nvml_library(library) { + Ok(version) => { + tracing::debug!(%version, "detected CUDA driver version via initialized NVML"); + Some(version) + } + Err(err) => { + tracing::debug!( + ?err, + "CUDA driver version query via initialized NVML failed" + ); + None + } + } + } else { + None + }; - // Iterate through all devices to find the minimum compute capability - let mut min_arch: Option = None; + // Enumerate devices to find the minimum compute capability. Wrapped in a closure so we always + // reach the `nvmlShutdown` call below regardless of the outcome. + let arch_info = query_arch + .then(|| { + tracing::trace!("querying CUDA compute capability via NVML"); + let nvml_device_get_count: Symbol<'_, unsafe extern "C" fn(*mut c_uint) -> c_int> = + match unsafe { + library + .get(b"nvmlDeviceGetCount_v2\0") + .or_else(|_| library.get(b"nvmlDeviceGetCount\0")) + } { + Ok(symbol) => symbol, + Err(err) => { + tracing::trace!(error = %err, "missing NVML device count symbol"); + return None; + } + }; - for device_idx in 0..device_count { - // Get device handle - let mut device = MaybeUninit::uninit(); - if unsafe { cu_device_get(device.as_mut_ptr(), device_idx) } != 0 { - continue; - } - let device = unsafe { device.assume_init() }; + let nvml_device_get_handle_by_index: Symbol< + '_, + unsafe extern "C" fn(c_uint, *mut NvmlDevice) -> c_int, + > = match unsafe { + library + .get(b"nvmlDeviceGetHandleByIndex_v2\0") + .or_else(|_| library.get(b"nvmlDeviceGetHandleByIndex\0")) + } { + Ok(symbol) => symbol, + Err(err) => { + tracing::trace!(error = %err, "missing NVML device handle symbol"); + return None; + } + }; - // Get compute capability major version - let mut cc_major = MaybeUninit::uninit(); - if unsafe { - cu_device_get_attribute( - cc_major.as_mut_ptr(), - CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, - device, - ) - } != 0 - { - continue; - } - let cc_major = unsafe { cc_major.assume_init() } as u32; + let nvml_device_get_cuda_compute_capability: Symbol< + '_, + unsafe extern "C" fn(NvmlDevice, *mut c_int, *mut c_int) -> c_int, + > = match unsafe { library.get(b"nvmlDeviceGetCudaComputeCapability\0") } { + Ok(symbol) => symbol, + Err(err) => { + tracing::trace!( + error = %err, + "missing NVML CUDA compute capability symbol" + ); + return None; + } + }; - // Get compute capability minor version - let mut cc_minor = MaybeUninit::uninit(); - if unsafe { - cu_device_get_attribute( - cc_minor.as_mut_ptr(), - CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, - device, - ) - } != 0 - { - continue; - } - let cc_minor = unsafe { cc_minor.assume_init() } as u32; + let mut device_count: c_uint = 0; + let device_count_result = unsafe { nvml_device_get_count(&mut device_count) }; + if device_count_result != NVML_SUCCESS { + tracing::trace!( + return_code = device_count_result, + "nvmlDeviceGetCount failed" + ); + return None; + } + tracing::trace!(device_count, "enumerating CUDA devices via NVML"); - // Check if this is the minimum compute capability so far - let is_new_minimum = min_arch.as_ref().is_none_or(|min| { - cc_major < min.major || (cc_major == min.major && cc_minor < min.minor) - }); + let mut min_arch: Option = None; + for device_idx in 0..device_count { + let mut device: NvmlDevice = ptr::null_mut(); + let handle_result = + unsafe { nvml_device_get_handle_by_index(device_idx, &mut device) }; + if handle_result != NVML_SUCCESS { + tracing::trace!( + device_idx, + return_code = handle_result, + "failed to get NVML device handle" + ); + continue; + } - if is_new_minimum { - min_arch = Some(CudaArchInfo { - major: cc_major, - minor: cc_minor, - }); - } + let mut cc_major: c_int = 0; + let mut cc_minor: c_int = 0; + let compute_capability_result = unsafe { + nvml_device_get_cuda_compute_capability(device, &mut cc_major, &mut cc_minor) + }; + if compute_capability_result != NVML_SUCCESS { + tracing::trace!( + device_idx, + return_code = compute_capability_result, + "failed to get CUDA compute capability via NVML" + ); + continue; + } + let cc_major = cc_major as u32; + let cc_minor = cc_minor as u32; + tracing::trace!( + device_idx, + major = cc_major, + minor = cc_minor, + "detected CUDA compute capability via NVML" + ); + + let is_new_minimum = min_arch.as_ref().is_none_or(|min| { + cc_major < min.major || (cc_major == min.major && cc_minor < min.minor) + }); + if is_new_minimum { + min_arch = Some(CudaArchInfo { + major: cc_major, + minor: cc_minor, + }); + } + } + if let Some(arch) = min_arch.as_ref() { + tracing::debug!( + major = arch.major, + minor = arch.minor, + "selected minimum CUDA compute capability" + ); + } else { + tracing::debug!("no CUDA compute capability detected via NVML"); + } + min_arch + }) + .flatten(); + + // Whatever happens, after initializing NVML we have to call `nvmlShutdown`. + let shutdown_result = unsafe { nvml_shutdown() }; + if shutdown_result != NVML_SUCCESS { + tracing::debug!(return_code = shutdown_result, "nvmlShutdown failed"); } - min_arch + (version, arch_info) } /// Attempts to detect the version of CUDA present in the current operating system by employing the @@ -319,8 +684,28 @@ pub fn detect_cuda_version() -> Option { // the nvidia-smi command. detect_cuda_version_via_nvidia_smi() } else { - detect_cuda_version_via_nvml() + detect_cuda_version_via_nvml().or_else(|| { + tracing::debug!( + "NVML did not detect a CUDA driver version; trying libcuda/nvidia-smi fallbacks" + ); + detect_cuda_version_fallbacks().map(|(version, _source)| version) + }) + } +} + +fn detect_cuda_version_fallbacks() -> Option<(Version, CudaDetectionMethod)> { + if cfg!(target_env = "musl") { + return detect_cuda_version_via_nvidia_smi() + .map(|version| (version, CudaDetectionMethod::NvidiaSmi)); + } + + let version = detect_cuda_version_via_libcuda(); + if let Some(version) = version { + return Some((version, CudaDetectionMethod::Libcuda)); } + + tracing::debug!("libcuda did not detect a CUDA driver version; trying nvidia-smi fallback"); + detect_cuda_version_via_nvidia_smi().map(|version| (version, CudaDetectionMethod::NvidiaSmi)) } /// Attempts to detect the version of CUDA present in the current operating system by loading the @@ -331,58 +716,41 @@ pub fn detect_cuda_version() -> Option { /// Although the required methods in the runtime are not implemented on much older machines it is /// considered old enough to be usable for our use case. Since Conda doesn't provide old versions of /// the CUDA SDK anyway this is considered a non-issue. +/// +/// Some drivers can answer `nvmlSystemGetCudaDriverVersion` without `nvmlInit`, avoiding the +/// expensive driver handshake / GPU attach that makes `nvmlInit` slow on Windows. If that no-init +/// query fails, this falls back to querying while NVML is initialized. pub fn detect_cuda_version_via_nvml() -> Option { // Try to open the library - let library = nvml_library_paths() - .iter() - .find_map(|path| unsafe { libloading::Library::new(*path).ok() })?; - - // Get the initialization function. We first try to get `nvmlInit_v2` but if we can't find that - // we use the `nvmlInit` function. - let nvml_init: Symbol<'_, unsafe extern "C" fn() -> c_int> = unsafe { - library - .get(b"nvmlInit_v2\0") - .or_else(|_| library.get(b"nvmlInit\0")) - } - .ok()?; - - // Find the shutdown function - let nvml_shutdown: Symbol<'_, unsafe extern "C" fn() -> c_int> = - unsafe { library.get(b"nvmlShutdown\0") }.ok()?; - - // Find the `nvmlSystemGetCudaDriverVersion_v2` function. If that function cannot be found, fall - // back to the `nvmlSystemGetCudaDriverVersion` function instead. - let nvml_system_get_cuda_driver_version: Symbol<'_, unsafe extern "C" fn(*mut c_int) -> c_int> = - unsafe { - library - .get(b"nvmlSystemGetCudaDriverVersion_v2\0") - .or_else(|_| library.get(b"nvmlSystemGetCudaDriverVersion\0")) + let mut library = None; + for path in nvml_library_paths() { + match unsafe { libloading::Library::new(*path) } { + Ok(loaded) => { + tracing::trace!(library_path = *path, "loaded NVML library"); + library = Some(loaded); + break; + } + Err(err) => { + tracing::trace!(library_path = *path, error = %err, "failed to load NVML library"); + } } - .ok()?; - - // Call the initialization function - if unsafe { nvml_init() } != 0 { - return None; } + let library = library?; - // Get the version - let mut cuda_driver_version = MaybeUninit::uninit(); - let result = unsafe { nvml_system_get_cuda_driver_version(cuda_driver_version.as_mut_ptr()) }; - - // Call the shutdown function (don't care about the result of the function). Whatever happens, - // after calling `nvmlInit` we have to call `nvmlShutdown`. - let _ = unsafe { nvml_shutdown() }; - - // If the call failed we dont have a version - if result != 0 { - return None; + match cuda_version_from_nvml_library(&library) { + Ok(version) => { + tracing::trace!(%version, "detected CUDA driver version via NVML without init"); + Some(version) + } + Err(err) if err.should_retry_after_init() => { + tracing::debug!(?err, "retrying CUDA driver version query after nvmlInit"); + detect_cuda_initialized_info_via_nvml(&library, true, false).0 + } + Err(err) => { + tracing::debug!(?err, "CUDA driver version query via NVML failed"); + None + } } - - // We can assume the value is initialized by the `nvmlSystemGetCudaDriverVersion` function. - let version = unsafe { cuda_driver_version.assume_init() }; - - // Convert the version integer to a version string - Version::from_str(&format!("{}.{}", version / 1000, (version % 1000) / 10)).ok() } /// Returns platform specific set of search paths for the CUDA library. @@ -430,30 +798,225 @@ fn nvml_library_paths() -> &'static [&'static str] { /// have this limitation. pub fn detect_cuda_version_via_libcuda() -> Option { // Try to open the library - let cuda_library = cuda_library_paths() - .iter() - .find_map(|path| unsafe { libloading::Library::new(*path).ok() })?; + let mut cuda_library = None; + for path in cuda_library_paths() { + match unsafe { libloading::Library::new(*path) } { + Ok(loaded) => { + tracing::trace!(library_path = *path, "loaded CUDA driver library"); + cuda_library = Some(loaded); + break; + } + Err(err) => { + tracing::trace!( + library_path = *path, + error = %err, + "failed to load CUDA driver library" + ); + } + } + } + let cuda_library = cuda_library?; - // Get entry points from the library - let cu_init: Symbol<'_, unsafe extern "C" fn(c_uint) -> c_ulong> = - unsafe { cuda_library.get(b"cuInit\0") }.ok()?; - let cu_driver_get_version: Symbol<'_, unsafe extern "C" fn(*mut c_int) -> c_ulong> = - unsafe { cuda_library.get(b"cuDriverGetVersion\0") }.ok()?; + // Get entry points from the library. `CUresult` is a 32-bit enum, so these are declared to + // return `c_int` (matching the NVML declarations); on some ABIs the upper bits of a wider + // return register are unspecified. + let cu_init: Symbol<'_, unsafe extern "C" fn(c_uint) -> c_int> = + match unsafe { cuda_library.get(b"cuInit\0") } { + Ok(symbol) => symbol, + Err(err) => { + tracing::debug!(error = %err, "missing cuInit symbol"); + return None; + } + }; + let cu_driver_get_version: Symbol<'_, unsafe extern "C" fn(*mut c_int) -> c_int> = + match unsafe { cuda_library.get(b"cuDriverGetVersion\0") } { + Ok(symbol) => symbol, + Err(err) => { + tracing::debug!(error = %err, "missing cuDriverGetVersion symbol"); + return None; + } + }; // Initialize the CUDA library - if unsafe { cu_init(0) } != 0 { + let init_result = unsafe { cu_init(0) }; + if init_result != 0 { + tracing::debug!(return_code = init_result, "cuInit failed"); return None; } - // Get the version from the library - let mut version_int = MaybeUninit::uninit(); - if unsafe { cu_driver_get_version(version_int.as_mut_ptr()) != 0 } { + // Get the version from the library. The out-parameter is zero-initialized so that a driver + // returning success without writing it yields a deterministic `0`, which is rejected below. + let mut version_int: c_int = 0; + let version_result = unsafe { cu_driver_get_version(&mut version_int) }; + if version_result != 0 { + tracing::debug!(return_code = version_result, "cuDriverGetVersion failed"); return None; } - let version = unsafe { version_int.assume_init() }; // Convert the version integer to a version string - Version::from_str(&format!("{}.{}", version / 1000, (version % 1000) / 10)).ok() + let version = parse_cuda_driver_version(version_int); + if let Some(version) = &version { + tracing::trace!(%version, "detected CUDA driver version via libcuda"); + } else { + tracing::trace!("failed to parse CUDA driver version reported by libcuda"); + } + version +} + +/// Attempts to detect the CUDA compute capability by loading the CUDA driver library and +/// enumerating all devices, returning the **minimum** compute capability across all devices. +/// +/// This is the fallback for platforms that ship libcuda but neither NVML nor nvidia-smi (e.g. +/// Jetson/Tegra). Device enumeration through libcuda is affected by `CUDA_VISIBLE_DEVICES`, so +/// the NVML and nvidia-smi probes are preferred. +fn detect_cuda_arch_via_libcuda() -> Option { + // CUDA device attribute constants for querying compute capability + const CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR: c_int = 75; + const CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR: c_int = 76; + + // Try to open the library + let mut cuda_library = None; + for path in cuda_library_paths() { + match unsafe { libloading::Library::new(*path) } { + Ok(loaded) => { + tracing::trace!(library_path = *path, "loaded CUDA driver library"); + cuda_library = Some(loaded); + break; + } + Err(err) => { + tracing::trace!( + library_path = *path, + error = %err, + "failed to load CUDA driver library" + ); + } + } + } + let cuda_library = cuda_library?; + + // Get entry points from the library. `CUresult` is a 32-bit enum, so these are declared to + // return `c_int`. + let cu_init: Symbol<'_, unsafe extern "C" fn(c_uint) -> c_int> = + match unsafe { cuda_library.get(b"cuInit\0") } { + Ok(symbol) => symbol, + Err(err) => { + tracing::debug!(error = %err, "missing cuInit symbol"); + return None; + } + }; + let cu_device_get_count: Symbol<'_, unsafe extern "C" fn(*mut c_int) -> c_int> = + match unsafe { cuda_library.get(b"cuDeviceGetCount\0") } { + Ok(symbol) => symbol, + Err(err) => { + tracing::debug!(error = %err, "missing cuDeviceGetCount symbol"); + return None; + } + }; + let cu_device_get: Symbol<'_, unsafe extern "C" fn(*mut c_int, c_int) -> c_int> = + match unsafe { cuda_library.get(b"cuDeviceGet\0") } { + Ok(symbol) => symbol, + Err(err) => { + tracing::debug!(error = %err, "missing cuDeviceGet symbol"); + return None; + } + }; + let cu_device_get_attribute: Symbol< + '_, + unsafe extern "C" fn(*mut c_int, c_int, c_int) -> c_int, + > = match unsafe { cuda_library.get(b"cuDeviceGetAttribute\0") } { + Ok(symbol) => symbol, + Err(err) => { + tracing::debug!(error = %err, "missing cuDeviceGetAttribute symbol"); + return None; + } + }; + + // Initialize the CUDA library + let init_result = unsafe { cu_init(0) }; + if init_result != 0 { + tracing::debug!(return_code = init_result, "cuInit failed"); + return None; + } + + // Get the number of CUDA devices + let mut device_count: c_int = 0; + let device_count_result = unsafe { cu_device_get_count(&mut device_count) }; + if device_count_result != 0 { + tracing::trace!(return_code = device_count_result, "cuDeviceGetCount failed"); + return None; + } + tracing::trace!(device_count, "enumerating CUDA devices via libcuda"); + + // Iterate through all devices to find the minimum compute capability + let mut min_arch: Option = None; + for device_idx in 0..device_count { + let mut device: c_int = 0; + let device_result = unsafe { cu_device_get(&mut device, device_idx) }; + if device_result != 0 { + tracing::trace!( + device_idx, + return_code = device_result, + "failed to get CUDA device handle" + ); + continue; + } + + let mut cc_major: c_int = 0; + let mut cc_minor: c_int = 0; + let major_result = unsafe { + cu_device_get_attribute( + &mut cc_major, + CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, + device, + ) + }; + let minor_result = unsafe { + cu_device_get_attribute( + &mut cc_minor, + CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, + device, + ) + }; + if major_result != 0 || minor_result != 0 { + tracing::trace!( + device_idx, + major_return_code = major_result, + minor_return_code = minor_result, + "failed to get CUDA compute capability via libcuda" + ); + continue; + } + let cc_major = cc_major as u32; + let cc_minor = cc_minor as u32; + tracing::trace!( + device_idx, + major = cc_major, + minor = cc_minor, + "detected CUDA compute capability via libcuda" + ); + + let is_new_minimum = min_arch.as_ref().is_none_or(|min| { + cc_major < min.major || (cc_major == min.major && cc_minor < min.minor) + }); + if is_new_minimum { + min_arch = Some(CudaArchInfo { + major: cc_major, + minor: cc_minor, + }); + } + } + + if let Some(arch) = min_arch.as_ref() { + tracing::debug!( + major = arch.major, + minor = arch.minor, + "selected minimum CUDA compute capability from libcuda" + ); + } else { + tracing::debug!("no CUDA compute capability detected via libcuda"); + } + + min_arch } /// Returns platform specific set of search paths for the CUDA library. @@ -502,14 +1065,10 @@ fn cuda_library_paths() -> &'static [&'static str] { /// dynamically load a library which might not be supported on all systems. The downside is that /// executing a subprocess is generally slower and more prone to errors. fn detect_cuda_version_via_nvidia_smi() -> Option { - static CUDA_VERSION_RE: once_cell::sync::Lazy = - once_cell::sync::Lazy::new(|| { - regex::Regex::new("(.*)<\\/cuda_version>").unwrap() - }); - + tracing::trace!("detecting CUDA driver version via nvidia-smi"); // Invoke the "nvidia-smi" command to query the driver version that is usually installed when // Cuda drivers are installed. - let nvidia_smi_output = Command::new("nvidia-smi") + let nvidia_smi_output = match Command::new("nvidia-smi") // Display GPU or unit info .arg("--query") // Show unit, rather than GPU, attributes @@ -524,25 +1083,371 @@ fn detect_cuda_version_via_nvidia_smi() -> Option { // environment. .env_remove("CUDA_VISIBLE_DEVICES") .output() - .ok()?; + { + Ok(output) => output, + Err(err) => { + tracing::debug!(error = %err, "failed to run nvidia-smi for CUDA driver version"); + return None; + } + }; + + // nvidia-smi can exit non-zero in degraded-but-parseable states (e.g. one GPU lost while others + // are healthy) where the XML still contains a usable ``. Log the failure but still + // attempt to parse stdout instead of bailing out on the exit status. + if !nvidia_smi_output.status.success() { + tracing::debug!( + status = %nvidia_smi_output.status, + stderr = %String::from_utf8_lossy(&nvidia_smi_output.stderr), + "nvidia-smi CUDA driver version query exited non-zero; attempting to parse output anyway" + ); + } // Convert the output to Utf8. The conversion is lossy so it might contain some illegal // characters. If that is the case we simply assume the version in the file also wont make sense // during parsing. let output = String::from_utf8_lossy(&nvidia_smi_output.stdout); + parse_nvidia_smi_cuda_version(&output) +} + +/// Extracts the CUDA driver version from the XML output produced by `nvidia-smi --query -u -x`. +/// +/// Returns `None` if the `` element is missing or cannot be parsed as a [`Version`]. +fn parse_nvidia_smi_cuda_version(output: &str) -> Option { + static CUDA_VERSION_RE: once_cell::sync::Lazy = + once_cell::sync::Lazy::new(|| { + regex::Regex::new("(.*)<\\/cuda_version>").unwrap() + }); // Extract the version from the XML - let version_match = CUDA_VERSION_RE.captures(&output)?; - let version_str = version_match.get(1)?.as_str(); + let Some(version_match) = CUDA_VERSION_RE.captures(output) else { + tracing::trace!("nvidia-smi output did not contain a CUDA driver version"); + return None; + }; + let Some(version_match) = version_match.get(1) else { + tracing::trace!("nvidia-smi CUDA driver version match was empty"); + return None; + }; + let version_str = version_match.as_str(); // Parse and return - Version::from_str(version_str).ok() + match Version::from_str(version_str) { + Ok(version) => { + tracing::trace!(%version, "detected CUDA driver version via nvidia-smi"); + Some(version) + } + Err(err) => { + tracing::trace!(version = version_str, error = %err, "failed to parse nvidia-smi CUDA driver version"); + None + } + } +} + +/// Attempts to detect the CUDA compute capability by executing the "nvidia-smi" command and +/// querying the `compute_cap` field of every GPU, returning the **minimum** across all devices. +/// +/// Like [`detect_cuda_version_via_nvidia_smi`] this does not dynamically load a library and thus +/// also works on musl systems. The `compute_cap` query field requires a reasonably modern driver +/// (roughly R510+); on older drivers the command fails and `None` is returned. +fn detect_cuda_arch_via_nvidia_smi() -> Option { + tracing::trace!("detecting CUDA compute capability via nvidia-smi"); + let nvidia_smi_output = match Command::new("nvidia-smi") + // Query the compute capability of every GPU as plain CSV, one line per GPU. + .arg("--query-gpu=compute_cap") + .arg("--format=csv,noheader") + // See `detect_cuda_version_via_nvidia_smi` for why this variable is removed. + .env_remove("CUDA_VISIBLE_DEVICES") + .output() + { + Ok(output) => output, + Err(err) => { + tracing::debug!(error = %err, "failed to run nvidia-smi for CUDA compute capability"); + return None; + } + }; + + // On drivers that do not support the `compute_cap` field the command exits with an error, but it + // can also exit non-zero while still reporting some healthy GPUs on stdout. Log the failure but + // still attempt to parse whatever was produced instead of bailing out on the exit status. + if !nvidia_smi_output.status.success() { + tracing::debug!( + status = %nvidia_smi_output.status, + stderr = %String::from_utf8_lossy(&nvidia_smi_output.stderr), + "nvidia-smi CUDA compute capability query exited non-zero; attempting to parse output anyway" + ); + } + + let output = String::from_utf8_lossy(&nvidia_smi_output.stdout); + parse_nvidia_smi_compute_capabilities(&output) +} + +/// Parses the CSV output of `nvidia-smi --query-gpu=compute_cap --format=csv,noheader` and returns +/// the **minimum** compute capability across all parseable GPU lines. +/// +/// Each line is expected to be a `major.minor` value. Lines that are not in that format (such as +/// `[N/A]` reported for a GPU whose capability is unknown, or other junk) are skipped while still +/// using the valid ones. Returns `None` if no line could be parsed. +fn parse_nvidia_smi_compute_capabilities(output: &str) -> Option { + // Find the minimum compute capability across all devices + let mut min_arch: Option = None; + for (device_idx, line) in output.lines().enumerate() { + let line = line.trim(); + let Some((major, minor)) = line.split_once('.') else { + tracing::trace!( + device_idx, + line, + "ignoring invalid nvidia-smi compute capability line" + ); + continue; + }; + let (Ok(major), Ok(minor)) = (major.parse::(), minor.parse::()) else { + tracing::trace!( + device_idx, + line, + "ignoring unparsable nvidia-smi compute capability line" + ); + continue; + }; + tracing::trace!( + device_idx, + major, + minor, + "detected CUDA compute capability via nvidia-smi" + ); + + let is_new_minimum = min_arch + .as_ref() + .is_none_or(|min| major < min.major || (major == min.major && minor < min.minor)); + if is_new_minimum { + min_arch = Some(CudaArchInfo { major, minor }); + } + } + + if let Some(arch) = min_arch.as_ref() { + tracing::debug!( + major = arch.major, + minor = arch.minor, + "selected minimum CUDA compute capability from nvidia-smi" + ); + } else { + tracing::debug!("no CUDA compute capability detected via nvidia-smi"); + } + + min_arch } #[cfg(test)] mod test { use super::*; + /// Times loading the NVML library only, as the first NVML use in this process. + /// + /// Compared against [`bench_cold_version`] this separates the cost of mapping the library from + /// the cost of the version query itself (which loads the CUDA driver library internally). + /// + /// Run on its own, see [`bench_cold_version`]. + #[test] + #[ignore = "benchmark, run manually and in isolation"] + fn bench_cold_library_load() { + let start = std::time::Instant::now(); + let loaded = nvml_library_paths() + .iter() + .find_map(|path| unsafe { Library::new(*path).ok() }) + .is_some(); + println!( + "cold NVML library load: {:?} -> loaded={loaded}", + start.elapsed() + ); + } + + /// Times a single cold `__cuda` detection, as the first NVML call in this process. + /// + /// The dominant factor is how long ago anything last touched the GPU, not the code path: the + /// driver powers down when idle, and the first query after that re-initializes it. Measured on + /// Windows with an NVIDIA GPU this is around 1.5s for an idle driver against roughly 460ms for + /// one that was used seconds earlier, so compare runs in the same driver state. + /// + /// Run on its own so nothing else has warmed up the driver: + /// + /// ```text + /// cargo test -p rattler_virtual_packages --release -- --ignored --nocapture --exact \ + /// cuda::test::bench_cold_version + /// ``` + #[test] + #[ignore = "benchmark, run manually and in isolation"] + fn bench_cold_version() { + let start = std::time::Instant::now(); + let version = detect_cuda_version_via_nvml(); + println!( + "cold __cuda (no init): {:?} -> {version:?}", + start.elapsed() + ); + } + + /// Times a single cold `__cuda_arch` detection, as the first NVML call in this process. + /// + /// Run on its own, see [`bench_cold_version`]. + #[test] + #[ignore = "benchmark, run manually and in isolation"] + fn bench_cold_arch() { + let start = std::time::Instant::now(); + let info = detect_cuda_info(); + println!( + "cold __cuda + __cuda_arch: {:?} -> {:?}", + start.elapsed(), + info.info.arch_info + ); + } + + /// Times a single cold detection through the on-disk cache, as the first NVML call in this + /// process. Run twice: the first run populates the cache, the second measures a cache hit. + /// + /// Run on its own, see [`bench_cold_version`]. + #[test] + #[ignore = "benchmark, run manually and in isolation"] + fn bench_cold_cached() { + let cache_dir = std::env::temp_dir().join("rattler-cuda-bench-cache"); + std::fs::create_dir_all(&cache_dir).unwrap(); + let start = std::time::Instant::now(); + let info = cuda_info(Some(&cache_dir)); + println!( + "cold detection via cache: {:?} -> {:?} / {:?}", + start.elapsed(), + info.version, + info.arch_info + ); + println!( + "(run again to measure a cache hit; cache dir: {})", + cache_dir.display() + ); + } + + /// Times each CUDA detection path so the cost of the on-disk cache can be + /// compared against actually talking to the driver. + /// + /// Ignored by default because it is a benchmark and because the interesting + /// numbers only show up on a machine with an NVIDIA GPU. Run it with: + /// + /// ```text + /// cargo test -p rattler_virtual_packages --release -- --ignored --nocapture bench_cuda_detection + /// ``` + #[test] + #[ignore = "benchmark, run manually on a machine with an NVIDIA GPU"] + fn bench_cuda_detection() { + /// Reproduction of the detection this crate used before the driver library was queried + /// without initializing NVML, so the two can be compared side by side. + fn legacy_detect_cuda_version_via_nvml() -> Option { + let library = nvml_library_paths() + .iter() + .find_map(|path| unsafe { Library::new(*path).ok() })?; + let nvml_init: Symbol<'_, unsafe extern "C" fn() -> c_int> = unsafe { + library + .get(b"nvmlInit_v2\0") + .or_else(|_| library.get(b"nvmlInit\0")) + } + .ok()?; + let nvml_shutdown: Symbol<'_, unsafe extern "C" fn() -> c_int> = + unsafe { library.get(b"nvmlShutdown\0") }.ok()?; + let get_version: Symbol<'_, unsafe extern "C" fn(*mut c_int) -> c_int> = unsafe { + library + .get(b"nvmlSystemGetCudaDriverVersion_v2\0") + .or_else(|_| library.get(b"nvmlSystemGetCudaDriverVersion\0")) + } + .ok()?; + + if unsafe { nvml_init() } != 0 { + return None; + } + let mut raw = std::mem::MaybeUninit::uninit(); + let result = unsafe { get_version(raw.as_mut_ptr()) }; + let _ = unsafe { nvml_shutdown() }; + if result != 0 { + return None; + } + parse_cuda_driver_version(unsafe { raw.assume_init() }) + } + + // The first call is the one that matters: every process pays it once, and the on-disk + // cache exists to avoid exactly that. The warm best-of-5 is reported next to it to show + // how much of the cost is one-time (library loading, driver wake-up) rather than + // intrinsic to the query. + fn time(label: &str, mut f: impl FnMut() -> T) -> T { + let start = std::time::Instant::now(); + let mut result = f(); + let first = start.elapsed(); + + let mut best = std::time::Duration::MAX; + for _ in 0..5 { + let start = std::time::Instant::now(); + result = f(); + best = best.min(start.elapsed()); + } + println!("{label:<48} {first:>12.3?} {best:>12.3?}"); + result + } + + println!("\n{:<48} {:>12} {:>12}", "", "first (cold)", "best of 5"); + println!("--- detection paths ---"); + let legacy = time("version via NVML (legacy, with init)", || { + legacy_detect_cuda_version_via_nvml() + }); + let version = time("version via NVML (no init)", detect_cuda_version_via_nvml); + let info = time("version + arch via NVML (one init)", detect_cuda_info); + let smi_version = time("version via nvidia-smi", detect_cuda_version_via_nvidia_smi); + let smi_arch = time("arch via nvidia-smi", detect_cuda_arch_via_nvidia_smi); + println!("legacy version: {legacy:?}"); + + println!("\n--- breakdown ---"); + if let Some(path) = nvml_library_paths() + .iter() + .copied() + .find(|path| unsafe { Library::new(*path) }.is_ok()) + { + time("load NVML library only", || { + drop(unsafe { Library::new(path) }); + }); + let library = unsafe { Library::new(path) }.expect("just loaded successfully"); + time("version query (library already loaded)", || { + cuda_version_from_nvml_library(&library).ok() + }); + time("init + arch (library already loaded)", || { + detect_cuda_initialized_info_via_nvml(&library, false, true) + }); + } + if let Some(path) = cuda_library_paths() + .iter() + .copied() + .find(|path| unsafe { Library::new(*path) }.is_ok()) + { + time("load CUDA driver library only", || { + drop(unsafe { Library::new(path) }); + }); + } + + println!("\n--- cache paths ---"); + let dir = tempfile::tempdir().unwrap(); + let env = time("cache env (boot + driver + device keys)", || { + cache::CacheEnv::current() + }); + time("cache write", || { + cache::write_with_env(&env, dir.path(), &info.info, info.sources); + }); + let cache_read = time("cache read (hit)", || { + cache::read_with_env(&env, dir.path()) + }); + + println!("\n--- results ---"); + println!("version: {version:?}"); + println!("arch: {:?}", info.info.arch_info); + println!("nvidia-smi: {smi_version:?} / {smi_arch:?}"); + println!("cache read: {:?}", cache_read.map(|c| c.info)); + if version.is_none() { + println!( + "\nNOTE: no CUDA driver found, so the driver paths short-circuit and their\n\ + timings are meaningless. Run this on a machine with an NVIDIA GPU." + ); + } + } + #[test] pub fn doesnt_crash() { let version = detect_cuda_version_via_nvml(); @@ -555,9 +1460,15 @@ mod test { println!("Cuda {version:?}"); } + #[test] + pub fn doesnt_crash_nvidia_smi_arch() { + let arch = detect_cuda_arch_via_nvidia_smi(); + println!("Cuda arch {arch:?}"); + } + #[test] pub fn test_cuda_info() { - let info = cuda_info(); + let info = cuda_info(None); println!("CUDA Info: {info:?}"); if let Some(ref arch) = info.arch_info { println!(" Compute capability: {}.{}", arch.major, arch.minor); @@ -566,10 +1477,280 @@ mod test { #[test] pub fn test_cuda_arch() { - let arch = cuda_arch(); + let arch = cuda_arch(None); println!("CUDA Arch: {arch:?}"); } + /// Builds a fully specified, deterministic cache environment for the tests. The `driver` + /// string becomes a kernel-module fingerprint and the `device` string a single GPU bus id. + fn fake_env( + boot: &str, + driver: Option<&str>, + device: Option<&str>, + now: u64, + ) -> cache::CacheEnv { + cache::CacheEnv { + boot_id: Some(cache::BootId::Uuid(boot.to_owned())), + driver_fingerprint: driver.map(|version| cache::DriverFingerprint::Module { + version: version.to_owned(), + }), + device_fingerprint: device.map(|gpu| cache::DeviceFingerprint { + gpus: vec![gpu.to_owned()], + device_nodes: Vec::new(), + }), + now, + } + } + + fn full_info() -> (CudaInfo, CudaInfoSources) { + ( + CudaInfo { + version: Some(Version::from_str("12.4").unwrap()), + arch_info: Some(CudaArchInfo { major: 8, minor: 6 }), + }, + CudaInfoSources { + version: Some(CudaDetectionMethod::NvmlInitialized), + arch: Some(CudaDetectionMethod::NvidiaSmi), + }, + ) + } + + #[test] + fn test_cache_roundtrip() { + let dir = tempfile::tempdir().unwrap(); + let env = fake_env("boot-1", Some("driver-1"), Some("dev-1"), 1_000); + + // Nothing cached yet. + assert!(cache::read_with_env(&env, dir.path()).is_none()); + + // Negative results are not cached. + cache::write_with_env( + &env, + dir.path(), + &CudaInfo { + version: None, + arch_info: None, + }, + CudaInfoSources::default(), + ); + assert!(cache::read_with_env(&env, dir.path()).is_none()); + + let (info, sources) = full_info(); + cache::write_with_env(&env, dir.path(), &info, sources); + let cached = cache::read_with_env(&env, dir.path()).unwrap(); + assert_eq!(cached.info.version, info.version); + assert_eq!(cached.info.arch_info, info.arch_info); + assert_eq!(cached.sources, sources); + + // Updating the cache should atomically replace the previous file. + let updated_info = CudaInfo { + version: Some(Version::from_str("12.5").unwrap()), + arch_info: None, + }; + let updated_sources = CudaInfoSources { + version: Some(CudaDetectionMethod::Libcuda), + arch: None, + }; + cache::write_with_env(&env, dir.path(), &updated_info, updated_sources); + let cached = cache::read_with_env(&env, dir.path()).unwrap(); + assert_eq!(cached.info.version, updated_info.version); + assert_eq!(cached.info.arch_info, updated_info.arch_info); + assert_eq!(cached.sources, updated_sources); + } + + #[test] + fn test_cache_ttl_version_only_vs_full() { + let dir = tempfile::tempdir().unwrap(); + let write_env = fake_env("boot-1", Some("driver-1"), Some("dev-1"), 1_000); + + // A version-only entry (transient arch failure) is readable immediately but expires after + // ten minutes. + let version_only = CudaInfo { + version: Some(Version::from_str("12.4").unwrap()), + arch_info: None, + }; + cache::write_with_env( + &write_env, + dir.path(), + &version_only, + CudaInfoSources::default(), + ); + assert!(cache::read_with_env(&write_env, dir.path()).is_some()); + let just_before = fake_env("boot-1", Some("driver-1"), Some("dev-1"), 1_000 + 600); + assert!(cache::read_with_env(&just_before, dir.path()).is_some()); + let after_10m = fake_env("boot-1", Some("driver-1"), Some("dev-1"), 1_000 + 601); + assert!(cache::read_with_env(&after_10m, dir.path()).is_none()); + + // A full entry is still readable at that same age (it uses the 24h TTL) but expires past a + // day. + let (info, sources) = full_info(); + cache::write_with_env(&write_env, dir.path(), &info, sources); + assert!(cache::read_with_env(&after_10m, dir.path()).is_some()); + let after_24h = fake_env( + "boot-1", + Some("driver-1"), + Some("dev-1"), + 1_000 + 24 * 3600 + 1, + ); + assert!(cache::read_with_env(&after_24h, dir.path()).is_none()); + } + + #[test] + fn test_cache_rejects_future_write_time() { + let dir = tempfile::tempdir().unwrap(); + let write_env = fake_env("boot-1", Some("driver-1"), Some("dev-1"), 10_000); + let (info, sources) = full_info(); + cache::write_with_env(&write_env, dir.path(), &info, sources); + + // Written more than five minutes in the future (clock stepped backwards): rejected. + let past = fake_env("boot-1", Some("driver-1"), Some("dev-1"), 10_000 - 301); + assert!(cache::read_with_env(&past, dir.path()).is_none()); + // Within the tolerated skew: still accepted. + let slight_past = fake_env("boot-1", Some("driver-1"), Some("dev-1"), 10_000 - 299); + assert!(cache::read_with_env(&slight_past, dir.path()).is_some()); + } + + #[test] + fn test_cache_invalidated_on_device_change() { + let dir = tempfile::tempdir().unwrap(); + let env = fake_env("boot-1", Some("driver-1"), Some("dev-a"), 1_000); + let (info, sources) = full_info(); + cache::write_with_env(&env, dir.path(), &info, sources); + + // A different device fingerprint (e.g. another container / hot-plugged GPU) invalidates. + let other_device = fake_env("boot-1", Some("driver-1"), Some("dev-b"), 1_000); + assert!(cache::read_with_env(&other_device, dir.path()).is_none()); + // `Some` cached versus `None` current also invalidates. + let no_device = fake_env("boot-1", Some("driver-1"), None, 1_000); + assert!(cache::read_with_env(&no_device, dir.path()).is_none()); + // The matching fingerprint still reads. + assert!(cache::read_with_env(&env, dir.path()).is_some()); + } + + #[test] + fn test_cache_requires_driver_fingerprint() { + let dir = tempfile::tempdir().unwrap(); + let (info, sources) = full_info(); + + // Without a current driver fingerprint nothing is written. + let no_driver = fake_env("boot-1", None, Some("dev-1"), 1_000); + cache::write_with_env(&no_driver, dir.path(), &info, sources); + assert!(!dir.path().join("cuda-info-v1.json").exists()); + + // A hand-written cache file is rejected when the current fingerprint is unavailable. + std::fs::write( + dir.path().join("cuda-info-v1.json"), + r#"{"boot_id":{"uuid":"boot-1"},"driver_fingerprint":{"module":{"version":"driver-1"}},"device_fingerprint":{"gpus":["dev-1"],"device_nodes":[]},"written_at":1000,"version":"12.4","arch":[8,6]}"#, + ) + .unwrap(); + assert!(cache::read_with_env(&no_driver, dir.path()).is_none()); + // With a driver fingerprint available it reads. + let with_driver = fake_env("boot-1", Some("driver-1"), Some("dev-1"), 1_000); + assert!(cache::read_with_env(&with_driver, dir.path()).is_some()); + } + + #[test] + fn test_cache_invalidated_after_driver_change() { + let dir = tempfile::tempdir().unwrap(); + // A cache file written with a different driver installed is ignored. + std::fs::write( + dir.path().join("cuda-info-v1.json"), + r#"{"boot_id":{"uuid":"boot-1"},"driver_fingerprint":{"module":{"version":"535.0"}},"device_fingerprint":null,"written_at":1000,"version":"12.4","arch":[8,6]}"#, + ) + .unwrap(); + let stale = fake_env("boot-1", Some("550.0"), None, 1_000); + assert!(cache::read_with_env(&stale, dir.path()).is_none()); + // The original driver still reads. + let current = fake_env("boot-1", Some("535.0"), None, 1_000); + assert!(cache::read_with_env(¤t, dir.path()).is_some()); + } + + #[test] + fn test_cache_invalidated_after_reboot() { + let dir = tempfile::tempdir().unwrap(); + // A cache file from a different boot session is ignored. + std::fs::write( + dir.path().join("cuda-info-v1.json"), + r#"{"boot_id":{"uuid":"boot-A"},"driver_fingerprint":{"module":{"version":"driver-1"}},"device_fingerprint":null,"written_at":1000,"version":"12.4","arch":[8,6]}"#, + ) + .unwrap(); + let other_boot = fake_env("boot-B", Some("driver-1"), None, 1_000); + assert!(cache::read_with_env(&other_boot, dir.path()).is_none()); + // The same boot session still reads. + let same_boot = fake_env("boot-A", Some("driver-1"), None, 1_000); + assert!(cache::read_with_env(&same_boot, dir.path()).is_some()); + } + + #[test] + fn test_boot_time_tolerance() { + // The extracted numeric comparison, testable on every platform. + assert!(cache::boot_times_within_tolerance(1_000, 1_000)); + assert!(cache::boot_times_within_tolerance(1_000, 1_120)); + assert!(cache::boot_times_within_tolerance(1_120, 1_000)); + assert!(!cache::boot_times_within_tolerance(1_000, 1_121)); + + // Two derived boot times match within tolerance. + let a = cache::BootId::BootTime(1000); + let b = cache::BootId::BootTime(1050); + assert!(a.matches(&b)); + let c = cache::BootId::BootTime(2000); + assert!(!a.matches(&c)); + + // Boot counters must match exactly, and a boot counter never matches a boot time. + let count = cache::BootId::BootCount(5); + assert!(count.matches(&cache::BootId::BootCount(5))); + assert!(!count.matches(&cache::BootId::BootCount(6))); + assert!(!count.matches(&a)); + } + + #[test] + fn test_late_persistence() { + let dir = tempfile::tempdir().unwrap(); + let env = fake_env("boot-1", Some("driver-1"), Some("dev-1"), 1_000); + + // Pre-seed the state with a known detection result, as if detection had already run. + let state: OnceCell = OnceCell::new(); + let (info, sources) = full_info(); + let _ = state.set(DetectedCudaInfo { info, sources }); + let persisted = AtomicBool::new(false); + + // A call with no cache directory does not persist to disk. + cuda_info_impl(&env, &state, &persisted, None); + assert!(!persisted.load(Ordering::Relaxed)); + assert!(cache::read_with_env(&env, dir.path()).is_none()); + + // A later call with a cache directory persists the already-detected result. + cuda_info_impl(&env, &state, &persisted, Some(dir.path())); + assert!(persisted.load(Ordering::Relaxed)); + let cached = cache::read_with_env(&env, dir.path()).unwrap(); + assert_eq!( + cached.info.version, + Some(Version::from_str("12.4").unwrap()) + ); + assert_eq!( + cached.info.arch_info, + Some(CudaArchInfo { major: 8, minor: 6 }) + ); + } + + #[test] + fn test_is_nvidia_pci_device_id() { + // Real NVIDIA device ids as Windows reports them + assert!(cache::is_nvidia_pci_device_id( + "PCI\\VEN_10DE&DEV_2484&SUBSYS_147D10DE&REV_A1\\4&2D2E5D1F&0&0008" + )); + assert!(cache::is_nvidia_pci_device_id("PCI\\VEN_10DE&DEV_1EB1")); + // Match case insensitively + assert!(cache::is_nvidia_pci_device_id("pci\\ven_10de&dev_2484")); + + // Other vendors: Intel, AMD + assert!(!cache::is_nvidia_pci_device_id( + "PCI\\VEN_8086&DEV_9A49&SUBSYS_00011025&REV_01\\3&11583659&0&10" + )); + assert!(!cache::is_nvidia_pci_device_id("PCI\\VEN_1002&DEV_73FF")); + assert!(!cache::is_nvidia_pci_device_id("")); + } + #[test] fn test_is_valid_cuda_version_format() { // Valid formats @@ -596,4 +1777,85 @@ mod test { assert!(!is_valid_cuda_version_format("8-6")); assert!(!is_valid_cuda_version_format("8_6")); } + + #[test] + fn test_parse_cuda_driver_version() { + // Valid values are decoded as `major.minor`. + assert_eq!( + parse_cuda_driver_version(12_040), + Some(Version::from_str("12.4").unwrap()) + ); + assert_eq!( + parse_cuda_driver_version(11_080), + Some(Version::from_str("11.8").unwrap()) + ); + // Smallest plausible value (CUDA major 1). + assert_eq!( + parse_cuda_driver_version(1_000), + Some(Version::from_str("1.0").unwrap()) + ); + // Largest plausible value (CUDA major 99). + assert_eq!( + parse_cuda_driver_version(99_990), + Some(Version::from_str("99.99").unwrap()) + ); + + // Implausible values are rejected rather than propagating garbage. + assert_eq!(parse_cuda_driver_version(0), None); + assert_eq!(parse_cuda_driver_version(-1), None); + assert_eq!(parse_cuda_driver_version(999), None); + assert_eq!(parse_cuda_driver_version(100_000), None); + assert_eq!(parse_cuda_driver_version(c_int::MAX), None); + assert_eq!(parse_cuda_driver_version(c_int::MIN), None); + } + + #[test] + fn test_parse_nvidia_smi_cuda_version() { + // A representative fragment of the `nvidia-smi --query -u -x` XML output. + let xml = "\n 12.4\n"; + assert_eq!( + parse_nvidia_smi_cuda_version(xml), + Some(Version::from_str("12.4").unwrap()) + ); + + // Missing the tag entirely. + assert_eq!( + parse_nvidia_smi_cuda_version(""), + None + ); + + // Empty input. + assert_eq!(parse_nvidia_smi_cuda_version(""), None); + } + + #[test] + fn test_parse_nvidia_smi_compute_capabilities() { + // Multiple GPUs: the minimum capability is selected. + assert_eq!( + parse_nvidia_smi_compute_capabilities("8.6\n7.5\n9.0"), + Some(CudaArchInfo { major: 7, minor: 5 }) + ); + + // Junk and `[N/A]`-style lines are skipped while the valid ones are still used. Here the + // valid minimum line (7.5) is retained even though another GPU reports `[N/A]`. + assert_eq!( + parse_nvidia_smi_compute_capabilities("8.6\n[N/A]\n7.5\ngarbage"), + Some(CudaArchInfo { major: 7, minor: 5 }) + ); + + // A line that is unparsable as numbers (`x.y`) is skipped as well. + assert_eq!( + parse_nvidia_smi_compute_capabilities("x.y\n8.0"), + Some(CudaArchInfo { major: 8, minor: 0 }) + ); + + // All lines are junk: nothing is detected. + assert_eq!( + parse_nvidia_smi_compute_capabilities("[N/A]\ngarbage\n"), + None + ); + + // Empty input. + assert_eq!(parse_nvidia_smi_compute_capabilities(""), None); + } } diff --git a/crates/rattler_virtual_packages/src/cuda/cache.rs b/crates/rattler_virtual_packages/src/cuda/cache.rs new file mode 100644 index 0000000000..e78216f7fd --- /dev/null +++ b/crates/rattler_virtual_packages/src/cuda/cache.rs @@ -0,0 +1,603 @@ +//! On-disk cache for detected CUDA information, valid for the current boot session. +//! +//! Detecting CUDA can be slow (initializing NVML attaches to every GPU), so the result is cached +//! on disk between processes. The cache is keyed on everything that can invalidate a previous +//! detection without the code noticing: +//! +//! * the current **boot session**, because a reboot re-enumerates hardware and drivers; +//! * a **driver fingerprint**, because drivers can be updated (and NVML reloaded) without a reboot; +//! * a **device fingerprint** of the host-visible GPUs, because two containers sharing a cache +//! volume see different GPU subsets and topology can change within a session (eGPU hot-plug, +//! PCI hot-plug into VMs, suspend/resume); +//! * a **TTL** as a staleness backstop for anything the fingerprints cannot catch (and to let +//! transient arch-detection failures self-heal quickly). +//! +//! Reads and writes are best-effort: any failure simply results in a fresh detection. Delete the +//! file or use the `CONDA_OVERRIDE_CUDA*` variables to bypass it. + +use super::{CudaArchInfo, CudaDetectionMethod, CudaInfo, CudaInfoSources, DetectedCudaInfo}; +use rattler_conda_types::Version; +use serde::{Deserialize, Serialize}; +use std::{ + io::Write, + path::{Path, PathBuf}, + str::FromStr, +}; + +const CACHE_FILE_NAME: &str = "cuda-info-v1.json"; + +/// Full detections (with compute capability) are trusted for a day; the fingerprints catch most +/// changes sooner, so this is only a staleness backstop. +const FULL_TTL_SECS: u64 = 24 * 60 * 60; +/// Entries without compute capability represent a transient arch-detection failure, so they expire +/// quickly and self-heal instead of lingering for the whole boot session. +const ARCH_MISSING_TTL_SECS: u64 = 10 * 60; +/// Reject entries whose write time is further than this into the future, which means the clock +/// stepped backwards and the recorded `written_at` can no longer be trusted for the TTL. +const MAX_CLOCK_SKEW_SECS: u64 = 5 * 60; + +/// Identifies a single boot session of the machine. +/// +/// All variants are compiled on every platform so the comparison logic can be unit-tested +/// anywhere; `current` only ever produces the variants that exist on the host platform. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(super) enum BootId { + /// The kernel's per-boot UUID from `/proc/sys/kernel/random/boot_id` (Linux). + Uuid(String), + /// The prefetcher boot counter from the registry, incremented once per boot (Windows). + BootCount(u32), + /// A boot time in unix seconds derived from the uptime (Windows fallback). The derivation + /// drifts a little between processes, which `matches` absorbs with a tolerance. + BootTime(u64), +} + +impl BootId { + /// Returns the identifier of the current boot session, or `None` if it cannot be determined + /// (in which case no caching takes place). + pub(super) fn current() -> Option { + #[cfg(target_os = "linux")] + { + // The kernel generates a fresh UUID on every boot. + let id = std::fs::read_to_string("/proc/sys/kernel/random/boot_id").ok()?; + Some(Self::Uuid(id.trim().to_owned())) + } + #[cfg(target_os = "windows")] + { + // Prefer the prefetcher boot counter: it increments exactly once per boot and involves + // no clock arithmetic, so it cannot be confused by reboots or clock steps. + if let Some(count) = windows_boot_count() { + return Some(Self::BootCount(count)); + } + // Fall back to deriving the boot time from the uptime. + let uptime_secs = + unsafe { windows_sys::Win32::System::SystemInformation::GetTickCount64() } / 1000; + let now_secs = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .ok()? + .as_secs(); + Some(Self::BootTime(now_secs.checked_sub(uptime_secs)?)) + } + #[cfg(not(any(target_os = "linux", target_os = "windows")))] + { + None + } + } + + /// Returns true if both identifiers refer to the same boot session. + pub(super) fn matches(&self, other: &Self) -> bool { + match (self, other) { + // Two derived boot times can drift a few seconds between processes; treat them as the + // same session when they are within tolerance. + (Self::BootTime(a), Self::BootTime(b)) => boot_times_within_tolerance(*a, *b), + // Everything else (boot UUIDs, boot counters, or mixed kinds) must match exactly; two + // different kinds never refer to the same session. + _ => self == other, + } + } +} + +/// Returns true if two `boottime:` second values are close enough to be the same boot session. +/// +/// Extracted as a plain function (not `cfg(windows)`-gated) so the tolerance logic is compiled and +/// unit-tested on every platform. A real reboot shifts the derived boot time by at least the +/// previous uptime, which is far larger than this tolerance. +pub(super) fn boot_times_within_tolerance(a: u64, b: u64) -> bool { + /// The derived boot time drifts a little between processes. + const BOOT_TIME_TOLERANCE_SECS: u64 = 120; + a.abs_diff(b) <= BOOT_TIME_TOLERANCE_SECS +} + +/// Reads the prefetcher boot counter from the registry, incremented once per boot. +#[cfg(target_os = "windows")] +fn windows_boot_count() -> Option { + use windows_sys::Win32::System::Registry::{ + HKEY_LOCAL_MACHINE, RRF_RT_REG_DWORD, RegGetValueW, + }; + + fn wide(s: &str) -> Vec { + s.encode_utf16().chain(std::iter::once(0)).collect() + } + + let subkey = wide( + "SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Memory Management\\PrefetchParameters", + ); + let value = wide("BootId"); + let mut data: u32 = 0; + let mut data_size: u32 = std::mem::size_of::() as u32; + let status = unsafe { + RegGetValueW( + HKEY_LOCAL_MACHINE, + subkey.as_ptr(), + value.as_ptr(), + RRF_RT_REG_DWORD, + std::ptr::null_mut(), + std::ptr::addr_of_mut!(data).cast::(), + &mut data_size, + ) + }; + // ERROR_SUCCESS + if status == 0 { Some(data) } else { None } +} + +/// Identifies the installed NVIDIA driver. +/// +/// Drivers can be updated without a reboot, so the boot session alone is not enough to key the +/// cache on. All variants are compiled on every platform so they can be unit-tested anywhere. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(super) enum DriverFingerprint { + /// The version of the loaded `nvidia` kernel module (Linux), which changes when the driver is + /// updated and the module is reloaded. + Module { version: String }, + /// The identity of the NVML library file on disk (Windows, WSL2), which driver updates + /// replace. + File { + path: PathBuf, + mtime_secs: u64, + len: u64, + }, +} + +/// Identifies the set of host-visible GPUs. +/// +/// This distinguishes containers that share a cache volume but see different GPU subsets, and it +/// changes when GPUs are hot-plugged or unplugged within a boot session. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub(super) struct DeviceFingerprint { + /// Sorted identifiers of the GPUs the host exposes: PCI bus ids from + /// `/proc/driver/nvidia/gpus` on Linux, Plug and Play device ids on Windows. + pub(super) gpus: Vec, + /// Sorted `/dev/nvidiaN` device nodes that are present. Always empty on Windows, which has no + /// device nodes. + pub(super) device_nodes: Vec, +} + +#[derive(Serialize, Deserialize)] +struct CacheFile { + /// The boot session during which detection ran. + boot_id: BootId, + /// Fingerprint of the installed driver when detection ran. Required: entries without a driver + /// fingerprint cannot be trusted, so they are neither written nor read. + driver_fingerprint: DriverFingerprint, + /// Fingerprint of the host-visible GPUs when detection ran, or `None` if it could not be + /// determined on this platform. + device_fingerprint: Option, + /// Unix seconds at which the entry was written, used for the TTL. + written_at: u64, + version: String, + arch: Option<(u32, u32)>, + #[serde(default)] + version_source: Option, + #[serde(default)] + arch_source: Option, +} + +/// The environment against which a cache entry is validated: everything that can invalidate a +/// previous detection, plus the current time for the TTL. +/// +/// [`read_with_env`] and [`write_with_env`] operate against an explicit `CacheEnv` so they can be +/// unit-tested deterministically on any machine; [`CacheEnv::current`] gathers the real values for +/// the call sites in `cuda.rs`. +pub(super) struct CacheEnv { + pub(super) boot_id: Option, + pub(super) driver_fingerprint: Option, + pub(super) device_fingerprint: Option, + pub(super) now: u64, +} + +impl CacheEnv { + /// Gathers the real cache environment from the current host. + pub(super) fn current() -> Self { + Self { + boot_id: BootId::current(), + driver_fingerprint: driver_fingerprint(), + device_fingerprint: device_fingerprint(), + now: now_unix_secs(), + } + } +} + +/// Returns the current time in unix seconds, or `0` if the clock is before the epoch. +fn now_unix_secs() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |d| d.as_secs()) +} + +/// Absolute libnvidia-ml paths probed by the detector, used as a driver-fingerprint fallback where +/// `/sys/module/nvidia` does not exist (notably WSL2). +/// +/// Keep in sync with the absolute entries of `nvml_library_paths()` in `cuda.rs`. +#[cfg(target_os = "linux")] +const LIBNVIDIA_ML_ABSOLUTE_PATHS: &[&str] = &[ + "/usr/lib64/nvidia/libnvidia-ml.so.1", // RHEL/Centos/Fedora + "/usr/lib64/nvidia/libnvidia-ml.so", + "/usr/lib/x86_64-linux-gnu/libnvidia-ml.so.1", // Ubuntu + "/usr/lib/x86_64-linux-gnu/libnvidia-ml.so", + "/usr/lib/wsl/lib/libnvidia-ml.so.1", // WSL + "/usr/lib/wsl/lib/libnvidia-ml.so", +]; + +/// Fingerprints a file by its path, modification time and length, or `None` if it does not exist. +#[cfg(any(target_os = "linux", target_os = "windows"))] +fn file_fingerprint(path: &Path) -> Option { + let metadata = std::fs::metadata(path).ok()?; + let mtime = metadata + .modified() + .ok()? + .duration_since(std::time::UNIX_EPOCH) + .ok()?; + Some(DriverFingerprint::File { + path: path.to_path_buf(), + mtime_secs: mtime.as_secs(), + len: metadata.len(), + }) +} + +/// Returns a fingerprint of the installed NVIDIA driver, or `None` if it cannot be determined. +fn driver_fingerprint() -> Option { + #[cfg(target_os = "linux")] + { + // The version of the loaded kernel module, which changes when the driver is updated and + // the module is reloaded. + if let Ok(version) = std::fs::read_to_string("/sys/module/nvidia/version") { + return Some(DriverFingerprint::Module { + version: version.trim().to_owned(), + }); + } + // WSL2 (and similar setups) has no `/sys/module/nvidia`, so fall back to fingerprinting the + // libnvidia-ml file the detector would load. + for path in LIBNVIDIA_ML_ABSOLUTE_PATHS { + if let Some(fingerprint) = file_fingerprint(Path::new(path)) { + return Some(fingerprint); + } + } + None + } + #[cfg(target_os = "windows")] + { + // Driver updates on Windows usually complete without a reboot but replace nvml.dll. The DLL + // does not always load from System32, so also consider the NVSMI install location. + let mut candidates = Vec::new(); + if let Some(windir) = std::env::var_os("WINDIR") { + candidates.push(Path::new(&windir).join("System32").join("nvml.dll")); + } + if let Some(program_files) = std::env::var_os("ProgramFiles") { + candidates.push( + Path::new(&program_files) + .join("NVIDIA Corporation") + .join("NVSMI") + .join("nvml.dll"), + ); + } + for path in candidates { + if let Some(fingerprint) = file_fingerprint(&path) { + return Some(fingerprint); + } + } + None + } + #[cfg(not(any(target_os = "linux", target_os = "windows")))] + { + None + } +} + +/// Returns true if a `/dev` entry name is an `nvidiaN` device node (all-digits suffix), which +/// excludes control nodes like `nvidiactl` and `nvidia-uvm`. +#[cfg(target_os = "linux")] +fn is_nvidia_device_node(name: &str) -> bool { + name.strip_prefix("nvidia") + .is_some_and(|rest| !rest.is_empty() && rest.bytes().all(|b| b.is_ascii_digit())) +} + +/// Returns a fingerprint of the host-visible GPUs, or `None` if it cannot be determined. +fn device_fingerprint() -> Option { + #[cfg(target_os = "linux")] + { + // PCI bus ids of the GPUs the driver exposes to us. + let gpus_dir = std::fs::read_dir("/proc/driver/nvidia/gpus").ok(); + let has_gpus_dir = gpus_dir.is_some(); + let mut gpus: Vec = gpus_dir + .map(|entries| { + entries + .filter_map(Result::ok) + .map(|entry| entry.file_name().to_string_lossy().into_owned()) + .collect() + }) + .unwrap_or_default(); + gpus.sort(); + + // The `/dev/nvidiaN` device nodes that are actually present. + let mut devices: Vec = std::fs::read_dir("/dev") + .map(|entries| { + entries + .filter_map(Result::ok) + .map(|entry| entry.file_name().to_string_lossy().into_owned()) + .filter(|name| is_nvidia_device_node(name)) + .collect() + }) + .unwrap_or_default(); + devices.sort(); + + if !has_gpus_dir && devices.is_empty() { + return None; + } + Some(DeviceFingerprint { + gpus, + device_nodes: devices, + }) + } + #[cfg(target_os = "windows")] + { + // Enumerate the NVIDIA PCI devices Windows knows about. This only reads Plug and Play metadata, so + // it does not initialize the driver or wake a powered-down GPU. + let mut gpus = windows_nvidia_device_ids()?; + gpus.sort(); + Some(DeviceFingerprint { + gpus, + // Device nodes are a Linux concept. + device_nodes: Vec::new(), + }) + } + #[cfg(not(any(target_os = "linux", target_os = "windows")))] + { + // No cheap per-container GPU enumeration is available; the TTL is the backstop here. + None + } +} + +/// Returns the Plug and Play device ids of all NVIDIA PCI devices, or `None` if they cannot be +/// enumerated. +/// +/// Uses the configuration manager rather than the driver so that hot-plugged GPUs (for instance an +/// external GPU over Thunderbolt) are noticed without paying for driver initialization. +#[cfg(target_os = "windows")] +fn windows_nvidia_device_ids() -> Option> { + use windows_sys::Win32::Devices::DeviceAndDriverInstallation::{ + CM_GETIDLIST_FILTER_ENUMERATOR, CM_GETIDLIST_FILTER_PRESENT, CM_Get_Device_ID_List_SizeW, + CM_Get_Device_ID_ListW, CR_SUCCESS, + }; + + // Only enumerate devices that are currently present on the PCI bus. + let filter: Vec = "PCI".encode_utf16().chain(std::iter::once(0)).collect(); + let flags = CM_GETIDLIST_FILTER_ENUMERATOR | CM_GETIDLIST_FILTER_PRESENT; + + let mut len: u32 = 0; + if unsafe { CM_Get_Device_ID_List_SizeW(&mut len, filter.as_ptr(), flags) } != CR_SUCCESS { + return None; + } + + let mut buffer = vec![0u16; len as usize]; + if unsafe { CM_Get_Device_ID_ListW(filter.as_ptr(), buffer.as_mut_ptr(), len, flags) } + != CR_SUCCESS + { + return None; + } + + // The buffer is a sequence of null terminated strings, terminated by an empty string. + Some( + buffer + .split(|&c| c == 0) + .filter(|segment| !segment.is_empty()) + .map(String::from_utf16_lossy) + .filter(|id| is_nvidia_pci_device_id(id)) + .collect(), + ) +} + +/// Returns true if `id` is the Plug and Play device id of an NVIDIA PCI device. +/// +/// NVIDIA's PCI vendor id is `10DE`. Windows reports device ids uppercased, but match +/// case-insensitively to be safe. +#[cfg(any(target_os = "windows", test))] +pub(super) fn is_nvidia_pci_device_id(id: &str) -> bool { + id.to_ascii_uppercase().contains("VEN_10DE") +} + +/// Reads a cached detection result, validating it against the given host environment. +pub(super) fn read_with_env(env: &CacheEnv, cache_dir: &Path) -> Option { + let path = cache_dir.join(CACHE_FILE_NAME); + let Ok(content) = std::fs::read_to_string(&path) else { + tracing::trace!("no CUDA info cache found at {}", path.display()); + return None; + }; + let Ok(cached) = serde_json::from_str::(&content) else { + tracing::debug!("ignoring invalid CUDA info cache at {}", path.display()); + return None; + }; + let Some(current_boot_id) = env.boot_id.as_ref() else { + tracing::debug!( + "ignoring CUDA info cache because the current boot id could not be determined" + ); + return None; + }; + if !cached.boot_id.matches(current_boot_id) { + tracing::info!( + cache_path = %path.display(), + cached_boot_id = ?cached.boot_id, + current_boot_id = ?current_boot_id, + "invalidating CUDA info cache from a previous boot session" + ); + return None; + } + let Some(current_driver_fingerprint) = env.driver_fingerprint.as_ref() else { + tracing::debug!( + "ignoring CUDA info cache because the current driver fingerprint could not be determined" + ); + return None; + }; + if &cached.driver_fingerprint != current_driver_fingerprint { + tracing::info!( + cache_path = %path.display(), + cached_driver_fingerprint = ?cached.driver_fingerprint, + current_driver_fingerprint = ?current_driver_fingerprint, + "invalidating CUDA info cache because the driver changed" + ); + return None; + } + if cached.device_fingerprint != env.device_fingerprint { + tracing::info!( + cache_path = %path.display(), + cached_device_fingerprint = ?cached.device_fingerprint, + current_device_fingerprint = ?env.device_fingerprint, + "invalidating CUDA info cache because the visible GPUs changed" + ); + return None; + } + // Reject entries written in the future: the clock stepped backwards and the TTL below can no + // longer be trusted. + if cached.written_at.saturating_sub(env.now) > MAX_CLOCK_SKEW_SECS { + tracing::debug!( + cache_path = %path.display(), + written_at = cached.written_at, + now = env.now, + "ignoring CUDA info cache written in the future" + ); + return None; + } + // Full detections are trusted for a day; transient arch failures self-heal within minutes. + let ttl = if cached.arch.is_some() { + FULL_TTL_SECS + } else { + ARCH_MISSING_TTL_SECS + }; + if env.now.saturating_sub(cached.written_at) > ttl { + tracing::info!( + cache_path = %path.display(), + written_at = cached.written_at, + now = env.now, + ttl, + "invalidating expired CUDA info cache" + ); + return None; + } + let version = match Version::from_str(&cached.version) { + Ok(version) => version, + Err(err) => { + tracing::debug!( + version = cached.version, + error = %err, + "ignoring CUDA info cache with invalid version" + ); + return None; + } + }; + tracing::trace!("using CUDA info cached at {}", path.display()); + Some(DetectedCudaInfo { + info: CudaInfo { + version: Some(version), + arch_info: cached + .arch + .map(|(major, minor)| CudaArchInfo { major, minor }), + }, + sources: CudaInfoSources { + version: cached.version_source, + arch: cached.arch_source, + }, + }) +} + +/// Writes a detection result to the cache, keyed on the given host environment. +pub(super) fn write_with_env( + env: &CacheEnv, + cache_dir: &Path, + info: &CudaInfo, + sources: CudaInfoSources, +) { + // Only cache when a driver was found: detection without a driver is fast anyway, and not + // caching the negative result means a freshly installed driver is picked up immediately. + let Some(version) = &info.version else { + tracing::trace!("not caching CUDA info because no CUDA driver version was detected"); + return; + }; + let Some(boot_id) = env.boot_id.clone() else { + tracing::debug!( + "not caching CUDA info because the current boot id could not be determined" + ); + return; + }; + // The driver fingerprint is required: without it the cache is keyed on the boot session alone + // (a host driver update would serve stale data on e.g. WSL2), so we simply do not cache. + let Some(driver_fingerprint) = env.driver_fingerprint.clone() else { + tracing::debug!( + "not caching CUDA info because the current driver fingerprint could not be determined" + ); + return; + }; + if let Err(err) = std::fs::create_dir_all(cache_dir) { + tracing::debug!( + cache_dir = %cache_dir.display(), + error = %err, + "failed to create CUDA info cache directory" + ); + return; + } + let cached = CacheFile { + boot_id, + driver_fingerprint, + device_fingerprint: env.device_fingerprint.clone(), + written_at: env.now, + version: version.to_string(), + arch: info.arch_info.as_ref().map(|arch| (arch.major, arch.minor)), + version_source: sources.version, + arch_source: sources.arch, + }; + let Ok(content) = serde_json::to_string(&cached) else { + tracing::debug!("failed to serialize CUDA info cache entry"); + return; + }; + // Write to a temporary file in the cache directory and persist it into place so concurrent + // readers never see a partial cache file. + let path = cache_dir.join(CACHE_FILE_NAME); + let mut tmp = match tempfile::NamedTempFile::new_in(cache_dir) { + Ok(tmp) => tmp, + Err(err) => { + tracing::debug!( + cache_dir = %cache_dir.display(), + error = %err, + "failed to create temporary CUDA info cache file" + ); + return; + } + }; + if let Err(err) = tmp.write_all(content.as_bytes()) { + tracing::debug!( + cache_path = %path.display(), + error = %err, + "failed to write temporary CUDA info cache file" + ); + return; + } + match tmp.persist(&path) { + Ok(_) => tracing::trace!("cached CUDA info at {}", path.display()), + Err(err) => { + tracing::debug!( + cache_path = %path.display(), + error = %err.error, + "failed to persist CUDA info cache" + ); + } + } +} diff --git a/crates/rattler_virtual_packages/src/lib.rs b/crates/rattler_virtual_packages/src/lib.rs index 915b28b7d7..db6844a8fe 100644 --- a/crates/rattler_virtual_packages/src/lib.rs +++ b/crates/rattler_virtual_packages/src/lib.rs @@ -44,6 +44,7 @@ use std::{ env, fmt, fmt::Display, hash::{Hash, Hasher}, + path::Path, str::FromStr, sync::Arc, }; @@ -159,6 +160,33 @@ pub trait EnvOverride: Sized { Self::detect_with_fallback(ov, Self::detect_from_host) }) } + + /// Detect the virtual package for the current system, using `cache_dir` for any on-disk + /// detection cache. + /// + /// The default implementation ignores `cache_dir` and defers to + /// [`EnvOverride::detect_from_host`]. Virtual packages with an expensive detection step (such + /// as CUDA) override this to thread the cache directory through. + fn detect_from_host_with_cache_dir( + cache_dir: Option<&Path>, + ) -> Result, DetectVirtualPackageError> { + let _ = cache_dir; + Self::detect_from_host() + } + + /// Like [`EnvOverride::detect`] but funnels the host detection through + /// [`EnvOverride::detect_from_host_with_cache_dir`] so an on-disk detection cache can be used. + fn detect_cached( + ov: Option<&Override>, + cache_dir: Option<&Path>, + ) -> Result, DetectVirtualPackageError> { + ov.map_or_else( + || Self::detect_from_host_with_cache_dir(cache_dir), + |ov| { + Self::detect_with_fallback(ov, || Self::detect_from_host_with_cache_dir(cache_dir)) + }, + ) + } } /// An enum that represents all virtual package types provided by this library. @@ -274,12 +302,32 @@ impl VirtualPackages { /// Detect the virtual packages of the current system with the given /// overrides. - pub fn detect(overrides: &VirtualPackageOverrides) -> Result { - let cuda = Cuda::detect(overrides.cuda.as_ref())?; - let mut cuda_arch = CudaArch::detect(overrides.cuda_arch.as_ref())?; + /// + /// `cache_dir` is used to cache expensive detection results (currently CUDA) on disk across + /// processes; pass `None` to disable the on-disk cache. + pub fn detect( + overrides: &VirtualPackageOverrides, + cache_dir: Option<&Path>, + ) -> Result { + tracing::trace!( + cache_dir = %cache_dir.map_or_else( + || "".to_string(), + |path| path.display().to_string() + ), + "detecting virtual packages" + ); + + let cuda = Cuda::detect_with_cache_dir(overrides.cuda.as_ref(), cache_dir)?; + tracing::trace!(?cuda, "detected CUDA virtual package"); + let mut cuda_arch = + CudaArch::detect_with_cache_dir(overrides.cuda_arch.as_ref(), cache_dir)?; + tracing::trace!(?cuda_arch, "detected CUDA architecture virtual package"); // Enforce CEP requirement: __cuda_arch must be absent when __cuda is absent if cuda.is_none() { + if cuda_arch.is_some() { + tracing::debug!("dropping __cuda_arch because __cuda was not detected"); + } cuda_arch = None; } @@ -322,8 +370,9 @@ impl VirtualPackages { pub fn detect_for_platform( platform: Platform, overrides: &VirtualPackageOverrides, + cache_dir: Option<&Path>, ) -> Result { - let virtual_packages = Self::detect(overrides)?; + let virtual_packages = Self::detect(overrides, cache_dir)?; if platform == Platform::current() { // If we're targeting the current platform, just return the detected packages Ok(virtual_packages) @@ -483,18 +532,21 @@ impl VirtualPackage { /// the versions could not be properly detected. #[deprecated( since = "1.1.0", - note = "Use `VirtualPackage::detect(&VirtualPackageOverrides::default())` instead." + note = "Use `VirtualPackage::detect(&VirtualPackageOverrides::default(), None)` instead." )] pub fn current() -> Result, DetectVirtualPackageError> { - Self::detect(&VirtualPackageOverrides::default()) + Self::detect(&VirtualPackageOverrides::default(), None) } /// Detect the virtual packages of the current system with the given /// overrides. + /// + /// See [`VirtualPackages::detect`] for the `cache_dir` semantics. pub fn detect( overrides: &VirtualPackageOverrides, + cache_dir: Option<&Path>, ) -> Result, DetectVirtualPackageError> { - Ok(VirtualPackages::detect(overrides)? + Ok(VirtualPackages::detect(overrides, cache_dir)? .into_virtual_packages() .collect()) } @@ -706,8 +758,19 @@ pub struct Cuda { impl Cuda { /// Returns the maximum Cuda version available on the current platform. - pub fn current() -> Option { - cuda::cuda_version().map(|version| Self { version }) + /// + /// See [`cuda::cuda_info`] for the `cache_dir` semantics. + pub fn current(cache_dir: Option<&Path>) -> Option { + cuda::cuda_version(cache_dir).map(|version| Self { version }) + } + + /// Detect the Cuda virtual package with the given override, using `cache_dir` for the on-disk + /// detection cache. + pub fn detect_with_cache_dir( + ov: Option<&Override>, + cache_dir: Option<&Path>, + ) -> Result, DetectVirtualPackageError> { + ::detect_cached(ov, cache_dir) } } @@ -724,7 +787,12 @@ impl EnvOverride for Cuda { }) } fn detect_from_host() -> Result, DetectVirtualPackageError> { - Ok(Self::current()) + Ok(Self::current(None)) + } + fn detect_from_host_with_cache_dir( + cache_dir: Option<&Path>, + ) -> Result, DetectVirtualPackageError> { + Ok(Self::current(cache_dir)) } const DEFAULT_ENV_NAME: &'static str = "CONDA_OVERRIDE_CUDA"; } @@ -776,13 +844,23 @@ impl CudaArch { /// * No CUDA drivers are installed /// * No CUDA devices are detected /// * Device enumeration fails - /// * The system is using musl libc (dynamic library loading not supported) - pub fn current() -> Option { - cuda::cuda_arch().map(|arch_info| Self { + /// + /// See [`cuda::cuda_info`] for the `cache_dir` semantics. + pub fn current(cache_dir: Option<&Path>) -> Option { + cuda::cuda_arch(cache_dir).map(|arch_info| Self { version: Version::from_str(&format!("{}.{}", arch_info.major, arch_info.minor)) .unwrap_or_else(|_| Version::major(u64::from(arch_info.major))), }) } + + /// Detect the CUDA compute capability virtual package with the given override, using + /// `cache_dir` for the on-disk detection cache. + pub fn detect_with_cache_dir( + ov: Option<&Override>, + cache_dir: Option<&Path>, + ) -> Result, DetectVirtualPackageError> { + ::detect_cached(ov, cache_dir) + } } impl EnvOverride for CudaArch { @@ -799,7 +877,13 @@ impl EnvOverride for CudaArch { } fn detect_from_host() -> Result, DetectVirtualPackageError> { - Ok(Self::current()) + Ok(Self::current(None)) + } + + fn detect_from_host_with_cache_dir( + cache_dir: Option<&Path>, + ) -> Result, DetectVirtualPackageError> { + Ok(Self::current(cache_dir)) } const DEFAULT_ENV_NAME: &'static str = "CONDA_OVERRIDE_CUDA_ARCH"; @@ -1265,7 +1349,7 @@ mod test { #[test] fn doesnt_crash() { let virtual_packages = - VirtualPackages::detect(&VirtualPackageOverrides::default()).unwrap(); + VirtualPackages::detect(&VirtualPackageOverrides::default(), None).unwrap(); println!("{virtual_packages:#?}"); } @@ -1437,7 +1521,7 @@ mod test { // Test Linux 64-bit let linux_packages = - VirtualPackages::detect_for_platform(Platform::Linux64, &overrides).unwrap(); + VirtualPackages::detect_for_platform(Platform::Linux64, &overrides, None).unwrap(); let linux_names: Vec = linux_packages .into_generic_virtual_packages() .map(|pkg| pkg.name.as_normalized().to_string()) @@ -1449,7 +1533,7 @@ mod test { // Test macOS ARM64 let osx_packages = - VirtualPackages::detect_for_platform(Platform::OsxArm64, &overrides).unwrap(); + VirtualPackages::detect_for_platform(Platform::OsxArm64, &overrides, None).unwrap(); let osx_names: Vec = osx_packages .into_generic_virtual_packages() .map(|pkg| pkg.name.as_normalized().to_string()) @@ -1460,7 +1544,7 @@ mod test { // Test Windows 64-bit let win_packages = - VirtualPackages::detect_for_platform(Platform::Win64, &overrides).unwrap(); + VirtualPackages::detect_for_platform(Platform::Win64, &overrides, None).unwrap(); let win_names: Vec = win_packages .into_generic_virtual_packages() .map(|pkg| pkg.name.as_normalized().to_string()) @@ -1481,7 +1565,7 @@ mod test { if !current.is_linux() { let packages = - VirtualPackages::detect_for_platform(Platform::Linux64, &overrides).unwrap(); + VirtualPackages::detect_for_platform(Platform::Linux64, &overrides, None).unwrap(); assert_eq!( packages.linux.expect("__linux should be present").version, defaults::default_linux_version() @@ -1493,7 +1577,7 @@ mod test { if !current.is_osx() { let packages = - VirtualPackages::detect_for_platform(Platform::OsxArm64, &overrides).unwrap(); + VirtualPackages::detect_for_platform(Platform::OsxArm64, &overrides, None).unwrap(); assert_eq!( packages.osx.expect("__osx should be present").version, defaults::default_mac_os_version(Platform::OsxArm64).unwrap() @@ -1502,7 +1586,7 @@ mod test { if !current.is_windows() { let packages = - VirtualPackages::detect_for_platform(Platform::Win64, &overrides).unwrap(); + VirtualPackages::detect_for_platform(Platform::Win64, &overrides, None).unwrap(); assert_eq!( packages.win.expect("__win should be present").version, Some(defaults::default_windows_version()) @@ -1517,6 +1601,7 @@ mod test { let ios_packages = VirtualPackages::detect_for_platform( Platform::IosArm64, &VirtualPackageOverrides::default(), + None, ) .unwrap(); let ios_names: Vec = ios_packages @@ -1533,7 +1618,8 @@ mod test { ..Default::default() }; let ios_packages = - VirtualPackages::detect_for_platform(Platform::IosSimulatorArm64, &overrides).unwrap(); + VirtualPackages::detect_for_platform(Platform::IosSimulatorArm64, &overrides, None) + .unwrap(); let ios = ios_packages .into_generic_virtual_packages() .find(|pkg| pkg.name.as_normalized() == "__ios") @@ -1545,6 +1631,7 @@ mod test { let android_packages = VirtualPackages::detect_for_platform( Platform::AndroidAarch64, &VirtualPackageOverrides::default(), + None, ) .unwrap(); let android_names: Vec = android_packages @@ -1561,7 +1648,8 @@ mod test { ..Default::default() }; let android_packages = - VirtualPackages::detect_for_platform(Platform::AndroidArmV7a, &overrides).unwrap(); + VirtualPackages::detect_for_platform(Platform::AndroidArmV7a, &overrides, None) + .unwrap(); let android = android_packages .into_generic_virtual_packages() .find(|pkg| pkg.name.as_normalized() == "__android") @@ -1630,7 +1718,7 @@ mod test { // Case 1: Both not present - cuda_arch should be None let overrides = VirtualPackageOverrides::default(); - let packages = VirtualPackages::detect(&overrides).unwrap(); + let packages = VirtualPackages::detect(&overrides, None).unwrap(); // If cuda is None, cuda_arch must also be None if packages.cuda.is_none() { assert!( @@ -1646,7 +1734,7 @@ mod test { cuda_arch: Some(cuda_arch_override), ..Default::default() }; - let packages = VirtualPackages::detect(&overrides).unwrap(); + let packages = VirtualPackages::detect(&overrides, None).unwrap(); if packages.cuda.is_none() { assert!( packages.cuda_arch.is_none(), @@ -1662,7 +1750,7 @@ mod test { cuda_arch: Some(cuda_arch_override), ..Default::default() }; - let packages = VirtualPackages::detect(&overrides).unwrap(); + let packages = VirtualPackages::detect(&overrides, None).unwrap(); assert!( packages.cuda.is_some(), "cuda should be present with override" @@ -1682,7 +1770,7 @@ mod test { cuda_arch: Some(cuda_arch_override), ..Default::default() }; - let packages = VirtualPackages::detect(&overrides).unwrap(); + let packages = VirtualPackages::detect(&overrides, None).unwrap(); assert!( packages.cuda.is_none(), "cuda should be None with empty string override" diff --git a/py-rattler/Cargo.lock b/py-rattler/Cargo.lock index e75284df14..d8e29e68f9 100644 --- a/py-rattler/Cargo.lock +++ b/py-rattler/Cargo.lock @@ -4446,8 +4446,11 @@ dependencies = [ "rattler_conda_types", "regex", "serde", + "serde_json", + "tempfile", "thiserror 2.0.18", "tracing", + "windows-sys 0.61.2", "winver", ] diff --git a/py-rattler/rattler/virtual_package/virtual_package.py b/py-rattler/rattler/virtual_package/virtual_package.py index 1abbdc14b5..ccaafe497b 100644 --- a/py-rattler/rattler/virtual_package/virtual_package.py +++ b/py-rattler/rattler/virtual_package/virtual_package.py @@ -1,5 +1,6 @@ from __future__ import annotations -from typing import List +import os +from typing import List, Optional, Union import warnings from rattler.rattler import PyVirtualPackage, PyOverride, PyVirtualPackageOverrides @@ -231,11 +232,20 @@ def current() -> List[VirtualPackage]: return VirtualPackage.detect() @staticmethod - def detect(overrides: VirtualPackageOverrides = VirtualPackageOverrides()) -> List[VirtualPackage]: + def detect( + overrides: VirtualPackageOverrides = VirtualPackageOverrides(), + cache_dir: Optional[Union[str, os.PathLike[str]]] = None, + ) -> List[VirtualPackage]: """ Returns virtual packages detected for the current system with the given overrides. + + If `cache_dir` is given, expensive detection results (currently CUDA) are cached in that + directory across processes until the next reboot. """ - return [VirtualPackage._from_py_virtual_package(vp) for vp in PyVirtualPackage.detect(overrides._overrides)] + return [ + VirtualPackage._from_py_virtual_package(vp) + for vp in PyVirtualPackage.detect(overrides._overrides, cache_dir) + ] def into_generic(self) -> GenericVirtualPackage: """ diff --git a/py-rattler/src/index_json.rs b/py-rattler/src/index_json.rs index 040eafe8a9..183403b40b 100644 --- a/py-rattler/src/index_json.rs +++ b/py-rattler/src/index_json.rs @@ -261,7 +261,7 @@ impl PyIndexJson { if let Some(ts) = timestamp { self.inner.timestamp = Some(TimestampMs::from_timestamp_millis( jiff::Timestamp::from_millisecond(ts) - .map_err(|_| PyValueError::new_err("Invalid timestamp"))?, + .map_err(|err| PyValueError::new_err(format!("Invalid timestamp: {err}")))?, )); } else { self.inner.timestamp = None; diff --git a/py-rattler/src/record.rs b/py-rattler/src/record.rs index d87363149f..b148019ce1 100644 --- a/py-rattler/src/record.rs +++ b/py-rattler/src/record.rs @@ -610,7 +610,7 @@ impl PyRecord { if let Some(ts) = timestamp { self.as_package_record_mut().timestamp = Some(TimestampMs::from_timestamp_millis( jiff::Timestamp::from_millisecond(ts) - .map_err(|_| PyValueError::new_err("Invalid timestamp"))?, + .map_err(|err| PyValueError::new_err(format!("Invalid timestamp: {err}")))?, )); } else { self.as_package_record_mut().timestamp = None; diff --git a/py-rattler/src/utils.rs b/py-rattler/src/utils.rs index 4558b0c495..bcce2dd85d 100644 --- a/py-rattler/src/utils.rs +++ b/py-rattler/src/utils.rs @@ -4,10 +4,10 @@ use rattler_digest::{Md5Hash, Sha256Hash}; pub fn sha256_from_pybytes(bytes: Bound<'_, PyBytes>) -> Result { Sha256Hash::try_from(bytes.as_bytes()) - .map_err(|_| PyValueError::new_err("Expected a 32 byte SHA256 digest")) + .map_err(|_err| PyValueError::new_err("Expected a 32 byte SHA256 digest")) } pub fn md5_from_pybytes(bytes: Bound<'_, PyBytes>) -> Result { Md5Hash::try_from(bytes.as_bytes()) - .map_err(|_| PyValueError::new_err("Expected a 16 byte MD5 digest")) + .map_err(|_err| PyValueError::new_err("Expected a 16 byte MD5 digest")) } diff --git a/py-rattler/src/virtual_package.rs b/py-rattler/src/virtual_package.rs index e75f860d37..b624c8ac39 100644 --- a/py-rattler/src/virtual_package.rs +++ b/py-rattler/src/virtual_package.rs @@ -1,3 +1,5 @@ +use std::path::PathBuf; + use pyo3::{PyResult, pyclass, pymethods}; use rattler_virtual_packages::{Override, VirtualPackage, VirtualPackageOverrides}; @@ -177,14 +179,23 @@ impl PyVirtualPackage { // we just warn directly from python. #[staticmethod] pub fn current() -> PyResult> { - Self::detect(&PyVirtualPackageOverrides::none()) + Self::detect(&PyVirtualPackageOverrides::none(), None) } + /// Returns virtual packages detected for the current system with the given overrides. If + /// `cache_dir` is given, expensive detection results (currently CUDA) are cached there across + /// processes until the next reboot. #[staticmethod] - pub fn detect(overrides: &PyVirtualPackageOverrides) -> PyResult> { - Ok(VirtualPackage::detect(&overrides.clone().into()) - .map(|vp| vp.iter().map(|v| v.clone().into()).collect::>()) - .map_err(PyRattlerError::from)?) + #[pyo3(signature = (overrides, cache_dir=None))] + pub fn detect( + overrides: &PyVirtualPackageOverrides, + cache_dir: Option, + ) -> PyResult> { + Ok( + VirtualPackage::detect(&overrides.clone().into(), cache_dir.as_deref()) + .map(|vp| vp.iter().map(|v| v.clone().into()).collect::>()) + .map_err(PyRattlerError::from)?, + ) } pub fn as_generic(&self) -> PyGenericVirtualPackage {