Revert "devices: xmupc1 xmupc2 -> srv2"

This reverts commit cac52c47af.

devices/one/default.nix

@@ -20,6 +20,7 @@ inputs:
           rollingRootfs = {};
         };
         nixpkgs.march = "tigerlake";
+        kernel.variant = "cachyos";
       };
       hardware = { cpus = [ "intel" ]; gpu.type = "intel"; };
       services =

devices/pc/default.nix

@@ -134,7 +134,7 @@ inputs:
             gpus."4060" = 1;
           };
           partitions.localhost = [ "pc" ];
-          tui = { cpuMpiThreads = 4; cpuOpenmpThreads = 4; gpus = [ "4060" ]; };
+          tui = { cpuQueues = [{ mpiThreads = 4; openmpThreads = 4; }]; gpuIds = [ "4060" ]; };
         };
         ollama = {};
         docker = {};

devices/srv1/default.nix

@@ -63,7 +63,11 @@ inputs:
             fdtd = [ "srv1-node2" ];
             all = [ "srv1-node0" "srv1-node1" "srv1-node2" "srv1-node3" ];
           };
-          tui = { cpuMpiThreads = 8; cpuOpenmpThreads = 10; };
+          tui.cpuQueues =
+          [
+            { mpiThreads = 8; openmpThreads = 10; }
+            { name = "old"; mpiThreads = 8; openmpThreads = 4; }
+          ];
           setupFirewall = true;
         };
       };

devices/xmupc1/default.nix

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

devices/xmupc2/default.nix

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

flake.lock

@@ -61,11 +61,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1726983918,
-        "narHash": "sha256-eRGl50mhbJRT9Mk2NbePr6yzxzZRygndx4O7kObVa3g=",
+        "lastModified": 1736213063,
+        "narHash": "sha256-Ux/KRFONsbp97N/O7vK1OC7huF2tFgrfVaScUTMpVYc=",
         "ref": "refs/heads/master",
-        "rev": "7370c1d1ef2421608e1355368361af92777a2c6c",
-        "revCount": 940,
+        "rev": "cd3730c07eb2358b81a5fc96706ba64ecf38d321",
+        "revCount": 941,
         "type": "git",
         "url": "https://git.chn.moe/chn/bscpkgs.git"
       },
@@ -221,6 +221,22 @@
         "type": "github"
       }
     },
+    "fancy-motd": {
+      "flake": false,
+      "locked": {
+        "lastModified": 1736229765,
+        "narHash": "sha256-EJGrEfuno6p7aDNcfo5SNGK0ytMZV7mA3zxrOHsK+Sw=",
+        "owner": "CHN-beta",
+        "repo": "fancy-motd",
+        "rev": "75e6dc700589ac43c7de5a8f92c6909bd6ca4eb7",
+        "type": "github"
+      },
+      "original": {
+        "owner": "CHN-beta",
+        "repo": "fancy-motd",
+        "type": "github"
+      }
+    },
     "flake-compat": {
       "flake": false,
       "locked": {
@@ -1185,22 +1201,6 @@
         "type": "github"
       }
     },
-    "qd": {
-      "flake": false,
-      "locked": {
-        "lastModified": 1452446912,
-        "narHash": "sha256-ftU28BMGUCq0or0nVkf9fryaLfqGnMonmHu7UuyjMTQ=",
-        "owner": "scibuilder",
-        "repo": "QD",
-        "rev": "a5dbb6136ac1739e7e6955bf7b0c0fa05f3b2cbd",
-        "type": "github"
-      },
-      "original": {
-        "owner": "scibuilder",
-        "repo": "QD",
-        "type": "github"
-      }
-    },
     "root": {
       "inputs": {
         "aagl": "aagl",
@@ -1214,6 +1214,7 @@
         "date": "date",
         "eigen": "eigen",
         "envfs": "envfs",
+        "fancy-motd": "fancy-motd",
         "git-lfs-transfer": "git-lfs-transfer",
         "gricad": "gricad",
         "hextra": "hextra",
@@ -1242,7 +1243,6 @@
         "pocketfft": "pocketfft",
         "py4vasp": "py4vasp",
         "qchem": "qchem",
-        "qd": "qd",
         "rsshub": "rsshub",
         "rycee": "rycee",
         "slate": "slate",

flake.nix

@@ -69,7 +69,7 @@
     ufo = { url = "git+https://git.chn.moe/chn/ufo.git"; flake = false; };
     highfive = { url = "git+https://github.com/CHN-beta/HighFive?submodules=1"; flake = false; };
     stickerpicker = { url = "github:maunium/stickerpicker"; flake = false; };
-    qd = { url = "github:scibuilder/QD"; flake = false; };
+    fancy-motd = { url = "github:CHN-beta/fancy-motd"; flake = false; };
   };
 
   outputs = inputs: let localLib = import ./flake/lib.nix inputs.nixpkgs.lib; in

modules/packages/server.nix

@@ -10,7 +10,7 @@ inputs:
       [
         # basic tools
         beep dos2unix gnugrep pv tmux screen parallel tldr cowsay jq yq zellij ipfetch localPackages.pslist
-        fastfetch reptyr duc ncdu progress libva-utils ksh neofetch
+        fastfetch reptyr duc ncdu progress libva-utils ksh neofetch dateutils
         # lsxx
         pciutils usbutils lshw util-linux lsof dmidecode lm_sensors hwloc acpica-tools
         # top

modules/services/slurm.nix

@@ -15,8 +15,6 @@ inputs:
         sockets = mkOption { type = types.ints.unsigned; default = 1; };
         cores = mkOption { type = types.ints.unsigned; default = 1; };
         threads = mkOption { type = types.ints.unsigned; default = 1; };
-        mpiThreads = mkOption { type = types.ints.unsigned; default = 1; };
-        openmpThreads = mkOption { type = types.ints.unsigned; default = 1; };
       };
       memoryMB = mkOption { type = types.ints.unsigned; default = 1024; };
       gpus = mkOption { type = types.nullOr (types.attrsOf types.ints.unsigned); default = null; };
@@ -25,9 +23,17 @@ inputs:
     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.listOf types.nonEmptyStr); default = null; };
+      cpuQueues = mkOption
+      {
+        type = types.nonEmptyListOf (types.submodule (submoduleInputs: { options =
+        {
+          name = mkOption { type = types.nonEmptyStr; default = "localhost"; };
+          mpiThreads = mkOption { type = types.ints.unsigned; default = 1; };
+          openmpThreads = mkOption { type = types.ints.unsigned; default = 1; };
+        };}));
+      };
+      gpuIds = mkOption { type = types.nullOr (types.listOf types.nonEmptyStr); default = null; };
+      gpuPartition = mkOption { type = types.nonEmptyStr; default = "localhost"; };
     };
     # 是否打开防火墙相应端口，对于多节点部署需要打开
     setupFirewall = mkOption { type = types.bool; default = false; };
@@ -125,6 +131,13 @@ inputs:
 
             # enable task plugins
             TaskPlugin=task/affinity,task/cgroup
+
+            # omit --mpi=pmix
+            MpiDefault=pmix
+
+            # record more info
+            JobAcctGatherType=jobacct_gather/cgroup
+            AccountingStorageTRES=gres/gpu
           '';
           extraConfigPaths =
             let gpus = slurm.node.${inputs.config.nixos.model.hostname}.gpus or null;
@@ -135,6 +148,14 @@ inputs:
                 (inputs.localLib.attrsToList gpus));
               in [(inputs.pkgs.writeTextDir "gres.conf" "AutoDetect=nvml\n${gpuString}")]
             );
+          extraCgroupConfig =
+          ''
+            ConstrainCores=yes
+            ConstrainDevices=yes
+            ConstrainRAMSpace=yes
+            ConstrainSwapSpace=yes
+            AllowedSwapSpace=20
+          '';
         };
         munge = { enable = true; password = inputs.config.sops.secrets."munge.key".path; };
       };
@@ -157,6 +178,7 @@ inputs:
       networking.firewall =
         let config = inputs.lib.mkIf slurm.setupFirewall [ 6818 ];
         in { allowedTCPPorts = config; allowedUDPPorts = config; };
+      environment.sessionVariables.SLURM_HINT = "nomultithread";
     }
     # master 配置
     (inputs.lib.mkIf (slurm.master == inputs.config.nixos.model.hostname)
@@ -196,9 +218,11 @@ inputs:
       };
       environment.etc."sbatch-tui.yaml".text = builtins.toJSON
       {
-        CpuMpiThreads = slurm.tui.cpuMpiThreads;
-        CpuOpenmpThreads = slurm.tui.cpuOpenmpThreads;
-        GpuIds = slurm.tui.gpus;
+        GpuIds = slurm.tui.gpuIds;
+        GpuPartition = slurm.tui.gpuPartition;
+        CpuQueues = builtins.map
+          (queue: [ queue.name [ queue.mpiThreads queue.openmpThreads ]])
+          slurm.tui.cpuQueues;
       };
       networking.firewall =
         let config = inputs.lib.mkIf slurm.setupFirewall [ 6817 ];

modules/services/sshd/banner-rendered.txt

@@ -1,19 +0,0 @@
-                   ░▒▓█▓▒░      ░▒▓█▓▒░▒▓███████▓▒░  
-                  ░▒▓█▓▒░      ░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░ 
-                 ░▒▓█▓▒░      ░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░ 
-                ░▒▓█▓▒░      ░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░ 
-               ░▒▓█▓▒░      ░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░ 
-              ░▒▓█▓▒░      ░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░ 
-             ░▒▓████████▓▒░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░
-
-      ░▒▓██████▓▒░░▒▓███████▓▒░ ░▒▓██████▓▒░░▒▓█▓▒░░▒▓█▓▒░▒▓███████▓▒░  
-    ░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░ 
-   ░▒▓█▓▒░      ░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░ 
-  ░▒▓█▓▒▒▓███▓▒░▒▓███████▓▒░░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓███████▓▒░  
- ░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░        
-░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░░▒▓█▓▒░▒▓█▓▒░        
-░▒▓██████▓▒░░▒▓█▓▒░░▒▓█▓▒░░▒▓██████▓▒░ ░▒▓██████▓▒░░▒▓█▓▒░        
-
-               ┌──────────────────────────┐
-               │ＩｎＡｌＧａＮ  ／  ＳｉＣ│
-               └──────────────────────────┘

modules/services/sshd/default.nix

@@ -9,25 +9,35 @@ inputs:
     };});
     default = null;
   };
-  config = let inherit (inputs.config.nixos.services) sshd; in inputs.lib.mkIf (sshd != null)
-  {
-    services.openssh =
+  config = let inherit (inputs.config.nixos.services) sshd; in inputs.lib.mkIf (sshd != null) (inputs.lib.mkMerge
+  [
     {
-      enable = true;
-      settings =
+      services.openssh =
       {
-        X11Forwarding = true;
-        ChallengeResponseAuthentication = false;
-        PasswordAuthentication = sshd.passwordAuthentication;
-        KbdInteractiveAuthentication = false;
-        UsePAM = true;
+        enable = true;
+        settings =
+        {
+          X11Forwarding = true;
+          ChallengeResponseAuthentication = false;
+          PasswordAuthentication = sshd.passwordAuthentication;
+          KbdInteractiveAuthentication = false;
+          UsePAM = true;
+        };
       };
-    };
-    nixos.services.xray.client.v2ray-forwarder.noproxyTcpPorts = [ 22 ];
-    # generate from https://patorjk.com/software/taag with font "BlurVision ASCII"
-    # generate using `toilet -f wideterm -F border "InAlGaN / SiC"`
-    # somehow lolcat could not run with these characters, use rendered directly
-    # TODO: move this settings to user
-    users.motdFile = inputs.lib.mkIf sshd.groupBanner ./banner-rendered.txt;
-  };
+      nixos.services.xray.client.v2ray-forwarder.noproxyTcpPorts = [ 22 ];
+    }
+    # 如果是服务器，那么启用 motd
+    (inputs.lib.mkIf (inputs.config.nixos.model.type == "server")
+    {
+      nixos =
+      {
+        packages.packages._packages =
+          [ (inputs.pkgs.fancy-motd.overrideAttrs { src = inputs.topInputs.fancy-motd; }) ];
+        user.sharedModules = [(home-inputs: { config.programs.zsh.loginExtra = "motd"; })];
+      };
+      # generate from https://patorjk.com/software/taag with font "BlurVision ASCII"
+      # generate using `toilet -f wideterm -F border "InAlGaN / SiC"`
+      environment.etc = inputs.lib.mkIf sshd.groupBanner { "fancy-motd/banner".source = ./banner.txt; };
+    })
+  ]);
 }

modules/system/kernel/default.nix

@@ -90,9 +90,9 @@ inputs:
         { boot.initrd = { systemd.enableTpm2 = false; includeDefaultModules = false; }; }
     )
     # enable scx when using cachyos
-    (
-      inputs.lib.mkIf (builtins.elem kernel.variant [ "cachyos" "cachyos-lto" "cachyos-server" ])
-        { services.scx = { enable = true; scheduler = "scx_lavd"; extraArgs = [ "--autopower" ]; }; }
-    )
+    # (
+    #   inputs.lib.mkIf (builtins.elem kernel.variant [ "cachyos" "cachyos-lto" "cachyos-server" ])
+    #     { services.scx = { enable = true; scheduler = "scx_lavd"; extraArgs = [ "--autopower" ]; }; }
+    # )
   ];
 }

modules/system/nixpkgs.nix

@@ -49,7 +49,7 @@ inputs:
           {
             # TODO: change znver4 after update oneapi
             # TODO: test znver3 do use AVX
-            oneapiArch = let match = { znver3 = "CORE-AVX2"; znver4 = "core-avx2"; };
+            oneapiArch = let match = {};
               in match.${nixpkgs.march} or nixpkgs.march;
             nvhpcArch = nixpkgs.march;
             # contentAddressedByDefault = true;

packages/nvhpcStdenv.nix

@@ -1,7 +1,7 @@
 {
-  src, stdenv, autoPatchelfHook, wrapCCWith, writeText, addAttrsToDerivation, config, overrideCC, symlinkJoin,
-  gcc, glibc, libz, zstd, libxml2, flock, numactl, ncurses, dpkg, cudaPackages, openssl, gmp,
-  libxcrypt-legacy, libfabric, rdma-core, gfortran, xorg, makeSetupHook, writeScript, bash, overrideInStdenv
+  src, stdenv, autoPatchelfHook, wrapCCWith, config, overrideCC, makeSetupHook, writeScript, overrideInStdenv,
+  gcc, glibc, libz, zstd, libxml2, flock, numactl, ncurses, openssl, gmp,
+  libxcrypt-legacy, libfabric, rdma-core, xorg, bash
 }:
 let
   nvhpc = stdenv.mkDerivation
@@ -10,7 +10,7 @@ let
     inherit (src) src version;
     buildInputs =
       [ libz libxml2 zstd numactl ncurses openssl gmp libxcrypt-legacy libfabric rdma-core xorg.libpciaccess ];
-    nativeBuildInputs = [ autoPatchelfHook dpkg flock ];
+    nativeBuildInputs = [ autoPatchelfHook flock ];
     langFortran = true;
     dontConfigure = true;
     dontBuild = true;

packages/sbatch-tui/src/main.cpp

@@ -11,36 +11,92 @@ int main()
 
   struct Device
   {
-    unsigned CpuMpiThreads, CpuOpenmpThreads;
+    // Queue : { CpuMpiThreads, CpuOpenmpThreads }
+    std::vector<std::pair<std::string, std::pair<int, int>>> CpuQueues;
     std::optional<std::vector<std::string>> GpuIds;
+    std::string GpuPartition;
   };
   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;
-    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";
+    // 第一行，要选择的程序
+    int program_selected = 0;
+    std::vector<std::string> program_entries; // 稍后初始化
+    int vasp_selected = 0;
+    std::vector<std::string> vasp_entries = { "std", "gam", "ncl" };
+
+    // 第二行，如果是CPU，要选择的队列，和队列的参数
+    int queue_selected = 0;
+    std::vector<std::string> queue_entries; // 稍后初始化
     std::string mpi_threads;
     std::string openmp_threads;
 
+    // 第二行，如果是GPU，要选择的方案和设备
+    int gpu_scheme_selected = 0;
+    std::vector<std::string> gpu_scheme_entries = { "manually select a GPU", "any single GPU" };
+    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 user_command;
+
+    // 生成的提交命令
     std::string submit_command;
-  } state;
-  if (device.GpuIds)
+  } state =
   {
-    state.device_type_entries = { "manually select GPU", "any single GPU", "CPU" };
-    state.gpu_entries = *device.GpuIds;
+    .program_entries = [&]
+    {
+      std::vector<std::string> entries;
+      if (device.GpuIds) entries.push_back("VASP(GPU)");
+      entries.push_back("VASP(CPU)");
+      return entries;
+    }(),
+    .queue_entries = device.CpuQueues | ranges::views::keys | ranges::to_vector,
+    .gpu_entries = device.GpuIds.value_or(std::vector<std::string>{})
+  };
+
+  // 尝试从配置文件中读取设置
+  try
+  {
+    auto config = YAML::LoadFile("{}/.config/sbatch-tui/config.yaml"_f(getenv("HOME")));
+    auto saved_state = config.as<decltype(state)>();
+    // 比较时会提升到 unsigned，所以不需要额外判断是否小于 0
+    if (saved_state.program_selected < state.program_entries.size())
+      state.program_selected = saved_state.program_selected;
+    if (saved_state.vasp_selected < state.vasp_entries.size())
+      state.vasp_selected = saved_state.vasp_selected;
+    if (saved_state.queue_selected < state.queue_entries.size())
+    {
+      state.queue_selected = saved_state.queue_selected;
+      state.mpi_threads = saved_state.mpi_threads;
+      state.openmp_threads = saved_state.openmp_threads;
+    }
+    if (saved_state.gpu_scheme_selected < state.gpu_scheme_entries.size())
+      state.gpu_scheme_selected = saved_state.gpu_scheme_selected;
+    if (saved_state.gpu_selected < state.gpu_entries.size())
+      state.gpu_selected = saved_state.gpu_selected;
   }
-  else state.device_type_entries = { "CPU" };
-  state.mpi_threads = std::to_string(device.CpuMpiThreads);
-  state.openmp_threads = std::to_string(device.CpuOpenmpThreads);
+  catch (...) {}
+
+  // 刷新状态
+  auto refresh_state = [&]
+  {
+    // 如果选择了 CPU 程序，那么按照选定的队列刷新 MPI 和 OpenMP 线程数
+    if (state.program_entries[state.program_selected] == "VASP(CPU)")
+    {
+      auto it = ranges::find_if(device.CpuQueues,
+        [&](auto &x){ return x.first == state.queue_entries[state.queue_selected]; });
+      auto [mpi_threads, openmp_threads] = it->second;
+      state.mpi_threads = std::to_string(mpi_threads);
+      state.openmp_threads = std::to_string(openmp_threads);
+    }
+  };
 
   // 为组件增加标题栏和分割线
   auto with_title = [](std::string title)
@@ -68,53 +124,88 @@ int main()
   auto screen = ftxui::ScreenInteractive::Fullscreen();
   auto request_interface = ftxui::Container::Vertical
   ({
-    ftxui::Menu(&state.vasp_version_entries, &state.vasp_version_selected)
-      | with_title("Select VASP version:"),
+    // 第一行：选择程序
     ftxui::Container::Horizontal
     ({
-      ftxui::Menu(&state.device_type_entries, &state.device_type_selected),
-      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::Menu(&state.program_entries, &state.program_selected, ftxui::MenuOption{.on_change = refresh_state}),
+      // 右侧：选择 VASP 版本
+      ftxui::Menu(&state.vasp_entries, &state.vasp_selected) | with_separator
+    }) | with_title("Program:"),
+    // 第二行
+    ftxui::Container::Horizontal
+    ({
+      // 如果是选择 CPU 程序
+      ftxui::Container::Horizontal
       ({
-        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_entries[state.device_type_selected] == "CPU"; }),
-    }) | with_title("Select device:"),
+        // 左侧：选择队列
+        ftxui::Menu(&state.queue_entries, &state.queue_selected, ftxui::MenuOption{.on_change = refresh_state}),
+        // 右侧：输入 MPI 和 OpenMP 线程数
+        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.program_entries[state.program_selected] == "VASP(CPU)"; }),
+      // 如果是选择 GPU 程序
+      ftxui::Container::Horizontal
+      ({
+        // 左侧：选择方案
+        ftxui::Menu(&state.gpu_scheme_entries, &state.gpu_scheme_selected,
+          ftxui::MenuOption{.on_change = refresh_state}),
+        // 右侧：选择 GPU
+        ftxui::Menu(&state.gpu_entries, &state.gpu_selected) | with_separator
+          | ftxui::Maybe([&]{ return state.gpu_scheme_entries[state.gpu_scheme_selected] == "manually select a GPU"; }),
+      }) | ftxui::Maybe([&]{ return state.program_entries[state.program_selected] == "VASP(GPU)"; }),
+    }) | with_title("Resource allocation parameters:"),
+    // 第三行：任务名
     ftxui::Input(&state.job_name) | with_title("Job name:"),
+    // 第四行：输出文件
     ftxui::Input(&state.output_file) | with_title("Output file:"),
+    // 操作按钮
     ftxui::Container::Horizontal
     ({
-      ftxui::Button("Continue",
+      ftxui::Button("Continue (Enter)",
         [&]{ state.user_command = "continue"; screen.ExitLoopClosure()(); }),
       ftxui::Button("Quit",
         [&]{ state.user_command = "quit"; screen.ExitLoopClosure()(); })
     })
-  }) | ftxui::borderHeavy | with_padding;
+  }) | ftxui::borderHeavy | with_padding | ftxui::CatchEvent([&](ftxui::Event event)
+  {
+    if (event == ftxui::Event::Return) { state.user_command = "continue"; screen.ExitLoopClosure()(); }
+    return event == ftxui::Event::Return;
+  });
   auto confirm_interface = ftxui::Container::Vertical
   ({
     ftxui::Input(&state.submit_command, "", ftxui::InputOption{.multiline = true})
       | with_title("Double check & modify submit command:"),
     ftxui::Container::Horizontal
     ({
-      ftxui::Button("Submit",
+      ftxui::Button("Submit (Enter)",
         [&]{state.user_command = "submit"; screen.ExitLoopClosure()();}),
       ftxui::Button("Quit",
         [&]{state.user_command = "quit"; screen.ExitLoopClosure()();}),
       ftxui::Button("Back",
         [&]{state.user_command = "back"; screen.ExitLoopClosure()();})
     })
-  }) | ftxui::borderHeavy | with_padding;
+  }) | ftxui::borderHeavy | with_padding | ftxui::CatchEvent([&](ftxui::Event event)
+  {
+    if (event == ftxui::Event::Return) { state.user_command = "submit"; screen.ExitLoopClosure()(); }
+    return event == ftxui::Event::Return;
+  });
 
   // 实际投递任务
-  auto submit = [](std::string submit_command)
+  auto submit = [&](std::string submit_command)
   {
+    // 保存设置
+    try
+    {
+      std::filesystem::create_directories("{}/.config/sbatch-tui"_f(std::getenv("HOME")));
+      std::ofstream("{}/.config/sbatch-tui/config.yaml"_f(std::getenv("HOME"))) << YAML::Node(state);
+    }
+    catch (...) {}
     // replace \n with space
     boost::replace_all(submit_command, "\n", " ");
     biu::exec<{.DirectStdout = true, .DirectStderr = true, .SearchPath = true}>
@@ -125,31 +216,38 @@ int main()
   while (true)
   {
     screen.Loop(request_interface);
-    if (state.user_command == "quit") return EXIT_FAILURE;
-    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_entries[state.device_type_selected] == "manually select GPU")
-      state.submit_command =
-        "sbatch --ntasks=1\n--gres=gpu:{}:1\n--job-name='{}'\n--output='{}'\nvasp-nvidia-{}"_f
+    if (state.user_command == "quit") return 0;
+    else if (state.user_command == "continue")
+    {
+      if (state.program_entries[state.program_selected] == "VASP(GPU)")
+        if (state.gpu_scheme_entries[state.gpu_scheme_selected] == "any single GPU")
+          state.submit_command =
+            "sbatch --partition={}\n--ntasks=1 --gpus=1\n--job-name='{}' --output='{}'\n"
+              "--wrap=\"vasp-nvidia mpirun vasp-{}\""_f
+            (device.GpuPartition, state.job_name, state.output_file, state.vasp_entries[state.vasp_selected]);
+        else
+          state.submit_command =
+            "sbatch --partition={}\n--ntasks=1 --gpus={}:1\n--job-name='{}' --output='{}'\n"
+              "--wrap=\"vasp-nvidia mpirun vasp-{}\""_f
+            (
+              device.GpuPartition, state.gpu_entries[state.gpu_selected],
+              state.job_name, state.output_file, state.vasp_entries[state.vasp_selected]
+            );
+      else state.submit_command =
+        "sbatch --partition={} --nodes=1-1\n--ntasks={} --cpus-per-task={}\n--job-name='{}' --output='{}'\n"
+          "--wrap=\"vasp-intel srun vasp-{}\""_f
         (
-          state.gpu_entries[state.gpu_selected],
-          state.job_name, state.output_file, state.vasp_version_entries[state.vasp_version_selected]
+          state.queue_entries[state.queue_selected],
+          state.mpi_threads, state.openmp_threads, state.job_name, state.output_file,
+          state.vasp_entries[state.vasp_selected]
         );
-    else state.submit_command =
-      "sbatch --ntasks={}\n--cpus-per-task={}\n"
-      "--export=ALL,OMP_NUM_THREADS={},OMP_STACKSIZE=512m\n--hint=nomultithread\n--job-name='{}'\n"
-      "--output='{}'\n--wrap=\"vasp-intel srun --mpi=pmix vasp-{}\""_f
-      (
-        state.mpi_threads, state.openmp_threads, state.openmp_threads, state.job_name, state.output_file,
-        state.vasp_version_entries[state.vasp_version_selected]
-      );
-
+      state.user_command.clear();
+    }
+    else return EXIT_FAILURE;
     screen.Loop(confirm_interface);
-    if (state.user_command == "quit") return EXIT_FAILURE;
-    else if (state.user_command == "back") continue;
-    submit(state.submit_command);
-    break;
+    if (state.user_command == "quit") return 0;
+    else if (state.user_command == "back") { state.user_command.clear(); continue; }
+    else if (state.user_command == "submit") { submit(state.submit_command); break; }
+    else return EXIT_FAILURE;
   }
 }

packages/vasp/intel/default.nix

@@ -12,7 +12,7 @@ let vasp = stdenv.mkDerivation
       cp ${./makefile.include} makefile.include
       chmod +w makefile.include
       cp ${../constr_cell_relax.F} src/constr_cell_relax.F
-      cp -r ${vtst}/vtstcode6.4/* src
+      cp -r ${vtst}/vtstcode6.4.3/* src
       chmod -R +w src
     '';
     buildInputs = [ hdf5 wannier90 mkl ];
@@ -41,5 +41,19 @@ let vasp = stdenv.mkDerivation
 in writeShellScriptBin "vasp-intel"
 ''
   export PATH=${vasp}/bin:${mpi}/bin''${PATH:+:$PATH}
+
+  # set OMP_NUM_THREADS if SLURM_CPUS_PER_TASK is set
+  if [ -z "$OMP_NUM_THREADS" ] && [ -n "$SLURM_CPUS_PER_TASK" ]; then
+    export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
+  fi
+  # set OMP_NUM_THREADS to 1 if not set
+  if [ -z "$OMP_NUM_THREADS" ]; then
+    export OMP_NUM_THREADS=1
+  fi
+  # set OMP_STACKSIZE to 512M if not set
+  if [ -z "$OMP_STACKSIZE" ]; then
+    export OMP_STACKSIZE=512m
+  fi
+
   exec "$@"
 ''

packages/vasp/nvidia/default.nix

@@ -12,7 +12,7 @@ let vasp = stdenv.mkDerivation
     cp ${./makefile.include} makefile.include
     chmod +w makefile.include
     cp ${../constr_cell_relax.F} src/constr_cell_relax.F
-    cp -r ${vtst}/vtstcode6.4/* src
+    cp -r ${vtst}/vtstcode6.4.3/* src
     chmod -R +w src
   '';
   buildInputs = [ hdf5 wannier90 mkl ];
@@ -39,5 +39,11 @@ let vasp = stdenv.mkDerivation
 in writeShellScriptBin "vasp-nvidia"
 ''
   export PATH=${vasp}/bin''${PATH:+:$PATH}
+
+  # set OMP_NUM_THREADS if SLURM_CPUS_PER_TASK is set
+  if [ -z "$OMP_NUM_THREADS" ] && [ -n "$SLURM_CPUS_PER_TASK" ]; then
+    export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
+  fi
+
   exec ${stdenv.cc.cc.runEnv} "$@"
 ''