feat(gpu): introduce GPU request spec

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
@@ -82,7 +85,9 @@ users.
 Custom sandbox images must include the agent runtime and any system
 dependencies, but they should not need to include the gateway. GPU-capable
 images must include the user-space libraries required by the workload. The
-runtime still owns GPU device injection.
+runtime still owns GPU device injection. GPU requests can include explicit
+driver-native device IDs or a requested count; the gateway validates the public
+request shape and each runtime enforces the GPU allocation modes it supports.
 
 ## Deployment Shape
 

crates/openshell-cli/src/main.rs

@@ -1215,10 +1215,15 @@ 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.
- /// Only valid with --gpu. When omitted with --gpu, the driver uses its default GPU selection.
- #[arg(long, requires = "gpu")]
+ /// Specifying --gpu-device also requests GPU resources.
+ /// When omitted with --gpu, the driver uses its default GPU selection.
+ #[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>,
@@ -2547,6 +2552,7 @@ async fn main() -> Result<()> {
  editor,
  gpu,
  gpu_device,
+ gpu_count,
  cpu,
  memory,
  driver_config_json,
@@ -2628,6 +2634,7 @@ async fn main() -> Result<()> {
  keep,
  gpu,
  gpu_device.as_deref(),
+ gpu_count,
  cpu.as_deref(),
  memory.as_deref(),
  driver_config_json.as_deref(),
@@ -4371,6 +4378,78 @@ mod tests {
  }
  }
 
+ #[test]
+ fn sandbox_create_gpu_device_parses_without_gpu_flag() {
+ let cli = Cli::try_parse_from([
+ "openshell",
+ "sandbox",
+ "create",
+ "--gpu-device",
+ "nvidia.com/gpu=0",
+ ])
+ .expect("sandbox create --gpu-device should parse without --gpu");
+
+ match cli.command {
+ Some(Commands::Sandbox {
+ command:
+ Some(SandboxCommands::Create {
+ gpu, gpu_device, ..
+ }),
+ ..
+ }) => {
+ assert!(!gpu);
+ assert_eq!(gpu_device.as_deref(), Some("nvidia.com/gpu=0"));
+ }
+ other => panic!("expected SandboxCommands::Create, got: {other:?}"),
+ }
+ }
+
+ #[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");
+
+ match cli.command {
+ Some(Commands::Sandbox {
+ command: Some(SandboxCommands::Create { gpu, gpu_count, .. }),
+ ..
+ }) => {
+ assert!(!gpu);
+ assert_eq!(gpu_count, Some(2));
+ }
+ other => panic!("expected SandboxCommands::Create, got: {other:?}"),
+ }
+ }
+
+ #[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",
+ "nvidia.com/gpu=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

@@ -39,17 +39,18 @@ use openshell_core::proto::{
  GetClusterInferenceRequest, GetDraftHistoryRequest, GetDraftPolicyRequest,
  GetGatewayConfigRequest, GetProviderProfileRequest, GetProviderRefreshStatusRequest,
  GetProviderRequest, GetSandboxConfigRequest, GetSandboxLogsRequest,
- GetSandboxPolicyStatusRequest, GetSandboxRequest, GetServiceRequest, 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,
+ 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,
+ 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};
@@ -1745,6 +1746,7 @@ pub async fn sandbox_create(
  keep: bool,
  gpu: bool,
  gpu_device: Option<&str>,
+ gpu_count: Option<u32>,
  cpu: Option<&str>,
  memory: Option<&str>,
  driver_config_json: Option<&str>,
@@ -1799,8 +1801,6 @@ pub async fn sandbox_create(
  }
  None => None,
  };
- let requested_gpu = gpu || 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()
@@ -1820,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 {
@@ -1834,8 +1839,7 @@ pub async fn sandbox_create(
 
  let request = CreateSandboxRequest {
  spec: Some(SandboxSpec {
- gpu: requested_gpu,
- gpu_device: gpu_device.unwrap_or_default().to_string(),
+ resource_requirements,
  policy,
  providers: configured_providers,
  template,
@@ -2270,6 +2274,29 @@ pub async fn sandbox_create(
  }
 }
 
+fn resource_requirements_from_cli(
+ image: Option<&str>,
+ gpu: bool,
+ gpu_device: Option<&str>,
+ gpu_count: Option<u32>,
+) -> 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,
+ }),
+ })
+}
+
 /// Resolved source for the `--from` flag on `sandbox create`.
 #[derive(Debug)]
 enum ResolvedSource {
@@ -7525,8 +7552,8 @@ mod tests {
  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;
@@ -8055,6 +8082,67 @@ mod tests {
  }
  }
 
+ #[test]
+ 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!(gpu.device_ids.is_empty());
+ assert_eq!(gpu.count, None);
+ }
+
+ #[test]
+ 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!(gpu.device_ids, vec!["0000:2d:00.0"]);
+ assert_eq!(gpu.count, None);
+ }
+
+ #[test]
+ 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!(gpu.device_ids.is_empty());
+ assert_eq!(gpu.count, Some(2));
+ }
+
+ #[test]
+ 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!(gpu.device_ids, vec!["nvidia.com/gpu=0"]);
+ assert_eq!(gpu.count, Some(2));
+ }
+
+ #[test]
+ 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]
  fn resolve_from_classifies_existing_dockerfile_path() {
  let temp = tempfile::tempdir().expect("failed to create tempdir");