From de3ce5ffa78eace4fadcd71b51026a37b56e609a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Fern=C3=A1ndez=20L=C3=B3pez?=
 <ereslibre@ereslibre.es>
Date: Mon, 22 Apr 2024 23:26:51 +0200
Subject: [PATCH] nixos/nvidia-container-toolkit: add
 `mount-nvidia-executables` and `mount-nvidia-docker-1-directories` options

- `mount-nvidia-binaries`: this option allows users to avoid mounting
nvidia binaries on the container.

- `mount-nvidia-docker-1-directories`: this option allows users to
avoid mounting `/usr/local/nvidia/lib{,64}` on containers.
---
 .../nvidia-container-toolkit/default.nix      | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/nixos/modules/services/hardware/nvidia-container-toolkit/default.nix b/nixos/modules/services/hardware/nvidia-container-toolkit/default.nix
index f4f3b69e1dc6..7b4973d3c6b0 100644
--- a/nixos/modules/services/hardware/nvidia-container-toolkit/default.nix
+++ b/nixos/modules/services/hardware/nvidia-container-toolkit/default.nix
@@ -42,6 +42,24 @@
         default = [];
         description = "Mounts to be added to every container under the Nvidia CDI profile.";
       };
+
+      mount-nvidia-executables = lib.mkOption {
+        default = true;
+        type = lib.types.bool;
+        description = ''
+          Mount executables nvidia-smi, nvidia-cuda-mps-control, nvidia-cuda-mps-server,
+          nvidia-debugdump, nvidia-powerd and nvidia-ctk on containers.
+        '';
+      };
+
+      mount-nvidia-docker-1-directories = lib.mkOption {
+        default = true;
+        type = lib.types.bool;
+        description = ''
+          Mount nvidia-docker-1 directories on containers: /usr/local/nvidia/lib and
+          /usr/local/nvidia/lib64.
+        '';
+      };
     };
 
   };
@@ -57,6 +75,26 @@
          containerPath = "${lib.getLib pkgs.glibc}/lib"; }
        { hostPath = "${lib.getLib pkgs.glibc}/lib64";
          containerPath = "${lib.getLib pkgs.glibc}/lib64"; }]
+      (lib.mkIf config.hardware.nvidia-container-toolkit.mount-nvidia-executables
+        [{ hostPath = lib.getExe' nvidia-driver "nvidia-cuda-mps-control";
+           containerPath = "/usr/bin/nvidia-cuda-mps-control"; }
+         { hostPath = lib.getExe' nvidia-driver "nvidia-cuda-mps-server";
+           containerPath = "/usr/bin/nvidia-cuda-mps-server"; }
+         { hostPath = lib.getExe' nvidia-driver "nvidia-debugdump";
+           containerPath = "/usr/bin/nvidia-debugdump"; }
+         { hostPath = lib.getExe' nvidia-driver "nvidia-powerd";
+           containerPath = "/usr/bin/nvidia-powerd"; }
+         { hostPath = lib.getExe' nvidia-driver "nvidia-smi";
+           containerPath = "/usr/bin/nvidia-smi"; }])
+      # nvidia-docker 1.0 uses /usr/local/nvidia/lib{,64}
+      #   e.g.
+      #     - https://gitlab.com/nvidia/container-images/cuda/-/blob/e3ff10eab3a1424fe394899df0e0f8ca5a410f0f/dist/12.3.1/ubi9/base/Dockerfile#L44
+      #     - https://github.com/NVIDIA/nvidia-docker/blob/01d2c9436620d7dde4672e414698afe6da4a282f/src/nvidia/volumes.go#L104-L173
+      (lib.mkIf config.hardware.nvidia-container-toolkit.mount-nvidia-docker-1-directories
+        [{ hostPath = "${lib.getLib nvidia-driver}/lib";
+           containerPath = "/usr/local/nvidia/lib"; }
+         { hostPath = "${lib.getLib nvidia-driver}/lib";
+           containerPath = "/usr/local/nvidia/lib64"; }])
     ]);
 
     systemd.services.nvidia-container-toolkit-cdi-generator = lib.mkIf config.hardware.nvidia-container-toolkit.enable {