From 24b76b7c45bd1d7e0b4ece00bdebdc94b46f8b88 Mon Sep 17 00:00:00 2001 From: Bas Zalmstra <4995967+baszalmstra@users.noreply.github.com> Date: Thu, 6 Aug 2026 08:08:30 +0000 Subject: [PATCH] perf(virtual_packages)!: cache CUDA detection on disk Detecting __cuda and __cuda_arch takes about 1.5 seconds on Windows when the GPU has been idle, which is the normal case for a command line tool. Almost all of it is the driver version query, which loads the CUDA driver library and starts the user-mode driver to read the version out of it. If the driver powered down in the meantime it has to come back up first. No detection API avoids this, so the only fix is to not do it again. Cache the result on disk, keyed on the boot session, a fingerprint of the installed driver and one of the visible GPUs, so reboots, driver updates and plugging in an eGPU all invalidate it. A TTL covers whatever the fingerprints miss. A cache hit takes about 1ms and never touches the driver, which also leaves an idle GPU asleep instead of waking it on every invocation. Callers pass the cache directory through a new cache_dir argument, and None disables it. The Python bindings take the same argument. This also fixes three problems in the detection itself: - Detection uses NVML instead of libcuda, so CUDA_VISIBLE_DEVICES no longer affects the result. Under a job scheduler or in CI, where that variable is often set, __cuda and __cuda_arch could come out wrong or disappear entirely. - __cuda_arch now works on musl, where it was always absent. - __cuda is detected even when the driver fails to initialize, since the version query no longer needs initialization. BREAKING CHANGE: VirtualPackages::detect, VirtualPackages::detect_for_platform, VirtualPackage::detect, Cuda::current, CudaArch::current and cuda::cuda_info, cuda::cuda_version and cuda::cuda_arch take an additional cache_dir argument. --- Cargo.lock | 3 + crates/rattler-bin/src/commands/create.rs | 1 + crates/rattler-bin/src/commands/exec.rs | 16 +- crates/rattler-bin/src/commands/prefix.rs | 1 + crates/rattler-bin/src/commands/solve.rs | 10 +- .../src/commands/virtual_packages.rs | 36 +- crates/rattler_virtual_packages/Cargo.toml | 7 + crates/rattler_virtual_packages/src/cuda.rs | 1726 ++++++++++++++--- .../src/cuda/cache.rs | 603 ++++++ crates/rattler_virtual_packages/src/lib.rs | 142 +- py-rattler/Cargo.lock | 3 + .../virtual_package/virtual_package.py | 16 +- py-rattler/src/index_json.rs | 2 +- py-rattler/src/record.rs | 2 +- py-rattler/src/utils.rs | 4 +- py-rattler/src/virtual_package.rs | 21 +- 16 files changed, 2307 insertions(+), 286 deletions(-) create mode 100644 crates/rattler_virtual_packages/src/cuda/cache.rs 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 {