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};
@@ -1801,11 +1801,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()
@@ -1825,6 +1820,11 @@ pub async fn sandbox_create(
     let driver_config = driver_config_json
         .map(parse_driver_config_json)
         .transpose()?;
+    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() || driver_config.is_some() {
         Some(SandboxTemplate {
@@ -1839,7 +1839,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,
@@ -2274,17 +2274,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,
+        }),
     })
 }
 
@@ -7537,15 +7546,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, parse_driver_config_json, 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,
+        refresh_status_row, resolve_from, 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;
@@ -8075,43 +8083,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

@@ -788,6 +788,7 @@ async fn sandbox_create_keeps_command_sessions_by_default() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -907,6 +908,7 @@ async fn sandbox_create_sends_driver_config_json() {
         None,
         None,
         None,
+        None,
         Some(r#"{"kubernetes":{"pod":{"priority_class_name":"batch-low"}}}"#),
         None,
         &[],
@@ -982,6 +984,7 @@ async fn sandbox_create_sends_gpu_device_request_without_gpu_flag() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -999,10 +1002,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);
 }
 
@@ -1020,14 +1024,15 @@ async fn sandbox_create_sends_gpu_count_request() {
         Some("gpu-count"),
         None,
         "openshell",
-        None,
+        &[],
         true,
         false,
         None,
         Some(2),
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -1045,10 +1050,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));
 }
 
@@ -1075,6 +1081,7 @@ async fn sandbox_create_does_not_infer_command_providers_when_v2_enabled() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -1133,6 +1140,7 @@ async fn sandbox_create_returns_vm_error_without_waiting_for_timeout() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -1187,6 +1195,7 @@ async fn sandbox_create_keeps_waiting_while_vm_progress_arrives() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -1233,6 +1242,7 @@ async fn sandbox_create_times_out_when_only_logs_arrive() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -1275,6 +1285,7 @@ async fn sandbox_create_deletes_command_sessions_with_no_keep() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -1321,6 +1332,7 @@ async fn sandbox_create_deletes_shell_sessions_with_no_keep() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -1367,6 +1379,7 @@ async fn sandbox_create_keeps_sandbox_with_hidden_keep_flag() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -1413,6 +1426,7 @@ async fn sandbox_create_keeps_sandbox_with_forwarding() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         Some(openshell_core::forward::ForwardSpec::new(forward_port)),