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

#!/usr/bin/env bash
# Submit 5 parallel pool sampling SLURM jobs.
# Run from repo root after generating shard manifests with pre_split_pool_shards.py.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
REPO_DIR="$(cd "$SCRIPT_DIR/../../.." && pwd)"
cd "$REPO_DIR"

NUM_JOBS="${NUM_JOBS:-5}"

echo "Submitting $NUM_JOBS pool sample jobs ..."

for i in $(seq 0 $((NUM_JOBS - 1))); do
    SHARD_FILE="$REPO_DIR/scripts/slurm/data/pool_shards_job_${i}.txt"
    if [ ! -f "$SHARD_FILE" ]; then
        echo "FAIL: Missing $SHARD_FILE"
        echo "Run: python3 scripts/manifests/pre_split_pool_shards.py --num-jobs $NUM_JOBS"
        exit 1
    fi
done

mkdir -p logs/pool_sample

for i in $(seq 0 $((NUM_JOBS - 1))); do
    JOB_ID=$(export JOB_INDEX=$i && sbatch --export=ALL,JOB_INDEX=$i scripts/slurm/sampling/sample_pool_150B.sh | grep -o '[0-9]*')
    echo "  Job $i: submitted (SLURM ID $JOB_ID)"
done

echo ""
echo "Monitor with: squeue -u \$USER"
echo "Logs in: logs/pool_sample/"