forked from rarias/bscpkgs
117 lines
3.3 KiB
Nix
117 lines
3.3 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 (
|
|
(
|
|
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;
|
|
}
|