modules.services.slurm: enable tui

devices/srv1/default.nix

@@ -62,6 +62,7 @@ inputs:
             localhost = [ "srv1-node0" ];
             old = [ "srv1-node1" "srv1-node2" "srv1-node3" ];
           };
+          tui = { cpuMpiThreads = 8; cpuOpenmpThreads = 10; };
         };
       };
       user.users = [ "chn" ];

devices/xmupc1/default.nix

@@ -82,6 +82,7 @@ inputs:
             gpus = { "p5000" = 1; "3090" = 1; "4090" = 1; };
           };
           partitions.localhost = [ "xmupc1" ];
+          tui = { cpuMpiThreads = 3; cpuOpenmpThreads = 4; gpus = [ "p5000" "3090" "4090" ]; };
         };
         xrdp = { enable = true; hostname = [ "xmupc1.chn.moe" ]; };
         samba =

devices/xmupc2/default.nix

@@ -81,6 +81,7 @@ inputs:
             gpus."4090" = 1;
           };
           partitions.localhost = [ "xmupc2" ];
+          tui = { cpuMpiThreads = 8; cpuOpenmpThreads = 10; gpus = [ "4090" ]; };
         };
         xrdp = { enable = true; hostname = [ "xmupc2.chn.moe" ]; };
         samba = { enable = true; hostsAllowed = ""; shares = { home.path = "/home"; root.path = "/"; }; };

flake/dev.nix

@@ -9,11 +9,13 @@
   hpcstat = pkgs.mkShell.override { stdenv = pkgs.clang18Stdenv; }
   {
     inputsFrom = [ (pkgs.localPackages.hpcstat.override { version = null; }) ];
+    packages = [ pkgs.clang-tools_18 ];
     CMAKE_EXPORT_COMPILE_COMMANDS = "1";
   };
   sbatch-tui = pkgs.mkShell.override { stdenv = pkgs.clang18Stdenv; }
   {
     inputsFrom = [ pkgs.localPackages.sbatch-tui ];
+    packages = [ pkgs.clang-tools_18 ];
     CMAKE_EXPORT_COMPILE_COMMANDS = "1";
   };
   ufo = pkgs.mkShell.override { stdenv = pkgs.clang18Stdenv; }

modules/packages/ssh.nix

@@ -82,6 +82,21 @@ inputs:
           ed25519 = "AAAAC3NzaC1lZDI1NTE5AAAAIDm6M1D7dBVhjjZtXYuzMj2P1fXNWN3O9wmwNssxEeDs";
           hostnames = [ "srv1.chn.moe" "node0.srv1.chn.moe" "wireguard.node0.srv1.chn.moe" ];
         };
+        srv1-node1 =
+        {
+          ed25519 = "AAAAC3NzaC1lZDI1NTE5AAAAIIFmG/ZzLDm23NeYa3SSI0a0uEyQWRFkaNRE9nB8egl7";
+          hostnames = [ "192.168.178.2" ];
+        };
+        srv1-node2 =
+        {
+          ed25519 = "AAAAC3NzaC1lZDI1NTE5AAAAIDhgEApzHhVPDvdVFPRuJ/zCDiR1K+rD4sZzH77imKPE";
+          hostnames = [ "192.168.178.3" ];
+        };
+        srv1-node3 =
+        {
+          ed25519 = "AAAAC3NzaC1lZDI1NTE5AAAAIO/4xbQNz6KNcEdjtBMGY8wUoFK1sCgamKl/r+kVjd7O";
+          hostnames = [ "192.168.178.4" ];
+        };
       };
       in builtins.listToAttrs (builtins.map
         (server:
@@ -147,9 +162,10 @@ inputs:
             forwardAgent = true;
             extraOptions.AddKeysToAgent = "yes";
           };
-          "wireguard.jykang" = jykang // { host = "internal.jykang"; proxyJump = "wireguard.xmupc1"; };
-          srv1-node0 = { host = "srv1-node0"; hostname = "srv1.chn.moe"; };
-          srv1-node1 = { host = "srv1-node1"; hostname = "192.168.178.2"; proxyJump = "srv1-node0"; };
+          "wireguard.jykang" = jykang // { host = "wireguard.jykang"; proxyJump = "wireguard.xmupc1"; };
+          srv1-node1 = { host = "srv1-node1"; hostname = "192.168.178.2"; proxyJump = "srv1"; };
+          srv1-node2 = { host = "srv1-node2"; hostname = "192.168.178.3"; proxyJump = "srv1"; };
+          srv1-node3 = { host = "srv1-node3"; hostname = "192.168.178.4"; proxyJump = "srv1"; };
         };
       };
     })];

modules/services/slurm.nix

@@ -23,6 +23,12 @@ inputs:
     };}));};
     partitions = mkOption { type = types.attrsOf (types.listOf types.nonEmptyStr); default = {}; };
     defaultPartition = mkOption { type = types.nonEmptyStr; default = "localhost"; };
+    tui =
+    {
+      cpuMpiThreads = mkOption { type = types.ints.unsigned; default = 1; };
+      cpuOpenmpThreads = mkOption { type = types.ints.unsigned; default = 1; };
+      gpus = mkOption { type = types.nullOr (types.attrsOf types.ints.unsigned); default = null; };
+    };
   };
   config = let inherit (inputs.config.nixos.services) slurm; in inputs.lib.mkIf slurm.enable (inputs.lib.mkMerge
   [
@@ -156,28 +162,26 @@ inputs:
       sops.secrets."slurm/db" = { owner = "slurm"; key = "mariadb/slurm"; };
       nixos =
       {
-        # TODO: rewrite
-        # packages.packages._packages = [(inputs.pkgs.localPackages.sbatch-tui.override { sbatchConfig =
-        # {
-        #   cpuMpiThreads = slurm.cpu.mpiThreads;
-        #   cpuOpenmpThreads = slurm.cpu.openmpThreads;
-        #   gpuIds =
-        #     if slurm.gpus == null then ""
-        #     else builtins.concatStringsSep ", " (builtins.map (gpu: ''"${gpu}"'') (builtins.attrNames slurm.gpus));
-        # };})];
-        # user.sharedModules = [{ home.packages =
-        # [
-        #   (inputs.pkgs.writeShellScriptBin "sbatch"
-        #   ''
-        #     if [ "$#" -eq 0 ]; then
-        #       sbatch-tui
-        #     else
-        #       /run/current-system/sw/bin/sbatch "$@"
-        #     fi
-        #   '')
-        # ];}];
+        packages.packages._packages = [ inputs.pkgs.localPackages.sbatch-tui ];
+        user.sharedModules = [{ home.packages =
+        [
+          (inputs.pkgs.writeShellScriptBin "sbatch"
+          ''
+            if [ "$#" -eq 0 ]; then
+              sbatch-tui
+            else
+              /run/current-system/sw/bin/sbatch "$@"
+            fi
+          '')
+        ];}];
         services.mariadb = { enable = true; instances.slurm = {}; };
       };
+      environment.etc."sbatch-tui.yaml".text = builtins.toJSON
+      {
+        CpuMpiThreads = slurm.tui.cpuMpiThreads;
+        CpuOpenmpThreads = slurm.tui.cpuOpenmpThreads;
+        GpuIds = slurm.tui.gpus;
+      };
     })
   ]);
 }

packages/sbatch-tui/.clangd

@@ -1,3 +0,0 @@
-CompileFlags:
-  Add: [ -Wall, -Wextra, -std=c++23 ]
-  Compiler: g++

packages/sbatch-tui/CMakeLists.txt

@@ -12,9 +12,8 @@ endif()
 find_package(ftxui REQUIRED)
 find_package(biu REQUIRED)
 
-add_executable(sbatch-tui src/main.cpp src/device.cpp)
+add_executable(sbatch-tui src/main.cpp)
 target_compile_features(sbatch-tui PUBLIC cxx_std_23)
-target_include_directories(sbatch-tui PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include)
 target_link_libraries(sbatch-tui PRIVATE ftxui::screen ftxui::dom ftxui::component biu::biu)
 
 install(TARGETS sbatch-tui RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})

packages/sbatch-tui/include/sbatch-tui/device.hpp

@@ -1,10 +0,0 @@
-# pragma once
-# include <string>
-# include <vector>
-
-struct Device_t
-{
-  unsigned CpuMpiThreads, CpuOpenmpThreads;
-  std::vector<std::string> GpuIds;
-};
-extern Device_t Device;

packages/sbatch-tui/src/device.cpp

@@ -1,8 +0,0 @@
-# include <sbatch-tui/device.hpp>
-
-Device_t Device
-{
-  .CpuMpiThreads = 1,
-  .CpuOpenmpThreads = 1,
-  .GpuIds = { "4060" }
-};

packages/sbatch-tui/src/device.cpp.template

@@ -1,8 +0,0 @@
-# include <sbatch-tui/device.hpp>
-
-Device_t Device
-{
-  .CpuMpiThreads = @cpuMpiThreads@,
-  .CpuOpenmpThreads = @cpuOpenmpThreads@,
-  .GpuIds = { @gpuIds@ }
-};

packages/sbatch-tui/src/main.cpp

@@ -3,31 +3,44 @@
 # include <ftxui/component/component_options.hpp>
 # include <ftxui/component/screen_interactive.hpp>
 # include <boost/algorithm/string.hpp>
-# include <range/v3/view.hpp>
-# include <sbatch-tui/device.hpp>
 # include <biu.hpp>
 
-using namespace biu::literals;
-
 int main()
 {
+  using namespace biu::literals;
+
+  struct Device
+  {
+    unsigned CpuMpiThreads, CpuOpenmpThreads;
+    std::optional<std::vector<std::string>> GpuIds;
+  };
+  auto device = YAML::LoadFile("/etc/sbatch-tui.yaml").as<Device>();
+
   // 需要绑定到界面上的变量
   struct
   {
     int vasp_version_selected = 0;
     std::vector<std::string> vasp_version_entries = { "std", "gam", "ncl" };
     int device_type_selected = 0;
-    std::vector<std::string> device_type_entries = { "manually select GPU", "any single GPU", "CPU" };
-    std::deque<bool> device_selected = std::deque<bool>(Device.GpuIds.size(), false);
-    std::vector<std::string> device_entries = Device.GpuIds;
+    std::vector<std::string> device_type_entries;
+    int gpu_selected = 0;
+    std::vector<std::string> gpu_entries;
     std::string job_name = std::filesystem::current_path().filename().string();
     std::string output_file = "output.txt";
-    std::string mpi_threads = std::to_string(Device.CpuMpiThreads);
-    std::string openmp_threads = std::to_string(Device.CpuOpenmpThreads);
+    std::string mpi_threads;
+    std::string openmp_threads;
 
     std::string user_command;
     std::string submit_command;
   } state;
+  if (device.GpuIds)
+  {
+    state.device_type_entries = { "manually select GPU", "any single GPU", "CPU" };
+    state.gpu_entries = *device.GpuIds;
+  }
+  else state.device_type_entries = { "CPU" };
+  state.mpi_threads = std::to_string(device.CpuMpiThreads);
+  state.openmp_threads = std::to_string(device.CpuOpenmpThreads);
 
   // 为组件增加标题栏和分割线
   auto with_title = [](std::string title)
@@ -60,30 +73,19 @@ int main()
     ftxui::Container::Horizontal
     ({
       ftxui::Menu(&state.device_type_entries, &state.device_type_selected),
-      ftxui::Container::Vertical([&]
-      {
-        std::vector<std::shared_ptr<ftxui::ComponentBase>> devices;
-        auto checkbox_option = ftxui::CheckboxOption::Simple();
-        checkbox_option.transform = [](const ftxui::EntryState& s)
-        {
-          auto prefix = ftxui::text(s.state ? "[X] " : "[ ] ");
-          auto t = ftxui::text(s.label);
-          if (s.active) t |= ftxui::bold;
-          if (s.focused) t |= ftxui::inverted;
-          return ftxui::hbox({prefix, t});
-        };
-        for (int i = 0; i < state.device_selected.size(); i++)
-          devices.push_back(ftxui::Checkbox
-            (state.device_entries[i], &state.device_selected[i], checkbox_option));
-        return devices;
-      }()) | with_separator | ftxui::Maybe([&]{ return state.device_type_selected == 0; }),
+      ftxui::Menu(&state.gpu_entries, &state.gpu_selected)
+      | with_separator
+      | ftxui::Maybe([&]
+        { return state.device_type_entries[state.device_type_selected] == "manually select GPU"; }),
       ftxui::Container::Vertical
       ({
         ftxui::Input(&state.mpi_threads) | ftxui::size(ftxui::WIDTH, ftxui::GREATER_THAN, 3)
           | with_subtitle("MPI threads: "),
         ftxui::Input(&state.openmp_threads) | ftxui::size(ftxui::WIDTH, ftxui::GREATER_THAN, 3)
           | with_subtitle("OpenMP threads: ")
-      }) | with_separator | ftxui::Maybe([&]{ return state.device_type_selected == 2; }),
+      })
+      | with_separator
+      | ftxui::Maybe([&]{ return state.device_type_entries[state.device_type_selected] == "CPU"; }),
     }) | with_title("Select device:"),
     ftxui::Input(&state.job_name) | with_title("Job name:"),
     ftxui::Input(&state.output_file) | with_title("Output file:"),
@@ -115,7 +117,8 @@ int main()
   {
     // replace \n with space
     boost::replace_all(submit_command, "\n", " ");
-    biu::exec<{.DirectStdout = true, .DirectStderr = true, .SearchPath = true}>({"sh", { "-c", submit_command }});
+    biu::exec<{.DirectStdout = true, .DirectStderr = true, .SearchPath = true}>
+      ({"sh", { "-c", submit_command }});
   };
 
   // 进入事件循环
@@ -123,11 +126,11 @@ int main()
   {
     screen.Loop(request_interface);
     if (state.user_command == "quit") return EXIT_FAILURE;
-    else if (state.device_type_selected == 1)
+    else if (state.device_type_entries[state.device_type_selected] == "any single GPU")
       state.submit_command =
         "sbatch --ntasks=1\n--gpus=1\n--job-name='{}'\n--output='{}'\nvasp-nvidia-{}"_f
         (state.job_name, state.output_file, state.vasp_version_entries[state.vasp_version_selected]);
-    else if (state.device_type_selected == 2)
+    else if (state.device_type_entries[state.device_type_selected] == "CPU")
       state.submit_command =
         "sbatch --ntasks={}\n--cpus-per-task={}\n--hint=nomultithread\n--job-name='{}'\n--output='{}'"
         "\nvasp-intel-{}"_f
@@ -135,21 +138,12 @@ int main()
           state.mpi_threads, state.openmp_threads, state.job_name, state.output_file,
           state.vasp_version_entries[state.vasp_version_selected]
         );
-    else
-    {
-      std::vector<std::string> selected_gpus;
-      for (int i = 0; i < state.device_selected.size(); i++)
-        if (state.device_selected[i]) selected_gpus.push_back(state.device_entries[i]);
-      state.submit_command =
-        "sbatch --ntasks={}\n--gres={}\n--job-name='{}'\n--output='{}'\nvasp-nvidia-{}"_f
-        (
-          selected_gpus.size(),
-          selected_gpus
-            | ranges::views::transform([](auto& entry) { return "gpu:{}:1"_f(entry); })
-            | ranges::views::join(',') | ranges::to<std::string>,
-          state.job_name, state.output_file, state.vasp_version_entries[state.vasp_version_selected]
-        );
-    }
+    else state.submit_command =
+      "sbatch --ntasks=1\n--gres=gpu:{}:1\n--job-name='{}'\n--output='{}'\nvasp-nvidia-{}"_f
+      (
+        state.gpu_entries[state.gpu_selected],
+        state.job_name, state.output_file, state.vasp_version_entries[state.vasp_version_selected]
+      );
     screen.Loop(confirm_interface);
     if (state.user_command == "quit") return EXIT_FAILURE;
     else if (state.user_command == "back") continue;