fix(gpu): validate exact device requests

Require exact driver GPU device lists to be tied to a GPU request, allow a single exact device to use the default countless request, and require explicit matching counts for multi-device lists.

Signed-off-by: Evan Lezar <elezar@nvidia.com>

Author: elezar

crates/openshell-core/src/gpu.rs

@@ -63,6 +63,47 @@ pub fn cdi_gpu_device_ids(
     }
 }
 
+/// Validate a compute-driver GPU request against driver-owned specific devices.
+///
+/// Drivers call this when a sandbox request combines portable GPU requirements
+/// with exact device identifiers in `driver_config`.
+///
+/// # Errors
+/// Returns an error when the sandbox GPU request is absent or when `gpu.count`
+/// does not equal the number of specific devices. A single exact device is
+/// compatible with the default sandbox GPU request where `gpu.count` is absent.
+pub fn validate_specific_gpu_device_request(
+    gpu: Option<&DriverGpuResourceRequirements>,
+    specific_devices: &[String],
+    driver_config_field: &str,
+) -> Result<(), String> {
+    let device_count = specific_devices.len();
+    if device_count == 0 {
+        return Ok(());
+    }
+
+    let Some(gpu) = gpu else {
+        return Err(format!("{driver_config_field} requires a gpu request"));
+    };
+
+    let Some(count) = gpu.count else {
+        if device_count == 1 {
+            return Ok(());
+        }
+        return Err(format!(
+            "{driver_config_field} requires an explicit gpu count matching its length ({device_count})"
+        ));
+    };
+
+    if usize::try_from(count).ok() != Some(device_count) {
+        return Err(format!(
+            "gpu count ({count}) must match {driver_config_field} length ({device_count})"
+        ));
+    }
+
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -100,4 +141,92 @@ mod tests {
             ])
         );
     }
+
+    #[test]
+    fn validate_specific_gpu_device_request_ignores_empty_devices() {
+        validate_specific_gpu_device_request(None, &[], "driver_config.cdi_devices")
+            .expect("empty exact device lists should not be validated");
+    }
+
+    #[test]
+    fn validate_specific_gpu_device_request_accepts_matching_count() {
+        let gpu = DriverGpuResourceRequirements { count: Some(2) };
+        let specific_devices = vec![
+            "nvidia.com/gpu=0".to_string(),
+            "nvidia.com/gpu=1".to_string(),
+        ];
+
+        validate_specific_gpu_device_request(
+            Some(&gpu),
+            &specific_devices,
+            "driver_config.cdi_devices",
+        )
+        .expect("matching count should be accepted");
+    }
+
+    #[test]
+    fn validate_specific_gpu_device_request_accepts_missing_count_for_one_device() {
+        let gpu = DriverGpuResourceRequirements { count: None };
+        let specific_devices = vec!["nvidia.com/gpu=0".to_string()];
+
+        validate_specific_gpu_device_request(
+            Some(&gpu),
+            &specific_devices,
+            "driver_config.cdi_devices",
+        )
+        .expect("single exact device should be compatible with a default GPU request");
+    }
+
+    #[test]
+    fn validate_specific_gpu_device_request_rejects_missing_gpu_request() {
+        let specific_devices = vec!["nvidia.com/gpu=0".to_string()];
+
+        let err = validate_specific_gpu_device_request(
+            None,
+            &specific_devices,
+            "driver_config.cdi_devices",
+        )
+        .expect_err("missing GPU request should be rejected");
+
+        assert_eq!(err, "driver_config.cdi_devices requires a gpu request");
+    }
+
+    #[test]
+    fn validate_specific_gpu_device_request_rejects_missing_count_for_multiple_devices() {
+        let gpu = DriverGpuResourceRequirements { count: None };
+        let specific_devices = vec![
+            "nvidia.com/gpu=0".to_string(),
+            "nvidia.com/gpu=1".to_string(),
+        ];
+
+        let err = validate_specific_gpu_device_request(
+            Some(&gpu),
+            &specific_devices,
+            "driver_config.cdi_devices",
+        )
+        .expect_err("missing count should be rejected for multiple devices");
+
+        assert_eq!(
+            err,
+            "driver_config.cdi_devices requires an explicit gpu count matching its length (2)"
+        );
+    }
+
+    #[test]
+    fn validate_specific_gpu_device_request_rejects_mismatch() {
+        let gpu = DriverGpuResourceRequirements { count: Some(2) };
+        let specific_devices = vec!["nvidia.com/gpu=0".to_string()];
+
+        let err = validate_specific_gpu_device_request(
+            Some(&gpu),
+            &specific_devices,
+            "driver_config.cdi_devices",
+        )
+        .expect_err("mismatched count should be rejected");
+
+        assert_eq!(
+            err,
+            "gpu count (2) must match driver_config.cdi_devices length (1)"
+        );
+    }
 }

crates/openshell-driver-docker/README.md

@@ -32,7 +32,7 @@ contract:
 | `apparmor=unconfined` | Avoids Docker's default profile blocking required mount operations. |
 | `restart_policy = unless-stopped` | Keeps managed sandboxes resumable across daemon or gateway restarts. |
 | `PidsLimit` | Enforces the sandbox PID budget at the Docker cgroup layer. Set `[openshell.drivers.docker].sandbox_pids_limit = 0` to inherit the Docker/runtime default. |
-| CDI GPU request | Uses `driver_config.cdi_devices` when set; otherwise requests all NVIDIA GPUs when the sandbox spec asks for GPU support and daemon CDI support is detected. Counted GPU requests are rejected. |
+| CDI GPU request | Uses `driver_config.cdi_devices` when set; otherwise requests all NVIDIA GPUs when the sandbox spec asks for GPU support and daemon CDI support is detected. Count-only GPU requests are rejected; exact CDI device lists with more than one entry require an explicit GPU count matching the device list length. |
 
 The agent child process does not retain these supervisor privileges.
 

crates/openshell-driver-docker/src/lib.rs

@@ -27,7 +27,9 @@ use openshell_core::driver_utils::{
     LABEL_MANAGED_BY, LABEL_MANAGED_BY_VALUE, LABEL_SANDBOX_ID, LABEL_SANDBOX_NAME,
     LABEL_SANDBOX_NAMESPACE, SUPERVISOR_IMAGE_BINARY_PATH, supervisor_image_should_refresh,
 };
-use openshell_core::gpu::{cdi_gpu_device_ids, driver_gpu_requirements};
+use openshell_core::gpu::{
+    cdi_gpu_device_ids, driver_gpu_requirements, validate_specific_gpu_device_request,
+};
 use openshell_core::progress::{
     PROGRESS_STEP_PULLING_IMAGE, PROGRESS_STEP_REQUESTING_SANDBOX, PROGRESS_STEP_STARTING_SANDBOX,
     format_bytes, mark_progress_active, mark_progress_complete, mark_progress_detail,
@@ -514,23 +516,25 @@ impl DockerComputeDriver {
         supports_gpu: bool,
         driver_config: &DockerSandboxDriverConfig,
     ) -> Result<(), Status> {
-        if gpu_requirements.is_none() && driver_config.cdi_devices.is_some() {
-            return Err(Status::invalid_argument(
-                "driver_config.cdi_devices requires gpu=true",
+        if gpu_requirements.is_some() && !supports_gpu {
+            return Err(Status::failed_precondition(
+                "docker GPU sandboxes require Docker CDI support. Enable CDI on the Docker daemon, then restart the OpenShell gateway/server so GPU capability is detected.",
             ));
         }
 
-        if gpu_requirements.and_then(|gpu| gpu.count).is_some() {
+        if let Some(cdi_devices) = driver_config.cdi_devices.as_deref() {
+            validate_specific_gpu_device_request(
+                gpu_requirements,
+                cdi_devices,
+                "driver_config.cdi_devices",
+            )
+            .map_err(Status::invalid_argument)?;
+        } else if gpu_requirements.and_then(|gpu| gpu.count).is_some() {
             return Err(Status::invalid_argument(
                 "docker GPU count requests are not supported; use --gpu without a count or driver_config.cdi_devices",
             ));
         }
 
-        if gpu_requirements.is_some() && !supports_gpu {
-            return Err(Status::failed_precondition(
-                "docker GPU sandboxes require Docker CDI support. Enable CDI on the Docker daemon, then restart the OpenShell gateway/server so GPU capability is detected.",
-            ));
-        }
         Ok(())
     }
 
@@ -2129,11 +2133,12 @@ fn build_device_requests(sandbox: &DriverSandbox) -> Result<Option<Vec<DeviceReq
         .cdi_devices
         .unwrap_or_default();
     let gpu_requirements = driver_gpu_requirements(spec.resource_requirements.as_ref());
-    if gpu_requirements.is_none() && !cdi_devices.is_empty() {
-        return Err(Status::invalid_argument(
-            "driver_config.cdi_devices requires gpu=true",
-        ));
-    }
+    validate_specific_gpu_device_request(
+        gpu_requirements,
+        &cdi_devices,
+        "driver_config.cdi_devices",
+    )
+    .map_err(Status::invalid_argument)?;
 
     Ok(
         cdi_gpu_device_ids(gpu_requirements, &cdi_devices).map(|device_ids| {

crates/openshell-driver-docker/src/tests.rs

@@ -1022,6 +1022,20 @@ fn validate_sandbox_rejects_gpu_when_cdi_unavailable() {
     assert!(err.message().contains("Docker CDI"));
 }
 
+#[test]
+fn validate_sandbox_rejects_missing_gpu_support_before_request_shape() {
+    let config = runtime_config();
+    let mut sandbox = test_sandbox();
+    let spec = sandbox.spec.as_mut().unwrap();
+    spec.resource_requirements = Some(gpu_resources(Some(2)));
+    spec.template.as_mut().unwrap().driver_config = Some(cdi_devices_config(&["nvidia.com/gpu=0"]));
+
+    let err = DockerComputeDriver::validate_sandbox(&sandbox, &config).unwrap_err();
+
+    assert_eq!(err.code(), tonic::Code::FailedPrecondition);
+    assert!(err.message().contains("Docker CDI"));
+}
+
 #[test]
 fn validate_sandbox_rejects_invalid_cdi_devices_before_gpu_capability() {
     let config = runtime_config();
@@ -1068,6 +1082,94 @@ fn validate_sandbox_rejects_gpu_count_request() {
     );
 }
 
+#[test]
+fn validate_sandbox_accepts_gpu_count_matching_cdi_devices() {
+    let mut config = runtime_config();
+    config.supports_gpu = true;
+    let mut sandbox = test_sandbox();
+    let spec = sandbox.spec.as_mut().unwrap();
+    spec.resource_requirements = Some(gpu_resources(Some(2)));
+    spec.template.as_mut().unwrap().driver_config = Some(cdi_devices_config(&[
+        "nvidia.com/gpu=0",
+        "nvidia.com/gpu=1",
+    ]));
+
+    DockerComputeDriver::validate_sandbox(&sandbox, &config)
+        .expect("matching explicit CDI device count should be accepted");
+}
+
+#[test]
+fn validate_sandbox_accepts_single_cdi_device_without_gpu_count() {
+    let mut config = runtime_config();
+    config.supports_gpu = true;
+    let mut sandbox = test_sandbox();
+    let spec = sandbox.spec.as_mut().unwrap();
+    spec.resource_requirements = Some(gpu_resources(None));
+    spec.template.as_mut().unwrap().driver_config = Some(cdi_devices_config(&["nvidia.com/gpu=0"]));
+
+    DockerComputeDriver::validate_sandbox(&sandbox, &config)
+        .expect("single exact CDI device should be compatible with a default GPU request");
+}
+
+#[test]
+fn validate_sandbox_rejects_multiple_cdi_devices_without_gpu_count() {
+    let mut config = runtime_config();
+    config.supports_gpu = true;
+    let mut sandbox = test_sandbox();
+    let spec = sandbox.spec.as_mut().unwrap();
+    spec.resource_requirements = Some(gpu_resources(None));
+    spec.template.as_mut().unwrap().driver_config = Some(cdi_devices_config(&[
+        "nvidia.com/gpu=0",
+        "nvidia.com/gpu=1",
+    ]));
+
+    let err = DockerComputeDriver::validate_sandbox(&sandbox, &config).unwrap_err();
+
+    assert_eq!(err.code(), tonic::Code::InvalidArgument);
+    assert!(
+        err.message()
+            .contains("requires an explicit gpu count matching its length (2)")
+    );
+}
+
+#[test]
+fn validate_sandbox_rejects_cdi_devices_without_gpu_request() {
+    let mut config = runtime_config();
+    config.supports_gpu = true;
+    let mut sandbox = test_sandbox();
+    sandbox
+        .spec
+        .as_mut()
+        .unwrap()
+        .template
+        .as_mut()
+        .unwrap()
+        .driver_config = Some(cdi_devices_config(&["nvidia.com/gpu=0"]));
+
+    let err = DockerComputeDriver::validate_sandbox(&sandbox, &config).unwrap_err();
+
+    assert_eq!(err.code(), tonic::Code::InvalidArgument);
+    assert!(err.message().contains("requires a gpu request"));
+}
+
+#[test]
+fn validate_sandbox_rejects_gpu_count_mismatched_cdi_devices() {
+    let mut config = runtime_config();
+    config.supports_gpu = true;
+    let mut sandbox = test_sandbox();
+    let spec = sandbox.spec.as_mut().unwrap();
+    spec.resource_requirements = Some(gpu_resources(Some(2)));
+    spec.template.as_mut().unwrap().driver_config = Some(cdi_devices_config(&["nvidia.com/gpu=0"]));
+
+    let err = DockerComputeDriver::validate_sandbox(&sandbox, &config).unwrap_err();
+
+    assert_eq!(err.code(), tonic::Code::InvalidArgument);
+    assert!(
+        err.message()
+            .contains("gpu count (2) must match driver_config.cdi_devices length (1)")
+    );
+}
+
 #[test]
 fn validate_sandbox_rejects_template_errors_before_device_config() {
     let config = runtime_config();
@@ -1153,7 +1255,25 @@ fn build_container_create_body_passes_explicit_cdi_device_id_through() {
 }
 
 #[test]
-fn build_container_create_body_rejects_cdi_devices_without_gpu() {
+fn build_container_create_body_rejects_gpu_count_mismatched_cdi_devices() {
+    let mut config = runtime_config();
+    config.supports_gpu = true;
+    let mut sandbox = test_sandbox();
+    let spec = sandbox.spec.as_mut().unwrap();
+    spec.resource_requirements = Some(gpu_resources(Some(2)));
+    spec.template.as_mut().unwrap().driver_config = Some(cdi_devices_config(&["nvidia.com/gpu=0"]));
+
+    let err = build_container_create_body(&sandbox, &config).unwrap_err();
+
+    assert_eq!(err.code(), tonic::Code::InvalidArgument);
+    assert!(
+        err.message()
+            .contains("gpu count (2) must match driver_config.cdi_devices length (1)")
+    );
+}
+
+#[test]
+fn build_container_create_body_rejects_cdi_devices_without_gpu_request() {
     let mut sandbox = test_sandbox();
     sandbox
         .spec
@@ -1166,7 +1286,7 @@ fn build_container_create_body_rejects_cdi_devices_without_gpu() {
 
     let err = build_container_create_body(&sandbox, &runtime_config()).unwrap_err();
     assert_eq!(err.code(), tonic::Code::InvalidArgument);
-    assert!(err.message().contains("requires gpu=true"));
+    assert!(err.message().contains("requires a gpu request"));
 }
 
 #[test]

crates/openshell-driver-podman/README.md

@@ -46,7 +46,7 @@ The container spec in `container.rs` sets these security-critical fields:
 | `no_new_privileges` | `true` | Prevents privilege escalation after exec. |
 | `seccomp_profile_path` | `unconfined` | The supervisor installs its own policy-aware BPF filter. A container-level profile can block Landlock/seccomp syscalls during setup. |
 | `mounts` | Private tmpfs at `/run/netns` | Lets the supervisor create named network namespaces in rootless Podman. |
-| CDI GPU devices | `driver_config.cdi_devices` when set, otherwise all NVIDIA GPUs | Exposes requested GPUs to GPU-enabled sandbox containers. Counted GPU requests are rejected. |
+| CDI GPU devices | `driver_config.cdi_devices` when set, otherwise all NVIDIA GPUs | Exposes requested GPUs to GPU-enabled sandbox containers. Count-only GPU requests are rejected; exact CDI device lists with more than one entry require an explicit GPU count matching the device list length. |
 
 The restricted agent child does not retain these supervisor privileges.