From 9b69ca7601c1fe1324affe48dbc161d668a428e0 Mon Sep 17 00:00:00 2001 From: Leechael Yim Date: Fri, 14 Aug 2026 19:49:34 +0800 Subject: [PATCH 1/4] fix(vmm): sanitize GPUs via VFIO hot reset instead of sysfs SBR The sanitize-on-attach path issued the Secondary Bus Reset by writing Bridge Control in the upstream bridge sysfs config space and re-probed devices through /sys/bus/pci/drivers_probe. Both files are writable by root only, so the feature could not be enabled in production where dstack-vmm runs as an unprivileged user with no sudo. Switch to the VFIO_DEVICE_PCI_HOT_RESET ioctl, which makes the kernel perform the same bus reset. The ioctl is authorized by device ownership rather than privilege: the caller presents fds for every VFIO group affected by the reset, and the /dev/vfio group nodes are the same ones QEMU opens to attach the GPU, so the VMM user already has access. A single group fd suffices because every sanitized GPU sits alone behind a dedicated PCIe bridge and alone in its IOMMU group. The bridge topology check is kept as defense, and the kernel-reported set of affected devices must all belong to the GPU own group or the launch is aborted. Devices stay bound to vfio-pci across the reset, so the drivers_probe re-probe logic is no longer needed and is removed. Not yet validated on GPU hardware; see plans/2026-08-14-vfio-gpu-hot-reset.md for the pending experiment. --- dstack/vmm/src/gpu_reset.rs | 352 ++++++++++++++++++++++++++++++------ dstack/vmm/vmm.toml | 5 +- 2 files changed, 295 insertions(+), 62 deletions(-) diff --git a/dstack/vmm/src/gpu_reset.rs b/dstack/vmm/src/gpu_reset.rs index e58876b9a..f35afb238 100644 --- a/dstack/vmm/src/gpu_reset.rs +++ b/dstack/vmm/src/gpu_reset.rs @@ -2,9 +2,36 @@ // // SPDX-License-Identifier: Apache-2.0 +//! GPU sanitization before QEMU attach. +//! +//! Terminating QEMU while GSP/SPDM initialization is in progress can leave +//! stale FSP/SPDM session state in a GPU. VFIO's attach-time FLR does not +//! clear that state; a PCIe Secondary Bus Reset does. The next guest +//! otherwise fails with an SPDM timeout followed by GSP/RmInitAdapter errors. +//! +//! The VMM runs as an unprivileged user, so it cannot issue the SBR by +//! writing Bridge Control in the upstream bridge's sysfs config space: that +//! file is writable by root only, as is /sys/bus/pci/drivers_probe. Instead +//! the VFIO_DEVICE_PCI_HOT_RESET ioctl asks the kernel to perform the same +//! bus reset. The ioctl is authorized by device ownership rather than +//! privilege: the caller must present an fd for every VFIO group affected by +//! the reset. The group nodes under /dev/vfio are the same ones QEMU opens +//! to attach the GPU, so the VMM user already has access to them. +//! +//! A single group fd suffices only because every sanitized GPU sits alone +//! behind a dedicated PCIe bridge and alone in its IOMMU group. The bridge +//! topology is validated first, and the kernel-reported set of devices +//! affected by the reset must all belong to the GPU's own group; anything +//! else aborts the launch rather than risking disruption to other devices. +//! +//! A VFIO group can be opened by only one process at a time, so every fd is +//! closed again before QEMU is spawned. + use std::{ collections::BTreeSet, fs::{File, OpenOptions}, + mem::size_of, + os::fd::{AsRawFd, FromRawFd, OwnedFd}, os::unix::fs::FileExt, path::{Path, PathBuf}, thread, @@ -18,12 +45,44 @@ use crate::{app::GpuConfig, config::GpuConfig as HostGpuConfig}; const PCI_SYSFS_DEVICES: &str = "/sys/bus/pci/devices"; const PCI_BRIDGE_CLASS: u32 = 0x0604; -const PCI_BRIDGE_CONTROL: u64 = 0x3e; -const PCI_BRIDGE_CTL_BUS_RESET: u16 = 1 << 6; -const SBR_ASSERT_TIME: Duration = Duration::from_millis(100); const SBR_POLL_INTERVAL: Duration = Duration::from_millis(100); const SBR_STABLE_TIME: Duration = Duration::from_millis(500); +const VFIO_CONTAINER: &str = "/dev/vfio/vfio"; +const VFIO_GROUP_DIR: &str = "/dev/vfio"; +const VFIO_API_VERSION: i32 = 0; +const VFIO_TYPE1_IOMMU: usize = 1; +const VFIO_TYPE1V2_IOMMU: usize = 3; +const VFIO_GROUP_FLAGS_VIABLE: u32 = 1; + +// _IO(';', 100 + n) from : no argument size is encoded. +const fn vfio_io(nr: u64) -> libc::c_ulong { + ((b';' as u64) << 8 | (100 + nr)) as libc::c_ulong +} +const VFIO_GET_API_VERSION: libc::c_ulong = vfio_io(0); +const VFIO_CHECK_EXTENSION: libc::c_ulong = vfio_io(1); +const VFIO_SET_IOMMU: libc::c_ulong = vfio_io(2); +const VFIO_GROUP_GET_STATUS: libc::c_ulong = vfio_io(3); +const VFIO_GROUP_SET_CONTAINER: libc::c_ulong = vfio_io(4); +const VFIO_GROUP_GET_DEVICE_FD: libc::c_ulong = vfio_io(6); +const VFIO_DEVICE_PCI_GET_HOT_RESET_INFO: libc::c_ulong = vfio_io(12); +const VFIO_DEVICE_PCI_HOT_RESET: libc::c_ulong = vfio_io(13); + +#[repr(C)] +struct VfioGroupStatus { + argsz: u32, + flags: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct VfioPciDependentDevice { + group_id: u32, + segment: u16, + bus: u8, + devfn: u8, +} + /// Clears device-internal state that can survive VFIO's attach-time FLR. /// /// Each selected GPU must sit alone behind a dedicated PCIe bridge. Resetting @@ -33,63 +92,252 @@ pub fn sanitize_on_attach(host: &HostGpuConfig, devices: &GpuConfig) -> Result<( if !host.enabled || !host.sanitize_on_attach || devices.gpus.is_empty() { return Ok(()); } - sanitize_at( - Path::new(PCI_SYSFS_DEVICES), - devices, - Duration::from_millis(host.sbr_timeout_ms), - ) + let slots = devices + .gpus + .iter() + .map(|gpu| gpu.slot.clone()) + .collect::>(); + sanitize_slots(&slots, Duration::from_millis(host.sbr_timeout_ms)) } -fn sanitize_at(sysfs_devices: &Path, devices: &GpuConfig, timeout: Duration) -> Result<()> { - let selected = devices - .gpus +/// Sanitizes the given GPU slots. Entry point for the `sanitize-gpu` +/// subcommand; `sanitize_on_attach` funnels here as well. +pub fn sanitize_slots(slots: &[String], timeout: Duration) -> Result<()> { + if slots.is_empty() { + bail!("no GPU slots specified"); + } + let selected = slots .iter() - .map(|gpu| normalize_slot(&gpu.slot)) + .map(|slot| normalize_slot(slot)) .collect::>(); - let mut bridges = BTreeSet::new(); + sanitize_at(Path::new(PCI_SYSFS_DEVICES), &selected, timeout) +} - for gpu in &selected { +fn sanitize_at(sysfs_devices: &Path, selected: &BTreeSet, timeout: Duration) -> Result<()> { + for gpu in selected { let bridge = upstream_bridge(sysfs_devices, gpu)?; ensure_dedicated_bridge(&bridge, gpu)?; - bridges.insert(bridge); } - for bridge in bridges { - info!(bridge = %bridge.display(), "sanitizing GPU with PCIe Secondary Bus Reset"); - secondary_bus_reset(&bridge) - .with_context(|| format!("failed to sanitize GPU using bridge {}", bridge.display()))?; + for gpu in selected { + info!(gpu = %gpu, "sanitizing GPU with VFIO PCI hot reset"); + hot_reset(sysfs_devices, gpu) + .with_context(|| format!("failed to sanitize GPU {gpu} with VFIO hot reset"))?; + } + wait_for_vfio_ready(sysfs_devices, selected, timeout)?; + Ok(()) +} + +/// Issues a Secondary Bus Reset on the GPU's upstream bridge through VFIO. +/// +/// All fds are dropped on return so that QEMU can open the group afterwards. +fn hot_reset(sysfs_devices: &Path, slot: &str) -> Result<()> { + let group_id = iommu_group_id(sysfs_devices, slot)?; + let container = open_container()?; + let group = open_group(&container, group_id)?; + let device = open_device(&group, slot)?; + + let dependents = hot_reset_dependents(&device)?; + let foreign = dependents + .iter() + .filter(|dep| dep.group_id != group_id) + .map(format_dependent) + .collect::>(); + if !foreign.is_empty() { + bail!( + "refusing hot reset of GPU {slot}: it would also reset devices outside \ + IOMMU group {group_id}: {}", + foreign.join(", ") + ); } - wait_for_vfio_ready(sysfs_devices, &selected, timeout)?; + info!( + gpu = %slot, + affected = %dependents.iter().map(format_dependent).collect::>().join(", "), + "issuing VFIO PCI hot reset" + ); + + // struct vfio_pci_hot_reset { argsz, flags, count, group_fds[] } + let mut request = [0_u32; 4]; + request[0] = (size_of::<[u32; 4]>()) as u32; + request[2] = 1; + request[3] = group.as_raw_fd() as u32; + vfio_ioctl( + &device, + VFIO_DEVICE_PCI_HOT_RESET, + request.as_mut_ptr() as usize, + "VFIO_DEVICE_PCI_HOT_RESET", + )?; Ok(()) } +fn open_container() -> Result { + let container = OpenOptions::new() + .read(true) + .write(true) + .open(VFIO_CONTAINER) + .with_context(|| format!("failed to open {VFIO_CONTAINER}"))?; + let version = vfio_ioctl(&container, VFIO_GET_API_VERSION, 0, "VFIO_GET_API_VERSION")?; + if version != VFIO_API_VERSION { + bail!("unsupported VFIO API version {version}"); + } + Ok(container) +} + +fn open_group(container: &File, group_id: u32) -> Result { + let path = format!("{VFIO_GROUP_DIR}/{group_id}"); + let group = OpenOptions::new() + .read(true) + .write(true) + .open(&path) + .with_context(|| format!("failed to open {path} (is the VMM user in its group?)"))?; + + let mut status = VfioGroupStatus { + argsz: size_of::() as u32, + flags: 0, + }; + vfio_ioctl( + &group, + VFIO_GROUP_GET_STATUS, + &mut status as *mut _ as usize, + "VFIO_GROUP_GET_STATUS", + )?; + if status.flags & VFIO_GROUP_FLAGS_VIABLE == 0 { + bail!("IOMMU group {group_id} is not viable; are all of its devices bound to vfio-pci?"); + } + + let container_fd = container.as_raw_fd(); + vfio_ioctl( + &group, + VFIO_GROUP_SET_CONTAINER, + &container_fd as *const _ as usize, + "VFIO_GROUP_SET_CONTAINER", + )?; + // An IOMMU backend must be set before device fds can be handed out. No + // DMA mappings are created; the container exists only for this reset. + let type1v2 = vfio_ioctl( + container, + VFIO_CHECK_EXTENSION, + VFIO_TYPE1V2_IOMMU, + "VFIO_CHECK_EXTENSION", + )?; + let iommu = if type1v2 > 0 { + VFIO_TYPE1V2_IOMMU + } else { + VFIO_TYPE1_IOMMU + }; + vfio_ioctl(container, VFIO_SET_IOMMU, iommu, "VFIO_SET_IOMMU")?; + Ok(group) +} + +fn open_device(group: &File, slot: &str) -> Result { + let name = std::ffi::CString::new(slot).context("invalid PCI slot name")?; + let fd = vfio_ioctl( + group, + VFIO_GROUP_GET_DEVICE_FD, + name.as_ptr() as usize, + "VFIO_GROUP_GET_DEVICE_FD", + )?; + Ok(unsafe { OwnedFd::from_raw_fd(fd) }) +} + +/// Returns the devices the kernel reports as affected by a hot reset. +fn hot_reset_dependents(device: &OwnedFd) -> Result> { + // struct vfio_pci_hot_reset_info { argsz, flags, count, devices[] } + const HEADER: usize = size_of::<[u32; 3]>(); + const ENTRY: usize = size_of::(); + + let mut probe = [HEADER as u32, 0, 0]; + let ret = unsafe { + libc::ioctl( + device.as_raw_fd(), + VFIO_DEVICE_PCI_GET_HOT_RESET_INFO, + probe.as_mut_ptr(), + ) + }; + if ret < 0 { + let err = std::io::Error::last_os_error(); + // ENOSPC is the expected way to learn the entry count. + if err.raw_os_error() != Some(libc::ENOSPC) { + return Err(err).context("VFIO_DEVICE_PCI_GET_HOT_RESET_INFO failed"); + } + } + let count = probe[2] as usize; + + // u64 storage keeps the buffer aligned for the header and entries. + let mut buffer = vec![0_u64; (HEADER + count * ENTRY).div_ceil(size_of::())]; + let header = buffer.as_mut_ptr() as *mut u32; + unsafe { + *header = (HEADER + count * ENTRY) as u32; + } + vfio_ioctl( + device, + VFIO_DEVICE_PCI_GET_HOT_RESET_INFO, + buffer.as_mut_ptr() as usize, + "VFIO_DEVICE_PCI_GET_HOT_RESET_INFO", + )?; + let filled = unsafe { *header.add(2) } as usize; + let entries = unsafe { + std::slice::from_raw_parts( + (buffer.as_ptr() as *const u8).add(HEADER) as *const VfioPciDependentDevice, + filled.min(count), + ) + }; + Ok(entries.to_vec()) +} + +fn format_dependent(dep: &VfioPciDependentDevice) -> String { + format!( + "{:04x}:{:02x}:{:02x}.{:x} (group {})", + dep.segment, + dep.bus, + dep.devfn >> 3, + dep.devfn & 0x7, + dep.group_id + ) +} + +fn vfio_ioctl(fd: &impl AsRawFd, request: libc::c_ulong, arg: usize, what: &str) -> Result { + let ret = unsafe { libc::ioctl(fd.as_raw_fd(), request, arg) }; + if ret < 0 { + return Err(std::io::Error::last_os_error()).with_context(|| format!("{what} failed")); + } + Ok(ret) +} + +fn iommu_group_id(sysfs_devices: &Path, slot: &str) -> Result { + let link = sysfs_devices.join(slot).join("iommu_group"); + let group = link + .canonicalize() + .with_context(|| format!("failed to resolve IOMMU group of {slot}"))?; + group + .file_name() + .and_then(|name| name.to_str()) + .and_then(|name| name.parse().ok()) + .with_context(|| format!("invalid IOMMU group path {}", group.display())) +} + fn wait_for_vfio_ready( sysfs_devices: &Path, devices: &BTreeSet, timeout: Duration, ) -> Result<()> { - let drivers_probe = sysfs_devices - .parent() - .context("PCI sysfs devices directory has no parent")? - .join("drivers_probe"); - let deadline = Instant::now() + timeout; + let started = Instant::now(); + let deadline = started + timeout; let mut stable_since = None; loop { let not_ready = devices .iter() - .filter(|slot| { - let device = sysfs_devices.join(slot); - if !driver_is_vfio(&device) { - let _ = fs_err::write(&drivers_probe, format!("{slot}\n")); - } - !is_vfio_ready(&device) - }) + .filter(|slot| !is_vfio_ready(&sysfs_devices.join(slot))) .cloned() .collect::>(); if not_ready.is_empty() { let since = *stable_since.get_or_insert_with(Instant::now); if since.elapsed() >= SBR_STABLE_TIME { - info!(count = devices.len(), "all sanitized GPUs are VFIO-ready"); + info!( + count = devices.len(), + elapsed_ms = started.elapsed().as_millis() as u64, + "all sanitized GPUs are VFIO-ready" + ); return Ok(()); } } else { @@ -194,39 +442,12 @@ fn is_pci_slot(name: &str) -> bool { .all(|(i, c)| matches!(i, 4 | 7 | 10) || c.is_ascii_hexdigit()) } -fn secondary_bus_reset(bridge: &Path) -> Result<()> { - let config_path = bridge.join("config"); - let config = OpenOptions::new() - .read(true) - .write(true) - .open(&config_path) - .with_context(|| format!("failed to open {}", config_path.display()))?; - let original = read_u16(&config, PCI_BRIDGE_CONTROL)?; - write_u16( - &config, - PCI_BRIDGE_CONTROL, - original | PCI_BRIDGE_CTL_BUS_RESET, - )?; - thread::sleep(SBR_ASSERT_TIME); - write_u16( - &config, - PCI_BRIDGE_CONTROL, - original & !PCI_BRIDGE_CTL_BUS_RESET, - )?; - Ok(()) -} - fn read_u16(file: &File, offset: u64) -> Result { let mut value = [0_u8; 2]; file.read_exact_at(&mut value, offset)?; Ok(u16::from_le_bytes(value)) } -fn write_u16(file: &File, offset: u64, value: u16) -> Result<()> { - file.write_all_at(&value.to_le_bytes(), offset)?; - Ok(()) -} - #[cfg(test)] mod tests { use super::*; @@ -247,6 +468,17 @@ mod tests { assert!(!is_pci_slot("0000:0g:00.0")); } + #[test] + fn formats_dependent_devices_with_pci_slot_and_function() { + let dep = VfioPciDependentDevice { + group_id: 46, + segment: 0, + bus: 0x0f, + devfn: (0x00 << 3) | 0x1, + }; + assert_eq!(format_dependent(&dep), "0000:0f:00.1 (group 46)"); + } + #[test] fn skips_sanitization_when_gpu_passthrough_is_disabled() { let host = HostGpuConfig { diff --git a/dstack/vmm/vmm.toml b/dstack/vmm/vmm.toml index 599c3e5c6..09954380d 100644 --- a/dstack/vmm/vmm.toml +++ b/dstack/vmm/vmm.toml @@ -166,8 +166,9 @@ include = [] # Allow attach all GPUs allow_attach_all = true # Issue a Secondary Bus Reset on each GPU's dedicated upstream PCIe bridge -# immediately before QEMU attaches the device. This clears stale GPU FSP/SPDM -# state left by a VM that terminated during GPU firmware initialization. +# immediately before QEMU attaches the device, using a VFIO PCI hot reset so +# that no root privileges are needed. This clears stale GPU FSP/SPDM state +# left by a VM that terminated during GPU firmware initialization. sanitize_on_attach = true # Maximum time for all reset GPUs to become continuously VFIO-ready. sbr_timeout_ms = 10000 From 0ca9646665a38081cf45c773327a467e82b7dfc1 Mon Sep 17 00:00:00 2001 From: Leechael Yim Date: Fri, 14 Aug 2026 19:49:41 +0800 Subject: [PATCH 2/4] feat(vmm): add sanitize-gpu subcommand for manual GPU reset Expose the sanitize path as "dstack-vmm sanitize-gpu ..." so operators can reset GPUs by hand and the pending hardware experiment can exercise exactly the code path used at VM launch, running as the unprivileged VMM user. The subcommand needs no server configuration, only /dev/vfio access, and is handled before config loading like the other special modes. --- dstack/vmm/src/main.rs | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/dstack/vmm/src/main.rs b/dstack/vmm/src/main.rs index f13489bfb..ca099a4ec 100644 --- a/dstack/vmm/src/main.rs +++ b/dstack/vmm/src/main.rs @@ -65,6 +65,8 @@ enum Command { CheckConfig, /// One-shot VM execution mode for debugging Run(RunArgs), + /// Sanitize GPUs with a VFIO PCI hot reset (debugging/ops) + SanitizeGpu(SanitizeGpuArgs), /// Run the privileged TAP and libvirt nwfilter broker. Netd(NetdArgs), /// Internal per-VM QEMU/swtpm launcher. @@ -94,6 +96,16 @@ struct RunArgs { dry_run: bool, } +#[derive(ClapArgs)] +struct SanitizeGpuArgs { + /// PCI slots of the GPUs to reset, e.g. 0000:0f:00.0 + #[arg(required = true)] + slots: Vec, + /// Maximum time in milliseconds for the GPUs to become VFIO-ready again + #[arg(long, default_value_t = 10_000)] + timeout_ms: u64, +} + #[derive(ClapArgs)] struct VmLauncherArgs { /// Path to the generated VM launch specification. @@ -211,6 +223,14 @@ async fn main() -> Result<()> { return vm_launcher::run(Path::new(&launcher_args.spec)).await; } + // Needs no server configuration; only /dev/vfio access. + if let Some(Command::SanitizeGpu(sanitize_args)) = &args.command { + return gpu_reset::sanitize_slots( + &sanitize_args.slots, + Duration::from_millis(sanitize_args.timeout_ms), + ); + } + let figment = config::load_config_figment(args.config.as_deref()); if let Some(Command::Netd(netd_args)) = &args.command { let mut netd_config: NetdConfig = figment @@ -263,6 +283,7 @@ async fn main() -> Result<()> { return Ok(()); } Command::Netd(_) => unreachable!("netd mode handled before server startup"), + Command::SanitizeGpu(_) => unreachable!("sanitize-gpu handled before config loading"), Command::Run(run_args) => { // One-shot VM execution mode return one_shot::run_one_shot( From 1dfe43dfbbbeddb47e47ba2e1937f66711204306 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 06:09:14 -0700 Subject: [PATCH 3/4] refactor(vmm): use vfio-ioctls for GPU hot reset --- dstack/Cargo.lock | 41 ++++++++ dstack/Cargo.toml | 1 + dstack/vmm/Cargo.toml | 1 + dstack/vmm/src/gpu_reset.rs | 199 ++++-------------------------------- 4 files changed, 62 insertions(+), 180 deletions(-) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 1fbf77bcb..43c3f6413 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -2354,6 +2354,7 @@ dependencies = [ "tracing-subscriber", "url", "uuid", + "vfio-ioctls", "wait-timeout", "which 7.0.3", ] @@ -8185,12 +8186,52 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "vfio-bindings" +version = "0.6.2" +source = "git+https://github.com/kvinwang/vfio.git?rev=d509db2#d509db2b403d258659c8dcf597a7b189d9512940" + +[[package]] +name = "vfio-ioctls" +version = "0.8.0" +source = "git+https://github.com/kvinwang/vfio.git?rev=d509db2#d509db2b403d258659c8dcf597a7b189d9512940" +dependencies = [ + "byteorder", + "libc", + "log", + "thiserror 2.0.18", + "vfio-bindings", + "vm-memory", + "vmm-sys-util", +] + [[package]] name = "virtue" version = "0.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "051eb1abcf10076295e815102942cc58f9d5e3b4560e46e53c21e8ff6f3af7b1" +[[package]] +name = "vm-memory" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b55e753c7725603745cb32b2287ef7ef3da05c03c7702cda3fa8abe25ae0465" +dependencies = [ + "libc", + "thiserror 2.0.18", + "winapi", +] + +[[package]] +name = "vmm-sys-util" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "506c62fdf617a5176827c2f9afbcf1be155b03a9b4bf9617a60dbc07e3a1642f" +dependencies = [ + "bitflags 1.3.2", + "libc", +] + [[package]] name = "void" version = "1.0.2" diff --git a/dstack/Cargo.toml b/dstack/Cargo.toml index 3f3fd93a0..b4e81c38b 100644 --- a/dstack/Cargo.toml +++ b/dstack/Cargo.toml @@ -122,6 +122,7 @@ dstack-mr = { path = "dstack-mr" } dstack-verifier = { path = "verifier", default-features = false } size-parser = { path = "size-parser" } wavekv = "1.0.0" +vfio-ioctls = { git = "https://github.com/kvinwang/vfio.git", rev = "d509db2", default-features = false } # Core dependencies anyhow = { version = "1.0.97", default-features = false } diff --git a/dstack/vmm/Cargo.toml b/dstack/vmm/Cargo.toml index 941802b05..6e4c5a988 100644 --- a/dstack/vmm/Cargo.toml +++ b/dstack/vmm/Cargo.toml @@ -64,6 +64,7 @@ flate2.workspace = true tar.workspace = true tempfile.workspace = true wait-timeout.workspace = true +vfio-ioctls.workspace = true [dev-dependencies] insta.workspace = true diff --git a/dstack/vmm/src/gpu_reset.rs b/dstack/vmm/src/gpu_reset.rs index f35afb238..0ce2efa18 100644 --- a/dstack/vmm/src/gpu_reset.rs +++ b/dstack/vmm/src/gpu_reset.rs @@ -29,17 +29,17 @@ use std::{ collections::BTreeSet, - fs::{File, OpenOptions}, - mem::size_of, - os::fd::{AsRawFd, FromRawFd, OwnedFd}, + fs::File, os::unix::fs::FileExt, path::{Path, PathBuf}, + sync::Arc, thread, time::{Duration, Instant}, }; use anyhow::{bail, Context, Result}; use tracing::info; +use vfio_ioctls::{PciHotResetDevice, VfioContainer, VfioDevice}; use crate::{app::GpuConfig, config::GpuConfig as HostGpuConfig}; @@ -48,41 +48,6 @@ const PCI_BRIDGE_CLASS: u32 = 0x0604; const SBR_POLL_INTERVAL: Duration = Duration::from_millis(100); const SBR_STABLE_TIME: Duration = Duration::from_millis(500); -const VFIO_CONTAINER: &str = "/dev/vfio/vfio"; -const VFIO_GROUP_DIR: &str = "/dev/vfio"; -const VFIO_API_VERSION: i32 = 0; -const VFIO_TYPE1_IOMMU: usize = 1; -const VFIO_TYPE1V2_IOMMU: usize = 3; -const VFIO_GROUP_FLAGS_VIABLE: u32 = 1; - -// _IO(';', 100 + n) from : no argument size is encoded. -const fn vfio_io(nr: u64) -> libc::c_ulong { - ((b';' as u64) << 8 | (100 + nr)) as libc::c_ulong -} -const VFIO_GET_API_VERSION: libc::c_ulong = vfio_io(0); -const VFIO_CHECK_EXTENSION: libc::c_ulong = vfio_io(1); -const VFIO_SET_IOMMU: libc::c_ulong = vfio_io(2); -const VFIO_GROUP_GET_STATUS: libc::c_ulong = vfio_io(3); -const VFIO_GROUP_SET_CONTAINER: libc::c_ulong = vfio_io(4); -const VFIO_GROUP_GET_DEVICE_FD: libc::c_ulong = vfio_io(6); -const VFIO_DEVICE_PCI_GET_HOT_RESET_INFO: libc::c_ulong = vfio_io(12); -const VFIO_DEVICE_PCI_HOT_RESET: libc::c_ulong = vfio_io(13); - -#[repr(C)] -struct VfioGroupStatus { - argsz: u32, - flags: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -struct VfioPciDependentDevice { - group_id: u32, - segment: u16, - bus: u8, - devfn: u8, -} - /// Clears device-internal state that can survive VFIO's attach-time FLR. /// /// Each selected GPU must sit alone behind a dedicated PCIe bridge. Resetting @@ -133,11 +98,16 @@ fn sanitize_at(sysfs_devices: &Path, selected: &BTreeSet, timeout: Durat /// All fds are dropped on return so that QEMU can open the group afterwards. fn hot_reset(sysfs_devices: &Path, slot: &str) -> Result<()> { let group_id = iommu_group_id(sysfs_devices, slot)?; - let container = open_container()?; - let group = open_group(&container, group_id)?; - let device = open_device(&group, slot)?; - - let dependents = hot_reset_dependents(&device)?; + let container = Arc::new(VfioContainer::new(None).context("failed to open VFIO container")?); + let group = container + .get_group(group_id) + .with_context(|| format!("failed to open VFIO group {group_id}"))?; + let device = VfioDevice::new(&sysfs_devices.join(slot), container) + .with_context(|| format!("failed to open VFIO device {slot}"))?; + + let dependents = device + .pci_hot_reset_info() + .context("VFIO_DEVICE_GET_PCI_HOT_RESET_INFO failed")?; let foreign = dependents .iter() .filter(|dep| dep.group_id != group_id) @@ -156,136 +126,13 @@ fn hot_reset(sysfs_devices: &Path, slot: &str) -> Result<()> { "issuing VFIO PCI hot reset" ); - // struct vfio_pci_hot_reset { argsz, flags, count, group_fds[] } - let mut request = [0_u32; 4]; - request[0] = (size_of::<[u32; 4]>()) as u32; - request[2] = 1; - request[3] = group.as_raw_fd() as u32; - vfio_ioctl( - &device, - VFIO_DEVICE_PCI_HOT_RESET, - request.as_mut_ptr() as usize, - "VFIO_DEVICE_PCI_HOT_RESET", - )?; + device + .pci_hot_reset(&[&group]) + .context("VFIO_DEVICE_PCI_HOT_RESET failed")?; Ok(()) } -fn open_container() -> Result { - let container = OpenOptions::new() - .read(true) - .write(true) - .open(VFIO_CONTAINER) - .with_context(|| format!("failed to open {VFIO_CONTAINER}"))?; - let version = vfio_ioctl(&container, VFIO_GET_API_VERSION, 0, "VFIO_GET_API_VERSION")?; - if version != VFIO_API_VERSION { - bail!("unsupported VFIO API version {version}"); - } - Ok(container) -} - -fn open_group(container: &File, group_id: u32) -> Result { - let path = format!("{VFIO_GROUP_DIR}/{group_id}"); - let group = OpenOptions::new() - .read(true) - .write(true) - .open(&path) - .with_context(|| format!("failed to open {path} (is the VMM user in its group?)"))?; - - let mut status = VfioGroupStatus { - argsz: size_of::() as u32, - flags: 0, - }; - vfio_ioctl( - &group, - VFIO_GROUP_GET_STATUS, - &mut status as *mut _ as usize, - "VFIO_GROUP_GET_STATUS", - )?; - if status.flags & VFIO_GROUP_FLAGS_VIABLE == 0 { - bail!("IOMMU group {group_id} is not viable; are all of its devices bound to vfio-pci?"); - } - - let container_fd = container.as_raw_fd(); - vfio_ioctl( - &group, - VFIO_GROUP_SET_CONTAINER, - &container_fd as *const _ as usize, - "VFIO_GROUP_SET_CONTAINER", - )?; - // An IOMMU backend must be set before device fds can be handed out. No - // DMA mappings are created; the container exists only for this reset. - let type1v2 = vfio_ioctl( - container, - VFIO_CHECK_EXTENSION, - VFIO_TYPE1V2_IOMMU, - "VFIO_CHECK_EXTENSION", - )?; - let iommu = if type1v2 > 0 { - VFIO_TYPE1V2_IOMMU - } else { - VFIO_TYPE1_IOMMU - }; - vfio_ioctl(container, VFIO_SET_IOMMU, iommu, "VFIO_SET_IOMMU")?; - Ok(group) -} - -fn open_device(group: &File, slot: &str) -> Result { - let name = std::ffi::CString::new(slot).context("invalid PCI slot name")?; - let fd = vfio_ioctl( - group, - VFIO_GROUP_GET_DEVICE_FD, - name.as_ptr() as usize, - "VFIO_GROUP_GET_DEVICE_FD", - )?; - Ok(unsafe { OwnedFd::from_raw_fd(fd) }) -} - -/// Returns the devices the kernel reports as affected by a hot reset. -fn hot_reset_dependents(device: &OwnedFd) -> Result> { - // struct vfio_pci_hot_reset_info { argsz, flags, count, devices[] } - const HEADER: usize = size_of::<[u32; 3]>(); - const ENTRY: usize = size_of::(); - - let mut probe = [HEADER as u32, 0, 0]; - let ret = unsafe { - libc::ioctl( - device.as_raw_fd(), - VFIO_DEVICE_PCI_GET_HOT_RESET_INFO, - probe.as_mut_ptr(), - ) - }; - if ret < 0 { - let err = std::io::Error::last_os_error(); - // ENOSPC is the expected way to learn the entry count. - if err.raw_os_error() != Some(libc::ENOSPC) { - return Err(err).context("VFIO_DEVICE_PCI_GET_HOT_RESET_INFO failed"); - } - } - let count = probe[2] as usize; - - // u64 storage keeps the buffer aligned for the header and entries. - let mut buffer = vec![0_u64; (HEADER + count * ENTRY).div_ceil(size_of::())]; - let header = buffer.as_mut_ptr() as *mut u32; - unsafe { - *header = (HEADER + count * ENTRY) as u32; - } - vfio_ioctl( - device, - VFIO_DEVICE_PCI_GET_HOT_RESET_INFO, - buffer.as_mut_ptr() as usize, - "VFIO_DEVICE_PCI_GET_HOT_RESET_INFO", - )?; - let filled = unsafe { *header.add(2) } as usize; - let entries = unsafe { - std::slice::from_raw_parts( - (buffer.as_ptr() as *const u8).add(HEADER) as *const VfioPciDependentDevice, - filled.min(count), - ) - }; - Ok(entries.to_vec()) -} - -fn format_dependent(dep: &VfioPciDependentDevice) -> String { +fn format_dependent(dep: &PciHotResetDevice) -> String { format!( "{:04x}:{:02x}:{:02x}.{:x} (group {})", dep.segment, @@ -296,14 +143,6 @@ fn format_dependent(dep: &VfioPciDependentDevice) -> String { ) } -fn vfio_ioctl(fd: &impl AsRawFd, request: libc::c_ulong, arg: usize, what: &str) -> Result { - let ret = unsafe { libc::ioctl(fd.as_raw_fd(), request, arg) }; - if ret < 0 { - return Err(std::io::Error::last_os_error()).with_context(|| format!("{what} failed")); - } - Ok(ret) -} - fn iommu_group_id(sysfs_devices: &Path, slot: &str) -> Result { let link = sysfs_devices.join(slot).join("iommu_group"); let group = link @@ -470,11 +309,11 @@ mod tests { #[test] fn formats_dependent_devices_with_pci_slot_and_function() { - let dep = VfioPciDependentDevice { + let dep = PciHotResetDevice { group_id: 46, segment: 0, bus: 0x0f, - devfn: (0x00 << 3) | 0x1, + devfn: 0x1, }; assert_eq!(format_dependent(&dep), "0000:0f:00.1 (group 46)"); } From 3d9c16185bf62b4e5f7847e0a63d55976f68526c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 06:42:58 -0700 Subject: [PATCH 4/4] refactor(vmm): keep VFIO group handling encapsulated --- dstack/Cargo.lock | 4 ++-- dstack/Cargo.toml | 2 +- dstack/vmm/src/gpu_reset.rs | 9 +++------ 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 43c3f6413..b81bc3107 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -8189,12 +8189,12 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" [[package]] name = "vfio-bindings" version = "0.6.2" -source = "git+https://github.com/kvinwang/vfio.git?rev=d509db2#d509db2b403d258659c8dcf597a7b189d9512940" +source = "git+https://github.com/kvinwang/vfio.git?rev=1b560c2#1b560c20c8c294db0df5c6d8e0346ee7f9d89b62" [[package]] name = "vfio-ioctls" version = "0.8.0" -source = "git+https://github.com/kvinwang/vfio.git?rev=d509db2#d509db2b403d258659c8dcf597a7b189d9512940" +source = "git+https://github.com/kvinwang/vfio.git?rev=1b560c2#1b560c20c8c294db0df5c6d8e0346ee7f9d89b62" dependencies = [ "byteorder", "libc", diff --git a/dstack/Cargo.toml b/dstack/Cargo.toml index b4e81c38b..02d61d59b 100644 --- a/dstack/Cargo.toml +++ b/dstack/Cargo.toml @@ -122,7 +122,7 @@ dstack-mr = { path = "dstack-mr" } dstack-verifier = { path = "verifier", default-features = false } size-parser = { path = "size-parser" } wavekv = "1.0.0" -vfio-ioctls = { git = "https://github.com/kvinwang/vfio.git", rev = "d509db2", default-features = false } +vfio-ioctls = { git = "https://github.com/kvinwang/vfio.git", rev = "1b560c2", default-features = false } # Core dependencies anyhow = { version = "1.0.97", default-features = false } diff --git a/dstack/vmm/src/gpu_reset.rs b/dstack/vmm/src/gpu_reset.rs index 0ce2efa18..5e76404d0 100644 --- a/dstack/vmm/src/gpu_reset.rs +++ b/dstack/vmm/src/gpu_reset.rs @@ -99,15 +99,12 @@ fn sanitize_at(sysfs_devices: &Path, selected: &BTreeSet, timeout: Durat fn hot_reset(sysfs_devices: &Path, slot: &str) -> Result<()> { let group_id = iommu_group_id(sysfs_devices, slot)?; let container = Arc::new(VfioContainer::new(None).context("failed to open VFIO container")?); - let group = container - .get_group(group_id) - .with_context(|| format!("failed to open VFIO group {group_id}"))?; let device = VfioDevice::new(&sysfs_devices.join(slot), container) .with_context(|| format!("failed to open VFIO device {slot}"))?; let dependents = device .pci_hot_reset_info() - .context("VFIO_DEVICE_GET_PCI_HOT_RESET_INFO failed")?; + .context("failed to query PCI hot reset scope")?; let foreign = dependents .iter() .filter(|dep| dep.group_id != group_id) @@ -127,8 +124,8 @@ fn hot_reset(sysfs_devices: &Path, slot: &str) -> Result<()> { ); device - .pci_hot_reset(&[&group]) - .context("VFIO_DEVICE_PCI_HOT_RESET failed")?; + .pci_hot_reset_own_group() + .context("failed to perform PCI hot reset")?; Ok(()) }