feat(qemu-acpi): model guest NUMA nodes without PXB expanders

dstack/crates/qemu-acpi/examples/dump.rs

@@ -29,11 +29,15 @@ fn main() -> Result<(), Box<dyn Error>> {
     let memory_size = optional(&mut args, 2u64 << 30)?;
     let volumes = optional(&mut args, 0)?;
     let pic = optional::<u8>(&mut args, 0)? == 1;
-    // Comma-separated PXB bus of each NUMA node; empty keeps the legacy layout.
+    // Comma-separated PXB bus of each NUMA node, `-` for none; empty keeps the
+    // legacy layout.
     let numa_nodes = optional::<String>(&mut args, String::new())?
         .split(',')
         .filter(|bus| !bus.is_empty())
-        .map(|bus| bus.parse().map(|bus| NumaNode { pxb_bus: Some(bus) }))
+        .map(|bus| match bus {
+            "-" => Ok(NumaNode { pxb_bus: None }),
+            bus => bus.parse().map(|bus| NumaNode { pxb_bus: Some(bus) }),
+        })
         .collect::<Result<_, _>>()?;
     // Devices appended after the root ports: `e` endpoint, `r` root port.
     let extra_root_devices = optional::<String>(&mut args, String::new())?

dstack/crates/qemu-acpi/fixtures/README.md

@@ -55,6 +55,10 @@ for each compatibility family cover the Cartesian product of:
 whose PXBs sit on buses 5, 6 and 8 (12 vCPUs, 6 GiB, two NICs), together with
 their loader and RSDP blobs.
 
+`qemu-11.1-q35-hotplug-off-numa2-plain-*.bin` capture two guest NUMA nodes
+without PXBs (8 vCPUs, 2 GiB, one NIC, two GPU root ports and one extra
+endpoint on the root bus).
+
 `qemu-11.1-q35-one-nic.bin` is an untrimmed byte-for-byte regression fixture.
 Fixture names identify the emulated compatibility version, not necessarily the
 source tree's release number.

dstack/crates/qemu-acpi/scripts/differential-qemu.sh

@@ -24,17 +24,19 @@ rust_dump="$workspace/target/debug/examples/dump"
 
 run_case() {
     local version=$1 topology=$2
-    # pxb_buses: one guest NUMA node per bus, GPUs split evenly across them.
-    local hugepages=0 gpus=0 pxb_buses=()
+    # pxb_buses: one guest NUMA node per bus, GPUs split evenly across them;
+    # plain_nodes: guest NUMA nodes without expanders, GPUs on the root bus.
+    local hugepages=0 gpus=0 pxb_buses=() plain_nodes=0 hotplug_off=0
     case "$topology" in
         normal) ;;
         numa) hugepages=1 ;;
         numa-pxb) hugepages=1; gpus=8; pxb_buses=(5) ;;
         numa2-pxb) hugepages=1; gpus=8; pxb_buses=(5 10) ;;
         numa4-pxb-adjacent) hugepages=1; gpus=4; pxb_buses=(5 6 8 9) ;;
+        numa2-plain) gpus=4; plain_nodes=2; hotplug_off=1 ;;
         *) echo "invalid internal topology: $topology" >&2; exit 2 ;;
     esac
-    local nodes=${#pxb_buses[@]}
+    local nodes=$((${#pxb_buses[@]} + plain_nodes))
     ((hugepages && nodes == 0)) && nodes=1
 
     local output="$tmp/$version-$topology"
@@ -58,34 +60,38 @@ run_case() {
             -numa "node,nodeid=$node,cpus=$((node * 8 / nodes))-$(((node + 1) * 8 / nodes - 1)),memdev=mem$node"
             -object "memory-backend-file,id=mem$node,size=$((2048 / nodes))M,mem-path=/dev/hugepages,share=on,prealloc=no"
         )
-        if ((gpus)); then
+        if ((gpus && !plain_nodes)); then
             args+=(-device "pxb-pcie,id=pcie.node$node,bus=pcie.0,addr=$((10 + node)),numa_node=$node,bus_nr=${pxb_buses[$node]}")
         fi
     done
     if ((gpus)); then
         args+=(-object "iommufd,id=iommufd0")
         for ((index = 0; index < gpus; index++)); do
+            local bus=pcie.0
+            ((plain_nodes)) || bus=pcie.node$((index * nodes / gpus))
             args+=(
-                -device "pcie-root-port,id=pci.$index,bus=pcie.node$((index * nodes / gpus)),chassis=$index"
+                -device "pcie-root-port,id=pci.$index,bus=$bus,chassis=$index"
                 -device "vfio-pci,host=00:00.0,bus=pci.$index,iommufd=iommufd0"
             )
         done
     fi
 
+    ((hotplug_off)) && args+=(-global ICH9-LPC.acpi-pci-hotplug-with-bridge-support=off)
     QEMU_ACPI_COMPAT_VER="$version" QEMU_ACPI_DUMP_DIR="$output" \
         "$qemu" "${args[@]}" -machine "$machine"
     local numa=""
     ((nodes > 1)) && numa=$(IFS=,; echo "${pxb_buses[*]}")
+    ((plain_nodes)) && numa=$(printf -- '-,%.0s' $(seq "$plain_nodes"))
     QEMU_ACPI_OUTPUT_DIR="$output/rust" \
-        "$rust_dump" 2 8 "$version" "$gpus" 0 "$hugepages" 1 0 0 0 $((2048 << 20)) 0 0 "$numa"
+        "$rust_dump" 2 8 "$version" "$gpus" 0 "$hugepages" 1 "$hotplug_off" 0 0 $((2048 << 20)) 0 0 "$numa"
     for blob in tables loader rsdp; do
         cmp "$output/$blob.bin" "$output/rust/$blob.bin"
     done
     printf '%-6s %-8s tables+loader+rsdp match\n' "$version" "$topology"
 }
 
 for version in 8.0.0 8.2.0 9.0.0 9.1.0 9.2.0 10.0.0 10.2.0 11.0.0 11.1.0 11.2.0; do
-    for topology in normal numa numa-pxb numa2-pxb numa4-pxb-adjacent; do
+    for topology in normal numa numa-pxb numa2-pxb numa4-pxb-adjacent numa2-plain; do
         run_case "$version" "$topology"
     done
 done

dstack/crates/qemu-acpi/scripts/differential-random.py

@@ -39,7 +39,8 @@ class Case:
     smm: bool = False
     pic: bool = False
     pci_hole64_size: int = 0
-    # PXB bus of each guest NUMA node; empty keeps the layout `hugepages` implies.
+    # PXB bus of each guest NUMA node, None for a node without one; empty
+    # keeps the layout `hugepages` implies.
     pxb_buses: tuple = ()
     # Devices a QEMU wrapper appends to the root bus: "e" endpoint, "r" root port.
     extra: str = ""
@@ -51,7 +52,7 @@ def numa_nodes(self):
     def pxbs(self):
         """PXB bus of each node that has one, explicit or implied."""
         if self.pxb_buses:
-            return self.pxb_buses
+            return tuple(bus for bus in self.pxb_buses if bus is not None)
         return (5,) if self.hugepages and self.gpus else ()
 
 
@@ -106,6 +107,28 @@ def fixed_cases():
             pxb_buses=(5, 7),
             extra="eree",
         ),
+        # Guest NUMA nodes without expanders: SRAT and _PXM only.
+        Case(cpus=2, pxb_buses=(None, None)),
+        Case(cpus=12, memory_mib=6144, nics=2, pxb_buses=(None,) * 3),
+        Case(
+            cpus=192,
+            memory_mib=8192,
+            gpus=4,
+            hotplug_off=True,
+            pci_hole64_size=1 << 43,
+            pxb_buses=(None, None),
+            extra="eeeee",
+        ),
+        Case(
+            version="9.2.0",
+            cpus=4,
+            hugepages=True,
+            gpus=2,
+            switches=1,
+            hotplug_off=True,
+            pxb_buses=(None, None),
+            extra="er",
+        ),
     ]
     return cases
 
@@ -114,11 +137,14 @@ def random_case(rng):
     hugepages = rng.choice([False, True])
     gpus = rng.randint(0, 8)
     switches = rng.randint(0, 4) if gpus else 0
-    nodes = rng.choice([0, 0, 2, 3, 4]) if hugepages else 0
+    nodes = rng.choice([0, 0, 2, 3, 4])
     pxb_buses = []
-    for _ in range(nodes):
-        pxb_buses.append(pxb_buses[-1] + rng.randint(1, 6) if pxb_buses else 5)
-    pxbs = len(pxb_buses) or int(hugepages and gpus > 0)
+    if hugepages:
+        for _ in range(nodes):
+            pxb_buses.append(pxb_buses[-1] + rng.randint(1, 6) if pxb_buses else 5)
+    else:
+        pxb_buses = [None] * nodes
+    pxbs = sum(bus is not None for bus in pxb_buses) or int(hugepages and gpus > 0)
     pxb = pxbs > 0
     available = 26 - pxbs
     passthrough = switches if pxb else gpus + switches
@@ -365,7 +391,7 @@ def run_case(case, image, rust_dump, output):
             str(case.memory_mib * 1024 * 1024),
             str(case.volumes),
             str(int(case.pic)),
-            ",".join(map(str, case.pxb_buses)),
+            ",".join("-" if bus is None else str(bus) for bus in case.pxb_buses),
             case.extra,
         ],
         check=True,

dstack/crates/qemu-acpi/src/golden_tests.rs

@@ -89,6 +89,31 @@ mod tests {
         Ok(())
     }
 
+    /// Two nodes without expanders: GPUs and a wrapper's endpoint stay on
+    /// the root bus, and only the SRAT and the CPUs' `_PXM` carry the nodes.
+    #[test]
+    fn numa_nodes_without_pxb_match_qemu_byte_for_byte() -> Result<(), Error> {
+        let mut c = config(1, 0);
+        c.cpu_count = 8;
+        c.numa_nodes = vec![NumaNode { pxb_bus: None }; 2];
+        c.num_gpus = 2;
+        c.extra_root_devices = vec![RootBusDevice::Endpoint];
+        c.hotplug_off = true;
+        let actual = build(&c)?;
+        let expected = include_bytes!("../fixtures/qemu-11.1-q35-hotplug-off-numa2-plain-base.bin");
+        assert_eq!(&actual.tables[..expected.len()], expected);
+        assert!(actual.tables[expected.len()..].iter().all(|b| *b == 0));
+        assert_eq!(
+            actual.loader,
+            include_bytes!("../fixtures/qemu-11.1-q35-hotplug-off-numa2-plain-loader.bin")
+        );
+        assert_eq!(
+            actual.rsdp,
+            include_bytes!("../fixtures/qemu-11.1-q35-hotplug-off-numa2-plain-rsdp.bin")
+        );
+        Ok(())
+    }
+
     #[test]
     fn numa_layouts_qemu_cannot_build_are_rejected() {
         let nodes = |buses: &[u8]| {
@@ -107,6 +132,7 @@ mod tests {
                 2,
                 vec![NumaNode { pxb_bus: None }, NumaNode { pxb_bus: Some(5) }],
             ),
+            (3, vec![NumaNode { pxb_bus: None }; 2]),
         ] {
             c.cpu_count = cpus;
             c.numa_nodes = buses;

dstack/crates/qemu-acpi/src/topology.rs

@@ -89,8 +89,8 @@ pub enum TopologyError {
     },
     #[error("{0} NUMA nodes exceed QEMU's limit of 128")]
     TooManyNumaNodes(usize),
-    #[error("every explicit NUMA node needs a PXB expander")]
-    NumaNodeWithoutPxb,
+    #[error("NUMA nodes must either all have a PXB expander or none")]
+    MixedPxbNumaNodes,
     #[error("{0} PXB expanders do not fit on the root bus")]
     TooManyPxbs(usize),
     #[error("PXB bus numbers must be nonzero and distinct: {0:?}")]
@@ -125,8 +125,13 @@ impl MachineConfig {
             return Err(TopologyError::TooManyNumaNodes(nodes));
         }
         // Node `i`'s expander sits on slot 0x10 + i only if every node has one.
-        if self.numa_nodes.iter().any(|node| node.pxb_bus.is_none()) {
-            return Err(TopologyError::NumaNodeWithoutPxb);
+        let pxbs = self
+            .numa_nodes
+            .iter()
+            .filter(|n| n.pxb_bus.is_some())
+            .count();
+        if pxbs != 0 && pxbs != nodes {
+            return Err(TopologyError::MixedPxbNumaNodes);
         }
         for (what, count) in [
             ("cpu_count", u64::from(self.cpu_count)),

dstack/dstack-mr/cli/src/main.rs

@@ -66,6 +66,11 @@ struct MachineConfig {
     #[arg(long, value_delimiter = ',')]
     numa_pxb_buses: Vec<u8>,
 
+    /// Number of guest NUMA nodes without PXB expanders (every device on the
+    /// root bus)
+    #[arg(long, conflicts_with = "numa_pxb_buses")]
+    numa_nodes: Option<usize>,
+
     /// Number of GPUs
     #[arg(long, default_value = "0")]
     num_gpus: u32,
@@ -132,11 +137,14 @@ fn main() -> Result<()> {
             // The image declares its OVMF layout. Older metadata.json files
             // predate the field, so fall back to the only layout that existed.
             let ovmf_variant = image_info.ovmf_variant.unwrap_or_default();
-            let numa_nodes: Vec<_> = config
-                .numa_pxb_buses
-                .iter()
-                .map(|&bus| dstack_types::NumaNodeConfig { pxb_bus: bus })
-                .collect();
+            let numa_nodes = match config.numa_nodes {
+                Some(nodes) => vec![dstack_types::NumaNodeConfig { pxb_bus: None }; nodes],
+                None => config
+                    .numa_pxb_buses
+                    .iter()
+                    .map(|&bus| dstack_types::NumaNodeConfig { pxb_bus: Some(bus) })
+                    .collect(),
+            };
 
             let machine = Machine::builder()
                 .cpu_count(config.cpu)

dstack/dstack-mr/src/acpi.rs

@@ -40,7 +40,7 @@ impl Machine<'_> {
                 .numa_nodes
                 .iter()
                 .map(|node| NumaNode {
-                    pxb_bus: Some(node.pxb_bus),
+                    pxb_bus: node.pxb_bus,
                 })
                 .collect(),
             num_gpus: self.num_gpus,

dstack/dstack-types/src/lib.rs

@@ -1427,8 +1427,10 @@ fn is_false(value: &bool) -> bool {
 /// node order.
 #[derive(Deserialize, Serialize, Debug, Clone, Copy, PartialEq, Eq)]
 pub struct NumaNodeConfig {
-    /// Bus number of the node's `pxb-pcie` expander (its `bus_nr`).
-    pub pxb_bus: u8,
+    /// Bus number of the node's `pxb-pcie` expander (its `bus_nr`). Either
+    /// every node has one or none does, leaving all devices on the root bus.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub pxb_bus: Option<u8>,
 }
 
 /// A device a host QEMU wrapper appends to the root PCIe bus.
@@ -2656,6 +2658,26 @@ mod vm_config_device_count_tests {
         );
     }
 
+    #[test]
+    fn numa_nodes_keep_their_legacy_serialization() {
+        let mut json = legacy_json();
+        json["numa_nodes"] = serde_json::json!([{ "pxb_bus": 5 }, { "pxb_bus": 10 }]);
+        let cfg: VmConfig = serde_json::from_value(json.clone()).unwrap();
+        assert_eq!(cfg.numa_nodes[1].pxb_bus, Some(10));
+        assert_eq!(
+            serde_json::to_value(&cfg).unwrap()["numa_nodes"],
+            json["numa_nodes"]
+        );
+
+        json["numa_nodes"] = serde_json::json!([{}, {}]);
+        let cfg: VmConfig = serde_json::from_value(json.clone()).unwrap();
+        assert_eq!(cfg.numa_nodes[0].pxb_bus, None);
+        assert_eq!(
+            serde_json::to_value(&cfg).unwrap()["numa_nodes"],
+            json["numa_nodes"]
+        );
+    }
+
     #[test]
     fn swtpm_is_serialized_only_when_enabled() {
         let mut cfg: VmConfig = serde_json::from_value(legacy_json()).unwrap();

dstack/vmm/src/app.rs

@@ -2213,7 +2213,7 @@ fn make_vm_config(
     } else {
         pxb_buses
             .into_iter()
-            .map(|pxb_bus| dstack_types::NumaNodeConfig { pxb_bus })
+            .map(|bus| dstack_types::NumaNodeConfig { pxb_bus: Some(bus) })
             .collect()
     };
     // Each resolved network interface becomes one virtio-net-pci device in the