bscpkgs/garlic/stdexp.nix
Rodrigo Arias Mallo 872ad1a289 stdexp: allow preSrun attribute in the srun stage
This option allows an experiment to inject commands before srun starts,
while keeping the standard srun stage options.
2021-03-29 17:46:19 +02:00

145 lines
3.9 KiB
Nix

{
stdenv
, config
, stages
, targetMachine
, garlicTools
, bsc
, writeTextFile
, runCommandLocal
, python
, pp
}:
with stdenv.lib;
with garlicTools;
let
machineConf = targetMachine.config;
in
rec {
/* Takes a list of units and builds an experiment, after executing the
trebuchet, runexp and isolate stages. Returns the trebuchet stage. */
buildTrebuchet = units:
let
trebuchet = stages.trebuchet {
inherit (machineConf) nixPrefix sshHost;
nextStage = stages.runexp {
inherit (machineConf) nixPrefix;
nextStage = stages.isolate {
inherit (machineConf) nixPrefix;
nextStage = stages.experiment {
inherit units;
};
};
};
};
in trebuchet // rec {
result = pp.store {
trebuchet=trebuchet;
experiment=trebuchet.experiment;
};
timetable = pp.timetable result;
};
/* Given an attrset of lists `varConf` and a function `genConf` that accepts a
attrset, computes the cartesian product of all combinations of `varConf` calls
genConf to produce the final list of configurations. */
buildConfigs = {varConf, genConf}:
map (c: genConf c) (genConfigs varConf);
stdStages = {
sbatch = {nextStage, conf, ...}: with conf; stages.sbatch (
# Allow a user to define a custom reservation for the job in MareNostrum4,
# by setting the garlic.sbatch.reservation attribute in the
# ~/.config/nixpkgs/config.nix file. If the attribute is not set, no
# reservation is used. The user reservation may be overwritten by the
# experiment, if the reservation is set like with nodes or ntasksPerNode.
optionalAttrs (config ? garlic.sbatch.reservation) {
inherit (config.garlic.sbatch) reservation;
} //
# However, if the experiment contains a reservation, that takes priority
# over the one set in the ~/.config/nixpkgs/config.nix file. Add other
# options if they are defined as well.
optionalInherit [ "reservation" "time" "qos" ] conf //
# Finally, add all the other required parameters
{
inherit nextStage nixPrefix;
# These sbatch options are mandatory
inherit cpusPerTask ntasksPerNode nodes jobName;
exclusive = true;
}
);
control = {nextStage, conf, ...}: stages.control {
inherit (conf) loops;
inherit nextStage;
};
srun = {nextStage, conf, preSrun ? "", ...}: (
assert (assertMsg (!(conf ? cpuBind))
"cpuBind is no longer available in the standard srun stage");
stages.srun {
inherit (conf) nixPrefix;
inherit nextStage preSrun;
# Binding is set to cores always
cpuBind = "cores,verbose";
}
);
isolate = {nextStage, conf, ...}: stages.isolate (
(
if (conf ? extraMounts) then { inherit (conf) extraMounts; }
else {}
) //
{
clusterName = machineConf.name;
inherit (conf) nixPrefix;
inherit nextStage;
}
);
};
stdPipelineOverride = {overrides ? {}}:
let
stages = stdStages // overrides;
in
with stages; [ sbatch isolate control srun isolate ];
stdPipeline = stdPipelineOverride {};
replaceMpi = customMpi: bsc.extend (self: super: {
mpi = customMpi;
});
# Generate the experimental units
genUnits = {configs, pipeline}: map (c: stages.unit {
conf = c;
stages = pipeline;
}) configs;
# Generate the complete experiment
genExperiment = {configs, pipeline}:
let
units = genUnits { inherit configs pipeline; };
in
buildTrebuchet units;
# Runs a python script and the standard output is directly imported as
# nix code
printPython = code:
let
p = writeTextFile {
name = "python-script";
text = ''
from math import *
${code}
'';
};
in
import (runCommandLocal "a" { buildInputs = [ python ]; } ''
python ${p} > $out'');
}