From 9b798444d181a306ebd1691549157b2710283424 Mon Sep 17 00:00:00 2001
From: Jonas Juselius <jonas.juselius@oceanbox.io>
Date: Thu, 25 Sep 2025 12:15:24 +0200
Subject: [PATCH] feat: enable slurm jwt and remove slocket proxy

---
 .gitignore            |  1 +
 modules/hpc/slurm.nix | 62 ++++++++++++++++++++++++-------------------
 2 files changed, 35 insertions(+), 28 deletions(-)

diff --git a/.gitignore b/.gitignore
index db3080e..534034a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,6 @@
 *.pem
 *.csr
+*.key
 result
 result-*
 gcroots/
diff --git a/modules/hpc/slurm.nix b/modules/hpc/slurm.nix
index 97bd829..9e3f32e 100644
--- a/modules/hpc/slurm.nix
+++ b/modules/hpc/slurm.nix
@@ -11,6 +11,12 @@ let
         uid = cfg.mungeUid;
         gid = 0;
     };
+    environment.etc."slurmctld/jwt_hs256.key" = {
+        source = cfg.jwtKey;
+        mode = "0400";
+        uid = cfg.slurmUid;
+        gid = 0;
+    };
 
     services.slurm = {
       controlMachine = cfg.controlMachine;
@@ -24,8 +30,8 @@ let
         MailProg=/run/wrappers/bin/sendmail
         SelectType=select/cons_tres
         SelectTypeParameters=CR_Core
-        # AuthAltTypes=auth/jwt
-        # AuthAltParameters=jwt_key=/var/spool/slurm/statesave/jwt_hs256.key
+        AuthAltTypes=auth/jwt
+        AuthAltParameters=jwt_key=/etc/slurmctld/jwt_hs256.key
       '';
     };
 
@@ -34,7 +40,6 @@ let
     nixpkgs.overlays = [ slurm-ucx ];
   };
 
-
   slurmServer = {
     services.mysql = {
       enable = true;
@@ -60,7 +65,7 @@ let
       # '';
       dbdserver = {
         enable = true;
-        # dbdHost = cfg.controlMachine;
+        dbdHost = cfg.dbdHost;
         # storagePass  = cfg.storagePass;
       };
     };
@@ -113,8 +118,14 @@ let
     };
   };
 
+
   slurmrestd = {
-    systemd.tmpfiles.rules = [ "d /run/slurmrestd 0750 hipster hipster -" ];
+    services.dnsmasq.enable = true;
+    services.dnsmasq.settings = {
+      domain = [ "cluster.local"  ];
+      address = [ "/slurmctld.cluster.local/127.0.0.1" ];
+      srv-host= "_slurmctld._tcp.cluster.local,slurmctld.cluster.local,6817,0,5";
+    };
 
     systemd.services.slurmrestd = {
         description = "Slurm REST API service";
@@ -126,42 +137,26 @@ let
             Group = "hipster";
         };
         environment = {
-            # SLURM_JWT = "daemon";
+            SLURM_JWT = "daemon";
         };
         script = ''
-            rm -f /run/slurmrestd/hipster.socket
-            /run/current-system/sw/bin/slurmrestd -v -a rest_auth/local unix:/run/slurmrestd/hipster.socket
+             /run/current-system/sw/bin/slurmrestd -v -a rest_auth/jwt :6822
         '';
         serviceConfig = {
           RuntimeDirectory = "slurmrestd";
         };
     };
-
-    systemd.sockets.slurm-http-proxy  = {
-        enable = true;
-        description = "Proxy slurmrestd unix socket to port 6822";
-        listenStreams = [ "0.0.0.0:6822" ];
-        wantedBy = [ "sockets.target" ];
-        # If 'true', allow multiple instances of corresponding service (fails)
-        socketConfig.Accept = false;
-    };
-
-    systemd.services."slurm-http-proxy@"  = {
-        enable = true;
-        description = "Proxy slurmrestd unix socket to port 6822";
-        serviceConfig = {
-            ExecStart = "-${pkgs.socat}/bin/socat STDIO UNIX-CONNECT:/run/slurmrestd/hipster.socket";
-            StandardInput="socket";
-            User = "hipster";
-            Group = "hipster";
-        };
-    };
   };
 in
 {
   options.features.hpc.slurm = {
     enable = mkEnableOption "Enable SLURM batch system";
 
+    jwtKey = mkOption {
+      type = types.path;
+      default = null;
+    };
+
     mungeKey = mkOption {
       type = types.path;
       default = null;
@@ -172,6 +167,11 @@ in
       default = 997;
     };
 
+    slurmUid = mkOption {
+      type = types.int;
+      default = 307;
+    };
+
     pkey = mkOption {
       type = types.str;
       default = "0x7fff";
@@ -181,6 +181,12 @@ in
       type = types.str;
       default = null;
     };
+
+    dbdHost = mkOption {
+      type = types.str;
+      default = null;
+    };
+
     server = mkOption {
       type = types.bool;
       default = false;