Skip to content

Commit fd2b5e8

Browse files
committed
refactor(sandbox): add ResourceSpec for sandbox resources
Signed-off-by: Ryan Angilly <rangilly@nvidia.com>
1 parent 44070a3 commit fd2b5e8

12 files changed

Lines changed: 595 additions & 269 deletions

File tree

architecture/compute-runtimes.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,8 +63,8 @@ Custom sandbox images must include the agent runtime and any system
6363
dependencies, but they should not need to include the gateway. GPU-capable
6464
images must include the user-space libraries required by the workload. The
6565
runtime still owns GPU device injection or resource scheduling. Kubernetes maps
66-
template resource limits such as `nvidia.com/gpu` into the sandbox pod when the
67-
cluster exposes those resources.
66+
portable `SandboxSpec.resources` CPU, memory, and GPU count requirements into
67+
pod resource requests and limits when the cluster exposes those resources.
6868

6969
## Deployment Shape
7070

crates/openshell-cli/src/main.rs

Lines changed: 81 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1010,6 +1010,7 @@ enum DoctorCommands {
10101010
}
10111011

10121012
#[derive(Subcommand, Debug)]
1013+
#[allow(clippy::large_enum_variant)]
10131014
enum SandboxCommands {
10141015
/// Create a sandbox.
10151016
#[command(help_template = LEAF_HELP_TEMPLATE, next_help_heading = "FLAGS")]
@@ -1068,21 +1069,33 @@ enum SandboxCommands {
10681069
///
10691070
/// This implies --gpu. Kubernetes-backed gateways schedule pods with
10701071
/// the corresponding nvidia.com/gpu resource limit.
1071-
#[arg(long, value_name = "COUNT", value_parser = clap::value_parser!(u32).range(1..))]
1072+
#[arg(long, value_name = "COUNT", value_parser = clap::value_parser!(u32).range(1..), conflicts_with = "gpu_device")]
10721073
gpu_count: Option<u32>,
10731074

1074-
/// Set compute resource requests and limits as JSON.
1075-
///
1076-
/// The JSON must be an object with the same shape as
1077-
/// SandboxTemplate.resources, for example:
1078-
/// {"requests":{"cpu":"2"},"limits":{"cpu":"16"}}
1079-
#[arg(long, value_name = "JSON")]
1080-
resources_json: Option<String>,
1075+
/// Minimum CPU cores requested, e.g. "500m" or "2".
1076+
#[arg(long, value_name = "QUANTITY")]
1077+
cpu_request: Option<String>,
1078+
1079+
/// Maximum CPU cores allowed, e.g. "2" or "4".
1080+
#[arg(long, value_name = "QUANTITY")]
1081+
cpu_limit: Option<String>,
1082+
1083+
/// Minimum memory requested, e.g. "512Mi" or "4Gi".
1084+
#[arg(long, value_name = "QUANTITY")]
1085+
memory_request: Option<String>,
1086+
1087+
/// Maximum memory allowed, e.g. "1Gi" or "8Gi".
1088+
#[arg(long, value_name = "QUANTITY")]
1089+
memory_limit: Option<String>,
1090+
1091+
/// Driver-specific resource configuration as KEY=VALUE.
1092+
#[arg(long = "resource-config", value_name = "KEY=VALUE")]
1093+
resource_config: Vec<String>,
10811094

10821095
/// Target a driver-specific GPU device. Docker and Podman use CDI device IDs
10831096
/// (for example "nvidia.com/gpu=0"); VM uses a PCI BDF or index.
10841097
/// Only valid with --gpu. When omitted with --gpu, the driver uses its default GPU selection.
1085-
#[arg(long, requires = "gpu")]
1098+
#[arg(long, requires = "gpu", conflicts_with = "gpu_count")]
10861099
gpu_device: Option<String>,
10871100

10881101
/// Provider names to attach to this sandbox.
@@ -2267,7 +2280,11 @@ async fn main() -> Result<()> {
22672280
editor,
22682281
gpu,
22692282
gpu_count,
2270-
resources_json,
2283+
cpu_request,
2284+
cpu_limit,
2285+
memory_request,
2286+
memory_limit,
2287+
resource_config,
22712288
gpu_device,
22722289
providers,
22732290
policy,
@@ -2334,8 +2351,14 @@ async fn main() -> Result<()> {
23342351
upload_spec.as_ref(),
23352352
keep,
23362353
gpu,
2337-
gpu_count,
2338-
resources_json.as_deref(),
2354+
run::SandboxResourceArgs {
2355+
cpu_request: cpu_request.as_deref(),
2356+
cpu_limit: cpu_limit.as_deref(),
2357+
memory_request: memory_request.as_deref(),
2358+
memory_limit: memory_limit.as_deref(),
2359+
gpu_count,
2360+
driver_config: &resource_config,
2361+
},
23392362
gpu_device.as_deref(),
23402363
editor,
23412364
&providers,
@@ -3552,27 +3575,66 @@ mod tests {
35523575
}
35533576

35543577
#[test]
3555-
fn sandbox_create_resources_json_parses() {
3578+
fn sandbox_create_resource_spec_flags_parse() {
35563579
let cli = Cli::try_parse_from([
35573580
"openshell",
35583581
"sandbox",
35593582
"create",
3560-
"--resources-json",
3561-
r#"{"requests":{"cpu":"2"},"limits":{"cpu":"16"}}"#,
3583+
"--cpu-request",
3584+
"2",
3585+
"--cpu-limit",
3586+
"4",
3587+
"--memory-request",
3588+
"8Gi",
3589+
"--memory-limit",
3590+
"16Gi",
3591+
"--resource-config",
3592+
"kubernetes.resource-name=nvidia.com/gpu",
35623593
])
3563-
.expect("sandbox create --resources-json should parse");
3594+
.expect("sandbox create resource flags should parse");
35643595

35653596
match cli.command {
35663597
Some(Commands::Sandbox {
3567-
command: Some(SandboxCommands::Create { resources_json, .. }),
3598+
command:
3599+
Some(SandboxCommands::Create {
3600+
cpu_request,
3601+
cpu_limit,
3602+
memory_request,
3603+
memory_limit,
3604+
resource_config,
3605+
..
3606+
}),
35683607
..
35693608
}) => {
3609+
assert_eq!(cpu_request.as_deref(), Some("2"));
3610+
assert_eq!(cpu_limit.as_deref(), Some("4"));
3611+
assert_eq!(memory_request.as_deref(), Some("8Gi"));
3612+
assert_eq!(memory_limit.as_deref(), Some("16Gi"));
35703613
assert_eq!(
3571-
resources_json.as_deref(),
3572-
Some(r#"{"requests":{"cpu":"2"},"limits":{"cpu":"16"}}"#)
3614+
resource_config,
3615+
vec!["kubernetes.resource-name=nvidia.com/gpu".to_string()]
35733616
);
35743617
}
35753618
other => panic!("expected sandbox create command, got: {other:?}"),
35763619
}
35773620
}
3621+
3622+
#[test]
3623+
fn sandbox_create_gpu_count_conflicts_with_gpu_device() {
3624+
let result = Cli::try_parse_from([
3625+
"openshell",
3626+
"sandbox",
3627+
"create",
3628+
"--gpu",
3629+
"--gpu-count",
3630+
"2",
3631+
"--gpu-device",
3632+
"nvidia.com/gpu=0",
3633+
]);
3634+
3635+
assert!(
3636+
result.is_err(),
3637+
"sandbox create should reject combining --gpu-count and --gpu-device"
3638+
);
3639+
}
35783640
}

0 commit comments

Comments
 (0)