dee/core/modal_client.py

"""Client for the Modal GPU backend (repo path ``modal/``).

Why this exists: the Space runs on a free CPU tier where a 650M forward pass
would monopolise both vCPUs behind ``scoring.SCORING_SEMAPHORE``, so
``scoring.effective_model`` currently downgrades a 650M request to 35M. A
picker built on that would be a lie — the user selects 650M and reads a number
produced by a different model. Routing the large models to rented GPUs is what
makes the choice honest instead of decorative.

Shape of the contract (see ``modal/orchestrate.py``):

    POST {MODAL_SUBMIT_URL}   {"wt_sequence": ..., "mutations": [...], ...}
      -> {"job_id": ...}                       returns immediately
    GET  {MODAL_STATUS_URL}?job_id=...
      -> {"status": "sieving"|"extracting"|"done"|"error", ...}

Both carry ``Authorization: Bearer <ORCHESTRATOR_API_KEY>``; unauthenticated
requests get a 401. Polling rather than one blocking call because a sieve plus
an ensemble pass runs for minutes, and holding an HTTP request open that long
loses to proxies and browsers alike.

**Not configured is not the same as broken.** When the required environment
variables are absent this module reports that plainly (:class:`NotConfigured`)
so a caller can say "the GPU backend isn't set up" instead of silently
substituting a smaller model — the failure this whole module exists to prevent.

**Two pipelines, one client.** ``orchestrate.py`` exposes a second, separate
submit/status pair for the Evo 2 DNA backend (``/submit_dna``/``/status_dna``
— see that module's docstring for why it doesn't chain into the protein
pipeline above). Every function here takes a ``pipeline`` kwarg
(``"protein"`` by default, or ``"dna"``) selecting which URL pair to use;
the API key is shared (same orchestrator, same secret). Existing callers
that never pass ``pipeline`` are unaffected — this is additive, not a
behavior change to the protein path.
"""

from __future__ import annotations

import logging
import os
import time
from typing import Any, Dict, Optional, Tuple

from dee.core import compute_budget as _budget

logger = logging.getLogger("dee.modal_client")

SUBMIT_URL_ENV = "MODAL_SUBMIT_URL"
STATUS_URL_ENV = "MODAL_STATUS_URL"
DNA_SUBMIT_URL_ENV = "MODAL_DNA_SUBMIT_URL"
DNA_STATUS_URL_ENV = "MODAL_DNA_STATUS_URL"
API_KEY_ENV = "ORCHESTRATOR_API_KEY"

# A cold container pulls its image and loads weights before the first token of
# work. Generous, because the alternative is reporting failure on a run that
# was merely starting.
DEFAULT_POLL_INTERVAL = 5.0
DEFAULT_TIMEOUT = 900.0
CONNECT_TIMEOUT = 30.0

# The DNA pipeline needs its own, LONGER budget. Both Evo 2 functions carry
# a 900s Modal timeout, so a client that also waits 900s gives up at exactly
# the moment the function would be killed — leaving zero room for the time
# before the function even starts running: queueing, pulling a multi-GB CUDA
# image, and downloading a 7B checkpoint into a cold Volume. A client
# timeout equal to the server's reports failure on a run that was merely
# starting, which is precisely the failure this module's docstring says it
# exists to prevent (and the same shape as the TimeoutError outage the
# orchestrator shipped — see modal/README.md).
#
# 1500s = the function's own 900s ceiling plus 600s of cold-start headroom.
# Nothing waits on this synchronously (orchestrator runs tools on a
# background thread and the client polls), so a long ceiling costs a parked
# thread, not a blocked user.
DNA_TIMEOUT = 1500.0


class NotConfigured(RuntimeError):
    """The GPU backend has no URLs/key in the environment.

    Distinct from a failure: nothing was attempted. Callers should surface this
    as "not set up" and MUST NOT quietly fall back to a smaller model, which
    would attribute one model's numbers to another.
    """


class ModalError(RuntimeError):
    """The backend was reachable but the run did not succeed."""


def _url_envs(pipeline: str) -> Tuple[str, str]:
    if pipeline == "dna":
        return DNA_SUBMIT_URL_ENV, DNA_STATUS_URL_ENV
    return SUBMIT_URL_ENV, STATUS_URL_ENV


def is_configured(pipeline: str = "protein") -> bool:
    submit_env, status_env = _url_envs(pipeline)
    return all(os.environ.get(v) for v in (submit_env, status_env, API_KEY_ENV))


def _config(pipeline: str = "protein") -> Dict[str, str]:
    submit_env, status_env = _url_envs(pipeline)
    missing = [v for v in (submit_env, status_env, API_KEY_ENV)
               if not os.environ.get(v)]
    if missing:
        raise NotConfigured(
            f"GPU backend ({pipeline}) not configured — missing " + ", ".join(missing) +
            ". Set them on the Space (see modal/README.md); until then only the "
            "35M model is available, and that is reported rather than substituted."
        )
    return {
        "submit": os.environ[submit_env].rstrip("/"),
        "status": os.environ[status_env].rstrip("/"),
        "key": os.environ[API_KEY_ENV],
    }


def _headers(key: str) -> Dict[str, str]:
    return {"Authorization": f"Bearer {key}", "Content-Type": "application/json"}


# ── Reachability, as distinct from configuration ─────────────────────────
# is_configured() answers "are the env vars set". That is NOT the same
# question as "would a call arrive", and treating them as one shipped a live
# lie: the DNA URLs were set on the Space before the Evo 2 endpoints had ever
# been deployed, so /api/models advertised both DNA tiers as available while
# both URLs returned 404. The agent was handed a tool that could only fail.
#
# The probe: GET the status URL with NO query string.
#   · 404 -> the web endpoint does not exist (undeployed)
#   · 422 -> it exists and rejected a missing job_id (FastAPI validation)
#   · 401 -> it exists and wanted credentials
# Anything that is not a 404 proves an endpoint is there. Verified live
# against both pipelines while the DNA one was undeployed.
#
# Deliberately unauthenticated: a liveness check must never put the API key
# on the wire, and 422 is returned before auth runs anyway.
# 5s, and two attempts. An earlier 3s single-shot produced a live FALSE
# NEGATIVE on a healthy endpoint — a cold TLS handshake to Modal ran past the
# deadline and the probe reported a deployed pipeline as missing. Hiding a
# working tool because a handshake was slow is its own kind of dishonesty,
# and the 300s cache means these numbers cost almost nothing in aggregate.
PROBE_TTL = 300.0
PROBE_TIMEOUT = 5.0
PROBE_ATTEMPTS = 2

_probe_cache: Dict[str, Tuple[float, bool]] = {}


def reachable(pipeline: str = "protein", *, force: bool = False) -> bool:
    """True when this pipeline's endpoints actually exist right now.

    Cached for PROBE_TTL because active_tool_specs() recomputes the tool list
    on EVERY agent step — an uncached probe would add a network round-trip to
    each one. Never raises: a tool list must not be breakable by a blip.

    Fails CLOSED on an unreachable backend, but only when nothing better is
    known: a cached True survives its TTL being refreshed by a transient
    error, so a momentary network fault cannot yank a working tool mid-run.
    """
    import requests

    if not is_configured(pipeline):
        return False

    now = time.monotonic()
    key = pipeline
    cached = _probe_cache.get(key)
    if not force and cached is not None and now - cached[0] < PROBE_TTL:
        return cached[1]

    _, status_env = _url_envs(pipeline)
    url = (os.environ.get(status_env) or "").rstrip("/")
    for attempt in range(PROBE_ATTEMPTS):
        try:
            code = requests.get(url, timeout=PROBE_TIMEOUT).status_code
            alive = code != 404
            _probe_cache[key] = (now, alive)
            if not alive:
                logger.warning(
                    "%s pipeline is configured but its status endpoint 404s "
                    "(%s) — reporting it as unavailable rather than offering a "
                    "tool that cannot work.", pipeline, url)
            return alive
        except Exception:  # noqa: BLE001 — availability must never raise
            logger.debug("reachability probe %d/%d failed for %s",
                         attempt + 1, PROBE_ATTEMPTS, pipeline, exc_info=True)
    # Keep a previous confirmation rather than flapping on a bad minute.
    return bool(cached[1]) if cached is not None else False


def submit(payload: Dict[str, Any], *, pipeline: str = "protein",
           timeout: float = CONNECT_TIMEOUT) -> str:
    """Start a run. Returns the job id; does not wait."""
    import requests

    cfg = _config(pipeline)
    resp = requests.post(cfg["submit"], json=payload,
                         headers=_headers(cfg["key"]), timeout=timeout)
    if resp.status_code == 401:
        raise ModalError("GPU backend rejected the API key (401). The value on "
                         "the Space and the Modal secret have diverged.")
    resp.raise_for_status()
    body = resp.json()
    job_id = str(body.get("job_id") or "")
    if not job_id:
        raise ModalError(f"Backend accepted the job but returned no job_id: {body!r}")
    return job_id


def status(job_id: str, *, pipeline: str = "protein",
           timeout: float = CONNECT_TIMEOUT) -> Dict[str, Any]:
    """One poll. Raises rather than inventing a state on a transport failure."""
    import requests

    cfg = _config(pipeline)
    resp = requests.get(cfg["status"], params={"job_id": job_id},
                        headers=_headers(cfg["key"]), timeout=timeout)
    if resp.status_code == 401:
        raise ModalError("GPU backend rejected the API key (401).")
    resp.raise_for_status()
    return dict(resp.json())


def run(payload: Dict[str, Any], *, pipeline: str = "protein",
        poll_interval: float = DEFAULT_POLL_INTERVAL,
        total_timeout: Optional[float] = None,
        on_progress: Optional[Any] = None) -> Dict[str, Any]:
    """Submit and poll to completion.

    ``on_progress(state)`` is called with each distinct status so a caller can
    show real progress — a multi-minute GPU run with no signal is
    indistinguishable from a hang, which is the complaint the elapsed-time work
    on the Interaction Radar already answered once.

    Raises :class:`NotConfigured` if the backend was never set up, or
    :class:`ModalError` on a reported error or a timeout. It never returns a
    partial result dressed as a complete one.
    """
    # Per-pipeline default rather than one number for both — see DNA_TIMEOUT.
    # An explicit argument still wins, so a caller that knows better can say so.
    if total_timeout is None:
        total_timeout = DNA_TIMEOUT if pipeline == "dna" else DEFAULT_TIMEOUT

    job_id = submit(payload, pipeline=pipeline)
    started = time.monotonic()
    deadline = started + total_timeout
    last_state = None

    def _meter() -> None:
        """Bill the run for the GPU it just rented.

        Metered on EVERY exit — done, error and timeout alike. A failed job
        still occupied a container, and a budget that only counts successes
        can be walked past by a caller whose calls keep failing.
        """
        try:
            _budget.record(pipeline or "gpu", time.monotonic() - started,
                           pipeline=pipeline)
        except Exception:  # noqa: BLE001 — accounting must never fail the work
            logger.debug("gpu metering failed; continuing", exc_info=True)

    while True:
        state = status(job_id, pipeline=pipeline)
        # The orchestrator calls this field "stage" (orchestrate_logic.
        # render_status), NOT "status". Reading the wrong key is why every
        # remote GPU call hung: phase was permanently "", so the loop below
        # never matched "done" or "error" and just polled until the deadline,
        # then reported a timeout on a job that had usually finished in
        # seconds. Prometheus scoring never completed end-to-end because of
        # this one word.
        #
        # Both names are accepted so a rename on either side cannot silently
        # reintroduce a 15-minute hang, and an unrecognised shape now fails
        # FAST and loudly (below) instead of looking like a slow GPU.
        phase = str(state.get("stage") or state.get("status") or "")
        if not phase:
            raise ModalError(
                "GPU backend returned a status with no stage/status field "
                f"(keys: {sorted(state)!r}). Refusing to poll a response this "
                "code cannot read — a shape mismatch must not present as a "
                "slow job.")

        if phase != last_state:
            last_state = phase
            logger.info("modal job %s: %s", job_id, phase)
            if on_progress is not None:
                try:
                    on_progress(state)
                except Exception:  # noqa: BLE001 — progress must never fail a run
                    logger.debug("on_progress raised; continuing", exc_info=True)

        if phase == "done":
            _meter()
            return state
        if phase == "error":
            _meter()
            raise ModalError(str(state.get("error") or "GPU run failed")[:400])

        if time.monotonic() >= deadline:
            _meter()
            # Say what was actually seen. "Timed out" alone leaves the caller
            # unable to tell a slow run from a stuck one.
            raise ModalError(
                f"GPU run {job_id} did not finish within {total_timeout:.0f}s "
                f"(last status: {phase or 'unknown'}). It may still be running "
                f"on Modal; the job id is preserved above.")
        time.sleep(poll_interval)