2020-08-18 00:50:18 +08:00
|
|
|
{
|
|
|
|
stdenv
|
2020-10-05 18:33:44 +08:00
|
|
|
, slurm
|
2020-10-09 21:55:37 +08:00
|
|
|
, garlicTools
|
2020-08-18 00:50:18 +08:00
|
|
|
}:
|
2020-09-02 16:44:13 +08:00
|
|
|
{
|
2020-10-09 21:55:37 +08:00
|
|
|
nextStage
|
|
|
|
, cpuBind
|
|
|
|
, nixPrefix
|
2020-09-02 16:44:13 +08:00
|
|
|
, srunOptions ? ""
|
2020-10-22 00:18:43 +08:00
|
|
|
, output ? "stdout.log"
|
|
|
|
, error ? "stderr.log"
|
2020-09-02 16:44:13 +08:00
|
|
|
}:
|
2020-08-18 00:50:18 +08:00
|
|
|
|
2020-10-09 21:55:37 +08:00
|
|
|
with garlicTools;
|
|
|
|
|
2020-08-18 00:50:18 +08:00
|
|
|
stdenv.mkDerivation rec {
|
2020-09-02 23:07:09 +08:00
|
|
|
name = "srun";
|
2020-08-18 00:50:18 +08:00
|
|
|
phases = [ "installPhase" ];
|
2020-10-09 21:55:37 +08:00
|
|
|
preferLocalBuild = true;
|
2020-08-18 00:50:18 +08:00
|
|
|
dontPatchShebangs = true;
|
|
|
|
installPhase = ''
|
2020-09-02 23:07:09 +08:00
|
|
|
cat > $out <<EOF
|
2020-10-14 22:29:22 +08:00
|
|
|
#!/bin/sh -e
|
2020-10-09 21:55:37 +08:00
|
|
|
exec ${slurm}/bin/srun \
|
|
|
|
--mpi=pmi2 \
|
|
|
|
--cpu-bind=${cpuBind} \
|
2020-10-22 00:18:43 +08:00
|
|
|
--output=${output} \
|
|
|
|
--error=${error} \
|
2020-10-09 21:55:37 +08:00
|
|
|
${srunOptions} \
|
|
|
|
${nixPrefix}${stageProgram nextStage}
|
2020-08-18 00:50:18 +08:00
|
|
|
EOF
|
2020-09-02 23:07:09 +08:00
|
|
|
chmod +x $out
|
2020-08-18 00:50:18 +08:00
|
|
|
'';
|
|
|
|
}
|