scripts/slurm/eval_ctt_generated_rollout.sbatch

#!/bin/bash
#SBATCH --job-name=ctt_ms_rollout
#SBATCH --account=def-yalda
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
#SBATCH --mem=24G
#SBATCH --time=01:00:00
#SBATCH --array=0-0
#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err

set -euo pipefail

PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
CHECKPOINT="${CHECKPOINT:?Set CHECKPOINT to a CTT model.pt}"
SOURCE_INDEX="${SOURCE_INDEX:-data/cil_charts/train/index.json}"
TARGET_INDEX="${TARGET_INDEX:-data/cil_charts/val/index.json}"
OUT_DIR="${OUT_DIR:-runs/ctt_residual_rollout_smoke_seed${SLURM_ARRAY_TASK_ID:-0}}"

SCRATCH_ROOT="/scratch/$USER/dovla"
SIF="${SIF:-$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif}"
PYTHON="${PYTHON:-$SCRATCH_ROOT/envs/maniskill/bin/python}"
NATIVE_LIBS="$SCRATCH_ROOT/native_libs/lib"
CPU_RENDER_LIBS="$SCRATCH_ROOT/cpu_render_libs"
CA_BUNDLE="$SCRATCH_ROOT/ca-bundle.crt"
VULKAN_ICD="$CPU_RENDER_LIBS/share/vulkan/icd.d/lvp_icd.x86_64.json"

K="${K:-8}"
POOL_SIZE="${POOL_SIZE:-0}"
NEIGHBORS="${NEIGHBORS:-8}"
MAX_TARGET_CHARTS="${MAX_TARGET_CHARTS:-8}"
GROUP_BATCH_SIZE="${GROUP_BATCH_SIZE:-1}"
DEVICE="${DEVICE:-cpu}"
SIM_BACKEND="${SIM_BACKEND:-physx_cpu}"
RENDER_BACKEND="${RENDER_BACKEND:-cpu}"
RESTORE_TOLERANCE="${RESTORE_TOLERANCE:-1e-5}"
DELTA_SCALE="${DELTA_SCALE:-1.0}"
BOOTSTRAP_SAMPLES="${BOOTSTRAP_SAMPLES:-200}"
INCLUDE_TARGETS_WITHOUT_POSITIVES="${INCLUDE_TARGETS_WITHOUT_POSITIVES:-0}"
EXCLUDE_SELF_SOURCE="${EXCLUDE_SELF_SOURCE:-0}"
SKIP_METRICS="${SKIP_METRICS:-0}"

module load StdEnv/2023 apptainer/1.4.5
cd "$PROJECT_DIR"
mkdir -p outputs/hpc/logs "$OUT_DIR"

RUNTIME_DIR="/tmp/$USER/ctt-ms-rollout-$SLURM_JOB_ID-${SLURM_ARRAY_TASK_ID:-0}"
CACHE_DIR="/tmp/$USER/ctt-ms-rollout-mesa-$SLURM_JOB_ID-${SLURM_ARRAY_TASK_ID:-0}"
mkdir -p "$RUNTIME_DIR" "$CACHE_DIR"
chmod 700 "$RUNTIME_DIR"

export OMP_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export MKL_NUM_THREADS=1
export DOVLA_TORCH_THREADS=1

EXTRA_ARGS=()
if [[ "$INCLUDE_TARGETS_WITHOUT_POSITIVES" == "1" ]]; then
  EXTRA_ARGS+=(--include-targets-without-positives)
fi
if [[ "$EXCLUDE_SELF_SOURCE" == "1" ]]; then
  EXTRA_ARGS+=(--exclude-self-source)
fi
if [[ "$SKIP_METRICS" == "1" ]]; then
  EXTRA_ARGS+=(--skip-metrics)
fi

apptainer exec \
  --env "LD_LIBRARY_PATH=$CPU_RENDER_LIBS/lib:$NATIVE_LIBS,VK_ICD_FILENAMES=$VULKAN_ICD,VK_DRIVER_FILES=$VULKAN_ICD,XDG_RUNTIME_DIR=$RUNTIME_DIR,MESA_SHADER_CACHE_DIR=$CACHE_DIR,LIBGL_ALWAYS_SOFTWARE=1,LP_NUM_THREADS=1,SSL_CERT_FILE=$CA_BUNDLE,REQUESTS_CA_BUNDLE=$CA_BUNDLE,OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1,DOVLA_TORCH_THREADS=1,MPLBACKEND=Agg,PYTHONDONTWRITEBYTECODE=1" \
  -B "$PROJECT_DIR:$PROJECT_DIR" \
  -B "/scratch/$USER:/scratch/$USER" \
  "$SIF" "$PYTHON" scripts/eval_ctt_generated_rollout.py \
    --checkpoint "$CHECKPOINT" \
    --source-index "$SOURCE_INDEX" \
    --target-index "$TARGET_INDEX" \
    --out-dir "$OUT_DIR" \
    --k "$K" \
    --pool-size "$POOL_SIZE" \
    --neighbors "$NEIGHBORS" \
    --max-target-charts "$MAX_TARGET_CHARTS" \
    --group-batch-size "$GROUP_BATCH_SIZE" \
    --device "$DEVICE" \
    --sim-backend "$SIM_BACKEND" \
    --render-backend "$RENDER_BACKEND" \
    --restore-tolerance "$RESTORE_TOLERANCE" \
    --delta-scale "$DELTA_SCALE" \
    --bootstrap-samples "$BOOTSTRAP_SAMPLES" \
    "${EXTRA_ARGS[@]}"