1use indexmap::IndexMap;
6use schemars::JsonSchema;
7use serde::{Deserialize, Serialize};
8
9use crate::Error;
10
11#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
15#[serde(try_from = "GpusInput")]
16pub struct Gpus {
17 #[serde(default, skip_serializing_if = "Option::is_none")]
19 pub count: Option<GpuCount>,
20 #[serde(default, skip_serializing_if = "Vec::is_empty")]
23 pub ids: Vec<String>,
24 #[serde(default, skip_serializing_if = "Option::is_none")]
26 pub driver: Option<String>,
27 #[serde(default, skip_serializing_if = "Vec::is_empty")]
31 pub capabilities: Vec<String>,
32 #[serde(default, skip_serializing_if = "IndexMap::is_empty")]
34 pub options: IndexMap<String, String>,
35}
36
37pub const DEFAULT_GPU_DRIVER: &str = "nvidia";
39
40impl Gpus {
41 #[must_use]
43 pub fn driver(&self) -> &str {
44 self.driver.as_deref().unwrap_or(DEFAULT_GPU_DRIVER)
45 }
46
47 #[must_use]
49 pub fn fixed_count(&self) -> Option<u32> {
50 match self.count {
51 Some(GpuCount::All) => None,
52 Some(GpuCount::Count(n)) => Some(n),
53 None if !self.ids.is_empty() => u32::try_from(self.ids.len()).ok(),
54 None => Some(1),
55 }
56 }
57}
58
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61pub enum GpuCount {
62 Count(u32),
64 All,
66}
67
68impl GpuCount {
69 const ALL: &'static str = "all";
71}
72
73impl Serialize for GpuCount {
74 fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
75 match self {
76 GpuCount::Count(n) => s.serialize_u32(*n),
77 GpuCount::All => s.serialize_str(GpuCount::ALL),
78 }
79 }
80}
81
82#[derive(Deserialize, JsonSchema)]
84#[serde(untagged)]
85pub enum GpusInput {
86 Count(u32),
88 Word(String),
90 Full(GpusFull),
92}
93
94#[derive(Deserialize, JsonSchema)]
96#[serde(deny_unknown_fields)]
97pub struct GpusFull {
98 #[serde(default)]
100 #[schemars(schema_with = "count_schema")]
101 count: Option<serde_json::Value>,
102 #[serde(default)]
105 ids: Vec<String>,
106 #[serde(default)]
108 driver: Option<String>,
109 #[serde(default)]
111 capabilities: Vec<String>,
112 #[serde(default)]
114 options: IndexMap<String, String>,
115}
116
117fn count_schema(_: &mut schemars::SchemaGenerator) -> schemars::Schema {
119 schemars::json_schema!({
120 "anyOf": [
121 { "type": "integer", "format": "uint32", "minimum": 0, "description": "That many." },
122 { "type": "string", "const": GpuCount::ALL, "description": "Every GPU of the host." }
123 ]
124 })
125}
126
127fn count(value: &serde_json::Value) -> Result<GpuCount, Error> {
128 match value {
129 serde_json::Value::String(s) if s == GpuCount::ALL => Ok(GpuCount::All),
130 serde_json::Value::Number(n) => n
131 .as_u64()
132 .and_then(|n| u32::try_from(n).ok())
133 .map(GpuCount::Count)
134 .ok_or_else(|| Error::GpuCount(n.to_string())),
135 other => Err(Error::GpuCount(other.to_string())),
136 }
137}
138
139impl TryFrom<GpusInput> for Gpus {
140 type Error = Error;
141
142 fn try_from(input: GpusInput) -> Result<Gpus, Error> {
143 match input {
144 GpusInput::Count(n) => Ok(Gpus {
145 count: Some(GpuCount::Count(n)),
146 ..Gpus::default()
147 }),
148 GpusInput::Word(w) if w == GpuCount::ALL => Ok(Gpus {
149 count: Some(GpuCount::All),
150 ..Gpus::default()
151 }),
152 GpusInput::Word(w) => Err(Error::GpusWord(w)),
153 GpusInput::Full(f) => Ok(Gpus {
154 count: f.count.as_ref().map(count).transpose()?,
155 ids: f.ids,
156 driver: f.driver,
157 capabilities: f.capabilities,
158 options: f.options,
159 }),
160 }
161 }
162}
163
164#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
166#[serde(try_from = "DeviceInput")]
167pub struct Device {
168 pub host: String,
170 #[serde(default, skip_serializing_if = "Option::is_none")]
172 pub guest: Option<String>,
173 #[serde(default, skip_serializing_if = "Option::is_none")]
176 pub permissions: Option<String>,
177}
178
179impl Device {
180 #[must_use]
182 pub fn guest(&self) -> &str {
183 self.guest.as_deref().unwrap_or(&self.host)
184 }
185}
186
187#[derive(Deserialize, JsonSchema)]
189#[serde(untagged)]
190pub enum DeviceInput {
191 Short(String),
193 Full(DeviceFull),
195}
196
197#[derive(Deserialize, JsonSchema)]
199#[serde(deny_unknown_fields)]
200pub struct DeviceFull {
201 host: String,
203 #[serde(default)]
205 guest: Option<String>,
206 #[serde(default)]
208 permissions: Option<String>,
209}
210
211impl TryFrom<DeviceInput> for Device {
212 type Error = Error;
213
214 fn try_from(input: DeviceInput) -> Result<Device, Error> {
215 match input {
216 DeviceInput::Full(DeviceFull { host, guest, permissions }) => Ok(Device { host, guest, permissions }),
217 DeviceInput::Short(s) => {
218 let mut parts = s.split(':').map(String::from);
219 let host = parts.next().unwrap_or_default();
220 let device = Device {
221 host,
222 guest: parts.next(),
223 permissions: parts.next(),
224 };
225 match parts.next() {
226 Some(_) => Err(Error::Device(s)),
227 None => Ok(device),
228 }
229 }
230 }
231 }
232}
233
234#[cfg(test)]
235mod tests {
236 use super::*;
237
238 fn gpus(yaml: &str) -> Result<Gpus, String> {
239 serde_yaml_ng::from_str(yaml).map_err(|e| e.to_string())
240 }
241
242 #[test]
243 fn gpus_read_as_a_count_all_or_in_full() {
244 assert_eq!(gpus("2").unwrap().fixed_count(), Some(2));
245 assert_eq!(gpus("all").unwrap().count, Some(GpuCount::All));
246 let full = gpus("{ ids: ['0', '1'], driver: amd, capabilities: [compute] }").unwrap();
247 assert_eq!((full.fixed_count(), full.driver()), (Some(2), "amd"));
248 assert_eq!(gpus("{ count: all }").unwrap().fixed_count(), None);
249 assert_eq!(gpus("{}").unwrap().fixed_count(), Some(1));
250 assert_eq!(gpus("{}").unwrap().driver(), DEFAULT_GPU_DRIVER);
251 assert!(gpus("some").unwrap_err().contains("a number, `all` or a mapping"));
252 assert!(gpus("{ count: many }").is_err());
253 }
254
255 #[test]
256 fn devices_read_short_or_in_full() {
257 for (yaml, guest, permissions) in [
258 ("/dev/fuse", "/dev/fuse", None),
259 ("/dev/ttyUSB0:/dev/ttyS1:r", "/dev/ttyS1", Some("r")),
260 ("{ host: /dev/kvm, permissions: rw }", "/dev/kvm", Some("rw")),
261 ] {
262 let d: Device = serde_yaml_ng::from_str(yaml).unwrap();
263 assert_eq!((d.guest(), d.permissions.as_deref()), (guest, permissions), "{yaml}");
264 }
265 assert!(serde_yaml_ng::from_str::<Device>("/a:/b:r:x").is_err());
266 }
267
268 #[test]
269 fn they_serialize_in_full() {
270 let d: Device = serde_yaml_ng::from_str("/dev/fuse").unwrap();
271 assert_eq!(serde_yaml_ng::to_string(&d).unwrap(), "host: /dev/fuse\n");
272 assert_eq!(serde_yaml_ng::to_string(&gpus("all").unwrap()).unwrap(), "count: all\n");
273 }
274}