Merge pull request #108341 from cpcloud/nomad-nvml

nomad: add optional nvml support
NixOS · Jan 4, 2021 · 91e2dd2 · 91e2dd2
2 parents 0efa9dd + 2aabde2
commit 91e2dd2
Show file tree

Hide file tree

Showing 5 changed files with 56 additions and 20 deletions.
diff --git a/pkgs/applications/networking/cluster/nomad/0.11.nix b/pkgs/applications/networking/cluster/nomad/0.11.nix
@@ -1,7 +1,11 @@
-{ callPackage, buildGoPackage }:
+{ callPackage
+, buildGoPackage
+, nvidia_x11
+, nvidiaGpuSupport
+}:
 
 callPackage ./generic.nix {
-  inherit buildGoPackage;
+  inherit buildGoPackage nvidia_x11 nvidiaGpuSupport;
   version = "0.11.8";
   sha256 = "1dhh07bifr02jh2lls8fv1d9ra67ymgh8qxqvpvm0cd0qdd469z1";
 }
diff --git a/pkgs/applications/networking/cluster/nomad/0.12.nix b/pkgs/applications/networking/cluster/nomad/0.12.nix
@@ -1,7 +1,11 @@
-{ callPackage, buildGoPackage }:
+{ callPackage
+, buildGoPackage
+, nvidia_x11
+, nvidiaGpuSupport
+}:
 
 callPackage ./generic.nix {
-  inherit buildGoPackage;
+  inherit buildGoPackage nvidia_x11 nvidiaGpuSupport;
   version = "0.12.9";
   sha256 = "1a0ig6pb0z3qp7zk4jgz3h241bifmjlyqsfikyy3sxdnzj7yha27";
 }
diff --git a/pkgs/applications/networking/cluster/nomad/1.0.nix b/pkgs/applications/networking/cluster/nomad/1.0.nix
@@ -1,7 +1,11 @@
-{ callPackage, buildGoPackage }:
+{ callPackage
+, buildGoPackage
+, nvidia_x11
+, nvidiaGpuSupport
+}:
 
 callPackage ./generic.nix {
-  inherit buildGoPackage;
+  inherit buildGoPackage nvidia_x11 nvidiaGpuSupport;
   version = "1.0.1";
   sha256 = "07k81csyxhgc7bgn297zlqyvc55qb5fmiavi7dk81rdpg5m2zjvv";
 }
diff --git a/pkgs/applications/networking/cluster/nomad/generic.nix b/pkgs/applications/networking/cluster/nomad/generic.nix
@@ -1,4 +1,12 @@
-{ stdenv, buildGoPackage, fetchFromGitHub, version, sha256 }:
+{ lib
+, buildGoPackage
+, fetchFromGitHub
+, version
+, sha256
+, nvidiaGpuSupport
+, patchelf
+, nvidia_x11
+}:
 
 buildGoPackage rec {
   pname = "nomad";
@@ -14,23 +22,33 @@ buildGoPackage rec {
     inherit rev sha256;
   };
 
+  nativeBuildInputs = lib.optionals nvidiaGpuSupport [
+    patchelf
+  ];
+
   # ui:
   #  Nomad release commits include the compiled version of the UI, but the file
   #  is only included if we build with the ui tag.
-  # nonvidia:
-  #  We disable Nvidia GPU scheduling on Linux, as it doesn't work there:
-  #  Ref: https://github.com/hashicorp/nomad/issues/5535
-  preBuild = let
-    tags = ["ui"]
-      ++ stdenv.lib.optional stdenv.isLinux "nonvidia";
-    tagsString = stdenv.lib.concatStringsSep " " tags;
-  in ''
-    export buildFlagsArray=(
-      -tags="${tagsString}"
-    )
- '';
+  preBuild =
+    let
+      tags = [ "ui" ] ++ lib.optional (!nvidiaGpuSupport) "nonvidia";
+      tagsString = lib.concatStringsSep " " tags;
+    in
+    ''
+      export buildFlagsArray=(
+        -tags="${tagsString}"
+      )
+    '';
+
+  # The dependency on NVML isn't explicit. We have to make it so otherwise the
+  # binary will not know where to look for the relevant symbols.
+  postFixup = lib.optionalString nvidiaGpuSupport ''
+    for bin in $out/bin/*; do
+      patchelf --add-needed "${nvidia_x11}/lib/libnvidia-ml.so" "$bin"
+    done
+  '';
 
-  meta = with stdenv.lib; {
+  meta = with lib; {
     homepage = "https://www.nomadproject.io/";
     description = "A Distributed, Highly Available, Datacenter-Aware Scheduler";
     platforms = platforms.unix;

diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix
@@ -6129,12 +6129,18 @@ in
   # with different versions we pin Go for all versions.
   nomad_0_11 = callPackage ../applications/networking/cluster/nomad/0.11.nix {
     buildGoPackage = buildGo114Package;
+    inherit (linuxPackages) nvidia_x11;
+    nvidiaGpuSupport = config.cudaSupport or (!stdenv.isLinux);
   };
   nomad_0_12 = callPackage ../applications/networking/cluster/nomad/0.12.nix {
     buildGoPackage = buildGo114Package;
+    inherit (linuxPackages) nvidia_x11;
+    nvidiaGpuSupport = config.cudaSupport or (!stdenv.isLinux);
   };
   nomad_1_0 = callPackage ../applications/networking/cluster/nomad/1.0.nix {
     buildGoPackage = buildGo115Package;
+    inherit (linuxPackages) nvidia_x11;
+    nvidiaGpuSupport = config.cudaSupport or (!stdenv.isLinux);
   };
 
   notable = callPackage ../applications/misc/notable { };