From 5e7c297ba25bdaaf2806a6dfa5f436f2a3afb7d7 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Wed, 10 Jan 2024 18:43:07 +0000
Subject: [PATCH] nvidia-docker: support config.toml as an attrset argument

---
 .../nvidia-container-toolkit/default.nix      | 10 ++++-
 .../nvidia-container-toolkit/packages.nix     | 41 ++++++++++++++++++-
 .../virtualization/nvidia-docker/config.toml  | 13 ------
 .../virtualization/nvidia-podman/config.toml  | 13 ------
 4 files changed, 48 insertions(+), 29 deletions(-)
 delete mode 100644 pkgs/applications/virtualization/nvidia-docker/config.toml
 delete mode 100644 pkgs/applications/virtualization/nvidia-podman/config.toml

diff --git a/pkgs/applications/virtualization/nvidia-container-toolkit/default.nix b/pkgs/applications/virtualization/nvidia-container-toolkit/default.nix
index 71fc18f0b327..7d0ecfab53e7 100644
--- a/pkgs/applications/virtualization/nvidia-container-toolkit/default.nix
+++ b/pkgs/applications/virtualization/nvidia-container-toolkit/default.nix
@@ -5,10 +5,16 @@
 , buildGoModule
 , linkFarm
 , writeShellScript
+, formats
 , containerRuntimePath
 , configTemplate
+, configTemplatePath ? null
 , libnvidia-container
 }:
+
+assert configTemplate != null -> (lib.isAttrs configTemplate && configTemplatePath == null);
+assert configTemplatePath != null -> (lib.isStringLike configTemplatePath && configTemplate == null);
+
 let
   isolatedContainerRuntimePath = linkFarm "isolated_container_runtime_path" [
     {
@@ -23,6 +29,8 @@ let
       echo >&2 "$(tput setaf 3)warning: \$XDG_CONFIG_HOME=$XDG_CONFIG_HOME$(tput sgr 0)"
     fi
   '';
+
+  configToml = if configTemplatePath != null then configTemplatePath else (formats.toml { }).generate "config.toml" configTemplate;
 in
 buildGoModule rec {
   pname = "container-toolkit/container-toolkit";
@@ -82,7 +90,7 @@ buildGoModule rec {
       --prefix PATH : ${isolatedContainerRuntimePath}:${libnvidia-container}/bin \
       --set-default XDG_CONFIG_HOME $out/etc
 
-    cp ${configTemplate} $out/etc/nvidia-container-runtime/config.toml
+    cp ${configToml} $out/etc/nvidia-container-runtime/config.toml
 
     substituteInPlace $out/etc/nvidia-container-runtime/config.toml \
       --subst-var-by glibcbin ${lib.getBin glibc}
diff --git a/pkgs/applications/virtualization/nvidia-container-toolkit/packages.nix b/pkgs/applications/virtualization/nvidia-container-toolkit/packages.nix
index ac7da5bd33ea..2d1058a00211 100644
--- a/pkgs/applications/virtualization/nvidia-container-toolkit/packages.nix
+++ b/pkgs/applications/virtualization/nvidia-container-toolkit/packages.nix
@@ -7,16 +7,53 @@
   symlinkJoin,
 }:
 
+# Note this scope isn't recursed into, at the time of writing.
 lib.makeScope newScope (
   self: {
 
+    # The config is only exposed as an attrset so that the user may reach the
+    # deafult values, for inspectability purposes.
+    dockerConfig = {
+      disable-require = false;
+      #swarm-resource = "DOCKER_RESOURCE_GPU"
+
+      nvidia-container-cli = {
+        #root = "/run/nvidia/driver";
+        #path = "/usr/bin/nvidia-container-cli";
+        environment = [ ];
+        #debug = "/var/log/nvidia-container-runtime-hook.log";
+        ldcache = "/tmp/ld.so.cache";
+        load-kmods = true;
+        #no-cgroups = false;
+        #user = "root:video";
+        ldconfig = "@@glibcbin@/bin/ldconfig";
+      };
+    };
     nvidia-container-toolkit-docker = self.callPackage ./. {
       containerRuntimePath = "${docker}/libexec/docker/docker";
-      configTemplate = ../nvidia-docker/config.toml;
+      configTemplate = self.dockerConfig;
+    };
+
+    podmanConfig = {
+      disable-require = true;
+      #swarm-resource = "DOCKER_RESOURCE_GPU";
+
+      nvidia-container-cli = {
+        #root = "/run/nvidia/driver";
+        #path = "/usr/bin/nvidia-container-cli";
+        environment = [ ];
+        #debug = "/var/log/nvidia-container-runtime-hook.log";
+        ldcache = "/tmp/ld.so.cache";
+        load-kmods = true;
+        no-cgroups = true;
+        #user = "root:video";
+        ldconfig = "@@glibcbin@/bin/ldconfig";
+      };
     };
     nvidia-container-toolkit-podman = self.nvidia-container-toolkit-docker.override {
       containerRuntimePath = lib.getExe runc;
-      configTemplate = ../nvidia-podman/config.toml;
+
+      configTemplate = self.podmanConfig;
     };
 
     nvidia-docker = symlinkJoin {
diff --git a/pkgs/applications/virtualization/nvidia-docker/config.toml b/pkgs/applications/virtualization/nvidia-docker/config.toml
deleted file mode 100644
index bbd166995f36..000000000000
--- a/pkgs/applications/virtualization/nvidia-docker/config.toml
+++ /dev/null
@@ -1,13 +0,0 @@
-disable-require = false
-#swarm-resource = "DOCKER_RESOURCE_GPU"
-
-[nvidia-container-cli]
-#root = "/run/nvidia/driver"
-#path = "/usr/bin/nvidia-container-cli"
-environment = []
-#debug = "/var/log/nvidia-container-runtime-hook.log"
-ldcache = "/tmp/ld.so.cache"
-load-kmods = true
-#no-cgroups = false
-#user = "root:video"
-ldconfig = "@@glibcbin@/bin/ldconfig"
diff --git a/pkgs/applications/virtualization/nvidia-podman/config.toml b/pkgs/applications/virtualization/nvidia-podman/config.toml
deleted file mode 100644
index eb39699b96b3..000000000000
--- a/pkgs/applications/virtualization/nvidia-podman/config.toml
+++ /dev/null
@@ -1,13 +0,0 @@
-disable-require = true
-#swarm-resource = "DOCKER_RESOURCE_GPU"
-
-[nvidia-container-cli]
-#root = "/run/nvidia/driver"
-#path = "/usr/bin/nvidia-container-cli"
-environment = []
-#debug = "/var/log/nvidia-container-runtime-hook.log"
-ldcache = "/tmp/ld.so.cache"
-load-kmods = true
-no-cgroups = true
-#user = "root:video"
-ldconfig = "@@glibcbin@/bin/ldconfig"