chn/nixos
1
0
Fork 0
mirror of https://github.com/CHN-beta/nixos.git synced 2024-10-23 04:18:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

localPackages.sbatch-tui: finish

Browse Source
This commit is contained in:
陈浩南 2024-06-09 12:57:01 +08:00
parent 50da5ffa0b
commit f1c8f8f1bb
6 changed files with 131 additions and 153 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
local/pkgs/sbatch-tui/default.nix
View File

@ -1,7 +1,10 @@
{ stdenv, cmake, pkg-config, fmt, ftxui, boost, range-v3 }: stdenv.mkDerivation
{ stdenv, cmake, pkg-config, fmt, ftxui, boost, range-v3, sbatchConfig ? null, substituteAll, runCommand, lib }:
stdenv.mkDerivation
{
name = "sbatch-tui";
src = ./.;
preConfigure = lib.optionalString (sbatchConfig != null)
"cp ${substituteAll ({ src = ./src/device.cpp.template; } // sbatchConfig)} src/device.cpp";
buildInputs = [ fmt ftxui boost range-v3 ];
nativeBuildInputs = [ cmake pkg-config ];
}

2
local/pkgs/sbatch-tui/include/sbatch-tui/device.hpp
View File

@ -4,7 +4,7 @@
struct Device_t
{
unsigned CpuMpiThreads, CpuOpenMPThreads;
unsigned CpuMpiThreads, CpuOpenmpThreads;
std::vector<std::string> GpuIds;
};
extern Device_t Device;

2
local/pkgs/sbatch-tui/src/device.cpp
View File

@ -3,6 +3,6 @@
Device_t Device
{
.CpuMpiThreads = 1,
.CpuOpenMPThreads = 1,
.CpuOpenmpThreads = 1,
.GpuIds = { "4060" }
};

8
local/pkgs/sbatch-tui/src/device.cpp.template
View File

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

251
local/pkgs/sbatch-tui/src/main.cpp
View File

@ -14,144 +14,123 @@ using namespace std::literals;
int main()
{
// 需要绑定到界面上的变量
struct
{
int vasp_version_selected = 0;
std::vector<std::string> vasp_version_entries = { "std", "gam", "ncl" };
int device_selected = 0;
std::vector<std::string> device_entries = []
{
std::vector<std::string> devices(Device.GpuIds.size() + 2);
for (std::size_t i = 0; i < Device.GpuIds.size(); ++i)
devices[i + 1] = Device.GpuIds[i];
devices[0] = "any single GPU";
devices.back() = "CPU";
return devices;
}();
std::string user_command;
std::string submit_command;
} state;
// 需要绑定到界面上的变量
struct
{
int vasp_version_selected = 0;
std::vector<std::string> vasp_version_entries = { "std", "gam", "ncl" };
int device_selected = 0;
std::vector<std::string> device_entries = []
{
std::vector<std::string> devices(Device.GpuIds.size() + 2);
for (std::size_t i = 0; i < Device.GpuIds.size(); ++i)
devices[i + 1] = Device.GpuIds[i];
devices[0] = "any single GPU";
devices.back() = "CPU";
return devices;
}();
std::string user_command;
std::string submit_command;
} state;
// 为组件增加标题栏
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)
{
return [title](ftxui::Element element)
{
return ftxui::vbox
(ftxui::text(title) | ftxui::bgcolor(ftxui::Color::Blue), element, ftxui::separatorLight());
};
};
// 为组件增加空白以填充界面
auto with_padding = [](ftxui::Element element) -> ftxui::Element
{
auto empty = ftxui::emptyElement() | ftxui::flex_grow;
return ftxui::vbox(empty, ftxui::hbox(empty, element | ftxui::center, empty), empty);
};
// 构建界面, 需要至少 25 行 47 列
auto screen = ftxui::ScreenInteractive::Fullscreen();
auto request_interface = [&state, &screen, &component_with_title]
{
auto vasp_version = component_with_title
(
"Select VASP version:",
ftxui::Menu(&state.vasp_version_entries, &state.vasp_version_selected)
)
| ftxui::size(ftxui::WIDTH, ftxui::EQUAL, 30);
auto device = component_with_title
(
"Select device:",
ftxui::Menu(&state.device_entries, &state.device_selected)
)
| ftxui::size(ftxui::WIDTH, ftxui::EQUAL, 30);
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
({
vasp_version, device,
ftxui::Container::Horizontal({continue_button, quit_button})
}) | ftxui::borderHeavy
| ftxui::size(ftxui::WIDTH, ftxui::EQUAL, 30)
| ftxui::size(ftxui::HEIGHT, ftxui::EQUAL, 18);
}();
auto confirm_interface = [&state, &screen, &component_with_title]
{
ftxui::InputOption input_option;
input_option.multiline = true;
return ftxui::Container::Vertical
({
component_with_title
(
"Double check & modify submit command:",
ftxui::Input(&state.submit_command, "", 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::borderHeavy
| ftxui::size(ftxui::WIDTH, ftxui::EQUAL, 30)
| ftxui::size(ftxui::HEIGHT, ftxui::EQUAL, 18);
}();
// 构建界面
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::Menu(&state.device_entries, &state.device_selected)
| with_title("Select device:"),
ftxui::Container::Horizontal
({
ftxui::Button("Continue",
[&]{ state.user_command = "continue"; screen.ExitLoopClosure()(); }),
ftxui::Button("Quit",
[&]{ state.user_command = "quit"; screen.ExitLoopClosure()(); })
})
}) | ftxui::borderHeavy | with_padding;
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",
[&]{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;
// 实际投递任务
auto submit = [](std::string submit_command)
{
// replace \n with space
boost::replace_all(submit_command, "\n", " ");
auto process = boost::process::child
(
boost::process::search_path("sh"), "-c", submit_command,
boost::process::std_in.close(),
boost::process::std_out > stdout,
boost::process::std_err > stderr
);
process.wait();
};
// 实际投递任务
auto submit = [](std::string submit_command)
{
// replace \n with space
boost::replace_all(submit_command, "\n", " ");
auto process = boost::process::child
(
boost::process::search_path("sh"), "-c", submit_command,
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");
else if (state.device_selected < 0 || state.device_selected >= state.device_entries.size())
throw std::runtime_error("device_selected is out of range");
else if (state.device_selected == 0) state.submit_command = fmt::format
(
"sbatch --ntasks=1\n--gpus=1\n--job-name='{}'\n--output=output.txt\nvasp-nvidia-{}",
std::filesystem::current_path().filename().string(), state.vasp_version_entries[state.vasp_version_selected]
);
else if (state.device_selected == state.device_entries.size() - 1) state.submit_command = fmt::format
(
"sbatch --ntasks={}\n--cpus-per-task={}\n--hint=nomultithread\n--job-name='{}'\n--output=output.txt"
"\nvasp-intel-{}",
Device.CpuMpiThreads, Device.CpuOpenMPThreads, std::filesystem::current_path().filename().string(),
state.vasp_version_entries[state.vasp_version_selected]
);
else state.submit_command = fmt::format
(
"sbatch --ntasks=1\n--gpus={}:1\n--job-name='{}'\n--output=output.txt\nvasp-nvidia-{}",
state.device_entries[state.device_selected], std::filesystem::current_path().filename().string(),
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;
else if (state.user_command != "submit")
throw std::runtime_error("user_command is not recognized");
submit(state.submit_command);
break;
}
// 进入事件循环
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");
else if (state.device_selected < 0 || state.device_selected >= state.device_entries.size())
throw std::runtime_error("device_selected is out of range");
else if (state.device_selected == 0) state.submit_command = fmt::format
(
"sbatch --ntasks=1\n--gpus=1\n--job-name='{}'\n--output=output.txt\nvasp-nvidia-{}",
std::filesystem::current_path().filename().string(), state.vasp_version_entries[state.vasp_version_selected]
);
else if (state.device_selected == state.device_entries.size() - 1) state.submit_command = fmt::format
(
"sbatch --ntasks={}\n--cpus-per-task={}\n--hint=nomultithread\n--job-name='{}'\n--output=output.txt"
"\nvasp-intel-{}",
Device.CpuMpiThreads, Device.CpuOpenmpThreads, std::filesystem::current_path().filename().string(),
state.vasp_version_entries[state.vasp_version_selected]
);
else state.submit_command = fmt::format
(
"sbatch --ntasks=1\n--gpus={}:1\n--job-name='{}'\n--output=output.txt\nvasp-nvidia-{}",
state.device_entries[state.device_selected], std::filesystem::current_path().filename().string(),
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;
else if (state.user_command != "submit")
throw std::runtime_error("user_command is not recognized");
submit(state.submit_command);
break;
}
}

16
modules/services/slurm.nix
View File

@ -116,16 +116,12 @@ inputs:
};
nixos =
{
packages._packages = [(inputs.pkgs.localPackages.sbatch-tui.overrideAttrs (prev: { src =
let device = inputs.pkgs.substituteAll
{
src = "${prev.src}/src/device.cpp.template";
CpuMpiThreads = slurm.cpu.mpiThreads;
CpuOpenmpThreads = slurm.cpu.openmpThreads;
GpuIds = builtins.concatStringsSep ", " (builtins.map (gpu: ''"${gpu}"'') (builtins.attrNames slurm.gpus));
};
in inputs.pkgs.runCommand "src" {}
''cp -r ${prev.src} $out; chmod +w -R $out; cp ${device} $out/src/device.cpp''; }))];
packages._packages = [(inputs.pkgs.localPackages.sbatch-tui.override { sbatchConfig =
{
cpuMpiThreads = slurm.cpu.mpiThreads;
cpuOpenmpThreads = slurm.cpu.openmpThreads;
gpuIds = builtins.concatStringsSep ", " (builtins.map (gpu: ''"${gpu}"'') (builtins.attrNames slurm.gpus));
};})];
user.sharedModules = [{ home.packages =
[
(inputs.pkgs.writeShellScriptBin "sbatch"