scripts/cerebrium_dispatch.py

"""Dispatch KOfam scans and/or ESM-2 embeddings to Cerebrium.

Each Cerebrium replica is a stateless HTTP endpoint that handles one genome at
a time (`replica_concurrency = 1`). We fan out across `max_replicas` parallel
in-flight requests; results stream to JSONL as they arrive.

Usage:
    # Smoke test: 5 KOfam scans
    uv run python scripts/cerebrium_dispatch.py kofam --limit 5

    # Smoke test: 5 embeddings
    uv run python scripts/cerebrium_dispatch.py embed --limit 5

    # Full corpus (defaults: concurrency = max_replicas of the deployed app)
    uv run python scripts/cerebrium_dispatch.py kofam
    uv run python scripts/cerebrium_dispatch.py embed
"""
from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
import time
from pathlib import Path
from typing import Any

import httpx
import pandas as pd
import yaml

PROJECT_ID = "p-58781999"
REGION_HOST = "https://api.aws.us-east-1.cerebrium.ai"

APP_CONFIG = {
    "kofam": {
        "function": "scan_genome",
        "concurrency": 10,
        "out_path": Path("data/kofam_hits.jsonl"),
        "id_field": "genome_accession",
        "request_timeout": 180,
        "ok_keys": ("ko_hits",),
    },
    "embed": {
        "function": "embed_genome",
        "concurrency": 3,
        "out_path": Path("data/per_marker_embeddings.jsonl"),
        "id_field": "bacdive_id",
        "request_timeout": 600,
        "ok_keys": ("row",),
    },
}


def _read_access_token() -> str:
    env_token = os.environ.get("CEREBRIUM_API_KEY") or os.environ.get("CEREBRIUM_INFERENCE_KEY")
    if env_token:
        return env_token
    sys.exit(
        "Set CEREBRIUM_API_KEY to a JWT from the dashboard's API Keys section. "
        "The CLI's accesstoken doesn't work for inference endpoints."
    )


def _load_pending_kofam(limit: int) -> list[dict[str, Any]]:
    feats = pd.read_parquet("data/features.parquet")
    accs = feats["genome_accession"].dropna().astype(str).unique().tolist()
    done: set[str] = set()
    out_path = APP_CONFIG["kofam"]["out_path"]
    if out_path.exists():
        with open(out_path) as fh:
            for line in fh:
                try:
                    row = json.loads(line)
                except Exception:
                    continue
                acc = row.get("genome_accession") or row.get("accession")
                if acc:
                    done.add(str(acc))
    pending = [a for a in accs if a not in done]
    if limit:
        pending = pending[:limit]
    return [{"accession": a} for a in pending]


def _load_pending_embed(limit: int) -> list[dict[str, Any]]:
    import microbe_model.config as cfg
    pheno = pd.read_parquet("data/bacdive_phenotypes.parquet")
    has_genome = pheno["genome_accession"].notna()
    label_cols = list(cfg.PHENOTYPE_TARGETS.keys())
    has_label = pheno[label_cols].notna().any(axis=1)
    ready = pheno[has_genome & has_label].copy()
    ready["bacdive_id"] = ready["bacdive_id"].astype(int)

    done: set[int] = set()
    out_path = APP_CONFIG["embed"]["out_path"]
    if out_path.exists():
        with open(out_path) as fh:
            for line in fh:
                try:
                    done.add(int(json.loads(line)["bacdive_id"]))
                except Exception:
                    continue
    pending = ready[~ready["bacdive_id"].isin(done)]
    if limit:
        pending = pending.head(limit)
    return [
        {"bacdive_id": int(row["bacdive_id"]), "accession": str(row["genome_accession"])}
        for _, row in pending.iterrows()
    ]


async def _call_once(
    client: httpx.AsyncClient, app: str, payload: dict[str, Any], token: str,
    timeout: float,
) -> dict[str, Any]:
    url = f"{REGION_HOST}/v4/{PROJECT_ID}/{app}/{APP_CONFIG[app]['function']}"
    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
    resp = await client.post(url, headers=headers, json=payload, timeout=timeout)
    resp.raise_for_status()
    data = resp.json()
    if isinstance(data, dict) and "result" in data:
        return data["result"]
    return data


async def _worker(
    app: str, queue: asyncio.Queue, log_fh, results: dict[str, int],
    token: str, timeout: float, sem: asyncio.Semaphore,
):
    async with httpx.AsyncClient() as client:
        while True:
            payload = await queue.get()
            if payload is None:
                queue.task_done()
                return
            async with sem:
                start = time.time()
                for attempt in range(3):
                    try:
                        out = await _call_once(client, app, payload, token, timeout)
                        elapsed = time.time() - start
                        if isinstance(out, dict) and out.get("ok"):
                            log_fh.write(json.dumps(out.get("row") if app == "embed" else out) + "\n")
                            log_fh.flush()
                            results["ok"] += 1
                            results["elapsed_sum"] += elapsed
                        else:
                            results["fail"] += 1
                            reason = out.get("reason", "?") if isinstance(out, dict) else "non-dict"
                            print(f"  fail {payload}: {reason}", flush=True)
                        break
                    except httpx.HTTPStatusError as exc:
                        if exc.response.status_code in (429, 502, 503, 504):
                            await asyncio.sleep(2 ** attempt)
                            continue
                        results["fail"] += 1
                        print(f"  http {exc.response.status_code} {payload}: {exc.response.text[:200]}",
                              flush=True)
                        break
                    except (httpx.TimeoutException, httpx.TransportError) as exc:
                        if attempt < 2:
                            await asyncio.sleep(2 ** attempt)
                            continue
                        results["fail"] += 1
                        print(f"  timeout {payload}: {exc}", flush=True)
            queue.task_done()


async def _run(app: str, jobs: list[dict[str, Any]], concurrency: int):
    cfg = APP_CONFIG[app]
    token = _read_access_token()
    out_path: Path = cfg["out_path"]
    out_path.parent.mkdir(parents=True, exist_ok=True)
    queue: asyncio.Queue = asyncio.Queue()
    for j in jobs:
        await queue.put(j)
    for _ in range(concurrency):
        await queue.put(None)

    results = {"ok": 0, "fail": 0, "elapsed_sum": 0.0}
    sem = asyncio.Semaphore(concurrency)
    t0 = time.time()
    with open(out_path, "a") as log_fh:
        workers = [
            asyncio.create_task(_worker(
                app, queue, log_fh, results, token, cfg["request_timeout"], sem,
            ))
            for _ in range(concurrency)
        ]
        last_report = t0
        while any(not w.done() for w in workers):
            await asyncio.sleep(15)
            now = time.time()
            done = results["ok"] + results["fail"]
            if done == 0:
                continue
            rate = done / (now - t0)
            remaining = len(jobs) - done
            eta = remaining / rate if rate > 0 else float("inf")
            if now - last_report >= 30:
                print(
                    f"  [{int(now - t0)}s] ok={results['ok']:,} fail={results['fail']:,} "
                    f"rate={rate:.2f}/s eta={int(eta/60)}min", flush=True,
                )
                last_report = now
        await asyncio.gather(*workers)
    elapsed = time.time() - t0
    avg_per_ok = results["elapsed_sum"] / max(results["ok"], 1)
    print(f"\nDone in {elapsed/60:.1f} min. ok={results['ok']:,} fail={results['fail']:,} "
          f"avg/ok={avg_per_ok:.1f}s")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("app", choices=list(APP_CONFIG.keys()))
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--concurrency", type=int, default=0,
                        help="Override default (default: app's max_replicas)")
    args = parser.parse_args()

    if args.app == "kofam":
        jobs = _load_pending_kofam(args.limit)
    else:
        jobs = _load_pending_embed(args.limit)
    if not jobs:
        print("Nothing to do.")
        return
    concurrency = args.concurrency or APP_CONFIG[args.app]["concurrency"]
    print(f"Dispatching {len(jobs):,} jobs to Cerebrium app '{args.app}' "
          f"at concurrency={concurrency}.")
    print(f"  Output: {APP_CONFIG[args.app]['out_path']}")
    asyncio.run(_run(args.app, jobs, concurrency))


if __name__ == "__main__":
    main()