Skip to main content

regent_sdk/state/attribute/gpu/
nvidia_container_toolkit.rs

1//! NVIDIA Container Toolkit attribute
2//!
3//! This module provides the `NvidiaContainerToolkitExpectedState` type for installing the
4//! [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
5//! from NVIDIA's repository (`https://nvidia.github.io/libnvidia-container/stable/`) and setting up
6//! what GPU containers need on the host: the `nvidia` runtime of Docker, containerd and CRI-O,
7//! and the CDI (Container Device Interface) specification used by Podman
8//! (`podman run --device nvidia.com/gpu=all ...`).
9//!
10//! The toolkit also installs `nvidia-container-cli` (`libnvidia-container-tools`), which is what
11//! the Enroot GPU hook relies on.
12//!
13//! **Compatible OS:** Linux distributions based on `apt-get`, `dnf` or `zypper` (detected from
14//! `/etc/os-release` at assessment time, like the `NvidiaDriver` attribute).
15//!
16//! **Privilege:** installing packages and configuring the engines needs root, use `!WithSudo`.
17//! Rootless Docker (`~/.config/docker/daemon.json`) is not managed.
18//!
19//! The NVIDIA driver itself is managed by the [`NvidiaDriver`](super::nvidia_driver) attribute:
20//! place it before this one.
21//!
22//! # Invalid states are unrepresentable
23//!
24//! - The toolkit configuration only exists in the `Present` state: an `Absent` toolkit has
25//!   nothing to configure.
26//! - Runtimes are a set of known engines: an engine can't be listed twice nor be misspelled.
27//! - The version is a full package version (`<major>.<minor>.<patch>-<release>`, e.g. `1.20.1-1`),
28//!   as needed to pin the four toolkit packages to the same version.
29//! - The automatic CDI specification refresh only exists from toolkit 1.18.0 on: `Cdi: Automatic`
30//!   can't be combined with an older `Version`.
31//!
32//! # How compliance is assessed
33//!
34//! 1. The NVIDIA Container Toolkit repository must be configured (`nvidia-container-toolkit.list`
35//!    and its signing key for apt, `nvidia-container-toolkit.repo` for dnf and zypper). It is never
36//!    removed by this attribute.
37//! 2. `nvidia-container-toolkit` must be installed. With `Version`, each of the four toolkit
38//!    packages (`nvidia-container-toolkit`, `nvidia-container-toolkit-base`,
39//!    `libnvidia-container-tools`, `libnvidia-container1`) must be installed in that version: the
40//!    ones in another version are upgraded or downgraded. Without it, the toolkit is installed
41//!    once and not upgraded afterwards.
42//! 3. Each engine of `Runtimes` must know the `nvidia` runtime: `runtimes.nvidia` in
43//!    `/etc/docker/daemon.json` for Docker, a runtime using `nvidia-container-runtime` in the
44//!    containerd (`/etc/containerd/config.toml`, `/etc/containerd/conf.d/`) or CRI-O
45//!    (`/etc/crio/`) configuration. If not, `nvidia-ctk runtime configure --runtime=<engine>` is
46//!    run, then the engine is restarted if it is running, which restarts its containers. Engines
47//!    left out of `Runtimes` are not managed.
48//! 4. `Cdi`, if set, manages the CDI specification of the GPUs. It embeds the paths of the
49//!    driver libraries, which contain the driver version, so it must be generated again when the
50//!    driver changes. The specification is current when it declares `nvidia.com/gpu` devices and
51//!    mentions the version of the loaded driver (`/proc/driver/nvidia/version`).
52//!    - `Automatic` (toolkit 1.18.0+): the `nvidia-cdi-refresh` systemd units generate
53//!      `/var/run/cdi/nvidia.yaml` at boot and when the driver or toolkit is installed or
54//!      upgraded. They must be enabled, and the specification current (it is refreshed if not).
55//!    - `Static`: `/etc/cdi/nvidia.yaml` is generated with `nvidia-ctk cdi generate`. The
56//!      automatic refresh is disabled and its `/var/run/cdi/nvidia.yaml` removed, as it would
57//!      take precedence over `/etc/cdi`.
58//!
59//!    The driver must be loaded to generate the specification: if it isn't, the remediation
60//!    fails and says so. MIG reconfigurations are not detected: regenerate the specification
61//!    after them.
62//!
63//! `Absent` removes the four toolkit packages, but keeps the repository, the engine
64//! configurations (which then point to a missing `nvidia-container-runtime`) and the CDI
65//! specifications.
66//!
67//! **Prerequisites:** `curl` (or `wget`) and, for apt, `gpg` to install the repository key.
68//!
69//! # Examples
70//!
71//! ## Rust API
72//!
73//! ```no_run
74//! use regent_sdk::state::attribute::gpu::nvidia_container_toolkit::{
75//!     NvidiaContainerToolkitConfig, NvidiaContainerToolkitExpectedState,
76//! };
77//! use regent_sdk::{Attribute, ExpectedState, Privilege, RegentError};
78//!
79//! fn gpu_containers() -> Result<ExpectedState, RegentError> {
80//!     let config = NvidiaContainerToolkitConfig::new()
81//!         .with_version("1.20.1-1".parse()?)?
82//!         .with_docker_runtime() // | with_containerd_runtime() | with_crio_runtime()
83//!         .with_automatic_cdi()?; // toolkit 1.18+ | with_static_cdi()
84//!
85//!     Ok(ExpectedState::new()
86//!         .with_attribute(Attribute::nvidia_container_toolkit(
87//!             NvidiaContainerToolkitExpectedState::present(config),
88//!             Privilege::WithSudo,
89//!             None,
90//!         ))
91//!         .build())
92//! }
93//! ```
94//!
95//! ## YAML API
96//!
97//! ```yaml
98//! Attributes:
99//!   - Name: NVIDIA Container Toolkit must be installed
100//!     Privilege: !WithSudo
101//!     Detail: !NvidiaContainerToolkit
102//!       State: !Present
103//!         Version: 1.20.1-1                # optional, pins the toolkit packages
104//!         Runtimes: [Docker, Containerd]   # Docker | Containerd | CriO
105//!         Cdi: Automatic                   # optional: Automatic (toolkit 1.18+) | Static
106//! ```
107//!
108//! For Podman only, the CDI specification is enough:
109//!
110//! ```yaml
111//! Attributes:
112//!   - Name: GPUs must be available to Podman
113//!     Privilege: !WithSudo
114//!     Detail: !NvidiaContainerToolkit
115//!       State: !Present
116//!         Cdi: Static
117//! ```
118//!
119//! To remove the toolkit:
120//!
121//! ```yaml
122//! Attributes:
123//!   - Name: NVIDIA Container Toolkit must be removed
124//!     Privilege: !WithSudo
125//!     Detail: !NvidiaContainerToolkit
126//!       State: Absent
127//! ```
128
129use crate::error::RegentError;
130use crate::hosts::managed_host::InternalApiCallOutcome;
131use crate::hosts::managed_host::{AssessCompliance, ReachCompliance, Timeout};
132use crate::hosts::properties::HostProperties;
133use crate::secrets::SecretProvidersPool;
134use crate::state::Check;
135use crate::state::attribute::HostHandler;
136use crate::state::attribute::Privilege;
137use crate::state::attribute::Remediation;
138use crate::state::attribute::RemediationsList;
139use crate::state::attribute::container::container_spec::quoted_command;
140use crate::state::compliance::AttributeComplianceAssessment;
141use serde::{Deserialize, Serialize};
142use serde_json::Value;
143use std::cmp::Ordering;
144use std::collections::BTreeSet;
145use std::fmt;
146use std::str::FromStr;
147use std::time::Duration;
148
149use crate::attribute::shared::{
150    PackageFamily, UnitEnablement, add_repo_file, check_linux_host, download, host_package_family,
151    installed_packages, is_repository_configured, is_success, loaded_driver, make_temp_dir, query,
152    remove_temp_dir, run_step, run_steps, unit_state,
153};
154use std::collections::BTreeMap;
155
156const ATTRIBUTE: &str = "NvidiaContainerToolkit";
157
158/// Packages of the toolkit, all pinned to the same version
159const TOOLKIT_PACKAGES: [&str; 4] = [
160    "nvidia-container-toolkit",
161    "nvidia-container-toolkit-base",
162    "libnvidia-container-tools",
163    "libnvidia-container1",
164];
165
166const REPOSITORY_URL_FRAGMENT: &str = "nvidia.github.io/libnvidia-container";
167const GPG_KEY_URL: &str = "https://nvidia.github.io/libnvidia-container/gpgkey";
168const APT_LIST_URL: &str =
169    "https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list";
170const RPM_REPO_URL: &str =
171    "https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo";
172const APT_KEYRING: &str = "/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg";
173const APT_LIST: &str = "/etc/apt/sources.list.d/nvidia-container-toolkit.list";
174
175/// CDI specification generated by the nvidia-cdi-refresh units
176const AUTOMATIC_CDI_SPEC: &str = "/var/run/cdi/nvidia.yaml";
177/// CDI specification generated by this attribute
178const STATIC_CDI_SPEC: &str = "/etc/cdi/nvidia.yaml";
179const CDI_REFRESH_PATH_UNIT: &str = "nvidia-cdi-refresh.path";
180const CDI_REFRESH_SERVICE_UNIT: &str = "nvidia-cdi-refresh.service";
181
182#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
183#[serde(deny_unknown_fields)]
184#[serde(rename_all = "PascalCase")]
185pub struct NvidiaContainerToolkitExpectedState {
186    state: NvidiaContainerToolkitState,
187}
188
189#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
190#[serde(rename_all = "PascalCase")]
191pub enum NvidiaContainerToolkitState {
192    /// Toolkit installed with this configuration
193    Present(Box<NvidiaContainerToolkitConfig>),
194    Absent,
195}
196
197impl NvidiaContainerToolkitExpectedState {
198    pub fn present(config: NvidiaContainerToolkitConfig) -> Self {
199        Self {
200            state: NvidiaContainerToolkitState::Present(Box::new(config)),
201        }
202    }
203
204    pub fn absent() -> Self {
205        Self {
206            state: NvidiaContainerToolkitState::Absent,
207        }
208    }
209}
210
211/// Optional fields left out are not managed.
212#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
213#[serde(try_from = "NvidiaContainerToolkitConfigFields")]
214#[serde(rename_all = "PascalCase")]
215pub struct NvidiaContainerToolkitConfig {
216    version: Option<ToolkitVersion>,
217    /// Container engines which must know the `nvidia` runtime
218    runtimes: BTreeSet<ContainerRuntime>,
219    cdi: Option<CdiMode>,
220}
221
222/// Fields of [`NvidiaContainerToolkitConfig`] before their consistency is checked
223#[derive(Deserialize)]
224#[serde(deny_unknown_fields)]
225#[serde(rename_all = "PascalCase")]
226struct NvidiaContainerToolkitConfigFields {
227    version: Option<ToolkitVersion>,
228    #[serde(default)]
229    runtimes: BTreeSet<ContainerRuntime>,
230    cdi: Option<CdiMode>,
231}
232
233impl TryFrom<NvidiaContainerToolkitConfigFields> for NvidiaContainerToolkitConfig {
234    type Error = RegentError;
235
236    fn try_from(fields: NvidiaContainerToolkitConfigFields) -> Result<Self, Self::Error> {
237        let config = Self {
238            version: fields.version,
239            runtimes: fields.runtimes,
240            cdi: fields.cdi,
241        };
242        config.validated()
243    }
244}
245
246impl NvidiaContainerToolkitConfig {
247    /// Latest toolkit, no runtime and no CDI specification managed
248    pub fn new() -> Self {
249        Self::default()
250    }
251
252    /// Pin the toolkit packages to this version
253    pub fn with_version(mut self, version: ToolkitVersion) -> Result<Self, RegentError> {
254        self.version = Some(version);
255        self.validated()
256    }
257
258    /// Configure the `nvidia` runtime of this engine
259    pub fn with_runtime(mut self, runtime: ContainerRuntime) -> Self {
260        self.runtimes.insert(runtime);
261        self
262    }
263
264    pub fn with_docker_runtime(self) -> Self {
265        self.with_runtime(ContainerRuntime::Docker)
266    }
267
268    pub fn with_containerd_runtime(self) -> Self {
269        self.with_runtime(ContainerRuntime::Containerd)
270    }
271
272    pub fn with_crio_runtime(self) -> Self {
273        self.with_runtime(ContainerRuntime::CriO)
274    }
275
276    pub fn with_cdi(mut self, cdi: CdiMode) -> Result<Self, RegentError> {
277        self.cdi = Some(cdi);
278        self.validated()
279    }
280
281    /// Keep the CDI specification current with the `nvidia-cdi-refresh` units: needs the
282    /// toolkit 1.18.0 or later
283    pub fn with_automatic_cdi(self) -> Result<Self, RegentError> {
284        self.with_cdi(CdiMode::Automatic)
285    }
286
287    /// Generate the CDI specification with `nvidia-ctk cdi generate`, automatic refresh
288    /// disabled (any toolkit version)
289    pub fn with_static_cdi(mut self) -> Self {
290        self.cdi = Some(CdiMode::Static);
291        self
292    }
293
294    fn validated(self) -> Result<Self, RegentError> {
295        if let (Some(CdiMode::Automatic), Some(version)) = (&self.cdi, &self.version)
296            && !version.has_automatic_cdi_refresh()
297        {
298            return Err(RegentError::IncoherentExpectedState(format!(
299                "Cdi: Automatic needs the nvidia-cdi-refresh units of the toolkit 1.18.0 or later, but Version is {}: use Cdi: Static",
300                version
301            )));
302        }
303        Ok(self)
304    }
305}
306
307/// Container engine whose `nvidia` runtime is configured with `nvidia-ctk runtime configure`
308#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
309#[serde(rename_all = "PascalCase")]
310pub enum ContainerRuntime {
311    Docker,
312    Containerd,
313    CriO,
314}
315
316impl ContainerRuntime {
317    /// Value of `nvidia-ctk runtime configure --runtime`
318    fn ctk_name(&self) -> &'static str {
319        match self {
320            ContainerRuntime::Docker => "docker",
321            ContainerRuntime::Containerd => "containerd",
322            ContainerRuntime::CriO => "crio",
323        }
324    }
325
326    /// systemd service of the engine
327    fn service(&self) -> &'static str {
328        self.ctk_name()
329    }
330
331    /// Command printing what [`has_nvidia_runtime`](Self::has_nvidia_runtime) inspects: the
332    /// Docker configuration, or the configuration lines of containerd and CRI-O mentioning
333    /// `nvidia-container-runtime`. Missing files make it fail: only the output matters.
334    fn config_command(&self) -> String {
335        match self {
336            ContainerRuntime::Docker => quoted_command(&["cat", "/etc/docker/daemon.json"]),
337            ContainerRuntime::Containerd => quoted_command(&[
338                "grep",
339                "-rhsF",
340                "nvidia-container-runtime",
341                "/etc/containerd/config.toml",
342                "/etc/containerd/conf.d",
343            ]),
344            ContainerRuntime::CriO => quoted_command(&[
345                "grep",
346                "-rhsF",
347                "nvidia-container-runtime",
348                "/etc/crio/crio.conf",
349                "/etc/crio/crio.conf.d",
350                "/etc/crio/conf.d",
351            ]),
352        }
353    }
354
355    fn has_nvidia_runtime(&self, config: &str) -> bool {
356        match self {
357            ContainerRuntime::Docker => serde_json::from_str::<Value>(config)
358                .ok()
359                .and_then(|config| config.pointer("/runtimes/nvidia").cloned())
360                .is_some(),
361            ContainerRuntime::Containerd | ContainerRuntime::CriO => config
362                .lines()
363                .map(str::trim)
364                .filter(|line| !line.starts_with('#'))
365                .any(|line| line.contains("nvidia-container-runtime")),
366        }
367    }
368}
369
370impl fmt::Display for ContainerRuntime {
371    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
372        match self {
373            ContainerRuntime::Docker => write!(f, "Docker"),
374            ContainerRuntime::Containerd => write!(f, "containerd"),
375            ContainerRuntime::CriO => write!(f, "CRI-O"),
376        }
377    }
378}
379
380/// How the CDI specification of the GPUs is generated
381#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
382#[serde(rename_all = "PascalCase")]
383pub enum CdiMode {
384    /// `nvidia-cdi-refresh` systemd units (toolkit 1.18.0+), `/var/run/cdi/nvidia.yaml`
385    Automatic,
386    /// `nvidia-ctk cdi generate`, `/etc/cdi/nvidia.yaml`, automatic refresh disabled
387    Static,
388}
389
390impl CdiMode {
391    fn spec_path(&self) -> &'static str {
392        match self {
393            CdiMode::Automatic => AUTOMATIC_CDI_SPEC,
394            CdiMode::Static => STATIC_CDI_SPEC,
395        }
396    }
397}
398
399/// Whether a CDI specification is current: it declares GPUs and mentions the loaded driver
400/// version (the driver libraries it lists are named after it)
401fn is_cdi_spec_current(spec: &str, driver_version: Option<&str>) -> bool {
402    let declares_gpus = spec.lines().any(|line| {
403        let line = line.trim();
404        line.starts_with("kind:") && line.contains("nvidia.com/gpu")
405    });
406    declares_gpus && driver_version.is_some_and(|version| spec.contains(version))
407}
408
409/// Toolkit package version: `<major>.<minor>.<patch>-<release>` (e.g. `1.20.1-1`)
410#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
411#[serde(try_from = "String", into = "String")]
412pub struct ToolkitVersion {
413    major: u32,
414    minor: u32,
415    patch: u32,
416    release: String,
417}
418
419impl TryFrom<String> for ToolkitVersion {
420    type Error = RegentError;
421
422    fn try_from(text: String) -> Result<Self, Self::Error> {
423        let invalid = || {
424            RegentError::IncoherentExpectedState(format!(
425                "Invalid NVIDIA Container Toolkit version {:?}: expected <major>.<minor>.<patch>-<release> such as 1.20.1-1",
426                text
427            ))
428        };
429        let (upstream, release) = text.split_once('-').ok_or_else(invalid)?;
430        let numbers: Vec<u32> = upstream
431            .split('.')
432            .map(|part| {
433                (!part.is_empty() && part.chars().all(|c| c.is_ascii_digit()))
434                    .then(|| part.parse().ok())
435                    .flatten()
436            })
437            .collect::<Option<Vec<u32>>>()
438            .ok_or_else(invalid)?;
439        let is_valid_release = release.starts_with(|c: char| c.is_ascii_digit())
440            && release
441                .chars()
442                .all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '~');
443        match numbers[..] {
444            [major, minor, patch] if is_valid_release => Ok(Self {
445                major,
446                minor,
447                patch,
448                release: release.to_string(),
449            }),
450            _ => Err(invalid()),
451        }
452    }
453}
454
455impl FromStr for ToolkitVersion {
456    type Err = RegentError;
457
458    fn from_str(text: &str) -> Result<Self, Self::Err> {
459        Self::try_from(text.to_string())
460    }
461}
462
463impl From<ToolkitVersion> for String {
464    fn from(version: ToolkitVersion) -> String {
465        version.to_string()
466    }
467}
468
469impl fmt::Display for ToolkitVersion {
470    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
471        write!(
472            f,
473            "{}.{}.{}-{}",
474            self.major, self.minor, self.patch, self.release
475        )
476    }
477}
478
479impl ToolkitVersion {
480    fn has_automatic_cdi_refresh(&self) -> bool {
481        (self.major, self.minor) >= (1, 18)
482    }
483
484    /// Order with an installed package version (`1.20.1-1`, possibly with an epoch)
485    fn compare_with_installed(&self, installed: &str) -> Option<Ordering> {
486        let installed: ToolkitVersion = installed
487            .split_once(':')
488            .map_or(installed, |(_, version)| version)
489            .parse()
490            .ok()?;
491        let release_key = |release: &str| {
492            release
493                .split(|c: char| !c.is_ascii_digit())
494                .next()
495                .and_then(|number| number.parse::<u64>().ok())
496                .unwrap_or_default()
497        };
498        Some(
499            (
500                self.major,
501                self.minor,
502                self.patch,
503                release_key(&self.release),
504            )
505                .cmp(&(
506                    installed.major,
507                    installed.minor,
508                    installed.patch,
509                    release_key(&installed.release),
510                )),
511        )
512    }
513
514    /// Package specification pinning `package` to this version
515    fn package_spec(&self, package: &str, family: PackageFamily) -> String {
516        match family {
517            PackageFamily::Apt => format!("{}={}", package, self),
518            PackageFamily::Dnf | PackageFamily::Zypper => format!("{}-{}", package, self),
519        }
520    }
521
522    /// Toolkit packages to install or upgrade, and to downgrade, to reach this version
523    fn packages_to_change(
524        &self,
525        installed: &BTreeMap<String, String>,
526    ) -> (Vec<String>, Vec<String>) {
527        let mut upgrades = Vec::new();
528        let mut downgrades = Vec::new();
529        for package in TOOLKIT_PACKAGES {
530            let ordering = installed
531                .get(package)
532                .map(|installed| self.compare_with_installed(installed));
533            match ordering {
534                Some(Some(Ordering::Equal)) => {}
535                Some(Some(Ordering::Less)) => downgrades.push(package.to_string()),
536                _ => upgrades.push(package.to_string()),
537            }
538        }
539        (upgrades, downgrades)
540    }
541}
542
543async fn setup_repository<Handler: HostHandler>(
544    host_handler: &mut Handler,
545    family: PackageFamily,
546    privilege: &Privilege,
547) -> Result<InternalApiCallOutcome, RegentError> {
548    if family != PackageFamily::Apt {
549        return add_repo_file(
550            host_handler,
551            family,
552            RPM_REPO_URL,
553            "nvidia-container-toolkit.repo",
554            privilege,
555        )
556        .await;
557    }
558    let tmp = match make_temp_dir(host_handler, privilege).await? {
559        Ok(tmp) => tmp,
560        Err(failure) => return Ok(failure),
561    };
562    let key = format!("{}/gpgkey", tmp);
563    let list = format!("{}/list", tmp);
564    let mut outcome = download(host_handler, GPG_KEY_URL, &key, privilege).await;
565    if is_success(&outcome) {
566        outcome = download(host_handler, APT_LIST_URL, &list, privilege).await;
567    }
568    if is_success(&outcome) {
569        outcome = run_steps(
570            host_handler,
571            &apt_repository_commands(&key, &list),
572            privilege,
573        )
574        .await;
575    }
576    remove_temp_dir(host_handler, &tmp, privilege).await;
577    Ok(outcome)
578}
579
580/// Commands installing the downloaded key and list of the apt repository
581fn apt_repository_commands(key: &str, list: &str) -> Vec<String> {
582    vec![
583        quoted_command(&[
584            "gpg",
585            "--batch",
586            "--yes",
587            "--dearmor",
588            "-o",
589            APT_KEYRING,
590            key,
591        ]),
592        quoted_command(&[
593            "sed",
594            "-i",
595            &format!("s#deb https://#deb [signed-by={}] https://#g", APT_KEYRING),
596            list,
597        ]),
598        quoted_command(&["install", "-m", "0644", list, APT_LIST]),
599        quoted_command(&["apt-get", "update"]),
600    ]
601}
602
603/// Failure message when a CDI specification is requested without a loaded driver
604fn driver_not_loaded_message() -> String {
605    "The NVIDIA driver is not loaded: the CDI specification can't be generated (install and load it first, see the NvidiaDriver attribute)".to_string()
606}
607
608impl Timeout for NvidiaContainerToolkitExpectedState {
609    fn default_timeout(&self) -> Duration {
610        Duration::from_secs(600)
611    }
612}
613
614impl Check for NvidiaContainerToolkitExpectedState {
615    fn check(&self) -> Result<(), RegentError> {
616        // Every value is validated when its type is built: an existing
617        // NvidiaContainerToolkitExpectedState is valid
618        Ok(())
619    }
620
621    fn check_host_compatibility(
622        &self,
623        host_properties: &HostProperties,
624    ) -> Result<(), RegentError> {
625        check_linux_host(host_properties, ATTRIBUTE)
626    }
627}
628
629impl<Handler: HostHandler> AssessCompliance<Handler> for NvidiaContainerToolkitExpectedState {
630    async fn assess_compliance(
631        &self,
632        host_handler: &mut Handler,
633        host_properties: &Option<HostProperties>,
634        privilege: &Privilege,
635        _optional_secret_provider: &Option<SecretProvidersPool>,
636    ) -> Result<AttributeComplianceAssessment, RegentError> {
637        if let Some(props) = host_properties {
638            self.check_host_compatibility(props)?;
639        }
640
641        let Some(family) = host_package_family(host_handler, privilege).await? else {
642            return Ok(AttributeComplianceAssessment::NonCompliantFatal(
643                "the distribution (/etc/os-release) is not based on apt-get, dnf or zypper: the NVIDIA Container Toolkit can't be installed".to_string(),
644            ));
645        };
646
647        let mut remediations: Vec<Remediation> = Vec::new();
648        let mut push = |api_call: NvidiaContainerToolkitInternalApiCall| {
649            remediations.push(Remediation::NvidiaContainerToolkit(
650                NvidiaContainerToolkitApiCall::from(api_call, privilege.clone()),
651            ));
652        };
653
654        let installed =
655            installed_packages(host_handler, family, &TOOLKIT_PACKAGES, privilege).await?;
656
657        match &self.state {
658            NvidiaContainerToolkitState::Absent => {
659                if !installed.is_empty() {
660                    push(NvidiaContainerToolkitInternalApiCall::RemovePackages {
661                        family,
662                        packages: installed.into_keys().collect(),
663                    });
664                }
665            }
666            NvidiaContainerToolkitState::Present(config) => {
667                if !is_repository_configured(
668                    host_handler,
669                    family,
670                    REPOSITORY_URL_FRAGMENT,
671                    privilege,
672                )
673                .await?
674                {
675                    push(NvidiaContainerToolkitInternalApiCall::AddRepository(family));
676                }
677
678                let installed_version = installed.get(TOOLKIT_PACKAGES[0]);
679                let (upgrades, downgrades) = match &config.version {
680                    Some(version) => version.packages_to_change(&installed),
681                    None if installed_version.is_none() => {
682                        (vec![TOOLKIT_PACKAGES[0].to_string()], Vec::new())
683                    }
684                    None => (Vec::new(), Vec::new()),
685                };
686                let is_installing = !upgrades.is_empty() || !downgrades.is_empty();
687                if is_installing {
688                    push(NvidiaContainerToolkitInternalApiCall::Install {
689                        family,
690                        version: config.version.clone(),
691                        upgrades,
692                        downgrades,
693                    });
694                }
695
696                for runtime in &config.runtimes {
697                    let result = query(host_handler, &runtime.config_command(), privilege).await?;
698                    if !runtime.has_nvidia_runtime(&result.stdout) {
699                        push(NvidiaContainerToolkitInternalApiCall::ConfigureRuntime(
700                            *runtime,
701                        ));
702                    }
703                }
704
705                if let Some(cdi) = config.cdi {
706                    let driver_version = loaded_driver(host_handler, privilege)
707                        .await?
708                        .map(|driver| driver.version);
709                    let spec = query(
710                        host_handler,
711                        &quoted_command(&["cat", cdi.spec_path()]),
712                        privilege,
713                    )
714                    .await?;
715                    let is_spec_current = spec.return_code == 0
716                        && is_cdi_spec_current(&spec.stdout, driver_version.as_deref());
717                    let refresh =
718                        unit_state(host_handler, CDI_REFRESH_PATH_UNIT, privilege).await?;
719
720                    match cdi {
721                        CdiMode::Automatic => {
722                            if !is_installing && !refresh.exists() {
723                                return Ok(AttributeComplianceAssessment::NonCompliantFatal(
724                                    format!(
725                                        "the installed NVIDIA Container Toolkit ({}) has no {} unit (toolkit 1.18.0 or later): use Cdi: Static or a newer Version",
726                                        installed_version.cloned().unwrap_or_default(),
727                                        CDI_REFRESH_PATH_UNIT
728                                    ),
729                                ));
730                            }
731                            let is_refresh_enabled = matches!(
732                                refresh.enablement,
733                                UnitEnablement::Enabled | UnitEnablement::Static
734                            );
735                            if !(is_refresh_enabled && refresh.active && is_spec_current) {
736                                push(NvidiaContainerToolkitInternalApiCall::RefreshCdiSpec);
737                            }
738                        }
739                        CdiMode::Static => {
740                            let automatic_spec = query(
741                                host_handler,
742                                &quoted_command(&["test", "-e", AUTOMATIC_CDI_SPEC]),
743                                privilege,
744                            )
745                            .await?;
746                            if refresh.enablement == UnitEnablement::Enabled
747                                || refresh.active
748                                || automatic_spec.return_code == 0
749                            {
750                                push(NvidiaContainerToolkitInternalApiCall::DisableCdiRefresh);
751                            }
752                            if !is_spec_current {
753                                push(NvidiaContainerToolkitInternalApiCall::GenerateCdiSpec);
754                            }
755                        }
756                    }
757                }
758            }
759        }
760
761        if remediations.is_empty() {
762            Ok(AttributeComplianceAssessment::Compliant)
763        } else {
764            Ok(AttributeComplianceAssessment::NonCompliant(
765                RemediationsList::from(remediations)?,
766            ))
767        }
768    }
769}
770
771#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
772#[serde(rename_all = "PascalCase", rename_all_fields = "PascalCase")]
773pub enum NvidiaContainerToolkitInternalApiCall {
774    /// Configure the NVIDIA Container Toolkit repository
775    AddRepository(PackageFamily),
776    /// Install the toolkit: the latest one, or its packages pinned to `version` (upgraded or
777    /// installed, and downgraded)
778    Install {
779        family: PackageFamily,
780        version: Option<ToolkitVersion>,
781        upgrades: Vec<String>,
782        downgrades: Vec<String>,
783    },
784    RemovePackages {
785        family: PackageFamily,
786        packages: Vec<String>,
787    },
788    /// `nvidia-ctk runtime configure`, then restart the engine if it is running
789    ConfigureRuntime(ContainerRuntime),
790    /// Enable the nvidia-cdi-refresh units and refresh `/var/run/cdi/nvidia.yaml`
791    RefreshCdiSpec,
792    /// Disable the nvidia-cdi-refresh units and remove `/var/run/cdi/nvidia.yaml`
793    DisableCdiRefresh,
794    /// Generate `/etc/cdi/nvidia.yaml`
795    GenerateCdiSpec,
796}
797
798#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
799pub struct NvidiaContainerToolkitApiCall {
800    pub api_call: NvidiaContainerToolkitInternalApiCall,
801    privilege: Privilege,
802}
803
804impl NvidiaContainerToolkitApiCall {
805    pub fn display(&self) -> String {
806        match &self.api_call {
807            NvidiaContainerToolkitInternalApiCall::AddRepository(family) => {
808                format!("Add the NVIDIA Container Toolkit repository ({})", family)
809            }
810            NvidiaContainerToolkitInternalApiCall::Install {
811                version: Some(version),
812                upgrades,
813                downgrades,
814                ..
815            } => format!(
816                "Install the NVIDIA Container Toolkit {}: {}",
817                version,
818                upgrades
819                    .iter()
820                    .map(|package| format!("install or upgrade {}", package))
821                    .chain(
822                        downgrades
823                            .iter()
824                            .map(|package| format!("downgrade {}", package))
825                    )
826                    .collect::<Vec<_>>()
827                    .join(", ")
828            ),
829            NvidiaContainerToolkitInternalApiCall::Install { version: None, .. } => {
830                "Install the NVIDIA Container Toolkit".to_string()
831            }
832            NvidiaContainerToolkitInternalApiCall::RemovePackages { packages, .. } => {
833                format!("Remove {}", packages.join(", "))
834            }
835            NvidiaContainerToolkitInternalApiCall::ConfigureRuntime(runtime) => format!(
836                "Configure the nvidia runtime of {} and restart it if running",
837                runtime
838            ),
839            NvidiaContainerToolkitInternalApiCall::RefreshCdiSpec => format!(
840                "Enable the automatic CDI specification refresh and regenerate {}",
841                AUTOMATIC_CDI_SPEC
842            ),
843            NvidiaContainerToolkitInternalApiCall::DisableCdiRefresh => format!(
844                "Disable the automatic CDI specification refresh and remove {}",
845                AUTOMATIC_CDI_SPEC
846            ),
847            NvidiaContainerToolkitInternalApiCall::GenerateCdiSpec => {
848                format!("Generate the CDI specification {}", STATIC_CDI_SPEC)
849            }
850        }
851    }
852
853    fn from(
854        api_call: NvidiaContainerToolkitInternalApiCall,
855        privilege: Privilege,
856    ) -> NvidiaContainerToolkitApiCall {
857        NvidiaContainerToolkitApiCall {
858            api_call,
859            privilege,
860        }
861    }
862
863    /// Commands of the remediations which don't need to inspect the host
864    fn commands(&self) -> Vec<String> {
865        match &self.api_call {
866            NvidiaContainerToolkitInternalApiCall::Install {
867                family,
868                version,
869                upgrades,
870                downgrades,
871            } => {
872                let specs = |packages: &[String]| -> Vec<String> {
873                    packages
874                        .iter()
875                        .map(|package| match version {
876                            Some(version) => version.package_spec(package, *family),
877                            None => package.clone(),
878                        })
879                        .collect()
880                };
881                let mut commands = Vec::new();
882                if *family == PackageFamily::Apt {
883                    commands.push(quoted_command(&["apt-get", "update"]));
884                }
885                if !upgrades.is_empty() {
886                    commands.push(family.install_command(&specs(upgrades), false));
887                }
888                if !downgrades.is_empty() {
889                    commands.push(family.install_command(&specs(downgrades), true));
890                }
891                commands
892            }
893            NvidiaContainerToolkitInternalApiCall::RemovePackages { family, packages } => {
894                vec![family.remove_command(packages)]
895            }
896            NvidiaContainerToolkitInternalApiCall::RefreshCdiSpec => vec![
897                quoted_command(&["systemctl", "enable", "--now", CDI_REFRESH_PATH_UNIT]),
898                quoted_command(&["systemctl", "enable", CDI_REFRESH_SERVICE_UNIT]),
899                quoted_command(&["systemctl", "restart", CDI_REFRESH_SERVICE_UNIT]),
900            ],
901            NvidiaContainerToolkitInternalApiCall::GenerateCdiSpec => vec![
902                quoted_command(&["mkdir", "-p", "/etc/cdi"]),
903                quoted_command(&[
904                    "nvidia-ctk",
905                    "cdi",
906                    "generate",
907                    &format!("--output={}", STATIC_CDI_SPEC),
908                ]),
909            ],
910            NvidiaContainerToolkitInternalApiCall::AddRepository(_)
911            | NvidiaContainerToolkitInternalApiCall::ConfigureRuntime(_)
912            | NvidiaContainerToolkitInternalApiCall::DisableCdiRefresh => Vec::new(),
913        }
914    }
915}
916
917impl Timeout for NvidiaContainerToolkitApiCall {
918    fn default_timeout(&self) -> Duration {
919        match &self.api_call {
920            NvidiaContainerToolkitInternalApiCall::Install { .. }
921            | NvidiaContainerToolkitInternalApiCall::AddRepository(_) => Duration::from_secs(600),
922            _ => Duration::from_secs(300),
923        }
924    }
925}
926
927impl Check for NvidiaContainerToolkitApiCall {
928    fn check(&self) -> Result<(), RegentError> {
929        Ok(())
930    }
931
932    fn check_host_compatibility(
933        &self,
934        host_properties: &HostProperties,
935    ) -> Result<(), RegentError> {
936        check_linux_host(host_properties, ATTRIBUTE)
937    }
938}
939
940impl<Handler: HostHandler> ReachCompliance<Handler> for NvidiaContainerToolkitApiCall {
941    async fn call(
942        &self,
943        host_handler: &mut Handler,
944        host_properties: &Option<HostProperties>,
945        _optional_secret_provider: &Option<SecretProvidersPool>,
946    ) -> Result<InternalApiCallOutcome, RegentError> {
947        if let Some(props) = host_properties {
948            self.check_host_compatibility(props)?;
949        }
950
951        let privilege = &self.privilege;
952        match &self.api_call {
953            NvidiaContainerToolkitInternalApiCall::AddRepository(family) => {
954                return setup_repository(host_handler, *family, privilege).await;
955            }
956            NvidiaContainerToolkitInternalApiCall::ConfigureRuntime(runtime) => {
957                return configure_runtime(host_handler, *runtime, privilege).await;
958            }
959            NvidiaContainerToolkitInternalApiCall::DisableCdiRefresh => {
960                return disable_cdi_refresh(host_handler, privilege).await;
961            }
962            _ => {}
963        }
964        let cdi_mode = match &self.api_call {
965            NvidiaContainerToolkitInternalApiCall::RefreshCdiSpec => Some(CdiMode::Automatic),
966            NvidiaContainerToolkitInternalApiCall::GenerateCdiSpec => Some(CdiMode::Static),
967            _ => None,
968        };
969
970        let Some(cdi_mode) = cdi_mode else {
971            return Ok(run_steps(host_handler, &self.commands(), privilege).await);
972        };
973
974        // A CDI specification generated without driver would be empty
975        let Some(driver) = loaded_driver(host_handler, privilege).await? else {
976            return Ok(InternalApiCallOutcome::Failure(driver_not_loaded_message()));
977        };
978        let outcome = run_steps(host_handler, &self.commands(), privilege).await;
979        if let InternalApiCallOutcome::Failure(_) = outcome {
980            return Ok(outcome);
981        }
982        let spec = query(
983            host_handler,
984            &quoted_command(&["cat", cdi_mode.spec_path()]),
985            privilege,
986        )
987        .await?;
988        if spec.return_code == 0 && is_cdi_spec_current(&spec.stdout, Some(&driver.version)) {
989            Ok(InternalApiCallOutcome::Success(Some(format!(
990                "CDI specification {} generated for driver {}",
991                cdi_mode.spec_path(),
992                driver.version
993            ))))
994        } else {
995            Ok(InternalApiCallOutcome::Failure(format!(
996                "{} doesn't declare the GPUs of driver {} after its generation",
997                cdi_mode.spec_path(),
998                driver.version
999            )))
1000        }
1001    }
1002}
1003
1004/// `nvidia-ctk runtime configure`, then restart the engine if it is running (a missing engine
1005/// keeps the configuration for when it is installed)
1006async fn configure_runtime<Handler: HostHandler>(
1007    host_handler: &mut Handler,
1008    runtime: ContainerRuntime,
1009    privilege: &Privilege,
1010) -> Result<InternalApiCallOutcome, RegentError> {
1011    let outcome = run_step(
1012        host_handler,
1013        &quoted_command(&[
1014            "nvidia-ctk",
1015            "runtime",
1016            "configure",
1017            &format!("--runtime={}", runtime.ctk_name()),
1018        ]),
1019        privilege,
1020    )
1021    .await;
1022    if !is_success(&outcome) {
1023        return Ok(outcome);
1024    }
1025    if !unit_state(host_handler, runtime.service(), privilege)
1026        .await?
1027        .active
1028    {
1029        return Ok(outcome);
1030    }
1031    Ok(run_step(
1032        host_handler,
1033        &quoted_command(&["systemctl", "restart", runtime.service()]),
1034        privilege,
1035    )
1036    .await)
1037}
1038
1039/// Disable and stop the nvidia-cdi-refresh units, then remove the specification they generated
1040async fn disable_cdi_refresh<Handler: HostHandler>(
1041    host_handler: &mut Handler,
1042    privilege: &Privilege,
1043) -> Result<InternalApiCallOutcome, RegentError> {
1044    let mut commands = Vec::new();
1045    for unit in [CDI_REFRESH_PATH_UNIT, CDI_REFRESH_SERVICE_UNIT] {
1046        let state = unit_state(host_handler, unit, privilege).await?;
1047        if state.enablement == UnitEnablement::Enabled {
1048            commands.push(quoted_command(&["systemctl", "disable", unit]));
1049        }
1050        if state.active {
1051            commands.push(quoted_command(&["systemctl", "stop", unit]));
1052        }
1053    }
1054    commands.push(quoted_command(&["rm", "-f", AUTOMATIC_CDI_SPEC]));
1055    Ok(run_steps(host_handler, &commands, privilege).await)
1056}
1057
1058#[cfg(test)]
1059mod tests {
1060    use super::*;
1061
1062    fn full_yaml() -> &'static str {
1063        "---
1064State: !Present
1065  Version: 1.20.1-1
1066  Runtimes: [Docker, Containerd, CriO]
1067  Cdi: Automatic
1068"
1069    }
1070
1071    fn full_config() -> NvidiaContainerToolkitConfig {
1072        let state: NvidiaContainerToolkitExpectedState = yaml_serde::from_str(full_yaml()).unwrap();
1073        match state.state {
1074            NvidiaContainerToolkitState::Present(config) => *config,
1075            other => panic!("unexpected state {:?}", other),
1076        }
1077    }
1078
1079    fn rejection(raw: &str) -> String {
1080        match yaml_serde::from_str::<NvidiaContainerToolkitExpectedState>(raw) {
1081            Ok(state) => panic!("should be rejected: {}\nparsed as {:?}", raw, state),
1082            Err(error) => error.to_string(),
1083        }
1084    }
1085
1086    #[test]
1087    fn parsing_nvidia_container_toolkit_module_block_from_yaml_str() {
1088        let config = full_config();
1089        assert_eq!(config.version.as_ref().unwrap().to_string(), "1.20.1-1");
1090        assert_eq!(config.runtimes.len(), 3);
1091        assert_eq!(config.cdi, Some(CdiMode::Automatic));
1092
1093        let raw = "---
1094- State: Absent
1095- State: !Present {}
1096- State: !Present
1097    Cdi: Static
1098- State: !Present
1099    Version: 1.17.8-1
1100    Runtimes: [Docker]
1101";
1102        let states: Vec<NvidiaContainerToolkitExpectedState> = yaml_serde::from_str(raw).unwrap();
1103        assert_eq!(states[0], NvidiaContainerToolkitExpectedState::absent());
1104        assert_eq!(
1105            states[1],
1106            NvidiaContainerToolkitExpectedState::present(NvidiaContainerToolkitConfig::new())
1107        );
1108        assert_eq!(
1109            states[2],
1110            NvidiaContainerToolkitExpectedState::present(
1111                NvidiaContainerToolkitConfig::new()
1112                    .with_cdi(CdiMode::Static)
1113                    .unwrap()
1114            )
1115        );
1116        assert_eq!(
1117            states[3],
1118            NvidiaContainerToolkitExpectedState::present(
1119                NvidiaContainerToolkitConfig::new()
1120                    .with_version("1.17.8-1".parse().unwrap())
1121                    .unwrap()
1122                    .with_runtime(ContainerRuntime::Docker)
1123            )
1124        );
1125    }
1126
1127    #[test]
1128    fn serialization_round_trip_preserves_state() {
1129        // Templating (Attribute::consider_context) relies on a JSON round trip
1130        for raw in [full_yaml(), "State: Absent\n", "State: !Present {}\n"] {
1131            let state: NvidiaContainerToolkitExpectedState = yaml_serde::from_str(raw).unwrap();
1132            let json = serde_json::to_string(&state).unwrap();
1133            assert_eq!(
1134                serde_json::from_str::<NvidiaContainerToolkitExpectedState>(&json).unwrap(),
1135                state,
1136                "{}",
1137                json
1138            );
1139        }
1140    }
1141
1142    #[test]
1143    fn invalid_states_are_rejected_while_parsing() {
1144        let rejected = [
1145            // Configuration makes no sense for an absent toolkit
1146            ("State: !Absent\n  Cdi: Static\n", "expected unit"),
1147            // Automatic CDI refresh appeared in 1.18.0
1148            (
1149                "State: !Present\n  Version: 1.17.8-1\n  Cdi: Automatic\n",
1150                "1.18.0 or later",
1151            ),
1152            // The release is needed to pin the four packages
1153            (
1154                "State: !Present\n  Version: 1.20.1\n",
1155                "<major>.<minor>.<patch>-<release>",
1156            ),
1157            (
1158                "State: !Present\n  Version: latest\n",
1159                "<major>.<minor>.<patch>-<release>",
1160            ),
1161            (
1162                "State: !Present\n  Version: 1.20-1\n",
1163                "<major>.<minor>.<patch>-<release>",
1164            ),
1165            (
1166                "State: !Present\n  Version: \"1.20.1-1; rm -rf /\"\n",
1167                "<major>.<minor>.<patch>-<release>",
1168            ),
1169            // Unknown engines and duplicated settings
1170            ("State: !Present\n  Runtimes: [Podman]\n", "unknown variant"),
1171            ("State: !Present\n  Cdi: Dynamic\n", "unknown variant"),
1172            ("State: !Present\n  Runtime: [Docker]\n", "unknown field"),
1173        ];
1174        for (raw, reason) in rejected {
1175            let error = rejection(raw);
1176            assert!(
1177                error.contains(reason),
1178                "{:?} rejected for another reason: {}",
1179                raw,
1180                error
1181            );
1182        }
1183
1184        // Same rules through the Rust API
1185        assert!(
1186            NvidiaContainerToolkitConfig::new()
1187                .with_cdi(CdiMode::Automatic)
1188                .unwrap()
1189                .with_version("1.17.8-1".parse().unwrap())
1190                .is_err()
1191        );
1192    }
1193
1194    #[test]
1195    fn runtime_and_cdi_shortcuts_match_the_explicit_form() {
1196        let new = NvidiaContainerToolkitConfig::new;
1197        assert_eq!(
1198            new()
1199                .with_docker_runtime()
1200                .with_containerd_runtime()
1201                .with_crio_runtime(),
1202            new()
1203                .with_runtime(ContainerRuntime::Docker)
1204                .with_runtime(ContainerRuntime::Containerd)
1205                .with_runtime(ContainerRuntime::CriO)
1206        );
1207        assert_eq!(
1208            new().with_automatic_cdi().unwrap(),
1209            new().with_cdi(CdiMode::Automatic).unwrap()
1210        );
1211        assert_eq!(
1212            new().with_static_cdi(),
1213            new().with_cdi(CdiMode::Static).unwrap()
1214        );
1215
1216        // The automatic refresh still needs the toolkit 1.18.0 or later
1217        let old_version = || new().with_version("1.17.8-1".parse().unwrap()).unwrap();
1218        assert!(old_version().with_automatic_cdi().is_err());
1219        assert_eq!(
1220            old_version().with_static_cdi(),
1221            old_version().with_cdi(CdiMode::Static).unwrap()
1222        );
1223    }
1224
1225    #[test]
1226    fn versions_are_compared_with_installed_packages() {
1227        let version: ToolkitVersion = "1.20.1-1".parse().unwrap();
1228        assert_eq!(
1229            version.compare_with_installed("1.20.1-1"),
1230            Some(Ordering::Equal)
1231        );
1232        assert_eq!(
1233            version.compare_with_installed("1.21.0-1"),
1234            Some(Ordering::Less)
1235        );
1236        assert_eq!(
1237            version.compare_with_installed("1.9.0-1"),
1238            Some(Ordering::Greater)
1239        );
1240        assert_eq!(
1241            version.compare_with_installed("0:1.20.1-1"),
1242            Some(Ordering::Equal)
1243        );
1244        assert_eq!(version.compare_with_installed("garbage"), None);
1245        assert_eq!(
1246            version.package_spec("nvidia-container-toolkit", PackageFamily::Apt),
1247            "nvidia-container-toolkit=1.20.1-1"
1248        );
1249        assert_eq!(
1250            version.package_spec("libnvidia-container1", PackageFamily::Dnf),
1251            "libnvidia-container1-1.20.1-1"
1252        );
1253
1254        let installed: BTreeMap<String, String> = [
1255            ("nvidia-container-toolkit", "1.20.1-1"),
1256            ("nvidia-container-toolkit-base", "1.21.0-1"),
1257            ("libnvidia-container1", "1.17.8-1"),
1258        ]
1259        .into_iter()
1260        .map(|(name, version)| (name.to_string(), version.to_string()))
1261        .collect();
1262        assert_eq!(
1263            version.packages_to_change(&installed),
1264            (
1265                vec![
1266                    "libnvidia-container-tools".to_string(),
1267                    "libnvidia-container1".to_string()
1268                ],
1269                vec!["nvidia-container-toolkit-base".to_string()]
1270            )
1271        );
1272    }
1273
1274    #[test]
1275    fn engine_configurations_are_inspected() {
1276        let docker =
1277            r#"{"runtimes": {"nvidia": {"args": [], "path": "nvidia-container-runtime"}}}"#;
1278        assert!(ContainerRuntime::Docker.has_nvidia_runtime(docker));
1279        assert!(!ContainerRuntime::Docker.has_nvidia_runtime(r#"{"log-driver": "journald"}"#));
1280        assert!(!ContainerRuntime::Docker.has_nvidia_runtime(""));
1281
1282        let containerd = "version = 3\n[plugins.\"io.containerd.cri.v1.runtime\".containerd.runtimes.nvidia.options]\n  BinaryName = \"/usr/bin/nvidia-container-runtime\"\n";
1283        assert!(ContainerRuntime::Containerd.has_nvidia_runtime(containerd));
1284        assert!(
1285            !ContainerRuntime::Containerd
1286                .has_nvidia_runtime("# BinaryName = \"nvidia-container-runtime\"\n")
1287        );
1288
1289        let crio = "[crio.runtime.runtimes.nvidia]\nruntime_path = \"/usr/bin/nvidia-container-runtime\"\n";
1290        assert!(ContainerRuntime::CriO.has_nvidia_runtime(crio));
1291    }
1292
1293    #[test]
1294    fn cdi_specifications_must_match_the_loaded_driver() {
1295        let spec = "---\ncdiVersion: 0.5.0\nkind: nvidia.com/gpu\ndevices:\n- name: \"0\"\ncontainerEdits:\n  mounts:\n  - hostPath: /usr/lib/x86_64-linux-gnu/libcuda.so.580.105.08\n";
1296        assert!(is_cdi_spec_current(spec, Some("580.105.08")));
1297        assert!(!is_cdi_spec_current(spec, Some("590.48.01")));
1298        assert!(!is_cdi_spec_current(spec, None));
1299        assert!(!is_cdi_spec_current(
1300            "cdiVersion: 0.5.0\n",
1301            Some("580.105.08")
1302        ));
1303    }
1304
1305    #[test]
1306    fn remediation_commands_are_quoted() {
1307        let api_call =
1308            |api_call| NvidiaContainerToolkitApiCall::from(api_call, Privilege::WithSudo);
1309        assert_eq!(
1310            api_call(NvidiaContainerToolkitInternalApiCall::Install {
1311                family: PackageFamily::Zypper,
1312                version: Some("1.20.1-1".parse().unwrap()),
1313                upgrades: vec!["nvidia-container-toolkit".to_string()],
1314                downgrades: vec!["libnvidia-container1".to_string()],
1315            })
1316            .commands(),
1317            vec![
1318                "'zypper' '--non-interactive' '--gpg-auto-import-keys' 'install' '--force-resolution' 'nvidia-container-toolkit-1.20.1-1'",
1319                "'zypper' '--non-interactive' '--gpg-auto-import-keys' 'install' '--force-resolution' '--oldpackage' 'libnvidia-container1-1.20.1-1'"
1320            ]
1321        );
1322        assert_eq!(
1323            api_call(NvidiaContainerToolkitInternalApiCall::Install {
1324                family: PackageFamily::Apt,
1325                version: None,
1326                upgrades: vec!["nvidia-container-toolkit".to_string()],
1327                downgrades: Vec::new(),
1328            })
1329            .commands(),
1330            vec![
1331                "'apt-get' 'update'",
1332                "DEBIAN_FRONTEND=noninteractive 'apt-get' 'install' '-y' 'nvidia-container-toolkit'"
1333            ]
1334        );
1335        assert_eq!(
1336            apt_repository_commands("/tmp/x/gpgkey", "/tmp/x/list"),
1337            vec![
1338                "'gpg' '--batch' '--yes' '--dearmor' '-o' '/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg' '/tmp/x/gpgkey'",
1339                "'sed' '-i' 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' '/tmp/x/list'",
1340                "'install' '-m' '0644' '/tmp/x/list' '/etc/apt/sources.list.d/nvidia-container-toolkit.list'",
1341                "'apt-get' 'update'"
1342            ]
1343        );
1344        assert!(
1345            api_call(NvidiaContainerToolkitInternalApiCall::RefreshCdiSpec)
1346                .display()
1347                .contains("/var/run/cdi/nvidia.yaml")
1348        );
1349    }
1350}