social-data-attribution-w2/scripts/slurm/sampling/sample_pool_150B.sh

#!/usr/bin/env bash
#SBATCH --job-name=pool_sample_150B
#SBATCH --partition=ice-cpu
#SBATCH --cpus-per-task=24
#SBATCH --mem=96G
#SBATCH --time=18:00:00
#SBATCH --output=logs/pool_sample/%x_%j.out
#SBATCH --error=logs/pool_sample/%x_%j.err

set -euo pipefail

JOB_INDEX="${JOB_INDEX:?Set JOB_INDEX (0-4) before submitting}"
TOKEN_BUDGET="${TOKEN_BUDGET:-35000000000}"
NUM_WORKERS="${NUM_WORKERS:-20}"
DOWNLOAD_WORKERS="${DOWNLOAD_WORKERS:-8}"

REPO_DIR="${SLURM_SUBMIT_DIR:-$(cd "$(dirname "$0")/../../.." && pwd)}"
cd "$REPO_DIR"
mkdir -p logs/pool_sample

SHARD_FILE="${REPO_DIR}/scripts/slurm/data/pool_shards_job_${JOB_INDEX}.txt"

if [ ! -f "$SHARD_FILE" ]; then
    echo "FAIL: Shard manifest not found: $SHARD_FILE"
    echo "Run scripts/manifests/pre_split_pool_shards.py first."
    exit 1
fi

if [ -f .venv/bin/activate ]; then
    source .venv/bin/activate
fi

export PYTHONPATH="${REPO_DIR}/src${PYTHONPATH:+:$PYTHONPATH}"
export HF_TOKEN="${HF_TOKEN:-$(cat ~/.hf_token 2>/dev/null || true)}"
export HF_HUB_DISABLE_XET=1

OUTPUT_DIR="${TMPDIR:-/tmp}/pool_sample_job_${JOB_INDEX}"
DOWNLOAD_DIR="${TMPDIR:-/tmp}/pool_download_job_${JOB_INDEX}"
FINAL_DIR="${HOME}/scratch/archive-dolma3-pool-150b/job_${JOB_INDEX}"

SHARD_COUNT=$(wc -l < "$SHARD_FILE" | tr -d ' ')

echo "=============================================="
echo "  Pool Sample 150B - Job ${JOB_INDEX}"
echo "=============================================="
echo "Shard file:  $SHARD_FILE ($SHARD_COUNT shards)"
echo "Output:      $OUTPUT_DIR"
echo "Download:    $DOWNLOAD_DIR"
echo "Final:       $FINAL_DIR"
echo "Budget:      $TOKEN_BUDGET tokens ($(echo "$TOKEN_BUDGET / 1000000000" | bc)B)"
echo "Workers:     $NUM_WORKERS"
echo "DL workers:  $DOWNLOAD_WORKERS"
echo "XET:         DISABLED"
echo "Start:       $(date)"
echo "Hostname:    $(hostname)"
echo "TMPDIR:      ${TMPDIR:-/tmp}"
echo "=============================================="

python3 -c "import zstandard; print(f'zstandard: {zstandard.__version__}')"
python3 -c "import huggingface_hub; print(f'huggingface_hub: {huggingface_hub.__version__}')"

SAMPLE_EXIT=0
python3 -m dolma.pool_sample.cli \
    --output-dir "$OUTPUT_DIR" \
    --download-dir "$DOWNLOAD_DIR" \
    --shard-paths-file "$SHARD_FILE" \
    --token-budget "$TOKEN_BUDGET" \
    --num-workers "$NUM_WORKERS" \
    --download-workers "$DOWNLOAD_WORKERS" \
    --verbose || SAMPLE_EXIT=$?

echo ""
echo "=============================================="
echo "  Copying output to scratch"
echo "=============================================="

if [ "$SAMPLE_EXIT" -ne 0 ]; then
    echo "WARNING: sampling exited with code $SAMPLE_EXIT, copying partial output"
fi

mkdir -p "$FINAL_DIR"
if ls "$OUTPUT_DIR"/* >/dev/null 2>&1; then
    cp -r "$OUTPUT_DIR"/* "$FINAL_DIR"/
    echo "Output copied to $FINAL_DIR"
else
    echo "No output files to copy"
fi

echo "Finished: $(date) (exit code: $SAMPLE_EXIT)"
echo "=============================================="

exit "$SAMPLE_EXIT"