Envío de trabajos al clúster con sbatch

Esta página explica cómo lanzar trabajos al clúster ANTS usando Slurm y una plantilla base de sbatch pensada para entrenamientos de modelos en GPU. La plantilla está optimizada para usar el scratch local del nodo durante la ejecución y persistir los resultados en tu home al terminar.

Antes de lanzar un trabajo conviene tener claro lo siguiente:

• sbatch envía un script a la cola de Slurm; el clúster decide cuándo y dónde se ejecuta.
• Las líneas que empiezan por #SBATCH son directivas que se leen antes de ejecutar nada y configuran el trabajo (partición, recursos, tiempo, etc.).
• El scratch de cada nodo ($SLURM_TMPDIR o /scratch/…) es almacenamiento rápido pero temporal: se borra al terminar el job. El home (/slurm/home/$USER) es persistente pero más lento.
• La buena práctica es: copiar al scratch → trabajar allí → copiar resultados al home al terminar.

La plantilla espera que tu proyecto siga esta estructura en el directorio desde el que lances sbatch:

mi-proyecto/
├── code/                  # Código fuente (main.py y módulos)
│   └── main.py            # Punto de entrada del entrenamiento
├── data/                  # Dataset
├── lib/
│   └── install_env.sh     # Script que crea/activa el entorno (venv, módulos, CUDA)
├── artifacts/             # (opcional) checkpoints previos para reanudar
└── train.sbatch           # La plantilla de este documento

Guarda el siguiente contenido como train.sbatch en la raíz de tu proyecto.

train.sbatch

#!/usr/bin/env bash
#SBATCH --job-name=base-train          # Nombre identificador del trabajo en la cola
#SBATCH --partition=gpu                # Particion (cola) a usar — CPU o GPU
#SBATCH --nodes=1                      # Numero de nodos del cluster a reservar
#SBATCH --ntasks=1                     # Numero de tareas MPI (procesos paralelos)
#SBATCH --cpus-per-task=19             # Hilos de CPU por tarea
#SBATCH --gres=gpu:1                   # Recurso generico: 1 GPU H100 NVL o gpu:1 para cualquiera
#SBATCH --mem=31G                      # Memoria RAM total reservada en el nodo
#SBATCH --time=24:00:00                # Tiempo maximo de ejecucion (HH:MM:SS)
#SBATCH --output=/slurm/home/%u/output/%j/terminal.out  # stdout
#SBATCH --error=/slurm/home/%u/output/%j/terminal.err   # stderr
 
set -euo pipefail
 
OUTDIR="/slurm/home/$USER/output/$SLURM_JOB_ID"
mkdir -p "$OUTDIR"
 
SCRATCH="${SLURM_TMPDIR:-/scratch/slurm/$USER/$SLURM_JOB_ID/tmp}"
JOBSCRATCH="${SCRATCH}/base-${SLURM_JOB_ID}"
mkdir -p "$JOBSCRATCH"
 
SUBMIT_DIR="${SLURM_SUBMIT_DIR:-$PWD}"
 
cleanup() {
  rsync -a --ignore-missing-args \
    "$JOBSCRATCH/code/artifacts/" "$SUBMIT_DIR/artifacts/" 2>/dev/null || true
}
trap cleanup EXIT INT TERM
 
rsync -a "$SUBMIT_DIR/code/" "$JOBSCRATCH/code/"
rsync -a "$SUBMIT_DIR/data/" "$JOBSCRATCH/data/"
 
if [[ -d "$SUBMIT_DIR/artifacts" ]]; then
  mkdir -p "$JOBSCRATCH/code/artifacts"
  rsync -a "$SUBMIT_DIR/artifacts/" "$JOBSCRATCH/code/artifacts/"
fi
 
rsync -a "$SUBMIT_DIR/lib/" "$JOBSCRATCH/lib/"
source "$JOBSCRATCH/lib/install_env.sh"
 
export ARTIFACTS_DIR="$JOBSCRATCH/code/artifacts"
export PERSISTENT_ARTIFACTS_DIR="/slurm/home/$USER/artifacts"
 
cd "$JOBSCRATCH/code"
python main.py 2>&1 | tee "$OUTDIR/training.log"

Estas líneas configuran el trabajo antes de empezar a ejecutarse. Slurm las usa para reservar recursos y planificar el job en la cola.

set -euo pipefail

Hace que el script aborte ante cualquier error: -e (sale si un comando falla), -u (error si usas una variable no definida), -o pipefail (un fallo dentro de un pipe propaga el código de error). Imprescindible para no continuar ejecutando si, por ejemplo, rsync del dataset falla.

OUTDIR="/slurm/home/$USER/output/$SLURM_JOB_ID"
mkdir -p "$OUTDIR"

Crea un directorio por Job ID en tu home donde se guardarán los logs (terminal.out, terminal.err y training.log). Sobrevive al borrado del scratch.

SCRATCH="${SLURM_TMPDIR:-/scratch/slurm/$USER/$SLURM_JOB_ID/tmp}"
JOBSCRATCH="${SCRATCH}/base-${SLURM_JOB_ID}"
mkdir -p "$JOBSCRATCH"

SLURM_TMPDIR lo define automáticamente Slurm cuando hay scratch local configurado; el ${VAR:-fallback} usa una ruta alternativa si no existe. Aquí trabajará el job: lectura/escritura rápidas porque suele ser NVMe local del nodo, no NFS.

cleanup() {
  rsync -a --ignore-missing-args \
    "$JOBSCRATCH/code/artifacts/" "$SUBMIT_DIR/artifacts/" 2>/dev/null || true
}
trap cleanup EXIT INT TERM

trap registra la función cleanup para que se ejecute siempre al terminar el script: tanto si finaliza con éxito (EXIT), si lo cancelas con Ctrl+C (INT), o si Slurm lo mata por tiempo (TERM). Así nunca pierdes los checkpoints del scratch.

rsync -a "$SUBMIT_DIR/code/" "$JOBSCRATCH/code/"
rsync -a "$SUBMIT_DIR/data/" "$JOBSCRATCH/data/"
 
if [[ -d "$SUBMIT_DIR/artifacts" ]]; then
  mkdir -p "$JOBSCRATCH/code/artifacts"
  rsync -a "$SUBMIT_DIR/artifacts/" "$JOBSCRATCH/code/artifacts/"
fi

Trae el código y el dataset al scratch local antes de empezar a entrenar. Si existen artefactos previos (checkpoint de un entrenamiento anterior) los restaura, permitiendo reanudar sin volver a empezar.

rsync -a "$SUBMIT_DIR/lib/" "$JOBSCRATCH/lib/"
source "$JOBSCRATCH/lib/install_env.sh"

install_env.sh es responsabilidad tuya: típicamente carga módulos (module load cuda/12.x), crea/activa un venv, e instala dependencias con pip. Debe ser idempotente.

Un comando o instrucción es idempotente si puedes ejecutarlo una o varias veces y el resultado final siempre será exactamente el mismo. No importa cuántas veces lo repitas, no causará efectos secundarios no deseados después de la primera ejecución ni depende del resultado de ejecuciones anteriores.

export ARTIFACTS_DIR="$JOBSCRATCH/code/artifacts"
export PERSISTENT_ARTIFACTS_DIR="/slurm/home/$USER/artifacts"

Tu main.py debe leerlas para decidir dónde escribir:

• ARTIFACTS_DIR: checkpoints frecuentes durante el entrenamiento (scratch, rápido).
• PERSISTENT_ARTIFACTS_DIR: el checkpoint final o copias periódicas de seguridad (home, persistente).

cd "$JOBSCRATCH/code"
python main.py 2>&1 | tee "$OUTDIR/training.log"

2>&1 fusiona stderr en stdout y tee escribe a la vez por pantalla (que va a terminal.out) y al log persistente en el home.

Desde la raíz de tu proyecto:

sbatch train.sbatch

Slurm devolverá algo como:

Submitted batch job 12345

Ese número es el Job ID ($SLURM_JOB_ID). Lo necesitarás para consultar el estado y leer los logs.

Cosas que cambiarás casi siempre:

• –job-name: pon un nombre descriptivo (resnet50-imagenet, llama-finetune-v3…).
• –cpus-per-task, –mem, –time: ajusta a tus necesidades reales.
• –gres: ver tabla siguiente.
• main.py: añade los argumentos que necesites después.

Lista las GPUs disponibles con:

sinfo -o "%P %N %G"

Modifica la última línea:

python main.py --epochs 50 --batch-size 64 --lr 1e-4 2>&1 | tee "$OUTDIR/training.log"

O mejor, parametriza el sbatch aceptando variables de entorno:

python main.py \
  --epochs "${EPOCHS:-50}" \
  --batch-size "${BATCH_SIZE:-64}" \
  2>&1 | tee "$OUTDIR/training.log"

Y lánzalo así:

sbatch --export=ALL,EPOCHS=100,BATCH_SIZE=128 train.sbatch

squeue -u $USER                # Solo tus trabajos
squeue --me                    # Equivalente moderno
squeue -p gpu                  # Todos los trabajos de la particion gpu

Códigos de estado más habituales: R (running), PD (pending), CG (completing), F (failed).

scontrol show job 12345

tail -f /slurm/home/$USER/output/12345/terminal.out
tail -f /slurm/home/$USER/output/12345/terminal.err

sacct -j 12345 --format=JobID,JobName,State,Elapsed,MaxRSS,ReqMem,ReqCPUS,AllocTRES%40

Útil para ajustar peticiones en el siguiente lanzamiento: si MaxRSS fue de 8G y pediste –mem=31G, estás desperdiciando memoria.

scancel 12345              # Cancela un job concreto
scancel -u $USER           # Cancela TODOS tus jobs (uso con cuidado)
scancel -n base-train      # Cancela por nombre

Para depurar sin pasar por la cola, abre una sesión interactiva con los mismos recursos:

srun --partition=gpu --gres=gpu:1 --cpus-per-task=4 --mem=8G \
     --time=01:00:00 --pty bash

Te dará una shell dentro de un nodo del clúster. Cuando termines, exit libera los recursos.

• Pide solo lo necesario: recursos sobrantes empeoran tu prioridad y bloquean a otros.
• Guarda checkpoints periódicos en PERSISTENT_ARTIFACTS_DIR, no solo al final.
• Usa el scratch para I/O intensivo; el home vía NFS es lento.
• Lanza un job corto de prueba (–time=00:15:00) antes de uno de 24 horas.
• Versiona tu sbatch junto al código en Git para reproducibilidad.
• No ejecutes entrenamientos directamente en el nodo de login; usa sbatch o srun.

Disponibles dentro del script:

• Documentación oficial de sbatch