diff --git a/devices/srv1/default.nix b/devices/srv1/default.nix index bce309ab..1903330a 100644 --- a/devices/srv1/default.nix +++ b/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" ]; diff --git a/devices/xmupc1/default.nix b/devices/xmupc1/default.nix index a79b8202..8df3a6e6 100644 --- a/devices/xmupc1/default.nix +++ b/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 = diff --git a/devices/xmupc2/default.nix b/devices/xmupc2/default.nix index 40e1b540..e0913279 100644 --- a/devices/xmupc2/default.nix +++ b/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 = "/"; }; }; diff --git a/flake/dev.nix b/flake/dev.nix index e1e92161..61d80898 100644 --- a/flake/dev.nix +++ b/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; } diff --git a/modules/services/slurm.nix b/modules/services/slurm.nix index a1418b82..fda95ab2 100644 --- a/modules/services/slurm.nix +++ b/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; + }; }) ]); } diff --git a/packages/sbatch-tui/.clangd b/packages/sbatch-tui/.clangd deleted file mode 100644 index a6585d7e..00000000 --- a/packages/sbatch-tui/.clangd +++ /dev/null @@ -1,3 +0,0 @@ -CompileFlags: - Add: [ -Wall, -Wextra, -std=c++23 ] - Compiler: g++ diff --git a/packages/sbatch-tui/CMakeLists.txt b/packages/sbatch-tui/CMakeLists.txt index ccaac4d8..b201f6ff 100644 --- a/packages/sbatch-tui/CMakeLists.txt +++ b/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}) diff --git a/packages/sbatch-tui/include/sbatch-tui/device.hpp b/packages/sbatch-tui/include/sbatch-tui/device.hpp deleted file mode 100644 index 6604f390..00000000 --- a/packages/sbatch-tui/include/sbatch-tui/device.hpp +++ /dev/null @@ -1,10 +0,0 @@ -# pragma once -# include -# include - -struct Device_t -{ - unsigned CpuMpiThreads, CpuOpenmpThreads; - std::vector GpuIds; -}; -extern Device_t Device; diff --git a/packages/sbatch-tui/src/device.cpp b/packages/sbatch-tui/src/device.cpp deleted file mode 100644 index f4f8569b..00000000 --- a/packages/sbatch-tui/src/device.cpp +++ /dev/null @@ -1,8 +0,0 @@ -# include - -Device_t Device -{ - .CpuMpiThreads = 1, - .CpuOpenmpThreads = 1, - .GpuIds = { "4060" } -}; diff --git a/packages/sbatch-tui/src/device.cpp.template b/packages/sbatch-tui/src/device.cpp.template deleted file mode 100644 index 59340e79..00000000 --- a/packages/sbatch-tui/src/device.cpp.template +++ /dev/null @@ -1,8 +0,0 @@ -# include - -Device_t Device -{ - .CpuMpiThreads = @cpuMpiThreads@, - .CpuOpenmpThreads = @cpuOpenmpThreads@, - .GpuIds = { @gpuIds@ } -}; diff --git a/packages/sbatch-tui/src/main.cpp b/packages/sbatch-tui/src/main.cpp index d3d6b753..4c082c59 100644 --- a/packages/sbatch-tui/src/main.cpp +++ b/packages/sbatch-tui/src/main.cpp @@ -3,31 +3,44 @@ # include # include # include -# include -# include # include -using namespace biu::literals; - int main() { + using namespace biu::literals; + + struct Device + { + unsigned CpuMpiThreads, CpuOpenmpThreads; + std::optional> GpuIds; + }; + auto device = YAML::LoadFile("/etc/sbatch-tui.yaml").as(); + // 需要绑定到界面上的变量 struct { int vasp_version_selected = 0; std::vector vasp_version_entries = { "std", "gam", "ncl" }; int device_type_selected = 0; - std::vector device_type_entries = { "manually select GPU", "any single GPU", "CPU" }; - std::deque device_selected = std::deque(Device.GpuIds.size(), false); - std::vector device_entries = Device.GpuIds; + std::vector device_type_entries; + int gpu_selected = 0; + std::vector 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> 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 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, - 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;