From a67790ec55e5cb04a32d87372a934660dfb73d47 Mon Sep 17 00:00:00 2001 From: Shridhar Damale Date: Wed, 29 Apr 2026 13:59:59 +0530 Subject: [PATCH] feat(cli): add CPU/memory resource limit flags to sandbox create Add four new optional CLI flags to openshell sandbox create: --cpu-request, --cpu-limit, --memory-request, --memory-limit These map to Kubernetes resource requests/limits via the existing DriverResourceRequirements proto fields and K8s driver. When omitted, behavior is unchanged (default K8s resource allocations). Includes ResourceArgs struct, prost_types::Struct serialization, unit tests, and integration test updates. Closes #1003 Signed-off-by: Shridhar Damale Made-with: Cursor --- crates/openshell-cli/src/main.rs | 29 ++++ crates/openshell-cli/src/run.rs | 143 +++++++++++++++++- .../sandbox_create_lifecycle_integration.rs | 5 + 3 files changed, 173 insertions(+), 4 deletions(-) diff --git a/crates/openshell-cli/src/main.rs b/crates/openshell-cli/src/main.rs index 3502c2b07..295624aa6 100644 --- a/crates/openshell-cli/src/main.rs +++ b/crates/openshell-cli/src/main.rs @@ -1138,6 +1138,22 @@ enum SandboxCommands { #[arg(long)] gpu: bool, + /// Minimum CPU cores requested (Kubernetes quantity, e.g. "500m", "2"). + #[arg(long, value_name = "QUANTITY", help_heading = "RESOURCE FLAGS")] + cpu_request: Option, + + /// Maximum CPU cores allowed (Kubernetes quantity, e.g. "500m", "4"). + #[arg(long, value_name = "QUANTITY", help_heading = "RESOURCE FLAGS")] + cpu_limit: Option, + + /// Minimum memory requested (Kubernetes quantity, e.g. "256Mi", "4Gi"). + #[arg(long, value_name = "QUANTITY", help_heading = "RESOURCE FLAGS")] + memory_request: Option, + + /// Maximum memory allowed (Kubernetes quantity, e.g. "512Mi", "8Gi"). + #[arg(long, value_name = "QUANTITY", help_heading = "RESOURCE FLAGS")] + memory_limit: Option, + /// SSH destination for remote bootstrap (e.g., user@hostname). /// Only used when no cluster exists yet; ignored if a cluster is /// already active. @@ -2307,6 +2323,10 @@ async fn main() -> Result<()> { no_keep, editor, gpu, + cpu_request, + cpu_limit, + memory_request, + memory_limit, remote, ssh_key, providers, @@ -2368,6 +2388,13 @@ async fn main() -> Result<()> { (local, remote, !no_git_ignore) }); + let resource_args = run::ResourceArgs { + cpu_request, + cpu_limit, + memory_request, + memory_limit, + }; + let editor = editor.map(Into::into); let forward = forward .map(|s| openshell_core::forward::ForwardSpec::parse(&s)) @@ -2402,6 +2429,7 @@ async fn main() -> Result<()> { upload_spec.as_ref(), keep, gpu, + &resource_args, editor, remote.as_deref(), ssh_key.as_deref(), @@ -2425,6 +2453,7 @@ async fn main() -> Result<()> { upload_spec.as_ref(), keep, gpu, + &resource_args, editor, remote.as_deref(), ssh_key.as_deref(), diff --git a/crates/openshell-cli/src/run.rs b/crates/openshell-cli/src/run.rs index 87489014a..f4cedb945 100644 --- a/crates/openshell-cli/src/run.rs +++ b/crates/openshell-cli/src/run.rs @@ -56,6 +56,72 @@ pub use openshell_core::forward::{ find_forward_by_port, list_forwards, stop_forward, stop_forwards_for_sandbox, }; +/// Compute resource arguments for sandbox creation. +#[derive(Debug, Default)] +pub struct ResourceArgs { + pub cpu_request: Option, + pub cpu_limit: Option, + pub memory_request: Option, + pub memory_limit: Option, +} + +impl ResourceArgs { + /// Build a `prost_types::Struct` matching the Kubernetes resources shape + /// expected by `extract_typed_resources` in the gateway server. + /// + /// Returns `None` when no CPU or memory fields are set (GPU is handled + /// separately via `SandboxSpec.gpu`). + pub fn to_resources_struct(&self) -> Option { + use prost_types::{Struct, Value, value::Kind}; + + let mut requests = std::collections::BTreeMap::new(); + let mut limits = std::collections::BTreeMap::new(); + + if let Some(ref v) = self.cpu_request { + requests.insert( + "cpu".to_string(), + Value { kind: Some(Kind::StringValue(v.clone())) }, + ); + } + if let Some(ref v) = self.memory_request { + requests.insert( + "memory".to_string(), + Value { kind: Some(Kind::StringValue(v.clone())) }, + ); + } + if let Some(ref v) = self.cpu_limit { + limits.insert( + "cpu".to_string(), + Value { kind: Some(Kind::StringValue(v.clone())) }, + ); + } + if let Some(ref v) = self.memory_limit { + limits.insert( + "memory".to_string(), + Value { kind: Some(Kind::StringValue(v.clone())) }, + ); + } + if requests.is_empty() && limits.is_empty() { + return None; + } + + let mut fields = std::collections::BTreeMap::new(); + if !requests.is_empty() { + fields.insert( + "requests".to_string(), + Value { kind: Some(Kind::StructValue(Struct { fields: requests })) }, + ); + } + if !limits.is_empty() { + fields.insert( + "limits".to_string(), + Value { kind: Some(Kind::StructValue(Struct { fields: limits })) }, + ); + } + Some(Struct { fields }) + } +} + /// Convert a sandbox phase integer to a human-readable string. fn phase_name(phase: i32) -> &'static str { match SandboxPhase::try_from(phase) { @@ -1923,6 +1989,7 @@ pub async fn sandbox_create_with_bootstrap( upload: Option<&(String, Option, bool)>, keep: bool, gpu: bool, + resource_args: &ResourceArgs, editor: Option, remote: Option<&str>, ssh_key: Option<&str>, @@ -1954,6 +2021,7 @@ pub async fn sandbox_create_with_bootstrap( upload, keep, gpu, + resource_args, editor, remote, ssh_key, @@ -2010,6 +2078,7 @@ pub async fn sandbox_create( upload: Option<&(String, Option, bool)>, keep: bool, gpu: bool, + resource_args: &ResourceArgs, editor: Option, remote: Option<&str>, ssh_key: Option<&str>, @@ -2109,10 +2178,19 @@ pub async fn sandbox_create( let policy = load_sandbox_policy(policy)?; - let template = image.map(|img| SandboxTemplate { - image: img, - ..SandboxTemplate::default() - }); + let resources_struct = resource_args.to_resources_struct(); + let template = match (image, &resources_struct) { + (Some(img), _) => Some(SandboxTemplate { + image: img, + resources: resources_struct, + ..SandboxTemplate::default() + }), + (None, Some(_)) => Some(SandboxTemplate { + resources: resources_struct, + ..SandboxTemplate::default() + }), + (None, None) => None, + }; let request = CreateSandboxRequest { spec: Some(SandboxSpec { @@ -6184,4 +6262,61 @@ mod tests { "should end with single-quote: {ssh_escaped}" ); } + + // ------------------------------------------------------------------ + // ResourceArgs tests + // ------------------------------------------------------------------ + + use super::ResourceArgs; + use prost_types::value::Kind; + + #[test] + fn resource_args_empty_returns_none() { + let args = ResourceArgs::default(); + assert!(args.to_resources_struct().is_none()); + } + + #[test] + fn resource_args_cpu_only_request() { + let args = ResourceArgs { + cpu_request: Some("500m".to_string()), + ..ResourceArgs::default() + }; + let s = args.to_resources_struct().expect("should return Some"); + let requests = &s.fields["requests"]; + let Kind::StructValue(inner) = requests.kind.as_ref().unwrap() else { + panic!("expected StructValue"); + }; + let Kind::StringValue(cpu) = inner.fields["cpu"].kind.as_ref().unwrap() else { + panic!("expected StringValue"); + }; + assert_eq!(cpu, "500m"); + assert!(!s.fields.contains_key("limits")); + } + + #[test] + fn resource_args_all_four_fields() { + let args = ResourceArgs { + cpu_request: Some("500m".to_string()), + cpu_limit: Some("4".to_string()), + memory_request: Some("256Mi".to_string()), + memory_limit: Some("8Gi".to_string()), + }; + let s = args.to_resources_struct().expect("should return Some"); + + let get = |section: &str, key: &str| -> String { + let Kind::StructValue(sec) = s.fields[section].kind.as_ref().unwrap() else { + panic!("expected StructValue for {section}"); + }; + let Kind::StringValue(val) = sec.fields[key].kind.as_ref().unwrap() else { + panic!("expected StringValue for {section}.{key}"); + }; + val.clone() + }; + + assert_eq!(get("requests", "cpu"), "500m"); + assert_eq!(get("requests", "memory"), "256Mi"); + assert_eq!(get("limits", "cpu"), "4"); + assert_eq!(get("limits", "memory"), "8Gi"); + } } diff --git a/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs b/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs index 79d482fdb..de637c8d9 100644 --- a/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs +++ b/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs @@ -571,6 +571,7 @@ async fn sandbox_create_keeps_command_sessions_by_default() { None, true, false, + &run::ResourceArgs::default(), None, None, None, @@ -612,6 +613,7 @@ async fn sandbox_create_deletes_command_sessions_with_no_keep() { None, false, false, + &run::ResourceArgs::default(), None, None, None, @@ -656,6 +658,7 @@ async fn sandbox_create_deletes_shell_sessions_with_no_keep() { None, false, false, + &run::ResourceArgs::default(), None, None, None, @@ -700,6 +703,7 @@ async fn sandbox_create_keeps_sandbox_with_hidden_keep_flag() { None, true, false, + &run::ResourceArgs::default(), None, None, None, @@ -741,6 +745,7 @@ async fn sandbox_create_keeps_sandbox_with_forwarding() { None, false, false, + &run::ResourceArgs::default(), None, None, None,