diff --git a/changelog.md b/changelog.md index 273ef19..9e75f40 100644 --- a/changelog.md +++ b/changelog.md @@ -2,6 +2,7 @@ ### [Latest] +- Add container image documentation and generic Slurm/HTCondor submission scripts with config-driven component enumeration [#160](https://github.com/umami-hep/umami-preprocessing/pull/160) ### [v0.3.2](https://github.com/umami-hep/umami-preprocessing/releases/tag/v0.3.2) (04.08.2026) - Add new configs for central dataset [#159](https://github.com/umami-hep/umami-preprocessing/pull/159) diff --git a/docs/hpc.md b/docs/hpc.md new file mode 100644 index 0000000..f96f1fe --- /dev/null +++ b/docs/hpc.md @@ -0,0 +1,207 @@ +# Running on HPC clusters + +The preprocessing stages can be parallelized over components, regions and splits using the +`--component`, `--region` and `--split` flags described in [Run](run.md). On a Slurm or HTCondor +cluster, each of these units of work can run as its own batch job inside the +[container image](setup.md#container-image). UPP ships a small set of submission scripts in +`scripts/batch/` that automate this. + +## Prerequisites + +- A cluster with Slurm or HTCondor and apptainer. On lxplus, HTCondor is provided by the CERN + batch service — see the [quickstart](https://batchdocs.web.cern.ch/local/quick.html) and + [job submission](https://batchdocs.web.cern.ch/local/submit.html) documentation. +- A shared filesystem between the submitting node and the workers — the input, intermediate and + output files must be visible to all jobs. This is the case on lxplus (AFS/EOS) and typical + institute clusters. + +!!!info "Input/output data on lxplus (AFS/EOS)" + + HTCondor on lxplus rejects submit files that reference EOS paths (executable, `output`, + `error`, `log`), so create the run directory — with the copied scripts and the `logs/` + directory — in your AFS work area and submit from there. The jobs themselves run with your + Kerberos credentials and can read and write `/afs` and `/eos` directly, so keep the large + input ntuples and outputs on EOS and bind both filesystems into the container: + + ```bash + export UPP_BINDS=/afs,/eos,/tmp + ``` + + For very I/O-heavy workflows the batch service recommends staging data through the local + pool space of the job instead of writing to EOS directly — see + [Data flows](https://batchdocs.web.cern.ch/concepts/dataflow.html) and + [EOS](https://batchdocs.web.cern.ch/troubleshooting/eos.html) in the CERN batch docs. + + If the run directory has to live on EOS, the experimental + [EosSubmit schedds](https://batchdocs.web.cern.ch/local/eossubmit.html) + (`module load lxbatch/eossubmit`) accept submit files with EOS paths, transferring all job + files via xrootd instead of using a shared filesystem. All submit file paths must then be on + EOS, and this mode has not been tested with these scripts. +- The UPP container image (see [Container image](setup.md#container-image)). The scripts default to + the CVMFS-unpacked image + `/cvmfs/unpacked.cern.ch/gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest` when + it exists (no download or conversion needed) and fall back to + `docker://gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest` otherwise. Running + from `docker://` directly is fine: apptainer caches the converted image and only downloads again + when a new version is published. Alternatively, pin a specific local file with + `export UPP_IMAGE=/path/to/upp.sif` after an `apptainer pull`. + +!!!info "Apptainer cache location" + + The apptainer cache defaults to `~/.apptainer/cache` and the conversion uses `/tmp` as + scratch space. On clusters with a small home quota (e.g. lxplus) point them to a larger + filesystem: + + ```bash + export APPTAINER_CACHEDIR=/path/to/big/storage/apptainer_cache + export APPTAINER_TMPDIR=/path/to/big/storage/apptainer_tmp + ``` + +## Interactive use + +For quick tests, run UPP inside the container on an interactive allocation: + +```bash +salloc --ntasks 1 --cpus-per-task 4 --time 2:00:00 +srun apptainer exec --contain --pwd "$PWD" -B "$PWD" -B /home -B /tmp \ + "$UPP_IMAGE" preprocess --config --prep +``` + +## Batch submission scripts + +The `scripts/batch/` directory contains: + +- `submit.sh` runs on the login node. It reads the components from your preprocessing config and + submits one batch job per unit of work. The scheduler is auto-detected (`sbatch` found → Slurm, + `condor_submit` found → HTCondor) and can be forced with `--scheduler slurm|condor`. +- `slurm_batch.sh` is the sbatch payload. It carries the `#SBATCH` resource header and starts the + container on the compute node. +- `condor_job.sub` and `condor_batch.sh` are the HTCondor equivalents: the submit description with + the resource requests, and the job executable starting the container. +- `run_stage.sh` runs inside the container and maps the submitted mode onto the `preprocess` + command line flags. + +To use them, create a run directory and copy the scripts. A clone of the repository is not +required — the image contains the repository at `/workspace`, so the scripts can be taken straight +from there: + +=== "CVMFS" + + ```bash + mkdir my_preprocessing && cd my_preprocessing + cp -r /cvmfs/unpacked.cern.ch/gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest/workspace/scripts/batch . + ``` + +=== "apptainer" + + ```bash + mkdir my_preprocessing && cd my_preprocessing + apptainer exec docker://gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest \ + cp -r /workspace/scripts/batch . + ``` + +=== "local clone" + + ```bash + mkdir my_preprocessing && cd my_preprocessing + cp -r /scripts/batch . + ``` + +Adapt the resource specifications (number of CPUs, memory allocation, time limit, partition or +account etc.) in the `slurm_batch.sh` (Slurm) or `condor_job.sub` (HTCondor) files to fit your +needs. The `+JobFlavour` walltime flavours available on lxplus are listed in the +[CERN batch documentation](https://batchdocs.web.cern.ch/local/submit.html). + +Then submit the stages in order, waiting for all jobs of one stage to finish before submitting the +next: + +```bash +./batch/submit.sh --config --dry-run prepare # preview only +./batch/submit.sh --config prepare +./batch/submit.sh --config resampling +./batch/submit.sh --config merge +./batch/submit.sh --config normalise +./batch/submit.sh --config plotting +``` + +Job logs are written to the `logs/output/` and `logs/error/` subdirectories of the current +directory. Running `submit.sh` without a mode enters +an interactive prompt for the mode and filters, and `./batch/submit.sh --help` prints all modes and +options. + +On Slurm each job is submitted with its own `sbatch` call and job name. On HTCondor all jobs of one +`submit.sh` invocation are submitted as a single cluster (one `condor_submit` with one process per +job), with the job arguments written to `logs/condor_.args`. + +The available modes and the jobs they submit: + +| Mode | Jobs | `preprocess` flags per job | +|------|------|----------------------------| +| `sequential` | 1 | full chain (`--prep`, `--resample`, `--merge`, `--norm`, `--plot`) | +| `prepare` | one per component and split | `--prep --component --split ` | +| `resampling` | one per region and split | `--resample --region --split ` | +| `fine_resampling` | one per component and split | `--resample --region --component --split ` | +| `merge` | one per split | `--merge --split ` | +| `normalise` | 1 | `--norm` | +| `plotting` | one per split | `--plot --split ` | + +!!!warning "Stage ordering and parallel h5py access" + + All jobs of a stage must finish before the next stage is submitted, e.g. all `prepare` jobs + before `resampling`. Also run the [initial sample check](run.md#additional-scripts-initial-sample-check) + once before submitting `prepare` jobs in parallel — it creates the virtual datasets which can + get corrupted when created by multiple jobs at once. + +## Config-driven job lists + +`submit.sh` never hardcodes which components exist. It calls the `list_components` script (part of +UPP) to enumerate the components defined in the `components:` block of your config. No local UPP +installation is needed for this: when `list_components` is not on the `PATH`, it is run inside the +container image automatically. Only the `prepare`, `fine_resampling` and `resampling` modes (and +the interactive mode) enumerate at all — the other modes submit without running UPP on the login +node. + +```bash +list_components --config +``` + +```text +lowpt ttbar bjets lowpt_ttbar_bjets +highpt zprime bjets highpt_zprime_bjets +... +``` + +Only combinations actually defined in the config are submitted. The selection can be narrowed with +filter flags, each taking a comma- or space-separated list: + +```bash +./batch/submit.sh --config --regions lowpt --splits train prepare +./batch/submit.sh --config --samples ttbar --flavs "bjets,cjets" fine_resampling +``` + +Note that enumerating the components fully validates the config, so a broken config fails directly +on the login node instead of inside the batch jobs. + +## Environment variables + +| Variable | Default | Purpose | +|----------|---------|---------| +| `UPP_IMAGE` | CVMFS-unpacked image if present, else `docker://...upp-images/upp:latest` | Container image (unpacked directory, local `.sif` or `docker://` URI) | +| `UPP_BINDS` | `/home,/tmp` | Comma-separated paths bound into the container (the run and script directories are always bound in addition) | +| `THROTTLE` | `30` | Seconds between `sbatch` calls (`0` disables; Slurm only) | +| `DRY_RUN` | `0` | Set to `1` to print the submission commands instead of submitting | + +Make sure `UPP_BINDS` covers your input ntuples and output directory if they live outside `/home` +(e.g. on a scratch filesystem), and export `UPP_IMAGE`/`UPP_BINDS` in your shell so they are also +picked up by the batch jobs. + +!!!warning "Keep the throttle enabled" + + The delay between `sbatch` calls avoids hammering the scheduler and gives jobs time to start + up without all of them hitting the shared filesystem at once. Only disable it for small + submissions. + +!!!info "Configs outside the repository" + + When you copy a config out of the repository, `!include` directives with relative paths no + longer resolve. Use absolute paths in `!include` lines of copied configs. diff --git a/docs/setup.md b/docs/setup.md index c9e04e8..bd8f2ae 100644 --- a/docs/setup.md +++ b/docs/setup.md @@ -128,6 +128,72 @@ python -m pip install . ``` +### Container image + +If you don't want to set up a Python environment at all, you can use the UPP container image. +The CI builds `gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest` on every merge to +`main` and a tagged image `upp:` (e.g. `upp:v0.3.1`) for every release. The image comes with UPP +and its command line scripts (`preprocess`, `check_input_samples`, `list_components`) pre-installed. + +=== "apptainer" + + On clusters (lxplus, HPC sites), apptainer can run the image directly from the registry: + + ```bash + apptainer exec docker://gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest \ + preprocess --config + ``` + + The first `docker://` invocation converts the image to apptainer's SIF format, which takes a + while. The result is cached (default `~/.apptainer/cache`), so later invocations start quickly + and only download again when a new image version is published. If your home quota is small + (e.g. on lxplus), point the cache and the conversion scratch space to a larger filesystem: + + ```bash + export APPTAINER_CACHEDIR=/path/to/big/storage/apptainer_cache + export APPTAINER_TMPDIR=/path/to/big/storage/apptainer_tmp + ``` + + On machines with CVMFS (lxplus, most HPC sites) the image is also distributed pre-unpacked + via [unpacked.cern.ch](https://gitlab.cern.ch/unpacked/sync), which skips the pull and + conversion entirely: + + ```bash + apptainer exec /cvmfs/unpacked.cern.ch/gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest \ + preprocess --config + ``` + + The release tags are available there as well (e.g. `upp:v0.3.2`). + + By default apptainer shares your home directory and working directory with the container. For a + cleaner environment use `--contain` and bind only the paths you need (your input ntuples and + output directory) with `-B`, keeping the working directory with `--pwd`: + + ```bash + apptainer exec --contain --pwd "$PWD" -B "$PWD" -B /home -B /tmp -B \ + preprocess --config + ``` + + where `` is any of the above (`docker://` URL, CVMFS path or local `.sif` file). + +=== "docker" + + With docker, mount your working directory and data paths into the container: + + ```bash + docker run --rm -it -v $PWD:$PWD -w $PWD \ + gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest \ + preprocess --config + ``` + +!!!info "Pin a release tag for production" + + `upp:latest` follows the `main` branch and changes over time. For reproducible production + preprocessing, use a tagged release image like `upp:v0.3.1` instead. + +For running UPP as batch jobs on Slurm or HTCondor clusters with the container image, see +[Running on HPC](hpc.md). + ### Run the tests (Optional) To ensure that the package is working correctly, you can run the tests using the pytest framework. diff --git a/pyproject.toml b/pyproject.toml index 93d975e..006731e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,7 @@ dev = [ [project.scripts] preprocess = "upp.main:main" check_input_samples = "upp.utils.check_input_samples:main" +list_components = "upp.utils.list_components:main" [tool.setuptools] packages = ["upp", "upp.classes", "upp.stages", "upp.utils"] diff --git a/scripts/batch/README.md b/scripts/batch/README.md new file mode 100644 index 0000000..857cbd1 --- /dev/null +++ b/scripts/batch/README.md @@ -0,0 +1,7 @@ +# Batch submission scripts + +Submit UPP preprocessing stages as Slurm or HTCondor batch jobs running inside the UPP container +image. + +See the [Running on HPC](https://umami-hep.github.io/umami-preprocessing/hpc/) documentation page +for usage instructions. diff --git a/scripts/batch/condor_batch.sh b/scripts/batch/condor_batch.sh new file mode 100755 index 0000000..02234a4 --- /dev/null +++ b/scripts/batch/condor_batch.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +# HTCondor job executable: starts the container on the worker node. + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# Prefer the CVMFS-unpacked image when available, fall back to the registry +DEFAULT_IMAGE="/cvmfs/unpacked.cern.ch/gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest" +if [[ ! -e "${DEFAULT_IMAGE}" ]]; then + DEFAULT_IMAGE="docker://gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest" +fi +IMAGE="${UPP_IMAGE:-${DEFAULT_IMAGE}}" +BINDS="${UPP_BINDS:-/home,/tmp}" + +apptainer exec --contain --pwd "${PWD}" -B "${BINDS}" -B "${PWD}" -B "${SCRIPT_DIR}" \ + "${IMAGE}" "${SCRIPT_DIR}/run_stage.sh" "$@" diff --git a/scripts/batch/condor_job.sub b/scripts/batch/condor_job.sub new file mode 100644 index 0000000..a0c735e --- /dev/null +++ b/scripts/batch/condor_job.sub @@ -0,0 +1,21 @@ +# HTCondor submit description used by submit.sh: +# condor_submit batch_dir=... batch_name=... upp_image=... upp_binds=... \ +# condor_job.sub -queue "args from " +# Edit the resources below for your pool. +request_cpus = 4 +request_memory = 32 GB +# CERN lxplus walltime flavour; ignored on other pools +# Available flavours: https://batchdocs.web.cern.ch/local/submit.html ++JobFlavour = "tomorrow" + +universe = vanilla +executable = $(batch_dir)/condor_batch.sh +arguments = $(args) +batch_name = $(batch_name) +environment = "UPP_IMAGE=$(upp_image) UPP_BINDS=$(upp_binds)" +output = logs/output/$(ClusterId).$(ProcId).out +error = logs/error/$(ClusterId).$(ProcId).err +log = logs/$(ClusterId).log + +# Assumes a shared filesystem between submit and worker nodes +should_transfer_files = NO diff --git a/scripts/batch/run_stage.sh b/scripts/batch/run_stage.sh new file mode 100755 index 0000000..dcba587 --- /dev/null +++ b/scripts/batch/run_stage.sh @@ -0,0 +1,100 @@ +#!/usr/bin/env bash +# Runs one preprocessing stage inside the container. Called by slurm_batch.sh / condor_batch.sh. +set -e + +usage() { + cat < + # full chain (prep+resample+merge+norm+plot) with split=all + + run_stage.sh normalise|normalize + run_stage.sh merge + run_stage.sh plotting + run_stage.sh prepare + run_stage.sh resampling + run_stage.sh fine_resampling +EOF +} + +if [ "$#" -eq 0 ] || [ "$1" = "-h" ] || [ "$1" = "--help" ]; then + usage + exit 2 +fi + +CONFIG="$1" +shift + +run_preprocess() { + preprocess --config "${CONFIG}" "$@" +} + +need_args() { + # Usage: need_args "$@" + local n_required="$1" + local mode="$2" + shift 2 + if [ "$#" -lt "$n_required" ]; then + echo "ERROR: '$mode' needs ${n_required} argument(s), got $#." >&2 + usage + exit 2 + fi +} + +if [ "$#" -eq 0 ]; then + echo "No mode given. Processing full chain." + run_preprocess --prep --split=all + run_preprocess --resample --split=all + run_preprocess --merge --split=all + run_preprocess --norm + run_preprocess --plot --split=all + exit 0 +fi + +MODE="$1" +shift + +case "${MODE}" in + normalise|normalize) + echo "Normalisation selected. Processing..." + run_preprocess --norm + ;; + + merge) + need_args 1 "merge" "$@" + echo "Start merging for $1. Processing..." + run_preprocess --merge --split "$1" + ;; + + plotting) + need_args 1 "plotting" "$@" + echo "Plotting selected ($1). Processing..." + run_preprocess --plot --split "$1" + ;; + + prepare) + need_args 2 "prepare" "$@" + echo "Start preparation for $1 ($2). Processing..." + run_preprocess --prep --component "$1" --split "$2" + ;; + + resampling) + need_args 2 "resampling" "$@" + echo "Start resampling for $1 ($2). Processing..." + run_preprocess --resample --region "$1" --split "$2" + ;; + + fine_resampling) + need_args 3 "fine_resampling" "$@" + echo "Start resampling for $2 ($3). Processing..." + run_preprocess --resample --region "$1" --component "$2" --split "$3" + ;; + + *) + echo "Step '${MODE}' not supported!" >&2 + usage + exit 2 + ;; +esac + +echo "Done!" diff --git a/scripts/batch/slurm_batch.sh b/scripts/batch/slurm_batch.sh new file mode 100755 index 0000000..3fbee28 --- /dev/null +++ b/scripts/batch/slurm_batch.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash +#SBATCH --ntasks 1 +#SBATCH --cpus-per-task 4 +#SBATCH --mem-per-cpu 8000 +#SBATCH --time 1-00:00:00 +# Edit the header above for your cluster, e.g. +# #SBATCH --partition +# #SBATCH --account + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# Prefer the CVMFS-unpacked image when available, fall back to the registry +DEFAULT_IMAGE="/cvmfs/unpacked.cern.ch/gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest" +if [[ ! -e "${DEFAULT_IMAGE}" ]]; then + DEFAULT_IMAGE="docker://gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest" +fi +IMAGE="${UPP_IMAGE:-${DEFAULT_IMAGE}}" +BINDS="${UPP_BINDS:-/home,/tmp}" + +srun apptainer exec --contain --pwd "${PWD}" -B "${BINDS}" -B "${PWD}" -B "${SCRIPT_DIR}" \ + "${IMAGE}" "${SCRIPT_DIR}/run_stage.sh" "$@" diff --git a/scripts/batch/submit.sh b/scripts/batch/submit.sh new file mode 100755 index 0000000..32799a6 --- /dev/null +++ b/scripts/batch/submit.sh @@ -0,0 +1,501 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# Prefer the CVMFS-unpacked image when available, fall back to the registry +DEFAULT_IMAGE="/cvmfs/unpacked.cern.ch/gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest" +if [[ ! -e "${DEFAULT_IMAGE}" ]]; then + DEFAULT_IMAGE="docker://gitlab-registry.cern.ch/aft/training-images/upp-images/upp:latest" +fi +IMAGE="${UPP_IMAGE:-${DEFAULT_IMAGE}}" +BINDS="${UPP_BINDS:-/home,/tmp}" + +# Batch scheduler: slurm or condor; empty means auto-detect +SCHEDULER="" + +# Throttle between sbatch calls (seconds, Slurm only). Set to 0 to disable. +THROTTLE="${THROTTLE:-30}" + +# Dry-run: if 1, print commands but do not execute sbatch. +DRY_RUN="${DRY_RUN:-0}" + +CONFIG="" + +# Filters on the components enumerated from the config; empty means "everything". +declare -a REGION_FILTER=() +declare -a SAMPLE_FILTER=() +declare -a FLAV_FILTER=() +declare -a SPLIT_FILTER=() + +# Argument lines collected for a single condor_submit call +declare -a CONDOR_JOBS=() + +# ---- Helpers -------------------------------------------------------------- +slugify() { + # keep alnum, dash, underscore; replace others with dash; squish repeats; trim + local s="${*:-}" + s="${s//[^[:alnum:]_-]/-}" + s="$(printf '%s' "$s" | sed -E 's/-+/-/g; s/^-+//; s/-+$//')" + # Slurm JobName limit is 128 chars; leave margin + printf '%.*s' 120 "$s" +} + +build_job_name() { + local mode="${1:-sequential}" + shift || true + local prefix="upp-$(basename "${CONFIG%.*}")" + case "$mode" in + prepare) slugify "${prefix}-prepare-${1:-component}-${2:-split}" ;; + fine_resampling) slugify "${prefix}-fres-${1:-region}-${2:-component}-${3:-split}" ;; + resampling) slugify "${prefix}-resampling-${1:-region}-${2:-split}" ;; + merge) slugify "${prefix}-merge-${1:-split}" ;; + normalise|normalize) slugify "${prefix}-normalise" ;; + plotting) slugify "${prefix}-plotting-${1:-split}" ;; + sequential|*) slugify "${prefix}-seq" ;; + esac +} + +detect_scheduler() { + if [[ -z "$SCHEDULER" ]]; then + if command -v sbatch >/dev/null 2>&1; then + SCHEDULER="slurm" + elif command -v condor_submit >/dev/null 2>&1; then + SCHEDULER="condor" + else + echo "ERROR: neither sbatch nor condor_submit found; use --scheduler slurm|condor" >&2 + exit 2 + fi + fi + case "$SCHEDULER" in + slurm|condor) ;; + *) + echo "ERROR: unknown scheduler: '$SCHEDULER' (expected slurm or condor)" >&2 + exit 2 + ;; + esac +} + +enumerate() { + # Use a local UPP installation when available, otherwise the container image + if command -v list_components >/dev/null 2>&1; then + list_components --config "${CONFIG}" "$@" + else + apptainer exec "${IMAGE}" list_components --config "${CONFIG}" "$@" + fi +} + +ensure_enumerated() { + # Fill ALL_ROWS/ALL_REGIONS/ALL_SAMPLES/ALL_FLAVS from the config (TSV rows), once + if [[ ${#ALL_ROWS[@]} -gt 0 ]]; then + return 0 + fi + mapfile -t ALL_ROWS < <(enumerate) + if [[ ${#ALL_ROWS[@]} -eq 0 ]]; then + echo "ERROR: no components found in $CONFIG" >&2 + exit 1 + fi + local row region sample flavour name + for row in "${ALL_ROWS[@]}"; do + IFS=$'\t' read -r region sample flavour name <<< "$row" + add_unique "$region" ALL_REGIONS + add_unique "$sample" ALL_SAMPLES + add_unique "$flavour" ALL_FLAVS + done +} + +append_list() { + # Append a comma/space separated list to the named array + local input="$1" + local out_name="$2" + local -a items=() + local x + IFS=', ' read -r -a items <<< "$input" + for x in "${items[@]}"; do + if [[ -n "$x" ]]; then + eval "$out_name+=(\"\$x\")" + fi + done +} + +add_unique() { + # Append the value to the named array if not already present + local x="$1" + local out_name="$2" + local -a arr=() + eval "arr=(\"\${${out_name}[@]}\")" + local i + for i in "${arr[@]}"; do + if [[ "$i" == "$x" ]]; then + return 0 + fi + done + eval "$out_name+=(\"\$x\")" +} + +in_list() { + # in_list [items...]; an empty list matches everything + local x="$1" + shift + if [[ $# -eq 0 ]]; then + return 0 + fi + local i + for i in "$@"; do + if [[ "$i" == "$x" ]]; then + return 0 + fi + done + return 1 +} + +print_resolved_config() { + printf '\nResolved configuration:\n' + printf ' CONFIG : %s\n' "$CONFIG" + printf ' SCHEDULER : %s\n' "$SCHEDULER" + printf ' IMAGE : %s\n' "$IMAGE" + printf ' THROTTLE : %s\n' "$THROTTLE" + if [[ ${#ALL_ROWS[@]} -gt 0 ]]; then + printf ' COMPONENTS : %s\n' "${#COMPONENTS[@]}" + fi + printf ' SPLITS : %s\n' "${SPLITS[*]}" + printf '\n' +} + +prompt_with_default() { + # Usage: prompt_with_default "Question" "default value" + local prompt="$1" + local default="$2" + local reply + + read -r -p "${prompt} [${default}]: " reply + if [[ -z "${reply}" ]]; then + printf '%s\n' "${default}" + else + printf '%s\n' "${reply}" + fi +} + +interactive_mode() { + echo + echo "No mode provided. Entering interactive mode." + echo + echo "Available modes:" + echo " sequential prepare fine_resampling resampling merge normalise plotting" + echo + + MODE="$(prompt_with_default "Select mode" "prepare")" + + local regions samples flavs splits + regions="$(prompt_with_default "Regions" "${ALL_REGIONS[*]}")" + samples="$(prompt_with_default "Samples" "${ALL_SAMPLES[*]}")" + flavs="$(prompt_with_default "Flavours" "${ALL_FLAVS[*]}")" + splits="$(prompt_with_default "Splits" "${SPLITS[*]}")" + + append_list "${regions}" REGION_FILTER + append_list "${samples}" SAMPLE_FILTER + append_list "${flavs}" FLAV_FILTER + SPLIT_FILTER=() + append_list "${splits}" SPLIT_FILTER +} + +submit() { + # Usage: submit [mode args...]; Slurm submits directly, condor collects for one condor_submit + if [[ "$SCHEDULER" == "condor" ]]; then + local line="$CONFIG" + if [[ $# -gt 0 ]]; then + line+=" $*" + fi + CONDOR_JOBS+=("$line") + return 0 + fi + + local jobname + jobname="$(build_job_name "$@")" + + local -a cmd=( + sbatch + --job-name="$jobname" + --output="${PWD}/logs/output/%j_%x.out" + --error="${PWD}/logs/error/%j_%x.err" + "${SCRIPT_DIR}/slurm_batch.sh" + "$CONFIG" + "$@" + ) + + echo "Submitting: $jobname" + + if [[ "$DRY_RUN" == "1" ]]; then + printf 'DRY-RUN: ' + printf '%q ' "${cmd[@]}" + printf '\n' + else + "${cmd[@]}" + if [[ "$THROTTLE" != "0" ]]; then + sleep "$THROTTLE" + fi + fi +} + +condor_submit_all() { + # Submit all collected jobs as a single condor cluster + if [[ ${#CONDOR_JOBS[@]} -eq 0 ]]; then + echo "Nothing to submit." + return 0 + fi + + local bname + bname="$(slugify "upp-$(basename "${CONFIG%.*}")-${MODE:-sequential}")" + local args_file="${PWD}/logs/condor_${MODE:-sequential}.args" + + local -a cmd=( + condor_submit + "batch_dir=${SCRIPT_DIR}" + "batch_name=${bname}" + "upp_image=${IMAGE}" + "upp_binds=${BINDS}" + "${SCRIPT_DIR}/condor_job.sub" + -queue "args from ${args_file}" + ) + + echo "Submitting ${#CONDOR_JOBS[@]} job(s) as batch $bname" + + if [[ "$DRY_RUN" == "1" ]]; then + printf 'DRY-RUN: job arguments:\n' + printf ' %s\n' "${CONDOR_JOBS[@]}" + printf 'DRY-RUN: ' + printf '%q ' "${cmd[@]}" + printf '\n' + else + printf '%s\n' "${CONDOR_JOBS[@]}" > "$args_file" + "${cmd[@]}" + fi +} + +usage() { + cat <<'EOF' +Usage: + submit.sh --config [options] [mode] + +Modes: + (no mode) # interactive mode: prompts for mode and filters + sequential # one job running the full chain + prepare # one job per component and split + fine_resampling # one job per component and split + resampling # one job per region and split + merge # one job per split + normalise|normalize # one job + plotting # one job per split + +Options: + --config # preprocessing config (required) + --scheduler # slurm or condor (default: auto-detect) + --dry-run # do not submit; print commands instead + --regions "" # only submit components in these regions + --samples "" # only submit components from these samples + --flavs "" # only submit components with these flavours + --splits "" # only submit these splits (default: train val test) + --throttle N # seconds between sbatch calls (default 30; 0 disables; Slurm only) + +Env (still supported): + UPP_IMAGE= # container image (unpacked dir, local .sif or docker:// URI) + UPP_BINDS= # comma-separated bind paths (default: /home,/tmp) + THROTTLE= # same as --throttle + DRY_RUN=1 # same as --dry-run + +List format: + Comma and/or space separated, e.g. "lowpt,highpt" or "lowpt highpt" + +Examples: + ./submit.sh --config configs/my-config.yaml --dry-run prepare + ./submit.sh --config configs/my-config.yaml --regions lowpt --splits train prepare + ./submit.sh --config configs/my-config.yaml resampling +EOF +} + +parse_args() { + local -a rest=() + while [[ $# -gt 0 ]]; do + case "$1" in + -h|--help) + usage + exit 0 + ;; + --config) + [[ $# -ge 2 ]] || { echo "ERROR: --config requires a value" >&2; exit 2; } + CONFIG="$2" + shift 2 + ;; + --scheduler) + [[ $# -ge 2 ]] || { echo "ERROR: --scheduler requires a value" >&2; exit 2; } + SCHEDULER="$2" + shift 2 + ;; + --dry-run) + DRY_RUN=1 + shift + ;; + --throttle) + [[ $# -ge 2 ]] || { echo "ERROR: --throttle requires a value" >&2; exit 2; } + THROTTLE="$2" + shift 2 + ;; + --regions) + [[ $# -ge 2 ]] || { echo "ERROR: --regions requires a value" >&2; exit 2; } + append_list "$2" REGION_FILTER + shift 2 + ;; + --samples) + [[ $# -ge 2 ]] || { echo "ERROR: --samples requires a value" >&2; exit 2; } + append_list "$2" SAMPLE_FILTER + shift 2 + ;; + --flavs|--flavors|--flavours) + [[ $# -ge 2 ]] || { echo "ERROR: --flavs requires a value" >&2; exit 2; } + append_list "$2" FLAV_FILTER + shift 2 + ;; + --splits) + [[ $# -ge 2 ]] || { echo "ERROR: --splits requires a value" >&2; exit 2; } + append_list "$2" SPLIT_FILTER + shift 2 + ;; + --) # end of options + shift + rest+=("$@") + break + ;; + -*) + echo "ERROR: unknown option: $1" >&2 + usage + exit 2 + ;; + *) + rest+=("$1") + shift + ;; + esac + done + + ARGS_REST=("${rest[@]}") +} + +# ---- Main ----------------------------------------------------------------- +main() { + declare -a ARGS_REST=() + parse_args "$@" + set -- "${ARGS_REST[@]}" + + if [[ -z "$CONFIG" ]]; then + echo "ERROR: --config is required" >&2 + usage + exit 2 + fi + + detect_scheduler + + declare -a ALL_ROWS=() ALL_REGIONS=() ALL_SAMPLES=() ALL_FLAVS=() + + MODE="${1:-}" + declare -a SPLITS=(train val test) + if [[ -z "$MODE" ]]; then + ensure_enumerated + interactive_mode + fi + if [[ ${#SPLIT_FILTER[@]} -gt 0 ]]; then + SPLITS=("${SPLIT_FILTER[@]}") + fi + + # Only the component/region level modes need the component list from the config + declare -a COMPONENTS=() REGIONS=() + local row region sample flavour name + case "$MODE" in prepare|fine_resampling|resampling) + ensure_enumerated + esac + for row in "${ALL_ROWS[@]}"; do + IFS=$'\t' read -r region sample flavour name <<< "$row" + if in_list "$region" "${REGION_FILTER[@]}" \ + && in_list "$sample" "${SAMPLE_FILTER[@]}" \ + && in_list "$flavour" "${FLAV_FILTER[@]}"; then + COMPONENTS+=("$row") + add_unique "$region" REGIONS + fi + done + + if [[ "$DRY_RUN" == "1" ]]; then + print_resolved_config + fi + + if [[ "$DRY_RUN" != "1" ]]; then + mkdir -p "${PWD}/logs/output" "${PWD}/logs/error" + fi + + local split + case "$MODE" in + sequential) + echo "Start submission for sequential processing..." + submit + ;; + + prepare) + echo "Start submission for preparation..." + for row in "${COMPONENTS[@]}"; do + IFS=$'\t' read -r region sample flavour name <<< "$row" + for split in "${SPLITS[@]}"; do + submit "prepare" "$name" "$split" + done + done + ;; + + fine_resampling) + echo "Start submission for fine resampling..." + for row in "${COMPONENTS[@]}"; do + IFS=$'\t' read -r region sample flavour name <<< "$row" + for split in "${SPLITS[@]}"; do + submit "fine_resampling" "$region" "$name" "$split" + done + done + ;; + + resampling) + echo "Start submission for resampling..." + for region in "${REGIONS[@]}"; do + for split in "${SPLITS[@]}"; do + submit "resampling" "$region" "$split" + done + done + ;; + + merge) + echo "Start submission for merging..." + for split in "${SPLITS[@]}"; do + submit "merge" "$split" + done + ;; + + normalise|normalize) + echo "Start submission for normalise..." + submit "normalise" + ;; + + plotting) + echo "Start submission for plotting..." + for split in "${SPLITS[@]}"; do + submit "plotting" "$split" + done + ;; + + *) + echo "Unsupported mode: '$MODE'" >&2 + usage + exit 2 + ;; + esac + + if [[ "$SCHEDULER" == "condor" ]]; then + condor_submit_all + fi + echo "Done!" +} + +main "$@" diff --git a/tests/unit/utils/test_list_components.py b/tests/unit/utils/test_list_components.py new file mode 100644 index 0000000..ab5f446 --- /dev/null +++ b/tests/unit/utils/test_list_components.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +from pathlib import Path + +from upp.utils.list_components import main + +CONFIG = Path(__file__).parents[3] / "upp/configs/test.yaml" + + +def test_list_components(capsys): + main(["--config", str(CONFIG)]) + rows = [line.split("\t") for line in capsys.readouterr().out.splitlines()] + assert len(rows) == 6 + assert all(len(row) == 4 for row in rows) + assert ["lowpt", "ttbar", "bjets", "lowpt_ttbar_bjets"] in rows + assert ["highpt", "zprime", "cjets", "highpt_zprime_cjets"] in rows + + +def test_list_components_regions(capsys): + main(["--config", str(CONFIG), "--regions"]) + assert capsys.readouterr().out.splitlines() == ["lowpt", "highpt"] diff --git a/upp/utils/list_components.py b/upp/utils/list_components.py new file mode 100644 index 0000000..85b9cf0 --- /dev/null +++ b/upp/utils/list_components.py @@ -0,0 +1,84 @@ +"""List the components defined in a preprocessing config.""" + +from __future__ import annotations + +import argparse +from typing import Any + +from ftag.cli_utils import HelpFormatter, valid_path + +from upp.classes.preprocessing_config import PreprocessingConfig + + +def parse_args(args: Any) -> argparse.Namespace: + """Parse the command line arguments. + + Parameters + ---------- + args : Any + Command line arguments. + + Returns + ------- + argparse.Namespace + Namespace with the parsed command line arguments + """ + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=HelpFormatter, + ) + parser.add_argument( + "--config", + required=True, + type=valid_path, + help="Path to config file", + ) + parser.add_argument( + "--split", + default="train", + choices=["train", "val", "test"], + help="Split to load the config for (component names are split-independent)", + ) + parser.add_argument( + "--regions", + action="store_true", + help="Only print the unique region names", + ) + + return parser.parse_args(args) + + +def main(args: Any | None = None) -> None: + """List components as tab-separated `region sample flavour name` rows. + + Parameters + ---------- + args : Any | None, optional + Command line arguments, by default None + """ + args = parse_args(args) + + config = PreprocessingConfig.from_file( + config_path=args.config, + split=args.split, + skip_checks=True, + skip_config_copy=True, + ) + + if args.regions: + for region in config.components.regions: + print(region.name) + return + + for component in config.components: + print( + component.region.name, + component.sample.name, + component.flavour.name, + component.name, + sep="\t", + ) + + +if __name__ == "__main__": + main() diff --git a/zensical.toml b/zensical.toml index 97af816..46ac56f 100644 --- a/zensical.toml +++ b/zensical.toml @@ -18,6 +18,7 @@ nav = [ { "Reweighting" = "reweighting.md" }, { "Configuration" = "configuration.md" }, { "Run" = "run.md" }, + { "Running on HPC" = "hpc.md" }, { "Umami integration" = "umami_int.md" }, { "Contributing" = "contributing.md" }, { "Docs development" = "docs_development.md" },