From 1f91cada6174abced1618c68c9d344d123d72231 Mon Sep 17 00:00:00 2001
From: Phillip Cloud <cloud@standard.ai>
Date: Wed, 16 Dec 2020 09:21:40 -0500
Subject: [PATCH] nomad: add optional nvml support

---
 .../networking/cluster/nomad/0.11.nix         |  8 +++-
 .../networking/cluster/nomad/0.12.nix         |  8 +++-
 .../networking/cluster/nomad/generic.nix      | 46 +++++++++++++------
 pkgs/top-level/all-packages.nix               |  4 ++
 4 files changed, 48 insertions(+), 18 deletions(-)

diff --git a/pkgs/applications/networking/cluster/nomad/0.11.nix b/pkgs/applications/networking/cluster/nomad/0.11.nix
index 2fd7c0403a3c..58bdb694eed5 100644
--- 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
index 0120887957a8..b6e053e757de 100644
--- 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/generic.nix b/pkgs/applications/networking/cluster/nomad/generic.nix
index 0bfdb0b27f6e..dc22ab41d8b6 100644
--- 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}"
+      )
+    '';
 
-  meta = with stdenv.lib; {
+  # 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 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
index ca6ec7381643..42871fe44bf0 100644
--- a/pkgs/top-level/all-packages.nix
+++ b/pkgs/top-level/all-packages.nix
@@ -6042,9 +6042,13 @@ 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);
   };
 
   notable = callPackage ../applications/misc/notable { };