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

"""Gate A verification: confirm Modal image, secrets, and R2 access."""

from __future__ import annotations

import logging
import os
from pathlib import Path


from .config import BLOOM_IMAGE_PATH, R2_PREFIX
from .soc127_app import app, hf_secret, image, r2_base_path, r2_mount, r2_secret

logger = logging.getLogger("verify_setup")


@app.function(
    image=image,
    secrets=[hf_secret, r2_secret],
    volumes={"/r2": r2_mount},
    timeout=600,
    memory=4096,
)
def run_checks() -> dict[str, object]:
    results: dict[str, object] = {}

    bloom_path = Path(BLOOM_IMAGE_PATH)
    results["bloom"] = {
        "exists": bloom_path.exists(),
        "size_bytes": bloom_path.stat().st_size if bloom_path.exists() else 0,
    }

    try:
        from huggingface_hub import HfApi

        api = HfApi(token=os.environ.get("HF_TOKEN"))
        user_info = api.whoami()
        results["hf_auth"] = {
            "ok": True,
            "username": user_info.get("name", "unknown"),
        }
    except Exception as exc:
        results["hf_auth"] = {"ok": False, "error": str(exc)}

    r2 = r2_base_path()
    sentinel = r2 / R2_PREFIX / "_verify_setup_sentinel.txt"
    try:
        sentinel.parent.mkdir(parents=True, exist_ok=True)
        sentinel.write_text("ok", encoding="utf-8")
        readback = sentinel.read_text(encoding="utf-8")
        sentinel.unlink()
        results["r2_mount"] = {"ok": readback == "ok"}
    except Exception as exc:
        results["r2_mount"] = {"ok": False, "error": str(exc)}

    try:
        from dolma.provenance import BloomIndex

        bloom = BloomIndex.load(bloom_path)
        results["bloom_load"] = {
            "ok": True,
            "expected_items": bloom.expected_items,
            "bit_count": bloom.bit_count,
        }
    except Exception as exc:
        results["bloom_load"] = {"ok": False, "error": str(exc)}

    return results


@app.local_entrypoint()
def main():
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )
    results = run_checks.remote()
    all_ok = True
    checks = [
        ("Image + Bloom file", "bloom", lambda r: r.get("exists", False)),
        ("HF authentication", "hf_auth", lambda r: r.get("ok", False)),
        ("R2 mount read/write", "r2_mount", lambda r: r.get("ok", False)),
        ("Bloom filter load", "bloom_load", lambda r: r.get("ok", False)),
    ]
    for label, key, check_fn in checks:
        result = results.get(key, {})
        passed = check_fn(result)
        status = "PASS" if passed else "FAIL"
        if not passed:
            all_ok = False
        logger.info("%s: %s  %s", status, label, result)
    if all_ok:
        logger.info("Gate A: all checks passed")
    else:
        logger.warning("Gate A: some checks failed")