gateway_client.py

"""Lazy token-authenticated client for the private AutoML Space."""

from __future__ import annotations

import os
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Mapping

from gradio_client import Client, handle_file
from huggingface_hub import HfApi

from gateway_contract import PROTOCOL_VERSION, validate_request
from gateway_routing import VALID_PROFILES


class BackendUnavailable(RuntimeError):
    """The private backend cannot safely accept this request."""


class BackendTimeout(BackendUnavailable):
    """The private backend did not become ready before the wake deadline."""


class BackendProtocolMismatch(BackendUnavailable):
    """Gateway and backend contract versions differ."""


@dataclass(frozen=True)
class BackendResult:
    envelope: dict[str, Any]
    package_path: str | None = None
    cleanup_files: tuple[str, ...] = ()


def _default_api_factory(token: str) -> HfApi:
    return HfApi(token=token)


def _default_client_factory(repo_id: str, *, token: str) -> Client:
    return Client(repo_id, token=token, verbose=False)


def _remote_file_path(value: Any) -> str | None:
    if not value:
        return None
    if isinstance(value, (str, os.PathLike)):
        return str(value)
    if isinstance(value, Mapping):
        for key in ("path", "name", "orig_name"):
            if value.get(key):
                return str(value[key])
    for attribute in ("path", "name"):
        candidate = getattr(value, attribute, None)
        if candidate:
            return str(candidate)
    return None


class GatewayBackendClient:
    """Call one private Space without contacting it during construction."""

    TERMINAL_OPERATOR_STAGES = {
        "PAUSED",
        "BUILD_ERROR",
        "RUNTIME_ERROR",
        "CONFIG_ERROR",
        "NO_APP_FILE",
    }
    READY_STAGES = {"RUNNING"}
    SPACE_ID_ENV = {
        "cpu": "AUTOML_CPU_BACKEND_SPACE_ID",
        "gpu": "AUTOML_GPU_BACKEND_SPACE_ID",
    }
    SPACE_ID_DEFAULT = {
        "cpu": "automl-team/AutoML-cpu",
        "gpu": "automl-team/AutoML-core",
    }

    def __init__(
        self,
        *,
        repo_id: str,
        token: str,
        expected_profile: str,
        wake_timeout_seconds: float = 900,
        poll_interval_seconds: float = 2,
        api_factory: Callable[[str], Any] = _default_api_factory,
        client_factory: Callable[..., Any] = _default_client_factory,
        file_wrapper: Callable[[str], Any] = handle_file,
        clock: Callable[[], float] = time.monotonic,
        sleeper: Callable[[float], None] = time.sleep,
    ) -> None:
        self.repo_id = str(repo_id or "").strip()
        self._token = str(token or "").strip()
        self.expected_profile = str(expected_profile or "").strip().lower()
        if self.expected_profile not in VALID_PROFILES:
            raise BackendUnavailable("Private backend execution profile is invalid.")
        self.wake_timeout_seconds = max(0.1, float(wake_timeout_seconds))
        self.poll_interval_seconds = max(0.05, float(poll_interval_seconds))
        self._api_factory = api_factory
        self._client_factory = client_factory
        self._file_wrapper = file_wrapper
        self._clock = clock
        self._sleeper = sleeper
        self._api_instance: Any = None
        self._client_instance: Any = None

    def __repr__(self) -> str:
        return (
            f"GatewayBackendClient(repo_id={self.repo_id!r}, "
            f"expected_profile={self.expected_profile!r})"
        )

    @classmethod
    def from_env(cls, profile: str) -> "GatewayBackendClient":
        normalized_profile = str(profile or "").strip().lower()
        if normalized_profile not in VALID_PROFILES:
            raise BackendUnavailable("Private backend execution profile is invalid.")
        repo_id = os.getenv(
            cls.SPACE_ID_ENV[normalized_profile],
            cls.SPACE_ID_DEFAULT[normalized_profile],
        ).strip()
        token = os.getenv("HF_BACKEND_TOKEN", "").strip()
        if not repo_id:
            raise BackendUnavailable("Private backend Space ID is not configured.")
        if not token:
            raise BackendUnavailable("Private backend access is not configured.")
        try:
            timeout = float(
                os.getenv("AUTOML_BACKEND_WAKE_TIMEOUT_SECONDS", "900")
            )
        except ValueError:
            timeout = 900
        return cls(
            repo_id=repo_id,
            token=token,
            expected_profile=normalized_profile,
            wake_timeout_seconds=timeout,
        )

    def _api(self) -> Any:
        if self._api_instance is None:
            self._api_instance = self._api_factory(self._token)
        return self._api_instance

    def _client(self) -> Any:
        if self._client_instance is None:
            self._client_instance = self._client_factory(
                self.repo_id, token=self._token
            )
        return self._client_instance

    @staticmethod
    def _stage(value: Any) -> str:
        normalized = str(value or "UNKNOWN").strip().upper()
        if "." in normalized:
            normalized = normalized.rsplit(".", 1)[-1]
        return normalized

    def _runtime_stage(self) -> str:
        try:
            runtime = self._api().get_space_runtime(
                self.repo_id, token=self._token
            )
        except Exception as exc:
            detail = str(exc).lower()
            if any(marker in detail for marker in ("401", "403", "404", "unauthorized", "forbidden")):
                raise BackendUnavailable(
                    "Private backend access was rejected. Check the gateway token and Space visibility."
                ) from None
            raise
        return self._stage(getattr(runtime, "stage", None))

    @staticmethod
    def _emit(progress: Callable[[str], None] | None, message: str) -> None:
        if progress is not None:
            progress(message)

    def _verify_health(self, value: Any) -> dict[str, Any]:
        if not isinstance(value, Mapping):
            raise BackendUnavailable("Private backend returned an invalid health response.")
        response = dict(value)
        if response.get("protocol_version") != PROTOCOL_VERSION:
            raise BackendProtocolMismatch(
                "Gateway and backend versions differ. Deploy matching revisions before running jobs."
            )
        if response.get("ok") is not True:
            raise BackendUnavailable("Private backend health verification failed.")
        if response.get("execution_profile") != self.expected_profile:
            raise BackendProtocolMismatch(
                "Private backend execution profile does not match the selected route."
            )
        capabilities = set(response.get("capabilities") or [])
        required = {"existing_training", "uploaded_training", "package_prediction"}
        if not required.issubset(capabilities):
            raise BackendProtocolMismatch(
                "Gateway and backend capabilities differ. Deploy matching revisions before running jobs."
            )
        return response

    def ensure_ready(
        self, progress: Callable[[str], None] | None = None
    ) -> dict[str, Any]:
        """Wake with health traffic and poll only before a model endpoint is submitted."""

        deadline = self._clock() + self.wake_timeout_seconds
        profile_label = self.expected_profile.upper()
        announced_wake = False
        while True:
            if self._clock() >= deadline:
                raise BackendTimeout(
                    f"Private {profile_label} backend did not become ready before the wake timeout."
                )
            try:
                stage = self._runtime_stage()
            except BackendUnavailable:
                raise
            except Exception:
                stage = "UNKNOWN"
            if stage in self.TERMINAL_OPERATOR_STAGES:
                raise BackendUnavailable(
                    "Private backend needs operator attention before it can accept jobs."
                )
            if stage not in self.READY_STAGES and not announced_wake:
                self._emit(progress, f"{profile_label} backend is waking up…")
                announced_wake = True
            try:
                response = self._client().predict(
                    PROTOCOL_VERSION, api_name="/v1/health"
                )
                verified = self._verify_health(response)
                self._emit(progress, f"{profile_label} backend ready—starting job…")
                return verified
            except BackendProtocolMismatch:
                raise
            except Exception:
                self._client_instance = None
            remaining = deadline - self._clock()
            if remaining <= 0:
                raise BackendTimeout(
                    f"Private {profile_label} backend did not become ready before the wake timeout."
                )
            self._sleeper(min(self.poll_interval_seconds, remaining))

    def _submit_once(self, api_name: str, *args: Any) -> Any:
        try:
            return self._client().predict(*args, api_name=api_name)
        except Exception:
            raise BackendUnavailable(
                "The private backend job failed. It was not automatically retried."
            ) from None

    def _verify_request_profile(self, payload: Mapping[str, Any]) -> None:
        if payload.get("execution_profile") != self.expected_profile:
            raise BackendProtocolMismatch(
                "Request execution profile does not match the selected private backend profile."
            )

    @staticmethod
    def _training_result(value: Any) -> BackendResult:
        if not isinstance(value, (tuple, list)) or len(value) != 2:
            raise BackendUnavailable("Private backend returned an invalid training response.")
        envelope, package = value
        if not isinstance(envelope, Mapping):
            raise BackendUnavailable("Private backend returned an invalid result envelope.")
        package_path = _remote_file_path(package)
        cleanup = (package_path,) if package_path else ()
        return BackendResult(dict(envelope), package_path, cleanup)

    def run_existing(
        self,
        request: Mapping[str, Any] | str,
        progress: Callable[[str], None] | None = None,
    ) -> BackendResult:
        payload = validate_request("existing", request)
        self._verify_request_profile(payload)
        self.ensure_ready(progress)
        result = self._submit_once("/v1/train/existing", payload)
        return self._training_result(result)

    def run_uploaded(
        self,
        request: Mapping[str, Any] | str,
        train_path: str,
        test_path: str | None = None,
        progress: Callable[[str], None] | None = None,
    ) -> BackendResult:
        payload = validate_request("uploaded", request)
        self._verify_request_profile(payload)
        self.ensure_ready(progress)
        result = self._submit_once(
            "/v1/train/uploaded",
            payload,
            self._file_wrapper(str(train_path)),
            self._file_wrapper(str(test_path)) if test_path else None,
        )
        return self._training_result(result)

    def predict_package(
        self,
        request: Mapping[str, Any] | str,
        package_path: str,
        *,
        tabular_path: str | None = None,
        image_path: str | None = None,
        progress: Callable[[str], None] | None = None,
    ) -> BackendResult:
        payload = validate_request("predict", request)
        self._verify_request_profile(payload)
        self.ensure_ready(progress)
        result = self._submit_once(
            "/v1/predict/package",
            payload,
            self._file_wrapper(str(package_path)),
            self._file_wrapper(str(tabular_path)) if tabular_path else None,
            self._file_wrapper(str(image_path)) if image_path else None,
        )
        if not isinstance(result, Mapping):
            raise BackendUnavailable("Private backend returned an invalid prediction response.")
        return BackendResult(dict(result))


__all__ = [
    "BackendUnavailable",
    "BackendTimeout",
    "BackendProtocolMismatch",
    "BackendResult",
    "GatewayBackendClient",
]