social-data-attribution-w2/scripts/modal/run_materialize_all.sh

#!/usr/bin/env bash
set -euo pipefail

# Materialize and upload all samples sequentially (one at a time).
#
# Usage:
#   bash scripts/modal/run_materialize_all.sh
#   bash scripts/modal/run_materialize_all.sh --sizes "500,1000,5000"
#   UPLOAD_TO=r2,hf-buckets bash scripts/modal/run_materialize_all.sh

SIZES="${1:-500,1000,5000,10000,50000,100000,200000}"
UPLOAD_TO="${UPLOAD_TO:-hf-buckets}"
CHUNK_COUNT="${CHUNK_COUNT:-128}"
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
VENV_PYTHON="${REPO_ROOT}/.venv/bin/python"

IFS=',' read -ra SIZE_ARRAY <<< "$SIZES"

echo "Sequential materialize pipeline"
echo "Sizes: ${SIZE_ARRAY[*]}"
echo "Upload to: $UPLOAD_TO"
echo "Chunks: $CHUNK_COUNT"
echo ""

for size in "${SIZE_ARRAY[@]}"; do
    sample_name="sample_${size}_docs"
    volume_name="soc149-mat-$(echo "$sample_name" | tr '_' '-')"
    bucket_id="HCAI-Lab/dolma3-6t-$(echo "$sample_name" | tr '_' '-')"
    manifest="data/samples/${sample_name}/working_sample_manifest.parquet"

    echo "========================================"
    echo "Processing: $sample_name"
    echo "  Volume: $volume_name"
    echo "  Bucket: $bucket_id"
    echo "========================================"

    # Use --manifest-from-volume for large samples, local for small
    if [ -f "$manifest" ]; then
        echo "  Using local manifest: $manifest"
        SOC134_OUTPUT_VOLUME="$volume_name" \
            "$VENV_PYTHON" -m modal run \
                scripts/modal/materialize_working_sample.py \
                --sample-manifest "$manifest" \
                --sample-name "$sample_name" \
                --hf-bucket-id "$bucket_id" \
                --chunk-count "$CHUNK_COUNT" \
                --upload-to "$UPLOAD_TO"
    else
        echo "  Using manifest from Modal volume"
        SOC134_OUTPUT_VOLUME="$volume_name" \
            "$VENV_PYTHON" -m modal run \
                scripts/modal/materialize_working_sample.py \
                --manifest-from-volume \
                --sample-name "$sample_name" \
                --hf-bucket-id "$bucket_id" \
                --chunk-count "$CHUNK_COUNT" \
                --upload-to "$UPLOAD_TO"
    fi

    echo ""
    echo "Completed: $sample_name"
    echo ""
done

echo "All samples processed."