social-data-attribution-w2/scripts/slurm/attribution/precond_cache_data.sbatch

#!/usr/bin/env bash
#SBATCH --job-name=precond_cache
#SBATCH --output=logs/trackstar_precond/%j_cache.out
#SBATCH --error=logs/trackstar_precond/%j_cache.err
#SBATCH --partition=ice-cpu
#SBATCH --cpus-per-task=4
#SBATCH --mem=32G
#SBATCH --time=02:00:00

# Consolidate preconditioner input data and cache in shared storage.
#
# 1. Downloads 100K doc sample from HF (if not already present)
# 2. Consolidates into single JSONL
# 3. Downloads base query data and concatenates
# 4. Copies to shared storage for reuse across worktrees
#
# After this runs once, launch_preconditioner.sh finds data automatically.

set -euo pipefail

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

PYTHONEXEC="${PYTHONEXEC:-$REPO_DIR/.venv/bin/python}"
export PATH="$REPO_DIR/.venv/bin:$PATH"
export PYTHONPATH="src"

export HF_TOKEN="${HF_TOKEN:-$(cat ~/.hf_token 2>/dev/null || true)}"
export HF_HOME="${TMPDIR:-/tmp}/hf_cache_$$"
mkdir -p "$HF_HOME"

STAGING="${STAGING_DIR:-$HOME/scratch/soc168/data}"
SHARED="/storage/ice-shared/cs7634/staff/TDA/preconditioner_data"
mkdir -p "$STAGING" "$SHARED"

echo "=============================================="
echo "  Preconditioner Data Cache"
echo "=============================================="
echo "Staging: $STAGING"
echo "Shared:  $SHARED"
echo "Start:   $(date)"
echo "=============================================="

# Step 1: Download 100K doc sample (skip if already consolidated)
OUTPUT_JSONL="$STAGING/preconditioner_100k.jsonl"
if [ -f "$SHARED/preconditioner_100k.jsonl" ]; then
    echo "Already cached in shared storage, skipping download."
    OUTPUT_JSONL="$SHARED/preconditioner_100k.jsonl"
elif [ -f "$OUTPUT_JSONL" ]; then
    echo "Consolidated JSONL exists in staging, skipping download."
else
    echo ""
    echo "--- Download 100K doc sample ---"
    RAW_DIR="$STAGING/preconditioner_100k_raw"
    if [ -d "$RAW_DIR/data" ]; then
        echo "Raw data already downloaded, skipping."
    else
        mkdir -p "$RAW_DIR"
        "$PYTHONEXEC" -c "
from huggingface_hub import snapshot_download
snapshot_download(
    'HCAI-Lab/dolma3-6t-preconditioner-100k',
    repo_type='dataset',
    local_dir='$RAW_DIR',
    token='$HF_TOKEN',
)
print('Download complete')
"
    fi

    echo ""
    echo "--- Consolidate to single JSONL ---"
    "$PYTHONEXEC" scripts/attribution/consolidate_preconditioner_sample.py \
        --input-dir "$RAW_DIR" \
        --output "$OUTPUT_JSONL"
    echo "Consolidated: $(wc -l < "$OUTPUT_JSONL") docs"
fi

# Step 2: Download and combine base query data
COMBINED="$STAGING/base_queries_combined.jsonl"
if [ -f "$SHARED/base_queries_combined.jsonl" ]; then
    echo "Query data already cached in shared storage."
    COMBINED="$SHARED/base_queries_combined.jsonl"
elif [ -f "$COMBINED" ]; then
    echo "Combined queries exist in staging."
else
    echo ""
    echo "--- Download and combine base query data ---"
    QUERY_DIR="$STAGING/base_queries"
    "$PYTHONEXEC" -m data_attribution.attribution.trackstar.queries \
        --variant base \
        --output-dir "$QUERY_DIR"
    cat "$QUERY_DIR"/*.jsonl > "$COMBINED"
    echo "Combined queries: $(wc -l < "$COMBINED") lines"
fi

# Step 3: Copy to shared storage
echo ""
echo "--- Copy to shared storage ---"
if [ ! -f "$SHARED/preconditioner_100k.jsonl" ] && [ -f "$STAGING/preconditioner_100k.jsonl" ]; then
    cp "$STAGING/preconditioner_100k.jsonl" "$SHARED/preconditioner_100k.jsonl"
    echo "Copied preconditioner_100k.jsonl to shared storage"
fi
if [ ! -f "$SHARED/base_queries_combined.jsonl" ] && [ -f "$STAGING/base_queries_combined.jsonl" ]; then
    cp "$STAGING/base_queries_combined.jsonl" "$SHARED/base_queries_combined.jsonl"
    echo "Copied base_queries_combined.jsonl to shared storage"
fi

# Step 4: Upload consolidated JSONL to HuggingFace dataset
HF_DATASET="HCAI-Lab/dolma3-6t-preconditioner-100k"
CONSOLIDATED="$SHARED/preconditioner_100k.jsonl"
if [ ! -f "$CONSOLIDATED" ]; then
    CONSOLIDATED="$STAGING/preconditioner_100k.jsonl"
fi

echo ""
echo "--- Upload consolidated JSONL to HF ---"
export HF_HUB_DISABLE_XET=1
export HF_HOME="${TMPDIR:-/tmp}/hf_upload_$$"
mkdir -p "$HF_HOME"

"$PYTHONEXEC" -c "
from huggingface_hub import HfApi
import os

api = HfApi(token=os.environ['HF_TOKEN'])
repo_id = '$HF_DATASET'
local_path = '$CONSOLIDATED'

existing = api.list_repo_files(repo_id, repo_type='dataset')
if 'preconditioner_100k.jsonl' in existing:
    print('preconditioner_100k.jsonl already exists in HF dataset, skipping upload.')
else:
    print(f'Uploading {local_path} to {repo_id}...')
    api.upload_file(
        path_or_fileobj=local_path,
        path_in_repo='preconditioner_100k.jsonl',
        repo_id=repo_id,
        repo_type='dataset',
        commit_message='Add consolidated preconditioner_100k.jsonl (id + text only)',
    )
    print('Upload complete.')
"

echo ""
echo "=============================================="
echo "  Data cache complete"
echo "=============================================="
echo "Shared storage:"
ls -lh "$SHARED/" 2>/dev/null || echo "(empty)"
echo "End: $(date)"
echo "=============================================="