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,
@@ -2441,6 +2447,7 @@ async fn main() -> Result<()> {
                         keep,
                         gpu,
                         gpu_device.as_deref(),
+                        gpu_count,
                         cpu.as_deref(),
                         memory.as_deref(),
                         editor,
@@ -3709,6 +3716,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

@@ -1564,6 +1564,7 @@ pub async fn sandbox_create(
     keep: bool,
     gpu: bool,
     gpu_device: Option<&str>,
+    gpu_count: Option<u32>,
     cpu: Option<&str>,
     memory: Option<&str>,
     editor: Option<Editor>,
@@ -1618,6 +1619,7 @@ pub async fn sandbox_create(
     };
     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 inferred_types: Vec<String> = inferred_provider_type(command).into_iter().collect();
@@ -1644,7 +1646,7 @@ pub async fn sandbox_create(
 
     let request = CreateSandboxRequest {
         spec: Some(SandboxSpec {
-            placement: resource_requirements_from_cli(requested_gpu, gpu_device),
+            placement: resource_requirements_from_cli(requested_gpu, gpu_device, gpu_count),
             policy,
             providers: configured_providers,
             template,
@@ -2079,13 +2081,22 @@ pub async fn sandbox_create(
 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,
         }),
     })
 }
@@ -6135,7 +6146,6 @@ mod tests {
         parse_credential_pairs, plaintext_gateway_is_remote, 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,
     };
     use crate::TEST_ENV_LOCK;
     use hyper::StatusCode;
@@ -6462,32 +6472,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());
+        assert!(resource_requirements_from_cli(false, None, None).is_none());
     }
 
     #[test]

crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs

@@ -601,6 +601,7 @@ async fn sandbox_create_keeps_command_sessions_by_default() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -639,6 +640,7 @@ async fn sandbox_create_sends_cpu_and_memory_limits_only() {
         true,
         false,
         None,
+        None,
         Some("500m"),
         Some("2Gi"),
         None,
@@ -717,6 +719,7 @@ async fn sandbox_create_deletes_command_sessions_with_no_keep() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -761,6 +764,7 @@ async fn sandbox_create_deletes_shell_sessions_with_no_keep() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -805,6 +809,7 @@ async fn sandbox_create_keeps_sandbox_with_hidden_keep_flag() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         None,
@@ -849,6 +854,7 @@ async fn sandbox_create_keeps_sandbox_with_forwarding() {
         None,
         None,
         None,
+        None,
         &[],
         None,
         Some(openshell_core::forward::ForwardSpec::new(forward_port)),

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 placement GPU device IDs when set; otherwise requests all NVIDIA GPUs for GPU placement 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.
 

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

@@ -336,6 +336,21 @@ impl DockerComputeDriver {
     }
 
     fn validate_gpu_request(gpu: Option<&GpuSpec>, supports_gpu: bool) -> Result<(), Status> {
+        if let Some(gpu) = gpu {
+            if gpu.count == Some(0) {
+                return Err(Status::invalid_argument("gpu.count must be greater than 0"));
+            }
+            if gpu.count.is_some() && !gpu.device_ids.is_empty() {
+                return Err(Status::invalid_argument(
+                    "gpu.count is mutually exclusive with gpu.device_ids",
+                ));
+            }
+            if gpu.count.is_some() {
+                return Err(Status::invalid_argument(
+                    "docker compute driver does not support GPU count requests",
+                ));
+            }
+        }
         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.",
@@ -1012,7 +1027,7 @@ fn build_container_create_body(
                     .as_ref()
                     .and_then(|placement| placement.gpu.as_ref()),
             ),
-            mounts: Some(build_mounts(config)),
+            binds: Some(build_binds(config)),
             restart_policy: Some(RestartPolicy {
                 name: Some(RestartPolicyNameEnum::UNLESS_STOPPED),
                 maximum_retry_count: None,

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

@@ -500,7 +500,10 @@ fn validate_sandbox_rejects_gpu_when_cdi_unavailable() {
     let config = runtime_config();
     let mut sandbox = test_sandbox();
     sandbox.spec.as_mut().unwrap().placement = Some(ResourceRequirements {
-        gpu: Some(GpuSpec { device_ids: vec![] }),
+        gpu: Some(GpuSpec {
+            device_ids: vec![],
+            count: None,
+        }),
     });
 
     let err = DockerComputeDriver::validate_sandbox(&sandbox, &config).unwrap_err();
@@ -509,13 +512,34 @@ fn validate_sandbox_rejects_gpu_when_cdi_unavailable() {
     assert!(err.message().contains("Docker CDI"));
 }
 
+#[test]
+fn validate_sandbox_rejects_gpu_count() {
+    let mut config = runtime_config();
+    config.supports_gpu = true;
+    let mut sandbox = test_sandbox();
+    sandbox.spec.as_mut().unwrap().placement = Some(ResourceRequirements {
+        gpu: Some(GpuSpec {
+            device_ids: vec![],
+            count: Some(2),
+        }),
+    });
+
+    let err = DockerComputeDriver::validate_sandbox(&sandbox, &config).unwrap_err();
+
+    assert_eq!(err.code(), tonic::Code::InvalidArgument);
+    assert!(err.message().contains("does not support GPU count"));
+}
+
 #[test]
 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().placement = Some(ResourceRequirements {
-        gpu: Some(GpuSpec { device_ids: vec![] }),
+        gpu: Some(GpuSpec {
+            device_ids: vec![],
+            count: None,
+        }),
     });
 
     let create_body = build_container_create_body(&sandbox, &config).unwrap();
@@ -544,6 +568,7 @@ fn build_container_create_body_passes_explicit_cdi_device_ids_through() {
                 "nvidia.com/gpu=0".to_string(),
                 "nvidia.com/gpu=1".to_string(),
             ],
+            count: None,
         }),
     });