Rodrigo Arias Mallo
b4a3bb0ede
Now the options cpusPerTask ntasksPerNode nodes and jobName are required for the sbatch stage. Also cpuBind has been removed and is always set to "cores,verbose" in the srun stage.
111 lines
3.1 KiB
Nix
111 lines
3.1 KiB
Nix
{
|
|
stdenv
|
|
, config
|
|
, stages
|
|
, targetMachine
|
|
, garlicTools
|
|
, bsc
|
|
}:
|
|
|
|
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: stages.trebuchet {
|
|
inherit (machineConf) nixPrefix sshHost;
|
|
nextStage = stages.runexp {
|
|
inherit (machineConf) nixPrefix;
|
|
nextStage = stages.isolate {
|
|
inherit (machineConf) nixPrefix;
|
|
nextStage = stages.experiment {
|
|
inherit units;
|
|
};
|
|
};
|
|
};
|
|
};
|
|
|
|
/* 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, ...}: (
|
|
assert (assertMsg (!(conf ? cpuBind))
|
|
"cpuBind is no longer available in the standard srun stage");
|
|
stages.srun {
|
|
inherit (conf) nixPrefix;
|
|
inherit nextStage;
|
|
|
|
# Binding is set to cores always
|
|
cpuBind = "cores,verbose";
|
|
}
|
|
);
|
|
|
|
isolate = {nextStage, conf, ...}: stages.isolate {
|
|
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;
|
|
}
|