Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions dstack/vmm/src/main_service.rs
Original file line number Diff line number Diff line change
Expand Up @@ -110,9 +110,30 @@ pub fn resolve_gpus_with_config(
if !cvm_config.gpu.allow_attach_all && gpus.attach_mode.is_all() {
bail!("Attaching all GPUs is not allowed");
}
if !gpus.gpus.is_empty() && !gpus.attach_mode.is_all() {
let offered: Vec<_> = cvm_config
.gpu
.list_devices()?
.into_iter()
.map(|dev| dev.slot)
.collect();
ensure_gpus_offered(&gpus.gpus, &offered)?;
}
Ok(gpus)
}

/// Listed slots must be among the GPUs `ListGpus` offers, i.e. pass
/// `cvm.gpu.listing`, `include` and `exclude`; otherwise any host PCI device
/// (or a string carrying QEMU option separators) reaches `-device vfio-pci`.
fn ensure_gpus_offered(requested: &[GpuSpec], offered: &[String]) -> Result<()> {
for gpu in requested {
if !offered.contains(&gpu.slot) {
bail!("GPU {} is not offered by this node", gpu.slot);
}
}
Ok(())
}

pub fn resolve_gpus(gpu_cfg: &rpc::GpuConfig) -> Result<GpuConfig> {
// Check the attach mode to determine how to handle GPUs
match gpu_cfg.attach_mode.as_str() {
Expand Down Expand Up @@ -2826,6 +2847,16 @@ mod tests {
Ok(())
}

#[test]
fn listed_gpus_must_be_offered_by_node() {
let offered = ["0000:0f:00.0".to_string()];
let gpu = |slot: &str| GpuSpec { slot: slot.into() };
assert!(ensure_gpus_offered(&[gpu("0000:0f:00.0")], &offered).is_ok());
for slot in ["0000:10:00.0", "0000:0f:00.0,romfile=/tmp/rom"] {
assert!(ensure_gpus_offered(&[gpu(slot)], &offered).is_err());
}
}

#[test]
fn resolve_volumes_attaches_duplicate_root_once() -> Result<()> {
let tmp = tempfile::tempdir()?;
Expand Down
Loading