Skip to content

Commit 6817cbc

Browse files
committed
feat(gpu): move device selection to driver config
1 parent 6416702 commit 6817cbc

23 files changed

Lines changed: 655 additions & 109 deletions

File tree

Cargo.lock

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

TESTING.md

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,9 @@ defaults to the image used by the `gateway` stage in
157157
`deploy/docker/Dockerfile.images`; set `OPENSHELL_E2E_GPU_PROBE_IMAGE` to
158158
override it. Per-device checks run only for NVIDIA CDI device IDs reported by
159159
the runtime's discovered devices list, so WSL2 hosts that expose only
160-
`nvidia.com/gpu=all` skip the index-based cases.
160+
`nvidia.com/gpu=all` skip the index-based cases. Exact CDI device selection is
161+
passed through `--driver-config-json` with the active Docker or Podman driver
162+
key.
161163

162164
Run the Docker-backed Rust CLI e2e suite:
163165

crates/openshell-cli/src/main.rs

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1213,12 +1213,6 @@ enum SandboxCommands {
12131213
#[arg(long)]
12141214
gpu: bool,
12151215

1216-
/// Target a driver-specific GPU device. Docker and Podman use CDI device IDs
1217-
/// (for example "nvidia.com/gpu=0"); VM uses a PCI BDF or index.
1218-
/// Only valid with --gpu. When omitted with --gpu, the driver uses its default GPU selection.
1219-
#[arg(long, requires = "gpu")]
1220-
gpu_device: Option<String>,
1221-
12221216
/// CPU limit for the sandbox (for example: 500m, 1, 2.5).
12231217
#[arg(long)]
12241218
cpu: Option<String>,
@@ -2554,7 +2548,6 @@ async fn main() -> Result<()> {
25542548
no_keep,
25552549
editor,
25562550
gpu,
2557-
gpu_device,
25582551
cpu,
25592552
memory,
25602553
driver_config_json,
@@ -2639,7 +2632,6 @@ async fn main() -> Result<()> {
26392632
&upload_specs,
26402633
keep,
26412634
gpu,
2642-
gpu_device.as_deref(),
26432635
cpu.as_deref(),
26442636
memory.as_deref(),
26452637
driver_config_json.as_deref(),

crates/openshell-cli/src/run.rs

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1725,7 +1725,6 @@ pub async fn sandbox_create(
17251725
uploads: &[(String, Option<String>, bool)],
17261726
keep: bool,
17271727
gpu: bool,
1728-
gpu_device: Option<&str>,
17291728
cpu: Option<&str>,
17301729
memory: Option<&str>,
17311730
driver_config_json: Option<&str>,
@@ -1817,7 +1816,6 @@ pub async fn sandbox_create(
18171816
let request = CreateSandboxRequest {
18181817
spec: Some(SandboxSpec {
18191818
gpu: requested_gpu,
1820-
gpu_device: gpu_device.unwrap_or_default().to_string(),
18211819
environment: environment.clone(),
18221820
policy,
18231821
providers: configured_providers,

crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -787,7 +787,6 @@ async fn sandbox_create_keeps_command_sessions_by_default() {
787787
None,
788788
None,
789789
None,
790-
None,
791790
&[],
792791
None,
793792
None,
@@ -827,7 +826,6 @@ async fn sandbox_create_sends_cpu_and_memory_limits_only() {
827826
&[],
828827
true,
829828
false,
830-
None,
831829
Some("500m"),
832830
Some("2Gi"),
833831
None,
@@ -907,7 +905,6 @@ async fn sandbox_create_sends_driver_config_json() {
907905
false,
908906
None,
909907
None,
910-
None,
911908
Some(r#"{"kubernetes":{"pod":{"priority_class_name":"batch-low"}}}"#),
912909
None,
913910
&[],
@@ -984,7 +981,6 @@ async fn sandbox_create_does_not_infer_command_providers_when_v2_enabled() {
984981
None,
985982
None,
986983
None,
987-
None,
988984
&[],
989985
None,
990986
None,
@@ -1043,7 +1039,6 @@ async fn sandbox_create_returns_vm_error_without_waiting_for_timeout() {
10431039
None,
10441040
None,
10451041
None,
1046-
None,
10471042
&[],
10481043
None,
10491044
None,
@@ -1098,7 +1093,6 @@ async fn sandbox_create_keeps_waiting_while_vm_progress_arrives() {
10981093
None,
10991094
None,
11001095
None,
1101-
None,
11021096
&[],
11031097
None,
11041098
None,
@@ -1145,7 +1139,6 @@ async fn sandbox_create_times_out_when_only_logs_arrive() {
11451139
None,
11461140
None,
11471141
None,
1148-
None,
11491142
&[],
11501143
None,
11511144
None,
@@ -1188,7 +1181,6 @@ async fn sandbox_create_deletes_command_sessions_with_no_keep() {
11881181
None,
11891182
None,
11901183
None,
1191-
None,
11921184
&[],
11931185
None,
11941186
None,
@@ -1235,7 +1227,6 @@ async fn sandbox_create_deletes_shell_sessions_with_no_keep() {
12351227
None,
12361228
None,
12371229
None,
1238-
None,
12391230
&[],
12401231
None,
12411232
None,
@@ -1282,7 +1273,6 @@ async fn sandbox_create_keeps_sandbox_with_hidden_keep_flag() {
12821273
None,
12831274
None,
12841275
None,
1285-
None,
12861276
&[],
12871277
None,
12881278
None,
@@ -1329,7 +1319,6 @@ async fn sandbox_create_keeps_sandbox_with_forwarding() {
13291319
None,
13301320
None,
13311321
None,
1332-
None,
13331322
&[],
13341323
None,
13351324
Some(openshell_core::forward::ForwardSpec::new(forward_port)),
@@ -1372,7 +1361,6 @@ async fn sandbox_create_sends_environment_variables() {
13721361
None,
13731362
None,
13741363
None,
1375-
None,
13761364
&[],
13771365
None,
13781366
None,

crates/openshell-core/src/error.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,9 @@ pub enum ComputeDriverError {
113113
/// The requested sandbox already exists.
114114
#[error("sandbox already exists")]
115115
AlreadyExists,
116+
/// The request contains an invalid argument.
117+
#[error("{0}")]
118+
InvalidArgument(String),
116119
/// A precondition for the operation was not met.
117120
#[error("{0}")]
118121
Precondition(String),
@@ -125,6 +128,7 @@ impl From<ComputeDriverError> for tonic::Status {
125128
fn from(err: ComputeDriverError) -> Self {
126129
match err {
127130
ComputeDriverError::AlreadyExists => Self::already_exists("sandbox already exists"),
131+
ComputeDriverError::InvalidArgument(m) => Self::invalid_argument(m),
128132
ComputeDriverError::Precondition(m) => Self::failed_precondition(m),
129133
ComputeDriverError::Message(m) => Self::internal(m),
130134
}

crates/openshell-core/src/gpu.rs

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5,18 +5,18 @@
55
66
use crate::config::CDI_GPU_DEVICE_ALL;
77

8-
/// Resolve the existing GPU request fields into CDI device identifiers.
8+
/// Resolve a GPU request into CDI device identifiers.
99
///
10-
/// `None` means no GPU was requested. A GPU request with no explicit device
11-
/// ID uses the CDI all-GPU request; otherwise the driver-native ID passes
12-
/// through unchanged.
10+
/// `None` means no GPU was requested. A GPU request with no explicit CDI
11+
/// devices uses the CDI all-GPU request; otherwise the driver-configured CDI
12+
/// devices pass through unchanged.
1313
#[must_use]
14-
pub fn cdi_gpu_device_ids(gpu: bool, gpu_device: &str) -> Option<Vec<String>> {
14+
pub fn cdi_gpu_device_ids(gpu: bool, cdi_devices: &[String]) -> Option<Vec<String>> {
1515
gpu.then(|| {
16-
if gpu_device.is_empty() {
16+
if cdi_devices.is_empty() {
1717
vec![CDI_GPU_DEVICE_ALL.to_string()]
1818
} else {
19-
vec![gpu_device.to_string()]
19+
cdi_devices.to_vec()
2020
}
2121
})
2222
}
@@ -27,22 +27,31 @@ mod tests {
2727

2828
#[test]
2929
fn cdi_gpu_device_ids_returns_none_when_absent() {
30-
assert_eq!(cdi_gpu_device_ids(false, ""), None);
30+
assert_eq!(cdi_gpu_device_ids(false, &[]), None);
3131
}
3232

3333
#[test]
3434
fn cdi_gpu_device_ids_defaults_empty_request_to_all_gpus() {
3535
assert_eq!(
36-
cdi_gpu_device_ids(true, ""),
36+
cdi_gpu_device_ids(true, &[]),
3737
Some(vec![CDI_GPU_DEVICE_ALL.to_string()])
3838
);
3939
}
4040

4141
#[test]
42-
fn cdi_gpu_device_ids_passes_explicit_device_id_through() {
42+
fn cdi_gpu_device_ids_passes_explicit_device_ids_through() {
4343
assert_eq!(
44-
cdi_gpu_device_ids(true, "nvidia.com/gpu=0"),
45-
Some(vec!["nvidia.com/gpu=0".to_string()])
44+
cdi_gpu_device_ids(
45+
true,
46+
&[
47+
"nvidia.com/gpu=0".to_string(),
48+
"nvidia.com/gpu=1".to_string()
49+
]
50+
),
51+
Some(vec![
52+
"nvidia.com/gpu=0".to_string(),
53+
"nvidia.com/gpu=1".to_string()
54+
])
4655
);
4756
}
4857
}

crates/openshell-core/src/proto_struct.rs

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33

44
//! Helpers for decoding `google.protobuf.Struct` values.
55
6+
use serde::{Deserialize, Deserializer, de::Error as _};
7+
68
/// Convert a protobuf Struct into a JSON object for typed serde decoding.
79
#[must_use]
810
pub fn struct_to_json_object(
@@ -43,3 +45,76 @@ pub fn value_to_json(value: &prost_types::Value) -> serde_json::Value {
4345
Some(prost_types::value::Kind::NullValue(_)) | None => serde_json::Value::Null,
4446
}
4547
}
48+
49+
/// Deserialize a present field as a non-empty list of non-empty strings.
50+
///
51+
/// Use with `#[serde(default, deserialize_with = "...")]` on
52+
/// `Option<Vec<String>>` fields. Missing fields use the option default; present
53+
/// fields must be arrays and cannot be empty.
54+
pub fn deserialize_optional_non_empty_string_list<'de, D>(
55+
deserializer: D,
56+
) -> Result<Option<Vec<String>>, D::Error>
57+
where
58+
D: Deserializer<'de>,
59+
{
60+
let values = Vec::<String>::deserialize(deserializer)?;
61+
if values.is_empty() {
62+
return Err(D::Error::custom("must be a non-empty list of strings"));
63+
}
64+
65+
for (idx, value) in values.iter().enumerate() {
66+
if value.trim().is_empty() {
67+
return Err(D::Error::custom(format!(
68+
"[{idx}] must be a non-empty string"
69+
)));
70+
}
71+
}
72+
73+
Ok(Some(values))
74+
}
75+
76+
#[cfg(test)]
77+
mod tests {
78+
use super::*;
79+
80+
#[derive(Debug, Default, Deserialize)]
81+
#[serde(default)]
82+
struct TestConfig {
83+
#[serde(
84+
default,
85+
deserialize_with = "deserialize_optional_non_empty_string_list"
86+
)]
87+
devices: Option<Vec<String>>,
88+
}
89+
90+
#[test]
91+
fn optional_non_empty_string_list_defaults_when_absent() {
92+
let config: TestConfig = serde_json::from_value(serde_json::json!({})).unwrap();
93+
94+
assert_eq!(config.devices, None);
95+
}
96+
97+
#[test]
98+
fn optional_non_empty_string_list_parses_present_list() {
99+
let config: TestConfig =
100+
serde_json::from_value(serde_json::json!({"devices": ["nvidia.com/gpu=0"]})).unwrap();
101+
102+
assert_eq!(config.devices, Some(vec!["nvidia.com/gpu=0".to_string()]));
103+
}
104+
105+
#[test]
106+
fn optional_non_empty_string_list_rejects_empty_list() {
107+
let err =
108+
serde_json::from_value::<TestConfig>(serde_json::json!({"devices": []})).unwrap_err();
109+
110+
assert!(err.to_string().contains("non-empty list"));
111+
}
112+
113+
#[test]
114+
fn optional_non_empty_string_list_rejects_empty_string() {
115+
let err =
116+
serde_json::from_value::<TestConfig>(serde_json::json!({"devices": [""]})).unwrap_err();
117+
118+
assert!(err.to_string().contains("non-empty string"));
119+
}
120+
}

crates/openshell-driver-docker/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ tempfile = "3"
2727
url = { workspace = true }
2828

2929
[dev-dependencies]
30+
prost-types = { workspace = true }
3031
temp-env = "0.3"
3132

3233
[lints]

crates/openshell-driver-docker/README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ contract:
3232
| `apparmor=unconfined` | Avoids Docker's default profile blocking required mount operations. |
3333
| `restart_policy = unless-stopped` | Keeps managed sandboxes resumable across daemon or gateway restarts. |
3434
| `PidsLimit` | Enforces the sandbox PID budget at the Docker cgroup layer. Set `[openshell.drivers.docker].sandbox_pids_limit = 0` to inherit the Docker/runtime default. |
35-
| CDI GPU request | Uses the sandbox `gpu_device` value when set; otherwise requests all NVIDIA GPUs when the sandbox spec asks for GPU support and daemon CDI support is detected. |
35+
| CDI GPU request | Uses `driver_config.cdi_devices` when set; otherwise requests all NVIDIA GPUs when the sandbox spec asks for GPU support and daemon CDI support is detected. |
3636

3737
The agent child process does not retain these supervisor privileges.
3838

0 commit comments

Comments
 (0)