feat(gpu): add GPU count resource requests

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

Author: elezar

crates/openshell-cli/src/main.rs

@@ -1029,6 +1029,7 @@ enum DoctorCommands {
     Check,
 }
 
+#[allow(clippy::large_enum_variant)]
 #[derive(Subcommand, Debug)]
 enum SandboxCommands {
     /// Create a sandbox.
@@ -1086,10 +1087,14 @@ enum SandboxCommands {
 
         /// Target a driver-specific GPU device. Docker and Podman use CDI device IDs
         /// (for example "nvidia.com/gpu=0"); VM uses a PCI BDF or index.
-        /// Implies a GPU request. When omitted with --gpu, the driver uses its default GPU selection.
-        #[arg(long)]
+        /// Implies a GPU request. Mutually exclusive with --gpu-count.
+        #[arg(long, conflicts_with = "gpu_count")]
         gpu_device: Option<String>,
 
+        /// Request a specific number of GPUs. Mutually exclusive with --gpu-device.
+        #[arg(long, value_parser = clap::value_parser!(u32).range(1..), conflicts_with = "gpu_device")]
+        gpu_count: Option<u32>,
+
         /// CPU limit for the sandbox (for example: 500m, 1, 2.5).
         #[arg(long)]
         cpu: Option<String>,
@@ -2373,6 +2378,7 @@ async fn main() -> Result<()> {
                     editor,
                     gpu,
                     gpu_device,
+                    gpu_count,
                     cpu,
                     memory,
                     providers,
@@ -2432,15 +2438,16 @@ async fn main() -> Result<()> {
                     let endpoint = &ctx.endpoint;
                     let mut tls = tls.with_gateway_name(&ctx.name);
                     apply_auth(&mut tls, &ctx.name);
+                    let resource_requirements =
+                        run::resource_requirements_from_cli(gpu, gpu_device.as_deref(), gpu_count);
                     Box::pin(run::sandbox_create(
                         endpoint,
                         name.as_deref(),
                         from.as_deref(),
                         &ctx.name,
                         upload_spec.as_ref(),
                         keep,
-                        gpu,
-                        gpu_device.as_deref(),
+                        resource_requirements,
                         cpu.as_deref(),
                         memory.as_deref(),
                         editor,
@@ -3713,6 +3720,52 @@ mod tests {
         }
     }
 
+    #[test]
+    fn sandbox_create_gpu_count_parses_without_gpu_flag() {
+        let cli = Cli::try_parse_from(["openshell", "sandbox", "create", "--gpu-count", "2"])
+            .expect("sandbox create --gpu-count should parse");
+
+        if let Some(Commands::Sandbox {
+            command: Some(SandboxCommands::Create { gpu, gpu_count, .. }),
+            ..
+        }) = cli.command
+        {
+            assert!(!gpu);
+            assert_eq!(gpu_count, Some(2));
+        } else {
+            panic!("expected SandboxCommands::Create");
+        }
+    }
+
+    #[test]
+    fn sandbox_create_gpu_count_rejects_zero() {
+        let result = Cli::try_parse_from(["openshell", "sandbox", "create", "--gpu-count", "0"]);
+
+        assert!(
+            result.is_err(),
+            "sandbox create --gpu-count 0 should be rejected"
+        );
+    }
+
+    #[test]
+    fn sandbox_create_gpu_count_conflicts_with_gpu_device() {
+        let result = Cli::try_parse_from([
+            "openshell",
+            "sandbox",
+            "create",
+            "--gpu",
+            "--gpu-device",
+            "0",
+            "--gpu-count",
+            "2",
+        ]);
+
+        assert!(
+            result.is_err(),
+            "sandbox create should reject --gpu-count with --gpu-device"
+        );
+    }
+
     #[test]
     fn service_expose_accepts_positional_target_port_and_service() {
         let cli = Cli::try_parse_from([

crates/openshell-cli/src/run.rs

@@ -1611,8 +1611,7 @@ pub async fn sandbox_create(
     gateway_name: &str,
     upload: Option<&(String, Option<String>, bool)>,
     keep: bool,
-    gpu: bool,
-    gpu_device: Option<&str>,
+    resource_requirements: Option<ResourceRequirements>,
     cpu: Option<&str>,
     memory: Option<&str>,
     editor: Option<Editor>,
@@ -1665,9 +1664,13 @@ pub async fn sandbox_create(
         }
         None => None,
     };
-    let requested_gpu = gpu
-        || gpu_device.is_some_and(|device_id| !device_id.is_empty())
-        || image.as_deref().is_some_and(image_requests_gpu);
+    let resource_requirements = resource_requirements.or_else(|| {
+        resource_requirements_from_cli(image.as_deref().is_some_and(image_requests_gpu), None, None)
+    });
+    let requested_gpu = resource_requirements
+        .as_ref()
+        .and_then(|requirements| requirements.gpu.as_ref())
+        .is_some();
 
     let inferred_types: Vec<String> = inferred_provider_type(command).into_iter().collect();
     let configured_providers = ensure_required_providers(
@@ -1693,7 +1696,7 @@ pub async fn sandbox_create(
 
     let request = CreateSandboxRequest {
         spec: Some(SandboxSpec {
-            placement: resource_requirements_from_cli(requested_gpu, gpu_device),
+            resource_requirements,
             policy,
             providers: configured_providers,
             template,
@@ -2083,16 +2086,25 @@ pub async fn sandbox_create(
     }
 }
 
-fn resource_requirements_from_cli(
+pub fn resource_requirements_from_cli(
     requested_gpu: bool,
     gpu_device: Option<&str>,
+    gpu_count: Option<u32>,
 ) -> Option<ResourceRequirements> {
+    let requested_gpu = requested_gpu
+        || gpu_device.is_some_and(|device_id| !device_id.is_empty())
+        || gpu_count.is_some();
     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(),
+            device_ids: if gpu_count.is_none() {
+                gpu_device
+                    .filter(|device_id| !device_id.is_empty())
+                    .map(|device_id| vec![device_id.to_string()])
+                    .unwrap_or_default()
+            } else {
+                Vec::new()
+            },
+            count: gpu_count,
         }),
     })
 }
@@ -6296,7 +6308,7 @@ mod tests {
         parse_credential_pairs, plaintext_gateway_is_remote, progress_step_from_metadata,
         provisioning_timeout_message, ready_false_condition_message, resolve_from,
         resource_requirements_from_cli, sandbox_should_persist, service_expose_status_error,
-        service_url_for_gateway, source_requests_gpu,
+        service_url_for_gateway,
     };
     use crate::TEST_ENV_LOCK;
     use hyper::StatusCode;
@@ -6644,32 +6656,54 @@ mod tests {
     }
 
     #[test]
-    fn source_requests_gpu_detects_known_community_gpu_name() {
-        assert!(source_requests_gpu("nvidia-gpu"));
-        assert!(!source_requests_gpu("base"));
+    fn image_requests_gpu_detects_known_community_gpu_name() {
+        assert!(image_requests_gpu("nvidia-gpu"));
+        assert!(!image_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)
+        let request = resource_requirements_from_cli(true, None, None)
             .expect("resource requirements should be present");
         let gpu = request.gpu.expect("gpu request 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 request = resource_requirements_from_cli(true, Some("0000:2d:00.0"))
+        let request = resource_requirements_from_cli(true, Some("0000:2d:00.0"), None)
+            .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"]);
+        assert_eq!(gpu.count, None);
+    }
+
+    #[test]
+    fn resource_requirements_from_cli_maps_gpu_device_without_explicit_gpu_flag() {
+        let request = resource_requirements_from_cli(false, Some("0000:2d:00.0"), None)
             .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"]);
+        assert_eq!(gpu.count, None);
+    }
+
+    #[test]
+    fn resource_requirements_from_cli_maps_gpu_count() {
+        let request = resource_requirements_from_cli(false, None, Some(2))
+            .expect("resource requirements should be present");
+        let gpu = request.gpu.expect("gpu request should be present");
+
+        assert!(gpu.device_ids.is_empty());
+        assert_eq!(gpu.count, Some(2));
     }
 
     #[test]
-    fn resource_requirements_from_cli_omits_placement_when_not_requested() {
-        assert!(resource_requirements_from_cli(false, Some("0")).is_none());
+    fn resource_requirements_from_cli_omits_requirements_when_not_requested() {
+        assert!(resource_requirements_from_cli(false, None, None).is_none());
     }
 
     #[test]

crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs

@@ -727,7 +727,6 @@ async fn sandbox_create_keeps_command_sessions_by_default() {
         "openshell",
         None,
         true,
-        false,
         None,
         None,
         None,
@@ -768,7 +767,6 @@ async fn sandbox_create_sends_cpu_and_memory_limits_only() {
         "openshell",
         None,
         true,
-        false,
         None,
         Some("500m"),
         Some("2Gi"),
@@ -853,7 +851,6 @@ async fn sandbox_create_returns_vm_error_without_waiting_for_timeout() {
         "openshell",
         None,
         true,
-        false,
         None,
         None,
         None,
@@ -905,7 +902,6 @@ async fn sandbox_create_keeps_waiting_while_vm_progress_arrives() {
         "openshell",
         None,
         true,
-        false,
         None,
         None,
         None,
@@ -949,7 +945,6 @@ async fn sandbox_create_times_out_when_only_logs_arrive() {
         "openshell",
         None,
         true,
-        false,
         None,
         None,
         None,
@@ -989,7 +984,6 @@ async fn sandbox_create_deletes_command_sessions_with_no_keep() {
         "openshell",
         None,
         false,
-        false,
         None,
         None,
         None,
@@ -1033,7 +1027,6 @@ async fn sandbox_create_deletes_shell_sessions_with_no_keep() {
         "openshell",
         None,
         false,
-        false,
         None,
         None,
         None,
@@ -1077,7 +1070,6 @@ async fn sandbox_create_keeps_sandbox_with_hidden_keep_flag() {
         "openshell",
         None,
         true,
-        false,
         None,
         None,
         None,
@@ -1121,7 +1113,6 @@ async fn sandbox_create_keeps_sandbox_with_forwarding() {
         "openshell",
         None,
         false,
-        false,
         None,
         None,
         None,

crates/openshell-core/src/gpu.rs

@@ -30,7 +30,10 @@ mod tests {
 
     #[test]
     fn cdi_gpu_device_ids_defaults_empty_request_to_all_gpus() {
-        let request = GpuSpec { device_ids: vec![] };
+        let request = GpuSpec {
+            device_ids: vec![],
+            count: None,
+        };
 
         assert_eq!(
             cdi_gpu_device_ids(Some(&request)),
@@ -42,6 +45,7 @@ mod tests {
     fn cdi_gpu_device_ids_passes_single_device_id_through() {
         let request = GpuSpec {
             device_ids: vec!["nvidia.com/gpu=0".to_string()],
+            count: None,
         };
 
         assert_eq!(
@@ -57,6 +61,7 @@ mod tests {
                 "nvidia.com/gpu=0".to_string(),
                 "nvidia.com/gpu=1".to_string(),
             ],
+            count: None,
         };
 
         assert_eq!(

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 explicit placement GPU device IDs when set; otherwise requests all NVIDIA GPUs for GPU placement requests when daemon CDI support is detected. |
+| CDI GPU request | Uses explicit GPU resource device IDs when set; otherwise requests all NVIDIA GPUs for GPU resource requests when daemon CDI support is detected. GPU count requests are rejected by this driver. |
 
 The agent child process does not retain these supervisor privileges.