scripts/biohub_esm_lib/routing.py

"""Explicit scientific and compute routing decision table."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from typing import Literal

from .constants import MODAL_BINDER_EXAMPLE_REVISION
from .errors import ValidationError

Task = Literal["atlas", "esmc", "fold", "binder-design"]
Route = Literal["atlas-api", "atlas-s3", "biohub", "modal", "self-hosted"]


@dataclass(frozen=True)
class RouteRequest:
    task: Task
    item_count: int = 1
    long_running: bool = False
    bulk_dataset: bool = False
    private: bool = False
    offline: bool = False
    data_residency: bool = False
    custom_model: bool = False
    fine_tune: bool = False
    sustained_workload: bool = False
    has_msa: bool = False
    accuracy_priority: bool = False
    owns_gpu: bool = False


@dataclass(frozen=True)
class RouteResult:
    route: Route
    model: str | None
    required_authentication: str | None
    rationale: tuple[str, ...]
    warnings: tuple[str, ...] = ()
    workflow: str | None = None
    source_revision: str | None = None

    def as_dict(self) -> dict[str, object]:
        return asdict(self)


def route_request(request: RouteRequest) -> RouteResult:
    if request.task not in {"atlas", "esmc", "fold", "binder-design"}:
        raise ValidationError(f"unsupported task: {request.task}")
    if request.item_count < 1:
        raise ValidationError("item_count must be at least 1")

    private_execution = any(
        (
            request.private,
            request.offline,
            request.data_residency,
            request.custom_model,
            request.fine_tune,
            request.sustained_workload,
        )
    )

    if request.task == "atlas":
        if private_execution or request.bulk_dataset:
            rationale = (
                "Private/offline Atlas work must stage the public dataset locally; do not send "
                "a private query to the public API."
                if private_execution
                else "Atlas bulk datasets are public through anonymous S3."
            )
            return RouteResult(
                route="atlas-s3",
                model=None,
                required_authentication=None,
                rationale=(rationale,),
                warnings=(
                    "Atlas data is CC-BY-4.0; preserve attribution and dataset provenance.",
                    "Atlas data is very large; plan the transfer, storage, and local-search footprint.",
                    (
                        "Before any multi-gigabyte or multi-terabyte Atlas S3 transfer, "
                        "freeze the exact source prefix, destination, estimated bytes, "
                        "storage and egress impact, and cost ceiling; then obtain separate "
                        "explicit current-turn confirmation."
                    ),
                ),
            )
        return RouteResult(
            route="atlas-api",
            model=None,
            required_authentication=None,
            rationale=(
                "Protein discovery, the Atlas feature catalog, and cluster exploration belong to Atlas.",
            ),
            warnings=("Atlas is a v1alpha1 API; validate schemas and preserve raw responses.",),
        )

    if request.task == "binder-design":
        if private_execution or request.owns_gpu:
            return RouteResult(
                route="self-hosted",
                model=None,
                required_authentication=None,
                rationale=(
                    "Binder design is not available through the Biohub managed API.",
                    "Private/custom execution or owned GPUs favor pinned open weights.",
                    (
                        "The workflow uses pinned experimental inversion and critic checkpoints, "
                        "not the standard ESMFold2 inference checkpoint."
                    ),
                ),
                warnings=(
                    (
                        "A one-seed smoke test is not representative; useful campaigns "
                        "usually need hundreds and often about 1,000 designs."
                    ),
                    "Estimate runtime and cost, then obtain explicit confirmation before material GPU spend.",
                ),
                workflow="self-hosted-esmfold2-binder-design",
                source_revision=MODAL_BINDER_EXAMPLE_REVISION,
            )
        return RouteResult(
            route="modal",
            model=None,
            required_authentication="MODAL_TOKEN_ID + MODAL_TOKEN_SECRET or Modal profile",
            rationale=(
                "Binder design is not available through the Biohub managed API.",
                "Modal supports parallel, durable execution of the released open-weight workflow.",
                (
                    "The workflow uses pinned experimental inversion and critic checkpoints, "
                    "not the standard ESMFold2 inference checkpoint."
                ),
            ),
            warnings=(
                (
                    "A one-seed smoke test is not representative; useful campaigns "
                    "usually need hundreds and often about 1,000 designs."
                ),
                "Estimate runtime and cost, then obtain explicit confirmation before material GPU spend.",
            ),
            workflow=f"modal-esmfold2-binder-design@{MODAL_BINDER_EXAMPLE_REVISION}",
            source_revision=MODAL_BINDER_EXAMPLE_REVISION,
        )

    if private_execution:
        hf_model = "biohub/ESMC-600M" if request.task == "esmc" else "biohub/ESMFold2"
        return RouteResult(
            route="self-hosted",
            model=hf_model,
            required_authentication=None,
            rationale=(
                (
                    "Private, offline, data-resident, customized, fine-tuned, or "
                    "sustained work belongs on user-owned compute."
                ),
                "Public weights do not require HF_TOKEN; it is optional for authenticated Hub access.",
            ),
        )

    if request.task == "esmc":
        if request.owns_gpu:
            return RouteResult(
                route="self-hosted",
                model="biohub/ESMC-600M",
                required_authentication=None,
                rationale=(
                    (
                        "Suitable user-owned GPUs favor pinned self-hosted public weights; "
                        "HF_TOKEN is optional for authenticated Hub access."
                    ),
                ),
            )
        return RouteResult(
            route="biohub",
            model="esmc-600m-2024-12",
            required_authentication="ESM_API_KEY",
            rationale=(
                (
                    "Biohub managed ESMC is the default for public inference, including "
                    "bounded concurrent calls that the managed backend can auto-batch."
                ),
                "The plugin does not ship an ESMC Modal function.",
            ),
        )

    # This is a folding orchestration heuristic, not an account quota. Account
    # credits and rate limits remain specific to the Biohub developer console at
    # https://biohub.ai/developer-console (keys: /developer-console/api-keys).
    scale_out = request.long_running or request.item_count > 32
    if scale_out:
        model = (
            "biohub/ESMFold2"
            if request.accuracy_priority or request.has_msa
            else "biohub/ESMFold2-Fast"
        )
        return RouteResult(
            route="self-hosted" if request.owns_gpu else "modal",
            model=model,
            required_authentication=(
                None if request.owns_gpu else "MODAL_TOKEN_ID + MODAL_TOKEN_SECRET or Modal profile"
            ),
            rationale=(
                "Independent bulk, parallel, or long-running inference should use durable scale-out compute.",
                (
                    "Suitable user-owned GPUs favor pinned self-hosted public weights; "
                    "HF_TOKEN is optional for authenticated Hub access."
                    if request.owns_gpu
                    else "Public weights on Modal do not require ESM_API_KEY."
                ),
            ),
            warnings=(
                ()
                if request.owns_gpu
                else (
                    (
                        "Before any Modal spawn, refresh current Modal pricing and "
                        "account/payment readiness; freeze and show the exact Modal "
                        "app/function, GPU type, job count, concurrency, timeout, "
                        "persistence plan, and cost ceiling."
                    ),
                    (
                        "Obtain separate explicit current-turn confirmation after that "
                        "frozen scope and current cost review. --confirm-cost is only an "
                        "execution backstop, not authorization."
                    ),
                )
            ),
        )

    model = (
        "esmfold2-2026-05"
        if request.accuracy_priority or request.has_msa
        else "esmfold2-fast-2026-05"
    )
    rationale = ["Biohub managed ESMFold2 is the default for modest interactive folding."]
    if request.has_msa:
        rationale.append("MSA conditioning requires the full ESMFold2 model, not Fast.")
    elif request.accuracy_priority:
        rationale.append("Accuracy-priority targets use full ESMFold2; Fast favors throughput.")
    else:
        rationale.append("ESMFold2-Fast favors single-sequence latency and throughput.")
    return RouteResult(
        route="biohub",
        model=model,
        required_authentication="ESM_API_KEY",
        rationale=tuple(rationale),
        warnings=(
            "Predicted structures are static hypotheses, not experimental truth or molecular dynamics.",
        ),
    )