packages.chn-bsub: fix, install on wlin

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 ];
+}

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 ]; }; };
   };

packages/chn-bsub/.envrc

@@ -0,0 +1 @@
+use flake .#chn-bsub

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})
 

packages/chn-bsub/bsub.yaml

@@ -0,0 +1,4 @@
+Queues:
+  normal: [ 4, 4, 20 ]
+  normal_1day: [ 4, 7, 28 ]
+

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}" ];
 }

packages/chn-bsub/src/main.cpp

@@ -1,202 +1,177 @@
-# include <map>
-# include <filesystem>
+# include <biu.hpp>
 # include <ftxui/component/component.hpp>
 # include <ftxui/component/component_options.hpp>
 # include <ftxui/component/screen_interactive.hpp>
 # include <boost/process.hpp>
 # include <boost/algorithm/string.hpp>
-# include <biu.hpp>
+
+# ifndef BSUB_CONFIG
+#   define BSUB_CONFIG "./bsub.yaml"
+# endif
 
 using namespace biu::literals;
 
 int main()
 {
-	// 需要绑定到界面上的变量
-	struct
-	{
-		std::array<int, 3> vasp_version_selected = {0, 0, 0};
-		std::vector<std::string> vasp_version_entries_level1 = {"640", "631"};
-		std::map<std::string, std::vector<std::string>> vasp_version_entries_level2 = 
+	biu::Logger::Guard log;
+	enum class UserCommandType { Continue, Back, Quit };
+  enum class InterfaceType { Request, Confirm };
+  struct
+  {
+    std::optional<UserCommandType> UserCommand;
+    std::string SubmitCommand;
+    InterfaceType CurrentInterface = InterfaceType::Request;
+		int VaspSelected = 0;
+		std::vector<std::string> VaspEntries = { "std", "gam", "ncl" };
+		int QueueSelected = 0;
+		std::vector<std::string> QueueEntries;
+    std::string JobName = []
 		{
-			{"640", {"(default)", "fixc", "optcell_vtst_wannier90", "shmem", "vtst"}},
-			{"631", {"shmem"}}
-		};
-		std::vector<std::string> vasp_version_entries_level3 = {"std", "gam", "ncl"};
-
-		int queue_selected = 0;
-		std::vector<std::string> queue_entries =
-		{
-			"normal_1day", "normal_1week", "normal",
-			"normal_1day_new", "ocean_530_1day", "ocean6226R_1day"
-		};
-		std::map<std::string, std::size_t> 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<std::string> 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<std::string, std::array<int, 3>> Queues; } QueueConfig =
+		YAML::LoadFile(BSUB_CONFIG).as<decltype(QueueConfig)>();
+	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<std::string>;
+  };
 
-	// 构建界面, 需要至少 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<ftxui::Component> 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<std::string>
+					{
+						"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<std::string>;
+				}();
+      }
+      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();
+    }
+  }
 }