sync: from ael backend/space/

rerank.py

@@ -78,9 +78,23 @@ DROP_BELOW = float(os.environ.get("AEL_RERANK_DROP", "0.5"))
 # one, not a target: a short page is fine, an empty one is not.
 MIN_RESULTS = 1
 
+# How many results the model is allowed to judge. A 421M cross-encoder on a
+# shared CPU does not score eight documents inside an 800ms budget, and the
+# first version's answer to that was to throw the whole set away — which is why
+# a reranker that was installed, loaded, healthy and simply too slow reported
+# itself identically to one that had never run. Judging the first N and leaving
+# the rest in engine order is strictly better than judging none: the survivors
+# are reordered by relevance and the rejects are dropped, and the untested tail
+# keeps the order the race gave it.
+MAX_SCORED = int(os.environ.get("AEL_RERANK_MAX", "6"))
+
 # Recent (query, scores) for threshold tuning. Bounded so a long-lived Space
 # cannot grow without limit.
 RECENT = deque(maxlen=int(os.environ.get("AEL_RERANK_LOG", "200")))
+# Wall-clock of recent reranks, same bound. This is the tuning data for
+# MAX_SCORED and the budget: the scores say what to keep, this says what it
+# costs to judge it.
+RECENT_MS = deque(maxlen=int(os.environ.get("AEL_RERANK_LOG", "200")))
 
 _router = None
 _router_lock = threading.Lock()
@@ -197,7 +211,10 @@ def status():
 def recent_scores(limit=20):
     """Recent (query, scores) pairs, newest first — the threshold-tuning data."""
     items = list(RECENT)[-limit:][::-1]
-    return [{"q": q, "scores": s} for q, s in items]
+    ms = list(RECENT_MS)[-limit:][::-1]
+    return [
+        {"q": q, "scores": s, "ms": m} for (q, s), m in zip(items, ms)
+    ]
 
 
 def _question(query):
@@ -270,7 +287,24 @@ def rerank(query, results, budget_ms=800):
 
     t0 = time.time()
     scored = []
+    judged = 0
     for i, r in enumerate(results):
+        if judged >= MAX_SCORED:
+            # Enough judged to reorder by. The tail is kept, unscored, at the
+            # end — we have no opinion about it and will not pretend to.
+            scored.append((None, i, r))
+            continue
+        if judged and (time.time() - t0) * 1000 > budget_ms:
+            # Over budget, but not empty-handed. A partial judgement is worth
+            # more than none: the documents we did score get reordered and the
+            # ones the model rejected get dropped. Only a run that scored
+            # *nothing* throws its work away.
+            print(
+                f"[rerank] over budget ({budget_ms}ms) after {judged} of {len(results)}"
+                f" — reranking what was judged",
+                flush=True,
+            )
+            break
         passage = _passage(r)
         if not passage:
             # No text to judge. Keep it, unscored, and let the ordering stand.
@@ -282,11 +316,18 @@ def rerank(query, results, budget_ms=800):
         except Exception as e:  # noqa: BLE001
             print(f"[rerank] score failed for result {i}: {type(e).__name__}: {e}", flush=True)
             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 _stamp(dict(out, results=results))
         scored.append((p, i, r))
+        judged += 1
+
+    # Anything the loop never reached keeps its place at the end, unscored.
+    if len(scored) < len(results):
+        seen_idx = {i for _, i, _ in scored}
+        for i, r in enumerate(results):
+            if i not in seen_idx:
+                scored.append((None, i, r))
+
+    out["ms"] = int((time.time() - t0) * 1000)
+    out["judged"] = judged
 
     if not any(p is not None for p, _, _ in scored):
         return _stamp(out)
@@ -303,6 +344,10 @@ def rerank(query, results, budget_ms=800):
         if p is not None
     ]
     RECENT.append((query[:120], [p for p, _, _ in scored if p is not None]))
+    # The duration is the number this stage actually needs tuning against, and
+    # it is the one that was missing: a reranker that is too slow for its own
+    # budget and a reranker that is not installed look identical from outside.
+    RECENT_MS.append(out["ms"])
 
     if dropped and len(kept) < MIN_RESULTS:
         # The model rejected everything, which is a model that did not