social-data-attribution-w2/scripts/slurm/sampling/extract_stratified_docs.sbatch

#!/usr/bin/env bash
#SBATCH --job-name=extract_strat_docs
#SBATCH --account=gts-mriedl3
#SBATCH --cpus-per-task=20
#SBATCH --mem=64G
#SBATCH --time=12:00:00
#SBATCH --array=0-4
#SBATCH --output=logs/extract_stratified_docs/%A_job%a.out
#SBATCH --error=logs/extract_stratified_docs/%A_job%a.err

# Extract stratified document text from archive-dolma3-pool-150b shards.
#
# 5 SLURM jobs (one per job_{0-4}), each spawning 20 parallel workers.
# Each worker handles one shard and writes:
#   - A .docs.jsonl.zst with matched documents
#   - A per-worker .manifest.parquet with extraction metadata
#
# Usage:
#   Single job test:  sbatch --array=0-0 scripts/slurm/extract_stratified_docs.sbatch
#   Full run:         sbatch scripts/slurm/extract_stratified_docs.sbatch

set -euo pipefail

REPO_DIR="${SLURM_SUBMIT_DIR:-$PWD}"
cd "$REPO_DIR"
mkdir -p logs/extract_stratified_docs

WORKERS_PER_JOB=20
JOB_ID="$SLURM_ARRAY_TASK_ID"

# ── Activate venv ────────────────────────────────────────────────────
if [ -f .venv/bin/activate ]; then
    source .venv/bin/activate
else
    echo "ERROR: .venv not found at $REPO_DIR/.venv"
    exit 1
fi

# ── HF auth ─────────────────────────────────────────────────────────
export HF_TOKEN="${HF_TOKEN:-$(cat ~/.hf_token 2>/dev/null || true)}"
export HF_HUB_DISABLE_XET=1
export PYTHONUNBUFFERED=1

# ── Configuration ───────────────────────────────────────────────────
ID_DIR="${ID_DIR:-stratified_data/shard_rows}"
SHARD_LIST="${SHARD_LIST:-scripts/slurm/data/stratified_shard_list.txt}"
OUTPUT_DIR="${OUTPUT_DIR:-stratified_data/merged_docs}"
CACHE_DIR="${TMPDIR:-/tmp}/hf_cache"
DATASET="${DATASET:-HCAI-Lab/archive-dolma3-pool-150b}"

echo "=============================================="
echo "  Extract Stratified Docs - Job $JOB_ID"
echo "=============================================="
echo "Job ID:       $JOB_ID"
echo "Workers:      $WORKERS_PER_JOB"
echo "ID dir:       $ID_DIR"
echo "Shard list:   $SHARD_LIST"
echo "Output dir:   $OUTPUT_DIR"
echo "Cache dir:    $CACHE_DIR"
echo "Dataset:      $DATASET"
echo "Start:        $(date)"
echo "Hostname:     $(hostname)"
echo "=============================================="

FAIL_COUNT=0

for WORKER_ID in $(seq 0 $((WORKERS_PER_JOB - 1))); do
    (
        WORKER_TAG="job${JOB_ID}_worker$(printf '%03d' $WORKER_ID)"
        WORKER_LOG="logs/extract_stratified_docs/${SLURM_ARRAY_JOB_ID}_${WORKER_TAG}.log"

        echo "[$WORKER_TAG] Starting at $(date)" > "$WORKER_LOG"

        python3 scripts/sampling/extract_stratified_docs.py \
            --job-id "$JOB_ID" \
            --worker-id "$WORKER_ID" \
            --workers-per-job "$WORKERS_PER_JOB" \
            --id-dir "$ID_DIR" \
            --shard-list "$SHARD_LIST" \
            --output-dir "$OUTPUT_DIR" \
            --cache-dir "$CACHE_DIR" \
            --dataset "$DATASET" \
            >> "$WORKER_LOG" 2>&1

        echo "[$WORKER_TAG] Finished at $(date)" >> "$WORKER_LOG"
    ) &
done

wait

echo "=============================================="
echo "  Job $JOB_ID complete"
echo "=============================================="
echo "Done: $(date)"
echo "=============================================="