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

"""Upload locally-drawn sample manifests to the soc149-drawn-samples Modal volume."""

from __future__ import annotations

from pathlib import Path

import modal

app = modal.App("soc149-upload-local-samples")

samples_volume = modal.Volume.from_name(
    "soc149-drawn-samples", create_if_missing=True
)

SAMPLES_MOUNT = "/samples"


@app.function(
    image=modal.Image.debian_slim(python_version="3.12"),
    volumes={SAMPLES_MOUNT: samples_volume},
    timeout=300,
    cpu=1,
)
def write_sample(
    sample_name: str,
    manifest_bytes: bytes,
    contract_json: str,
    bin_summary_csv: str,
) -> dict[str, object]:
    sample_dir = Path(SAMPLES_MOUNT) / sample_name
    sample_dir.mkdir(parents=True, exist_ok=True)
    (sample_dir / "working_sample_manifest.parquet").write_bytes(manifest_bytes)
    (sample_dir / "sample_contract.json").write_text(contract_json)
    (sample_dir / "bin_summary.csv").write_text(bin_summary_csv)
    samples_volume.commit()
    return {"sample_name": sample_name, "manifest_bytes": len(manifest_bytes)}


@app.local_entrypoint()
def main():
    import logging

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

    samples_dir = Path("data/samples")
    sample_dirs = sorted(samples_dir.iterdir()) if samples_dir.exists() else []

    if not sample_dirs:
        print("No samples found in data/samples/")
        return

    for sample_dir in sample_dirs:
        manifest_path = sample_dir / "working_sample_manifest.parquet"
        if not manifest_path.exists():
            continue
        sample_name = sample_dir.name
        manifest_bytes = manifest_path.read_bytes()
        contract_json = (sample_dir / "sample_contract.json").read_text()
        bin_summary_csv = (sample_dir / "bin_summary.csv").read_text()
        print(f"Uploading {sample_name} ({len(manifest_bytes):,} bytes)...")
        result = write_sample.remote(sample_name, manifest_bytes, contract_json, bin_summary_csv)
        print(f"  Done: {result}")

    print("All local samples uploaded to Modal volume: soc149-drawn-samples")