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

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

set -euo pipefail

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

module load python/3.12.5

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

OUTPUT_DIR="${TMPDIR:-/tmp}/smoke_pool_sample"
DOWNLOAD_DIR="${TMPDIR:-/tmp}/smoke_download"

echo "=============================================="
echo "  Smoke Test: Pool Sample"
echo "=============================================="
echo "Output:      $OUTPUT_DIR"
echo "Download:    $DOWNLOAD_DIR"
echo "Start:       $(date)"
echo "Hostname:    $(hostname)"
echo "TMPDIR:      ${TMPDIR:-/tmp}"
echo "Python:      $(python3 --version)"
echo "=============================================="

python3 -m dolma.pool_sample.cli \
    --output-dir "$OUTPUT_DIR" \
    --download-dir "$DOWNLOAD_DIR" \
    --token-budget 1000000 \
    --shard-fraction 0.001 \
    --num-workers 2 \
    --lines-per-shard 100000 \
    --log-every 1000 \
    --verbose

echo ""
echo "=============================================="
echo "  Output files"
echo "=============================================="
find "$OUTPUT_DIR" -type f | head -20
echo ""

if [ -f "$OUTPUT_DIR/manifest.json" ]; then
    echo "manifest.json:"
    cat "$OUTPUT_DIR/manifest.json"
fi

echo ""
echo "Finished: $(date)"
echo "=============================================="