Studio: print reranker stage milestones to stderr for diagnostic visibility
When the reranker hung on rerank=True there were zero log lines after 'retrieved=N (no threshold)', which made it impossible to tell whether the hang was in _load (CrossEncoder construction), in get_reranker's lock acquisition, or in predict. Structlog routing may also be the culprit since we never saw the 'Loading RAG reranker' info line. Add unconditional stderr prints at each milestone — entered, device resolved, before CrossEncoder, after CrossEncoder, rerank entered, predict starting, predict done. These bypass any logger config and show up directly in /tmp/studio.log next to the rest of the captured stdout/stderr. Leaving structlog logger.info calls in place too so the structured stream still gets the same data when routing works.
This commit is contained in:
parent
8fb2fb9e2a
commit
1db654abb1
1 changed files with 49 additions and 2 deletions
|
|
@ -6,6 +6,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from typing import Any
|
||||
|
|
@ -37,21 +38,45 @@ def _resolve_device() -> str:
|
|||
|
||||
|
||||
def _load(model_name: str) -> Any:
|
||||
# Stderr print is unconditional so we can see this line even when
|
||||
# structlog routing is misbehaving — diagnostics for a previously
|
||||
# invisible hang.
|
||||
print(
|
||||
f"[rag.reranker] _load entered: model={model_name}",
|
||||
file = sys.stderr,
|
||||
flush = True,
|
||||
)
|
||||
from sentence_transformers import CrossEncoder
|
||||
|
||||
device = _resolve_device()
|
||||
print(
|
||||
f"[rag.reranker] device resolved: {device}",
|
||||
file = sys.stderr,
|
||||
flush = True,
|
||||
)
|
||||
logger.info(
|
||||
"Loading RAG reranker",
|
||||
model = model_name,
|
||||
device = device,
|
||||
)
|
||||
started = time.perf_counter()
|
||||
print(
|
||||
f"[rag.reranker] calling CrossEncoder(...) on {device}",
|
||||
file = sys.stderr,
|
||||
flush = True,
|
||||
)
|
||||
model = CrossEncoder(model_name, device = device)
|
||||
elapsed = round(time.perf_counter() - started, 2)
|
||||
print(
|
||||
f"[rag.reranker] CrossEncoder returned in {elapsed}s",
|
||||
file = sys.stderr,
|
||||
flush = True,
|
||||
)
|
||||
logger.info(
|
||||
"RAG reranker loaded",
|
||||
model = model_name,
|
||||
device = device,
|
||||
elapsed_seconds = round(time.perf_counter() - started, 2),
|
||||
elapsed_seconds = elapsed,
|
||||
)
|
||||
return model
|
||||
|
||||
|
|
@ -129,9 +154,25 @@ def rerank(
|
|||
text_pairs = [(h, t) for h, t in pairs if h.kind != "image"]
|
||||
image_hits = [h for h, _t in pairs if h.kind == "image"]
|
||||
|
||||
print(
|
||||
f"[rag.reranker] rerank entered: n_pairs={len(text_pairs)}",
|
||||
file = sys.stderr,
|
||||
flush = True,
|
||||
)
|
||||
model = get_reranker(model_name)
|
||||
print(
|
||||
"[rag.reranker] reranker model in hand",
|
||||
file = sys.stderr,
|
||||
flush = True,
|
||||
)
|
||||
if text_pairs:
|
||||
inputs = [(query, text) for _, text in text_pairs]
|
||||
print(
|
||||
f"[rag.reranker] predict starting: n_inputs={len(inputs)} "
|
||||
f"batch_size={RAG_RERANK_BATCH_SIZE}",
|
||||
file = sys.stderr,
|
||||
flush = True,
|
||||
)
|
||||
logger.info(
|
||||
"RAG reranker predict starting",
|
||||
n_inputs = len(inputs),
|
||||
|
|
@ -143,10 +184,16 @@ def rerank(
|
|||
batch_size = RAG_RERANK_BATCH_SIZE,
|
||||
show_progress_bar = False,
|
||||
)
|
||||
elapsed = round(time.perf_counter() - started, 2)
|
||||
print(
|
||||
f"[rag.reranker] predict done in {elapsed}s",
|
||||
file = sys.stderr,
|
||||
flush = True,
|
||||
)
|
||||
logger.info(
|
||||
"RAG reranker predict done",
|
||||
n_inputs = len(inputs),
|
||||
elapsed_seconds = round(time.perf_counter() - started, 2),
|
||||
elapsed_seconds = elapsed,
|
||||
)
|
||||
ranked = sorted(
|
||||
zip(text_pairs, scores),
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue