services/hardware: add nvidia-container-toolkit

2 files changed, 77 insertions, 0 deletions

diff --git a/nixos/modules/services/hardware/nvidia-container-toolkit-cdi-generator/cdi-generate.nix b/nixos/modules/services/hardware/nvidia-container-toolkit-cdi-generator/cdi-generate.nix
new file mode 100644
index 000000000000..a90d234f65c0
--- /dev/null
+++ b/nixos/modules/services/hardware/nvidia-container-toolkit-cdi-generator/cdi-generate.nix
@@ -0,0 +1,39 @@
+{ config, lib, pkgs }: let
+  mountOptions = { options = ["ro" "nosuid" "nodev" "bind"]; };
+  mounts = [
+    { hostPath = "${lib.getBin config.hardware.nvidia.package}/bin/nvidia-cuda-mps-control";
+      containerPath = "/usr/bin/nvidia-cuda-mps-control"; }
+    { hostPath = "${lib.getBin config.hardware.nvidia.package}/bin/nvidia-cuda-mps-server";
+      containerPath = "/usr/bin/nvidia-cuda-mps-server"; }
+    { hostPath = "${lib.getBin config.hardware.nvidia.package}/bin/nvidia-debugdump";
+      containerPath = "/usr/bin/nvidia-debugdump"; }
+    { hostPath = "${lib.getBin config.hardware.nvidia.package}/bin/nvidia-powerd";
+      containerPath = "/usr/bin/nvidia-powerd"; }
+    { hostPath = "${lib.getBin config.hardware.nvidia.package}/bin/nvidia-smi";
+      containerPath = "/usr/bin/nvidia-smi"; }
+    { hostPath = "${pkgs.nvidia-container-toolkit}/bin/nvidia-ctk";
+      containerPath = "/usr/bin/nvidia-ctk"; }
+    { hostPath = "${pkgs.glibc}/lib";
+      containerPath = "${pkgs.glibc}/lib"; }
+    { hostPath = "${pkgs.glibc}/lib64";
+      containerPath = "${pkgs.glibc}/lib64"; }
+  ];
+  jqAddMountExpression = ".containerEdits.mounts[.containerEdits.mounts | length] |= . +";
+  mountsToJq = lib.concatMap
+    (mount:
+      ["${pkgs.jq}/bin/jq '${jqAddMountExpression} ${builtins.toJSON (mount // mountOptions)}'"])
+    mounts;
+in ''
+#! ${pkgs.runtimeShell}
+
+function cdiGenerate {
+  ${pkgs.nvidia-container-toolkit}/bin/nvidia-ctk cdi generate \
+    --format json \
+    --ldconfig-path ${pkgs.glibc.bin}/bin/ldconfig \
+    --library-search-path ${config.hardware.nvidia.package}/lib \
+    --nvidia-ctk-path ${pkgs.nvidia-container-toolkit}/bin/nvidia-ctk
+}
+
+cdiGenerate | \
+  ${lib.concatStringsSep " | " mountsToJq} > $RUNTIME_DIRECTORY/nvidia-container-toolkit.json
+''
diff --git a/nixos/modules/services/hardware/nvidia-container-toolkit-cdi-generator/default.nix b/nixos/modules/services/hardware/nvidia-container-toolkit-cdi-generator/default.nix
new file mode 100644
index 000000000000..3c96e9c41be5
--- /dev/null
+++ b/nixos/modules/services/hardware/nvidia-container-toolkit-cdi-generator/default.nix
@@ -0,0 +1,38 @@
+{ config, lib, pkgs, ... }:
+
+{
+
+  options = {
+
+    hardware.nvidia-container-toolkit-cdi-generator.enable = lib.mkOption {
+      default = false;
+      internal = true;
+      visible = false;
+      type = lib.types.bool;
+      description = lib.mdDoc ''
+        Enable dynamic CDI configuration for NVidia devices by running
+        nvidia-container-toolkit on boot.
+      '';
+    };
+
+  };
+
+  config = {
+
+    systemd.services.nvidia-container-toolkit-cdi-generator = lib.mkIf config.hardware.nvidia-container-toolkit-cdi-generator.enable {
+      description = "Container Device Interface (CDI) for Nvidia generator";
+      wantedBy = [ "multi-user.target" ];
+      after = [ "systemd-udev-settle.service" ];
+      serviceConfig = {
+        RuntimeDirectory = "cdi";
+        RemainAfterExit = true;
+        ExecStart = let
+          script = (pkgs.writeScriptBin "nvidia-cdi-generator"
+            (import ./cdi-generate.nix { inherit config lib pkgs; })); in (lib.getExe script);
+        Type = "oneshot";
+      };
+    };
+
+  };
+
+}