Running on SLURM

We maintain a tested and approved SLURM batch script for running the pipeline on the CSD3 cluster with multiple icelake nodes. It uses dask distribution for the calibration stage, and MPI distribution for the imaging stage.

Feel free to adapt it to your needs.

#!/bin/bash
#
# Example file for submitting Mid pipeline jobs on CSD3
# Adjust requested resources and job variables as necessary
#
#SBATCH --job-name=MSCP
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=76
#SBATCH --time=02:30:00
#SBATCH --exclusive
#SBATCH --partition=icelake
#SBATCH --account=SKA-SDHP-SL2-CPU
#SBATCH --signal=B:TERM@180

##### User defined variables
BASE_JOB_DIR=/home/hpcmore1/rds/rds-sdhp-S7lLL7eOZIg/hpcmore1/PI22/test_wsclean_mp_integration_2

##### Job variables that MUST be defined
INPUT_MS=/rds/project/rds-S7lLL7eOZIg/hpcmore1/MeerKAT/MKT_CDFS25_4hours_950_1378MHz.ms
SINGULARITY_IMAGE=/rds/project/rds-S7lLL7eOZIg/singularity-images/dp3-wsclean-mpi-icelake-20240321.sif
BASE_OUTDIR=$BASE_JOB_DIR/output
SKYMODEL=$BASE_JOB_DIR/initial_skymodel_950_1378MHz.txt
CONFIG=$BASE_JOB_DIR/config.yml
NPIX=10240
PIXEL_SCALE_ASEC=1.1
DASK_WORKERS_PER_NODE=8
DASK_PORT=8786

# -----------------------------------------------------------------------------
# Anything below should not be edited
# -----------------------------------------------------------------------------

##### Setup modules
# Source bashrc to make module command available
source ~/.bashrc

# Load required base modules for icelake, and then OpenMPI 4.x
module purge
module load rhel8/default-icl
module load openmpi-4.0.5-gcc-8.4.1-l7ihwk3

set -x

##### Create output directories
PIPELINE_DIR=$BASE_OUTDIR/pipeline
DASK_LOGS_DIR=$BASE_OUTDIR/dask_logs
SYSMON_DIR=$BASE_OUTDIR/sysmon

mkdir -p $BASE_OUTDIR
mkdir -p $PIPELINE_DIR
mkdir -p $DASK_LOGS_DIR
mkdir -p $SYSMON_DIR

##### Fetch list of nodes
NODES=($(scontrol show hostnames))
HEAD_NODE="$(hostname)"

# Join array into space-separated string
NODES_SPACE_SEPARATED="${NODES[*]}"

echo "Allocated nodes: $NODES_SPACE_SEPARATED"
echo "Head node: $HEAD_NODE"

##### Start system monitor on ALL nodes
for node in "${NODES[@]}"; do
    outfile=$SYSMON_DIR/"system_usage_$node.jsonl"
    ssh $node mid-selfcal-system-monitor >$outfile &
    echo "Started system monitor on $node"
done

##### Start dask scheduler on head node
DASK_SCHEDULER_ADDR=$HEAD_NODE:$DASK_PORT

dask scheduler --port ${DASK_PORT} >$DASK_LOGS_DIR/scheduler_$HEAD_NODE.log 2>&1 &
echo "Started dask scheduler on $DASK_SCHEDULER_ADDR"

##### Start dask workers on all nodes via ssh (even on head node, it's fine)
for node in "${NODES[@]}"; do
    logfile=$DASK_LOGS_DIR/worker_$node.log
    # NOTE: the Mid pipeline expects the custom resource "subprocess_slots"
    # and it must be equal to 1.
    # Implicit assumption: Python environment for pipeline is loaded upon login
    ssh $node dask worker $DASK_SCHEDULER_ADDR --name $node \
        --nworkers $DASK_WORKERS_PER_NODE --resources subprocess_slots=1 \
        >$logfile 2>&1 &
    echo "Started dask worker on $node"
done

##### Start pipeline
# "exec" so that SIGTERM propagates to the pipeline executable
exec mid-selfcal-dd --dask-scheduler $DASK_SCHEDULER_ADDR \
    --mpi-hosts $NODES_SPACE_SEPARATED \
    --singularity-image $SINGULARITY_IMAGE \
    --config $CONFIG \
    --num-pixels $NPIX --pixel-scale $PIXEL_SCALE_ASEC \
    --sky-model $SKYMODEL \
    --outdir $PIPELINE_DIR \
    $INPUT_MS