social-data-attribution-w2/scripts/slurm/dedup/materialize_pool_docs.sbatch

#!/usr/bin/env bash
#SBATCH --job-name=materialize_pool
#SBATCH --array=0-0%200
#SBATCH --cpus-per-task=4
#SBATCH --mem=16G
#SBATCH --time=04:00:00
#SBATCH -C localNVMe
#SBATCH --output=logs/dedup/%x_%A_%a.out
#SBATCH --error=logs/dedup/%x_%A_%a.err

set -euo pipefail

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

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

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

python3 scripts/slurm/dedup/materialize_pool_docs.py \
    --manifest "${POOL_MANIFEST:-scripts/slurm/shard_manifest_pool_shared.txt}" \
    --bloom-file "${BLOOM_FILE:-/storage/ice-shared/cs7634/staff/TDA/soc-90/unique_6t_ids.bloom}" \
    --output-dir "${POOL_OUTPUT_DIR:-/storage/ice-shared/cs7634/staff/TDA/soc-90/unique_docs_work/pool}" \
    --download-dir "$DOWNLOAD_DIR_DEFAULT" \
    --task-id "${SLURM_ARRAY_TASK_ID}" \
    --task-count "${SLURM_ARRAY_TASK_COUNT}" \
    --shards-per-task "${SHARDS_PER_TASK:-5}" \
    --startup-delay "${STARTUP_DELAY:-60}" \
    --skip-existing \
    --verbose