vmm: sanitize GPUs through vfio-ioctls hot reset

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 = "1b560c2", default-features = false }
 
 # Core dependencies
 anyhow = { version = "1.0.97", default-features = false }

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

dstack/vmm/src/gpu_reset.rs

@@ -2,25 +2,49 @@
 //
 // 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},
+    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};
 
 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);
 
@@ -33,63 +57,123 @@ 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::<Vec<_>>();
+    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::<BTreeSet<_>>();
-    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<String>, 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 = Arc::new(VfioContainer::new(None).context("failed to open VFIO container")?);
+    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("failed to query PCI hot reset scope")?;
+    let foreign = dependents
+        .iter()
+        .filter(|dep| dep.group_id != group_id)
+        .map(format_dependent)
+        .collect::<Vec<_>>();
+    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::<Vec<_>>().join(", "),
+        "issuing VFIO PCI hot reset"
+    );
+
+    device
+        .pci_hot_reset_own_group()
+        .context("failed to perform PCI hot reset")?;
     Ok(())
 }
 
+fn format_dependent(dep: &PciHotResetDevice) -> String {
+    format!(
+        "{:04x}:{:02x}:{:02x}.{:x} (group {})",
+        dep.segment,
+        dep.bus,
+        dep.devfn >> 3,
+        dep.devfn & 0x7,
+        dep.group_id
+    )
+}
+
+fn iommu_group_id(sysfs_devices: &Path, slot: &str) -> Result<u32> {
+    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<String>,
     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::<Vec<_>>();
         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 +278,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<u16> {
     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 +304,17 @@ mod tests {
         assert!(!is_pci_slot("0000:0g:00.0"));
     }
 
+    #[test]
+    fn formats_dependent_devices_with_pci_slot_and_function() {
+        let dep = PciHotResetDevice {
+            group_id: 46,
+            segment: 0,
+            bus: 0x0f,
+            devfn: 0x1,
+        };
+        assert_eq!(format_dependent(&dep), "0000:0f:00.1 (group 46)");
+    }
+
     #[test]
     fn skips_sanitization_when_gpu_passthrough_is_disabled() {
         let host = HostGpuConfig {

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<String>,
+    /// 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(