From 896b2c7ff834055fe549fee381e68a4e577c3790 Mon Sep 17 00:00:00 2001 From: Drew Newberry Date: Thu, 30 Apr 2026 00:05:09 -0700 Subject: [PATCH] feat(docker): support GPU sandboxes Signed-off-by: Drew Newberry --- .agents/skills/openshell-cli/cli-reference.md | 2 + architecture/gateway-single-node.md | 9 +- architecture/gateway.md | 2 +- crates/openshell-bootstrap/src/docker.rs | 140 +++++++++++++++++- crates/openshell-bootstrap/src/lib.rs | 41 +++-- crates/openshell-cli/src/main.rs | 36 +++++ crates/openshell-cli/src/run.rs | 2 + crates/openshell-driver-docker/src/lib.rs | 129 ++++++++++++++-- crates/openshell-driver-docker/src/tests.rs | 95 ++++++++++++ docs/sandboxes/manage-gateways.mdx | 3 +- docs/sandboxes/manage-sandboxes.mdx | 2 + 11 files changed, 432 insertions(+), 29 deletions(-) diff --git a/.agents/skills/openshell-cli/cli-reference.md b/.agents/skills/openshell-cli/cli-reference.md index bdbb35572..9252a2ed8 100644 --- a/.agents/skills/openshell-cli/cli-reference.md +++ b/.agents/skills/openshell-cli/cli-reference.md @@ -85,6 +85,8 @@ Provision or start a cluster (local or remote). | `--port ` | 8080 | Host port mapped to gateway | | `--gateway-host ` | none | Override gateway host in metadata | | `--recreate` | false | Destroy and recreate from scratch if a gateway already exists (skips interactive prompt) | +| `--gpu` | false | Force NVIDIA GPU passthrough | +| `--no-gpu` | false | Disable automatic NVIDIA GPU passthrough detection | ### `openshell gateway stop` diff --git a/architecture/gateway-single-node.md b/architecture/gateway-single-node.md index 01b69b2f5..0bcafe979 100644 --- a/architecture/gateway-single-node.md +++ b/architecture/gateway-single-node.md @@ -297,7 +297,8 @@ When environment variables are set, the entrypoint modifies the HelmChart manife GPU support is part of the single-node gateway bootstrap path rather than a separate architecture. -- `openshell gateway start --gpu` threads GPU device options through `crates/openshell-cli`, `crates/openshell-bootstrap`, and `crates/openshell-bootstrap/src/docker.rs`. +- `openshell gateway start` auto-detects GPU support and threads GPU device options through `crates/openshell-cli`, `crates/openshell-bootstrap`, and `crates/openshell-bootstrap/src/docker.rs`. Users can force passthrough with `--gpu` or disable auto-detection with `--no-gpu`. +- Auto-detection enables passthrough when Docker reports NVIDIA CDI devices. For local non-CDI hosts, it also enables passthrough when `/dev/nvidia*` devices exist and Docker reports the NVIDIA runtime. Remote legacy-runtime hosts still require explicit `--gpu`. - When enabled, the cluster container is created with Docker `DeviceRequests`. The injection mechanism is selected based on whether CDI is enabled on the daemon (`SystemInfo.CDISpecDirs` via `GET /info`): - **CDI enabled** (daemon reports non-empty `CDISpecDirs`): CDI device injection — `driver="cdi"` with `nvidia.com/gpu=all`. Specs are expected to be pre-generated on the host (e.g. automatically by the `nvidia-cdi-refresh.service` or manually via `nvidia-ctk generate`). - **CDI not enabled**: `--gpus all` device request — `driver="nvidia"`, `count=-1`, which relies on the NVIDIA Container Runtime hook. @@ -317,9 +318,11 @@ Host GPU drivers & NVIDIA Container Toolkit └─ Pods: request nvidia.com/gpu in resource limits (CDI injection — no runtimeClassName needed) ``` -### `--gpu` flag +### GPU flags -The `--gpu` flag on `gateway start` enables GPU passthrough. OpenShell auto-selects CDI when enabled on the daemon and falls back to Docker's NVIDIA GPU request path (`--gpus all`) otherwise. +`gateway start` enables GPU passthrough automatically when it detects NVIDIA GPU support. The `--gpu` flag forces GPU passthrough even when auto-detection does not find a device. The `--no-gpu` flag disables auto-detection. + +OpenShell auto-selects CDI when enabled on the daemon and falls back to Docker's NVIDIA GPU request path (`--gpus all`) otherwise. Device injection uses CDI (`deviceListStrategy: cdi-cri`): the device plugin injects devices via direct CDI device requests in the CRI. Sandbox pods only need `nvidia.com/gpu: 1` in their resource limits, and GPU pods do not set `runtimeClassName`. diff --git a/architecture/gateway.md b/architecture/gateway.md index e83640a43..54ddfd72f 100644 --- a/architecture/gateway.md +++ b/architecture/gateway.md @@ -606,7 +606,7 @@ The gateway reaches the sandbox exclusively through the supervisor-initiated `Co The Docker driver (`crates/openshell-driver-docker/src/lib.rs`) is an in-process compute backend for local standalone gateways. It creates one Docker container per sandbox, labels each container with `openshell.ai/managed-by=openshell`, `openshell.ai/sandbox-id`, `openshell.ai/sandbox-name`, and `openshell.ai/sandbox-namespace`, and bind-mounts a Linux `openshell-sandbox` supervisor binary into the container. -- **Create**: Pulls or validates the sandbox image according to `sandbox_image_pull_policy`, creates a labeled container, mounts the supervisor binary and optional TLS material, and starts the container with the supervisor as entrypoint. +- **Create**: Pulls or validates the sandbox image according to `sandbox_image_pull_policy`, creates a labeled container, mounts the supervisor binary and optional TLS material, and starts the container with the supervisor as entrypoint. When the sandbox spec requests GPU and Docker exposes NVIDIA CDI devices or the NVIDIA runtime, the driver adds a Docker `DeviceRequest` for those GPUs. - **List/Get/Watch**: Reads labeled containers in the configured sandbox namespace and derives driver-native sandbox status from Docker state plus supervisor relay readiness. - **Stop**: Stops the matching labeled container without deleting it. - **Delete**: Force-removes the matching labeled container. diff --git a/crates/openshell-bootstrap/src/docker.rs b/crates/openshell-bootstrap/src/docker.rs index 65482739f..cf6e74407 100644 --- a/crates/openshell-bootstrap/src/docker.rs +++ b/crates/openshell-bootstrap/src/docker.rs @@ -8,9 +8,10 @@ use bollard::API_DEFAULT_VERSION; use bollard::Docker; use bollard::errors::Error as BollardError; use bollard::models::{ - ContainerCreateBody, DeviceRequest, EndpointSettings, HostConfig, HostConfigCgroupnsModeEnum, - NetworkConnectRequest, NetworkCreateRequest, NetworkDisconnectRequest, PortBinding, - RestartPolicy, RestartPolicyNameEnum, VolumeCreateRequest, + ContainerCreateBody, DeviceInfo, DeviceRequest, EndpointSettings, HostConfig, + HostConfigCgroupnsModeEnum, NetworkConnectRequest, NetworkCreateRequest, + NetworkDisconnectRequest, PortBinding, RestartPolicy, RestartPolicyNameEnum, Runtime, + SystemInfo, VolumeCreateRequest, }; use bollard::query_parameters::{ CreateContainerOptions, CreateImageOptions, InspectContainerOptions, InspectNetworkOptions, @@ -19,6 +20,7 @@ use bollard::query_parameters::{ }; use futures::StreamExt; use miette::{IntoDiagnostic, Result, WrapErr}; +use openshell_core::config::CDI_GPU_DEVICE_ALL; use std::collections::HashMap; const REGISTRY_NAMESPACE_DEFAULT: &str = "openshell"; @@ -64,6 +66,92 @@ pub(crate) fn resolve_gpu_device_ids(gpu: &[String], cdi_enabled: bool) -> Vec, + local_nvidia_devices_present: bool, +) -> Vec { + let Some(info) = info else { + return Vec::new(); + }; + + let cdi_device_ids = nvidia_cdi_device_ids(info); + if !cdi_device_ids.is_empty() { + return cdi_device_ids; + } + + if local_nvidia_devices_present && docker_info_has_nvidia_runtime(info) { + return vec!["legacy".to_string()]; + } + + Vec::new() +} + +pub(crate) fn docker_info_cdi_enabled(info: Option<&SystemInfo>) -> bool { + info.and_then(|info| info.cdi_spec_dirs.as_ref()) + .is_some_and(|dirs| !dirs.is_empty()) +} + +pub(crate) fn local_nvidia_devices_present() -> bool { + ["/dev/nvidia0", "/dev/nvidiactl", "/proc/driver/nvidia/gpus"] + .iter() + .any(|path| std::path::Path::new(path).exists()) +} + +fn nvidia_cdi_device_ids(info: &SystemInfo) -> Vec { + let Some(devices) = info.discovered_devices.as_ref() else { + return Vec::new(); + }; + + let mut ids = devices + .iter() + .filter_map(nvidia_cdi_device_id) + .collect::>(); + ids.sort(); + ids.dedup(); + + if ids.iter().any(|id| id == CDI_GPU_DEVICE_ALL) { + vec![CDI_GPU_DEVICE_ALL.to_string()] + } else { + ids + } +} + +fn nvidia_cdi_device_id(device: &DeviceInfo) -> Option { + let id = device.id.as_ref()?; + if id.contains("nvidia.com/gpu=") + || (id.contains("/gpu=") && device.source.as_deref().is_some_and(contains_nvidia)) + { + return Some(id.clone()); + } + None +} + +fn docker_info_has_nvidia_runtime(info: &SystemInfo) -> bool { + info.runtimes.as_ref().is_some_and(|runtimes| { + runtimes + .iter() + .any(|(name, runtime)| is_nvidia_runtime(name, runtime)) + }) +} + +fn is_nvidia_runtime(name: &str, runtime: &Runtime) -> bool { + contains_nvidia(name) + || runtime + .path + .as_deref() + .is_some_and(|path| path.contains("nvidia-container-runtime")) +} + +fn contains_nvidia(value: &str) -> bool { + value.to_ascii_lowercase().contains("nvidia") +} + const REGISTRY_MODE_EXTERNAL: &str = "external"; fn env_non_empty(key: &str) -> Option { @@ -1466,4 +1554,50 @@ mod tests { ]; assert_eq!(resolve_gpu_device_ids(&multi, true), multi); } + + #[test] + fn auto_detect_gpu_prefers_reported_nvidia_cdi_devices() { + let info = SystemInfo { + discovered_devices: Some(vec![DeviceInfo { + source: Some("cdi".to_string()), + id: Some("nvidia.com/gpu=all".to_string()), + }]), + runtimes: Some(HashMap::from([("nvidia".to_string(), Runtime::default())])), + ..Default::default() + }; + + assert_eq!( + auto_detect_gpu_device_ids(Some(&info), false), + vec!["nvidia.com/gpu=all".to_string()] + ); + } + + #[test] + fn auto_detect_gpu_uses_legacy_runtime_only_when_local_devices_exist() { + let info = SystemInfo { + runtimes: Some(HashMap::from([("nvidia".to_string(), Runtime::default())])), + ..Default::default() + }; + + assert_eq!( + auto_detect_gpu_device_ids(Some(&info), false), + Vec::::new() + ); + assert_eq!( + auto_detect_gpu_device_ids(Some(&info), true), + vec!["legacy".to_string()] + ); + } + + #[test] + fn docker_info_cdi_enabled_requires_cdi_dirs() { + assert!(!docker_info_cdi_enabled(None)); + assert!(!docker_info_cdi_enabled(Some(&SystemInfo::default()))); + + let info = SystemInfo { + cdi_spec_dirs: Some(vec!["/etc/cdi".to_string()]), + ..Default::default() + }; + assert!(docker_info_cdi_enabled(Some(&info))); + } } diff --git a/crates/openshell-bootstrap/src/lib.rs b/crates/openshell-bootstrap/src/lib.rs index 53f659fc6..bd537ea0e 100644 --- a/crates/openshell-bootstrap/src/lib.rs +++ b/crates/openshell-bootstrap/src/lib.rs @@ -119,6 +119,9 @@ pub struct DeployOptions { /// - `["auto"]` — resolved at deploy time: CDI if enabled on the daemon, else the non-CDI fallback /// - `[cdi-ids…]` — CDI DeviceRequest with the given device IDs pub gpu: Vec, + /// Detect NVIDIA GPU support during deploy and enable passthrough when no + /// explicit GPU device IDs were supplied. + pub gpu_auto_detect: bool, /// When true, destroy any existing gateway resources before deploying. /// When false, an existing gateway is left as-is and deployment is /// skipped (the caller is responsible for prompting the user first). @@ -138,6 +141,7 @@ impl DeployOptions { registry_username: None, registry_token: None, gpu: vec![], + gpu_auto_detect: false, recreate: false, } } @@ -202,6 +206,13 @@ impl DeployOptions { self } + /// Enable or disable automatic GPU passthrough detection. + #[must_use] + pub fn with_gpu_auto_detect(mut self, auto_detect: bool) -> Self { + self.gpu_auto_detect = auto_detect; + self + } + /// Set whether to destroy and recreate existing gateway resources. #[must_use] pub fn with_recreate(mut self, recreate: bool) -> Self { @@ -270,7 +281,8 @@ where let disable_gateway_auth = options.disable_gateway_auth; let registry_username = options.registry_username; let registry_token = options.registry_token; - let gpu = options.gpu; + let mut gpu = options.gpu; + let gpu_auto_detect = options.gpu_auto_detect; let recreate = options.recreate; // Wrap on_log in Arc> so we can share it with pull_remote_image @@ -296,17 +308,22 @@ where (preflight.docker, None) }; - // CDI is considered enabled when the daemon reports at least one CDI spec - // directory via `GET /info` (`SystemInfo.CDISpecDirs`). An empty list or - // missing field means CDI is not configured and we fall back to the legacy - // NVIDIA `DeviceRequest` (driver="nvidia"). Detection is best-effort — - // failure to query daemon info is non-fatal. - let cdi_supported = target_docker - .info() - .await - .ok() - .and_then(|info| info.cdi_spec_dirs) - .is_some_and(|dirs| !dirs.is_empty()); + // GPU discovery is best-effort. Explicit `--gpu` still uses the legacy + // CDI-enabled check below, while auto-detection only enables GPU when the + // daemon reports NVIDIA CDI devices or the local host has NVIDIA devices + // plus the NVIDIA Docker runtime. + let docker_info = target_docker.info().await.ok(); + let cdi_supported = docker::docker_info_cdi_enabled(docker_info.as_ref()); + if gpu_auto_detect && gpu.is_empty() { + let detected_gpu = docker::auto_detect_gpu_device_ids( + docker_info.as_ref(), + remote_opts.is_none() && docker::local_nvidia_devices_present(), + ); + if !detected_gpu.is_empty() { + log("[status] Detected NVIDIA GPU support".to_string()); + gpu = detected_gpu; + } + } // If an existing gateway is found, decide how to proceed: // - recreate: destroy everything and start fresh diff --git a/crates/openshell-cli/src/main.rs b/crates/openshell-cli/src/main.rs index 3502c2b07..5214eaca8 100644 --- a/crates/openshell-cli/src/main.rs +++ b/crates/openshell-cli/src/main.rs @@ -821,6 +821,14 @@ enum GatewayCommands { /// (`--gpus all`) otherwise. #[arg(long)] gpu: bool, + + /// Disable automatic NVIDIA GPU passthrough detection. + /// + /// By default, `gateway start` enables GPU passthrough when Docker + /// reports NVIDIA CDI devices, or when the local host exposes NVIDIA + /// devices and the NVIDIA Docker runtime is configured. + #[arg(long, conflicts_with = "gpu")] + no_gpu: bool, }, /// Stop the gateway (preserves state). @@ -1733,12 +1741,14 @@ async fn main() -> Result<()> { registry_username, registry_token, gpu, + no_gpu, } => { let gpu = if gpu { vec!["auto".to_string()] } else { vec![] }; + let gpu_auto_detect = !no_gpu && gpu.is_empty(); run::gateway_admin_deploy( &name, remote.as_deref(), @@ -1751,6 +1761,7 @@ async fn main() -> Result<()> { registry_username.as_deref(), registry_token.as_deref(), gpu, + gpu_auto_detect, ) .await?; } @@ -3408,6 +3419,31 @@ mod tests { } } + #[test] + fn gateway_start_parses_no_gpu_flag() { + let cli = Cli::try_parse_from(["openshell", "gateway", "start", "--no-gpu"]) + .expect("gateway start --no-gpu should parse"); + + match cli.command { + Some(Commands::Gateway { + command: Some(GatewayCommands::Start { no_gpu, gpu, .. }), + }) => { + assert!(no_gpu); + assert!(!gpu); + } + other => panic!("expected gateway start command, got: {other:?}"), + } + } + + #[test] + fn gateway_start_rejects_gpu_and_no_gpu_together() { + let result = Cli::try_parse_from(["openshell", "gateway", "start", "--gpu", "--no-gpu"]); + assert!( + result.is_err(), + "gateway start should reject conflicting GPU flags" + ); + } + // ── sandbox create arg-shape tests ─────────────────────────────────── /// Verify that `sandbox create --name ` still parses as a named diff --git a/crates/openshell-cli/src/run.rs b/crates/openshell-cli/src/run.rs index 87489014a..2c3d8e06f 100644 --- a/crates/openshell-cli/src/run.rs +++ b/crates/openshell-cli/src/run.rs @@ -1436,6 +1436,7 @@ pub async fn gateway_admin_deploy( registry_username: Option<&str>, registry_token: Option<&str>, gpu: Vec, + gpu_auto_detect: bool, ) -> Result<()> { let location = if remote.is_some() { "remote" } else { "local" }; @@ -1489,6 +1490,7 @@ pub async fn gateway_admin_deploy( .with_disable_tls(disable_tls) .with_disable_gateway_auth(disable_gateway_auth) .with_gpu(gpu) + .with_gpu_auto_detect(gpu_auto_detect) .with_recreate(recreate); if let Some(opts) = remote_opts { options = options.with_remote(opts); diff --git a/crates/openshell-driver-docker/src/lib.rs b/crates/openshell-driver-docker/src/lib.rs index b4371ce7d..64ec3ccf6 100644 --- a/crates/openshell-driver-docker/src/lib.rs +++ b/crates/openshell-driver-docker/src/lib.rs @@ -8,8 +8,8 @@ use bollard::Docker; use bollard::errors::Error as BollardError; use bollard::models::{ - ContainerCreateBody, ContainerSummary, ContainerSummaryStateEnum, HostConfig, Mount, - MountTypeEnum, RestartPolicy, RestartPolicyNameEnum, + ContainerCreateBody, ContainerSummary, ContainerSummaryStateEnum, DeviceInfo, DeviceRequest, + HostConfig, Mount, MountTypeEnum, RestartPolicy, RestartPolicyNameEnum, Runtime, SystemInfo, }; use bollard::query_parameters::{ CreateContainerOptionsBuilder, CreateImageOptions, DownloadFromContainerOptionsBuilder, @@ -17,7 +17,7 @@ use bollard::query_parameters::{ }; use bytes::Bytes; use futures::{Stream, StreamExt}; -use openshell_core::config::DEFAULT_STOP_TIMEOUT_SECS; +use openshell_core::config::{CDI_GPU_DEVICE_ALL, DEFAULT_STOP_TIMEOUT_SECS}; use openshell_core::proto::compute::v1::{ CreateSandboxRequest, CreateSandboxResponse, DeleteSandboxRequest, DeleteSandboxResponse, DriverCondition, DriverSandbox, DriverSandboxStatus, DriverSandboxTemplate, @@ -159,6 +159,7 @@ struct DockerDriverRuntimeConfig { supervisor_bin: PathBuf, guest_tls: Option, daemon_version: String, + gpu_device_request: Option, } #[derive(Clone)] @@ -175,6 +176,12 @@ struct DockerResourceLimits { memory_bytes: Option, } +#[derive(Debug, Clone, PartialEq, Eq)] +enum DockerGpuDeviceRequest { + Cdi(Vec), + Nvidia, +} + type WatchStream = Pin> + Send + 'static>>; @@ -195,6 +202,12 @@ impl DockerComputeDriver { let version = docker.version().await.map_err(|err| { Error::execution(format!("failed to query Docker daemon version: {err}")) })?; + let gpu_device_request = docker + .info() + .await + .ok() + .as_ref() + .and_then(docker_gpu_device_request_from_info); let daemon_arch = normalize_docker_arch(version.arch.as_deref().unwrap_or_default()); let supervisor_bin = resolve_supervisor_bin(&docker, docker_config, &daemon_arch).await?; let guest_tls = docker_guest_tls_paths(config, docker_config)?; @@ -212,6 +225,7 @@ impl DockerComputeDriver { supervisor_bin, guest_tls, daemon_version: version.version.unwrap_or_else(|| "unknown".to_string()), + gpu_device_request, }, events: broadcast::channel(WATCH_BUFFER).0, supervisor_readiness, @@ -230,11 +244,11 @@ impl DockerComputeDriver { driver_name: "docker".to_string(), driver_version: self.config.daemon_version.clone(), default_image: self.config.default_image.clone(), - supports_gpu: false, + supports_gpu: self.config.gpu_device_request.is_some(), } } - fn validate_sandbox(sandbox: &DriverSandbox) -> Result<(), Status> { + fn validate_sandbox(sandbox: &DriverSandbox, gpu_supported: bool) -> Result<(), Status> { let spec = sandbox .spec .as_ref() @@ -249,9 +263,9 @@ impl DockerComputeDriver { "docker sandboxes require a template image", )); } - if spec.gpu { + if spec.gpu && !gpu_supported { return Err(Status::failed_precondition( - "docker compute driver does not support gpu sandboxes", + "GPU sandbox requested, but the Docker daemon does not expose NVIDIA GPU support", )); } if !template.agent_socket_path.trim().is_empty() { @@ -299,7 +313,7 @@ impl DockerComputeDriver { } async fn create_sandbox_inner(&self, sandbox: &DriverSandbox) -> Result<(), Status> { - Self::validate_sandbox(sandbox)?; + Self::validate_sandbox(sandbox, self.config.gpu_device_request.is_some())?; if self .find_managed_container_summary(&sandbox.id, &sandbox.name) @@ -673,7 +687,7 @@ impl ComputeDriver for DockerComputeDriver { .into_inner() .sandbox .ok_or_else(|| Status::invalid_argument("sandbox is required"))?; - Self::validate_sandbox(&sandbox)?; + Self::validate_sandbox(&sandbox, self.config.gpu_device_request.is_some())?; Ok(Response::new(ValidateSandboxCreateResponse {})) } @@ -888,6 +902,7 @@ fn build_container_create_body( .as_ref() .ok_or_else(|| Status::invalid_argument("sandbox.spec.template is required"))?; let resource_limits = docker_resource_limits(template)?; + let device_requests = docker_gpu_device_requests(spec.gpu, config)?; let mut labels = template.labels.clone(); labels.insert( MANAGED_BY_LABEL_KEY.to_string(), @@ -917,6 +932,7 @@ fn build_container_create_body( nano_cpus: resource_limits.nano_cpus, memory: resource_limits.memory_bytes, mounts: Some(build_mounts(config)), + device_requests, restart_policy: Some(RestartPolicy { name: Some(RestartPolicyNameEnum::UNLESS_STOPPED), maximum_retry_count: None, @@ -1004,6 +1020,101 @@ fn docker_resource_limits( }) } +fn docker_gpu_device_requests( + gpu_requested: bool, + config: &DockerDriverRuntimeConfig, +) -> Result>, Status> { + if !gpu_requested { + return Ok(None); + } + + let Some(request) = config.gpu_device_request.as_ref() else { + return Err(Status::failed_precondition( + "GPU sandbox requested, but the Docker daemon does not expose NVIDIA GPU support", + )); + }; + + Ok(Some(vec![match request { + DockerGpuDeviceRequest::Cdi(device_ids) => DeviceRequest { + driver: Some("cdi".to_string()), + device_ids: Some(device_ids.clone()), + ..Default::default() + }, + DockerGpuDeviceRequest::Nvidia => DeviceRequest { + driver: Some("nvidia".to_string()), + count: Some(-1), + capabilities: Some(vec![vec![ + "gpu".to_string(), + "utility".to_string(), + "compute".to_string(), + ]]), + ..Default::default() + }, + }])) +} + +fn docker_gpu_device_request_from_info(info: &SystemInfo) -> Option { + let cdi_device_ids = nvidia_cdi_device_ids(info); + if !cdi_device_ids.is_empty() { + return Some(DockerGpuDeviceRequest::Cdi(cdi_device_ids)); + } + + if docker_info_has_nvidia_runtime(info) { + return Some(DockerGpuDeviceRequest::Nvidia); + } + + None +} + +fn nvidia_cdi_device_ids(info: &SystemInfo) -> Vec { + let Some(devices) = info.discovered_devices.as_ref() else { + return Vec::new(); + }; + + let mut ids = devices + .iter() + .filter_map(nvidia_cdi_device_id) + .collect::>(); + ids.sort(); + ids.dedup(); + + if ids.iter().any(|id| id == CDI_GPU_DEVICE_ALL) { + vec![CDI_GPU_DEVICE_ALL.to_string()] + } else { + ids + } +} + +fn nvidia_cdi_device_id(device: &DeviceInfo) -> Option { + let id = device.id.as_ref()?; + if id.contains("nvidia.com/gpu=") + || (id.contains("/gpu=") && device.source.as_deref().is_some_and(contains_nvidia)) + { + return Some(id.clone()); + } + None +} + +fn docker_info_has_nvidia_runtime(info: &SystemInfo) -> bool { + info.runtimes.as_ref().is_some_and(|runtimes| { + runtimes + .iter() + .any(|(name, runtime)| is_nvidia_runtime(name, runtime)) + }) +} + +fn is_nvidia_runtime(name: &str, runtime: &Runtime) -> bool { + contains_nvidia(name) + || runtime + .path + .as_deref() + .is_some_and(|path| path.contains("nvidia-container-runtime")) +} + +fn contains_nvidia(value: &str) -> bool { + value.to_ascii_lowercase().contains("nvidia") +} + #[allow(clippy::cast_possible_truncation)] fn parse_cpu_limit(value: &str) -> Result, Status> { let value = value.trim(); diff --git a/crates/openshell-driver-docker/src/tests.rs b/crates/openshell-driver-docker/src/tests.rs index b20fbf5ce..ae53474e7 100644 --- a/crates/openshell-driver-docker/src/tests.rs +++ b/crates/openshell-driver-docker/src/tests.rs @@ -51,6 +51,7 @@ fn runtime_config() -> DockerDriverRuntimeConfig { key: PathBuf::from("/tmp/tls.key"), }), daemon_version: "28.0.0".to_string(), + gpu_device_request: None, } } @@ -171,6 +172,100 @@ fn build_container_create_body_clears_inherited_cmd() { ); } +#[test] +fn validate_sandbox_rejects_gpu_when_daemon_has_no_gpu_support() { + let mut sandbox = test_sandbox(); + sandbox.spec.as_mut().unwrap().gpu = true; + + let err = DockerComputeDriver::validate_sandbox(&sandbox, false).unwrap_err(); + + assert_eq!(err.code(), tonic::Code::FailedPrecondition); + assert!(err.message().contains("NVIDIA GPU support")); +} + +#[test] +fn build_container_create_body_adds_legacy_gpu_device_request() { + let mut config = runtime_config(); + config.gpu_device_request = Some(DockerGpuDeviceRequest::Nvidia); + let mut sandbox = test_sandbox(); + sandbox.spec.as_mut().unwrap().gpu = true; + + let create_body = build_container_create_body(&sandbox, &config).unwrap(); + let device_requests = create_body + .host_config + .unwrap() + .device_requests + .expect("GPU sandbox should request Docker devices"); + + assert_eq!(device_requests.len(), 1); + assert_eq!(device_requests[0].driver.as_deref(), Some("nvidia")); + assert_eq!(device_requests[0].count, Some(-1)); + assert_eq!( + device_requests[0].capabilities, + Some(vec![vec![ + "gpu".to_string(), + "utility".to_string(), + "compute".to_string() + ]]) + ); +} + +#[test] +fn build_container_create_body_adds_cdi_gpu_device_request() { + let mut config = runtime_config(); + config.gpu_device_request = Some(DockerGpuDeviceRequest::Cdi(vec![ + "nvidia.com/gpu=all".to_string(), + ])); + let mut sandbox = test_sandbox(); + sandbox.spec.as_mut().unwrap().gpu = true; + + let create_body = build_container_create_body(&sandbox, &config).unwrap(); + let device_requests = create_body + .host_config + .unwrap() + .device_requests + .expect("GPU sandbox should request Docker devices"); + + assert_eq!(device_requests.len(), 1); + assert_eq!(device_requests[0].driver.as_deref(), Some("cdi")); + assert_eq!( + device_requests[0].device_ids, + Some(vec!["nvidia.com/gpu=all".to_string()]) + ); +} + +#[test] +fn docker_gpu_device_request_prefers_cdi_devices_over_legacy_runtime() { + let info = SystemInfo { + discovered_devices: Some(vec![DeviceInfo { + source: Some("cdi".to_string()), + id: Some("nvidia.com/gpu=all".to_string()), + }]), + runtimes: Some(HashMap::from([("nvidia".to_string(), Runtime::default())])), + ..Default::default() + }; + + assert_eq!( + docker_gpu_device_request_from_info(&info), + Some(DockerGpuDeviceRequest::Cdi(vec![ + "nvidia.com/gpu=all".to_string() + ])) + ); +} + +#[test] +fn docker_gpu_device_request_detects_legacy_nvidia_runtime() { + let info = SystemInfo { + runtimes: Some(HashMap::from([("nvidia".to_string(), Runtime::default())])), + ..Default::default() + }; + + assert_eq!( + docker_gpu_device_request_from_info(&info), + Some(DockerGpuDeviceRequest::Nvidia) + ); +} + #[test] fn require_sandbox_identifier_rejects_when_id_and_name_are_empty() { // Regression test: `delete_sandbox` (and the other identifier-keyed diff --git a/docs/sandboxes/manage-gateways.mdx b/docs/sandboxes/manage-gateways.mdx index a20a27c20..6058e53b3 100644 --- a/docs/sandboxes/manage-gateways.mdx +++ b/docs/sandboxes/manage-gateways.mdx @@ -160,7 +160,8 @@ openshell gateway info --name my-remote-cluster | Flag | Purpose | |---|---| -| `--gpu` | Enable NVIDIA GPU passthrough. Requires NVIDIA drivers and the Container Toolkit on the host. OpenShell auto-selects CDI when enabled on the daemon and falls back to Docker's NVIDIA GPU request path (`--gpus all`) otherwise. | +| `--gpu` | Force NVIDIA GPU passthrough. Requires NVIDIA drivers and the Container Toolkit on the host. OpenShell auto-selects CDI when enabled on the daemon and falls back to Docker's NVIDIA GPU request path (`--gpus all`) otherwise. | +| `--no-gpu` | Disable automatic GPU passthrough detection. By default, `gateway start` enables GPU passthrough when Docker reports NVIDIA CDI devices, or when the local host exposes NVIDIA devices and the NVIDIA Docker runtime is configured. | | `--plaintext` | Listen on HTTP instead of mTLS. Use behind a TLS-terminating reverse proxy. | | `--disable-gateway-auth` | Skip mTLS client certificate checks. Use when a reverse proxy cannot forward client certs. | | `--registry-username` | Username for registry authentication. Defaults to `__token__` when `--registry-token` is set. Only needed for private registries. Also configurable with `OPENSHELL_REGISTRY_USERNAME`. | diff --git a/docs/sandboxes/manage-sandboxes.mdx b/docs/sandboxes/manage-sandboxes.mdx index fb24bae9b..76d938b40 100644 --- a/docs/sandboxes/manage-sandboxes.mdx +++ b/docs/sandboxes/manage-sandboxes.mdx @@ -39,6 +39,8 @@ To request GPU resources, add `--gpu`: openshell sandbox create --gpu -- claude ``` +On GPU-capable Docker hosts, `openshell gateway start` enables gateway GPU passthrough automatically. Use `openshell gateway start --no-gpu` to disable auto-detection, or pass `openshell gateway start --gpu` to force GPU passthrough. + ### Custom Containers Use `--from` to create a sandbox from a pre-built community package, a local directory, or a container image: