social-data-attribution-w2/scripts/modal/materialize_working_sample.py

"""Modal runner for SOC-134 working sample materialization."""

from __future__ import annotations

import json
import logging
import os
import time
import traceback
from datetime import datetime, timezone
from pathlib import Path

import modal

from config import (
    R2_BUCKET,
    R2_ENDPOINT_URL,
    R2_SECRET_NAME,
)

logger = logging.getLogger("materialize_working_sample")

_local_path = Path(__file__).resolve()
if len(_local_path.parents) > 2:
    _REPO_ROOT = _local_path.parents[2]
    _SRC_ROOT = str(_REPO_ROOT / "src")
    _CONFIG_PY = str(_REPO_ROOT / "scripts" / "modal" / "config.py")
    _materialize_image = (
        modal.Image.debian_slim(python_version="3.12")
        .pip_install(
            "boto3>=1.37.0",
            "pyarrow>=18.0.0",
            "zstandard>=0.24.0",
            "pandas>=2.0.0",
            "huggingface-hub>=0.25",
        )
        .env({"PYTHONPATH": "/root/src:/root"})
        .add_local_file(_CONFIG_PY, remote_path="/root/config.py", copy=True)
        .add_local_dir(_SRC_ROOT, remote_path="/root/src", copy=True)
    )
    _hf_upload_image = (
        modal.Image.debian_slim(python_version="3.12")
        .pip_install(
            "huggingface-hub>=1.5.0",
            "hf-xet",
            "zstandard>=0.24.0",
        )
        .env(
            {
                "PYTHONPATH": "/root/src:/root",
                "HF_XET_HIGH_PERFORMANCE": "1",
            }
        )
        .add_local_file(_CONFIG_PY, remote_path="/root/config.py", copy=True)
        .add_local_dir(_SRC_ROOT, remote_path="/root/src", copy=True)
    )
    _hf_buckets_image = (
        modal.Image.debian_slim(python_version="3.12")
        .pip_install(
            "huggingface-hub>=1.7.0",
            "hf-xet",
        )
        .env({"PYTHONPATH": "/root/src:/root", "HF_XET_HIGH_PERFORMANCE": "1"})
        .add_local_file(_CONFIG_PY, remote_path="/root/config.py", copy=True)
        .add_local_dir(_SRC_ROOT, remote_path="/root/src", copy=True)
    )
else:
    _materialize_image = modal.Image.debian_slim(python_version="3.12")
    _hf_upload_image = modal.Image.debian_slim(python_version="3.12")
    _hf_buckets_image = modal.Image.debian_slim(python_version="3.12")

app = modal.App("soc134-materialize-working-sample")

r2_secret = modal.Secret.from_name(R2_SECRET_NAME)
hf_secret = modal.Secret.from_name("huggingface-secret")

_OUTPUT_VOLUME_ENV = "SOC134_OUTPUT_VOLUME"
_DEFAULT_OUTPUT_VOLUME = "soc134-materialize-cache"
_output_volume_name = os.environ.get(_OUTPUT_VOLUME_ENV, _DEFAULT_OUTPUT_VOLUME)

materialize_volume = modal.Volume.from_name(_output_volume_name, create_if_missing=True)
corpus_volume = modal.Volume.from_name("soc134-corpus-cache", create_if_missing=True)
samples_volume = modal.Volume.from_name("soc149-drawn-samples", create_if_missing=True)

VOLUME_MOUNT = "/cache"
SAMPLES_MOUNT = "/drawn-samples"
CORPUS_MOUNT = "/corpus"
WORKER_CPU = 2
WORKER_MEMORY = 16384
WORKER_TIMEOUT = 7200
R2_SAMPLES_PREFIX = "soc134-samples"


def _ensure_r2_env() -> None:
    env_aliases = {
        "R2_ACCESS_KEY_ID": ("R2_ACCESS_KEY_ID", "AWS_ACCESS_KEY_ID", "access_key_id"),
        "R2_SECRET_ACCESS_KEY": (
            "R2_SECRET_ACCESS_KEY",
            "AWS_SECRET_ACCESS_KEY",
            "secret_access_key",
        ),
    }
    for target, aliases in env_aliases.items():
        if target in os.environ and os.environ[target]:
            continue
        for alias in aliases:
            value = os.environ.get(alias)
            if value:
                os.environ[target] = value
                break
        if target not in os.environ:
            raise KeyError(target)

    if "R2_ENDPOINT_URL" not in os.environ:
        os.environ["R2_ENDPOINT_URL"] = R2_ENDPOINT_URL
    if "R2_BUCKET" not in os.environ:
        os.environ["R2_BUCKET"] = R2_BUCKET


def _create_r2_client():
    from dolma.quality.r2 import R2Config, create_r2_client

    config = R2Config.from_env(output_prefix="soc134-materialized")
    return create_r2_client(config), config


@app.function(
    image=_materialize_image,
    volumes={VOLUME_MOUNT: materialize_volume},
    timeout=600,
    cpu=1,
)
def cleanup_volume(current_run_id: str = "") -> dict[str, object]:
    import shutil

    materialize_volume.reload()
    cache_root = Path(VOLUME_MOUNT)
    entries = [p for p in cache_root.iterdir() if p.is_dir()]

    if not entries:
        return {"deleted": 0, "kept": 0, "message": "Volume already empty"}

    deleted = 0
    kept = 0
    removed_names = []
    for entry in entries:
        if entry.name.startswith("."):
            continue
        if current_run_id and entry.name == current_run_id:
            kept += 1
            continue
        shutil.rmtree(entry, ignore_errors=True)
        removed_names.append(entry.name)
        deleted += 1

    materialize_volume.commit()
    return {
        "deleted": deleted,
        "kept": kept,
        "dirs_removed": removed_names[:10],
    }


@app.function(
    image=_materialize_image,
    volumes={
        VOLUME_MOUNT: materialize_volume,
        SAMPLES_MOUNT: samples_volume,
    },
    timeout=600,
    cpu=1,
    memory=8192,
)
def stage_manifest_from_volume(
    run_id: str,
    sample_name: str,
) -> dict[str, object]:
    import shutil

    samples_volume.reload()
    src_dir = Path(SAMPLES_MOUNT) / sample_name
    if not src_dir.exists():
        raise FileNotFoundError(f"Sample not found on volume: {sample_name}")

    manifest_dir = Path(VOLUME_MOUNT) / run_id
    manifest_dir.mkdir(parents=True, exist_ok=True)

    src_manifest = src_dir / "working_sample_manifest.parquet"
    dst_manifest = manifest_dir / "manifest.parquet"
    shutil.copy2(src_manifest, dst_manifest)
    manifest_size = dst_manifest.stat().st_size

    for name in ("sample_contract.json", "bin_summary.csv"):
        src = src_dir / name
        if src.exists():
            shutil.copy2(src, manifest_dir / name)

    materialize_volume.commit()
    return {
        "run_id": run_id,
        "sample_name": sample_name,
        "manifest_size": manifest_size,
        "source": "volume",
    }


@app.function(
    image=_materialize_image,
    volumes={VOLUME_MOUNT: materialize_volume},
    timeout=300,
    cpu=1,
)
def stage_manifest(
    run_id: str,
    manifest_bytes: bytes,
    contract_json: str = "",
    bin_summary_csv: str = "",
) -> dict[str, object]:
    manifest_dir = Path(VOLUME_MOUNT) / run_id
    manifest_dir.mkdir(parents=True, exist_ok=True)
    manifest_path = manifest_dir / "manifest.parquet"
    manifest_path.write_bytes(manifest_bytes)
    if contract_json:
        (manifest_dir / "sample_contract.json").write_text(contract_json)
    if bin_summary_csv:
        (manifest_dir / "bin_summary.csv").write_text(bin_summary_csv)
    materialize_volume.commit()
    return {
        "run_id": run_id,
        "manifest_path": str(manifest_path),
        "manifest_size": len(manifest_bytes),
    }


@app.function(
    image=_materialize_image,
    volumes={
        VOLUME_MOUNT: materialize_volume,
        CORPUS_MOUNT: corpus_volume,
    },
    secrets=[r2_secret],
    timeout=WORKER_TIMEOUT,
    cpu=WORKER_CPU,
    memory=WORKER_MEMORY,
    retries=2,
)
def run_chunk(
    run_id: str,
    chunk_index: int,
    chunk_count: int,
) -> dict[str, object]:
    t0 = time.monotonic()
    worker_dir = Path(VOLUME_MOUNT) / run_id / f"worker_{chunk_index:04d}"
    worker_dir.mkdir(parents=True, exist_ok=True)

    try:
        _ensure_r2_env()
        client, config = _create_r2_client()

        from dolma.materialize_sample import (
            materialize_all,
            write_materialize_stats,
        )

        materialize_volume.reload()
        corpus_volume.reload()
        manifest_path = Path(VOLUME_MOUNT) / run_id / "manifest.parquet"

        corpus_dir = Path(CORPUS_MOUNT)

        result = materialize_all(
            manifest_path=manifest_path,
            output_dir=worker_dir,
            r2_client=client,
            bucket=config.bucket,
            chunk_index=chunk_index,
            chunk_count=chunk_count,
            skip_existing=True,
            corpus_dir=corpus_dir,
        )

        write_materialize_stats(result, worker_dir)
        materialize_volume.commit()

        elapsed = time.monotonic() - t0
        return {
            "run_id": run_id,
            "chunk_index": chunk_index,
            "found_docs": result.total_found,
            "expected_docs": result.total_expected,
            "missing_docs": result.total_missing,
            "shards_processed": len(result.shard_stats),
            "bytes_written": result.total_bytes,
            "elapsed_seconds": round(elapsed, 1),
        }
    except Exception as exc:
        error_msg = f"{type(exc).__name__}: {exc}"
        logger.error(
            "FATAL chunk %d: %s\n%s",
            chunk_index,
            error_msg,
            traceback.format_exc(),
        )
        return {
            "run_id": run_id,
            "chunk_index": chunk_index,
            "status": "error",
            "error": error_msg,
        }


def _build_readme(hf_repo: str, contract: dict) -> str:
    repo_name = hf_repo.split("/")[-1] if "/" in hf_repo else hf_repo
    dpb = contract.get("WORKING_SAMPLE_DOCS_PER_BIN", 0)
    tfpb = contract.get("WORKING_SAMPLE_TOKEN_FLOOR_PER_BIN", 0)
    mode = f"{dpb:,} docs/bin" if dpb else f"{tfpb:,} token floor/bin"
    return (
        "\n".join(
            [
                "---",
                "license: odc-by",
                "task_categories:",
                "  - text-generation",
                "language:",
                "  - en",
                "---",
                "",
                f"# {repo_name}",
                "",
                "Materialized working sample from the Dolma3 6T deduplicated corpus.",
                "",
                "## Parameters",
                "",
                "| Parameter | Value |",
                "|-----------|-------|",
                f"| Sampling mode | {mode} |",
                f"| Seed | {contract.get('WORKING_SAMPLE_SAMPLING_SEED', 'N/A')} |",
                "",
                "## Counts",
                "",
                "| Metric | Value |",
                "|--------|-------|",
                f"| Total documents | {contract.get('WORKING_SAMPLE_REALIZED_DOC_COUNT', 0):,} |",
                f"| Total tokens | {contract.get('WORKING_SAMPLE_REALIZED_TOKEN_TOTAL', 0):,} |",
                f"| Bins covered | {contract.get('WORKING_SAMPLE_COVERED_BIN_COUNT', 0)}/{contract.get('WORKING_SAMPLE_TOTAL_BIN_COUNT', 576)} |",
                f"| Bins underfilled | {contract.get('WORKING_SAMPLE_UNDERFILLED_BIN_COUNT', 0)} |",
                "",
                "## Format",
                "",
                "JSONL files compressed with zstandard under `data/`.",
                "Each record contains the original Dolma document fields (id, text, metadata).",
                "",
                "Sampling metadata files:",
                "- `working_sample_manifest.parquet` - sampled doc_ids with bin assignments",
                "- `sample_contract.json` - aggregate sampling contract",
                "- `bin_summary.csv` - per-bin fill rates (576 bins = 24 topics x 24 formats)",
            ]
        )
        + "\n"
    )


def _r2_upload_file(client, local_path: Path, r2_key: str, content_type: str) -> bool:
    try:
        client.head_object(Bucket=R2_BUCKET, Key=r2_key)
        return False
    except Exception:
        pass
    max_retries = 3
    for attempt in range(max_retries):
        try:
            data = local_path.read_bytes()
            client.put_object(
                Bucket=R2_BUCKET, Key=r2_key, Body=data, ContentType=content_type
            )
            return True
        except Exception:
            if attempt == max_retries - 1:
                raise
            time.sleep((attempt + 1) * 5)
    return False


@app.function(
    image=_materialize_image,
    volumes={VOLUME_MOUNT: materialize_volume},
    secrets=[r2_secret],
    timeout=7200,
    cpu=2,
    memory=4096,
)
def upload_worker_dir_to_r2(
    sample_name: str,
    run_id: str,
    worker_index: int,
) -> dict[str, object]:
    t0 = time.monotonic()
    _ensure_r2_env()
    client, _ = _create_r2_client()

    worker_name = f"worker_{worker_index:04d}"
    worker_dir = Path(VOLUME_MOUNT) / run_id / worker_name
    r2_prefix = f"{R2_SAMPLES_PREFIX}/{sample_name}"

    materialize_volume.reload()

    if not worker_dir.exists():
        return {"worker_index": worker_index, "uploaded": 0, "skipped": 0, "failed": 0}

    shard_files = sorted(worker_dir.glob("*.jsonl.zst"))
    uploaded = 0
    skipped = 0
    failed = 0
    total_bytes = 0

    for shard_file in shard_files:
        r2_key = f"{r2_prefix}/{worker_name}/{shard_file.name}"
        try:
            if _r2_upload_file(client, shard_file, r2_key, "application/zstd"):
                uploaded += 1
                total_bytes += shard_file.stat().st_size
            else:
                skipped += 1
        except Exception as exc:
            failed += 1
            logger.error("R2 upload failed %s: %s", r2_key, exc)

    return {
        "worker_index": worker_index,
        "uploaded": uploaded,
        "skipped": skipped,
        "failed": failed,
        "total_bytes": total_bytes,
        "elapsed_seconds": round(time.monotonic() - t0, 1),
    }


@app.function(
    image=_hf_buckets_image,
    volumes={VOLUME_MOUNT: materialize_volume},
    secrets=[hf_secret],
    timeout=3600,
    cpu=2,
    memory=4096,
    retries=2,
)
def sync_worker_to_hf_bucket(
    run_id: str,
    hf_bucket_id: str,
    worker_index: int,
) -> dict[str, object]:
    from huggingface_hub import sync_bucket

    materialize_volume.reload()
    worker_name = f"worker_{worker_index:04d}"
    worker_dir = Path(VOLUME_MOUNT) / run_id / worker_name
    token = os.environ.get("HF_TOKEN")

    if not worker_dir.exists():
        return {"worker_index": worker_index, "status": "skipped", "files": 0}

    dest = f"hf://buckets/{hf_bucket_id}/{worker_name}"
    max_retries = 3
    for attempt in range(max_retries):
        try:
            sync_bucket(
                str(worker_dir),
                dest,
                include=["*.jsonl.zst"],
                token=token,
            )
            file_count = len(list(worker_dir.glob("*.jsonl.zst")))
            return {"worker_index": worker_index, "status": "ok", "files": file_count}
        except Exception as exc:
            if attempt == max_retries - 1:
                return {
                    "worker_index": worker_index,
                    "status": "error",
                    "error": f"{type(exc).__name__}: {exc}",
                }
            time.sleep((attempt + 1) * 5)
    return {"worker_index": worker_index, "status": "error", "error": "unreachable"}


@app.function(
    image=_hf_buckets_image,
    volumes={VOLUME_MOUNT: materialize_volume},
    secrets=[hf_secret],
    timeout=600,
    cpu=1,
    memory=4096,
)
def upload_hf_bucket_metadata(
    run_id: str,
    hf_bucket_id: str,
) -> dict[str, object]:
    from huggingface_hub import batch_bucket_files, create_bucket

    materialize_volume.reload()
    run_root = Path(VOLUME_MOUNT) / run_id
    token = os.environ.get("HF_TOKEN")

    create_bucket(hf_bucket_id, exist_ok=True, token=token)

    metadata_files: list[tuple[bytes, str]] = []
    for name in ("sample_contract.json", "bin_summary.csv"):
        path = run_root / name
        if path.exists():
            metadata_files.append((path.read_bytes(), name))
    manifest_path = run_root / "manifest.parquet"
    if manifest_path.exists():
        metadata_files.append(
            (manifest_path.read_bytes(), "working_sample_manifest.parquet")
        )
    if metadata_files:
        batch_bucket_files(hf_bucket_id, add=metadata_files, token=token)

    return {"bucket": hf_bucket_id, "metadata_files": len(metadata_files)}


def upload_to_hf_buckets(
    run_id: str,
    hf_bucket_id: str,
    chunk_count: int,
) -> dict[str, object]:
    logger.info("Uploading metadata to %s", hf_bucket_id)
    meta_result = upload_hf_bucket_metadata.remote(run_id, hf_bucket_id)
    logger.info("Metadata: %s", meta_result)

    logger.info("Syncing %d workers to %s in parallel", chunk_count, hf_bucket_id)
    worker_results = list(
        sync_worker_to_hf_bucket.starmap(
            [(run_id, hf_bucket_id, i) for i in range(chunk_count)]
        )
    )

    ok = sum(1 for r in worker_results if r.get("status") == "ok")
    skipped = sum(1 for r in worker_results if r.get("status") == "skipped")
    errors = [r for r in worker_results if r.get("status") == "error"]
    total_files = sum(r.get("files", 0) for r in worker_results)

    logger.info(
        "Bucket upload: %d ok, %d skipped, %d errors, %d files",
        ok, skipped, len(errors), total_files,
    )
    if errors:
        for e in errors[:5]:
            logger.error("Worker %d: %s", e["worker_index"], e.get("error"))

    return {
        "hf_bucket_id": hf_bucket_id,
        "workers_ok": ok,
        "workers_skipped": skipped,
        "workers_errored": len(errors),
        "total_files": total_files,
    }


@app.function(
    image=_hf_upload_image,
    volumes={VOLUME_MOUNT: materialize_volume},
    secrets=[hf_secret],
    timeout=86400,
    cpu=2,
    memory=16384,
    ephemeral_disk=524288,
)
def upload_to_hf(
    run_id: str,
    hf_repo: str,
    chunk_count: int,
) -> dict[str, object]:
    from huggingface_hub import CommitOperationAdd, HfApi

    materialize_volume.reload()
    run_root = Path(VOLUME_MOUNT) / run_id
    tracker_dir = run_root / ".hf_upload"
    tracker_dir.mkdir(parents=True, exist_ok=True)

    def _new_api():
        return HfApi(token=os.environ.get("HF_TOKEN"))

    def _batch_done(batch_num: int) -> bool:
        return (tracker_dir / f"batch_{batch_num:04d}.done").exists()

    def _mark_batch_done(batch_num: int) -> None:
        (tracker_dir / f"batch_{batch_num:04d}.done").write_text("done\n")
        materialize_volume.commit()

    api = _new_api()
    api.create_repo(repo_id=hf_repo, repo_type="dataset", exist_ok=True)

    if not (tracker_dir / "metadata.done").exists():
        metadata_ops = []
        contract = {}
        contract_path = run_root / "sample_contract.json"
        if contract_path.exists():
            contract_text = contract_path.read_text()
            contract = json.loads(contract_text)
            metadata_ops.append(
                CommitOperationAdd(
                    path_in_repo="sample_contract.json",
                    path_or_fileobj=contract_text.encode(),
                )
            )

        bin_summary_path = run_root / "bin_summary.csv"
        if bin_summary_path.exists():
            metadata_ops.append(
                CommitOperationAdd(
                    path_in_repo="bin_summary.csv",
                    path_or_fileobj=bin_summary_path.read_bytes(),
                )
            )

        manifest_path = run_root / "manifest.parquet"
        if manifest_path.exists():
            metadata_ops.append(
                CommitOperationAdd(
                    path_in_repo="working_sample_manifest.parquet",
                    path_or_fileobj=str(manifest_path),
                )
            )

        readme_text = _build_readme(hf_repo, contract)
        metadata_ops.append(
            CommitOperationAdd(
                path_in_repo="README.md",
                path_or_fileobj=readme_text.encode(),
            )
        )

        if metadata_ops:
            api.preupload_lfs_files(
                repo_id=hf_repo,
                additions=[
                    op for op in metadata_ops if op.path_in_repo.endswith(".parquet")
                ],
                repo_type="dataset",
            )
            api.create_commit(
                repo_id=hf_repo,
                operations=metadata_ops,
                commit_message="Add sampling metadata and data card",
                repo_type="dataset",
            )

        (tracker_dir / "metadata.done").write_text("done\n")
        materialize_volume.commit()
        logger.info("Metadata committed to HF")
    else:
        logger.info("Metadata already uploaded, skipping")

    max_per_dir = 5000
    all_files: list[tuple[Path, str]] = []
    for chunk_idx in range(chunk_count):
        worker_dir = run_root / f"worker_{chunk_idx:04d}"
        if not worker_dir.exists():
            continue
        for jsonl_file in sorted(worker_dir.glob("*.jsonl.zst")):
            file_idx = len(all_files)
            part = f"part_{file_idx // max_per_dir:03d}"
            hf_path = f"data/{part}/{jsonl_file.name}"
            all_files.append((jsonl_file, hf_path))

    if not all_files:
        return {"status": "error", "error": "No JSONL.zst files found"}

    batch_size = 100
    total_uploaded = 0
    total_skipped = 0
    total_bytes = 0
    max_retries = 3

    for batch_start in range(0, len(all_files), batch_size):
        batch_num = batch_start // batch_size
        batch = all_files[batch_start : batch_start + batch_size]

        if _batch_done(batch_num):
            total_skipped += len(batch)
            continue

        for attempt in range(max_retries):
            try:
                api = _new_api()
                operations = [
                    CommitOperationAdd(
                        path_in_repo=hf_path,
                        path_or_fileobj=str(local_path),
                    )
                    for local_path, hf_path in batch
                ]

                api.preupload_lfs_files(
                    repo_id=hf_repo,
                    additions=operations,
                    repo_type="dataset",
                )
                api.create_commit(
                    repo_id=hf_repo,
                    operations=operations,
                    commit_message=f"Add materialized shards batch {batch_num}",
                    repo_type="dataset",
                )

                batch_bytes = sum(p.stat().st_size for p, _ in batch)
                total_uploaded += len(batch)
                total_bytes += batch_bytes
                _mark_batch_done(batch_num)

                if (batch_num + 1) % 10 == 0:
                    total_batches = (len(all_files) + batch_size - 1) // batch_size
                    logger.info(
                        "Upload progress: batch %d/%d (%d files)",
                        batch_num + 1,
                        total_batches,
                        total_uploaded,
                    )
                break
            except Exception as exc:
                if attempt == max_retries - 1:
                    error_msg = f"Batch {batch_num}: {type(exc).__name__}: {exc}"
                    logger.error("FATAL batch %d: %s", batch_num, error_msg)
                    return {
                        "status": "error",
                        "error": error_msg,
                        "batches_completed": batch_num,
                        "files_uploaded": total_uploaded,
                        "files_skipped": total_skipped,
                    }
                wait = (attempt + 1) * 10
                logger.warning(
                    "Batch %d failed (attempt %d), retrying in %ds: %s",
                    batch_num,
                    attempt + 1,
                    wait,
                    exc,
                )
                time.sleep(wait)

    return {
        "hf_repo": hf_repo,
        "files_uploaded": total_uploaded,
        "files_skipped": total_skipped,
        "bytes_uploaded": total_bytes,
    }


@app.function(
    image=_materialize_image,
    volumes={
        VOLUME_MOUNT: materialize_volume,
        CORPUS_MOUNT: corpus_volume,
    },
    secrets=[r2_secret, hf_secret],
    timeout=43200,
    cpu=2,
    memory=8192,
)
def run_materialize_pipeline(
    run_id: str,
    sample_name: str,
    chunk_count: int,
    upload_targets_json: str,
    hf_repo: str = "",
    hf_bucket_id: str = "",
    upload_only: bool = False,
    force_upload: bool = False,
) -> str:
    upload_targets = json.loads(upload_targets_json)
    skip_upload = False
    report_path = Path(VOLUME_MOUNT) / run_id / "materialize_report.json"

    if not upload_only:
        logger.info("Launching %d materialize workers...", chunk_count)
        worker_results = list(
            run_chunk.starmap(
                [(run_id, i, chunk_count) for i in range(chunk_count)]
            )
        )

        total_found = sum(r.get("found_docs", 0) for r in worker_results)
        total_expected = sum(r.get("expected_docs", 0) for r in worker_results)
        total_missing = sum(r.get("missing_docs", 0) for r in worker_results)
        errors = [r for r in worker_results if r.get("status") == "error"]

        report = {
            "run_id": run_id,
            "sample_name": sample_name,
            "total_found": total_found,
            "total_expected": total_expected,
            "total_missing": total_missing,
            "workers_completed": len(worker_results) - len(errors),
            "workers_failed": len(errors),
            "completeness_ratio": (
                total_found / total_expected if total_expected > 0 else 0.0
            ),
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "worker_summary": [
                {
                    "chunk": r.get("chunk_index"),
                    "found": r.get("found_docs", 0),
                    "expected": r.get("expected_docs", 0),
                    "status": r.get("status", "ok"),
                }
                for r in worker_results
            ],
        }
        report_path.write_text(json.dumps(report, indent=2) + "\n")
        materialize_volume.commit()

        logger.info(
            "Materialization: %d/%d found (%.1f%%), %d missing, %d errors",
            total_found,
            total_expected,
            100 * report["completeness_ratio"],
            total_missing,
            len(errors),
        )

        if errors:
            for e in errors[:5]:
                logger.error("Chunk %s: %s", e.get("chunk_index"), e.get("error"))
            skip_upload = True
            logger.error("BLOCKING UPLOAD: %d worker errors", len(errors))

        if total_found < total_expected and not force_upload:
            skip_upload = True
            logger.error(
                "BLOCKING UPLOAD: incomplete materialization "
                "(%d/%d docs, %d missing). Use --force-upload to override.",
                total_found,
                total_expected,
                total_missing,
            )
    else:
        materialize_volume.reload()
        if report_path.exists():
            report = json.loads(report_path.read_text())
            total_found = report["total_found"]
            total_expected = report["total_expected"]
            total_missing = report["total_missing"]
            errors = []
            if total_found < total_expected and not force_upload:
                skip_upload = True
                logger.error(
                    "BLOCKING UPLOAD: report shows incomplete materialization "
                    "(%d/%d docs). Use --force-upload to override.",
                    total_found,
                    total_expected,
                )
            else:
                logger.info(
                    "Report verified: %d/%d docs (%.1f%%)",
                    total_found,
                    total_expected,
                    100 * total_found / total_expected if total_expected > 0 else 0,
                )
        elif not force_upload:
            skip_upload = True
            total_found = 0
            total_expected = 0
            total_missing = 0
            errors = []
            logger.error(
                "BLOCKING UPLOAD: no materialize_report.json found. "
                "Cannot verify completeness. Use --force-upload to override."
            )
        else:
            total_found = 0
            total_expected = 0
            total_missing = 0
            errors = []
            logger.warning("No report found, proceeding with --force-upload")

    results = {
        "run_id": run_id,
        "sample_name": sample_name,
        "total_found": total_found,
        "total_expected": total_expected,
        "total_missing": total_missing,
        "workers_failed": len(errors),
        "upload_blocked": skip_upload,
    }

    if not skip_upload and "r2" in upload_targets:
        logger.info("Uploading to R2: %s/%s/", R2_SAMPLES_PREFIX, sample_name)
        r2_results = list(
            upload_worker_dir_to_r2.starmap(
                [(sample_name, run_id, i) for i in range(chunk_count)]
            )
        )
        results["r2_uploaded"] = sum(r.get("uploaded", 0) for r in r2_results)
        results["r2_failed"] = sum(r.get("failed", 0) for r in r2_results)

    if not skip_upload and "hf" in upload_targets:
        logger.info("Uploading to HF Dataset: %s", hf_repo)
        hf_result = upload_to_hf.remote(run_id, hf_repo, chunk_count)
        results["hf_result"] = hf_result

    if not skip_upload and "hf-buckets" in upload_targets:
        logger.info("Uploading to HF Bucket: %s", hf_bucket_id)
        bucket_result = upload_to_hf_buckets(run_id, hf_bucket_id, chunk_count)
        results["hf_bucket_result"] = bucket_result

        if "hf-buckets" in upload_targets and hf_bucket_id:
            from huggingface_hub import list_bucket_tree

            token = os.environ.get("HF_TOKEN")
            bucket_files = sum(
                1
                for item in list_bucket_tree(
                    hf_bucket_id, recursive=True, token=token
                )
                if hasattr(item, "type") and item.type == "file"
                and item.path.endswith(".jsonl.zst")
            )
            results["bucket_verified_files"] = bucket_files
            logger.info(
                "Bucket verification: %d .jsonl.zst files on %s",
                bucket_files,
                hf_bucket_id,
            )

    return json.dumps(results)


@app.local_entrypoint()
def main(
    sample_manifest: str = "",
    sample_name: str = "",
    hf_repo: str = "HCAI-Lab/dolma3-6t-working-sample",
    hf_bucket_id: str = "",
    chunk_count: int = 128,
    run_id: str = "",
    upload_only: bool = False,
    upload_to: str = "hf-buckets",
    keep_old_runs: bool = False,
    manifest_from_volume: bool = False,
    force_upload: bool = False,
) -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )
    use_volume_manifest = manifest_from_volume
    if not sample_manifest and not use_volume_manifest:
        logger.error("--sample-manifest or --manifest-from-volume is required")
        raise SystemExit(1)

    if use_volume_manifest and not sample_name:
        logger.error("--sample-name is required with --manifest-from-volume")
        raise SystemExit(1)

    if sample_manifest:
        manifest_path = Path(sample_manifest)
        if not manifest_path.exists():
            logger.error("manifest not found: %s", manifest_path)
            raise SystemExit(1)
        resolved_sample_name = sample_name or manifest_path.parent.name
    else:
        resolved_sample_name = sample_name

    upload_targets = [t.strip() for t in upload_to.split(",")]

    current_run_id = run_id or datetime.now(timezone.utc).strftime(
        "soc134_materialize_%Y%m%d_%H%M%S"
    )

    resolved_bucket_id = hf_bucket_id or (
        "HCAI-Lab/dolma3-6t-" + resolved_sample_name.replace("_", "-")
    )

    logger.info("Run ID: %s", current_run_id)
    logger.info("Sample: %s", resolved_sample_name)
    logger.info("Volume: %s", _output_volume_name)
    logger.info("Chunks: %d", chunk_count)
    logger.info("Upload to: %s", ", ".join(upload_targets))
    if "hf-buckets" in upload_targets:
        logger.info("HF bucket: %s", resolved_bucket_id)

    if not keep_old_runs and not upload_only:
        logger.info("Cleaning volume...")
        cleanup_result = cleanup_volume.remote(current_run_id=current_run_id)
        logger.info("Removed %d old run(s)", cleanup_result["deleted"])

    if use_volume_manifest:
        logger.info("Staging manifest from volume (%s)...", resolved_sample_name)
        stage_result = stage_manifest_from_volume.remote(
            current_run_id, resolved_sample_name
        )
        logger.info("Staged %s bytes (volume-to-volume)", f"{stage_result['manifest_size']:,}")
    else:
        manifest_path_obj = Path(sample_manifest)
        sample_dir = manifest_path_obj.parent
        contract_path = sample_dir / "sample_contract.json"
        bin_summary_path = sample_dir / "bin_summary.csv"
        contract_json = contract_path.read_text() if contract_path.exists() else ""
        bin_summary_csv = bin_summary_path.read_text() if bin_summary_path.exists() else ""

        logger.info("Staging manifest from local...")
        manifest_bytes = manifest_path_obj.read_bytes()
        stage_manifest.remote(
            current_run_id,
            manifest_bytes,
            contract_json=contract_json,
            bin_summary_csv=bin_summary_csv,
        )
        logger.info("Staged %s bytes", f"{len(manifest_bytes):,}")

    logger.info("Dispatching pipeline to Modal...")
    result_json = run_materialize_pipeline.remote(
        run_id=current_run_id,
        sample_name=resolved_sample_name,
        chunk_count=chunk_count,
        upload_targets_json=json.dumps(upload_targets),
        hf_repo=hf_repo,
        hf_bucket_id=resolved_bucket_id,
        upload_only=upload_only,
        force_upload=force_upload,
    )

    result = json.loads(result_json)
    logger.info("Result:\n%s", json.dumps(result, indent=2))