Skip to main content

isoloom_spec/model/
hardware.rs

1//! What a machine asks of the hardware under it: GPUs, the host's devices, and nested
2//! virtualization. Each target either gives it or refuses the spec with the reason: a machine
3//! never silently runs without the hardware it asked for.
4
5use indexmap::IndexMap;
6use schemars::JsonSchema;
7use serde::{Deserialize, Serialize};
8
9use crate::Error;
10
11/// The GPUs a machine is given. A container gets the host's GPUs through its runtime; a VM gets
12/// them passed through by its hypervisor (Proxmox) or sized with them (a cloud's GPU instance
13/// types).
14#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
15#[serde(try_from = "GpusInput")]
16pub struct Gpus {
17    /// How many GPUs, or all the host has. Default: one, unless `ids` names them.
18    #[serde(default, skip_serializing_if = "Option::is_none")]
19    pub count: Option<GpuCount>,
20    /// Which GPUs, by the host's ids: a device index or UUID for a container, a PCI address
21    /// (`0000:01:00.0`) or a resource mapping's name on Proxmox. Not with `count`.
22    #[serde(default, skip_serializing_if = "Vec::is_empty")]
23    pub ids: Vec<String>,
24    /// The GPUs' vendor driver (`nvidia`, `amd`, `intel`). Default: `nvidia`.
25    #[serde(default, skip_serializing_if = "Option::is_none")]
26    pub driver: Option<String>,
27    /// What the container may use the GPUs for, as NVIDIA's runtime names it (`compute`,
28    /// `utility`, `video`, `graphics`, `display`). Default: the runtime's own. A VM gets the
29    /// whole device, so every capability.
30    #[serde(default, skip_serializing_if = "Vec::is_empty")]
31    pub capabilities: Vec<String>,
32    /// Options for the container runtime's GPU driver, passed as they are (Compose only).
33    #[serde(default, skip_serializing_if = "IndexMap::is_empty")]
34    pub options: IndexMap<String, String>,
35}
36
37/// The driver a machine's GPUs use when the spec names none.
38pub const DEFAULT_GPU_DRIVER: &str = "nvidia";
39
40impl Gpus {
41    /// The driver: the spec's, else [`DEFAULT_GPU_DRIVER`].
42    #[must_use]
43    pub fn driver(&self) -> &str {
44        self.driver.as_deref().unwrap_or(DEFAULT_GPU_DRIVER)
45    }
46
47    /// How many GPUs a fixed count gives: the count, the ids', else one. `None` for `all`.
48    #[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/// A number of GPUs, or every one the host has.
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61pub enum GpuCount {
62    /// That many.
63    Count(u32),
64    /// `all`: every GPU of the host.
65    All,
66}
67
68impl GpuCount {
69    /// The word for every GPU.
70    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/// `gpus` as written: a count, `all`, or the full form.
83#[derive(Deserialize, JsonSchema)]
84#[serde(untagged)]
85pub enum GpusInput {
86    /// That many GPUs.
87    Count(u32),
88    /// `all`: every GPU of the host.
89    Word(String),
90    /// The full form.
91    Full(GpusFull),
92}
93
94/// `gpus` in full, as written (`count` may be a number or `all`).
95#[derive(Deserialize, JsonSchema)]
96#[serde(deny_unknown_fields)]
97pub struct GpusFull {
98    /// How many GPUs (a number, or `all`). Default: one, unless `ids` names them.
99    #[serde(default)]
100    #[schemars(schema_with = "count_schema")]
101    count: Option<serde_json::Value>,
102    /// Which GPUs, by the host's ids (a device index or UUID; a PCI address or a resource
103    /// mapping on Proxmox). Not with `count`.
104    #[serde(default)]
105    ids: Vec<String>,
106    /// The GPUs' vendor driver (`nvidia`, `amd`, `intel`). Default: `nvidia`.
107    #[serde(default)]
108    driver: Option<String>,
109    /// What the container may use the GPUs for (`compute`, `utility`, `video`, `graphics`).
110    #[serde(default)]
111    capabilities: Vec<String>,
112    /// Options for the container runtime's GPU driver (Compose only).
113    #[serde(default)]
114    options: IndexMap<String, String>,
115}
116
117/// The schema of `gpus.count`, read as JSON to give a precise error: a number, or `all`.
118fn 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/// A device of the host given to a container, at a path inside it.
165#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
166#[serde(try_from = "DeviceInput")]
167pub struct Device {
168    /// The device's path on the host (`/dev/fuse`).
169    pub host: String,
170    /// Its path inside the machine. Default: the host's path.
171    #[serde(default, skip_serializing_if = "Option::is_none")]
172    pub guest: Option<String>,
173    /// What the machine may do with it: `r` (read), `w` (write), `m` (create the device node),
174    /// in any combination. Default: `rwm`.
175    #[serde(default, skip_serializing_if = "Option::is_none")]
176    pub permissions: Option<String>,
177}
178
179impl Device {
180    /// Its path inside the machine.
181    #[must_use]
182    pub fn guest(&self) -> &str {
183        self.guest.as_deref().unwrap_or(&self.host)
184    }
185}
186
187/// A device as written: `host[:guest[:permissions]]`, or a mapping.
188#[derive(Deserialize, JsonSchema)]
189#[serde(untagged)]
190pub enum DeviceInput {
191    /// `/dev/fuse`, `/dev/ttyUSB0:/dev/ttyS1`, `/dev/sda:/dev/xvda:r`.
192    Short(String),
193    /// The full form.
194    Full(DeviceFull),
195}
196
197/// A device in full, as written.
198#[derive(Deserialize, JsonSchema)]
199#[serde(deny_unknown_fields)]
200pub struct DeviceFull {
201    /// The device's path on the host.
202    host: String,
203    /// Its path inside the machine. Default: the host's path.
204    #[serde(default)]
205    guest: Option<String>,
206    /// `r`, `w` and `m`, in any combination. Default: `rwm`.
207    #[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}