sync: from ael backend/space/

rerank.py

@@ -86,6 +86,14 @@ _router = None
 _router_lock = threading.Lock()
 _load_error = None
 _load_message = None
+_loading = False
+_load_failed_at = 0.0
+
+# A failed load used to latch for the lifetime of the container, which turned one
+# flaky checkpoint download into a permanently disabled reranker — the Space
+# would report "no opinion" for the rest of its life and nobody would know why.
+# Failures are retried on a cooldown instead.
+RETRY_AFTER_S = float(os.environ.get("AEL_RERANK_RETRY_S", "60"))
 
 
 def _get_router():
@@ -95,28 +103,62 @@ def _get_router():
     serves search, and the route needs to be able to ask "any opinion?"
     without handling an exception.
     """
-    global _router, _load_error, _load_message
-    if _router is not None or _load_error is not None:
+    global _router, _load_error, _load_message, _loading, _load_failed_at
+    if _router is not None:
         return _router
+    if _load_error is not None and (time.time() - _load_failed_at) < RETRY_AFTER_S:
+        return None
+    if _loading:
+        # Never block a search on a checkpoint download. The caller's deadline is
+        # ~1s; loading 421M parameters takes minutes the first time and is not
+        # something to discover inside a user's request. `warm()` exists so this
+        # is the rare path rather than the normal one.
+        return None
     with _router_lock:
-        if _router is not None or _load_error is not None:
+        if _router is not None or _loading:
             return _router
+        _loading = True
         try:
             from laya import Router
 
             # max_loaded=1: one checkpoint resident, which is all this uses.
             _router = Router(max_loaded=1)
             _router.preload([MODEL])
+            _load_error = None
+            _load_message = None
+            print(f"[rerank] {MODEL} ready", flush=True)
         except Exception as e:  # noqa: BLE001 - reported, never raised
             _load_error = f"{type(e).__name__}"
+            _load_failed_at = time.time()
             # The message matters and nothing else about this exception does:
             # "the reranker is off" is a decision, "the reranker broke because
             # <this>" is a bug report. The caller redacts before logging it.
             _load_message = f"{type(e).__name__}: {e}"[:300]
             print(f"[rerank] disabled: cannot load laya ({_load_error}): {e}", flush=True)
+        finally:
+            _loading = False
     return _router
 
 
+def warm():
+    """Start loading the checkpoint now, in the background.
+
+    Called once when the Space boots. Without this the first search request of
+    every container lifetime paid the full download-and-load cost inside an
+    800ms budget, so the reranker could never once have applied on a cold
+    container — and the only symptom was a timeout, which reads like a slow
+    network. Warming moves the cost to the ~2 minutes after a deploy, where
+    nobody is waiting.
+
+    Never raises and never blocks: a Space that cannot afford the model still
+    serves search.
+    """
+    if not ENABLED or _router is not None or _loading:
+        return
+    t = threading.Thread(target=_get_router, name="rerank-warm", daemon=True)
+    t.start()
+
+
 def _no_opinion():
     """Why this module has nothing to say, for the caller.
 
@@ -130,6 +172,8 @@ def _no_opinion():
     """
     if not ENABLED:
         return "disabled by config"
+    if _loading:
+        return "still loading"
     if _load_error:
         return _load_message or _load_error
     return None
@@ -141,6 +185,7 @@ def status():
     return {
         "enabled": ENABLED,
         "ready": router is not None,
+        "loading": _loading,
         "model": MODEL,
         "drop_below": DROP_BELOW,
         "min_results": MIN_RESULTS,
@@ -186,6 +231,20 @@ def _passage(result):
     return (title + "\n" + snippet).strip()
 
 
+def _stamp(out):
+    """Attach the *current* reason at the moment of return.
+
+    Stamping once, at the top, is wrong in both directions and both happened:
+    a reason captured before the load was attempted is `None` on exactly the
+    path that needs one, and a reason captured before a *successful* retry
+    survives into an `ok: true` response. A field that is always populated is a
+    field everyone learns to ignore, which is the failure this whole mechanism
+    exists to prevent.
+    """
+    out["no_opinion"] = _no_opinion()
+    return out
+
+
 def rerank(query, results, budget_ms=800):
     """Score, reorder and optionally drop `results` for `query`.
 
@@ -201,19 +260,13 @@ def rerank(query, results, budget_ms=800):
         "dropped": 0,
         "starved": False,
         "scores": [],
-        "no_opinion": _no_opinion(),
+        "no_opinion": None,
     }
     if not ENABLED or not results:
-        return out
+        return _stamp(out)
     router = _get_router()
     if router is None:
-        # Recomputed here, not reused from above: `_no_opinion()` reads
-        # `_load_error`, and until this line nothing had *tried* to load the
-        # model, so the value captured in `out` was always None on exactly the
-        # path that needed it. The first version of this reported "no opinion"
-        # for a reranker that had never once been installed.
-        out["no_opinion"] = _no_opinion()
-        return out
+        return _stamp(out)
 
     t0 = time.time()
     scored = []
@@ -228,15 +281,15 @@ def rerank(query, results, budget_ms=800):
             p = float(res["answers"]["relevant"]["noul"])
         except Exception as e:  # noqa: BLE001
             print(f"[rerank] score failed for result {i}: {type(e).__name__}: {e}", flush=True)
-            return dict(out, results=results)  # fail open, whole set untouched
+            return _stamp(dict(out, results=results))  # fail open, whole set untouched
         if (time.time() - t0) * 1000 > budget_ms:
             # Over budget: return what we have not yet reordered, in order.
             print(f"[rerank] over budget ({budget_ms}ms) after {i} results", flush=True)
-            return dict(out, results=results)
+            return _stamp(dict(out, results=results))
         scored.append((p, i, r))
 
     if not any(p is not None for p, _, _ in scored):
-        return out
+        return _stamp(out)
 
     # Highest score first. Unscored results keep their relative position at the
     # end rather than being dropped: we have no opinion about them.
@@ -255,7 +308,7 @@ def rerank(query, results, budget_ms=800):
         # The model rejected everything, which is a model that did not
         # understand the query rather than a set with nothing in it. Say so
         # instead of pretending.
-        return dict(out, results=results, starved=True)
+        return _stamp(dict(out, results=results, starved=True))
 
     out.update(ok=True, results=kept, dropped=dropped)
-    return out
+    return _stamp(out)

searx_bridge.py

@@ -24,6 +24,13 @@ from fastapi.responses import JSONResponse
 import rerank
 import searx_runtime
 
+# Load the reranker's checkpoint as soon as the container boots, not on the
+# first search request. The first request has a ~1s deadline and the load takes
+# minutes, so without this the reranker can never apply on a cold container and
+# the only symptom is a timeout. A daemon thread, so a failure or a slow mirror
+# cannot hold the Space's startup hostage.
+rerank.warm()
+
 
 def register_routes(fa_app):
     @fa_app.get("/respite/v2/searx/status")
@@ -33,8 +40,7 @@ def register_routes(fa_app):
             "installed": searx_runtime.searxng_ready(),
             "boot_error": searx_runtime.boot_error(),
             "port": searx_runtime.SEARXNG_PORT,
-            "reranker": rerank.status(),
-        })
+            "reranker": rerank.status(),        })
 
     @fa_app.get("/respite/v2/searx/search")
     async def _search(request: Request):