refactor(gpu): use resource requirements for GPU requests

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

Author: elezar

architecture/compute-runtimes.md

@@ -45,6 +45,9 @@ through the driver configuration. The Helm chart defaults sandbox agents to
 `Unconfined` so runtime/default AppArmor profiles do not block supervisor
 network namespace setup on AppArmor-enabled nodes.
 
+GPU requests enter the driver layer through
+`SandboxSpec.resource_requirements.gpu`.
+
 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
 private `run/` directory plus Unix peer UID/PID checks. Standalone

crates/openshell-cli/src/run.rs

@@ -39,18 +39,18 @@ use openshell_core::proto::{
     GetClusterInferenceRequest, GetDraftHistoryRequest, GetDraftPolicyRequest,
     GetGatewayConfigRequest, GetProviderProfileRequest, GetProviderRefreshStatusRequest,
     GetProviderRequest, GetSandboxConfigRequest, GetSandboxLogsRequest,
-    GetSandboxPolicyStatusRequest, GetSandboxRequest, GetServiceRequest, GpuRequestSpec,
+    GetSandboxPolicyStatusRequest, GetSandboxRequest, GetServiceRequest, GpuResourceRequirement,
     HealthRequest, ImportProviderProfilesRequest, LintProviderProfilesRequest,
     ListProviderProfilesRequest, ListProvidersRequest, ListSandboxPoliciesRequest,
     ListSandboxProvidersRequest, ListSandboxesRequest, ListServicesRequest, PlatformEvent,
     PolicySource, PolicyStatus, Provider, ProviderCredentialRefreshStatus,
     ProviderCredentialRefreshStrategy, ProviderProfile, ProviderProfileDiagnostic,
     ProviderProfileImportItem, RejectDraftChunkRequest, RevokeSshSessionRequest,
-    RotateProviderCredentialRequest, Sandbox, SandboxPhase, SandboxPolicy, SandboxSpec,
-    SandboxTemplate, ServiceEndpointResponse, SetClusterInferenceRequest, SettingScope,
-    SettingValue, TcpForwardFrame, TcpForwardInit, TcpRelayTarget, UpdateConfigRequest,
-    UpdateProviderRequest, WatchSandboxRequest, exec_sandbox_event, setting_value,
-    tcp_forward_init,
+    RotateProviderCredentialRequest, Sandbox, SandboxPhase, SandboxPolicy,
+    SandboxResourceRequirements, SandboxSpec, SandboxTemplate, ServiceEndpointResponse,
+    SetClusterInferenceRequest, SettingScope, SettingValue, TcpForwardFrame, TcpForwardInit,
+    TcpRelayTarget, UpdateConfigRequest, UpdateProviderRequest, WatchSandboxRequest,
+    exec_sandbox_event, setting_value, tcp_forward_init,
 };
 use openshell_core::settings::{self, SettingValueKind};
 use openshell_core::{ObjectId, ObjectName};
@@ -1754,11 +1754,6 @@ pub async fn sandbox_create(
         }
         None => None,
     };
-    let requested_gpu = gpu
-        || gpu_device.is_some_and(|device_id| !device_id.is_empty())
-        || gpu_count.is_some()
-        || 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 {
         Vec::new()
@@ -1775,6 +1770,11 @@ pub async fn sandbox_create(
 
     let policy = load_sandbox_policy(policy)?;
     let resource_limits = build_sandbox_resource_limits(cpu, memory)?;
+    let resource_requirements =
+        resource_requirements_from_cli(image.as_deref(), gpu, gpu_device, gpu_count);
+    let requested_gpu = resource_requirements
+        .as_ref()
+        .is_some_and(|requirements| requirements.gpu.is_some());
 
     let template = if image.is_some() || resource_limits.is_some() {
         Some(SandboxTemplate {
@@ -1788,7 +1788,7 @@ pub async fn sandbox_create(
 
     let request = CreateSandboxRequest {
         spec: Some(SandboxSpec {
-            gpu: gpu_request_from_cli(requested_gpu, gpu_device, gpu_count),
+            resource_requirements,
             policy,
             providers: configured_providers,
             template,
@@ -2223,17 +2223,26 @@ pub async fn sandbox_create(
     }
 }
 
-fn gpu_request_from_cli(
-    requested_gpu: bool,
+fn resource_requirements_from_cli(
+    image: Option<&str>,
+    gpu: bool,
     gpu_device: Option<&str>,
     gpu_count: Option<u32>,
-) -> Option<GpuRequestSpec> {
-    requested_gpu.then(|| GpuRequestSpec {
-        device_id: gpu_device
-            .filter(|device_id| !device_id.is_empty())
-            .map(|device_id| vec![device_id.to_string()])
-            .unwrap_or_default(),
-        count: gpu_count,
+) -> Option<SandboxResourceRequirements> {
+    let device_ids = gpu_device
+        .filter(|device_id| !device_id.is_empty())
+        .map(|device_id| vec![device_id.to_string()])
+        .unwrap_or_default();
+    let requested_gpu = gpu
+        || gpu_count.is_some()
+        || !device_ids.is_empty()
+        || image.is_some_and(image_requests_gpu);
+
+    requested_gpu.then_some(SandboxResourceRequirements {
+        gpu: Some(GpuResourceRequirement {
+            device_ids,
+            count: gpu_count,
+        }),
     })
 }
 
@@ -7486,15 +7495,14 @@ 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, gpu_request_from_cli, http_health_check, image_requests_gpu,
-        import_local_package_mtls_bundle, inferred_provider_type, mtls_certs_exist_for_gateway,
-        package_managed_tls_dirs,
+        git_sync_files, http_health_check, image_requests_gpu, import_local_package_mtls_bundle,
+        inferred_provider_type, mtls_certs_exist_for_gateway, 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,
-        sandbox_should_persist, sandbox_upload_plan, service_expose_status_error,
-        service_url_for_gateway,
+        resource_requirements_from_cli, sandbox_should_persist, sandbox_upload_plan,
+        service_expose_status_error, service_url_for_gateway,
     };
     use crate::TEST_ENV_LOCK;
     use hyper::StatusCode;
@@ -7974,43 +7982,64 @@ mod tests {
     }
 
     #[test]
-    fn gpu_request_from_cli_uses_presence_with_empty_device_ids_for_default_gpu() {
-        let request =
-            gpu_request_from_cli(true, None, None).expect("gpu request should be present");
+    fn resource_requirements_from_cli_uses_presence_for_default_gpu() {
+        let requirements = resource_requirements_from_cli(None, true, None, None)
+            .expect("resource requirements should be present");
+        let gpu = requirements.gpu.expect("GPU requirement should be present");
 
-        assert!(request.device_id.is_empty());
-        assert_eq!(request.count, None);
+        assert!(gpu.device_ids.is_empty());
+        assert_eq!(gpu.count, None);
     }
 
     #[test]
-    fn gpu_request_from_cli_maps_gpu_device_to_one_device_id() {
-        let request = gpu_request_from_cli(true, Some("0000:2d:00.0"), None)
-            .expect("gpu request should be present");
+    fn resource_requirements_from_cli_maps_gpu_device_to_one_device_id() {
+        let requirements = resource_requirements_from_cli(None, false, Some("0000:2d:00.0"), None)
+            .expect("resource requirements should be present");
+        let gpu = requirements.gpu.expect("GPU requirement should be present");
 
-        assert_eq!(request.device_id, vec!["0000:2d:00.0"]);
-        assert_eq!(request.count, None);
+        assert_eq!(gpu.device_ids, vec!["0000:2d:00.0"]);
+        assert_eq!(gpu.count, None);
     }
 
     #[test]
-    fn gpu_request_from_cli_maps_gpu_count() {
-        let request = gpu_request_from_cli(true, None, Some(2)).expect("gpu request should exist");
+    fn resource_requirements_from_cli_maps_gpu_count() {
+        let requirements = resource_requirements_from_cli(None, false, None, Some(2))
+            .expect("requirements should exist");
+        let gpu = requirements.gpu.expect("GPU requirement should be present");
 
-        assert!(request.device_id.is_empty());
-        assert_eq!(request.count, Some(2));
+        assert!(gpu.device_ids.is_empty());
+        assert_eq!(gpu.count, Some(2));
     }
 
     #[test]
-    fn gpu_request_from_cli_preserves_device_and_gpu_count_for_gateway_validation() {
-        let request = gpu_request_from_cli(true, Some("nvidia.com/gpu=0"), Some(2))
-            .expect("gpu request should exist");
+    fn resource_requirements_from_cli_preserves_device_and_gpu_count_for_gateway_validation() {
+        let requirements =
+            resource_requirements_from_cli(None, false, Some("nvidia.com/gpu=0"), Some(2))
+                .expect("requirements should exist");
+        let gpu = requirements.gpu.expect("GPU requirement should be present");
 
-        assert_eq!(request.device_id, vec!["nvidia.com/gpu=0"]);
-        assert_eq!(request.count, Some(2));
+        assert_eq!(gpu.device_ids, vec!["nvidia.com/gpu=0"]);
+        assert_eq!(gpu.count, Some(2));
     }
 
     #[test]
-    fn gpu_request_from_cli_omits_gpu_request_when_not_requested() {
-        assert!(gpu_request_from_cli(false, Some("0"), None).is_none());
+    fn resource_requirements_from_cli_omits_gpu_request_when_not_requested() {
+        assert!(resource_requirements_from_cli(None, false, None, None).is_none());
+    }
+
+    #[test]
+    fn resource_requirements_from_cli_infers_gpu_from_image() {
+        let requirements = resource_requirements_from_cli(
+            Some("ghcr.io/nvidia/openshell-community/sandboxes/nvidia-gpu:latest"),
+            false,
+            None,
+            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]

crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs

@@ -907,6 +907,7 @@ async fn sandbox_create_sends_gpu_device_request_without_gpu_flag() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -924,10 +925,11 @@ async fn sandbox_create_sends_gpu_device_request_without_gpu_flag() {
     let gpu = requests[0]
         .spec
         .as_ref()
-        .and_then(|spec| spec.gpu.as_ref())
+        .and_then(|spec| spec.resource_requirements.as_ref())
+        .and_then(|requirements| requirements.gpu.as_ref())
         .expect("GPU request should be sent");
 
-    assert_eq!(gpu.device_id, vec!["nvidia.com/gpu=0"]);
+    assert_eq!(gpu.device_ids, vec!["nvidia.com/gpu=0"]);
     assert_eq!(gpu.count, None);
 }
 
@@ -970,10 +972,11 @@ async fn sandbox_create_sends_gpu_count_request() {
     let gpu = requests[0]
         .spec
         .as_ref()
-        .and_then(|spec| spec.gpu.as_ref())
+        .and_then(|spec| spec.resource_requirements.as_ref())
+        .and_then(|requirements| requirements.gpu.as_ref())
         .expect("GPU request should be sent");
 
-    assert!(gpu.device_id.is_empty());
+    assert!(gpu.device_ids.is_empty());
     assert_eq!(gpu.count, Some(2));
 }
 

crates/openshell-core/src/gpu.rs

@@ -4,23 +4,26 @@
 //! Shared GPU request helpers.
 
 use crate::config::CDI_GPU_DEVICE_ALL;
-use crate::proto::compute::v1::{DriverSandboxSpec, GpuRequestSpec};
+use crate::proto::compute::v1::{DriverGpuResourceRequirement, DriverSandboxSpec};
 
-/// Extract the driver GPU request from a sandbox spec.
+/// Extract the driver GPU requirement from a sandbox spec.
 #[must_use]
-pub fn driver_gpu_request(spec: &DriverSandboxSpec) -> Option<&GpuRequestSpec> {
-    spec.gpu.as_ref()
+pub fn driver_gpu_requirement(spec: &DriverSandboxSpec) -> Option<&DriverGpuResourceRequirement> {
+    spec.resource_requirements
+        .as_ref()
+        .and_then(|requirements| requirements.gpu.as_ref())
 }
 
 /// Resolve a driver GPU request into CDI device identifiers.
 ///
 /// `None` means no GPU was requested. Presence with no explicit device IDs
-/// uses the CDI all-GPU request; otherwise the driver-native IDs pass through.
+/// uses the CDI all-GPU request, preserving the current default GPU behavior;
+/// otherwise the driver-native IDs pass through.
 #[must_use]
-pub fn cdi_gpu_device_ids(gpu: Option<&GpuRequestSpec>) -> Option<Vec<String>> {
+pub fn cdi_gpu_device_ids(gpu: Option<&DriverGpuResourceRequirement>) -> Option<Vec<String>> {
     match gpu {
-        Some(gpu) if gpu.device_id.is_empty() => Some(vec![CDI_GPU_DEVICE_ALL.to_string()]),
-        Some(gpu) => Some(gpu.device_id.clone()),
+        Some(gpu) if gpu.device_ids.is_empty() => Some(vec![CDI_GPU_DEVICE_ALL.to_string()]),
+        Some(gpu) => Some(gpu.device_ids.clone()),
         None => None,
     }
 }
@@ -36,8 +39,8 @@ mod tests {
 
     #[test]
     fn cdi_gpu_device_ids_defaults_empty_request_to_all_gpus() {
-        let request = GpuRequestSpec {
-            device_id: vec![],
+        let request = DriverGpuResourceRequirement {
+            device_ids: vec![],
             count: None,
         };
 
@@ -49,8 +52,8 @@ mod tests {
 
     #[test]
     fn cdi_gpu_device_ids_passes_single_device_id_through() {
-        let request = GpuRequestSpec {
-            device_id: vec!["nvidia.com/gpu=0".to_string()],
+        let request = DriverGpuResourceRequirement {
+            device_ids: vec!["nvidia.com/gpu=0".to_string()],
             count: None,
         };
 
@@ -62,8 +65,8 @@ mod tests {
 
     #[test]
     fn cdi_gpu_device_ids_passes_multiple_device_ids_through() {
-        let request = GpuRequestSpec {
-            device_id: vec![
+        let request = DriverGpuResourceRequirement {
+            device_ids: vec![
                 "nvidia.com/gpu=0".to_string(),
                 "nvidia.com/gpu=1".to_string(),
             ],

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 explicit GPU request device IDs when set; otherwise requests all NVIDIA GPUs when the sandbox spec asks for GPU support and daemon CDI support is detected. Count-based GPU requests are rejected until Docker CDI selection can map counts to concrete devices. |
+| CDI GPU request | Uses explicit `resource_requirements.gpu.device_ids` when set; otherwise requests all NVIDIA GPUs when `resource_requirements.gpu` is present and daemon CDI support is detected. Count-based GPU requests are rejected until Docker CDI selection can map counts to concrete devices. |
 
 The agent child process does not retain these supervisor privileges.
 

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

@@ -25,16 +25,16 @@ 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_request};
+use openshell_core::gpu::{cdi_gpu_device_ids, driver_gpu_requirement};
 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,
 };
 use openshell_core::proto::compute::v1::{
     CreateSandboxRequest, CreateSandboxResponse, DeleteSandboxRequest, DeleteSandboxResponse,
-    DriverCondition, DriverPlatformEvent, DriverSandbox, DriverSandboxStatus,
-    DriverSandboxTemplate, GetCapabilitiesRequest, GetCapabilitiesResponse, GetSandboxRequest,
-    GetSandboxResponse, GpuRequestSpec, ListSandboxesRequest, ListSandboxesResponse,
+    DriverCondition, DriverGpuResourceRequirement, DriverPlatformEvent, DriverSandbox,
+    DriverSandboxStatus, DriverSandboxTemplate, GetCapabilitiesRequest, GetCapabilitiesResponse,
+    GetSandboxRequest, GetSandboxResponse, ListSandboxesRequest, ListSandboxesResponse,
     StopSandboxRequest, StopSandboxResponse, ValidateSandboxCreateRequest,
     ValidateSandboxCreateResponse, WatchSandboxesDeletedEvent, WatchSandboxesEvent,
     WatchSandboxesPlatformEvent, WatchSandboxesRequest, WatchSandboxesSandboxEvent,
@@ -375,7 +375,7 @@ impl DockerComputeDriver {
                 "docker sandboxes require a template image",
             ));
         }
-        Self::validate_gpu_request(driver_gpu_request(spec), config.supports_gpu)?;
+        Self::validate_gpu_request(driver_gpu_requirement(spec), config.supports_gpu)?;
         if !template.agent_socket_path.trim().is_empty() {
             return Err(Status::failed_precondition(
                 "docker compute driver does not support template.agent_socket_path",
@@ -410,7 +410,7 @@ impl DockerComputeDriver {
     }
 
     fn validate_gpu_request(
-        gpu: Option<&GpuRequestSpec>,
+        gpu: Option<&DriverGpuResourceRequirement>,
         supports_gpu: bool,
     ) -> Result<(), Status> {
         if gpu.is_some_and(|gpu| gpu.count.is_some()) {
@@ -1721,7 +1721,9 @@ fn build_environment(sandbox: &DriverSandbox, config: &DockerDriverRuntimeConfig
         .collect()
 }
 
-fn docker_gpu_device_requests(gpu: Option<&GpuRequestSpec>) -> Option<Vec<DeviceRequest>> {
+fn docker_gpu_device_requests(
+    gpu: Option<&DriverGpuResourceRequirement>,
+) -> Option<Vec<DeviceRequest>> {
     cdi_gpu_device_ids(gpu).map(|device_ids| {
         vec![DeviceRequest {
             driver: Some("cdi".to_string()),
@@ -1773,7 +1775,7 @@ fn build_container_create_body(
             nano_cpus: resource_limits.nano_cpus,
             memory: resource_limits.memory_bytes,
             pids_limit: docker_pids_limit(config.sandbox_pids_limit)?,
-            device_requests: docker_gpu_device_requests(driver_gpu_request(spec)),
+            device_requests: docker_gpu_device_requests(driver_gpu_requirement(spec)),
             binds: Some(build_binds(sandbox, config)?),
             restart_policy: Some(RestartPolicy {
                 name: Some(RestartPolicyNameEnum::UNLESS_STOPPED),