From d1d27ce1943476f6567d18d309708f9feb4aa055 Mon Sep 17 00:00:00 2001 From: Haonan Chen Date: Sat, 3 Jan 2026 12:17:01 +0800 Subject: [PATCH] packages.chn-bsub: fix, install on wlin --- devices/wlin/bsub.nix | 9 + devices/wlin/default.nix | 5 +- packages/chn-bsub/.envrc | 1 + packages/chn-bsub/CMakeLists.txt | 8 +- packages/chn-bsub/bsub.yaml | 4 + packages/chn-bsub/default.nix | 3 +- packages/chn-bsub/src/main.cpp | 341 ++++++++++++++----------------- 7 files changed, 182 insertions(+), 189 deletions(-) create mode 100644 devices/wlin/bsub.nix create mode 100644 packages/chn-bsub/.envrc create mode 100644 packages/chn-bsub/bsub.yaml diff --git a/devices/wlin/bsub.nix b/devices/wlin/bsub.nix new file mode 100644 index 00000000..8f36e3e7 --- /dev/null +++ b/devices/wlin/bsub.nix @@ -0,0 +1,9 @@ +{ + normal = [ 4 4 20 ]; + normal_1day = [ 4 7 28 ]; + normal_1week = [ 4 7 28 ]; + normal_2week = [ 6 8 48 ]; + normal_1day_new = [ 4 6 24 ]; + ocean_530_1day = [ 4 6 24 ]; + ocean6226R_1day = [ 4 8 32 ]; +} diff --git a/devices/wlin/default.nix b/devices/wlin/default.nix index 8809efc3..327ab656 100644 --- a/devices/wlin/default.nix +++ b/devices/wlin/default.nix @@ -9,10 +9,13 @@ let nixpkgs = { march = "haswell"; nixRoot = "/data/gpfs01/wlin/.nix"; nixos = false; }; }); python = pkgs.python3.withPackages (ps: with ps; [ phonopy ]); + chn-bsub = pkgs.localPackages.chn-bsub.overrideAttrs + (prev: { bsubConfig = builtins.toFile "bsub.yaml" (builtins.toJSON (import ./bsub.nix)); }); wlin = pkgs.symlinkJoin { name = "wlin"; - paths = with pkgs; [ gnuplot localPackages.vaspkit pv python localPackages.vasp.intel ]; + paths = with pkgs; + [ gnuplot localPackages.vaspkit pv python localPackages.vasp.intel chn-bsub hwloc ]; postBuild = "echo ${inputs.self.rev or "dirty"} > $out/.version"; passthru = { inherit pkgs; archive = pkgs.closureInfo { rootPaths = [ wlin.drvPath ]; }; }; }; diff --git a/packages/chn-bsub/.envrc b/packages/chn-bsub/.envrc new file mode 100644 index 00000000..eb9f01ea --- /dev/null +++ b/packages/chn-bsub/.envrc @@ -0,0 +1 @@ +use flake .#chn-bsub diff --git a/packages/chn-bsub/CMakeLists.txt b/packages/chn-bsub/CMakeLists.txt index c5c37c7f..dc3c355e 100644 --- a/packages/chn-bsub/CMakeLists.txt +++ b/packages/chn-bsub/CMakeLists.txt @@ -10,14 +10,14 @@ if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) endif() find_package(ftxui REQUIRED) -find_package(Boost REQUIRED COMPONENTS filesystem iostreams) -find_package(range-v3 REQUIRED) find_package(biu REQUIRED) add_executable(chn-bsub src/main.cpp) target_compile_features(chn-bsub PUBLIC cxx_std_23) -target_link_libraries(chn-bsub PRIVATE fmt::fmt ftxui::screen ftxui::dom ftxui::component Boost::filesystem - range-v3::range-v3 biu::biu) +target_link_libraries(chn-bsub PRIVATE ftxui::screen ftxui::dom ftxui::component biu::biu) +if(BSUB_CONFIG) + target_compile_definitions(chn-bsub PRIVATE BSUB_CONFIG="${BSUB_CONFIG}") +endif() install(TARGETS chn-bsub RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/packages/chn-bsub/bsub.yaml b/packages/chn-bsub/bsub.yaml new file mode 100644 index 00000000..4f6faab4 --- /dev/null +++ b/packages/chn-bsub/bsub.yaml @@ -0,0 +1,4 @@ +Queues: + normal: [ 4, 4, 20 ] + normal_1day: [ 4, 7, 28 ] + diff --git a/packages/chn-bsub/default.nix b/packages/chn-bsub/default.nix index 5aeb425e..61f2a34e 100644 --- a/packages/chn-bsub/default.nix +++ b/packages/chn-bsub/default.nix @@ -1,8 +1,9 @@ -{ stdenv, lib, cmake, pkg-config, ftxui, biu }: stdenv.mkDerivation +{ stdenv, cmake, pkg-config, ftxui, biu, bsubConfig ? null, lib }: stdenv.mkDerivation { name = "chn-bsub"; src = ./.; buildInputs = [ ftxui biu ]; nativeBuildInputs = [ cmake pkg-config ]; postInstall = "ln -s chn-bsub $out/bin/chn_bsub"; + cmakeFlags = lib.optional (bsubConfig != null) [ "-DBSUB_CONFIG=${bsubConfig}" ]; } diff --git a/packages/chn-bsub/src/main.cpp b/packages/chn-bsub/src/main.cpp index 3b9f00ff..b25cdfae 100644 --- a/packages/chn-bsub/src/main.cpp +++ b/packages/chn-bsub/src/main.cpp @@ -1,202 +1,177 @@ -# include -# include +# include # include # include # include # include # include -# include + +# ifndef BSUB_CONFIG +# define BSUB_CONFIG "./bsub.yaml" +# endif using namespace biu::literals; int main() { - // 需要绑定到界面上的变量 - struct - { - std::array vasp_version_selected = {0, 0, 0}; - std::vector vasp_version_entries_level1 = {"640", "631"}; - std::map> vasp_version_entries_level2 = + biu::Logger::Guard log; + enum class UserCommandType { Continue, Back, Quit }; + enum class InterfaceType { Request, Confirm }; + struct + { + std::optional UserCommand; + std::string SubmitCommand; + InterfaceType CurrentInterface = InterfaceType::Request; + int VaspSelected = 0; + std::vector VaspEntries = { "std", "gam", "ncl" }; + int QueueSelected = 0; + std::vector QueueEntries; + std::string JobName = [] { - {"640", {"(default)", "fixc", "optcell_vtst_wannier90", "shmem", "vtst"}}, - {"631", {"shmem"}} - }; - std::vector vasp_version_entries_level3 = {"std", "gam", "ncl"}; - - int queue_selected = 0; - std::vector queue_entries = - { - "normal_1day", "normal_1week", "normal", - "normal_1day_new", "ocean_530_1day", "ocean6226R_1day" - }; - std::map max_cores = - { - {"normal_1day", 28}, {"normal_1week", 28}, {"normal", 20}, - {"normal_1day_new", 24}, {"ocean_530_1day", 24}, {"ocean6226R_1day", 32} - }; - std::string ncores = ""; - std::string job_name = [] - { - // /data/gpfs01/jykang/linwei/chn/lammps-SiC + // /data/gpfs01/wlin/chn/lammps-SiC std::vector paths; - boost::split(paths, std::filesystem::current_path().string(), - boost::is_any_of("/")); - if (paths.size() < 7) - return "my-great-job"s; - else - return paths[5] + "_" + paths.back(); + boost::split(paths, std::filesystem::current_path().string(), boost::is_any_of("/")); + if (paths.size() < 6) return "my-great-job"s; + else return paths[4] + "_" + paths.back(); }(); - std::string bsub = ""; - std::string user_command = ""; - } state; + std::string OutputFile = "output.txt"; + } State; + struct { std::map> Queues; } QueueConfig = + YAML::LoadFile(BSUB_CONFIG).as(); + State.QueueEntries = QueueConfig.Queues + | ranges::views::transform([](auto const& item) { return item.first; }) + | ranges::to_vector; - // 为组件增加标题栏 - auto component_with_title = [](std::string title, ftxui::Component component) - { - return ftxui::Renderer(component, [title, component] - { - return ftxui::vbox - ({ - ftxui::text(title) | ftxui::bgcolor(ftxui::Color::Blue), - component->Render(), - ftxui::separator() - }); - }); - }; + // 为组件增加标题栏 + auto with_title = [](std::string title, ftxui::Color bgcolor = ftxui::Color::Blue) + { + return [=](ftxui::Element element) + { return ftxui::vbox(ftxui::text(title) | ftxui::bgcolor(bgcolor), element); }; + }; + // 为组件增加下边框 + auto with_bottom = [](ftxui::Element element) + { return ftxui::vbox(element, ftxui::separatorLight()); }; + // 为组件增加比较粗的下边框 + auto with_bottom_heavy = [](ftxui::Element element) + { return ftxui::vbox(element, ftxui::separatorHeavy()); }; + // 在组件左边增加小标题 + auto with_subtitle = [](std::string title) + { return [title](ftxui::Element element) { return ftxui::hbox(ftxui::text(title), element); }; }; + // 带标题的文本输入框 + auto input = [with_subtitle](std::string* content, std::string title) + { + return ftxui::Input(content) | ftxui::underlined + | ftxui::size(ftxui::WIDTH, ftxui::GREATER_THAN, 3) + | ftxui::size(ftxui::HEIGHT, ftxui::EQUAL, 1) + | with_subtitle(title); + }; + // 在组件左边增加分割线 + auto with_separator = [](ftxui::Element element) + { return ftxui::hbox(ftxui::separatorLight(), element); }; + // 为组件增加空白以填充界面 + auto with_padding = [](ftxui::Element element) + { + auto empty = ftxui::emptyElement() | ftxui::flex_grow; + return ftxui::vbox(empty, ftxui::hbox(empty, element | ftxui::center, empty), empty); + }; + // 转义字符 + auto escape = [](std::string str) + { + return str | ranges::views::transform([](char c) + { + // only the following characters need to be escaped: $ ` \ " newline * @ space tab + if (std::set{'$', '`', '\\', '\"', '\n', '*', '@', ' ', '\t'}.contains(c)) + return '\\' + std::string(1, c); + else return std::string(1, c); + }) | ranges::views::join("") | ranges::to; + }; - // 构建界面, 需要至少 25 行 47 列 - auto screen = ftxui::ScreenInteractive::Fullscreen(); - auto request_interface = [&state, &screen, &component_with_title] - { - auto vasp_version_level1 = ftxui::Menu - (&state.vasp_version_entries_level1, &state.vasp_version_selected[0]) - | ftxui::size(ftxui::WIDTH, ftxui::EQUAL, 8); - std::vector vasp_version_level2_children; - for (auto& i : state.vasp_version_entries_level1) - vasp_version_level2_children.push_back(ftxui::Menu - ( - &state.vasp_version_entries_level2[i], - &state.vasp_version_selected[1] - )); - auto vasp_version_level2 = ftxui::Container::Tab - ( - vasp_version_level2_children, - &state.vasp_version_selected[0] - ) | ftxui::size(ftxui::WIDTH, ftxui::EQUAL, 27); - auto vasp_version_level3 = ftxui::Menu - (&state.vasp_version_entries_level3, &state.vasp_version_selected[2]) - | ftxui::size(ftxui::WIDTH, ftxui::EQUAL, 8); - auto vasp_version = component_with_title("Select vasp version:", - ftxui::Container::Horizontal - ({vasp_version_level1, vasp_version_level2, vasp_version_level3}) - | ftxui::size(ftxui::HEIGHT, ftxui::EQUAL, 5)); - auto queue = component_with_title("Select queue:", - ftxui::Menu(&state.queue_entries, &state.queue_selected) - | ftxui::size(ftxui::HEIGHT, ftxui::EQUAL, 6)); - auto ncores = component_with_title("Input cores you want to use:", - ftxui::Input(&state.ncores, "(leave blank to use all cores)")) - | ftxui::size(ftxui::HEIGHT, ftxui::EQUAL, 3); - auto job_name = component_with_title("Job name:", - ftxui::Input(&state.job_name, "")) - | ftxui::size(ftxui::HEIGHT, ftxui::EQUAL, 3); - auto continue_button = ftxui::Button("Continue", - [&]{state.user_command = "continue"; screen.ExitLoopClosure()();}); - auto quit_button = ftxui::Button("Quit", - [&]{state.user_command = "quit"; screen.ExitLoopClosure()();}); - return ftxui::Container::Vertical + // 构建界面 + auto Screen = ftxui::ScreenInteractive::Fullscreen(); + auto key_event_handler = [&](ftxui::Event event) + { + if (event == ftxui::Event::Return) + { State.UserCommand = UserCommandType::Continue; Screen.ExitLoopClosure()(); return true; } + else return false; + }; + auto InterfaceRequest = ftxui::Container::Vertical + ({ + ftxui::Container::Horizontal ({ - vasp_version, queue, ncores, job_name, - ftxui::Container::Horizontal({continue_button, quit_button}) - }) | ftxui::borderHeavy - | ftxui::size(ftxui::WIDTH, ftxui::EQUAL, 47) - | ftxui::size(ftxui::HEIGHT, ftxui::EQUAL, 24); - }(); - auto confirm_interface = [&state, &screen, &component_with_title] - { - ftxui::InputOption input_option; - input_option.multiline = true; - return ftxui::Container::Vertical + ftxui::Menu(&State.VaspEntries, &State.VaspSelected) | with_title("VASP variant:"), + ftxui::Menu(&State.QueueEntries, &State.QueueSelected) + | with_title("Queue:") | with_separator + }) | with_bottom_heavy, + input(&State.JobName, "Job name: "), + input(&State.OutputFile, "Output file: "), + // 操作按钮 + ftxui::Container::Horizontal ({ - component_with_title - ( - "Double check & modify submit command:", - ftxui::Input(&state.bsub, "", input_option) - ) - | ftxui::size(ftxui::HEIGHT, ftxui::EQUAL, 7), - ftxui::Container::Horizontal - ({ - ftxui::Button("Submit", - [&]{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::Renderer([]{return ftxui::vbox - ({ - ftxui::separator(), - ftxui::text("Source code:"), - ftxui::text("https://github.com/CHN-beta/chn_bsub.git"), - ftxui::text("Star & PR are welcome!"), - });}) - }) | ftxui::borderHeavy - | ftxui::size(ftxui::WIDTH, ftxui::EQUAL, 47) - | ftxui::size(ftxui::HEIGHT, ftxui::EQUAL, 14); - }(); + ftxui::Button("Continue (Enter)", + [&]{ State.UserCommand = UserCommandType::Continue; Screen.ExitLoopClosure()(); }), + ftxui::Button("Quit", + [&]{ State.UserCommand = UserCommandType::Quit; Screen.ExitLoopClosure()(); }) + }), + }) | ftxui::borderHeavy | with_padding | ftxui::CatchEvent(key_event_handler); + auto InterfaceConfirm = ftxui::Container::Vertical + ({ + ftxui::Input(&State.SubmitCommand, "", ftxui::InputOption{.multiline = true}) + | with_title("Double check & modify submit command:") | with_bottom_heavy, + ftxui::Container::Horizontal + ({ + ftxui::Button("Submit (Enter)", + [&]{State.UserCommand = UserCommandType::Continue; Screen.ExitLoopClosure()();}), + ftxui::Button("Back", + [&]{State.UserCommand = UserCommandType::Back; Screen.ExitLoopClosure()();}), + ftxui::Button("Quit", + [&]{State.UserCommand = UserCommandType::Quit; Screen.ExitLoopClosure()();}) + }) + }) | ftxui::borderHeavy | with_padding | ftxui::CatchEvent(key_event_handler); - // 实际投递任务 - auto submit = [](std::string bsub) - { - // replace \n with space - boost::replace_all(bsub, "\n", " "); - auto process = boost::process::child - ( - boost::process::search_path("sh"), "-c", bsub, - boost::process::std_in.close(), - boost::process::std_out > stdout, - boost::process::std_err > stderr - ); - process.wait(); - }; - - // 进入事件循环 - while (true) - { - screen.Loop(request_interface); - if (state.user_command == "quit") - return EXIT_FAILURE; - else if (state.user_command != "continue") - throw std::runtime_error("user_command is not recognized"); - state.bsub = fmt::format - ( - "bsub -J '{}'\n-q {}\n-n {}\n-R 'span[hosts=1]'\n-o 'output.txt'\nchn_vasp.sh {}", - state.job_name, - state.queue_entries[state.queue_selected], - state.ncores.empty() ? state.max_cores[state.queue_entries[state.queue_selected]] : - std::stoi(state.ncores), - [&] - { - auto version_level1 = state.vasp_version_entries_level1[state.vasp_version_selected[0]]; - auto version_level2 = state.vasp_version_entries_level2[version_level1] - [state.vasp_version_selected[1]]; - auto version_level3 = state.vasp_version_entries_level3[state.vasp_version_selected[2]]; - return fmt::format - ( - "{}{}_{}", - version_level1, - version_level2 == "(default)" ? ""s : "_" + version_level2, - version_level3 - ); - }() - ); - screen.Loop(confirm_interface); - if (state.user_command == "quit") - return EXIT_FAILURE; - else if (state.user_command == "back") - continue; - else if (state.user_command != "submit") - throw std::runtime_error("user_command is not recognized"); - submit(state.bsub); - break; - } + // 进入事件循环 + while (true) + { + if (State.CurrentInterface == InterfaceType::Request) + { + State.UserCommand.reset(); + Screen.Loop(InterfaceRequest); + if (State.UserCommand == UserCommandType::Quit) return 0; + else if (State.UserCommand == UserCommandType::Continue) + { + State.CurrentInterface = InterfaceType::Confirm; + State.SubmitCommand = [&] + { + auto [nproc, nthr, ncpu] = QueueConfig.Queues.at(State.QueueEntries[State.QueueSelected]); + auto args = std::vector + { + "bsub", + "-J {} -o {}"_f(escape(State.JobName), escape(State.OutputFile)), + "-q {} -n {} -R 'span[hosts=1]'"_f(escape(State.QueueEntries[State.QueueSelected]), ncpu), + "vasp-intel mpirun -n {} -x OMP_NUM_THREADS={} vasp-{}"_f + (nproc, nthr, State.VaspEntries[State.VaspSelected]) + }; + return args | ranges::views::join(" \\\n ") | ranges::to; + }(); + } + else if (!State.UserCommand) return EXIT_FAILURE; + else std::unreachable(); + } + else if (State.CurrentInterface == InterfaceType::Confirm) + { + State.UserCommand.reset(); + Screen.Loop(InterfaceConfirm); + if (State.UserCommand == UserCommandType::Quit) return 0; + else if (State.UserCommand == UserCommandType::Back) { State.CurrentInterface = InterfaceType::Request; } + else if (State.UserCommand == UserCommandType::Continue) + { + // 提交任务 + log.debug("submit command: {}"_f(State.SubmitCommand)); + // -c 对 \\n 的处理与通常情况下不同，我们需要用 -s 然后将命令通过标准输入传入 + biu::exec<{.SearchPath = true, .Stdin = biu::IoType::String}> + ({.Program = "sh", .Args = { "-s"}, .Stdin = State.SubmitCommand}); + break; + } + else if (!State.UserCommand) return EXIT_FAILURE; + else std::unreachable(); + } + } }