feat(gpu): route device selection through driver config

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

Author: elezar

Cargo.lock

@@ -42,7 +42,11 @@ but currently ignores them.
 
 GPU requests enter the driver layer through
 `SandboxSpec.resource_requirements.gpu`. The compact interim shape supports a
-default GPU request, GPU count, and driver-specific device IDs.
+default GPU request and GPU count. Exact driver-native device selection is
+passed through the selected runtime's `driver_config` block; the gateway
+selects that block but does not interpret the nested driver schema. Drivers
+that support exact selection validate that the unique `gpu_device_ids` entry
+count matches the portable GPU count.
 
 VM runtime state paths are derived only from driver-validated sandbox IDs
 matching `[A-Za-z0-9._-]{1,128}`. The gateway-owned VM driver socket uses a
@@ -81,9 +85,7 @@ users.
 Custom sandbox images must include the agent runtime and any system
 dependencies, but they should not need to include the gateway. GPU-capable
 images must include the user-space libraries required by the workload. The
-runtime still owns GPU device injection. GPU requests can include explicit
-driver-native device IDs or a requested count; the gateway validates the public
-request shape and each runtime enforces the GPU allocation modes it supports.
+runtime still owns GPU device injection.
 
 ## Deployment Shape
 

architecture/compute-runtimes.md

@@ -1734,8 +1734,10 @@ pub async fn sandbox_create(
         }
         None => None,
     };
-    let requested_gpu =
-        gpu || gpu_count.is_some() || image.as_deref().is_some_and(image_requests_gpu);
+    let requested_gpu = gpu
+        || gpu_count.is_some()
+        || gpu_device.is_some_and(|device_id| !device_id.trim().is_empty())
+        || image.as_deref().is_some_and(image_requests_gpu);
 
     let providers_v2_enabled = gateway_providers_v2_enabled(&mut client).await?;
     let inferred_types: Vec<String> = if providers_v2_enabled {
@@ -1753,11 +1755,13 @@ pub async fn sandbox_create(
 
     let policy = load_sandbox_policy(policy)?;
     let resource_limits = build_sandbox_resource_limits(cpu, memory)?;
+    let driver_config = gpu_driver_config_from_cli(gpu_device);
 
-    let template = if image.is_some() || resource_limits.is_some() {
+    let template = if image.is_some() || resource_limits.is_some() || driver_config.is_some() {
         Some(SandboxTemplate {
             image: image.unwrap_or_default(),
             resources: resource_limits,
+            driver_config,
             ..SandboxTemplate::default()
         })
     } else {
@@ -1766,11 +1770,7 @@ pub async fn sandbox_create(
 
     let request = CreateSandboxRequest {
         spec: Some(SandboxSpec {
-            resource_requirements: resource_requirements_from_cli(
-                requested_gpu,
-                gpu_device,
-                gpu_count,
-            ),
+            resource_requirements: resource_requirements_from_cli(requested_gpu, gpu_count),
             policy,
             providers: configured_providers,
             template,
@@ -2197,17 +2197,48 @@ pub async fn sandbox_create(
 
 fn resource_requirements_from_cli(
     requested_gpu: bool,
-    gpu_device: Option<&str>,
     gpu_count: Option<u32>,
 ) -> Option<SandboxResourceRequirements> {
-    requested_gpu.then(|| SandboxResourceRequirements {
-        gpu: Some(GpuResourceRequirement {
-            device_ids: gpu_device
-                .filter(|device_id| !device_id.is_empty())
-                .map(|device_id| vec![device_id.to_string()])
-                .unwrap_or_default(),
-            count: gpu_count,
-        }),
+    requested_gpu.then_some(SandboxResourceRequirements {
+        gpu: Some(GpuResourceRequirement { count: gpu_count }),
+    })
+}
+
+fn gpu_driver_config_from_cli(gpu_device: Option<&str>) -> Option<prost_types::Struct> {
+    use prost_types::{ListValue, Struct, Value, value::Kind};
+
+    fn string_value(value: &str) -> Value {
+        Value {
+            kind: Some(Kind::StringValue(value.to_string())),
+        }
+    }
+
+    fn driver_block(gpu_device: &str) -> Value {
+        Value {
+            kind: Some(Kind::StructValue(Struct {
+                fields: std::iter::once((
+                    "gpu_device_ids".to_string(),
+                    Value {
+                        kind: Some(Kind::ListValue(ListValue {
+                            values: vec![string_value(gpu_device)],
+                        })),
+                    },
+                ))
+                .collect(),
+            })),
+        }
+    }
+
+    let gpu_device = gpu_device.filter(|device_id| !device_id.trim().is_empty())?;
+
+    Some(Struct {
+        fields: [
+            ("docker".to_string(), driver_block(gpu_device)),
+            ("podman".to_string(), driver_block(gpu_device)),
+            ("vm".to_string(), driver_block(gpu_device)),
+        ]
+        .into_iter()
+        .collect(),
     })
 }
 
@@ -7460,10 +7491,10 @@ mod tests {
         dockerfile_sources_supported_for_gateway, format_endpoint, format_gateway_select_header,
         format_gateway_select_items, format_provider_attachment_table, gateway_add,
         gateway_auth_label, gateway_env_override_warning, gateway_select_with, gateway_type_label,
-        git_sync_files, http_health_check, image_requests_gpu, import_local_package_mtls_bundle,
-        inferred_provider_type, package_managed_tls_dirs, parse_cli_setting_value,
-        parse_credential_expiry_cli_value, parse_credential_expiry_pairs, parse_credential_pairs,
-        plaintext_gateway_is_remote, progress_step_from_metadata,
+        git_sync_files, gpu_driver_config_from_cli, http_health_check, image_requests_gpu,
+        import_local_package_mtls_bundle, inferred_provider_type, package_managed_tls_dirs,
+        parse_cli_setting_value, parse_credential_expiry_cli_value, parse_credential_expiry_pairs,
+        parse_credential_pairs, plaintext_gateway_is_remote, progress_step_from_metadata,
         provider_profile_allows_refresh_bootstrap, provisioning_timeout_message,
         ready_false_condition_message, refresh_status_header, refresh_status_row, resolve_from,
         resource_requirements_from_cli, sandbox_should_persist, sandbox_upload_plan,
@@ -7948,47 +7979,40 @@ mod tests {
 
     #[test]
     fn resource_requirements_from_cli_uses_presence_for_default_gpu() {
-        let requirements = resource_requirements_from_cli(true, None, None)
+        let requirements = resource_requirements_from_cli(true, None)
             .expect("resource requirements should be present");
         let gpu = requirements.gpu.expect("GPU requirement should be present");
 
-        assert!(gpu.device_ids.is_empty());
         assert_eq!(gpu.count, None);
     }
 
     #[test]
-    fn resource_requirements_from_cli_maps_gpu_device_to_one_device_id() {
-        let requirements = resource_requirements_from_cli(true, Some("0000:2d:00.0"), None)
-            .expect("resource requirements should be present");
-        let gpu = requirements.gpu.expect("GPU requirement should be present");
+    fn gpu_driver_config_from_cli_maps_gpu_device_to_driver_blocks() {
+        let config = gpu_driver_config_from_cli(Some("nvidia.com/gpu=0"))
+            .expect("driver config should be present");
 
-        assert_eq!(gpu.device_ids, vec!["0000:2d:00.0"]);
-        assert_eq!(gpu.count, None);
+        assert!(config.fields.contains_key("docker"));
+        assert!(config.fields.contains_key("podman"));
+        assert!(config.fields.contains_key("vm"));
     }
 
     #[test]
     fn resource_requirements_from_cli_maps_gpu_count() {
         let requirements =
-            resource_requirements_from_cli(true, None, Some(2)).expect("requirements should exist");
+            resource_requirements_from_cli(true, Some(2)).expect("requirements should exist");
         let gpu = requirements.gpu.expect("GPU requirement should be present");
 
-        assert!(gpu.device_ids.is_empty());
         assert_eq!(gpu.count, Some(2));
     }
 
     #[test]
-    fn resource_requirements_from_cli_preserves_device_and_gpu_count_for_gateway_validation() {
-        let requirements = resource_requirements_from_cli(true, Some("nvidia.com/gpu=0"), Some(2))
-            .expect("requirements should exist");
-        let gpu = requirements.gpu.expect("GPU requirement should be present");
-
-        assert_eq!(gpu.device_ids, vec!["nvidia.com/gpu=0"]);
-        assert_eq!(gpu.count, Some(2));
+    fn gpu_driver_config_from_cli_omits_empty_device() {
+        assert!(gpu_driver_config_from_cli(Some("")).is_none());
     }
 
     #[test]
     fn resource_requirements_from_cli_omits_gpu_request_when_not_requested() {
-        assert!(resource_requirements_from_cli(false, Some("0"), None).is_none());
+        assert!(resource_requirements_from_cli(false, None).is_none());
     }
 
     #[test]

crates/openshell-cli/src/run.rs

@@ -922,15 +922,23 @@ async fn sandbox_create_sends_gpu_count_request() {
     .expect("sandbox create should succeed");
 
     let requests = create_requests(&server).await;
-    let gpu = requests[0]
+    let spec = requests[0]
         .spec
         .as_ref()
-        .and_then(|spec| spec.resource_requirements.as_ref())
+        .expect("sandbox spec should be sent");
+    let gpu = spec
+        .resource_requirements
+        .as_ref()
         .and_then(|requirements| requirements.gpu.as_ref())
         .expect("GPU request should be sent");
 
-    assert!(gpu.device_ids.is_empty());
     assert_eq!(gpu.count, Some(2));
+    assert!(
+        spec.template
+            .as_ref()
+            .and_then(|template| template.driver_config.as_ref())
+            .is_none()
+    );
 }
 
 #[tokio::test]