diff --git a/crates/openshell-cli/src/main.rs b/crates/openshell-cli/src/main.rs index aaabf26625..f671e3eeba 100644 --- a/crates/openshell-cli/src/main.rs +++ b/crates/openshell-cli/src/main.rs @@ -19,7 +19,7 @@ use openshell_bootstrap::{ use openshell_cli::completers; use openshell_cli::run; use openshell_cli::tls::TlsOptions; -use openshell_core::proto::GpuResourceRequirements; +use openshell_core::proto::{GpuResourceRequirements, ResourceRequirements}; /// Resolved gateway context: name + gateway endpoint. struct GatewayContext { @@ -3066,6 +3066,21 @@ async fn run_async() -> Result<()> { .transpose()?; let keep = keep || !no_keep || editor.is_some() || forward.is_some(); let gpu_requirements: Option = gpu.map(Into::into); + let cpu_requirements = run::build_cpu_resource_requirements(cpu.as_deref())?; + let memory_requirements = + run::build_memory_resource_requirements(memory.as_deref())?; + let resource_requirements = if gpu_requirements.is_some() + || cpu_requirements.is_some() + || memory_requirements.is_some() + { + Some(ResourceRequirements { + gpu: gpu_requirements, + cpu: cpu_requirements, + memory: memory_requirements, + }) + } else { + None + }; let ctx = resolve_gateway(&cli.gateway, &cli.gateway_endpoint)?; let endpoint = &ctx.endpoint; @@ -3079,9 +3094,7 @@ async fn run_async() -> Result<()> { from: from.as_deref(), uploads: &upload_specs, keep, - gpu_requirements, - cpu: cpu.as_deref(), - memory: memory.as_deref(), + resource_requirements, driver_config_json: driver_config_json.as_deref(), editor, providers: &providers, diff --git a/crates/openshell-cli/src/run.rs b/crates/openshell-cli/src/run.rs index 0a0b21a7f4..7944550d80 100644 --- a/crates/openshell-cli/src/run.rs +++ b/crates/openshell-cli/src/run.rs @@ -35,27 +35,27 @@ use openshell_core::net::set_tcp_nodelay_best_effort; use openshell_core::proto::ProviderProfileCategory; use openshell_core::proto::{ ApproveAllDraftChunksRequest, ApproveDraftChunkRequest, AttachSandboxProviderRequest, - ClearDraftChunksRequest, ConfigureProviderRefreshRequest, CreateProviderRequest, - CreateSandboxRequest, CreateSshSessionRequest, DeleteInferenceRouteRequest, - DeleteProviderProfileRequest, DeleteProviderRefreshRequest, DeleteProviderRequest, - DeleteSandboxRequest, DeleteServiceRequest, DetachSandboxProviderRequest, ExecSandboxRequest, - ExposeServiceRequest, GetCurrentUserRequest, GetDraftHistoryRequest, GetDraftPolicyRequest, - GetGatewayConfigRequest, GetInferenceRouteRequest, GetProviderProfileRequest, - GetProviderRefreshStatusRequest, GetProviderRequest, GetSandboxConfigRequest, - GetSandboxConfigResponse, GetSandboxLogsRequest, GetSandboxPolicyStatusRequest, - GetSandboxRequest, GetServiceRequest, GpuResourceRequirements, ImportProviderProfilesRequest, - LintProviderProfilesRequest, ListProviderProfilesRequest, ListProvidersRequest, - ListSandboxPoliciesRequest, ListSandboxProvidersRequest, ListSandboxesRequest, - ListServicesRequest, PolicySource, PolicyStatus, Provider, - ProviderCredentialRefreshRecoveryAction, ProviderCredentialRefreshStatus, - ProviderCredentialRefreshStrategy, ProviderCredentialTokenGrantType, ProviderProfile, - ProviderProfileDiagnostic, ProviderProfileImportItem, RejectDraftChunkRequest, - ResourceRequirements, RevokeSshSessionRequest, RotateProviderCredentialRequest, Sandbox, - SandboxPhase, SandboxPolicy, SandboxSpec, SandboxTemplate, ServiceEndpointResponse, - SetInferenceRouteRequest, SettingScope, StartSandboxRequest, StopSandboxRequest, - TcpForwardFrame, TcpForwardInit, TcpRelayTarget, UpdateConfigRequest, - UpdateProviderProfilesRequest, UpdateProviderRequest, WatchSandboxRequest, exec_sandbox_event, - setting_value, tcp_forward_init, + ClearDraftChunksRequest, ConfigureProviderRefreshRequest, CpuResourceRequirements, + CreateProviderRequest, CreateSandboxRequest, CreateSshSessionRequest, + DeleteInferenceRouteRequest, DeleteProviderProfileRequest, DeleteProviderRefreshRequest, + DeleteProviderRequest, DeleteSandboxRequest, DeleteServiceRequest, + DetachSandboxProviderRequest, ExecSandboxRequest, ExposeServiceRequest, GetCurrentUserRequest, + GetDraftHistoryRequest, GetDraftPolicyRequest, GetGatewayConfigRequest, + GetInferenceRouteRequest, GetProviderProfileRequest, GetProviderRefreshStatusRequest, + GetProviderRequest, GetSandboxConfigRequest, GetSandboxConfigResponse, GetSandboxLogsRequest, + GetSandboxPolicyStatusRequest, GetSandboxRequest, GetServiceRequest, + ImportProviderProfilesRequest, LintProviderProfilesRequest, ListProviderProfilesRequest, + ListProvidersRequest, ListSandboxPoliciesRequest, ListSandboxProvidersRequest, + ListSandboxesRequest, ListServicesRequest, MemoryResourceRequirements, PolicySource, + PolicyStatus, Provider, ProviderCredentialRefreshRecoveryAction, + ProviderCredentialRefreshStatus, ProviderCredentialRefreshStrategy, + ProviderCredentialTokenGrantType, ProviderProfile, ProviderProfileDiagnostic, + ProviderProfileImportItem, RejectDraftChunkRequest, ResourceRequirements, + RevokeSshSessionRequest, RotateProviderCredentialRequest, Sandbox, SandboxPhase, SandboxPolicy, + SandboxSpec, SandboxTemplate, ServiceEndpointResponse, SetInferenceRouteRequest, SettingScope, + StartSandboxRequest, StopSandboxRequest, TcpForwardFrame, TcpForwardInit, TcpRelayTarget, + UpdateConfigRequest, UpdateProviderProfilesRequest, UpdateProviderRequest, WatchSandboxRequest, + exec_sandbox_event, setting_value, tcp_forward_init, }; use openshell_core::settings; use openshell_core::{ObjectId, ObjectName, ObjectWorkspace}; @@ -224,41 +224,28 @@ fn sandbox_should_persist(keep: bool, forward: Option<&ForwardSpec>) -> bool { keep || forward.is_some() } -fn build_sandbox_resource_limits( +pub fn build_cpu_resource_requirements( cpu: Option<&str>, - memory: Option<&str>, -) -> Result> { - use prost_types::{Struct, Value, value::Kind}; - - fn string_value(value: String) -> Value { - Value { - kind: Some(Kind::StringValue(value)), - } - } +) -> Result> { + let Some(cpu) = cpu else { + return Ok(None); + }; - let mut limits = std::collections::BTreeMap::new(); - if let Some(cpu) = cpu { - limits.insert("cpu".to_string(), string_value(validate_cpu_quantity(cpu)?)); - } - if let Some(memory) = memory { - limits.insert( - "memory".to_string(), - string_value(validate_memory_quantity(memory)?), - ); - } + Ok(Some(CpuResourceRequirements { + limit: validate_cpu_quantity(cpu)?, + })) +} - if limits.is_empty() { +pub fn build_memory_resource_requirements( + memory: Option<&str>, +) -> Result> { + let Some(memory) = memory else { return Ok(None); - } + }; - let mut fields = std::collections::BTreeMap::new(); - fields.insert( - "limits".to_string(), - Value { - kind: Some(Kind::StructValue(Struct { fields: limits })), - }, - ); - Ok(Some(Struct { fields })) + Ok(Some(MemoryResourceRequirements { + limit: validate_memory_quantity(memory)?, + })) } fn parse_driver_config_json(value: &str) -> Result { @@ -370,9 +357,7 @@ pub struct SandboxCreateConfig<'a> { pub from: Option<&'a str>, pub uploads: &'a [(String, Option, bool)], pub keep: bool, - pub gpu_requirements: Option, - pub cpu: Option<&'a str>, - pub memory: Option<&'a str>, + pub resource_requirements: Option, pub driver_config_json: Option<&'a str>, pub editor: Option, pub providers: &'a [String], @@ -395,9 +380,7 @@ impl Default for SandboxCreateConfig<'_> { from: None, uploads: &[], keep: false, - gpu_requirements: None, - cpu: None, - memory: None, + resource_requirements: None, driver_config_json: None, editor: None, providers: &[], @@ -428,9 +411,7 @@ pub async fn sandbox_create( from, uploads, keep, - gpu_requirements, - cpu, - memory, + resource_requirements, driver_config_json, editor, providers, @@ -513,15 +494,13 @@ pub async fn sandbox_create( .await?; let policy = load_sandbox_policy(policy)?; - let resource_limits = build_sandbox_resource_limits(cpu, memory)?; let driver_config = driver_config_json .map(parse_driver_config_json) .transpose()?; - let template = if image.is_some() || resource_limits.is_some() || driver_config.is_some() { + let template = if image.is_some() || driver_config.is_some() { Some(SandboxTemplate { image: image.unwrap_or_default(), - resources: resource_limits, driver_config, ..SandboxTemplate::default() }) @@ -529,8 +508,6 @@ pub async fn sandbox_create( None }; - let resource_requirements = gpu_requirements.map(|gpu| ResourceRequirements { gpu: Some(gpu) }); - let main_terminal = tty_override .unwrap_or_else(|| std::io::stdin().is_terminal() && std::io::stdout().is_terminal()); let main_command = if command.is_empty() { @@ -540,7 +517,7 @@ pub async fn sandbox_create( }; let request = CreateSandboxRequest { spec: Some(SandboxSpec { - resource_requirements, + resource_requirements: resource_requirements.clone(), environment, policy, providers: configured_providers, @@ -7404,9 +7381,9 @@ fn format_endpoint(endpoint: &openshell_core::proto::NetworkEndpoint) -> String #[cfg(test)] mod tests { use super::{ - PolicyGetView, ProvisioningStep, build_sandbox_resource_limits, - dockerfile_sources_supported_for_gateway, format_endpoint, - format_provider_attachment_table, git_sync_files, inferred_provider_type, + PolicyGetView, ProvisioningStep, build_cpu_resource_requirements, + build_memory_resource_requirements, dockerfile_sources_supported_for_gateway, + format_endpoint, format_provider_attachment_table, git_sync_files, inferred_provider_type, parse_cli_setting_value, parse_credential_expiry_cli_value, parse_credential_expiry_pairs, parse_credential_pairs, parse_driver_config_json, parse_secret_material_env_pairs, policy_revision_to_json, provider_profile_allows_empty_credentials, @@ -7812,52 +7789,29 @@ mod tests { } #[test] - fn build_sandbox_resource_limits_sets_limits_only() { - let resources = build_sandbox_resource_limits(Some("500m"), Some("2Gi")) - .expect("resource limits should parse") - .expect("resource limits should be present"); + fn build_cpu_resource_requirements_sets_typed_limit() { + let cpu = build_cpu_resource_requirements(Some("500m")) + .expect("CPU limit should parse") + .expect("CPU requirements should be present"); - let limits = resources - .fields - .get("limits") - .and_then(|value| value.kind.as_ref()) - .and_then(|kind| match kind { - prost_types::value::Kind::StructValue(inner) => Some(inner), - _ => None, - }) - .expect("limits should be a struct"); + assert_eq!(cpu.limit, "500m"); + } - assert_eq!( - limits - .fields - .get("cpu") - .and_then(|value| value.kind.as_ref()) - .and_then(|kind| match kind { - prost_types::value::Kind::StringValue(value) => Some(value.as_str()), - _ => None, - }), - Some("500m") - ); - assert_eq!( - limits - .fields - .get("memory") - .and_then(|value| value.kind.as_ref()) - .and_then(|kind| match kind { - prost_types::value::Kind::StringValue(value) => Some(value.as_str()), - _ => None, - }), - Some("2Gi") - ); - assert!(!resources.fields.contains_key("requests")); + #[test] + fn build_memory_resource_requirements_sets_typed_limit() { + let memory = build_memory_resource_requirements(Some("2Gi")) + .expect("memory limit should parse") + .expect("memory requirements should be present"); + + assert_eq!(memory.limit, "2Gi"); } #[test] - fn build_sandbox_resource_limits_rejects_invalid_quantities() { - assert!(build_sandbox_resource_limits(Some("0"), None).is_err()); - assert!(build_sandbox_resource_limits(Some("half"), None).is_err()); - assert!(build_sandbox_resource_limits(None, Some("0Gi")).is_err()); - assert!(build_sandbox_resource_limits(None, Some("1.5Gi")).is_err()); + fn build_cpu_and_memory_resource_requirements_reject_invalid_quantities() { + assert!(build_cpu_resource_requirements(Some("0")).is_err()); + assert!(build_cpu_resource_requirements(Some("half")).is_err()); + assert!(build_memory_resource_requirements(Some("0Gi")).is_err()); + assert!(build_memory_resource_requirements(Some("1.5Gi")).is_err()); } #[test] @@ -8157,6 +8111,8 @@ mod tests { fn provisioning_timeout_message_includes_condition_and_gpu_hint() { let resource_requirements = ResourceRequirements { gpu: Some(GpuResourceRequirements { count: None }), + cpu: None, + memory: None, }; let message = provisioning_timeout_message( 120, @@ -8178,7 +8134,11 @@ mod tests { #[test] fn provisioning_timeout_message_omits_gpu_hint_without_gpu_requirements() { - let resource_requirements = ResourceRequirements { gpu: None }; + let resource_requirements = ResourceRequirements { + gpu: None, + cpu: None, + memory: None, + }; let message = provisioning_timeout_message(120, Some(&resource_requirements), None); assert_eq!(message, "sandbox provisioning timed out after 120s"); diff --git a/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs b/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs index bcc07619ee..0f67a4c02d 100644 --- a/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs +++ b/crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs @@ -24,9 +24,9 @@ use openshell_core::proto::{ GetSandboxProviderEnvironmentResponse, GetSandboxRequest, GpuResourceRequirements, HealthRequest, HealthResponse, ListProvidersRequest, ListProvidersResponse, ListSandboxProvidersRequest, ListSandboxProvidersResponse, ListSandboxesRequest, - ListSandboxesResponse, PlatformEvent, ProviderResponse, RevokeSshSessionRequest, - RevokeSshSessionResponse, Sandbox, SandboxCondition, SandboxLogLine, SandboxPhase, - SandboxResponse, SandboxStatus, SandboxStreamEvent, ServiceStatus, SettingValue, + ListSandboxesResponse, PlatformEvent, ProviderResponse, ResourceRequirements, + RevokeSshSessionRequest, RevokeSshSessionResponse, Sandbox, SandboxCondition, SandboxLogLine, + SandboxPhase, SandboxResponse, SandboxStatus, SandboxStreamEvent, ServiceStatus, SettingValue, SupervisorMessage, UpdateProviderRequest, WatchSandboxRequest, sandbox_stream_event, setting_value, }; @@ -1160,6 +1160,14 @@ fn gpu_requirements(count: Option) -> GpuResourceRequirements { GpuResourceRequirements { count } } +fn resource_requirements( + gpu: Option, + cpu: Option, + memory: Option, +) -> ResourceRequirements { + ResourceRequirements { gpu, cpu, memory } +} + /// Shared defaults for integration tests. Note: `keep` is `true` here (most /// tests expect persistent sandboxes) while `SandboxCreateConfig::default()` /// sets `keep: false` (the safe production default). Tests that exercise @@ -1239,7 +1247,7 @@ async fn sandbox_create_without_inferred_provider_skips_gateway_config() { } #[tokio::test] -async fn sandbox_create_sends_cpu_and_memory_limits_only() { +async fn sandbox_create_sends_typed_cpu_and_memory_requirements() { let server = run_server().await; let fake_ssh_dir = tempfile::tempdir().unwrap(); let xdg_dir = tempfile::tempdir().unwrap(); @@ -1252,8 +1260,15 @@ async fn sandbox_create_sends_cpu_and_memory_limits_only() { "openshell", run::SandboxCreateConfig { name: Some("resources"), - cpu: Some("500m"), - memory: Some("2Gi"), + resource_requirements: Some(resource_requirements( + None, + Some(openshell_core::proto::CpuResourceRequirements { + limit: "500m".to_string(), + }), + Some(openshell_core::proto::MemoryResourceRequirements { + limit: "2Gi".to_string(), + }), + )), command: &["echo".into(), "OK".into()], ..test_config() }, @@ -1264,45 +1279,31 @@ async fn sandbox_create_sends_cpu_and_memory_limits_only() { .expect("sandbox create should succeed"); let requests = create_requests(&server).await; - let resources = requests[0] + let requirements = requests[0] .spec .as_ref() - .and_then(|spec| spec.template.as_ref()) - .and_then(|template| template.resources.as_ref()) - .expect("resource limits should be sent"); - let limits = resources - .fields - .get("limits") - .and_then(|value| value.kind.as_ref()) - .and_then(|kind| match kind { - prost_types::value::Kind::StructValue(inner) => Some(inner), - _ => None, - }) - .expect("limits should be a struct"); + .and_then(|spec| spec.resource_requirements.as_ref()) + .expect("resource requirements should be sent"); assert_eq!( - limits - .fields - .get("cpu") - .and_then(|value| value.kind.as_ref()) - .and_then(|kind| match kind { - prost_types::value::Kind::StringValue(value) => Some(value.as_str()), - _ => None, - }), + requirements.cpu.as_ref().map(|cpu| cpu.limit.as_str()), Some("500m") ); assert_eq!( - limits - .fields - .get("memory") - .and_then(|value| value.kind.as_ref()) - .and_then(|kind| match kind { - prost_types::value::Kind::StringValue(value) => Some(value.as_str()), - _ => None, - }), + requirements + .memory + .as_ref() + .map(|memory| memory.limit.as_str()), Some("2Gi") ); - assert!(!resources.fields.contains_key("requests")); + assert!( + requests[0] + .spec + .as_ref() + .and_then(|spec| spec.template.as_ref()) + .is_none(), + "resource-only create should not synthesize a template" + ); } #[tokio::test] @@ -1421,7 +1422,11 @@ async fn sandbox_create_sends_gpu_default_request() { "openshell", run::SandboxCreateConfig { name: Some("gpu-default"), - gpu_requirements: Some(gpu_requirements(None)), + resource_requirements: Some(resource_requirements( + Some(gpu_requirements(None)), + None, + None, + )), command: &["echo".into(), "OK".into()], ..test_config() }, @@ -1456,7 +1461,11 @@ async fn sandbox_create_sends_gpu_count_request() { "openshell", run::SandboxCreateConfig { name: Some("gpu-two"), - gpu_requirements: Some(gpu_requirements(Some(2))), + resource_requirements: Some(resource_requirements( + Some(gpu_requirements(Some(2))), + None, + None, + )), command: &["echo".into(), "OK".into()], ..test_config() }, diff --git a/crates/openshell-core/src/gpu.rs b/crates/openshell-core/src/gpu.rs index f5ff67cd35..749a477487 100644 --- a/crates/openshell-core/src/gpu.rs +++ b/crates/openshell-core/src/gpu.rs @@ -12,7 +12,7 @@ use crate::config::CDI_GPU_DEVICE_ALL; use crate::proto::ResourceRequirements as SandboxResourceRequirements; use crate::proto::compute::v1::{ GpuResourceRequirements as DriverGpuResourceRequirements, - ResourceRequirements as DriverResourceRequirements, + ResourceRequirements as DriverSandboxResourceRequirements, }; /// Return whether sandbox resource requirements request a GPU. @@ -53,7 +53,7 @@ pub fn effective_driver_gpu_count( /// Return the requested compute-driver GPU requirements, if present. #[must_use] pub fn driver_gpu_requirements( - resources: Option<&DriverResourceRequirements>, + resources: Option<&DriverSandboxResourceRequirements>, ) -> Option<&DriverGpuResourceRequirements> { resources.and_then(|resources| resources.gpu.as_ref()) } diff --git a/crates/openshell-driver-docker/src/lib.rs b/crates/openshell-driver-docker/src/lib.rs index b1859f54cc..4b4eb2e52b 100644 --- a/crates/openshell-driver-docker/src/lib.rs +++ b/crates/openshell-driver-docker/src/lib.rs @@ -46,12 +46,12 @@ use openshell_core::proto::compute::v1::{ EnsureWorkspaceResponse, GatewayListenerRequirement, GetCapabilitiesRequest, GetCapabilitiesResponse, GetGatewayListenerRequirementsRequest, GetGatewayListenerRequirementsResponse, GetSandboxRequest, GetSandboxResponse, - GpuResourceRequirements, ListSandboxesRequest, ListSandboxesResponse, StartSandboxRequest, - StartSandboxResponse, StopSandboxRequest, StopSandboxResponse, ValidateSandboxCreateRequest, - ValidateSandboxCreateResponse, WatchSandboxesDeletedEvent, WatchSandboxesEvent, - WatchSandboxesPlatformEvent, WatchSandboxesRequest, WatchSandboxesSandboxEvent, - compute_driver_server::ComputeDriver, gateway_listener_requirement::Selector, - watch_sandboxes_event, + GpuResourceRequirements, ListSandboxesRequest, ListSandboxesResponse, ResourceRequirements, + StartSandboxRequest, StartSandboxResponse, StopSandboxRequest, StopSandboxResponse, + ValidateSandboxCreateRequest, ValidateSandboxCreateResponse, WatchSandboxesDeletedEvent, + WatchSandboxesEvent, WatchSandboxesPlatformEvent, WatchSandboxesRequest, + WatchSandboxesSandboxEvent, compute_driver_server::ComputeDriver, + gateway_listener_requirement::Selector, watch_sandboxes_event, }; use openshell_core::proto_struct::{ deserialize_optional_non_empty_string_list, struct_to_json_value, @@ -661,7 +661,7 @@ impl DockerComputeDriver { .ok_or_else(|| Status::invalid_argument("sandbox.spec.template is required"))?; Self::validate_sandbox_template_base(template)?; - let _ = docker_resource_limits(template)?; + let _ = docker_resource_limits(spec.resource_requirements.as_ref())?; let driver_config = DockerSandboxDriverConfig::from_template(template).map_err(Status::invalid_argument)?; validate_docker_driver_mounts(&driver_config.mounts, config.enable_bind_mounts)?; @@ -2956,7 +2956,7 @@ fn build_container_create_body_for_image( .template .as_ref() .ok_or_else(|| Status::invalid_argument("sandbox.spec.template is required"))?; - let resource_limits = docker_resource_limits(template)?; + let resource_limits = docker_resource_limits(spec.resource_requirements.as_ref())?; let workspace_root = driver_mounts::resolve_oci_workspace_root(&image.working_dir) .map_err(Status::failed_precondition)?; driver_mounts::validate_workspace_control_path(&workspace_root, &config.ssh_socket_path) @@ -3320,26 +3320,20 @@ fn docker_bridge_gateway_ip( } fn docker_resource_limits( - template: &DriverSandboxTemplate, + resources: Option<&ResourceRequirements>, ) -> Result { - let Some(resources) = template.resources.as_ref() else { + let Some(resources) = resources else { return Ok(DockerResourceLimits::default()); }; - if !resources.cpu_request.trim().is_empty() { - return Err(Status::failed_precondition( - "docker compute driver does not support resources.requests.cpu", - )); - } - if !resources.memory_request.trim().is_empty() { - return Err(Status::failed_precondition( - "docker compute driver does not support resources.requests.memory", - )); - } - Ok(DockerResourceLimits { - nano_cpus: parse_cpu_limit(&resources.cpu_limit)?, - memory_bytes: parse_memory_limit(&resources.memory_limit)?, + nano_cpus: parse_cpu_limit(resources.cpu.as_ref().map_or("", |cpu| cpu.limit.as_str()))?, + memory_bytes: parse_memory_limit( + resources + .memory + .as_ref() + .map_or("", |memory| memory.limit.as_str()), + )?, }) } @@ -3374,12 +3368,12 @@ fn parse_cpu_limit(value: &str) -> Result, Status> { if let Some(millicores) = value.strip_suffix('m') { let millicores = millicores.parse::().map_err(|_| { Status::failed_precondition(format!( - "invalid docker cpu_limit '{value}'; expected an integer or millicore quantity", + "invalid docker cpu.limit '{value}'; expected an integer or millicore quantity", )) })?; if millicores <= 0 { return Err(Status::failed_precondition( - "docker cpu_limit must be greater than zero", + "docker cpu.limit must be greater than zero", )); } return Ok(Some(millicores.saturating_mul(1_000_000))); @@ -3387,12 +3381,12 @@ fn parse_cpu_limit(value: &str) -> Result, Status> { let cores = value.parse::().map_err(|_| { Status::failed_precondition(format!( - "invalid docker cpu_limit '{value}'; expected an integer or millicore quantity", + "invalid docker cpu.limit '{value}'; expected an integer or millicore quantity", )) })?; if !cores.is_finite() || cores <= 0.0 { return Err(Status::failed_precondition( - "docker cpu_limit must be greater than zero", + "docker cpu.limit must be greater than zero", )); } @@ -3412,12 +3406,12 @@ fn parse_memory_limit(value: &str) -> Result, Status> { let (number, suffix) = value.split_at(number_end); let amount = number.parse::().map_err(|_| { Status::failed_precondition(format!( - "invalid docker memory_limit '{value}'; expected a Kubernetes-style quantity", + "invalid docker memory.limit '{value}'; expected a Kubernetes-style quantity", )) })?; if !amount.is_finite() || amount <= 0.0 { return Err(Status::failed_precondition( - "docker memory_limit must be greater than zero", + "docker memory.limit must be greater than zero", )); } @@ -3437,7 +3431,7 @@ fn parse_memory_limit(value: &str) -> Result, Status> { "E" => 1000_f64.powi(6), _ => { return Err(Status::failed_precondition(format!( - "invalid docker memory_limit suffix '{suffix}'", + "invalid docker memory.limit suffix '{suffix}'", ))); } }; diff --git a/crates/openshell-driver-docker/src/tests.rs b/crates/openshell-driver-docker/src/tests.rs index b52cb87836..af03fd9d19 100644 --- a/crates/openshell-driver-docker/src/tests.rs +++ b/crates/openshell-driver-docker/src/tests.rs @@ -13,9 +13,9 @@ use openshell_core::progress::{ PROGRESS_STEP_STARTING_SANDBOX, }; use openshell_core::proto::compute::v1::{ - DriverResourceRequirements, DriverSandboxSpec, DriverSandboxTemplate, - GetGatewayListenerRequirementsRequest, GpuResourceRequirements, ResourceRequirements, - gateway_listener_requirement::Selector, + CpuResourceRequirements, DriverSandboxSpec, DriverSandboxTemplate, + GetGatewayListenerRequirementsRequest, GpuResourceRequirements, MemoryResourceRequirements, + ResourceRequirements, gateway_listener_requirement::Selector, }; use std::fs; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; @@ -87,6 +87,20 @@ fn list_string_driver_config(field: &str, values: &[&str]) -> prost_types::Struc fn gpu_resources(count: Option) -> ResourceRequirements { ResourceRequirements { gpu: Some(GpuResourceRequirements { count }), + cpu: None, + memory: None, + } +} + +fn cpu_memory_resources(cpu: Option<&str>, memory: Option<&str>) -> ResourceRequirements { + ResourceRequirements { + gpu: None, + cpu: cpu.map(|limit| CpuResourceRequirements { + limit: limit.to_string(), + }), + memory: memory.map(|limit| MemoryResourceRequirements { + limit: limit.to_string(), + }), } } @@ -1004,43 +1018,11 @@ fn parse_memory_limit_supports_binary_quantities() { assert!(parse_memory_limit("12XB").is_err()); } -#[test] -fn docker_resource_limits_rejects_requests() { - let template = DriverSandboxTemplate { - image: "img".to_string(), - agent_socket_path: String::new(), - labels: HashMap::new(), - environment: HashMap::new(), - resources: Some(DriverResourceRequirements { - cpu_request: "250m".to_string(), - cpu_limit: String::new(), - memory_request: String::new(), - memory_limit: String::new(), - }), - ..Default::default() - }; - - let err = docker_resource_limits(&template).unwrap_err(); - assert_eq!(err.code(), tonic::Code::FailedPrecondition); - assert!(err.message().contains("resources.requests.cpu")); -} - #[test] fn docker_resource_limits_applies_cpu_and_memory_limits() { - let template = DriverSandboxTemplate { - image: "img".to_string(), - agent_socket_path: String::new(), - labels: HashMap::new(), - environment: HashMap::new(), - resources: Some(DriverResourceRequirements { - cpu_limit: "500m".to_string(), - memory_limit: "2Gi".to_string(), - ..Default::default() - }), - ..Default::default() - }; + let resources = cpu_memory_resources(Some("500m"), Some("2Gi")); - let limits = docker_resource_limits(&template).unwrap(); + let limits = docker_resource_limits(Some(&resources)).unwrap(); assert_eq!(limits.nano_cpus, Some(500_000_000)); assert_eq!(limits.memory_bytes, Some(2_147_483_648)); } diff --git a/crates/openshell-driver-kubernetes/src/driver.rs b/crates/openshell-driver-kubernetes/src/driver.rs index c1885870e8..c2f10ff85b 100644 --- a/crates/openshell-driver-kubernetes/src/driver.rs +++ b/crates/openshell-driver-kubernetes/src/driver.rs @@ -43,9 +43,9 @@ use openshell_core::proto::compute::v1::{ DriverCondition as SandboxCondition, DriverPlatformEvent as PlatformEvent, DriverSandbox as Sandbox, DriverSandboxSpec as SandboxSpec, DriverSandboxStatus as SandboxStatus, DriverSandboxTemplate as SandboxTemplate, - GetCapabilitiesResponse, GpuResourceRequirements, WatchSandboxesDeletedEvent, - WatchSandboxesEvent, WatchSandboxesPlatformEvent, WatchSandboxesSandboxEvent, - watch_sandboxes_event, + GetCapabilitiesResponse, GpuResourceRequirements, ResourceRequirements, + WatchSandboxesDeletedEvent, WatchSandboxesEvent, WatchSandboxesPlatformEvent, + WatchSandboxesSandboxEvent, watch_sandboxes_event, }; use openshell_core::proto_struct::{struct_to_json_object, value_to_json}; use serde::Deserialize; @@ -3422,7 +3422,7 @@ fn sandbox_to_k8s_spec( "podTemplate".to_string(), sandbox_template_to_k8s_with_validated_config( template, - driver_gpu_requirements(spec.resource_requirements.as_ref()), + spec.resource_requirements.as_ref(), &pod_env, Some(spec), &driver_config, @@ -3456,7 +3456,7 @@ fn sandbox_to_k8s_spec( "podTemplate".to_string(), sandbox_template_to_k8s_with_validated_config( &SandboxTemplate::default(), - driver_gpu_requirements(spec.and_then(|s| s.resource_requirements.as_ref())), + spec.and_then(|s| s.resource_requirements.as_ref()), &pod_env, spec, &driver_config, @@ -3479,12 +3479,16 @@ fn sandbox_template_to_k8s( inject_workspace: bool, params: &SandboxPodParams<'_>, ) -> serde_json::Value { - let gpu_requirements = gpu.then_some(GpuResourceRequirements { count: None }); + let resource_requirements = gpu.then_some(ResourceRequirements { + gpu: Some(GpuResourceRequirements { count: None }), + cpu: None, + memory: None, + }); let driver_config = KubernetesSandboxDriverConfig::from_template(template) .expect("test Kubernetes driver_config should be valid"); sandbox_template_to_k8s_with_validated_config( template, - gpu_requirements.as_ref(), + resource_requirements.as_ref(), spec_environment, None, &driver_config, @@ -3501,11 +3505,37 @@ fn sandbox_template_to_k8s_with_gpu_requirements( inject_workspace: bool, params: &SandboxPodParams<'_>, ) -> serde_json::Value { + let resource_requirements = gpu_requirements.map(|gpu| ResourceRequirements { + gpu: Some(*gpu), + cpu: None, + memory: None, + }); let driver_config = KubernetesSandboxDriverConfig::from_template(template) .expect("test Kubernetes driver_config should be valid"); sandbox_template_to_k8s_with_validated_config( template, - gpu_requirements, + resource_requirements.as_ref(), + spec_environment, + None, + &driver_config, + inject_workspace, + params, + ) +} + +#[cfg(test)] +fn sandbox_template_to_k8s_with_resource_requirements( + template: &SandboxTemplate, + resource_requirements: Option<&ResourceRequirements>, + spec_environment: &std::collections::HashMap, + inject_workspace: bool, + params: &SandboxPodParams<'_>, +) -> serde_json::Value { + let driver_config = KubernetesSandboxDriverConfig::from_template(template) + .expect("test Kubernetes driver_config should be valid"); + sandbox_template_to_k8s_with_validated_config( + template, + resource_requirements, spec_environment, None, &driver_config, @@ -3516,13 +3546,14 @@ fn sandbox_template_to_k8s_with_gpu_requirements( fn sandbox_template_to_k8s_with_validated_config( template: &SandboxTemplate, - gpu_requirements: Option<&GpuResourceRequirements>, + resource_requirements: Option<&ResourceRequirements>, spec_environment: &std::collections::HashMap, sandbox_spec: Option<&openshell_core::proto::compute::v1::DriverSandboxSpec>, driver_config: &KubernetesSandboxDriverConfig, inject_workspace: bool, params: &SandboxPodParams<'_>, ) -> serde_json::Value { + let gpu_requirements = driver_gpu_requirements(resource_requirements); let mut metadata = serde_json::Map::new(); let mut pod_labels = template .labels @@ -3720,7 +3751,7 @@ fn sandbox_template_to_k8s_with_validated_config( serde_json::Value::Array(volume_mounts), ); - if let Some(resources) = container_resources(template, gpu_requirements) { + if let Some(resources) = container_resources(template, resource_requirements) { container.insert("resources".to_string(), resources); } apply_agent_driver_resources(&mut container, &driver_config.containers.agent.resources); @@ -3958,16 +3989,18 @@ fn app_armor_profile_to_k8s(profile: &AppArmorProfile) -> serde_json::Value { fn container_resources( template: &SandboxTemplate, - gpu_requirements: Option<&GpuResourceRequirements>, + resource_requirements: Option<&ResourceRequirements>, ) -> Option { + let gpu_requirements = driver_gpu_requirements(resource_requirements); + // Start from the raw resources passthrough in platform_config (preserves // custom resource types like GPU limits that users set via the public API - // Struct), then overlay the typed DriverResourceRequirements on top. + // Struct), then overlay the typed resource requirements on top. let mut resources = platform_config_struct(template, "resources_raw").unwrap_or_else(|| serde_json::json!({})); - // Overlay typed CPU/memory from DriverResourceRequirements. - if let Some(ref req) = template.resources { + // Overlay typed CPU/memory from ResourceRequirements. + if let Some(requirements) = resource_requirements { let obj = resources.as_object_mut().unwrap(); let mut apply = |section: &str, key: &str, value: &str| { if !value.is_empty() { @@ -3975,21 +4008,14 @@ fn container_resources( sec[key] = serde_json::json!(value); } }; - apply("limits", "cpu", &req.cpu_limit); - apply("limits", "memory", &req.memory_limit); - - let cpu_request = if req.cpu_request.is_empty() { - &req.cpu_limit - } else { - &req.cpu_request - }; - let memory_request = if req.memory_request.is_empty() { - &req.memory_limit - } else { - &req.memory_request - }; - apply("requests", "cpu", cpu_request); - apply("requests", "memory", memory_request); + if let Some(cpu) = requirements.cpu.as_ref() { + apply("limits", "cpu", &cpu.limit); + apply("requests", "cpu", &cpu.limit); + } + if let Some(memory) = requirements.memory.as_ref() { + apply("limits", "memory", &memory.limit); + apply("requests", "memory", &memory.limit); + } } if let Some(gpu) = gpu_requirements { @@ -4642,7 +4668,10 @@ mod tests { PROGRESS_ACTIVE_DETAIL_KEY, PROGRESS_ACTIVE_STEP_KEY, PROGRESS_COMPLETE_LABEL_KEY, PROGRESS_COMPLETE_STEP_KEY, }; - use openshell_core::proto::compute::v1::{GpuResourceRequirements, ResourceRequirements}; + use openshell_core::proto::compute::v1::{ + CpuResourceRequirements, GpuResourceRequirements, MemoryResourceRequirements, + ResourceRequirements, + }; use prost_types::{Struct, Value, value::Kind}; static ENV_LOCK: std::sync::LazyLock> = @@ -5505,6 +5534,8 @@ mod tests { spec: Some(SandboxSpec { resource_requirements: Some(ResourceRequirements { gpu: Some(GpuResourceRequirements { count: Some(0) }), + cpu: None, + memory: None, }), ..SandboxSpec::default() }), @@ -6585,22 +6616,26 @@ mod tests { } #[test] - fn gpu_sandbox_preserves_existing_resource_limits() { - use openshell_core::proto::compute::v1::DriverResourceRequirements; - let template = SandboxTemplate { - resources: Some(DriverResourceRequirements { - cpu_limit: "2".to_string(), - ..Default::default() + fn gpu_sandbox_preserves_typed_cpu_resource_limits() { + let template = SandboxTemplate::default(); + let resource_requirements = ResourceRequirements { + gpu: Some(GpuResourceRequirements { count: None }), + cpu: Some(CpuResourceRequirements { + limit: "2".to_string(), }), - ..SandboxTemplate::default() + memory: None, }; let pod_template = { let params = SandboxPodParams::default(); - sandbox_template_to_k8s( + let driver_config = KubernetesSandboxDriverConfig::from_template(&template) + .expect("test Kubernetes driver_config should be valid"); + sandbox_template_to_k8s_with_validated_config( &template, - true, + Some(&resource_requirements), &std::collections::HashMap::new(), + None, + &driver_config, true, ¶ms, ) @@ -6613,21 +6648,22 @@ mod tests { #[test] fn cpu_and_memory_limits_are_mirrored_to_requests() { - use openshell_core::proto::compute::v1::DriverResourceRequirements; - let template = SandboxTemplate { - resources: Some(DriverResourceRequirements { - cpu_limit: "500m".to_string(), - memory_limit: "2Gi".to_string(), - ..Default::default() + let template = SandboxTemplate::default(); + let resource_requirements = ResourceRequirements { + gpu: None, + cpu: Some(CpuResourceRequirements { + limit: "500m".to_string(), + }), + memory: Some(MemoryResourceRequirements { + limit: "2Gi".to_string(), }), - ..SandboxTemplate::default() }; let pod_template = { let params = SandboxPodParams::default(); - sandbox_template_to_k8s( + sandbox_template_to_k8s_with_resource_requirements( &template, - false, + Some(&resource_requirements), &std::collections::HashMap::new(), true, ¶ms, diff --git a/crates/openshell-driver-mxc/src/driver.rs b/crates/openshell-driver-mxc/src/driver.rs index d8e981a46d..f32ef82862 100644 --- a/crates/openshell-driver-mxc/src/driver.rs +++ b/crates/openshell-driver-mxc/src/driver.rs @@ -319,6 +319,17 @@ impl MxcComputeBackend { "mxc driver does not support GPU sandboxes", )); } + if spec + .resource_requirements + .as_ref() + .is_some_and(|requirements| { + requirements.cpu.is_some() || requirements.memory.is_some() + }) + { + return Err(tonic::Status::failed_precondition( + "mxc driver does not support spec.resource_requirements.cpu or spec.resource_requirements.memory", + )); + } if let Some(tmpl) = &spec.template && !tmpl.agent_socket_path.is_empty() { diff --git a/crates/openshell-driver-podman/src/container.rs b/crates/openshell-driver-podman/src/container.rs index abb9d69dd2..add90102ce 100644 --- a/crates/openshell-driver-podman/src/container.rs +++ b/crates/openshell-driver-podman/src/container.rs @@ -8,7 +8,9 @@ use openshell_core::ComputeDriverError; use openshell_core::driver_mounts::SelinuxLabel; #[cfg(test)] use openshell_core::gpu::{driver_gpu_requirements, validate_specific_gpu_device_request}; -use openshell_core::proto::compute::v1::{DriverSandbox, DriverSandboxTemplate}; +use openshell_core::proto::compute::v1::{ + DriverSandbox, DriverSandboxTemplate, ResourceRequirements, +}; use openshell_core::proto_struct::deserialize_optional_non_empty_string_list; use openshell_core::{driver_mounts, proto_struct}; use serde::Serialize; @@ -628,32 +630,57 @@ fn build_labels(sandbox: &DriverSandbox) -> BTreeMap { labels } -/// Parse resource limits from the sandbox template, falling back to defaults. -fn build_resource_limits(sandbox: &DriverSandbox, config: &PodmanComputeConfig) -> ResourceLimits { +/// Parse resource limits from typed sandbox requirements, falling back to defaults. +fn build_resource_limits( + sandbox: &DriverSandbox, + config: &PodmanComputeConfig, +) -> Result { let resources = sandbox .spec .as_ref() - .and_then(|s| s.template.as_ref()) - .and_then(|t| t.resources.as_ref()); + .and_then(|s| s.resource_requirements.as_ref()); - let cpu_micros = resources - .filter(|r| !r.cpu_limit.is_empty()) - .and_then(|r| parse_cpu_to_microseconds(&r.cpu_limit)) - .unwrap_or(DEFAULT_CPU_QUOTA); + let cpu_micros = parse_podman_cpu_limit(resources)?.unwrap_or(DEFAULT_CPU_QUOTA); + let mem_bytes = parse_podman_memory_limit(resources)?.unwrap_or(DEFAULT_MEMORY_LIMIT); - let mem_bytes = resources - .filter(|r| !r.memory_limit.is_empty()) - .and_then(|r| parse_memory_to_bytes(&r.memory_limit)) - .unwrap_or(DEFAULT_MEMORY_LIMIT); - - ResourceLimits { + Ok(ResourceLimits { cpu: CpuLimits { quota: cpu_micros, period: DEFAULT_CPU_PERIOD, }, memory: MemoryLimits { limit: mem_bytes }, pids_limit: podman_pids_limit(config.sandbox_pids_limit), - } + }) +} + +fn parse_podman_cpu_limit( + resources: Option<&ResourceRequirements>, +) -> Result, ComputeDriverError> { + let Some(cpu) = resources.and_then(|resources| resources.cpu.as_ref()) else { + return Ok(None); + }; + parse_cpu_to_microseconds(&cpu.limit) + .map(Some) + .ok_or_else(|| { + ComputeDriverError::Precondition(format!( + "invalid podman cpu limit '{}'; expected positive cores or millicores", + cpu.limit + )) + }) +} + +fn parse_podman_memory_limit( + resources: Option<&ResourceRequirements>, +) -> Result, ComputeDriverError> { + let Some(memory) = resources.and_then(|resources| resources.memory.as_ref()) else { + return Ok(None); + }; + parse_memory_to_bytes(&memory.limit).map(Some).ok_or_else(|| { + ComputeDriverError::Precondition(format!( + "invalid podman memory limit '{}'; expected positive bytes or a Kubernetes-style quantity", + memory.limit + )) + }) } fn podman_pids_limit(value: i64) -> Option { @@ -1022,7 +1049,7 @@ pub fn build_container_spec_for_image( let env = build_env(sandbox, config, requested_image, oci_user); let labels = build_labels(sandbox); - let resource_limits = build_resource_limits(sandbox, config); + let resource_limits = build_resource_limits(sandbox, config)?; let user_mounts = podman_user_mounts(sandbox, config.enable_bind_mounts) .map_err(ComputeDriverError::InvalidArgument)?; if sandbox @@ -1490,7 +1517,10 @@ fn parse_memory_to_bytes(quantity: &str) -> Option { #[cfg(test)] mod tests { use super::*; - use openshell_core::proto::compute::v1::{GpuResourceRequirements, ResourceRequirements}; + use openshell_core::proto::compute::v1::{ + CpuResourceRequirements, GpuResourceRequirements, MemoryResourceRequirements, + ResourceRequirements, + }; static ENV_LOCK: std::sync::LazyLock> = std::sync::LazyLock::new(|| std::sync::Mutex::new(())); @@ -1506,6 +1536,8 @@ mod tests { fn gpu_resources(count: Option) -> ResourceRequirements { ResourceRequirements { gpu: Some(GpuResourceRequirements { count }), + cpu: None, + memory: None, } } @@ -1544,19 +1576,19 @@ mod tests { #[test] fn container_spec_applies_cpu_and_memory_limits() { - use openshell_core::proto::compute::v1::{ - DriverResourceRequirements, DriverSandboxSpec, DriverSandboxTemplate, - }; + use openshell_core::proto::compute::v1::{DriverSandboxSpec, DriverSandboxTemplate}; let mut sandbox = test_sandbox("test-id", "test-name"); sandbox.spec = Some(DriverSandboxSpec { - template: Some(DriverSandboxTemplate { - resources: Some(DriverResourceRequirements { - cpu_limit: "500m".to_string(), - memory_limit: "2Gi".to_string(), - ..Default::default() + template: Some(DriverSandboxTemplate::default()), + resource_requirements: Some(ResourceRequirements { + gpu: None, + cpu: Some(CpuResourceRequirements { + limit: "500m".to_string(), + }), + memory: Some(MemoryResourceRequirements { + limit: "2Gi".to_string(), }), - ..Default::default() }), ..Default::default() }); diff --git a/crates/openshell-driver-podman/src/driver.rs b/crates/openshell-driver-podman/src/driver.rs index a11189cbc6..257ca9dafb 100644 --- a/crates/openshell-driver-podman/src/driver.rs +++ b/crates/openshell-driver-podman/src/driver.rs @@ -1643,6 +1643,8 @@ mod tests { fn gpu_resources(count: Option) -> ResourceRequirements { ResourceRequirements { gpu: Some(GpuResourceRequirements { count }), + cpu: None, + memory: None, } } diff --git a/crates/openshell-driver-vm/src/driver.rs b/crates/openshell-driver-vm/src/driver.rs index 13e57f546d..15258edeea 100644 --- a/crates/openshell-driver-vm/src/driver.rs +++ b/crates/openshell-driver-vm/src/driver.rs @@ -3530,6 +3530,7 @@ fn validate_vm_sandbox(sandbox: &Sandbox, gpu_enabled: bool) -> Result<(), Statu if let Some(template) = spec.template.as_ref() { validate_vm_sandbox_template(template)?; } + validate_cpu_memory_request(spec)?; validate_gpu_request(sandbox, gpu_enabled)?; Ok(()) @@ -3550,6 +3551,22 @@ fn validate_vm_sandbox_template(template: &SandboxTemplate) -> Result<(), Status Ok(()) } +#[allow(clippy::result_large_err)] +fn validate_cpu_memory_request( + spec: &openshell_core::proto::compute::v1::DriverSandboxSpec, +) -> Result<(), Status> { + let resources = spec.resource_requirements.as_ref(); + if resources + .is_some_and(|requirements| requirements.cpu.is_some() || requirements.memory.is_some()) + { + return Err(Status::failed_precondition( + "vm sandboxes do not support spec.resource_requirements.cpu or spec.resource_requirements.memory yet; configure VM driver vcpus and mem_mib instead", + )); + } + + Ok(()) +} + #[allow(clippy::result_large_err)] fn validate_gpu_request(sandbox: &Sandbox, gpu_enabled: bool) -> Result<(), Status> { let spec = sandbox @@ -5579,8 +5596,9 @@ mod tests { PROGRESS_COMPLETE_STEP_KEY, }; use openshell_core::proto::compute::v1::{ - DriverSandboxSpec as SandboxSpec, DriverSandboxTemplate as SandboxTemplate, - GpuResourceRequirements, ResourceRequirements, + CpuResourceRequirements, DriverSandboxSpec as SandboxSpec, + DriverSandboxTemplate as SandboxTemplate, GpuResourceRequirements, + MemoryResourceRequirements, ResourceRequirements, }; use prost_types::{Struct, Value, value::Kind}; use std::fs; @@ -6203,6 +6221,8 @@ mod tests { fn gpu_resources(count: Option) -> ResourceRequirements { ResourceRequirements { gpu: Some(GpuResourceRequirements { count }), + cpu: None, + memory: None, } } @@ -6574,26 +6594,28 @@ mod tests { } #[test] - fn validate_vm_sandbox_accepts_template_resources_as_noop() { - use openshell_core::proto::compute::v1::DriverResourceRequirements; - + fn validate_vm_sandbox_rejects_typed_cpu_and_memory_resources() { let sandbox = Sandbox { id: "sandbox-123".to_string(), spec: Some(SandboxSpec { - template: Some(SandboxTemplate { - resources: Some(DriverResourceRequirements { - cpu_limit: "2".to_string(), - memory_limit: "4Gi".to_string(), - ..Default::default() + resource_requirements: Some(ResourceRequirements { + gpu: None, + cpu: Some(CpuResourceRequirements { + limit: "2".to_string(), + }), + memory: Some(MemoryResourceRequirements { + limit: "4Gi".to_string(), }), - ..Default::default() }), ..Default::default() }), ..Default::default() }; - validate_vm_sandbox(&sandbox, false) - .expect("template.resources should be accepted and ignored"); + let err = validate_vm_sandbox(&sandbox, false).expect_err( + "typed CPU/memory resources should be rejected until VM sizing is supported", + ); + assert_eq!(err.code(), Code::FailedPrecondition); + assert!(err.message().contains("spec.resource_requirements.cpu")); } #[test] diff --git a/crates/openshell-sdk/src/client.rs b/crates/openshell-sdk/src/client.rs index f95b7ee111..825c4b2919 100644 --- a/crates/openshell-sdk/src/client.rs +++ b/crates/openshell-sdk/src/client.rs @@ -808,6 +808,8 @@ fn create_sandbox_request(spec: SandboxSpec) -> proto::CreateSandboxRequest { }); let resource_requirements = gpu.then_some(proto::ResourceRequirements { gpu: Some(proto::GpuResourceRequirements { count: None }), + cpu: None, + memory: None, }); proto::CreateSandboxRequest { spec: Some(proto::SandboxSpec { diff --git a/crates/openshell-server/src/compute/mod.rs b/crates/openshell-server/src/compute/mod.rs index b85641c986..a8ac22e564 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -34,13 +34,14 @@ use openshell_core::ComputeDriverKind; #[cfg(target_os = "windows")] use openshell_core::proto::SandboxPolicy; use openshell_core::proto::compute::v1::{ - CreateSandboxRequest, DeleteSandboxRequest, DeleteWorkspaceRequest, DeleteWorkspaceResponse, - DriverCondition, DriverPlatformEvent, DriverResourceRequirements, DriverSandbox, - DriverSandboxSpec, DriverSandboxStatus, DriverSandboxTemplate, EnsureWorkspaceRequest, - EnsureWorkspaceResponse, GatewayListenerRequirement as ProtoGatewayListenerRequirement, - GetCapabilitiesRequest, GetGatewayListenerRequirementsRequest, - GetGatewayListenerRequirementsResponse, GetSandboxRequest, - GpuResourceRequirements as DriverGpuResourceRequirements, ListSandboxesRequest, + CpuResourceRequirements as DriverCpuResourceRequirements, CreateSandboxRequest, + DeleteSandboxRequest, DeleteWorkspaceRequest, DeleteWorkspaceResponse, DriverCondition, + DriverPlatformEvent, DriverSandbox, DriverSandboxSpec, DriverSandboxStatus, + DriverSandboxTemplate, EnsureWorkspaceRequest, EnsureWorkspaceResponse, + GatewayListenerRequirement as ProtoGatewayListenerRequirement, GetCapabilitiesRequest, + GetGatewayListenerRequirementsRequest, GetGatewayListenerRequirementsResponse, + GetSandboxRequest, GpuResourceRequirements as DriverGpuResourceRequirements, + ListSandboxesRequest, MemoryResourceRequirements as DriverMemoryResourceRequirements, ResourceRequirements as DriverSandboxResourceRequirements, StartSandboxRequest, StopSandboxRequest, ValidateSandboxCreateRequest, WatchSandboxesEvent, WatchSandboxesRequest, compute_driver_client::ComputeDriverClient, compute_driver_server::ComputeDriver, @@ -3481,6 +3482,17 @@ fn driver_sandbox_spec_from_public( .gpu .as_ref() .map(|gpu| DriverGpuResourceRequirements { count: gpu.count }), + cpu: requirements + .cpu + .as_ref() + .map(|cpu| DriverCpuResourceRequirements { + limit: cpu.limit.clone(), + }), + memory: requirements.memory.as_ref().map(|memory| { + DriverMemoryResourceRequirements { + limit: memory.limit.clone(), + } + }), } }), sandbox_token: String::new(), @@ -3498,7 +3510,6 @@ fn driver_sandbox_template_from_public( agent_socket_path: template.agent_socket.clone(), labels: template.labels.clone(), environment: template.environment.clone(), - resources: extract_typed_resources(&template.resources), platform_config: build_platform_config(template), driver_config: select_driver_config(&template.driver_config, driver_name)?, }) @@ -3522,50 +3533,6 @@ fn select_driver_config( } } -/// Extract typed CPU/memory quantities from the public `resources` Struct. -/// -/// The public API exposes resources as an untyped `google.protobuf.Struct` -/// with the Kubernetes limits/requests shape. We pull out the well-known -/// keys into the typed `DriverResourceRequirements` message. -fn extract_typed_resources( - resources: &Option, -) -> Option { - fn get_quantity(s: &prost_types::Struct, section: &str, key: &str) -> String { - s.fields - .get(section) - .and_then(|v| match v.kind.as_ref() { - Some(prost_types::value::Kind::StructValue(inner)) => inner.fields.get(key), - _ => None, - }) - .and_then(|v| match v.kind.as_ref() { - Some(prost_types::value::Kind::StringValue(val)) => Some(val.clone()), - _ => None, - }) - .unwrap_or_default() - } - - let s = resources.as_ref()?; - - let req = DriverResourceRequirements { - cpu_request: get_quantity(s, "requests", "cpu"), - cpu_limit: get_quantity(s, "limits", "cpu"), - memory_request: get_quantity(s, "requests", "memory"), - memory_limit: get_quantity(s, "limits", "memory"), - }; - - // Return None when all fields are empty so drivers can distinguish - // "no resource requirements" from "zero requirements". - if req.cpu_request.is_empty() - && req.cpu_limit.is_empty() - && req.memory_request.is_empty() - && req.memory_limit.is_empty() - { - None - } else { - Some(req) - } -} - /// Build the opaque `platform_config` Struct from platform-specific public /// template fields (`runtime_class_name`, annotations) plus any resource fields /// beyond CPU/memory. @@ -3619,9 +3586,8 @@ fn build_platform_config(template: &SandboxTemplate) -> Option Result<(), Status> { Ok(()) } -fn validate_gpu_request_fields(spec: &openshell_core::proto::SandboxSpec) -> Result<(), Status> { +fn validate_resource_requirement_fields( + spec: &openshell_core::proto::SandboxSpec, +) -> Result<(), Status> { if openshell_core::gpu::sandbox_gpu_count(spec.resource_requirements.as_ref()) == Some(0) { return Err(Status::invalid_argument("gpu count must be greater than 0")); } + if let Some(cpu) = spec + .resource_requirements + .as_ref() + .and_then(|requirements| requirements.cpu.as_ref()) + { + validate_cpu_quantity(&cpu.limit, "spec.resource_requirements.cpu.limit")?; + } + + if let Some(memory) = spec + .resource_requirements + .as_ref() + .and_then(|requirements| requirements.memory.as_ref()) + { + validate_memory_quantity(&memory.limit, "spec.resource_requirements.memory.limit")?; + } + + Ok(()) +} + +fn validate_cpu_quantity(value: &str, field_name: &str) -> Result<(), Status> { + let value = value.trim(); + if value.is_empty() { + return Err(Status::invalid_argument(format!( + "{field_name} must not be empty" + ))); + } + + if let Some(millicores) = value.strip_suffix('m') { + if millicores.is_empty() || !millicores.bytes().all(|b| b.is_ascii_digit()) { + return Err(Status::invalid_argument(format!( + "invalid {field_name} value '{value}': expected positive cores or millicores, for example 2, 0.5, or 500m" + ))); + } + let millicores = millicores.parse::().map_err(|_| { + Status::invalid_argument(format!( + "invalid {field_name} value '{value}': expected positive cores or millicores, for example 2, 0.5, or 500m" + )) + })?; + if millicores == 0 { + return Err(Status::invalid_argument(format!( + "{field_name} must be greater than zero" + ))); + } + return Ok(()); + } + + let cores = value.parse::().map_err(|_| { + Status::invalid_argument(format!( + "invalid {field_name} value '{value}': expected positive cores or millicores, for example 2, 0.5, or 500m" + )) + })?; + if !cores.is_finite() || cores <= 0.0 { + return Err(Status::invalid_argument(format!( + "{field_name} must be greater than zero" + ))); + } + Ok(()) +} + +fn validate_memory_quantity(value: &str, field_name: &str) -> Result<(), Status> { + let value = value.trim(); + if value.is_empty() { + return Err(Status::invalid_argument(format!( + "{field_name} must not be empty" + ))); + } + + let number_end = value + .find(|ch: char| !ch.is_ascii_digit()) + .unwrap_or(value.len()); + let (number, suffix) = value.split_at(number_end); + if number.is_empty() + || !matches!( + suffix, + "" | "Ki" | "Mi" | "Gi" | "Ti" | "Pi" | "Ei" | "K" | "M" | "G" | "T" | "P" | "E" + ) + { + return Err(Status::invalid_argument(format!( + "invalid {field_name} value '{value}': expected positive bytes or a quantity such as 512Mi, 4Gi, or 8G" + ))); + } + + let amount = number.parse::().map_err(|_| { + Status::invalid_argument(format!( + "invalid {field_name} value '{value}': expected positive bytes or a quantity such as 512Mi, 4Gi, or 8G" + )) + })?; + if amount == 0 { + return Err(Status::invalid_argument(format!( + "{field_name} must be greater than zero" + ))); + } Ok(()) } @@ -299,6 +393,7 @@ fn validate_sandbox_template(tmpl: &SandboxTemplate) -> Result<(), Status> { "template.resources serialized size exceeds maximum ({size} > {MAX_TEMPLATE_STRUCT_SIZE})" ))); } + reject_legacy_template_cpu_memory_resources(s)?; } if let Some(ref s) = tmpl.driver_config { let size = s.encoded_len(); @@ -312,6 +407,34 @@ fn validate_sandbox_template(tmpl: &SandboxTemplate) -> Result<(), Status> { Ok(()) } +fn reject_legacy_template_cpu_memory_resources( + resources: &prost_types::Struct, +) -> Result<(), Status> { + for section_name in ["limits", "requests"] { + let Some(section) = + resources + .fields + .get(section_name) + .and_then(|value| match value.kind.as_ref() { + Some(prost_types::value::Kind::StructValue(section)) => Some(section), + _ => None, + }) + else { + continue; + }; + + for resource_name in ["cpu", "memory"] { + if section.fields.contains_key(resource_name) { + return Err(Status::invalid_argument(format!( + "template.resources.{section_name}.{resource_name} is no longer supported; use spec.resource_requirements.{resource_name}.limit" + ))); + } + } + } + + Ok(()) +} + /// Validate a `map` field: entry count, key length, value length. pub(super) fn validate_string_map( map: &std::collections::HashMap, @@ -1004,6 +1127,8 @@ mod tests { let spec = SandboxSpec { resource_requirements: Some(openshell_core::proto::ResourceRequirements { gpu: Some(openshell_core::proto::GpuResourceRequirements { count: None }), + cpu: None, + memory: None, }), ..Default::default() }; @@ -1015,6 +1140,8 @@ mod tests { let spec = SandboxSpec { resource_requirements: Some(openshell_core::proto::ResourceRequirements { gpu: Some(openshell_core::proto::GpuResourceRequirements { count: Some(2) }), + cpu: None, + memory: None, }), ..Default::default() }; @@ -1026,6 +1153,8 @@ mod tests { let spec = SandboxSpec { resource_requirements: Some(openshell_core::proto::ResourceRequirements { gpu: Some(openshell_core::proto::GpuResourceRequirements { count: Some(0) }), + cpu: None, + memory: None, }), ..Default::default() }; @@ -1034,6 +1163,82 @@ mod tests { assert!(err.message().contains("gpu count must be greater than 0")); } + #[test] + fn validate_sandbox_spec_accepts_cpu_and_memory_requirements() { + let spec = SandboxSpec { + resource_requirements: Some(openshell_core::proto::ResourceRequirements { + gpu: None, + cpu: Some(openshell_core::proto::CpuResourceRequirements { + limit: "500m".to_string(), + }), + memory: Some(openshell_core::proto::MemoryResourceRequirements { + limit: "2Gi".to_string(), + }), + }), + ..Default::default() + }; + + assert!(validate_sandbox_spec("compute-sandbox", &spec).is_ok()); + } + + #[test] + fn validate_sandbox_spec_rejects_invalid_cpu_requirements() { + let spec = SandboxSpec { + resource_requirements: Some(openshell_core::proto::ResourceRequirements { + gpu: None, + cpu: Some(openshell_core::proto::CpuResourceRequirements { + limit: "0".to_string(), + }), + memory: Some(openshell_core::proto::MemoryResourceRequirements { + limit: "2Gi".to_string(), + }), + }), + ..Default::default() + }; + + let err = validate_sandbox_spec("compute-sandbox", &spec).unwrap_err(); + assert_eq!(err.code(), Code::InvalidArgument); + assert!( + err.message() + .contains("spec.resource_requirements.cpu.limit") + ); + } + + #[test] + fn validate_sandbox_spec_rejects_legacy_template_cpu_memory_resources() { + use prost_types::{Struct, Value, value::Kind}; + + let mut limits = std::collections::BTreeMap::new(); + limits.insert( + "cpu".to_string(), + Value { + kind: Some(Kind::StringValue("500m".to_string())), + }, + ); + let mut fields = std::collections::BTreeMap::new(); + fields.insert( + "limits".to_string(), + Value { + kind: Some(Kind::StructValue(Struct { fields: limits })), + }, + ); + let spec = SandboxSpec { + template: Some(SandboxTemplate { + resources: Some(Struct { fields }), + ..Default::default() + }), + ..Default::default() + }; + + let err = validate_sandbox_spec("legacy-resources", &spec).unwrap_err(); + assert_eq!(err.code(), Code::InvalidArgument); + assert!(err.message().contains("template.resources.limits.cpu")); + assert!( + err.message() + .contains("spec.resource_requirements.cpu.limit") + ); + } + #[test] fn validate_sandbox_spec_accepts_empty_defaults() { assert!(validate_sandbox_spec("", &default_spec()).is_ok()); diff --git a/proto/compute_driver.proto b/proto/compute_driver.proto index afa93f1b18..ba66d0beec 100644 --- a/proto/compute_driver.proto +++ b/proto/compute_driver.proto @@ -157,6 +157,10 @@ message DriverSandboxSpec { message ResourceRequirements { // GPU requirements for the sandbox. Presence indicates a GPU request. GpuResourceRequirements gpu = 1; + // CPU requirements for the sandbox workload. + CpuResourceRequirements cpu = 2; + // Memory requirements for the sandbox workload. + MemoryResourceRequirements memory = 3; } // Driver GPU resource requirements. @@ -166,6 +170,18 @@ message GpuResourceRequirements { optional uint32 count = 1; } +// Driver CPU resource requirements. +message CpuResourceRequirements { + // CPU limit for the sandbox workload (e.g. "500m", "2"). + string limit = 1; +} + +// Driver memory resource requirements. +message MemoryResourceRequirements { + // Memory limit for the sandbox workload (e.g. "512Mi", "4Gi"). + string limit = 1; +} + // Driver-owned runtime template consumed by the compute platform. // // This message describes the sandbox workload in backend-neutral terms. @@ -182,8 +198,6 @@ message DriverSandboxTemplate { map labels = 4; // Additional environment variables injected into the sandbox runtime. map environment = 6; - // Typed compute-resource requirements for the sandbox workload. - DriverResourceRequirements resources = 10; // Opaque, platform-specific configuration passed through to the driver. // The gateway does not inspect this; each driver defines its own schema. // For the Kubernetes driver this carries fields such as runtimeClassName, @@ -193,22 +207,8 @@ message DriverSandboxTemplate { // This is the inner block selected from public SandboxTemplate.driver_config. // The selected driver owns nested schema validation. google.protobuf.Struct driver_config = 12; -} - -// Typed compute-resource requirements. -// -// Values use Kubernetes-style quantity strings (e.g. "500m", "2", "4Gi") -// because they are a well-known, widely-adopted notation. Drivers for -// non-Kubernetes platforms must parse these strings into their native units. -message DriverResourceRequirements { - // Minimum CPU cores requested (e.g. "500m", "2"). - string cpu_request = 1; - // Maximum CPU cores allowed (e.g. "500m", "4"). - string cpu_limit = 2; - // Minimum memory requested (e.g. "256Mi", "4Gi"). - string memory_request = 3; - // Maximum memory allowed (e.g. "512Mi", "8Gi"). - string memory_limit = 4; + reserved 10; + reserved "resources"; } // Raw status observed directly from the compute platform. diff --git a/proto/openshell.proto b/proto/openshell.proto index 246fe0626f..4ca9959021 100644 --- a/proto/openshell.proto +++ b/proto/openshell.proto @@ -845,6 +845,10 @@ message SandboxSpec { message ResourceRequirements { // GPU requirements for the sandbox. Presence indicates a GPU request. GpuResourceRequirements gpu = 1; + // CPU requirements for the sandbox workload. + CpuResourceRequirements cpu = 2; + // Memory requirements for the sandbox workload. + MemoryResourceRequirements memory = 3; } // Public GPU resource requirements. @@ -854,6 +858,18 @@ message GpuResourceRequirements { optional uint32 count = 1; } +// Public CPU resource requirements. +message CpuResourceRequirements { + // CPU limit for the sandbox workload (e.g. "500m", "2"). + string limit = 1; +} + +// Public memory resource requirements. +message MemoryResourceRequirements { + // Memory limit for the sandbox workload (e.g. "512Mi", "4Gi"). + string limit = 1; +} + // Public sandbox template mapped onto compute-driver template inputs. message SandboxTemplate { // Fully-qualified OCI image reference used to boot the sandbox. @@ -868,7 +884,8 @@ message SandboxTemplate { map annotations = 5; // Additional environment variables injected by the template. map environment = 6; - // Platform-specific compute resource requirements and limits. + // Platform-specific resource passthrough. CPU and memory under + // limits/requests are rejected; use ResourceRequirements.cpu/memory instead. google.protobuf.Struct resources = 7; reserved 9; reserved "volume_claim_templates";