Spaces:
Running
Running
| """Model discovery and inference smoke-test logic for the StockAnalysis Space. | |
| The module deliberately keeps inference cheap: it resolves the most-trending | |
| conversational models that currently have at least one live provider, then sends | |
| one tiny chat request per model. All tunables are applied through | |
| ``configure()``, which ``app.py`` calls with its SETTINGS block. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import re | |
| import threading | |
| import time | |
| import uuid | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| from typing import Any, Callable | |
| import requests | |
| from huggingface_hub import InferenceClient, InferenceTimeoutError | |
| from huggingface_hub.errors import HfHubHTTPError | |
| HUB_MODELS_API = "https://huggingface.co/api/models" | |
| WHOAMI_API = "https://huggingface.co/api/whoami-v2" | |
| RESULTS_ROOT = Path("/data/hf-inference-smoke") | |
| RUNS_DIR = RESULTS_ROOT / "runs" | |
| MODELS_DIR = RESULTS_ROOT / "models" | |
| DEFAULT_MODEL_LIMIT = 30 | |
| DEFAULT_PAGE_SIZE = 100 | |
| DEFAULT_MAX_SCAN = 500 | |
| RUN_ID_RE = re.compile(r"^[0-9TZ]+-[0-9a-f]{6}$") | |
| PROMPT = "ping" | |
| MAX_TOKENS = 8 | |
| DEFAULT_TIMEOUT = 60.0 | |
| DEFAULT_CONCURRENCY = 3 | |
| ATTEMPT_POLICIES = ("cheapest", "fastest") | |
| def configure( | |
| *, | |
| results_root: str | Path | None = None, | |
| model_limit: int | None = None, | |
| concurrency: int | None = None, | |
| max_tokens: int | None = None, | |
| request_timeout: float | None = None, | |
| page_size: int | None = None, | |
| max_scan: int | None = None, | |
| ) -> dict[str, Any]: | |
| """Apply the settings declared in ``app.py``. | |
| Every value is optional so tests can override just the storage location. | |
| The app calls this once at import time, which keeps all user-facing knobs | |
| in a single place instead of scattered environment variables. | |
| """ | |
| global RESULTS_ROOT, RUNS_DIR, MODELS_DIR | |
| global DEFAULT_MODEL_LIMIT, DEFAULT_CONCURRENCY, MAX_TOKENS | |
| global DEFAULT_TIMEOUT, DEFAULT_PAGE_SIZE, DEFAULT_MAX_SCAN | |
| if results_root is not None: | |
| RESULTS_ROOT = Path(results_root) | |
| RUNS_DIR = RESULTS_ROOT / "runs" | |
| MODELS_DIR = RESULTS_ROOT / "models" | |
| if model_limit is not None: | |
| DEFAULT_MODEL_LIMIT = int(model_limit) | |
| if concurrency is not None: | |
| DEFAULT_CONCURRENCY = int(concurrency) | |
| if max_tokens is not None: | |
| MAX_TOKENS = int(max_tokens) | |
| if request_timeout is not None: | |
| DEFAULT_TIMEOUT = float(request_timeout) | |
| if page_size is not None: | |
| DEFAULT_PAGE_SIZE = int(page_size) | |
| if max_scan is not None: | |
| DEFAULT_MAX_SCAN = int(max_scan) | |
| return { | |
| "results_root": str(RESULTS_ROOT), | |
| "model_limit": DEFAULT_MODEL_LIMIT, | |
| "concurrency": DEFAULT_CONCURRENCY, | |
| "max_tokens": MAX_TOKENS, | |
| "request_timeout": DEFAULT_TIMEOUT, | |
| "page_size": DEFAULT_PAGE_SIZE, | |
| "max_scan": DEFAULT_MAX_SCAN, | |
| } | |
| def utc_now() -> str: | |
| """Return an ISO-8601 UTC timestamp with a trailing Z.""" | |
| return datetime.now(timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z") | |
| def new_run_id() -> str: | |
| """Return a sortable, path-safe run identifier.""" | |
| stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") | |
| return f"{stamp}-{uuid.uuid4().hex[:6]}" | |
| def ensure_storage() -> None: | |
| """Create the persistent-storage directories on demand.""" | |
| RUNS_DIR.mkdir(parents=True, exist_ok=True) | |
| MODELS_DIR.mkdir(parents=True, exist_ok=True) | |
| def _write_json(path: Path, payload: Any) -> None: | |
| """Write JSON atomically so a crashed run cannot leave a partial file.""" | |
| ensure_storage() | |
| tmp_path = path.with_suffix(path.suffix + ".tmp") | |
| tmp_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") | |
| tmp_path.replace(path) | |
| def _auth_headers() -> dict[str, str]: | |
| token = os.getenv("HF_TOKEN") | |
| return {"Authorization": f"Bearer {token}"} if token else {} | |
| def _parse_provider_mapping(raw_mapping: Any) -> tuple[list[str], dict[str, str]]: | |
| """Return live providers and every provider status from a mapping payload.""" | |
| live: list[str] = [] | |
| statuses: dict[str, str] = {} | |
| if not isinstance(raw_mapping, list): | |
| return live, statuses | |
| for item in raw_mapping: | |
| if not isinstance(item, dict): | |
| continue | |
| provider = item.get("provider") | |
| status = item.get("status") | |
| if not provider or not status: | |
| continue | |
| statuses[provider] = status | |
| if status == "live" and provider not in live: | |
| live.append(provider) | |
| return live, statuses | |
| def fetch_trending_models( | |
| limit: int | None = None, | |
| *, | |
| page_size: int | None = None, | |
| max_scan: int | None = None, | |
| save: bool = True, | |
| snapshot_name: str | None = None, | |
| session: requests.Session | None = None, | |
| ) -> dict[str, Any]: | |
| """Resolve the top ``limit`` trending conversational models with live providers. | |
| The Hub API can return models whose only provider is staging or whose | |
| provider mapping is empty, so we scan more than ``limit`` rows and keep | |
| only candidates that have at least one provider marked ``live``. | |
| """ | |
| if limit is None: | |
| limit = DEFAULT_MODEL_LIMIT | |
| if page_size is None: | |
| page_size = DEFAULT_PAGE_SIZE | |
| if max_scan is None: | |
| max_scan = DEFAULT_MAX_SCAN | |
| if limit < 1: | |
| raise ValueError("limit must be >= 1") | |
| client = session or requests.Session() | |
| headers = _auth_headers() | |
| models: list[dict[str, Any]] = [] | |
| skipped: list[dict[str, Any]] = [] | |
| scanned = 0 | |
| pages = 0 | |
| offset = 0 | |
| while len(models) < limit and scanned < max_scan: | |
| params = { | |
| "inference_provider": "all", | |
| "filter": "conversational", | |
| "sort": "trendingScore", | |
| "direction": "-1", | |
| "limit": page_size, | |
| "offset": offset, | |
| "expand[]": "inferenceProviderMapping", | |
| } | |
| response = client.get(HUB_MODELS_API, params=params, headers=headers, timeout=30) | |
| response.raise_for_status() | |
| page = response.json() | |
| if not isinstance(page, list) or not page: | |
| break | |
| pages += 1 | |
| for raw in page: | |
| if not isinstance(raw, dict): | |
| continue | |
| scanned += 1 | |
| model_id = raw.get("id") | |
| if not model_id: | |
| continue | |
| live, statuses = _parse_provider_mapping(raw.get("inferenceProviderMapping")) | |
| if not live: | |
| skipped.append({"model": model_id, "provider_statuses": statuses}) | |
| continue | |
| models.append( | |
| { | |
| "rank": len(models) + 1, | |
| "model": model_id, | |
| "trending_score": raw.get("trendingScore"), | |
| "downloads": raw.get("downloads"), | |
| "pipeline_tag": raw.get("pipeline_tag"), | |
| "live_providers": live, | |
| "provider_statuses": statuses, | |
| } | |
| ) | |
| if len(models) >= limit: | |
| break | |
| if len(page) < page_size: | |
| break | |
| offset += page_size | |
| if len(models) < limit: | |
| raise RuntimeError( | |
| f"Only found {len(models)} models with a live provider after scanning " | |
| f"{scanned} rows; requested {limit}. Increase SMOKE_MAX_SCAN or lower " | |
| "SMOKE_MODEL_LIMIT." | |
| ) | |
| snapshot = { | |
| "fetched_at": utc_now(), | |
| "source": HUB_MODELS_API, | |
| "query": { | |
| "inference_provider": "all", | |
| "filter": "conversational", | |
| "sort": "trendingScore", | |
| "direction": "-1", | |
| "page_size": page_size, | |
| "max_scan": max_scan, | |
| }, | |
| "requested_limit": limit, | |
| "returned": len(models), | |
| "scanned": scanned, | |
| "pages": pages, | |
| "skipped_without_live_provider": skipped, | |
| "models": models, | |
| } | |
| if save: | |
| name = snapshot_name or f"snapshot-{uuid.uuid4().hex[:8]}" | |
| _write_json(MODELS_DIR / f"{name}.json", snapshot) | |
| _write_json(MODELS_DIR / "latest.json", snapshot) | |
| return snapshot | |
| def get_token_status(token: str | None = None, *, session: requests.Session | None = None) -> dict[str, Any]: | |
| """Validate HF_TOKEN and report the authenticated account and plan.""" | |
| resolved = token or os.getenv("HF_TOKEN") | |
| if not resolved: | |
| return {"ok": False, "name": None, "plan": None, "error": "HF_TOKEN is not set."} | |
| client = session or requests.Session() | |
| try: | |
| response = client.get( | |
| WHOAMI_API, | |
| headers={"Authorization": f"Bearer {resolved}"}, | |
| timeout=15, | |
| ) | |
| except requests.RequestException as exc: | |
| return {"ok": False, "name": None, "plan": None, "error": f"Token check failed: {exc}"} | |
| if response.status_code != 200: | |
| detail = response.text.strip().replace("\n", " ")[:200] | |
| return { | |
| "ok": False, | |
| "name": None, | |
| "plan": None, | |
| "error": f"Token check returned HTTP {response.status_code}: {detail}", | |
| } | |
| data = response.json() | |
| return { | |
| "ok": True, | |
| "name": data.get("name"), | |
| "plan": data.get("plan"), | |
| "error": None, | |
| } | |
| def load_latest_model_snapshot() -> dict[str, Any] | None: | |
| """Load the most recently saved discovery snapshot, if one exists.""" | |
| path = MODELS_DIR / "latest.json" | |
| if not path.exists(): | |
| return None | |
| try: | |
| return json.loads(path.read_text(encoding="utf-8")) | |
| except (OSError, json.JSONDecodeError): | |
| return None | |
| _thread_state = threading.local() | |
| def _client_for(token: str, timeout: float) -> InferenceClient: | |
| """Return one InferenceClient per worker thread. | |
| ``InferenceClient`` wraps a requests session, so sharing a single instance | |
| across worker threads is avoidable risk. | |
| """ | |
| key = (token, timeout) | |
| if getattr(_thread_state, "key", None) != key: | |
| _thread_state.key = key | |
| _thread_state.client = InferenceClient(token=token, timeout=timeout) | |
| return _thread_state.client | |
| def _format_exception(exc: BaseException, token: str) -> str: | |
| if isinstance(exc, HfHubHTTPError): | |
| response = getattr(exc, "response", None) | |
| status = getattr(response, "status_code", None) | |
| detail = getattr(response, "text", "") or str(exc) | |
| message = f"HTTP {status}: {detail}" if status else str(exc) | |
| elif isinstance(exc, InferenceTimeoutError): | |
| message = "Request timed out." | |
| else: | |
| message = f"{type(exc).__name__}: {exc}" | |
| message = " ".join(message.split()) | |
| if token: | |
| message = message.replace(token, "hf_***") | |
| return message[:500] | |
| def _http_status(exc: BaseException) -> int | None: | |
| response = getattr(exc, "response", None) | |
| return getattr(response, "status_code", None) | |
| def _run_attempt( | |
| model_id: str, | |
| policy: str, | |
| token: str, | |
| timeout: float, | |
| ) -> dict[str, Any]: | |
| """Run a single chat-completion attempt and normalize its outcome.""" | |
| started = time.perf_counter() | |
| try: | |
| response = _client_for(token, timeout).chat_completion( | |
| messages=[{"role": "user", "content": PROMPT}], | |
| model=f"{model_id}:{policy}", | |
| max_tokens=MAX_TOKENS, | |
| temperature=0.0, | |
| stream=False, | |
| ) | |
| except Exception as exc: # noqa: BLE001 - the point is to classify any failure | |
| return { | |
| "policy": policy, | |
| "ok": False, | |
| "latency_ms": round((time.perf_counter() - started) * 1000), | |
| "http_status": _http_status(exc), | |
| "error": _format_exception(exc, token), | |
| "usage": None, | |
| } | |
| latency_ms = round((time.perf_counter() - started) * 1000) | |
| choices = getattr(response, "choices", None) or [] | |
| usage = getattr(response, "usage", None) | |
| usage_payload = ( | |
| { | |
| "prompt_tokens": getattr(usage, "prompt_tokens", None), | |
| "completion_tokens": getattr(usage, "completion_tokens", None), | |
| "total_tokens": getattr(usage, "total_tokens", None), | |
| } | |
| if usage is not None | |
| else None | |
| ) | |
| if not choices: | |
| return { | |
| "policy": policy, | |
| "ok": False, | |
| "latency_ms": latency_ms, | |
| "http_status": 200, | |
| "error": "Response contained no choices.", | |
| "usage": usage_payload, | |
| } | |
| return { | |
| "policy": policy, | |
| "ok": True, | |
| "latency_ms": latency_ms, | |
| "http_status": 200, | |
| "error": None, | |
| "usage": usage_payload, | |
| } | |
| def test_model( | |
| candidate: dict[str, Any], | |
| *, | |
| token: str, | |
| timeout: float | None = None, | |
| ) -> dict[str, Any]: | |
| """Test one model with cheapest routing, then fastest on failure.""" | |
| if timeout is None: | |
| timeout = DEFAULT_TIMEOUT | |
| attempts: list[dict[str, Any]] = [] | |
| for policy in ATTEMPT_POLICIES: | |
| try: | |
| attempt = _run_attempt(candidate["model"], policy, token, timeout) | |
| except Exception as exc: # noqa: BLE001 - defensive boundary for worker threads | |
| attempt = { | |
| "policy": policy, | |
| "ok": False, | |
| "latency_ms": 0, | |
| "http_status": None, | |
| "error": _format_exception(exc, token), | |
| "usage": None, | |
| } | |
| attempts.append(attempt) | |
| if attempt["ok"]: | |
| return { | |
| "rank": candidate["rank"], | |
| "model": candidate["model"], | |
| "ok": True, | |
| "attempt": policy, | |
| "fallback": policy == "fastest", | |
| "latency_ms": sum(item["latency_ms"] for item in attempts), | |
| "http_status": attempt["http_status"], | |
| "error": None, | |
| "attempts": attempts, | |
| } | |
| last = attempts[-1] | |
| return { | |
| "rank": candidate["rank"], | |
| "model": candidate["model"], | |
| "ok": False, | |
| "attempt": last["policy"], | |
| "fallback": True, | |
| "latency_ms": sum(item["latency_ms"] for item in attempts), | |
| "http_status": last["http_status"], | |
| "error": last["error"], | |
| "attempts": attempts, | |
| } | |
| def run_smoke( | |
| models: list[dict[str, Any]], | |
| *, | |
| token: str | None = None, | |
| concurrency: int | None = None, | |
| timeout: float | None = None, | |
| run_id: str | None = None, | |
| progress: Callable[[int, int, dict[str, Any]], None] | None = None, | |
| ) -> dict[str, Any]: | |
| """Run the smoke test and persist a JSON result set.""" | |
| if concurrency is None: | |
| concurrency = DEFAULT_CONCURRENCY | |
| if timeout is None: | |
| timeout = DEFAULT_TIMEOUT | |
| resolved_token = token or os.getenv("HF_TOKEN") | |
| if not resolved_token: | |
| raise RuntimeError("HF_TOKEN is not set; refusing to make inference calls.") | |
| if not models: | |
| raise ValueError("No models supplied for the smoke test.") | |
| resolved_run_id = run_id or new_run_id() | |
| if not RUN_ID_RE.match(resolved_run_id): | |
| raise ValueError("Invalid run_id.") | |
| started_at = utc_now() | |
| started_clock = time.perf_counter() | |
| results: list[dict[str, Any]] = [] | |
| total = len(models) | |
| with ThreadPoolExecutor(max_workers=max(1, concurrency)) as pool: | |
| futures = { | |
| pool.submit(test_model, candidate, token=resolved_token, timeout=timeout): candidate | |
| for candidate in models | |
| } | |
| for done, future in enumerate(as_completed(futures), start=1): | |
| result = future.result() | |
| results.append(result) | |
| if progress is not None: | |
| progress(done, total, result) | |
| results.sort(key=lambda item: item["rank"]) | |
| finished_at = utc_now() | |
| payload = { | |
| "run_id": resolved_run_id, | |
| "started_at": started_at, | |
| "finished_at": finished_at, | |
| "duration_seconds": round(time.perf_counter() - started_clock, 2), | |
| "model_limit": total, | |
| "passed": sum(1 for item in results if item["ok"]), | |
| "failed": sum(1 for item in results if not item["ok"]), | |
| "fallbacks": sum(1 for item in results if item["ok"] and item["fallback"]), | |
| "concurrency": concurrency, | |
| "max_tokens": MAX_TOKENS, | |
| "request_timeout_seconds": timeout, | |
| "results": results, | |
| } | |
| _write_json(RUNS_DIR / f"{resolved_run_id}.json", payload) | |
| _write_json(RESULTS_ROOT / "latest.json", payload) | |
| return payload | |
| def list_runs(limit: int = 25) -> list[dict[str, Any]]: | |
| """Return recent run metadata, newest first.""" | |
| ensure_storage() | |
| paths = sorted(RUNS_DIR.glob("*.json"), key=lambda path: path.stat().st_mtime, reverse=True) | |
| runs: list[dict[str, Any]] = [] | |
| for path in paths[:limit]: | |
| try: | |
| payload = json.loads(path.read_text(encoding="utf-8")) | |
| except (OSError, json.JSONDecodeError): | |
| continue | |
| runs.append( | |
| { | |
| "run_id": payload.get("run_id", path.stem), | |
| "started_at": payload.get("started_at"), | |
| "passed": payload.get("passed"), | |
| "failed": payload.get("failed"), | |
| "model_limit": payload.get("model_limit"), | |
| "duration_seconds": payload.get("duration_seconds"), | |
| } | |
| ) | |
| return runs | |
| def load_run(run_id: str) -> dict[str, Any] | None: | |
| """Load one persisted run, rejecting path-traversal attempts.""" | |
| if not RUN_ID_RE.match(run_id): | |
| raise ValueError("Invalid run_id.") | |
| path = RUNS_DIR / f"{run_id}.json" | |
| if not path.exists(): | |
| return None | |
| try: | |
| return json.loads(path.read_text(encoding="utf-8")) | |
| except (OSError, json.JSONDecodeError): | |
| return None | |
| def load_latest_run() -> dict[str, Any] | None: | |
| """Load the most recent completed run, if one exists.""" | |
| path = RESULTS_ROOT / "latest.json" | |
| if not path.exists(): | |
| return None | |
| try: | |
| return json.loads(path.read_text(encoding="utf-8")) | |
| except (OSError, json.JSONDecodeError): | |
| return None | |