Spaces:
Sleeping
Sleeping
| """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 | |
| from ui_contracts import DecisionTrace, ResearchPlan | |
| 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.""" | |
| class BackendResearchError(BackendUnavailable): | |
| """The private research API returned a stable public failure.""" | |
| def __init__(self, code: str, message: str) -> None: | |
| self.code = str(code or "RESEARCH_UNAVAILABLE") | |
| super().__init__(str(message or "Research is unavailable.")) | |
| class BackendResult: | |
| envelope: dict[str, Any] | |
| package_path: str | None = None | |
| cleanup_files: tuple[str, ...] = () | |
| class ScheduledBackendResult: | |
| metadata: dict[str, Any] | |
| result: BackendResult | None = None | |
| 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})" | |
| ) | |
| 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 | |
| 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)) | |
| 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", | |
| "hf_jobs_overflow", | |
| } | |
| 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 _research_envelope(value: Any) -> dict[str, Any]: | |
| if not isinstance(value, Mapping): | |
| raise BackendUnavailable( | |
| "Private backend returned an invalid research response." | |
| ) | |
| result = dict(value) | |
| if result.get("ok") is not True: | |
| error = result.get("error") | |
| if not isinstance(error, Mapping): | |
| raise BackendUnavailable( | |
| "Private backend returned an invalid research failure." | |
| ) | |
| raise BackendResearchError( | |
| str(error.get("code") or "RESEARCH_UNAVAILABLE"), | |
| str(error.get("message") or "Research is unavailable."), | |
| ) | |
| if result.get("protocol_version") != PROTOCOL_VERSION: | |
| raise BackendProtocolMismatch( | |
| "Gateway and backend versions differ. Deploy matching revisions before researching." | |
| ) | |
| return result | |
| def _require_research_capability( | |
| health: Mapping[str, Any], capability: str | |
| ) -> None: | |
| if capability not in set(health.get("capabilities") or ()): | |
| raise BackendProtocolMismatch( | |
| "Private backend does not expose the required research capability." | |
| ) | |
| 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." | |
| ) | |
| 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 _scheduled_result(cls, value: Any) -> ScheduledBackendResult: | |
| if not isinstance(value, (tuple, list)) or len(value) != 3: | |
| raise BackendUnavailable( | |
| "Private backend returned an invalid scheduler response." | |
| ) | |
| metadata, envelope, package = value | |
| if not isinstance(metadata, Mapping): | |
| raise BackendUnavailable( | |
| "Private backend returned invalid scheduler metadata." | |
| ) | |
| if envelope is None: | |
| result = None | |
| else: | |
| result = cls._training_result((envelope, package)) | |
| return ScheduledBackendResult(dict(metadata), result) | |
| 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 research_dataset_profile(self, selection: Mapping[str, Any], train_path: str | None = None) -> dict[str, Any]: | |
| if self.expected_profile != "cpu": | |
| raise BackendProtocolMismatch("Dataset profiling must use the private CPU backend.") | |
| health = self.ensure_ready() | |
| self._require_research_capability(health, "research_profiling") | |
| envelope = self._research_envelope(self._submit_once( | |
| "/v1/research/profile", dict(selection), self._file_wrapper(train_path) if train_path else None, | |
| )) | |
| profile = envelope.get("dataset_profile") | |
| allowed = {"dataset_id", "dataset_name", "modality", "data_source", "row_count", "sample_count", | |
| "feature_count", "class_count", "shape", "dtype_counts", "feature_type_counts", "missing_fraction", "target"} | |
| if not isinstance(profile, Mapping) or set(profile) - allowed or profile.get("modality") not in {"image", "tabular"}: | |
| raise BackendUnavailable("Private backend returned an invalid dataset profile.") | |
| # The private profiling endpoint sanitizes aggregates; the planning | |
| # endpoint validates them again before any external request. Keep the | |
| # public gateway independent from the private research implementation. | |
| return dict(profile) | |
| def research_plan( | |
| self, | |
| request: Mapping[str, Any], | |
| progress: Callable[[str], None] | None = None, | |
| ) -> dict[str, Any]: | |
| if self.expected_profile != "cpu": | |
| raise BackendProtocolMismatch( | |
| "Research requests must use the private CPU backend." | |
| ) | |
| health = self.ensure_ready(progress) | |
| self._require_research_capability(health, "research_planning") | |
| envelope = self._research_envelope( | |
| self._submit_once("/v1/research/plan", dict(request)) | |
| ) | |
| raw_plan = envelope.get("research_plan") | |
| if not isinstance(raw_plan, Mapping): | |
| raise BackendUnavailable( | |
| "Private backend returned an invalid research plan." | |
| ) | |
| try: | |
| plan = ResearchPlan.from_mapping(raw_plan) | |
| except ValueError as exc: | |
| raise BackendUnavailable( | |
| "Private backend returned an invalid research plan." | |
| ) from exc | |
| if not 1 <= len(plan.candidates) <= 3: | |
| raise BackendUnavailable( | |
| "Private backend returned an invalid candidate count." | |
| ) | |
| return plan.to_dict() | |
| def approve_research_candidate( | |
| self, | |
| plan: Mapping[str, Any], | |
| candidate_id: str, | |
| progress: Callable[[str], None] | None = None, | |
| ) -> dict[str, Any]: | |
| if self.expected_profile != "cpu": | |
| raise BackendProtocolMismatch( | |
| "Research approval must use the private CPU backend." | |
| ) | |
| canonical_plan = ResearchPlan.from_mapping(plan).to_dict() | |
| health = self.ensure_ready(progress) | |
| self._require_research_capability(health, "research_approval") | |
| envelope = self._research_envelope( | |
| self._submit_once( | |
| "/v1/research/approve", canonical_plan, str(candidate_id) | |
| ) | |
| ) | |
| config = envelope.get("pipeline_config") | |
| evidence_ids = envelope.get("evidence_ids") | |
| trace = envelope.get("decision_trace") | |
| if ( | |
| not isinstance(config, Mapping) | |
| or not isinstance(evidence_ids, list) | |
| or not all(isinstance(item, str) for item in evidence_ids) | |
| or not isinstance(trace, Mapping) | |
| ): | |
| raise BackendUnavailable( | |
| "Private backend returned an invalid candidate approval." | |
| ) | |
| try: | |
| canonical_trace = DecisionTrace.from_mapping(trace) | |
| except ValueError as exc: | |
| raise BackendUnavailable( | |
| "Private backend returned an invalid candidate approval." | |
| ) from exc | |
| approved_id = str(envelope.get("approved_candidate_id") or "") | |
| if approved_id != str(candidate_id): | |
| raise BackendUnavailable( | |
| "Private backend approved a different candidate." | |
| ) | |
| return { | |
| "schema_version": 1, | |
| "plan_id": str(envelope.get("plan_id") or ""), | |
| "approved_candidate_id": approved_id, | |
| "pipeline_config": dict(config), | |
| "evidence_ids": list(evidence_ids), | |
| "decision_trace": canonical_trace.to_dict(), | |
| } | |
| def schedule_existing( | |
| self, | |
| request: Mapping[str, Any] | str, | |
| *, | |
| request_id: str, | |
| idempotency_key: str, | |
| account_key: str, | |
| progress: Callable[[str], None] | None = None, | |
| ) -> ScheduledBackendResult: | |
| payload = validate_request("existing", request) | |
| self._verify_request_profile(payload) | |
| self.ensure_ready(progress) | |
| result = self._submit_once( | |
| "/v1/schedule/existing", | |
| payload, | |
| str(request_id), | |
| str(idempotency_key), | |
| str(account_key), | |
| ) | |
| return self._scheduled_result(result) | |
| def inspect_overflow( | |
| self, | |
| *, | |
| request_id: str, | |
| external_job_id: str, | |
| result_location: str, | |
| output_prefix: str = "", | |
| ) -> ScheduledBackendResult: | |
| self.ensure_ready() | |
| result = self._submit_once( | |
| "/v1/overflow/status", | |
| str(request_id), | |
| str(external_job_id), | |
| str(result_location), | |
| str(output_prefix), | |
| ) | |
| return self._scheduled_result(result) | |
| def cancel_overflow(self, *, external_job_id: str) -> dict[str, Any]: | |
| self.ensure_ready() | |
| result = self._submit_once( | |
| "/v1/overflow/cancel", str(external_job_id) | |
| ) | |
| if not isinstance(result, Mapping): | |
| raise BackendUnavailable( | |
| "Private backend returned invalid cancellation metadata." | |
| ) | |
| return dict(result) | |
| def cleanup_overflow( | |
| self, *, request_id: str, result_location: str | |
| ) -> dict[str, Any]: | |
| self.ensure_ready() | |
| result = self._submit_once( | |
| "/v1/overflow/cleanup", | |
| str(request_id), | |
| str(result_location), | |
| ) | |
| if not isinstance(result, Mapping): | |
| raise BackendUnavailable( | |
| "Private backend returned invalid cleanup metadata." | |
| ) | |
| return dict(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", | |
| "BackendResearchError", | |
| "BackendResult", | |
| "ScheduledBackendResult", | |
| "GatewayBackendClient", | |
| ] | |