api/sim_backend.py

"""Where jobs execute: Modal (default) or RunPod serverless.

Both backends expose the same two calls, so the API does not care which one is configured:

    submit(job_input, execution_timeout_s, function) -> {"id": ..., "status": "IN_QUEUE"}
    status(job_id) -> {"status", "output"?, "error"?, "executionTime"?, "delayTime"?}

Times are in milliseconds. The worker also reports its own wall time as ``output["worker_seconds"]``,
which is used when a backend does not report execution time itself (Modal).

`function` names the worker function (see worker/modal_app.py). Modal deploys all of them; RunPod's single
endpoint image only has the MuJoCo ones (``RunPodClient.FUNCTIONS``).
"""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import Protocol

from .config import Settings

TERMINAL = {"COMPLETED", "FAILED", "CANCELLED", "TIMED_OUT"}
BACKENDS = ("modal", "runpod")

log = logging.getLogger("rosdiff.runs")


class BackendError(RuntimeError):
    """The backend could not be reached or refused the request. The run's status is left unchanged."""


class SimBackend(Protocol):
    name: str

    def supports(self, function: str) -> bool: ...

    def submit(self, job_input: dict, execution_timeout_s: int, function: str) -> dict: ...

    def status(self, job_id: str) -> dict: ...


def modal_credentials_present() -> bool:
    if os.environ.get("MODAL_TOKEN_ID") and os.environ.get("MODAL_TOKEN_SECRET"):
        return True
    return Path(os.environ.get("MODAL_CONFIG_PATH", Path.home() / ".modal.toml")).is_file()


def make_backend(settings: Settings) -> SimBackend | None:
    """The configured backend, or None (with a log line saying why) when runs are not set up."""
    if settings.sim_backend == "modal":
        if not modal_credentials_present():
            log.warning("simulation runs disabled: set MODAL_TOKEN_ID/MODAL_TOKEN_SECRET or run `modal token new`")
            return None
        from .modal_client import ModalClient

        return ModalClient(settings.modal_app)
    if settings.sim_backend == "runpod":
        if not (settings.runpod_api_key and settings.runpod_endpoint_id):
            log.warning("simulation runs disabled: RUNPOD_API_KEY and RUNPOD_ENDPOINT_ID are not set")
            return None
        from .runpod_client import RunPodClient

        return RunPodClient(settings.runpod_api_key, settings.runpod_endpoint_id)
    raise ValueError(f"SIM_BACKEND must be one of {', '.join(BACKENDS)}, got {settings.sim_backend!r}")