social-data-attribution-w2/scripts/slurm/eda/eda_array.sbatch

#!/usr/bin/env bash
#SBATCH --cpus-per-task=8
#SBATCH --mem=32G
#SBATCH --time=12:00:00
#SBATCH --partition=cpu-small
#SBATCH --array=0-0
#SBATCH --output=logs/dolma_eda/%A_%a.out
#SBATCH --error=logs/dolma_eda/%A_%a.err

set -euo pipefail

mkdir -p logs/dolma_eda

PROJECT_ROOT=${PROJECT_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}}
cd "$PROJECT_ROOT"

PYTHONEXEC=${PYTHONEXEC:-.venv/bin/python}
PYTHONPATH=${PYTHONPATH:-$PROJECT_ROOT/src}
export PYTHONPATH

MANIFEST=${MANIFEST:?}
OUTDIR=${OUTDIR:-runs/dolma_enriched/eda}
WORKERS_DIR=${WORKERS_DIR:-$OUTDIR/workers}
VALUE_FIELDS=${VALUE_FIELDS:-"lang warc_content_type cc_dump"}
TOP_K=${TOP_K:-50}
JOINT_TOP_K=${JOINT_TOP_K:-200}
HEATMAP_K=${HEATMAP_K:-25}

mkdir -p "$WORKERS_DIR"

LINE=$(sed -n "$((SLURM_ARRAY_TASK_ID+1))p" "$MANIFEST")
if [[ -z "$LINE" ]]; then
  echo "No manifest entry for task $SLURM_ARRAY_TASK_ID"
  exit 1
fi

"$PYTHONEXEC" -m dolma.eda.core \
  --mode worker \
  --manifest "$MANIFEST" \
  --manifest-index "$SLURM_ARRAY_TASK_ID" \
  --output-dir "$OUTDIR" \
  --worker-output-dir "$WORKERS_DIR/$SLURM_ARRAY_TASK_ID" \
  --value-fields $VALUE_FIELDS \
  --top-k "$TOP_K" \
  --joint-top-k "$JOINT_TOP_K" \
  --heatmap-k "$HEATMAP_K" \
  --no-report \
  --no-write-png \
  --resume