feat(gpu): introduce GPU request spec

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

Author: elezar

crates/openshell-cli/src/run.rs

@@ -37,15 +37,16 @@ use openshell_core::proto::{
     GetClusterInferenceRequest, GetDraftHistoryRequest, GetDraftPolicyRequest,
     GetGatewayConfigRequest, GetProviderProfileRequest, GetProviderRequest,
     GetSandboxConfigRequest, GetSandboxLogsRequest, GetSandboxPolicyStatusRequest,
-    GetSandboxRequest, GetServiceRequest, HealthRequest, ImportProviderProfilesRequest,
+    GetSandboxRequest, GetServiceRequest, GpuSpec, HealthRequest, ImportProviderProfilesRequest,
     LintProviderProfilesRequest, ListProviderProfilesRequest, ListProvidersRequest,
     ListSandboxPoliciesRequest, ListSandboxProvidersRequest, ListSandboxesRequest,
     ListServicesRequest, PlatformEvent, PolicySource, PolicyStatus, Provider, ProviderProfile,
     ProviderProfileDiagnostic, ProviderProfileImportItem, RejectDraftChunkRequest,
-    RevokeSshSessionRequest, Sandbox, SandboxPhase, SandboxPolicy, SandboxSpec, SandboxTemplate,
-    ServiceEndpointResponse, SetClusterInferenceRequest, SettingScope, SettingValue,
-    TcpForwardFrame, TcpForwardInit, TcpRelayTarget, UpdateConfigRequest, UpdateProviderRequest,
-    WatchSandboxRequest, exec_sandbox_event, setting_value, tcp_forward_init,
+    ResourceRequirements, RevokeSshSessionRequest, Sandbox, SandboxPhase, SandboxPolicy,
+    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};
@@ -1692,8 +1693,7 @@ pub async fn sandbox_create(
 
     let request = CreateSandboxRequest {
         spec: Some(SandboxSpec {
-            gpu: requested_gpu,
-            gpu_device: gpu_device.unwrap_or_default().to_string(),
+            placement: resource_requirements_from_cli(requested_gpu, gpu_device),
             policy,
             providers: configured_providers,
             template,
@@ -2083,6 +2083,20 @@ pub async fn sandbox_create(
     }
 }
 
+fn resource_requirements_from_cli(
+    requested_gpu: bool,
+    gpu_device: Option<&str>,
+) -> Option<ResourceRequirements> {
+    requested_gpu.then(|| ResourceRequirements {
+        gpu: Some(GpuSpec {
+            device_ids: gpu_device
+                .filter(|device_id| !device_id.is_empty())
+                .map(|device_id| vec![device_id.to_string()])
+                .unwrap_or_default(),
+        }),
+    })
+}
+
 /// Resolved source for the `--from` flag on `sandbox create`.
 #[derive(Debug)]
 enum ResolvedSource {
@@ -6281,7 +6295,8 @@ mod tests {
         inferred_provider_type, package_managed_tls_dirs, parse_cli_setting_value,
         parse_credential_pairs, plaintext_gateway_is_remote, progress_step_from_metadata,
         provisioning_timeout_message, ready_false_condition_message, resolve_from,
-        sandbox_should_persist, service_expose_status_error, service_url_for_gateway,
+        resource_requirements_from_cli, sandbox_should_persist, service_expose_status_error,
+        service_url_for_gateway, source_requests_gpu,
     };
     use crate::TEST_ENV_LOCK;
     use hyper::StatusCode;
@@ -6628,6 +6643,35 @@ mod tests {
         }
     }
 
+    #[test]
+    fn source_requests_gpu_detects_known_community_gpu_name() {
+        assert!(source_requests_gpu("nvidia-gpu"));
+        assert!(!source_requests_gpu("base"));
+    }
+
+    #[test]
+    fn resource_requirements_from_cli_uses_presence_with_empty_device_ids_for_default_gpu() {
+        let request = resource_requirements_from_cli(true, None)
+            .expect("resource requirements should be present");
+        let gpu = request.gpu.expect("gpu request should be present");
+
+        assert!(gpu.device_ids.is_empty());
+    }
+
+    #[test]
+    fn resource_requirements_from_cli_maps_gpu_device_to_one_device_id() {
+        let request = resource_requirements_from_cli(true, Some("0000:2d:00.0"))
+            .expect("resource requirements should be present");
+        let gpu = request.gpu.expect("gpu request should be present");
+
+        assert_eq!(gpu.device_ids, vec!["0000:2d:00.0"]);
+    }
+
+    #[test]
+    fn resource_requirements_from_cli_omits_placement_when_not_requested() {
+        assert!(resource_requirements_from_cli(false, Some("0")).is_none());
+    }
+
     #[test]
     fn resolve_from_classifies_existing_dockerfile_path() {
         let temp = tempfile::tempdir().expect("failed to create tempdir");

crates/openshell-core/src/gpu.rs

@@ -4,21 +4,19 @@
 //! Shared GPU request helpers.
 
 use crate::config::CDI_GPU_DEVICE_ALL;
+use crate::proto::compute::v1::GpuSpec;
 
-/// Resolve the existing GPU request fields into CDI device identifiers.
+/// Resolve a driver GPU request into CDI device identifiers.
 ///
-/// `None` means no GPU was requested. A GPU request with no explicit device
-/// ID uses the CDI all-GPU request; otherwise the driver-native ID passes
-/// through unchanged.
+/// `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.
 #[must_use]
-pub fn cdi_gpu_device_ids(gpu: bool, gpu_device: &str) -> Option<Vec<String>> {
-    gpu.then(|| {
-        if gpu_device.is_empty() {
-            vec![CDI_GPU_DEVICE_ALL.to_string()]
-        } else {
-            vec![gpu_device.to_string()]
-        }
-    })
+pub fn cdi_gpu_device_ids(gpu: Option<&GpuSpec>) -> Option<Vec<String>> {
+    match gpu {
+        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,
+    }
 }
 
 #[cfg(test)]
@@ -27,22 +25,46 @@ mod tests {
 
     #[test]
     fn cdi_gpu_device_ids_returns_none_when_absent() {
-        assert_eq!(cdi_gpu_device_ids(false, ""), None);
+        assert_eq!(cdi_gpu_device_ids(None), None);
     }
 
     #[test]
     fn cdi_gpu_device_ids_defaults_empty_request_to_all_gpus() {
+        let request = GpuSpec { device_ids: vec![] };
+
         assert_eq!(
-            cdi_gpu_device_ids(true, ""),
+            cdi_gpu_device_ids(Some(&request)),
             Some(vec![CDI_GPU_DEVICE_ALL.to_string()])
         );
     }
 
     #[test]
-    fn cdi_gpu_device_ids_passes_explicit_device_id_through() {
+    fn cdi_gpu_device_ids_passes_single_device_id_through() {
+        let request = GpuSpec {
+            device_ids: vec!["nvidia.com/gpu=0".to_string()],
+        };
+
         assert_eq!(
-            cdi_gpu_device_ids(true, "nvidia.com/gpu=0"),
+            cdi_gpu_device_ids(Some(&request)),
             Some(vec!["nvidia.com/gpu=0".to_string()])
         );
     }
+
+    #[test]
+    fn cdi_gpu_device_ids_passes_multiple_device_ids_through() {
+        let request = GpuSpec {
+            device_ids: vec![
+                "nvidia.com/gpu=0".to_string(),
+                "nvidia.com/gpu=1".to_string(),
+            ],
+        };
+
+        assert_eq!(
+            cdi_gpu_device_ids(Some(&request)),
+            Some(vec![
+                "nvidia.com/gpu=0".to_string(),
+                "nvidia.com/gpu=1".to_string()
+            ])
+        );
+    }
 }

crates/openshell-driver-docker/README.md

@@ -30,7 +30,7 @@ contract:
 | `cap_add` | Grants supervisor-only capabilities required for namespace setup and process inspection. |
 | `apparmor=unconfined` | Avoids Docker's default profile blocking required mount operations. |
 | `restart_policy = unless-stopped` | Keeps managed sandboxes resumable across daemon or gateway restarts. |
-| CDI GPU request | Uses the sandbox `gpu_device` value when set; otherwise requests all NVIDIA GPUs when the sandbox spec asks for GPU support and daemon CDI support is detected. |
+| CDI GPU request | Uses explicit placement GPU device IDs when set; otherwise requests all NVIDIA GPUs for GPU placement requests when daemon CDI support is detected. |
 
 The agent child process does not retain these supervisor privileges.
 

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

@@ -24,7 +24,7 @@ use openshell_core::proto::compute::v1::{
     CreateSandboxRequest, CreateSandboxResponse, DeleteSandboxRequest, DeleteSandboxResponse,
     DriverCondition, DriverSandbox, DriverSandboxStatus, DriverSandboxTemplate,
     GetCapabilitiesRequest, GetCapabilitiesResponse, GetSandboxRequest, GetSandboxResponse,
-    ListSandboxesRequest, ListSandboxesResponse, StopSandboxRequest, StopSandboxResponse,
+    GpuSpec, ListSandboxesRequest, ListSandboxesResponse, StopSandboxRequest, StopSandboxResponse,
     ValidateSandboxCreateRequest, ValidateSandboxCreateResponse, WatchSandboxesDeletedEvent,
     WatchSandboxesEvent, WatchSandboxesRequest, WatchSandboxesSandboxEvent,
     compute_driver_server::ComputeDriver, watch_sandboxes_event,
@@ -355,7 +355,12 @@ impl DockerComputeDriver {
                 "docker sandboxes require a template image",
             ));
         }
-        Self::validate_gpu_request(spec.gpu, config.supports_gpu)?;
+        Self::validate_gpu_request(
+            spec.placement
+                .as_ref()
+                .and_then(|placement| placement.gpu.as_ref()),
+            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",
@@ -375,8 +380,8 @@ impl DockerComputeDriver {
         Ok(())
     }
 
-    fn validate_gpu_request(gpu: bool, supports_gpu: bool) -> Result<(), Status> {
-        if gpu && !supports_gpu {
+    fn validate_gpu_request(gpu: Option<&GpuSpec>, supports_gpu: bool) -> Result<(), Status> {
+        if gpu.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.",
             ));
@@ -996,8 +1001,8 @@ fn build_environment(sandbox: &DriverSandbox, config: &DockerDriverRuntimeConfig
         .collect()
 }
 
-fn docker_gpu_device_requests(gpu: bool, gpu_device: &str) -> Option<Vec<DeviceRequest>> {
-    cdi_gpu_device_ids(gpu, gpu_device).map(|device_ids| {
+fn docker_gpu_device_requests(gpu: Option<&GpuSpec>) -> Option<Vec<DeviceRequest>> {
+    cdi_gpu_device_ids(gpu).map(|device_ids| {
         vec![DeviceRequest {
             driver: Some("cdi".to_string()),
             device_ids: Some(device_ids),
@@ -1047,8 +1052,12 @@ fn build_container_create_body(
         host_config: Some(HostConfig {
             nano_cpus: resource_limits.nano_cpus,
             memory: resource_limits.memory_bytes,
-            device_requests: docker_gpu_device_requests(spec.gpu, &spec.gpu_device),
-            binds: Some(build_binds(config)),
+            device_requests: docker_gpu_device_requests(
+                spec.placement
+                    .as_ref()
+                    .and_then(|placement| placement.gpu.as_ref()),
+            ),
+            mounts: Some(build_mounts(config)),
             restart_policy: Some(RestartPolicy {
                 name: Some(RestartPolicyNameEnum::UNLESS_STOPPED),
                 maximum_retry_count: None,

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

@@ -4,7 +4,8 @@
 use super::*;
 use openshell_core::config::{CDI_GPU_DEVICE_ALL, DEFAULT_SERVER_PORT};
 use openshell_core::proto::compute::v1::{
-    DriverResourceRequirements, DriverSandboxSpec, DriverSandboxTemplate,
+    DriverResourceRequirements, DriverSandboxSpec, DriverSandboxTemplate, GpuSpec,
+    ResourceRequirements,
 };
 use std::fs;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
@@ -31,8 +32,7 @@ fn test_sandbox() -> DriverSandbox {
                 resources: None,
                 platform_config: None,
             }),
-            gpu: false,
-            gpu_device: String::new(),
+            placement: None,
         }),
         status: None,
     }
@@ -499,7 +499,9 @@ fn build_container_create_body_clears_inherited_cmd() {
 fn validate_sandbox_rejects_gpu_when_cdi_unavailable() {
     let config = runtime_config();
     let mut sandbox = test_sandbox();
-    sandbox.spec.as_mut().unwrap().gpu = true;
+    sandbox.spec.as_mut().unwrap().placement = Some(ResourceRequirements {
+        gpu: Some(GpuSpec { device_ids: vec![] }),
+    });
 
     let err = DockerComputeDriver::validate_sandbox(&sandbox, &config).unwrap_err();
 
@@ -512,7 +514,9 @@ fn build_container_create_body_maps_gpu_to_all_cdi_device() {
     let mut config = runtime_config();
     config.supports_gpu = true;
     let mut sandbox = test_sandbox();
-    sandbox.spec.as_mut().unwrap().gpu = true;
+    sandbox.spec.as_mut().unwrap().placement = Some(ResourceRequirements {
+        gpu: Some(GpuSpec { device_ids: vec![] }),
+    });
 
     let create_body = build_container_create_body(&sandbox, &config).unwrap();
     let request = create_body
@@ -530,13 +534,18 @@ fn build_container_create_body_maps_gpu_to_all_cdi_device() {
 }
 
 #[test]
-fn build_container_create_body_passes_explicit_cdi_device_id_through() {
+fn build_container_create_body_passes_explicit_cdi_device_ids_through() {
     let mut config = runtime_config();
     config.supports_gpu = true;
     let mut sandbox = test_sandbox();
-    let spec = sandbox.spec.as_mut().unwrap();
-    spec.gpu = true;
-    spec.gpu_device = "nvidia.com/gpu=0".to_string();
+    sandbox.spec.as_mut().unwrap().placement = Some(ResourceRequirements {
+        gpu: Some(GpuSpec {
+            device_ids: vec![
+                "nvidia.com/gpu=0".to_string(),
+                "nvidia.com/gpu=1".to_string(),
+            ],
+        }),
+    });
 
     let create_body = build_container_create_body(&sandbox, &config).unwrap();
     let request = create_body
@@ -549,7 +558,10 @@ fn build_container_create_body_passes_explicit_cdi_device_id_through() {
     assert_eq!(request.driver.as_deref(), Some("cdi"));
     assert_eq!(
         request.device_ids.as_ref().unwrap(),
-        &vec!["nvidia.com/gpu=0".to_string()]
+        &vec![
+            "nvidia.com/gpu=0".to_string(),
+            "nvidia.com/gpu=1".to_string()
+        ]
     );
 }