Create endpoint_utils.py

endpoint_utils.py

@@ -0,0 +1,69 @@
+# endpoint_utils.py
+from __future__ import annotations
+from typing import Optional, Tuple, Callable, Dict, Any
+from urllib.parse import urlparse
+import os, time, requests
+
+
+def _valid_uri(uri: Optional[str]) -> bool:
+    if not uri:
+        return False
+    p = urlparse(uri)
+    return p.scheme in {"http", "https"} and bool(p.netloc)
+
+
+def wake_endpoint(
+    uri: Optional[str],
+    *,
+    token: Optional[str] = None,
+    max_wait: int = 180,
+    poll_every: float = 5.0,
+    warm_payload: Optional[Dict[str, Any]] = None,
+    log: Callable[[str], None] = lambda _: None,
+) -> Tuple[bool, Optional[str]]:
+    """
+    Nudge a scale-to-zero Hugging Face Inference Endpoint and poll until it responds.
+
+    Returns:
+        (True, None) on success, or (False, "reason") on timeout / invalid input.
+
+    Notes:
+      - Expects endpoints that accept POST JSON bodies like {"inputs": "..."}.
+      - Treats 429/503/504 as "warming" signals while polling.
+    """
+    if not _valid_uri(uri):
+        return False, "invalid or missing URI (expect http(s)://...)"
+
+    headers: Dict[str, str] = {}
+    tok = token or os.environ.get("HF_TOKEN")
+    if tok:
+        headers["Authorization"] = f"Bearer {tok}"
+
+    # 1) Initial nudge (ignore any errors)
+    payload = warm_payload if warm_payload is not None else {"inputs": "wake"}
+    try:
+        requests.post(uri, headers=headers, json=payload, timeout=5)
+    except requests.RequestException:
+        pass
+
+    # 2) Poll until healthy or timeout
+    deadline = time.time() + max_wait
+    while time.time() < deadline:
+        try:
+            r = requests.post(uri, headers=headers, json={"inputs": "ping"}, timeout=5)
+            if r.ok:
+                log("✅ Endpoint is awake and responsive.")
+                return True, None
+
+            # Common warmup statuses: throttle/warming
+            if r.status_code in (429, 503, 504):
+                log(f"Endpoint warming (HTTP {r.status_code}); retrying in {poll_every:.0f}s…")
+            else:
+                log(f"Unexpected response (HTTP {r.status_code}); retrying in {poll_every:.0f}s…")
+
+        except requests.RequestException as e:
+            log(f"{type(e).__name__}; retrying in {poll_every:.0f}s…")
+
+        time.sleep(poll_every)
+
+    return False, f"timed out after {max_wait}s waiting for endpoint"