From 62e290f9fa351b6b9af4186094a6bb6d1967f8ab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Fern=C3=A1ndez=20L=C3=B3pez?=
 <ereslibre@ereslibre.es>
Date: Fri, 25 Apr 2025 18:22:00 +0200
Subject: [PATCH] nixos/nvidia-container-toolkit: allow to provide CSV files

(cherry picked from commit 1948ad4bb5aa16054e07cbfd0b585bd58763965b)
---
 .../nvidia-container-toolkit/cdi-generate.nix | 15 +++-
 .../nvidia-container-toolkit/default.nix      | 75 +++++++++++++------
 2 files changed, 65 insertions(+), 25 deletions(-)

diff --git a/nixos/modules/services/hardware/nvidia-container-toolkit/cdi-generate.nix b/nixos/modules/services/hardware/nvidia-container-toolkit/cdi-generate.nix
index bf4e5632271a..13595c320833 100644
--- a/nixos/modules/services/hardware/nvidia-container-toolkit/cdi-generate.nix
+++ b/nixos/modules/services/hardware/nvidia-container-toolkit/cdi-generate.nix
@@ -1,9 +1,11 @@
 {
-  deviceNameStrategy,
+  csv-files,
+  device-name-strategy,
+  discovery-mode,
+  mounts,
   glibc,
   jq,
   lib,
-  mounts,
   nvidia-container-toolkit,
   nvidia-driver,
   runtimeShell,
@@ -36,7 +38,14 @@ writeScriptBin "nvidia-cdi-generator" ''
   function cdiGenerate {
     ${lib.getExe' nvidia-container-toolkit "nvidia-ctk"} cdi generate \
       --format json \
-      --device-name-strategy ${deviceNameStrategy} \
+      ${
+        if (builtins.length csv-files) > 0 then
+          lib.concatMapStringsSep "\n" (file: "--csv.file ${file} \\") csv-files
+        else
+          "\\"
+      }
+      --discovery-mode ${discovery-mode} \
+      --device-name-strategy ${device-name-strategy} \
       --ldconfig-path ${lib.getExe' glibc "ldconfig"} \
       --library-search-path ${lib.getLib nvidia-driver}/lib \
       --nvidia-cdi-hook-path ${lib.getExe' nvidia-container-toolkit.tools "nvidia-cdi-hook"}
diff --git a/nixos/modules/services/hardware/nvidia-container-toolkit/default.nix b/nixos/modules/services/hardware/nvidia-container-toolkit/default.nix
index af94fbe6a45b..d9596f754c6b 100644
--- a/nixos/modules/services/hardware/nvidia-container-toolkit/default.nix
+++ b/nixos/modules/services/hardware/nvidia-container-toolkit/default.nix
@@ -50,12 +50,39 @@
           '';
         };
 
-        suppressNvidiaDriverAssertion = lib.mkOption {
-          default = false;
-          type = lib.types.bool;
+        device-name-strategy = lib.mkOption {
+          default = "index";
+          type = lib.types.enum [
+            "index"
+            "uuid"
+            "type-index"
+          ];
           description = ''
-            Suppress the assertion for installing Nvidia driver.
-            Useful in WSL where drivers are mounted from Windows, not provided by NixOS.
+            Specify the strategy for generating device names,
+            passed to `nvidia-ctk cdi generate`. This will affect how
+            you reference the device using `nvidia.com/gpu=` in
+            the container runtime.
+          '';
+        };
+
+        discovery-mode = lib.mkOption {
+          default = "auto";
+          type = lib.types.enum [
+            "auto"
+            "csv"
+            "nvml"
+            "wsl"
+          ];
+          description = ''
+            The mode to use when discovering the available entities.
+          '';
+        };
+
+        csv-files = lib.mkOption {
+          default = [ ];
+          type = lib.types.listOf lib.types.path;
+          description = ''
+            The path to the list of CSV files to use when generating the CDI specification in CSV mode.
           '';
         };
 
@@ -74,21 +101,6 @@
           '';
         };
 
-        device-name-strategy = lib.mkOption {
-          default = "index";
-          type = lib.types.enum [
-            "index"
-            "uuid"
-            "type-index"
-          ];
-          description = ''
-            Specify the strategy for generating device names,
-            passed to `nvidia-ctk cdi generate`. This will affect how
-            you reference the device using `nvidia.com/gpu=` in
-            the container runtime.
-          '';
-        };
-
         mount-nvidia-docker-1-directories = lib.mkOption {
           default = true;
           type = lib.types.bool;
@@ -98,6 +110,15 @@
           '';
         };
 
+        suppressNvidiaDriverAssertion = lib.mkOption {
+          default = false;
+          type = lib.types.bool;
+          description = ''
+            Suppress the assertion for installing Nvidia driver.
+            Useful in WSL where drivers are mounted from Windows, not provided by NixOS.
+          '';
+        };
+
         package = lib.mkPackageOption pkgs "nvidia-container-toolkit" { };
       };
 
@@ -112,6 +133,12 @@
           || config.hardware.nvidia-container-toolkit.suppressNvidiaDriverAssertion;
         message = ''`nvidia-container-toolkit` requires nvidia drivers: set `hardware.nvidia.datacenter.enable`, add "nvidia" to `services.xserver.videoDrivers`, or set `hardware.nvidia-container-toolkit.suppressNvidiaDriverAssertion` if the driver is provided by another NixOS module (e.g. from NixOS-WSL)'';
       }
+      {
+        assertion =
+          ((builtins.length config.hardware.nvidia-container-toolkit.csv-files) > 0)
+          -> config.hardware.nvidia-container-toolkit.discovery-mode == "csv";
+        message = ''When CSV files are provided, `config.hardware.nvidia-container-toolkit.discovery-mode` has to be set to `csv`.'';
+      }
     ];
 
     virtualisation.docker = {
@@ -209,10 +236,14 @@
         ExecStart =
           let
             script = pkgs.callPackage ./cdi-generate.nix {
-              inherit (config.hardware.nvidia-container-toolkit) mounts;
+              inherit (config.hardware.nvidia-container-toolkit)
+                csv-files
+                device-name-strategy
+                discovery-mode
+                mounts
+                ;
               nvidia-container-toolkit = config.hardware.nvidia-container-toolkit.package;
               nvidia-driver = config.hardware.nvidia.package;
-              deviceNameStrategy = config.hardware.nvidia-container-toolkit.device-name-strategy;
             };
           in
           lib.getExe script;