diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 25c0e21e2..51b807abf 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -2299,6 +2299,7 @@ dependencies = [ "luks2", "nvattest", "nvml-wrapper", + "nvml-wrapper-sys", "parity-scale-codec", "ra-rpc", "ra-tls", diff --git a/dstack/Cargo.toml b/dstack/Cargo.toml index d80088f92..67d687732 100644 --- a/dstack/Cargo.toml +++ b/dstack/Cargo.toml @@ -143,6 +143,8 @@ log = "0.4.26" moka = { version = "0.12.15", default-features = false, features = ["sync"] } notify = "8.0.0" nvml-wrapper = "0.12.1" +# Raw NVML entry points the safe wrapper does not bind (fabric + P2P status). +nvml-wrapper-sys = "0.9.1" rand = "0.8.5" regorus = { version = "0.10.1", default-features = false, features = ["full-opa", "arc"] } tracing = "0.1.40" diff --git a/dstack/dstack-util/Cargo.toml b/dstack/dstack-util/Cargo.toml index ef4066831..3736fef0d 100644 --- a/dstack/dstack-util/Cargo.toml +++ b/dstack/dstack-util/Cargo.toml @@ -60,6 +60,7 @@ sodiumbox.workspace = true libc.workspace = true luks2.workspace = true nvml-wrapper.workspace = true +nvml-wrapper-sys.workspace = true scopeguard.workspace = true tempfile.workspace = true ez-hash.workspace = true diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index e15fe7890..55355eccf 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -1352,6 +1352,8 @@ async fn do_sys_setup(stage0: Stage0<'_>) -> Result<()> { /// unverified GPU either. mod gpu { use super::*; + use nvml_wrapper_sys::bindings as ffi; + use std::time::Instant; const EVENT_VERSION: u32 = 2; const POLICY_ENTRYPOINT: &str = "data.policy.nv_match"; @@ -1544,6 +1546,267 @@ mod gpu { /// Set the system-wide GPU ready state without appraisal for the explicit /// `gpu_policy.attest_gpu: false` compatibility path. + /// Fabric registration is asynchronous: the GPU probes the fabric manager + /// over NVLink inband once the driver initialises it, so a single read + /// taken right after boot legitimately catches + /// NVML_GPU_FABRIC_STATE_IN_PROGRESS. Poll instead, and log how long it + /// actually took so this bound can be tuned from deployments rather than + /// guessed at a second time. + const FABRIC_READY_TIMEOUT: Duration = Duration::from_secs(30); + const FABRIC_POLL_INTERVAL: Duration = Duration::from_millis(500); + + /// nvmlGpuFabricState_t values, from nvml.h. + const FABRIC_STATE_NOT_SUPPORTED: u8 = 0; + const FABRIC_STATE_COMPLETED: u8 = 3; + + #[derive(Clone, Copy)] + pub(super) struct FabricInfo { + state: u8, + status: u32, + clique_id: u32, + } + + impl FabricInfo { + /// A GPU counts as registered only when the probe finished *and* + /// reported success. The two are separate fields and the second is the + /// one everything else drops: NCCL logs `state` and never `status`, so + /// a GPU sitting at Completed with a failed status looks healthy in + /// every NCCL log while its P2P capability is gone. + pub(super) fn registered(&self) -> bool { + self.state == FABRIC_STATE_COMPLETED && self.status == ffi::nvmlReturn_enum_NVML_SUCCESS + } + } + + /// Read one GPU's NVLink fabric registration. + /// + /// `Ok(None)` means this platform has no fabric to register with -- a + /// single-GPU tenant, or any host without NVSwitch -- which is not a + /// fault. nvml-wrapper 0.12 exposes no safe binding for + /// `nvmlDeviceGetGpuFabricInfo`, so this reaches through the sys layer + /// using the handle and the loaded library the safe wrapper already owns. + fn fabric_info(nvml: &nvml_wrapper::Nvml, index: u32) -> Result> { + let device = nvml + .device_by_index(index) + .with_context(|| format!("failed to get NVML device {index}"))?; + let get = nvml + .lib() + .nvmlDeviceGetGpuFabricInfo + .as_ref() + .map_err(|err| anyhow!("NVML exports no nvmlDeviceGetGpuFabricInfo: {err}"))?; + let mut raw = ffi::nvmlGpuFabricInfo_t { + clusterUuid: [0; 16], + status: 0, + cliqueId: 0, + state: 0, + }; + // SAFETY: `device` is a live handle owned by `nvml`, which outlives + // this call, and `raw` is a correctly sized nvmlGpuFabricInfo_t that + // NVML only writes into. + let ret = unsafe { get(device.handle(), &mut raw) }; + if ret == ffi::nvmlReturn_enum_NVML_ERROR_NOT_SUPPORTED { + return Ok(None); + } + if ret != ffi::nvmlReturn_enum_NVML_SUCCESS { + bail!("nvmlDeviceGetGpuFabricInfo failed for GPU {index}: {ret}"); + } + if raw.state == FABRIC_STATE_NOT_SUPPORTED { + return Ok(None); + } + Ok(Some(FabricInfo { + state: raw.state, + status: raw.status, + clique_id: raw.cliqueId, + })) + } + + /// Pairs whose NVLink peer-to-peer access NVML reports as anything but OK, + /// as `(from, to, direction, status)`. + /// + /// This is the check NCCL itself makes -- its "P2P is disabled between + /// NVLINK connected GPUs" warning comes from `nvmlDeviceGetP2PStatus` -- + /// and it is the one that catches a fabric which registered cleanly but + /// still cannot move traffic. Both directions are read because NCCL reads + /// both. + fn blocked_p2p_pairs( + nvml: &nvml_wrapper::Nvml, + devices: u32, + ) -> Result> { + let get = nvml + .lib() + .nvmlDeviceGetP2PStatus + .as_ref() + .map_err(|err| anyhow!("NVML exports no nvmlDeviceGetP2PStatus: {err}"))?; + let mut blocked = Vec::new(); + for a in 0..devices { + for b in 0..devices { + if a == b { + continue; + } + let from = nvml.device_by_index(a)?; + let to = nvml.device_by_index(b)?; + for (name, index) in [ + ( + "read", + ffi::nvmlGpuP2PCapsIndex_enum_NVML_P2P_CAPS_INDEX_READ, + ), + ( + "write", + ffi::nvmlGpuP2PCapsIndex_enum_NVML_P2P_CAPS_INDEX_WRITE, + ), + ] { + let mut status: ffi::nvmlGpuP2PStatus_t = 0; + // SAFETY: both handles are live and owned by `nvml`, and + // `status` is a plain u32 NVML only writes into. + let ret = unsafe { get(from.handle(), to.handle(), index, &mut status) }; + if ret != ffi::nvmlReturn_enum_NVML_SUCCESS { + bail!("nvmlDeviceGetP2PStatus failed for {a}->{b} ({name}): {ret}"); + } + if status != ffi::nvmlGpuP2PStatus_enum_NVML_P2P_STATUS_OK { + blocked.push((a, b, name, status)); + } + } + } + } + Ok(blocked) + } + + /// Warn when the GPUs never finish joining one working NVLink fabric. + /// + /// Deliberately advisory. The failure this reports is a host-side one -- + /// a partition that was never activated, a stale Fabric Manager, a + /// tenant's GPUs landing in different cliques -- and none of it is + /// something the guest can fix or should refuse to boot over. What it + /// cannot be allowed to do is stay silent: links train in hardware through + /// ALI whatever the fabric is doing, so `nvidia-smi nvlink -s` reports + /// every link Active either way and a collective quietly falls back to + /// PCIe at a fraction of the bandwidth the tenant is paying for. + pub(super) fn warn_unless_fabric_ready(expected_devices: u32) { + if expected_devices < 2 { + return; + } + if let Err(err) = fabric_ready_inner(expected_devices) { + warn!("could not determine NVLink fabric readiness: {err:?}"); + } + } + + fn fabric_ready_inner(expected_devices: u32) -> Result<()> { + let nvml = init_nvml(expected_devices)?; + let started = Instant::now(); + // Assigned on every path that reaches the break below, so the timeout + // arm always reports the readings it actually gave up on. + let mut last; + loop { + let mut infos = Vec::with_capacity(expected_devices as usize); + for index in 0..expected_devices { + match fabric_info(&nvml, index)? { + // The driver says this platform has no fabric at all, so + // there is nothing to wait for and nothing to warn about. + None => { + info!("no NVLink fabric on this platform; skipping fabric readiness check"); + return Ok(()); + } + Some(info) => infos.push(info), + } + } + if infos.iter().all(FabricInfo::registered) { + report_ready(&nvml, &infos, started.elapsed(), expected_devices); + return Ok(()); + } + last = infos; + if started.elapsed() >= FABRIC_READY_TIMEOUT { + break; + } + std::thread::sleep(FABRIC_POLL_INTERVAL); + } + for (index, info) in last.iter().enumerate() { + warn!( + "GPU {index} has not joined the NVLink fabric: state={} status={} cliqueId={}", + info.state, info.status, info.clique_id + ); + } + warn!( + "NVLink fabric not ready after {:?}; NCCL will fall back to PCIe. \ + This is a host-side condition: check that the partition is active \ + (`fmpm -l`) and that /var/log/fabricmanager.log reports every GPU notified", + FABRIC_READY_TIMEOUT + ); + Ok(()) + } + + /// The clique every GPU shares, or `None` when they disagree. Two GPUs + /// registered into different cliques cannot reach each other over NVLink + /// however healthy each one looks on its own. + pub(super) fn shared_clique(infos: &[FabricInfo]) -> Option { + let first = infos.first()?.clique_id; + infos + .iter() + .all(|info| info.clique_id == first) + .then_some(first) + } + + /// Everything registered. Two things can still be wrong, and both are + /// invisible without asking: the GPUs may have registered into *different* + /// cliques, which means they cannot reach each other however healthy each + /// one looks, and P2P may be blocked for a reason the fabric state does + /// not express. + fn report_ready( + nvml: &nvml_wrapper::Nvml, + infos: &[FabricInfo], + elapsed: Duration, + expected_devices: u32, + ) { + match shared_clique(infos) { + Some(clique) => { + info!("NVLink fabric ready in {elapsed:?}, all GPUs in clique {clique}") + } + None => { + let ids: Vec<_> = infos.iter().map(|info| info.clique_id).collect(); + warn!( + "GPUs registered with the NVLink fabric but landed in different cliques: \ + {ids:?}. They cannot reach each other over NVLink; the host partition is \ + probably wrong" + ); + } + } + match blocked_p2p_pairs(nvml, expected_devices) { + Ok(blocked) if blocked.is_empty() => { + info!("NVLink peer-to-peer access is available between all GPU pairs"); + } + Ok(blocked) => { + for (from, to, direction, status) in &blocked { + warn!( + "NVLink peer-to-peer {direction} is blocked from GPU {from} to GPU {to} \ + (nvmlGpuP2PStatus={status})" + ); + } + warn!( + "{} GPU pair directions cannot use NVLink despite a registered fabric; \ + collectives will fall back to PCIe", + blocked.len() + ); + } + Err(err) => warn!("could not read NVLink peer-to-peer status: {err:?}"), + } + } + + #[cfg(test)] + pub(super) mod test_support { + use super::{FabricInfo, FABRIC_STATE_COMPLETED}; + + pub const COMPLETED: u8 = FABRIC_STATE_COMPLETED; + pub const NOT_STARTED: u8 = 1; + pub const IN_PROGRESS: u8 = 2; + pub const SUCCESS: u32 = super::ffi::nvmlReturn_enum_NVML_SUCCESS; + + pub fn info(state: u8, status: u32, clique_id: u32) -> FabricInfo { + FabricInfo { + state, + status, + clique_id, + } + } + } + pub(super) fn set_gpu_ready_state(expected_devices: u32) -> Result<()> { let nvml = init_nvml(expected_devices)?; set_gpu_ready_state_with_nvml(&nvml) @@ -2162,6 +2425,12 @@ impl Stage0<'_> { gpu::evaluate_rego_policy(&gpu_policy, attestation.claims())?; info!("application GPU policy accepted the attestation claims and state"); + // Advisory, and deliberately after attestation so a fabric problem can + // never be mistaken for a security one. Runs on a blocking thread + // because it polls for up to FABRIC_READY_TIMEOUT. + let fabric_devices = expected_devices; + let _ = tokio::task::spawn_blocking(move || gpu::warn_unless_fabric_ready(fabric_devices)) + .await; gpu_state.set_ready()?; let devtools = gpu_state.any_devtools(); let event = attestation.event(devtools)?; @@ -4252,3 +4521,34 @@ Endpoint = [2001:db8::1]:51822 assert!(wireguard_endpoint_hosts("Endpoint = missing-port").is_err()); } } + +#[cfg(test)] +mod fabric_tests { + use super::gpu::{shared_clique, test_support::*}; + + #[test] + fn a_gpu_counts_as_registered_only_when_state_and_status_both_agree() { + // The pairing that matters: Completed with a failed status is a real, + // observed condition, and it is the one NCCL cannot see because it + // logs state and never status. + assert!(info(COMPLETED, SUCCESS, 7).registered()); + assert!(!info(COMPLETED, 999, 7).registered()); + assert!(!info(IN_PROGRESS, SUCCESS, 7).registered()); + assert!(!info(NOT_STARTED, SUCCESS, 7).registered()); + } + + #[test] + fn a_shared_clique_is_only_reported_when_every_gpu_agrees() { + assert_eq!(shared_clique(&[]), None); + assert_eq!(shared_clique(&[info(COMPLETED, SUCCESS, 7)]), Some(7)); + assert_eq!( + shared_clique(&[info(COMPLETED, SUCCESS, 7), info(COMPLETED, SUCCESS, 7)]), + Some(7) + ); + // Both GPUs registered successfully and still cannot reach each other. + assert_eq!( + shared_clique(&[info(COMPLETED, SUCCESS, 7), info(COMPLETED, SUCCESS, 9)]), + None + ); + } +}