1use 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
158const 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
175const AUTOMATIC_CDI_SPEC: &str = "/var/run/cdi/nvidia.yaml";
177const 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 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#[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 runtimes: BTreeSet<ContainerRuntime>,
219 cdi: Option<CdiMode>,
220}
221
222#[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 pub fn new() -> Self {
249 Self::default()
250 }
251
252 pub fn with_version(mut self, version: ToolkitVersion) -> Result<Self, RegentError> {
254 self.version = Some(version);
255 self.validated()
256 }
257
258 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 pub fn with_automatic_cdi(self) -> Result<Self, RegentError> {
284 self.with_cdi(CdiMode::Automatic)
285 }
286
287 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#[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 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 fn service(&self) -> &'static str {
328 self.ctk_name()
329 }
330
331 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
382#[serde(rename_all = "PascalCase")]
383pub enum CdiMode {
384 Automatic,
386 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
399fn 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#[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 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 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 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
580fn 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
603fn 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 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 "ed_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 "ed_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 AddRepository(PackageFamily),
776 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 ConfigureRuntime(ContainerRuntime),
790 RefreshCdiSpec,
792 DisableCdiRefresh,
794 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 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 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 "ed_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
1004async 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 "ed_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 "ed_command(&["systemctl", "restart", runtime.service()]),
1034 privilege,
1035 )
1036 .await)
1037}
1038
1039async 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 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 ("State: !Absent\n Cdi: Static\n", "expected unit"),
1147 (
1149 "State: !Present\n Version: 1.17.8-1\n Cdi: Automatic\n",
1150 "1.18.0 or later",
1151 ),
1152 (
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 ("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 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 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}