This commit is contained in:
2025-01-06 19:58:10 +08:00
parent 8bd7746ca6
commit 783f5f6754
4 changed files with 27 additions and 5 deletions

View File

@@ -125,6 +125,9 @@ inputs:
# enable task plugins
TaskPlugin=task/affinity,task/cgroup
# omit --mpi=pmix
MpiDefault=pmix
'';
extraConfigPaths =
let gpus = slurm.node.${inputs.config.nixos.model.hostname}.gpus or null;
@@ -157,6 +160,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)

View File

@@ -128,19 +128,17 @@ int main()
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
"sbatch --ntasks=1\n--gpus=1\n--job-name='{}'\n--output='{}'\n--wrap=\"vasp-nvidia mpirun vasp-{}\""_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
"sbatch --ntasks=1\n--gres=gpu:{}:1\n--job-name='{}'\n--output='{}'\n--wrap=\"vasp-nvidia mpirun vasp-{}\""_f
(
state.gpu_entries[state.gpu_selected],
state.job_name, state.output_file, state.vasp_version_entries[state.vasp_version_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
"sbatch --ntasks={}\n--cpus-per-task={}\n--job-name='{}'\n--output='{}'\n--wrap=\"vasp-intel srun 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]

View File

@@ -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 "$@"
''

View File

@@ -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} "$@"
''