From c2603f99bd2efc81a72a2184409b8c6ede70e11f Mon Sep 17 00:00:00 2001 From: Piotr Mlocek Date: Thu, 12 Mar 2026 01:32:22 -0700 Subject: [PATCH] feat(sandbox): add gpu sandbox scheduling support Allow sandbox creation to request GPU resources explicitly or infer them from GPU image names. This wires GPU intent through bootstrap, validates gateway support, and adds dedicated GPU E2E coverage for follow-up cluster testing. --- architecture/gateway-single-node.md | 1 + architecture/sandbox-custom-containers.md | 4 + crates/openshell-cli/src/bootstrap.rs | 2 + crates/openshell-cli/src/main.rs | 13 + crates/openshell-cli/src/run.rs | 78 ++++- crates/openshell-server/src/grpc.rs | 21 ++ crates/openshell-server/src/sandbox/mod.rs | 339 ++++++++++++++++++++- docs/reference/cli.md | 258 ++++++++++++++++ docs/sandboxes/create-and-manage.md | 10 + e2e/python/conftest.py | 13 + e2e/python/test_sandbox_gpu.py | 30 ++ e2e/rust/tests/cli_smoke.rs | 6 +- mise.toml | 2 +- proto/datamodel.proto | 2 + pyproject.toml | 5 +- tasks/test.toml | 12 +- 16 files changed, 779 insertions(+), 17 deletions(-) create mode 100644 docs/reference/cli.md create mode 100644 e2e/python/test_sandbox_gpu.py diff --git a/architecture/gateway-single-node.md b/architecture/gateway-single-node.md index f708de69f..c2f6f96e0 100644 --- a/architecture/gateway-single-node.md +++ b/architecture/gateway-single-node.md @@ -377,6 +377,7 @@ When `openshell sandbox create` cannot connect to a gateway (connection refused, 1. `should_attempt_bootstrap()` in `crates/openshell-cli/src/bootstrap.rs` checks the error type. It returns `true` for connectivity errors and missing default TLS materials, but `false` for TLS handshake/auth errors. 2. If running in a terminal, the user is prompted to confirm. 3. `run_bootstrap()` deploys a gateway named `"openshell"`, sets it as active, and returns fresh `TlsOptions` pointing to the newly-written mTLS certs. +4. When `sandbox create` requests GPU explicitly (`--gpu`) or infers it from an image whose final name component contains `gpu` (such as `nvidia-gpu`), the bootstrap path enables gateway GPU support before retrying sandbox creation. ## Container Environment Variables diff --git a/architecture/sandbox-custom-containers.md b/architecture/sandbox-custom-containers.md index 72d7e141b..27ee2a916 100644 --- a/architecture/sandbox-custom-containers.md +++ b/architecture/sandbox-custom-containers.md @@ -24,6 +24,10 @@ The CLI classifies the value in this order: The community registry prefix defaults to `ghcr.io/nvidia/openshell-community/sandboxes` and can be overridden with the `OPENSHELL_COMMUNITY_REGISTRY` environment variable. +### GPU image-name detection + +`sandbox create` also infers GPU intent from the final image name. The current rule matches when the last image name component contains `gpu` (for example `ghcr.io/nvidia/openshell-community/sandboxes/nvidia-gpu:latest` or `registry.example.com/team/my-gpu-image:latest`). When that rule matches, the sandbox request is treated the same as passing `--gpu`. + ### Dockerfile build flow When `--from` points to a Dockerfile or directory, the CLI: diff --git a/crates/openshell-cli/src/bootstrap.rs b/crates/openshell-cli/src/bootstrap.rs index 7a50ed9a5..2caa20b96 100644 --- a/crates/openshell-cli/src/bootstrap.rs +++ b/crates/openshell-cli/src/bootstrap.rs @@ -122,6 +122,7 @@ fn resolve_bootstrap_name() -> String { pub async fn run_bootstrap( remote: Option<&str>, ssh_key: Option<&str>, + gpu: bool, ) -> Result<(TlsOptions, String, String)> { let gateway_name = resolve_bootstrap_name(); let location = if remote.is_some() { "remote" } else { "local" }; @@ -159,6 +160,7 @@ pub async fn run_bootstrap( { options = options.with_registry_token(token); } + options = options.with_gpu(gpu); let handle = deploy_gateway_with_panel(options, &gateway_name, location).await?; let server = handle.gateway_endpoint().to_string(); diff --git a/crates/openshell-cli/src/main.rs b/crates/openshell-cli/src/main.rs index 868dee5db..419c7e753 100644 --- a/crates/openshell-cli/src/main.rs +++ b/crates/openshell-cli/src/main.rs @@ -1051,6 +1051,14 @@ enum SandboxCommands { #[arg(long, value_enum, conflicts_with = "no_keep")] editor: Option, + /// Request GPU resources for the sandbox. + /// + /// When no gateway is running, auto-bootstrap starts a GPU-enabled + /// gateway. GPU intent is also inferred automatically for known + /// GPU-designated image names such as `nvidia-gpu`. + #[arg(long)] + gpu: bool, + /// SSH destination for remote bootstrap (e.g., user@hostname). /// Only used when no cluster exists yet; ignored if a cluster is /// already active. @@ -1791,6 +1799,7 @@ async fn main() -> Result<()> { keep, no_keep, editor, + gpu, remote, ssh_key, providers, @@ -1869,6 +1878,8 @@ async fn main() -> Result<()> { upload_spec.as_ref(), keep, editor, + gpu, + editor, remote.as_deref(), ssh_key.as_deref(), &providers, @@ -1890,6 +1901,8 @@ async fn main() -> Result<()> { upload_spec.as_ref(), keep, editor, + gpu, + editor, remote.as_deref(), ssh_key.as_deref(), &providers, diff --git a/crates/openshell-cli/src/run.rs b/crates/openshell-cli/src/run.rs index 1270e0597..3d0c7ebdb 100644 --- a/crates/openshell-cli/src/run.rs +++ b/crates/openshell-cli/src/run.rs @@ -1700,6 +1700,7 @@ pub async fn sandbox_create_with_bootstrap( from: Option<&str>, upload: Option<&(String, Option, bool)>, keep: bool, + gpu: bool, editor: Option, remote: Option<&str>, ssh_key: Option<&str>, @@ -1718,7 +1719,9 @@ pub async fn sandbox_create_with_bootstrap( Or deploy a new gateway: openshell gateway start" )); } - let (tls, server, gateway_name) = crate::bootstrap::run_bootstrap(remote, ssh_key).await?; + let requested_gpu = gpu || from.is_some_and(source_requests_gpu); + let (tls, server, gateway_name) = + crate::bootstrap::run_bootstrap(remote, ssh_key, requested_gpu).await?; // Disable bootstrap inside sandbox_create so that a transient connection // failure right after deploy does not trigger a second bootstrap attempt. sandbox_create( @@ -1728,6 +1731,7 @@ pub async fn sandbox_create_with_bootstrap( &gateway_name, upload, keep, + gpu, editor, remote, ssh_key, @@ -1778,6 +1782,7 @@ pub async fn sandbox_create( gateway_name: &str, upload: Option<&(String, Option, bool)>, keep: bool, + gpu: bool, editor: Option, remote: Option<&str>, ssh_key: Option<&str>, @@ -1829,7 +1834,9 @@ pub async fn sandbox_create( eprintln!(); return Err(err); } - let (new_tls, new_server, _) = crate::bootstrap::run_bootstrap(remote, ssh_key).await?; + let requested_gpu = gpu || from.is_some_and(source_requests_gpu); + let (new_tls, new_server, _) = + crate::bootstrap::run_bootstrap(remote, ssh_key, requested_gpu).await?; let c = grpc_client(&new_server, &new_tls) .await .wrap_err("bootstrap succeeded but failed to connect to gateway")?; @@ -1855,6 +1862,7 @@ pub async fn sandbox_create( } None => None, }; + let requested_gpu = gpu || image.as_deref().is_some_and(image_requests_gpu); let inferred_types: Vec = inferred_provider_type(command).into_iter().collect(); let configured_providers = ensure_required_providers( @@ -1874,6 +1882,7 @@ pub async fn sandbox_create( let request = CreateSandboxRequest { spec: Some(SandboxSpec { + gpu: requested_gpu, policy, providers: configured_providers, template, @@ -2343,6 +2352,30 @@ fn resolve_from(value: &str) -> Result { Ok(ResolvedSource::Image(format!("{prefix}/{value}:latest"))) } +fn source_requests_gpu(source: &str) -> bool { + if let Ok(resolved) = resolve_from(source) { + match resolved { + ResolvedSource::Image(image) => image_requests_gpu(&image), + ResolvedSource::Dockerfile { .. } => false, + } + } else { + false + } +} + +fn image_requests_gpu(image: &str) -> bool { + let image_name = image + .rsplit('/') + .next() + .unwrap_or(image) + .split([':', '@']) + .next() + .unwrap_or(image) + .to_ascii_lowercase(); + + image_name.contains("gpu") +} + /// Build a Dockerfile and push the resulting image into the gateway. /// /// Returns the image tag that was built so the caller can use it for sandbox @@ -4187,8 +4220,9 @@ mod tests { use super::{ GatewayControlTarget, TlsOptions, format_gateway_select_header, format_gateway_select_items, gateway_auth_label, gateway_select_with, gateway_type_label, - git_sync_files, http_health_check, inferred_provider_type, parse_credential_pairs, - resolve_gateway_control_target_from, sandbox_should_persist, + git_sync_files, http_health_check, image_requests_gpu, inferred_provider_type, + parse_credential_pairs, resolve_gateway_control_target_from, sandbox_should_persist, + source_requests_gpu, }; use crate::TEST_ENV_LOCK; use hyper::StatusCode; @@ -4357,6 +4391,42 @@ mod tests { assert!(sandbox_should_persist(false, Some(8080))); } + #[test] + fn image_requests_gpu_matches_known_gpu_image_names() { + for image in [ + "ghcr.io/nvidia/openshell-community/sandboxes/nvidia-gpu:latest", + "registry.example.com/team/gpu:dev", + "nvcr.io/example/my-gpu-image@sha256:deadbeef", + ] { + assert!( + image_requests_gpu(image), + "expected GPU detection for {image}" + ); + } + } + + #[test] + fn image_requests_gpu_ignores_non_gpu_image_names() { + for image in [ + "ghcr.io/nvidia/openshell-community/sandboxes/base:latest", + "registry.example.com/gpu/team/base:latest", + "registry.example.com/team/openclaw:latest", + "cuda-toolkit:latest", + "registry.example.com/team/graphics:latest", + ] { + assert!( + !image_requests_gpu(image), + "did not expect GPU detection for {image}" + ); + } + } + + #[test] + fn source_requests_gpu_detects_known_community_gpu_name() { + assert!(source_requests_gpu("nvidia-gpu")); + assert!(!source_requests_gpu("base")); + } + fn init_git_repo(path: &Path) { let status = Command::new("git") .args(["init"]) diff --git a/crates/openshell-server/src/grpc.rs b/crates/openshell-server/src/grpc.rs index 8846d60c9..615318821 100644 --- a/crates/openshell-server/src/grpc.rs +++ b/crates/openshell-server/src/grpc.rs @@ -171,6 +171,17 @@ impl OpenShell for OpenShellService { template.image = self.state.sandbox_client.default_image().to_string(); } + if spec.gpu { + self.state + .sandbox_client + .validate_gpu_support() + .await + .map_err(|status| { + warn!(error = %status, "Rejecting GPU sandbox request"); + status + })?; + } + // Ensure process identity defaults to "sandbox" when missing or // empty, then validate policy safety before persisting. if let Some(ref mut policy) = spec.policy { @@ -4268,6 +4279,16 @@ mod tests { SandboxSpec::default() } + #[test] + fn validate_sandbox_spec_accepts_gpu_flag() { + let spec = SandboxSpec { + gpu: true, + ..Default::default() + }; + + assert!(validate_sandbox_spec("gpu-sandbox", &spec).is_ok()); + } + #[test] fn validate_sandbox_spec_accepts_empty_defaults() { assert!(validate_sandbox_spec("", &default_spec()).is_ok()); diff --git a/crates/openshell-server/src/sandbox/mod.rs b/crates/openshell-server/src/sandbox/mod.rs index 69429cacc..78896e0fe 100644 --- a/crates/openshell-server/src/sandbox/mod.rs +++ b/crates/openshell-server/src/sandbox/mod.rs @@ -5,8 +5,8 @@ use crate::persistence::{ObjectId, ObjectName, ObjectType, Store}; use futures::{StreamExt, TryStreamExt}; -use k8s_openapi::api::core::v1::Pod; -use kube::api::{Api, ApiResource, DeleteParams, PostParams}; +use k8s_openapi::api::core::v1::{Node, Pod}; +use kube::api::{Api, ApiResource, DeleteParams, ListParams, PostParams}; use kube::core::gvk::GroupVersionKind; use kube::core::{DynamicObject, ObjectMeta}; use kube::runtime::watcher::{self, Event}; @@ -31,6 +31,9 @@ pub const SANDBOX_KIND: &str = "Sandbox"; const SANDBOX_ID_LABEL: &str = "openshell.ai/sandbox-id"; const SANDBOX_MANAGED_LABEL: &str = "openshell.ai/managed-by"; const SANDBOX_MANAGED_VALUE: &str = "openshell"; +const GPU_RUNTIME_CLASS_NAME: &str = "nvidia"; +const GPU_RESOURCE_NAME: &str = "nvidia.com/gpu"; +const GPU_RESOURCE_QUANTITY: &str = "1"; #[derive(Clone)] pub struct SandboxClient { @@ -118,6 +121,47 @@ impl SandboxClient { Api::namespaced_with(self.client.clone(), &self.namespace, &resource) } + pub async fn validate_gpu_support(&self) -> Result<(), tonic::Status> { + let runtime_classes: Api = Api::all_with( + self.client.clone(), + &ApiResource::from_gvk(&GroupVersionKind::gvk("node.k8s.io", "v1", "RuntimeClass")), + ); + + let runtime_class_exists = runtime_classes + .get_opt(GPU_RUNTIME_CLASS_NAME) + .await + .map_err(|err| { + tonic::Status::internal(format!("check GPU runtime class failed: {err}")) + })? + .is_some(); + + if !runtime_class_exists { + return Err(tonic::Status::failed_precondition( + "GPU sandbox requested, but the active gateway is not GPU-enabled. To start a gateway with GPU support run: `openshell gateway start --gpu`", + )); + } + + let nodes: Api = Api::all(self.client.clone()); + let node_list = nodes.list(&ListParams::default()).await.map_err(|err| { + tonic::Status::internal(format!("check GPU node capacity failed: {err}")) + })?; + + let has_gpu_capacity = node_list.items.into_iter().any(|node| { + node.status + .and_then(|status| status.allocatable) + .and_then(|allocatable| allocatable.get(GPU_RESOURCE_NAME).cloned()) + .is_some_and(|quantity| quantity.0 != "0") + }); + + if !has_gpu_capacity { + return Err(tonic::Status::failed_precondition( + "GPU sandbox requested, but the active gateway has no allocatable GPUs. Please refer to documentation and use `openshell doctor` commands to inspect GPU support and gateway configuration.", + )); + } + + Ok(()) + } + pub async fn agent_pod_ip(&self, pod_name: &str) -> Result, KubeError> { let api: Api = Api::namespaced(self.client.clone(), &self.namespace); match api.get(pod_name).await { @@ -467,14 +511,15 @@ async fn handle_applied( .unwrap_or_else(|| client.namespace().to_string()); let deletion_timestamp = obj.metadata.deletion_timestamp.is_some(); - let status = status_from_object(&obj); - let phase = derive_phase(&status, deletion_timestamp); - let existing = store .get_message::(&id) .await .map_err(|e| e.to_string())?; + let mut status = status_from_object(&obj); + rewrite_user_facing_conditions(&mut status, existing.as_ref().and_then(|sandbox| sandbox.spec.as_ref())); + let phase = derive_phase(&status, deletion_timestamp); + // If the record doesn't exist yet, the `create_sandbox` handler may // still be in-flight (it creates the k8s resource first, then writes // to the store). Build a minimal placeholder but never overwrite an @@ -728,6 +773,7 @@ fn sandbox_to_k8s_spec( "podTemplate".to_string(), sandbox_template_to_k8s( template, + spec.gpu, default_image, image_pull_policy, sandbox_id, @@ -760,6 +806,7 @@ fn sandbox_to_k8s_spec( "podTemplate".to_string(), sandbox_template_to_k8s( &SandboxTemplate::default(), + spec.as_ref().is_some_and(|s| s.gpu), default_image, image_pull_policy, sandbox_id, @@ -782,6 +829,7 @@ fn sandbox_to_k8s_spec( #[allow(clippy::too_many_arguments)] fn sandbox_template_to_k8s( template: &SandboxTemplate, + gpu: bool, default_image: &str, image_pull_policy: &str, sandbox_id: &str, @@ -797,6 +845,8 @@ fn sandbox_template_to_k8s( return inject_pod_template( pod_template, template, + gpu, + default_image, image_pull_policy, sandbox_id, sandbox_name, @@ -824,7 +874,12 @@ fn sandbox_template_to_k8s( } let mut spec = serde_json::Map::new(); - if !template.runtime_class_name.is_empty() { + if gpu { + spec.insert( + "runtimeClassName".to_string(), + serde_json::json!(GPU_RUNTIME_CLASS_NAME), + ); + } else if !template.runtime_class_name.is_empty() { spec.insert( "runtimeClassName".to_string(), serde_json::json!(template.runtime_class_name), @@ -891,7 +946,7 @@ fn sandbox_template_to_k8s( ); } - if let Some(resources) = struct_to_json(&template.resources) { + if let Some(resources) = container_resources(template, gpu) { container.insert("resources".to_string(), resources); } spec.insert( @@ -928,6 +983,8 @@ fn sandbox_template_to_k8s( fn inject_pod_template( mut pod_template: serde_json::Value, template: &SandboxTemplate, + gpu: bool, + default_image: &str, image_pull_policy: &str, sandbox_id: &str, sandbox_name: &str, @@ -945,6 +1002,13 @@ fn inject_pod_template( return pod_template; }; + if gpu { + spec.insert( + "runtimeClassName".to_string(), + serde_json::json!(GPU_RUNTIME_CLASS_NAME), + ); + } + // Inject TLS volume at the pod spec level. if !client_tls_secret_name.is_empty() { let volumes = spec @@ -1015,6 +1079,10 @@ fn inject_pod_template( })); } } + + if gpu { + apply_gpu_to_container(container); + } } // Always side-load the supervisor binary from the node filesystem @@ -1023,6 +1091,51 @@ fn inject_pod_template( pod_template } +fn container_resources(template: &SandboxTemplate, gpu: bool) -> Option { + let mut resources = + struct_to_json(&template.resources).unwrap_or_else(|| serde_json::json!({})); + if gpu { + apply_gpu_limit(&mut resources); + } + if resources + .as_object() + .is_some_and(|object| object.is_empty()) + { + None + } else { + Some(resources) + } +} + +fn apply_gpu_to_container(container: &mut serde_json::Value) { + if let Some(container_obj) = container.as_object_mut() { + let resources = container_obj + .entry("resources") + .or_insert_with(|| serde_json::json!({})); + apply_gpu_limit(resources); + } +} + +fn apply_gpu_limit(resources: &mut serde_json::Value) { + let Some(resources_obj) = resources.as_object_mut() else { + *resources = serde_json::json!({}); + return apply_gpu_limit(resources); + }; + + let limits = resources_obj + .entry("limits") + .or_insert_with(|| serde_json::json!({})); + let Some(limits_obj) = limits.as_object_mut() else { + *limits = serde_json::json!({}); + return apply_gpu_limit(resources); + }; + + limits_obj.insert( + GPU_RESOURCE_NAME.to_string(), + serde_json::json!(GPU_RESOURCE_QUANTITY), + ); +} + #[allow(clippy::too_many_arguments)] fn update_container_env( container: &mut serde_json::Value, @@ -1241,6 +1354,24 @@ fn condition_from_value(value: &serde_json::Value) -> Option { }) } +fn rewrite_user_facing_conditions(status: &mut Option, spec: Option<&SandboxSpec>) { + let gpu_requested = spec.is_some_and(|sandbox_spec| sandbox_spec.gpu); + if !gpu_requested { + return; + } + + if let Some(status) = status { + for condition in &mut status.conditions { + if condition.r#type == "Ready" + && condition.status.eq_ignore_ascii_case("false") + && condition.reason.eq_ignore_ascii_case("Unschedulable") + { + condition.message = "GPU sandbox could not be scheduled on the active gateway. Another GPU sandbox may already be using the available GPU, or the gateway may not currently be able to satisfy GPU placement. Please refer to documentation and use `openshell doctor` commands to inspect GPU support and gateway configuration.".to_string(); + } + } + } +} + fn derive_phase(status: &Option, deleting: bool) -> SandboxPhase { if deleting { return SandboxPhase::Deleting; @@ -1287,6 +1418,7 @@ fn is_terminal_failure_condition(condition: &SandboxCondition) -> bool { #[cfg(test)] mod tests { use super::*; + use prost_types::{Struct, Value, value::Kind}; fn make_condition(reason: &str, message: &str) -> SandboxCondition { SandboxCondition { @@ -1402,6 +1534,65 @@ mod tests { assert_eq!(derive_phase(&status, false), SandboxPhase::Error); } + #[test] + fn rewrite_user_facing_conditions_rewrites_gpu_unschedulable_message() { + let mut status = Some(SandboxStatus { + sandbox_name: "test".to_string(), + agent_pod: "test-pod".to_string(), + agent_fd: String::new(), + sandbox_fd: String::new(), + conditions: vec![SandboxCondition { + r#type: "Ready".to_string(), + status: "False".to_string(), + reason: "Unschedulable".to_string(), + message: "0/1 nodes are available: 1 Insufficient nvidia.com/gpu.".to_string(), + last_transition_time: String::new(), + }], + }); + + rewrite_user_facing_conditions( + &mut status, + Some(&SandboxSpec { + gpu: true, + ..Default::default() + }), + ); + + let message = &status.unwrap().conditions[0].message; + assert_eq!( + message, + "GPU sandbox could not be scheduled on the active gateway. Another GPU sandbox may already be using the available GPU, or the gateway may not currently be able to satisfy GPU placement. Please refer to documentation and use `openshell doctor` commands to inspect GPU support and gateway configuration." + ); + } + + #[test] + fn rewrite_user_facing_conditions_leaves_non_gpu_unschedulable_message_unchanged() { + let original = "0/1 nodes are available: 1 Insufficient cpu."; + let mut status = Some(SandboxStatus { + sandbox_name: "test".to_string(), + agent_pod: "test-pod".to_string(), + agent_fd: String::new(), + sandbox_fd: String::new(), + conditions: vec![SandboxCondition { + r#type: "Ready".to_string(), + status: "False".to_string(), + reason: "Unschedulable".to_string(), + message: original.to_string(), + last_transition_time: String::new(), + }], + }); + + rewrite_user_facing_conditions( + &mut status, + Some(&SandboxSpec { + gpu: false, + ..Default::default() + }), + ); + + assert_eq!(status.unwrap().conditions[0].message, original); + } + #[test] fn derive_phase_returns_ready_for_ready_true() { let status = Some(SandboxStatus { @@ -1590,4 +1781,138 @@ mod tests { "OPENSHELL_TLS_KEY path '{tls_key}' must start with mount path '{TLS_MOUNT_PATH}'" ); } + + fn string_value(value: &str) -> Value { + Value { + kind: Some(Kind::StringValue(value.to_string())), + } + } + + #[test] + fn gpu_sandbox_adds_runtime_class_and_gpu_limit() { + let pod_template = sandbox_template_to_k8s( + &SandboxTemplate::default(), + true, + "openshell/sandbox:latest", + "sandbox-id", + "sandbox-name", + "https://gateway.example.com", + "0.0.0.0:2222", + "secret", + 300, + &std::collections::HashMap::new(), + "", + ); + + assert_eq!( + pod_template["spec"]["runtimeClassName"], + serde_json::json!(GPU_RUNTIME_CLASS_NAME) + ); + assert_eq!( + pod_template["spec"]["containers"][0]["resources"]["limits"][GPU_RESOURCE_NAME], + serde_json::json!(GPU_RESOURCE_QUANTITY) + ); + } + + #[test] + fn gpu_sandbox_preserves_existing_resource_limits() { + let template = SandboxTemplate { + resources: Some(Struct { + fields: [( + "limits".to_string(), + Value { + kind: Some(Kind::StructValue(Struct { + fields: [("cpu".to_string(), string_value("2"))] + .into_iter() + .collect(), + })), + }, + )] + .into_iter() + .collect(), + }), + ..SandboxTemplate::default() + }; + + let pod_template = sandbox_template_to_k8s( + &template, + true, + "openshell/sandbox:latest", + "sandbox-id", + "sandbox-name", + "https://gateway.example.com", + "0.0.0.0:2222", + "secret", + 300, + &std::collections::HashMap::new(), + "", + ); + + let limits = &pod_template["spec"]["containers"][0]["resources"]["limits"]; + assert_eq!(limits["cpu"], serde_json::json!("2")); + assert_eq!( + limits[GPU_RESOURCE_NAME], + serde_json::json!(GPU_RESOURCE_QUANTITY) + ); + } + + #[test] + fn gpu_sandbox_updates_custom_pod_template() { + let template = SandboxTemplate { + pod_template: Some(Struct { + fields: [( + "spec".to_string(), + Value { + kind: Some(Kind::StructValue(Struct { + fields: [( + "containers".to_string(), + Value { + kind: Some(Kind::ListValue(prost_types::ListValue { + values: vec![Value { + kind: Some(Kind::StructValue(Struct { + fields: [( + "name".to_string(), + string_value("agent"), + )] + .into_iter() + .collect(), + })), + }], + })), + }, + )] + .into_iter() + .collect(), + })), + }, + )] + .into_iter() + .collect(), + }), + ..SandboxTemplate::default() + }; + + let pod_template = sandbox_template_to_k8s( + &template, + true, + "openshell/sandbox:latest", + "sandbox-id", + "sandbox-name", + "https://gateway.example.com", + "0.0.0.0:2222", + "secret", + 300, + &std::collections::HashMap::new(), + "", + ); + + assert_eq!( + pod_template["spec"]["runtimeClassName"], + serde_json::json!(GPU_RUNTIME_CLASS_NAME) + ); + assert_eq!( + pod_template["spec"]["containers"][0]["resources"]["limits"][GPU_RESOURCE_NAME], + serde_json::json!(GPU_RESOURCE_QUANTITY) + ); + } } diff --git a/docs/reference/cli.md b/docs/reference/cli.md new file mode 100644 index 000000000..8be263bb2 --- /dev/null +++ b/docs/reference/cli.md @@ -0,0 +1,258 @@ +--- +title: + page: CLI Reference + nav: CLI +description: Complete command reference for the openshell CLI including all subcommands, flags, and environment variables. +topics: +- Generative AI +- Cybersecurity +tags: +- CLI +- Reference +- Commands +- AI Agents +content: + type: reference + difficulty: technical_beginner + audience: + - engineer + - data_scientist +--- + + + +# CLI Reference + +Complete command reference for the `openshell` CLI. Every subcommand, flag, and option is documented here. + +## Command Tree + +The full hierarchy of `openshell` subcommands: + +```text +openshell +├── status +├── logs [name] +├── term +├── forward +│ ├── start +│ ├── stop +│ └── list +├── gateway +│ ├── start +│ ├── stop +│ ├── destroy +│ ├── info +│ ├── add +│ └── select [name] +├── sandbox +│ ├── create +│ ├── get [name] +│ ├── list +│ ├── delete [--all] +│ ├── connect [name] +│ ├── upload [name] +│ ├── download [name] +│ └── ssh-config +├── policy +│ ├── set +│ ├── get +│ └── list +├── provider +│ ├── create +│ ├── get +│ ├── list +│ ├── update +│ └── delete +├── inference +│ ├── set +│ ├── update +│ └── get +├── doctor +│ ├── logs +│ └── exec +└── completions +``` + +:::{tip} +Commands that accept an optional `[name]` argument — such as `get`, `connect`, `upload`, `download`, and `logs` — fall back to the last-used sandbox when the name is omitted. The CLI records the sandbox name each time you create or connect, and prints a hint showing which sandbox was selected. +::: + +## Status, Logs, and Terminal + +Check gateway health, stream sandbox logs, and launch the terminal dashboard: + +| Command | Description | +|---|---| +| `openshell status` | Show the health and status of the active gateway. | +| `openshell logs [name]` | View sandbox logs. Use `--tail` for streaming, `--source` and `--level` to filter. | +| `openshell term` | Launch the OpenShell Terminal — a dashboard showing sandbox status, live logs, and policy decisions in a single view. Navigate with `j`/`k`, press `f` to follow live output, `s` to filter by source, and `q` to quit. | + +Refer to {doc}`/sandboxes/create-and-manage` for more on monitoring sandboxes and reading log entries. + +## Gateway Commands + +Manage the OpenShell runtime cluster. + +| Command | Description | +|---|---| +| `openshell gateway start` | Deploy a new cluster. Add `--remote user@host` for remote deployment. | +| `openshell gateway stop` | Stop the active cluster, preserving state. | +| `openshell gateway destroy` | Permanently remove the cluster and all its data. | +| `openshell gateway info` | Show detailed information about the cluster. | +| `openshell gateway add ` | Register an existing remote gateway by URL. | +| `openshell gateway select ` | Set the active cluster. All subsequent commands target this cluster. | +| `openshell gateway select` | List all registered clusters (when called without a name). | + +## Diagnostic Commands + +Troubleshoot gateway issues. + +| Command | Description | +|---|---| +| `openshell doctor logs` | Fetch logs from the gateway Docker container. Use `--tail` for streaming, `-n` to limit line count. | +| `openshell doctor exec -- ` | Run a command inside the gateway container (e.g., `kubectl get pods -A`, `k9s`, `sh`). | + +## Sandbox Commands + +Create and manage isolated agent execution environments. + +| Command | Description | +|---|---| +| `openshell sandbox create` | Create a new sandbox. Refer to the flag reference below. | +| `openshell sandbox get [name]` | Show detailed information about a sandbox. | +| `openshell sandbox list` | List all sandboxes in the active cluster. | +| `openshell sandbox delete ` | Delete one or more sandboxes by name. Use `--all` to delete every sandbox in the active gateway. | +| `openshell sandbox connect [name]` | Open an interactive SSH session into a running sandbox. | +| `openshell sandbox upload [name]` | Upload files from the host into a sandbox. | +| `openshell sandbox download [name]` | Download files from a sandbox to the host. | +| `openshell sandbox ssh-config ` | Print the generated SSH config block for a sandbox. | + +### Sandbox Create Flags + +The following flags control sandbox creation: + +| Flag | Description | +|---|---| +| `--name` | Assign a human-readable name to the sandbox. Auto-generated if omitted. | +| `--provider` | Attach a credential provider. Repeatable for multiple providers. | +| `--policy` | Path to a policy YAML file to apply at creation time. | +| `--upload` | Upload local files into the sandbox before running. | +| `--no-keep` | Delete the sandbox after the initial command or shell exits. | +| `--gpu` | Request a GPU-enabled sandbox. If no gateway is running, auto-bootstrap starts a GPU-enabled gateway. | +| `--editor` | Launch `vscode` or `cursor` into `/sandbox`; installs OpenShell-managed SSH config and keeps the sandbox alive. | +| `--forward` | Forward a local port into the sandbox at startup. Keeps the sandbox alive. | +| `--from` | Build from a community sandbox name, local Dockerfile directory, or container image reference. | +| `-- ` | The command to run inside the sandbox. Everything after `--` is passed as the agent command. | + +## Policy Commands + +Apply and inspect sandbox policies at runtime. + +| Command | Description | +|---|---| +| `openshell policy set ` | Apply or update a policy on a running sandbox. Pass `--policy `. | +| `openshell policy get ` | Show the active policy for a sandbox. Add `--full` for the complete policy with metadata. | +| `openshell policy list ` | List all policy versions applied to a sandbox, with status. | + +## Port Forwarding Commands + +Forward sandbox ports to the host for local access. + +| Command | Description | +|---|---| +| `openshell forward start ` | Forward a sandbox port to the host. Add `-d` for background mode. | +| `openshell forward stop ` | Stop an active port forward. | +| `openshell forward list` | List all active port forwards. | + +## Provider Commands + +Manage credential providers that inject secrets into sandboxes. + +| Command | Description | +|---|---| +| `openshell provider create` | Create a new credential provider. Refer to the flag reference below. | +| `openshell provider get ` | Show details of a provider. | +| `openshell provider list` | List all providers in the active cluster. | +| `openshell provider update ` | Update a provider's credentials or configuration. | +| `openshell provider delete ` | Delete a provider. | + +### Provider Create Flags + +The following flags control provider creation: + +| Flag | Description | +|---|---| +| `--name` | Name for the provider. | +| `--type` | Provider type: `claude`, `codex`, `opencode`, `github`, `gitlab`, `nvidia`, `generic`, `outlook`. | +| `--from-existing` | Discover credentials from your current shell environment variables. | +| `--credential` | Set a credential explicitly. Format: `KEY=VALUE` or bare `KEY` to read from env. Repeatable. | +| `--config` | Set a configuration value. Format: `KEY=VALUE`. Repeatable. | + +## Inference Commands + +Configure the backend used by `https://inference.local`. + +### `openshell inference set` + +Set the provider and model for managed inference. Both flags are required. + +| Flag | Description | +|---|---| +| `--provider` | Provider record name to use for injected credentials. | +| `--model` | Model identifier to force on generation requests. | + +### `openshell inference update` + +Update only the fields you specify. + +| Flag | Description | +|---|---| +| `--provider` | Replace the current provider record. | +| `--model` | Replace the current model ID. | + +### `openshell inference get` + +Show the current inference configuration, including provider, model, and version. + +## Environment Variables + +The following environment variables override CLI defaults: + +| Variable | Description | +|---|---| +| `OPENSHELL_GATEWAY` | Name of the gateway to operate on. Overrides the active gateway set by `openshell gateway select`. | +| `OPENSHELL_SANDBOX_POLICY` | Default path to a policy YAML file. When set, `openshell sandbox create` uses this policy if no `--policy` flag is provided. | + +## Shell Completions + +Generate shell completion scripts for tab completion: + +```console +$ openshell completions bash +$ openshell completions zsh +$ openshell completions fish +``` + +Pipe the output to your shell's config file: + +```console +$ openshell completions zsh >> ~/.zshrc +$ source ~/.zshrc +``` + +## Built-in Help + +Every command and subcommand includes built-in help. Use `--help` at any level to see available subcommands, flags, and usage examples: + +```console +$ openshell --help +$ openshell sandbox --help +$ openshell sandbox create --help +$ openshell gateway --help +``` + +Help output groups flags under distinct headings: `FLAGS` for command-specific options, `GATEWAY FLAGS` for `--gateway` and `--gateway-endpoint`, and `GLOBAL FLAGS` for `--verbose`, `--help`, and `--version`. diff --git a/docs/sandboxes/create-and-manage.md b/docs/sandboxes/create-and-manage.md index 14dc5f31c..195c8bb01 100644 --- a/docs/sandboxes/create-and-manage.md +++ b/docs/sandboxes/create-and-manage.md @@ -44,6 +44,14 @@ $ openshell sandbox create -- claude If you have an existing gateway, the sandbox is created in it. Otherwise, a gateway is created automatically. +To request GPU resources explicitly, add `--gpu`: + +```console +$ openshell sandbox create --gpu -- claude +``` + +If no gateway is running, the auto-bootstrap path starts a GPU-enabled gateway first. + A fully specified creation command might look like: @@ -79,6 +87,8 @@ $ openshell sandbox create --from ./my-sandbox-dir $ openshell sandbox create --from my-registry.example.com/my-image:latest ``` +Images whose final name component contains `gpu` also trigger GPU sandbox requests automatically. For example, `--from nvidia-gpu` behaves like a GPU sandbox request even without `--gpu`. + ## List and Inspect Sandboxes Check the status of your sandboxes and retrieve detailed information about individual ones. diff --git a/e2e/python/conftest.py b/e2e/python/conftest.py index 07b7feb34..dc95e91eb 100644 --- a/e2e/python/conftest.py +++ b/e2e/python/conftest.py @@ -11,6 +11,7 @@ import grpc import pytest from openshell import InferenceRouteClient, Sandbox, SandboxClient +from openshell._proto import datamodel_pb2 if TYPE_CHECKING: from collections.abc import Callable, Iterator @@ -91,3 +92,15 @@ def run_python() -> Callable[[Sandbox, str], tuple[int, str, str]]: return result.exit_code, result.stdout, result.stderr return _run + + +@pytest.fixture(scope="session") +def gpu_sandbox_spec() -> datamodel_pb2.SandboxSpec: + image = os.environ.get( + "OPENSHELL_E2E_GPU_IMAGE", + "ghcr.io/nvidia/openshell-community/sandboxes/nvidia-gpu:latest", + ) + return datamodel_pb2.SandboxSpec( + gpu=True, + template=datamodel_pb2.SandboxTemplate(image=image), + ) diff --git a/e2e/python/test_sandbox_gpu.py b/e2e/python/test_sandbox_gpu.py new file mode 100644 index 000000000..510b3d92d --- /dev/null +++ b/e2e/python/test_sandbox_gpu.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from openshell._proto import datamodel_pb2 + +if TYPE_CHECKING: + from collections.abc import Callable + + from openshell import Sandbox + + +@pytest.mark.gpu +def test_gpu_sandbox_reports_available_gpu( + sandbox: Callable[..., Sandbox], + gpu_sandbox_spec: datamodel_pb2.SandboxSpec, +) -> None: + with sandbox(spec=gpu_sandbox_spec, delete_on_exit=True) as sb: + result = sb.exec( + ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + timeout_seconds=30, + ) + + assert result.exit_code == 0, result.stderr + assert result.stdout.strip() diff --git a/e2e/rust/tests/cli_smoke.rs b/e2e/rust/tests/cli_smoke.rs index 876c93ce2..35b2801c9 100644 --- a/e2e/rust/tests/cli_smoke.rs +++ b/e2e/rust/tests/cli_smoke.rs @@ -83,8 +83,8 @@ async fn sandbox_help_shows_upload_download() { } } -/// `openshell sandbox create --help` must show `--upload`, `--no-git-ignore`, -/// `--no-bootstrap`, `--editor`, and +/// `openshell sandbox create --help` must show `--gpu`, `--upload`, +/// `--no-git-ignore`, `--no-bootstrap`, `--editor`, and /// `--auto-providers`/`--no-auto-providers`. /// Note: `--bootstrap` is intentionally hidden (it's the default behaviour). #[tokio::test] @@ -94,6 +94,7 @@ async fn sandbox_create_help_shows_new_flags() { let clean = strip_ansi(&output); for flag in [ + "--gpu", "--upload", "--no-git-ignore", "--no-bootstrap", @@ -158,4 +159,3 @@ async fn status_without_gateway_prints_friendly_message() { "expected hint to run 'openshell gateway start':\n{clean}" ); } - diff --git a/mise.toml b/mise.toml index 5b6ce8353..02116620a 100644 --- a/mise.toml +++ b/mise.toml @@ -44,7 +44,7 @@ DOCKER_BUILDKIT = "1" [vars] # Python paths to include in formatting/linting -python_paths = "python/ tasks/scripts/*.py deploy/sbom/*.py" +python_paths = "python/openshell python/navigator_models tasks/scripts/*.py deploy/sbom/*.py" [task_config] includes = ["tasks/*.toml"] diff --git a/proto/datamodel.proto b/proto/datamodel.proto index ef8b07da5..157afb51d 100644 --- a/proto/datamodel.proto +++ b/proto/datamodel.proto @@ -31,6 +31,8 @@ message SandboxSpec { openshell.sandbox.v1.SandboxPolicy policy = 7; // Provider names to attach to this sandbox. repeated string providers = 8; + // Request NVIDIA GPU resources for this sandbox. + bool gpu = 9; } // Sandbox template mapped onto Kubernetes pod template inputs. diff --git a/pyproject.toml b/pyproject.toml index c6f946af7..d9d370624 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,7 +64,7 @@ module-name = "openshell" [tool.ruff] target-version = "py312" line-length = 88 -exclude = ["python/openshell/_proto"] +exclude = ["python/openshell/_proto", "python/navigator/_proto"] [tool.ruff.lint] select = [ @@ -100,6 +100,9 @@ python_files = ["*_test.py"] python_functions = ["test_*"] addopts = "-v --tb=short" asyncio_mode = "auto" +markers = [ + "gpu: requires a GPU-capable gateway and GPU-enabled sandbox image", +] [tool.coverage.run] source = ["python"] diff --git a/tasks/test.toml b/tasks/test.toml index 3fba1fcfe..f53f9152f 100644 --- a/tasks/test.toml +++ b/tasks/test.toml @@ -11,6 +11,10 @@ depends = ["test:rust", "test:python"] description = "Run all end-to-end tests (Rust + Python)" depends = ["e2e:rust", "e2e:python"] +["e2e:gpu"] +description = "Run GPU end-to-end tests" +depends = ["e2e:python:gpu"] + ["test:rust"] description = "Run Rust tests" run = "cargo test --workspace" @@ -32,4 +36,10 @@ run = ["cargo build -p openshell-cli", "cargo test --manifest-path e2e/rust/Carg description = "Run Python e2e tests (E2E_PARALLEL=N or 'auto'; default 5)" depends = ["python:proto", "cluster"] env = { UV_NO_SYNC = "1", PYTHONPATH = "python" } -run = "uv run pytest -o python_files='test_*.py' -n ${E2E_PARALLEL:-5} e2e/python" +run = "uv run pytest -o python_files='test_*.py' -m 'not gpu' -n ${E2E_PARALLEL:-5} e2e/python" + +["e2e:python:gpu"] +description = "Run Python GPU e2e tests" +depends = ["python:proto", "cluster"] +env = { UV_NO_SYNC = "1", PYTHONPATH = "python" } +run = "uv run pytest -o python_files='test_*.py' -m gpu -n ${E2E_PARALLEL:-1} e2e/python"