Update endpoint_utils.py

endpoint_utils.py

@@ -16,20 +16,14 @@ def wake_endpoint(
     uri: Optional[str],
     *,
     token: Optional[str] = None,
-    max_wait: int = 180,
+    max_wait: int = 600,          # was 180 — bump to 10 minutes
     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.
+    Wake a scale-to-zero HF Inference Endpoint by nudging it, then polling until ready.
+    Returns (True, None) if ready; otherwise (False, "<last status/message>").
     """
     if not _valid_uri(uri):
         return False, "invalid or missing URI (expect http(s)://...)"
@@ -39,7 +33,23 @@ def wake_endpoint(
     if tok:
         headers["Authorization"] = f"Bearer {tok}"
 
-    # 1) Initial nudge (ignore any errors)
+    # 0) Try a quick health check first (cheap)
+    last_detail = "no response"
+    try:
+        hr = requests.get(f"{uri.rstrip('/')}/health", headers=headers, timeout=5)
+        if hr.ok:
+            log("✅ /health reports ready.")
+            return True, None
+        try:
+            last_detail = (hr.json().get("error") or hr.json().get("message"))  # type: ignore
+        except Exception:
+            last_detail = (hr.text or "").strip()
+        log(f"[health] HTTP {hr.status_code} – {last_detail or 'warming?'}")
+    except requests.RequestException as e:
+        last_detail = type(e).__name__
+        log(f"[health] {last_detail}")
+
+    # 1) Initial nudge (ignore errors)
     payload = warm_payload if warm_payload is not None else {"inputs": "wake"}
     try:
         requests.post(uri, headers=headers, json=payload, timeout=5)
@@ -50,20 +60,29 @@ def wake_endpoint(
     deadline = time.time() + max_wait
     while time.time() < deadline:
         try:
-            r = requests.post(uri, headers=headers, json={"inputs": "ping"}, timeout=5)
+            r = requests.post(uri, headers=headers, json={"inputs": "ping"}, timeout=8)
             if r.ok:
                 log("✅ Endpoint is awake and responsive.")
                 return True, None
 
-            # Common warmup statuses: throttle/warming
+            # extract any helpful server message
+            detail = ""
+            try:
+                data = r.json()
+                detail = data.get("error") or data.get("message") or ""
+            except ValueError:
+                detail = (r.text or "").strip()
+
+            last_detail = f"HTTP {r.status_code}" + (f" – {detail}" if detail else "")
             if r.status_code in (429, 503, 504):
-                log(f"Endpoint warming (HTTP {r.status_code}); retrying in {poll_every:.0f}s…")
+                log(f"[server] {detail or 'warming up'} (HTTP {r.status_code}); retrying in {int(poll_every)}s…")
             else:
-                log(f"Unexpected response (HTTP {r.status_code}); retrying in {poll_every:.0f}s…")
+                log(f"[server] {detail or 'unexpected response'} (HTTP {r.status_code}); retrying in {int(poll_every)}s…")
 
         except requests.RequestException as e:
-            log(f"{type(e).__name__}; retrying in {poll_every:.0f}s…")
+            last_detail = type(e).__name__
+            log(f"[client] {last_detail}; retrying in {int(poll_every)}s…")
 
         time.sleep(poll_every)
 
-    return False, f"timed out after {max_wait}s waiting for endpoint"
+    return False, f"Timed out after {max_wait}s — last status: {last_detail}"