feat: add lifecycle-owned rerank runtime

This commit is contained in:
Patrick Buckley
2026-08-15 16:11:54 -07:00
parent 911468c799
commit 05d75fa2f8
17 changed files with 2081 additions and 172 deletions
+38 -23
View File
@@ -846,23 +846,23 @@ boundary.
```
ChatSession
|
+-- ResolvedModelBinding
| +-- ModelLane (provider, client, model, capabilities, params)
| +-- immutable ModelConfig snapshot
| +-- registry generation
| |
| v
| model_turn(ModelLane, list[Turn])
| +-- lowering.py: Turn IR -> repaired provider-neutral wire dicts
| +-- LLMProvider.create_streaming() (single model transport call site)
| +--- OpenAIProvider --- OpenAI/vLLM/llama.cpp
| +--- AnthropicProvider --- Anthropic Messages API
| +--- GoogleProvider --- Gemini via /v1beta/openai/
|
v
model_turn(ModelLane, list[Turn])
|
+-- lowering.py: Turn IR -> repaired provider-neutral wire dicts
|
v
LLMProvider.create_streaming() (the single transport call site)
|
+--- OpenAIProvider --- OpenAI, vLLM, llama.cpp, any /v1/chat/completions API
+--- AnthropicProvider --- Anthropic Messages API (native streaming, thinking)
+--- GoogleProvider --- Google Gemini via /v1beta/openai/ (extends OpenAIProvider)
+-- RerankLane (runtime, alias admission, registry/config witnesses)
|
v
rerank(RerankLane, query, documents) --> Cohere/Jina-compatible endpoint
```
**Protocol methods:**
@@ -882,6 +882,7 @@ LLMProvider.create_streaming() (the single transport call site)
| `StreamChunk` | `content_delta`, `reasoning_delta`, `tool_call_deltas`, `info_delta`, `usage`, `finish_reason`, `provider_blocks` |
| `CompletionResult` | `content`, `tool_calls`, `finish_reason`, `usage`, `provider_blocks` |
| `ModelLane` | Frozen per-loop provider/client/model binding, capabilities, sampling knobs, registry reference, and backend-auth seam |
| `RerankLane` | Frozen per-batch binding to a registry-owned rerank runtime, stable alias admission gate, and registry/config version witnesses |
| `ResolvedModelBinding` | A `ModelLane`, its immutable `ModelConfig`, and the registry generation read in the same snapshot |
| `ModelTurnResult` | Canonical assistant `Turn`, tool-call dispatch mirror, serving-lane provenance, usage, and exact lowered wire facts |
| `ModelCapabilities` | `context_window`, `max_output_tokens`, `supports_temperature`, `token_param`, `thinking_mode`, `supports_effort`, `supports_web_search`, `supports_tool_search`, `supports_vision`, `supports_reasoning_replay`, `supports_verbosity`, `verbosity`, `supports_pro_mode`, `reasoning_mode` |
@@ -895,6 +896,17 @@ old endpoint/client state and new capabilities/configuration. Per-call operator
toggles that are intentionally live, such as reasoning replay, are re-read by
`model_turn()` through the lane's registry reference.
`RerankLane` is a deliberately narrow sibling rather than a subtype of
`ModelLane`: endpoint reranking does not construct an LLM provider or SDK
client. `ModelRegistry` owns the active rerank runtime, whose HTTP pool,
per-endpoint circuit breaker, and active-call retirement state are shared by
all sessions. Each retrieval batch resolves the Reranker alias and instruction
from one ConfigStore snapshot and binds the matching registry generation.
Relevant configuration changes retire the old runtime, reject new work on its
stale lanes, and close its transport after active calls drain; cap-only reloads
resize the stable alias admission gate without discarding the pool. Calibration
remains an isolated one-shot client and always closes it after the probe.
**OpenAIProvider** (`_openai.py`): passes messages through unchanged (they are
already in OpenAI format), including multi-part content blocks (text + images)
in tool results. Model capability lookup covers GPT-5 through GPT-5.6,
@@ -1001,17 +1013,20 @@ remain warm, but the session still receives a new coherent config/lane.
simultaneous generations for that alias in one process. Every registry-backed
role carries the same gate on its `ModelLane`, so main turns, judges, task
agents, perception, compaction, and background generation coordinate through
one FIFO. Two aliases never share a gate implicitly, even when their URLs are
identical. For calls with context-first request admission, `model_turn()` admits
the request before materializing attachment fallbacks; an oversized refusal
therefore invokes neither attachment resolution nor nested perception/audio.
Ordinary calls preserve their lowering cadence. Both paths finish attachment
materialization before acquiring model capacity, then hold one lease across
eager stream creation and the complete drain, releasing before retry backoff.
Capacity wait is credited out of deadline accounting, preventing queued judges
from spending their request budget before dispatch. The gate survives cap-only
reloads in place; the field is excluded from semantic `ModelConfig` equality so
a capacity edit does not reset judges or output-guard state.
one FIFO. Endpoint reranking carries the same gate on its `RerankLane`, so a
reranker alias cannot exceed that cap or bypass capacity shared with another
role. Two aliases never share a gate implicitly, even when their URLs are
identical. For calls with context-first request admission, `model_turn()`
admits the request before materializing attachment fallbacks; an oversized
refusal therefore invokes neither attachment resolution nor nested
perception/audio. Ordinary calls preserve their lowering cadence. Both paths
finish attachment materialization before acquiring model capacity, then hold
one lease across eager stream creation and the complete drain, releasing before
retry backoff. Capacity wait is credited out of deadline accounting, preventing
queued judges from spending their request budget before dispatch. The gate
survives cap-only reloads in place; the field is excluded from semantic
`ModelConfig` equality so a capacity edit does not reset judges or output-guard
state.
Primary loops, recursive compaction, judges, title generation, audio, and task
agents all consume `ModelLane` rather than inspecting provider/client handles.
+9 -2
View File
@@ -68,8 +68,9 @@ point to the same URL; Turnstone does not infer shared capacity from endpoint
text. Queue time is excluded from judge/output-guard deadline accounting, and
each retry releases its slot before backoff and reacquires for the next wire
attempt. The cap is local to each process, not cluster-wide; account for the
number of nodes targeting the same inference server. Direct STT/TTS protocol
calls and Cohere/Jina reranking do not currently consume this generation cap.
number of nodes targeting the same inference server. Cohere/Jina reranking
selected through the Reranker role consumes the same gate as every other use of
that alias. Direct STT/TTS protocol calls do not consume this generation cap.
### Judge batch parallelism
@@ -479,6 +480,12 @@ The alias's admission gate is retained and resized in place, so a concurrency
edit preserves in-flight accounting and does not reset cached judges or the
output-guard rate limiter.
The Reranker role is resolved from one coherent ConfigStore snapshot for each
retrieval batch, so existing sessions observe alias or instruction changes
without reconstruction. A relevant model-definition edit retires the old
reranker transport and lets active requests drain on their immutable lanes;
cap-only and unrelated edits keep the pooled transport warm.
Sampling and other saved workstream configuration remain workstream state. A
model-definition edit does not silently rewrite a live workstream's chosen
temperature, reasoning effort, max tokens, skill, or persona. Use
+10
View File
@@ -360,6 +360,16 @@ Search the web using a text query.
The `rerank_web_search` toggle defaults on once a reranker is selected. If the endpoint is unreachable or errors, web_search falls back silently to the backend's native result order — reranking never makes a search fail.
The selected alias's `max_concurrency` limit applies to reranking. Turnstone
keeps one process-owned HTTP connection pool for the active endpoint, shares it
across sessions, and rotates it when relevant model or Reranker-role settings
change. Three consecutive endpoint failures open a 30-second circuit so later
retrievals preserve native order immediately instead of repeatedly waiting for
the network timeout; one probe is allowed after cooldown. Queued work observes
generation cancellation. A synchronous HTTP request already dispatched cannot
be interrupted safely, but its per-call timeout still bounds eventual drain
and fallback latency.
When `rerank_bm25` is enabled, Turnstone also sends the current query and BM25 candidate metadata to the rerank endpoint. Live memory-pointer candidates contain only the memory name and authored description—never the body. Tool and skill candidates contain their names and descriptive metadata. A self-hosted endpoint (vLLM/TEI/llama.cpp) keeps this on your infrastructure; a hosted provider (Cohere/Jina/Voyage) sends it off-box.
**Serving a Qwen3-Reranker with vLLM.** The model is instruction-aware, so vLLM **must** apply its chat template — pass `--chat-template` explicitly. Without it the bare query produces near-random scores and reranking actively *hurts* retrieval (verified: an irrelevant passage outscored the correct one):
@@ -70,6 +70,36 @@ def _bootstrap_app(**overrides: Any) -> Any:
return SimpleNamespace(state=SimpleNamespace(**state_kwargs))
def test_partial_coord_teardown_shuts_registry_after_adapter() -> None:
from turnstone.console import server as server_module
calls: list[str] = []
adapter = MagicMock()
adapter.shutdown.side_effect = lambda: calls.append("adapter")
registry = MagicMock()
registry.shutdown.side_effect = lambda: calls.append("registry")
app = _bootstrap_app(coord_adapter=adapter, coord_registry=registry)
server_module._teardown_partial_coord_subsystem(app)
assert calls == ["adapter", "registry"]
assert app.state.coord_registry is None
def test_console_lifespan_shuts_coord_registry_after_adapter() -> None:
"""Pin the normal teardown ordering without booting the full console."""
import inspect
from turnstone.console import server as server_module
source = inspect.getsource(server_module._lifespan)
adapter = source.find('getattr(app.state, "coord_adapter", None)')
registry = source.find('getattr(app.state, "coord_registry", None)', adapter)
state_writer = source.find('getattr(app.state, "coord_state_writer", None)', registry)
assert adapter != -1 and registry != -1 and state_writer != -1
assert adapter < registry < state_writer
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
+127
View File
@@ -3978,6 +3978,133 @@ class TestCancelledSendCleanupOwnership:
session.cancel()
raise GenerationCancelled()
def test_cancel_cleanup_skips_remote_rerank_and_finishes(self, tmp_db) -> None:
"""A set cancel event cannot re-enter reranking from its own repair path."""
from turnstone.core.admission import ModelAdmission
from turnstone.core.memory import save_structured_memory_strict
from turnstone.core.rerank import RerankLane, RerankRuntime
class _NeverBackend:
calls = 0
def rerank(self, *args: Any, **kwargs: Any):
self.calls += 1
raise AssertionError("cancel cleanup dispatched reranking")
ui = NullUI()
session = _make_registered_session(ui=ui, user_id="owner")
session._title_generated = True
save_structured_memory_strict(
"kafka_runbook",
"private body",
description="queued for next seam",
scope="global",
)
backend = _NeverBackend()
lane = RerankLane(
RerankRuntime(backend, alias="rr", model="m"),
"rr",
"m",
ModelAdmission("rr"),
0,
)
try:
with (
patch.object(session, "_resolve_rerank_lane", return_value=lane) as resolve,
patch.object(
session,
"_stream_response",
side_effect=lambda _generation: self._cancel_with_pending_cleanup(session),
),
):
session.send("unrelated opening request")
finally:
lane.runtime.retire()
resolve.assert_not_called()
assert backend.calls == 0
assert any(turn.role is Role.TOOL for turn in session.messages)
assert session.messages[-2].role is Role.USER
assert session.messages[-2].text == "queued for next seam"
assert session.messages[-1].role is Role.SYSTEM
assert session.messages[-1].source == "memory_pointer"
assert "kafka_runbook" in session.messages[-1].text
assert ui.states[-1] == "idle"
def test_queued_turn_rerank_is_fenced_by_generation_owner(self, tmp_db) -> None:
"""A force successor prevents the abandoned queue planner from dispatching."""
from turnstone.core.admission import ModelAdmission
from turnstone.core.memory import save_structured_memory_strict
from turnstone.core.rerank import RerankHit, RerankLane, RerankRuntime
class _CountingBackend:
def __init__(self) -> None:
self.calls = 0
def rerank(self, _query: str, documents: list[str], **_kwargs: Any):
self.calls += 1
return [RerankHit(index=i, score=1.0) for i in range(len(documents))]
session = _make_registered_session(user_id="owner")
save_structured_memory_strict(
"kafka_runbook",
"private body",
description="restart kafka brokers",
scope="global",
)
backend = _CountingBackend()
lane = RerankLane(
RerankRuntime(backend, alias="rr", model="m"),
"rr",
"m",
ModelAdmission("rr"),
0,
)
resolve_started = threading.Event()
release_resolve = threading.Event()
send_errors: list[BaseException] = []
def resolve_lane() -> RerankLane:
resolve_started.set()
if not release_resolve.wait(2):
raise RuntimeError("rerank lane resolution was not released")
return lane
def stream():
yield StreamChunk(content_delta="done")
with session._queued_lock:
session._queued_messages["queued-old"] = (
"restart kafka brokers",
"normal",
)
yield StreamChunk(finish_reason="stop")
arm_session(session, stream())
def run_send() -> None:
try:
session.send("zzzxxyy")
except BaseException as exc:
send_errors.append(exc)
sender = threading.Thread(target=run_send)
try:
with patch.object(session, "_resolve_rerank_lane", side_effect=resolve_lane):
sender.start()
assert resolve_started.wait(2)
assert session._claim_generation() == 2
release_resolve.set()
sender.join(2)
finally:
release_resolve.set()
sender.join(2)
lane.runtime.retire()
assert not sender.is_alive()
assert send_errors == []
assert backend.calls == 0
def test_successor_waits_for_complete_cancel_cleanup_transaction(self, tmp_db):
"""A claim already waiting on the lock observes every cleanup effect."""
ui = NullUI()
+221
View File
@@ -1019,6 +1019,227 @@ class TestResolveEnvVars:
assert _resolve_env_vars("") == ""
# ---------------------------------------------------------------------------
# Registry-owned rerank lanes
# ---------------------------------------------------------------------------
class TestRerankLaneRegistry:
@staticmethod
def _cfg(
alias: str = "rr",
*,
url: str = "http://rerank.example/rerank",
model: str = "bge",
key: str = "secret",
max_concurrency: int = 2,
) -> ModelConfig:
return ModelConfig(
alias,
url,
key,
model,
capabilities={"supports_rerank": True},
max_concurrency=max_concurrency,
)
def test_resolve_reuses_runtime_and_stable_admission_without_llm_client(self) -> None:
cfg = self._cfg()
reg = ModelRegistry({"rr": cfg}, "rr")
first = reg.resolve_rerank_lane("rr", instruction="rank", config_version=4)
second = reg.resolve_rerank_lane("rr", instruction="rank", config_version=5)
assert first.runtime is second.runtime
assert first.admission is second.admission is reg.get_admission("rr")
assert first.config_version == 4
assert second.config_version == 5
assert first.admission.limit == 2
assert reg._clients == {}
assert reg._providers == {}
reg.shutdown()
def test_instruction_change_rotates_and_closes_old_runtime(self) -> None:
reg = ModelRegistry({"rr": self._cfg()}, "rr")
old = reg.resolve_rerank_lane("rr", instruction="old")
new = reg.resolve_rerank_lane("rr", instruction="new")
assert new.runtime is not old.runtime
assert old.runtime.snapshot().retired
assert old.runtime.snapshot().closed
assert not new.runtime.snapshot().retired
reg.shutdown()
def test_cap_only_reload_preserves_runtime_and_resizes_gate(self) -> None:
reg = ModelRegistry({"rr": self._cfg(max_concurrency=1)}, "rr")
old = reg.resolve_rerank_lane("rr", instruction="rank")
reg.reload(
{"rr": self._cfg(max_concurrency=4)},
"rr",
app_state=_KEYED_STATE,
)
new = reg.resolve_rerank_lane("rr", instruction="rank")
assert new.runtime is old.runtime
assert new.admission is old.admission
assert new.admission.limit == 4
assert new.registry_generation == 1
assert not old.runtime.snapshot().retired
reg.shutdown()
@pytest.mark.parametrize(
"replacement",
[
{"url": "http://other.example/rerank"},
{"model": "other"},
{"key": "different"},
],
)
def test_relevant_reload_eagerly_retires_runtime(self, replacement: dict[str, str]) -> None:
reg = ModelRegistry({"rr": self._cfg()}, "rr")
old = reg.resolve_rerank_lane("rr")
reg.reload(
{"rr": self._cfg(**replacement)},
"rr",
app_state=_KEYED_STATE,
)
assert old.runtime.snapshot().retired
assert old.runtime.snapshot().closed
assert reg._rerank_runtimes == {}
new = reg.resolve_rerank_lane("rr")
assert new.runtime is not old.runtime
reg.shutdown()
def test_relevant_reload_lets_active_call_drain_before_close(self) -> None:
from turnstone.core.rerank import RerankHit, rerank
class _BlockingClient:
def __init__(self) -> None:
self.entered = threading.Event()
self.release = threading.Event()
self.close_calls = 0
def rerank(self, query: str, documents: list[str], **kwargs: Any) -> list[RerankHit]:
del query, kwargs
self.entered.set()
assert self.release.wait(5)
return [RerankHit(i, 1.0) for i in range(len(documents))]
def close(self) -> None:
self.close_calls += 1
client = _BlockingClient()
reg = ModelRegistry({"rr": self._cfg()}, "rr")
with patch("turnstone.core.rerank.resolve_rerank_client", return_value=client):
old = reg.resolve_rerank_lane("rr")
outcome: list[list[RerankHit]] = []
worker = threading.Thread(
target=lambda: outcome.append(rerank(old, "q", ["d"], timeout=2.0)),
daemon=True,
)
worker.start()
assert client.entered.wait(2)
reg.reload(
{"rr": self._cfg(url="http://other.example/rerank")},
"rr",
app_state=_KEYED_STATE,
)
assert old.runtime.snapshot().retired
assert not old.runtime.snapshot().closed
assert client.close_calls == 0
client.release.set()
worker.join(2)
assert not worker.is_alive()
assert outcome and outcome[0][0].index == 0
assert old.runtime.snapshot().closed
assert client.close_calls == 1
reg.shutdown()
def test_resolving_new_role_alias_retires_previous_alias(self) -> None:
a = self._cfg("a", url="http://a.example/rerank")
b = self._cfg("b", url="http://b.example/rerank")
reg = ModelRegistry({"a": a, "b": b}, "a")
old = reg.resolve_rerank_lane("a")
current = reg.resolve_rerank_lane("b")
assert old.runtime.snapshot().closed
assert set(reg._rerank_runtimes) == {"b"}
assert current.admission is reg.get_admission("b")
reg.shutdown()
def test_deactivate_and_shutdown_are_idempotent(self) -> None:
reg = ModelRegistry({"rr": self._cfg()}, "rr")
lane = reg.resolve_rerank_lane("rr")
reg.deactivate_rerank_runtime()
reg.deactivate_rerank_runtime()
reg.shutdown()
reg.shutdown()
assert lane.runtime.snapshot().retired
assert lane.runtime.snapshot().closed
assert reg._rerank_runtimes == {}
@pytest.mark.parametrize("operation", ["reload", "shutdown"])
def test_rerank_transport_closes_when_llm_client_close_raises(
self,
operation: str,
) -> None:
class _RerankClient:
def __init__(self) -> None:
self.close_calls = 0
def close(self) -> None:
self.close_calls += 1
class _FailingLLMClient:
def close(self) -> None:
raise RuntimeError("llm close failed")
rr_cfg = self._cfg()
llm_cfg = ModelConfig("llm", "http://old.example/v1", "key", "model")
reg = ModelRegistry({"llm": llm_cfg, "rr": rr_cfg}, "llm")
rerank_client = _RerankClient()
with patch(
"turnstone.core.rerank.resolve_rerank_client",
return_value=rerank_client,
):
lane = reg.resolve_rerank_lane("rr")
reg._clients["llm"] = _FailingLLMClient()
with pytest.raises(RuntimeError, match="llm close failed"):
if operation == "reload":
reg.reload(
{
"llm": dataclasses.replace(
llm_cfg,
base_url="http://new.example/v1",
),
"rr": dataclasses.replace(
rr_cfg,
base_url="http://new-rerank.example/rerank",
),
},
"llm",
app_state=_KEYED_STATE,
)
else:
reg.shutdown()
assert rerank_client.close_calls == 1
assert lane.runtime.snapshot().retired
assert lane.runtime.snapshot().closed
assert reg._rerank_runtimes == {}
# ---------------------------------------------------------------------------
# ModelRegistry.reload
# ---------------------------------------------------------------------------
+130 -54
View File
@@ -19,17 +19,11 @@ from turnstone.core.rerank import (
from turnstone.core.session import ChatSession
def _mock_httpx_post(handler):
"""Patch target for ``rerank.httpx.post`` that routes the call through a real
``httpx.MockTransport``. The request flows through genuine httpx JSON/header
encoding and response parsing a true boundary, not a bare MagicMock.
"""
def _install_mock_httpx_client(monkeypatch, handler):
"""Install the lifecycle-owned client over a real MockTransport boundary."""
client = httpx.Client(transport=httpx.MockTransport(handler))
def _post(url, **kwargs):
return client.post(url, **kwargs)
return _post
monkeypatch.setattr("turnstone.core.rerank.httpx.Client", lambda: client)
return client
# A Cohere/Jina/vLLM-shaped response: results wrapper + relevance_score, returned
@@ -112,7 +106,7 @@ class TestCohereJinaRerankClient:
captured["auth"] = request.headers.get("authorization")
return httpx.Response(200, json=RESULTS_WRAPPED)
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
_install_mock_httpx_client(monkeypatch, handler)
client = CohereJinaRerankClient(
"http://vllm:8000/rerank", model="bge", api_key="secret", timeout=10
)
@@ -137,7 +131,7 @@ class TestCohereJinaRerankClient:
captured["auth"] = request.headers.get("authorization")
return httpx.Response(200, json={"results": []})
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
_install_mock_httpx_client(monkeypatch, handler)
CohereJinaRerankClient("http://h/rerank").rerank("q", ["a"])
assert captured["body"] == {"query": "q", "documents": ["a"]}
@@ -150,7 +144,7 @@ class TestCohereJinaRerankClient:
called["n"] += 1
return httpx.Response(200, json={"results": []})
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
_install_mock_httpx_client(monkeypatch, handler)
assert CohereJinaRerankClient("http://h/rerank").rerank("q", []) == []
assert called["n"] == 0 # short-circuits before any HTTP call
@@ -158,7 +152,7 @@ class TestCohereJinaRerankClient:
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, json=BARE_LIST)
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
_install_mock_httpx_client(monkeypatch, handler)
hits = CohereJinaRerankClient("http://tei/rerank").rerank("q", ["a", "b"])
assert [h.index for h in hits] == [1, 0]
@@ -166,7 +160,7 @@ class TestCohereJinaRerankClient:
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(401, text="unauthorized")
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
_install_mock_httpx_client(monkeypatch, handler)
with pytest.raises(httpx.HTTPStatusError):
CohereJinaRerankClient("http://h/rerank").rerank("q", ["a"])
@@ -208,6 +202,7 @@ class _FakeConfigStore:
def __init__(self, values: dict, stored: set[str] | None = None) -> None:
self._values = dict(values)
self._stored = frozenset(stored if stored is not None else values.keys())
self.version = 1
def stored_keys(self) -> frozenset[str]:
return self._stored
@@ -215,11 +210,14 @@ class _FakeConfigStore:
def get(self, key: str):
return self._values.get(key)
def effective_snapshot(self):
return self.version, dict(self._values)
class TestSessionRerankWiring:
def test_disabled_when_no_config_store(self):
stub = SimpleNamespace(_config_store=None, _registry=None, tool_timeout=30)
assert ChatSession._resolve_rerank_client(stub) is None
assert ChatSession._resolve_rerank_lane(stub) is None
def test_disabled_when_no_reranker_alias(self):
# No Reranker role selected -> reranking off. There is no global endpoint
@@ -230,7 +228,7 @@ class TestSessionRerankWiring:
tool_timeout=30,
_registry=SimpleNamespace(get_config=lambda a: None),
)
assert ChatSession._resolve_rerank_client(stub) is None
assert ChatSession._resolve_rerank_lane(stub) is None
def test_enabled_for_defaults_true_without_store(self):
stub = SimpleNamespace(_config_store=None)
@@ -254,16 +252,18 @@ class TestSessionRerankWiring:
capabilities={"supports_rerank": True},
)
cs = _FakeConfigStore({"tools.reranker_alias": "rr"})
stub = SimpleNamespace(
_config_store=cs,
tool_timeout=30,
_registry=SimpleNamespace(get_config=lambda a: cfg),
)
client = ChatSession._resolve_rerank_client(stub)
from turnstone.core.model_registry import ModelRegistry
registry = ModelRegistry({"rr": cfg}, "rr")
stub = SimpleNamespace(_config_store=cs, tool_timeout=30, _registry=registry)
lane = ChatSession._resolve_rerank_lane(stub)
assert lane is not None
client = lane.runtime.client
assert isinstance(client, CohereJinaRerankClient)
assert client._url == "http://rr:8000/rerank"
assert client._model == "bge"
assert client._api_key == "k"
registry.shutdown()
def test_ignores_alias_without_rerank_capability(self):
# A non-reranker model (no supports_rerank) must NOT be used as a reranker,
@@ -274,12 +274,58 @@ class TestSessionRerankWiring:
alias="chat", base_url="http://chat/v1", api_key="k", model="gpt", capabilities={}
)
cs = _FakeConfigStore({"tools.reranker_alias": "chat"})
stub = SimpleNamespace(
_config_store=cs,
tool_timeout=30,
_registry=SimpleNamespace(get_config=lambda a: cfg),
from turnstone.core.model_registry import ModelRegistry
registry = ModelRegistry({"chat": cfg}, "chat")
stub = SimpleNamespace(_config_store=cs, tool_timeout=30, _registry=registry)
assert ChatSession._resolve_rerank_lane(stub) is None
registry.shutdown()
def test_unrelated_settings_version_reuses_runtime(self):
from turnstone.core.model_registry import ModelConfig, ModelRegistry
cfg = ModelConfig(
"rr",
"http://rr/rerank",
"k",
"bge",
capabilities={"supports_rerank": True},
)
assert ChatSession._resolve_rerank_client(stub) is None
cs = _FakeConfigStore({"tools.reranker_alias": "rr", "tools.rerank_instruction": "rank"})
registry = ModelRegistry({"rr": cfg}, "rr")
stub = SimpleNamespace(_config_store=cs, _registry=registry)
first = ChatSession._resolve_rerank_lane(stub)
cs._values["unrelated.setting"] = True
cs.version += 1
second = ChatSession._resolve_rerank_lane(stub)
assert first is not None and second is not None
assert first.runtime is second.runtime
assert first.config_version != second.config_version
registry.shutdown()
def test_clearing_role_retires_previous_runtime_on_next_resolution(self):
from turnstone.core.model_registry import ModelConfig, ModelRegistry
cfg = ModelConfig(
"rr",
"http://rr/rerank",
"k",
"bge",
capabilities={"supports_rerank": True},
)
cs = _FakeConfigStore({"tools.reranker_alias": "rr"})
registry = ModelRegistry({"rr": cfg}, "rr")
stub = SimpleNamespace(_config_store=cs, _registry=registry)
lane = ChatSession._resolve_rerank_lane(stub)
assert lane is not None
cs._values["tools.reranker_alias"] = ""
cs.version += 1
assert ChatSession._resolve_rerank_lane(stub) is None
assert lane.runtime.snapshot().closed
registry.shutdown()
# ---------------------------------------------------------------------------
@@ -287,18 +333,6 @@ class TestSessionRerankWiring:
# ---------------------------------------------------------------------------
class _FakeRerankClient:
"""In-process RerankClient stub returning fixed hits (no HTTP)."""
def __init__(self, hits: list[RerankHit]) -> None:
self._hits = hits
def rerank(
self, query: str, documents: list[str], *, top_n: int | None = None
) -> list[RerankHit]:
return self._hits
class TestRerankInstruction:
"""`rerank_instruction` wraps the query for instruction-aware rerankers.
@@ -318,13 +352,13 @@ class TestRerankInstruction:
def test_no_instruction_sends_bare_query(self, monkeypatch):
sent, handler = self._capture()
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
_install_mock_httpx_client(monkeypatch, handler)
CohereJinaRerankClient("http://x/rerank").rerank("capital of France", ["d0"])
assert sent["body"]["query"] == "capital of France"
def test_instruction_wraps_query(self, monkeypatch):
sent, handler = self._capture()
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
_install_mock_httpx_client(monkeypatch, handler)
CohereJinaRerankClient("http://x/rerank", instruction="Find relevant passages").rerank(
"capital of France", ["d0"]
)
@@ -374,31 +408,35 @@ class TestNormalizeScores:
class TestSessionBM25Reranker:
"""``_bm25_reranker`` / ``_bm25_rerank_threshold`` — the BM25 seam adapters.
Drives the real closure through a fake ``RerankClient`` (the in-process
callable seam), never by patching internal state. The HTTP boundary lives in
``TestCohereJinaRerankClient`` above.
Drives the real closure through fixed hits at the session dispatch seam.
The HTTP boundary lives in ``TestCohereJinaRerankClient`` above.
"""
def test_none_when_disabled_for_bm25(self):
# Per-tool toggle off -> no reranker even with an endpoint configured.
def test_disabled_bm25_is_identity_order(self):
# Long-lived closures stay installed across settings changes; disabled
# state is an identity rerank and must not resolve a runtime.
stub = SimpleNamespace(
_rerank_enabled_for=lambda tool: False,
_resolve_rerank_client=lambda: _FakeRerankClient([]),
_rerank_hits=lambda *args, **kwargs: (_ for _ in ()).throw(
AssertionError("disabled reranker dispatched")
),
)
assert ChatSession._bm25_reranker(stub) is None
rank = ChatSession._bm25_reranker(stub)
assert rank("q", ["d0", "d1"]) == [0, 1]
def test_none_when_no_client(self):
# Enabled, but no endpoint resolves -> None.
def test_no_lane_is_identity_order(self):
# Enabled, but no endpoint resolves: preserve BM25 pool order.
stub = SimpleNamespace(
_rerank_enabled_for=lambda tool: True,
_resolve_rerank_client=lambda: None,
_rerank_hits=lambda *args, **kwargs: None,
)
assert ChatSession._bm25_reranker(stub) is None
rank = ChatSession._bm25_reranker(stub)
assert rank("q", ["d0", "d1"]) == [0, 1]
def _enabled_stub(self, hits: list[RerankHit]) -> SimpleNamespace:
return SimpleNamespace(
_rerank_enabled_for=lambda tool: True,
_resolve_rerank_client=lambda: _FakeRerankClient(hits),
_rerank_hits=lambda query, docs, origin_generation=0: hits,
)
def test_no_threshold_returns_all_hit_indices(self):
@@ -461,6 +499,44 @@ class TestSessionBM25Reranker:
assert rank is not None
assert rank("q", ["d0", "d1"]) == [0]
def test_generation_cancellation_bypasses_bm25_exception_fallback(self):
import threading
from types import MethodType
from turnstone.core.admission import ModelAdmission
from turnstone.core.bm25 import BM25Index
from turnstone.core.rerank import RerankLane, RerankRuntime
from turnstone.core.session import GenerationCancelled
class _NeverBackend:
def rerank(self, *args, **kwargs):
raise AssertionError("superseded generation dispatched")
lane = RerankLane(
RerankRuntime(_NeverBackend(), alias="rr", model="m"),
"rr",
"m",
ModelAdmission("rr"),
0,
)
stub = SimpleNamespace(
_generation=2,
_generation_lock=threading.Lock(),
_cancel_event=threading.Event(),
_publication_shutdown=False,
_rerank_enabled_for=lambda tool: True,
_resolve_rerank_lane=lambda: lane,
tool_timeout=10,
)
stub._check_cancelled = MethodType(ChatSession._check_cancelled, stub)
stub._rerank_hits = MethodType(ChatSession._rerank_hits, stub)
rank = ChatSession._bm25_reranker(stub, origin_generation=1)
with pytest.raises(GenerationCancelled):
BM25Index(["alpha"], reranker=rank).search("alpha")
assert lane.runtime.snapshot().circuit.consecutive_failures == 0
lane.runtime.retire()
def test_threshold_reads_setting(self):
cs = _FakeConfigStore({"tools.rerank_bm25_threshold": 0.42})
stub = SimpleNamespace(_config_store=cs)
+45 -1
View File
@@ -2,8 +2,16 @@
from __future__ import annotations
import pytest
from turnstone.core.rerank import RerankHit
from turnstone.core.rerank_calibrate import _GAP_FRACTION, _PROBE_SET, _build_result, calibrate
from turnstone.core.rerank_calibrate import (
_GAP_FRACTION,
_PROBE_SET,
_build_result,
calibrate,
calibrate_model,
)
class _ScriptedClient:
@@ -95,6 +103,42 @@ class TestCalibrate:
assert res.raw_scale == "unknown (no scores)"
class TestCalibrateModelLifecycle:
class _ClosableClient(_ScriptedClient):
def __init__(self, *, fail: bool = False) -> None:
super().__init__(0.9, 0.1)
self.fail = fail
self.close_calls = 0
def rerank(
self, query: str, documents: list[str], *, top_n: int | None = None
) -> list[RerankHit]:
if self.fail:
raise RuntimeError("probe failed")
return super().rerank(query, documents, top_n=top_n)
def close(self) -> None:
self.close_calls += 1
def test_one_shot_client_closes_after_success(self, monkeypatch) -> None:
client = self._ClosableClient()
monkeypatch.setattr("turnstone.core.rerank.resolve_rerank_client", lambda *a, **k: client)
result = calibrate_model("http://rr/rerank", "m", "k")
assert result.separated
assert client.close_calls == 1
def test_one_shot_client_closes_after_failure(self, monkeypatch) -> None:
client = self._ClosableClient(fail=True)
monkeypatch.setattr("turnstone.core.rerank.resolve_rerank_client", lambda *a, **k: client)
with pytest.raises(RuntimeError, match="probe failed"):
calibrate_model("http://rr/rerank", "m", "k")
assert client.close_calls == 1
class TestCalibrationCapsConfinement:
def test_calibrate_merge_confined_to_calibration_fields(self):
"""The merge touches only the three probe-derived keys and preserves
+256
View File
@@ -0,0 +1,256 @@
"""Live lifecycle and admission checks for the managed rerank runtime.
The local counting proxy preserves real endpoint responses while making the
client-side pool and admission boundary observable. Run explicitly against a
Cohere/Jina-compatible endpoint::
TURNSTONE_LIVE_RERANK_URL=http://127.0.0.1:8000/rerank \
TURNSTONE_LIVE_RERANK_MODEL=my-reranker \
pytest tests/test_rerank_live.py -m live -v
``TURNSTONE_LIVE_RERANK_API_KEY`` is optional.
"""
from __future__ import annotations
import math
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import TYPE_CHECKING, Any
import httpx
import pytest
from turnstone.core.model_registry import ModelConfig, ModelRegistry
from turnstone.core.rerank import RerankHit, RerankLane, rerank
if TYPE_CHECKING:
from types import TracebackType
_LIVE_URL = os.environ.get("TURNSTONE_LIVE_RERANK_URL", "").strip()
_LIVE_MODEL = os.environ.get("TURNSTONE_LIVE_RERANK_MODEL", "").strip()
_LIVE_API_KEY = os.environ.get("TURNSTONE_LIVE_RERANK_API_KEY", "").strip()
_HOP_BY_HOP_HEADERS = frozenset(
{
"connection",
"content-encoding",
"content-length",
"host",
"keep-alive",
"proxy-authenticate",
"proxy-authorization",
"te",
"trailer",
"transfer-encoding",
"upgrade",
}
)
@dataclass(frozen=True)
class _Observed:
requests: int
active: int
peak: int
accepted_connections: int
class _Counter:
"""Count proxy requests and optionally hold the first admitted cohort."""
def __init__(self, *, rendezvous_size: int = 1, hold_seconds: float = 0.0) -> None:
self._lock = threading.Lock()
self._rendezvous = threading.Event()
self._rendezvous_size = rendezvous_size
self._hold_seconds = hold_seconds
self._requests = 0
self._active = 0
self._peak = 0
self._accepted_connections = 0
if rendezvous_size <= 1:
self._rendezvous.set()
def accepted(self) -> None:
with self._lock:
self._accepted_connections += 1
def enter(self) -> None:
with self._lock:
self._requests += 1
self._active += 1
self._peak = max(self._peak, self._active)
if self._active >= self._rendezvous_size:
self._rendezvous.set()
assert self._rendezvous.wait(10), "live rerank requests did not overlap"
if self._hold_seconds:
time.sleep(self._hold_seconds)
def leave(self) -> None:
with self._lock:
self._active -= 1
def snapshot(self) -> _Observed:
with self._lock:
return _Observed(
requests=self._requests,
active=self._active,
peak=self._peak,
accepted_connections=self._accepted_connections,
)
class _LiveProxyServer(ThreadingHTTPServer):
daemon_threads = True
def __init__(self, counter: _Counter) -> None:
self.counter = counter
headers = {"Authorization": f"Bearer {_LIVE_API_KEY}"} if _LIVE_API_KEY else {}
self.upstream = httpx.Client(
headers=headers,
timeout=httpx.Timeout(60.0, connect=5.0),
)
super().__init__(("127.0.0.1", 0), _LiveProxyHandler)
def get_request(self):
request, address = super().get_request()
self.counter.accepted()
return request, address
class _LiveProxyHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
@property
def _proxy(self) -> _LiveProxyServer:
if not isinstance(self.server, _LiveProxyServer):
raise TypeError("live rerank handler requires _LiveProxyServer")
return self.server
def do_POST(self) -> None: # noqa: N802 - stdlib handler contract
length = int(self.headers.get("content-length", "0"))
body = self.rfile.read(length)
self._proxy.counter.enter()
try:
headers = {
name: value
for name, value in self.headers.items()
if name.lower() not in _HOP_BY_HOP_HEADERS and name.lower() != "authorization"
}
response = self._proxy.upstream.post(_LIVE_URL, content=body, headers=headers)
payload = response.content
self.send_response(response.status_code)
content_type = response.headers.get("content-type")
if content_type:
self.send_header("Content-Type", content_type)
self.send_header("Content-Length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
self.wfile.flush()
finally:
self._proxy.counter.leave()
def log_message(self, _format: str, *args: Any) -> None:
del args
class _LiveProxy:
def __init__(self, counter: _Counter) -> None:
self._server = _LiveProxyServer(counter)
self._thread = threading.Thread(
target=self._server.serve_forever,
name="rerank-live-proxy",
daemon=True,
)
@property
def url(self) -> str:
host, port = self._server.server_address
return f"http://{host}:{port}/rerank"
def __enter__(self) -> _LiveProxy:
self._thread.start()
return self
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> None:
del exc_type, exc_value, traceback
self._server.shutdown()
self._server.server_close()
self._server.upstream.close()
self._thread.join(timeout=5)
if self._thread.is_alive():
raise AssertionError("live rerank proxy did not stop")
def _registry_lane(url: str, *, limit: int) -> tuple[ModelRegistry, RerankLane]:
cfg = ModelConfig(
alias="live-reranker",
base_url=url,
api_key="",
model=_LIVE_MODEL,
max_concurrency=limit,
capabilities={"supports_rerank": True},
)
registry = ModelRegistry({cfg.alias: cfg}, default=cfg.alias)
return registry, registry.resolve_rerank_lane(cfg.alias)
def _dispatch(lane: RerankLane) -> list[RerankHit]:
return rerank(
lane,
"What is the capital of France?",
["Paris is the capital of France.", "Whales are mammals."],
timeout=30.0,
)
@pytest.mark.live
@pytest.mark.skipif(
not _LIVE_URL or not _LIVE_MODEL,
reason="TURNSTONE_LIVE_RERANK_URL and TURNSTONE_LIVE_RERANK_MODEL are required",
)
class TestLiveRerankRuntime:
def test_real_scores_and_reuses_one_client_connection(self) -> None:
counter = _Counter()
with _LiveProxy(counter) as proxy:
registry, first_lane = _registry_lane(proxy.url, limit=2)
try:
second_lane = registry.resolve_rerank_lane("live-reranker")
assert second_lane.runtime is first_lane.runtime
first = _dispatch(first_lane)
second = _dispatch(second_lane)
finally:
registry.shutdown()
assert first[0].index == second[0].index == 0
assert all(math.isfinite(hit.score) for hit in [*first, *second])
observed = counter.snapshot()
assert observed.requests == 2
assert observed.active == 0
assert observed.accepted_connections == 1
def test_real_endpoint_never_exceeds_alias_cap(self) -> None:
counter = _Counter(rendezvous_size=2, hold_seconds=0.25)
with _LiveProxy(counter) as proxy:
registry, lane = _registry_lane(proxy.url, limit=2)
try:
with ThreadPoolExecutor(max_workers=4) as pool:
futures = [pool.submit(_dispatch, lane) for _ in range(4)]
results = [future.result(timeout=75) for future in futures]
finally:
registry.shutdown()
assert all(result and result[0].index == 0 for result in results)
observed = counter.snapshot()
assert observed.requests == 4
assert observed.active == 0
assert observed.peak == 2
+408
View File
@@ -0,0 +1,408 @@
"""Lifecycle, admission, circuit, and real keep-alive tests for RerankLane."""
from __future__ import annotations
import json
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any
import pytest
from turnstone.core.admission import ModelAdmission
from turnstone.core.deadline import DeadlineCancelledError, StreamAbortRef
from turnstone.core.rerank import (
CohereJinaRerankClient,
RerankCircuitOpenError,
RerankHit,
RerankLane,
RerankRuntime,
RerankRuntimeRetiredError,
rerank,
)
class _Backend:
"""Thread-safe scripted backend with close and concurrency accounting."""
def __init__(self) -> None:
self.calls = 0
self.active = 0
self.max_active = 0
self.close_calls = 0
self.fail = False
self.entered = threading.Event()
self.release: threading.Event | None = None
self.on_call: Any = None
self._lock = threading.Lock()
def rerank(
self,
query: str,
documents: list[str],
*,
top_n: int | None = None,
timeout: float | None = None,
) -> list[RerankHit]:
del query, top_n, timeout
with self._lock:
self.calls += 1
self.active += 1
self.max_active = max(self.max_active, self.active)
self.entered.set()
try:
if self.on_call is not None:
self.on_call()
if self.release is not None:
assert self.release.wait(5), "test backend release timed out"
if self.fail:
raise RuntimeError("rerank endpoint failed")
return [RerankHit(index=i, score=1.0 - i / 100) for i in range(len(documents))]
finally:
with self._lock:
self.active -= 1
def close(self) -> None:
with self._lock:
self.close_calls += 1
def _lane(
backend: _Backend,
*,
limit: int = 0,
failure_threshold: int = 3,
cooldown_seconds: float = 30.0,
clock: Any = time.monotonic,
) -> RerankLane:
runtime = RerankRuntime(
backend,
alias="rr",
model="m",
failure_threshold=failure_threshold,
cooldown_seconds=cooldown_seconds,
clock=clock,
)
return RerankLane(runtime, "rr", "m", ModelAdmission("rr", limit), 0)
def _call(lane: RerankLane, cancel_ref: Any = None) -> list[RerankHit]:
return rerank(lane, "q", ["d0", "d1"], timeout=2.0, cancel_ref=cancel_ref)
class TestRuntimeLifecycle:
def test_empty_batch_touches_no_runtime_state(self) -> None:
backend = _Backend()
lane = _lane(backend)
assert rerank(lane, "q", [], timeout=1.0) == []
snapshot = lane.runtime.snapshot()
assert backend.calls == 0
assert snapshot.active_calls == 0
assert snapshot.circuit.state == "closed"
assert snapshot.circuit.consecutive_failures == 0
def test_retire_idle_closes_exactly_once(self) -> None:
backend = _Backend()
runtime = _lane(backend).runtime
runtime.retire()
runtime.retire()
assert backend.close_calls == 1
assert runtime.snapshot().retired
assert runtime.snapshot().closed
with pytest.raises(RerankRuntimeRetiredError):
runtime.acquire_call()
def test_active_call_drains_before_close_and_new_call_is_refused(self) -> None:
backend = _Backend()
backend.release = threading.Event()
lane = _lane(backend)
outcome: list[Any] = []
worker = threading.Thread(target=lambda: outcome.append(_call(lane)), daemon=True)
worker.start()
assert backend.entered.wait(2)
lane.runtime.retire()
during = lane.runtime.snapshot()
assert during.retired
assert not during.closed
assert backend.close_calls == 0
with pytest.raises(RerankRuntimeRetiredError):
_call(lane)
assert backend.calls == 1
backend.release.set()
worker.join(2)
assert not worker.is_alive()
assert len(outcome) == 1
after = lane.runtime.snapshot()
assert after.closed
assert backend.close_calls == 1
class TestAdmission:
def test_shared_alias_cap_bounds_peak_dispatch(self) -> None:
backend = _Backend()
backend.release = threading.Event()
lane = _lane(backend, limit=2)
with ThreadPoolExecutor(max_workers=4) as pool:
futures = [pool.submit(_call, lane) for _ in range(4)]
deadline = time.monotonic() + 2
while backend.active < 2 and time.monotonic() < deadline:
time.sleep(0.01)
assert backend.active == 2
assert lane.admission.snapshot().queued == 2
backend.release.set()
assert all(len(f.result(timeout=2)) == 2 for f in futures)
assert backend.calls == 4
assert backend.max_active == 2
assert lane.admission.snapshot().in_flight == 0
def test_cancelled_waiter_is_removed_without_dispatch_or_circuit_failure(self) -> None:
backend = _Backend()
lane = _lane(backend, limit=1)
held = lane.admission.acquire()
cancel_ref = StreamAbortRef()
errors: list[BaseException] = []
def _waiter() -> None:
try:
_call(lane, cancel_ref)
except BaseException as exc: # capture the exact cancellation type
errors.append(exc)
worker = threading.Thread(target=_waiter, daemon=True)
worker.start()
deadline = time.monotonic() + 2
while lane.admission.snapshot().queued != 1 and time.monotonic() < deadline:
time.sleep(0.01)
assert lane.admission.snapshot().queued == 1
cancel_ref.abort()
worker.join(2)
held.release()
assert not worker.is_alive()
assert len(errors) == 1
assert isinstance(errors[0], DeadlineCancelledError)
assert backend.calls == 0
snapshot = lane.runtime.snapshot()
assert snapshot.circuit.state == "closed"
assert snapshot.circuit.consecutive_failures == 0
def test_post_request_cancellation_is_not_endpoint_failure(self) -> None:
cancel_ref = StreamAbortRef()
backend = _Backend()
backend.on_call = cancel_ref.abort
lane = _lane(backend)
with pytest.raises(DeadlineCancelledError):
_call(lane, cancel_ref)
assert backend.calls == 1
snapshot = lane.runtime.snapshot()
assert snapshot.circuit.state == "closed"
assert snapshot.circuit.consecutive_failures == 0
def test_cancellation_wins_when_dispatched_request_also_fails(self) -> None:
cancel_ref = StreamAbortRef()
backend = _Backend()
backend.on_call = cancel_ref.abort
backend.fail = True
lane = _lane(backend)
with pytest.raises(DeadlineCancelledError):
_call(lane, cancel_ref)
assert backend.calls == 1
snapshot = lane.runtime.snapshot()
assert snapshot.circuit.state == "closed"
assert snapshot.circuit.consecutive_failures == 0
class _Clock:
def __init__(self) -> None:
self.now = 100.0
def __call__(self) -> float:
return self.now
class TestCircuit:
def test_open_fast_fallback_and_successful_half_open_probe(self) -> None:
clock = _Clock()
backend = _Backend()
backend.fail = True
lane = _lane(backend, failure_threshold=3, cooldown_seconds=30, clock=clock)
for _ in range(3):
with pytest.raises(RuntimeError, match="endpoint failed"):
_call(lane)
assert backend.calls == 3
assert lane.runtime.snapshot().circuit.state == "open"
with pytest.raises(RerankCircuitOpenError):
_call(lane)
assert backend.calls == 3 # no HTTP/backend dispatch while open
clock.now += 30
backend.fail = False
assert len(_call(lane)) == 2
assert backend.calls == 4
recovered = lane.runtime.snapshot().circuit
assert recovered.state == "closed"
assert recovered.consecutive_failures == 0
def test_open_circuit_preserves_bm25_order_without_another_dispatch(self) -> None:
from turnstone.core.bm25 import BM25Index
backend = _Backend()
backend.fail = True
lane = _lane(backend, failure_threshold=1)
documents = ["alpha alpha", "alpha beta", "beta"]
expected = BM25Index(documents).search("alpha", k=3)
def _rank(query: str, candidates: list[str]) -> list[int]:
return [hit.index for hit in rerank(lane, query, candidates, timeout=2.0)]
index = BM25Index(documents, reranker=_rank)
assert index.search("alpha", k=3) == expected
assert backend.calls == 1
assert index.search("alpha", k=3) == expected
assert backend.calls == 1
def test_only_one_half_open_probe_dispatches(self) -> None:
clock = _Clock()
backend = _Backend()
backend.fail = True
lane = _lane(backend, failure_threshold=1, cooldown_seconds=5, clock=clock)
with pytest.raises(RuntimeError):
_call(lane)
clock.now += 5
backend.fail = False
backend.entered.clear()
backend.release = threading.Event()
outcome: list[Any] = []
probe = threading.Thread(target=lambda: outcome.append(_call(lane)), daemon=True)
probe.start()
assert backend.entered.wait(2)
with pytest.raises(RerankCircuitOpenError):
_call(lane)
assert backend.calls == 2
backend.release.set()
probe.join(2)
assert not probe.is_alive()
assert len(outcome) == 1
assert lane.runtime.snapshot().circuit.state == "closed"
def test_failed_half_open_probe_starts_fresh_cooldown(self) -> None:
clock = _Clock()
backend = _Backend()
backend.fail = True
lane = _lane(backend, failure_threshold=1, cooldown_seconds=10, clock=clock)
with pytest.raises(RuntimeError):
_call(lane)
clock.now += 10
with pytest.raises(RuntimeError):
_call(lane)
assert backend.calls == 2
clock.now += 9
with pytest.raises(RerankCircuitOpenError):
_call(lane)
assert backend.calls == 2
def test_waiters_admitted_after_open_do_not_dispatch(self) -> None:
backend = _Backend()
backend.fail = True
backend.release = threading.Event()
lane = _lane(backend, limit=1, failure_threshold=3)
with ThreadPoolExecutor(max_workers=4) as pool:
futures = [pool.submit(_call, lane) for _ in range(4)]
assert backend.entered.wait(2)
deadline = time.monotonic() + 2
while lane.admission.snapshot().queued != 3 and time.monotonic() < deadline:
time.sleep(0.01)
assert lane.admission.snapshot().queued == 3
backend.release.set()
errors = [future.exception(timeout=2) for future in futures]
assert sum(type(error) is RuntimeError for error in errors) == 3
assert sum(isinstance(error, RerankCircuitOpenError) for error in errors) == 1
assert backend.calls == 3
assert lane.admission.snapshot().in_flight == 0
assert lane.runtime.snapshot().circuit.state == "open"
class _ConnectionCountingServer(ThreadingHTTPServer):
daemon_threads = True
def __init__(self) -> None:
self.accepted_connections = 0
self.requests = 0
super().__init__(("127.0.0.1", 0), _KeepAliveHandler)
def get_request(self):
request, address = super().get_request()
self.accepted_connections += 1
return request, address
class _KeepAliveHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def do_POST(self) -> None: # noqa: N802 - stdlib handler contract
length = int(self.headers.get("content-length", "0"))
json.loads(self.rfile.read(length))
server = self.server
assert isinstance(server, _ConnectionCountingServer)
server.requests += 1
body = json.dumps(
{"results": [{"index": 0, "relevance_score": 0.9}]},
separators=(",", ":"),
).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
self.wfile.flush()
def log_message(self, fmt: str, *args: Any) -> None:
del fmt, args
def test_repeated_reranks_reuse_one_real_tcp_connection() -> None:
server = _ConnectionCountingServer()
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
url = f"http://127.0.0.1:{server.server_port}/rerank"
client = CohereJinaRerankClient(url, model="m")
lane = RerankLane(
RerankRuntime(client, alias="rr", model="m"),
"rr",
"m",
ModelAdmission("rr"),
0,
)
try:
assert [hit.index for hit in _call(lane)] == [0]
assert [hit.index for hit in _call(lane)] == [0]
assert server.requests == 2
assert server.accepted_connections == 1
finally:
lane.runtime.retire()
server.shutdown()
server.server_close()
thread.join(2)
+1 -1
View File
@@ -106,7 +106,7 @@ def _make_session(*, kind: str = "interactive", user_id: str = "test-user") -> A
# Truncation budget — required by _truncate_output on every exec.
session.tool_truncation = 100_000
# skills(action='find') ranks via BM25Index(..., reranker=self._bm25_reranker()),
# which reaches _resolve_rerank_client -> self.tool_timeout. No _config_store/
# which reaches _resolve_rerank_lane -> self.tool_timeout. No _config_store/
# _registry here -> no endpoint -> reranker is None -> pure-BM25 path.
session.tool_timeout = 30
+30
View File
@@ -5382,6 +5382,10 @@ def _bootstrap_coord_subsystem(
coord_adapter.shutdown()
except Exception:
log.warning("console.coord_bootstrap_rollback_adapter_failed", exc_info=True)
try:
coord_registry.shutdown()
except Exception:
log.warning("console.coord_bootstrap_rollback_registry_failed", exc_info=True)
try:
shutdown_idle_nudge_watchers(app)
except Exception:
@@ -5470,6 +5474,7 @@ def _load_and_bootstrap_coord_subsystem(app: Starlette, storage: Any, config_sto
"""
from turnstone.core.model_registry import load_model_registry
coord_registry: Any | None = None
try:
try:
coord_registry = load_model_registry(storage=storage)
@@ -5489,6 +5494,11 @@ def _load_and_bootstrap_coord_subsystem(app: Starlette, storage: Any, config_sto
_bootstrap_coord_subsystem(app, storage, config_store, coord_registry)
except Exception:
log.warning("console.coordinator_init_failed", exc_info=True)
if coord_registry is not None:
try:
coord_registry.shutdown()
except Exception:
log.warning("console.coord_startup_registry_shutdown_failed", exc_info=True)
# ``_bootstrap_coord_subsystem`` rolls back its own partial
# side-effects from locals before re-raising, so this helper
# is normally redundant — kept as defence-in-depth in case
@@ -5556,6 +5566,13 @@ def _teardown_partial_coord_subsystem(app: Any) -> None:
except Exception:
log.warning("console.coord_partial_adapter_shutdown_failed", exc_info=True)
registry = getattr(state, "coord_registry", None)
if registry is not None:
try:
registry.shutdown()
except Exception:
log.warning("console.coord_partial_registry_shutdown_failed", exc_info=True)
# Idle nudge watchers are tracked in a list on app.state; the
# console only ever installs one (the coord watcher), so a blanket
# shutdown is safe — there is no other watcher to tear down by
@@ -5875,6 +5892,15 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
coord_adapter_shutdown.shutdown()
except Exception:
log.debug("console.coord_adapter_shutdown_failed", exc_info=True)
coord_registry_shutdown = getattr(app.state, "coord_registry", None)
if coord_registry_shutdown is not None:
try:
# Coordinator sessions stop through the adapter first; the registry
# can then retire pooled model and rerank transports without a new
# session resolving behind the teardown.
await asyncio.to_thread(coord_registry_shutdown.shutdown)
except Exception:
log.debug("console.coord_registry_shutdown_failed", exc_info=True)
coord_state_writer_shutdown = getattr(app.state, "coord_state_writer", None)
if coord_state_writer_shutdown is not None:
try:
@@ -12852,6 +12878,10 @@ def _maybe_bootstrap_coord_subsystem(app: Any, storage: Any) -> None:
_bootstrap_coord_subsystem(app, storage, config_store, coord_registry)
except Exception as exc:
log.warning("console.coord_bootstrap_failed", exc_info=True)
try:
coord_registry.shutdown()
except Exception:
log.warning("console.coord_bootstrap_registry_shutdown_failed", exc_info=True)
# Tear down any partially-stamped handles so a later retry
# via the same CRUD path doesn't spawn duplicate daemons.
_teardown_partial_coord_subsystem(app)
+159 -3
View File
@@ -7,6 +7,8 @@ resilience when the primary model is unreachable.
from __future__ import annotations
import contextlib
import hashlib
import re
import threading
from dataclasses import dataclass, field
@@ -14,7 +16,9 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Mapping
from collections.abc import Callable, Iterator, Mapping
from turnstone.core.rerank import RerankLane, RerankRuntime
from turnstone.core.admission import ModelAdmission
from turnstone.core.config import load_config
@@ -23,6 +27,20 @@ from turnstone.core.providers import LLMProvider, create_client, create_provider
log = get_logger(__name__)
@contextlib.contextmanager
def _deferred_close_actions(actions: list[Callable[[], None]]) -> Iterator[None]:
"""Run collected transport closes after the enclosing registry lock exits."""
try:
yield
finally:
for close in actions:
try:
close()
except Exception:
log.warning("rerank.deferred_close_failed", exc_info=True)
MODEL_AUTH_MODES = frozenset({"static", "entra_obo", "entra_app", "rfc8693_obo"})
# One bound for the backend-auth text columns (obo_audience, obo_scopes),
@@ -448,6 +466,37 @@ def _validate_registry_args(
raise ValueError(f"Task model '{task_model}' not found in registry")
def _rerank_runtime_fingerprint(cfg: ModelConfig, instruction: str) -> bytes:
"""Return a non-secret digest of configuration that changes rerank behavior."""
caps = cfg.capabilities
values = (
cfg.base_url.strip(),
cfg.api_key,
cfg.model.strip(),
cfg.provider,
cfg.auth_mode,
cfg.obo_audience,
cfg.obo_scopes,
str(caps.get("rerank_mode") or "endpoint"),
str(bool(caps.get("supports_rerank"))),
str(bool(caps.get("supports_prefill_rerank"))),
instruction.strip(),
)
digest = hashlib.sha256()
for value in values:
encoded = value.encode("utf-8")
digest.update(len(encoded).to_bytes(8, "big"))
digest.update(encoded)
return digest.digest()
@dataclass(frozen=True)
class _RerankRuntimeEntry:
runtime: RerankRuntime
fingerprint: bytes = field(repr=False)
instruction: str = field(repr=False)
class ModelRegistry:
"""Holds named model configurations with thread-safe lazy client creation.
@@ -482,6 +531,7 @@ class ModelRegistry:
self.task_effort = task_effort
self._clients: dict[str, Any] = {}
self._providers: dict[str, LLMProvider] = {}
self._rerank_runtimes: dict[str, _RerankRuntimeEntry] = {}
self._admissions = {
alias: ModelAdmission(alias, int(getattr(cfg, "max_concurrency", 0)))
for alias, cfg in self._models.items()
@@ -577,6 +627,88 @@ class ModelRegistry:
raise UnknownModelAliasError(alias)
return gate
def resolve_rerank_lane(
self,
alias: str,
*,
instruction: str = "",
config_version: int = 0,
) -> RerankLane:
"""Resolve one coherent endpoint-backed rerank binding.
Unlike :meth:`resolve_binding`, this path deliberately constructs no
LLM provider or SDK client: a Cohere/Jina rerank route is not a chat
model surface. It shares only the alias's stable admission object and
registry generation with normal model lanes.
"""
from turnstone.core.rerank import RerankLane, RerankRuntime, resolve_rerank_client
close_actions: list[Callable[[], None]] = []
instruction = instruction.strip()
with _deferred_close_actions(close_actions), self._client_lock:
cfg = self._models.get(alias)
if cfg is None:
raise UnknownModelAliasError(alias)
if not cfg.base_url.strip() or not cfg.capabilities.get("supports_rerank"):
raise ValueError(f"Model alias {alias!r} is not a configured reranker")
fingerprint = _rerank_runtime_fingerprint(cfg, instruction)
entry = self._rerank_runtimes.get(alias)
if entry is None or entry.fingerprint != fingerprint:
client = resolve_rerank_client(
cfg.base_url,
model=cfg.model,
api_key=cfg.api_key,
instruction=instruction,
)
if client is None: # guarded above; retain a fail-closed boundary
raise ValueError(f"Model alias {alias!r} has no rerank endpoint")
replacement = _RerankRuntimeEntry(
runtime=RerankRuntime(client, alias=alias, model=cfg.model),
fingerprint=fingerprint,
instruction=instruction,
)
if entry is not None:
close = entry.runtime.begin_retirement()
if close is not None:
close_actions.append(close)
log.info("rerank.runtime_retired alias=%s reason=config", alias)
self._rerank_runtimes[alias] = replacement
entry = replacement
# The Reranker role selects one alias per process. Retire a
# previously selected alias when a settings change resolves its
# replacement; active calls drain on their old immutable lanes.
for other_alias, other in list(self._rerank_runtimes.items()):
if other_alias == alias:
continue
close = other.runtime.begin_retirement()
if close is not None:
close_actions.append(close)
del self._rerank_runtimes[other_alias]
log.info("rerank.runtime_retired alias=%s reason=role_change", other_alias)
lane = RerankLane(
runtime=entry.runtime,
alias=alias,
model=cfg.model,
admission=self._admissions[alias],
registry_generation=self._generation,
config_version=config_version,
)
return lane
def deactivate_rerank_runtime(self) -> None:
"""Retire any selected rerank runtime after the role becomes empty/invalid."""
close_actions: list[Callable[[], None]] = []
with _deferred_close_actions(close_actions), self._client_lock:
for alias, entry in list(self._rerank_runtimes.items()):
close = entry.runtime.begin_retirement()
if close is not None:
close_actions.append(close)
log.info("rerank.runtime_retired alias=%s reason=disabled", alias)
self._rerank_runtimes.clear()
def has_alias(self, alias: str) -> bool:
"""Check if *alias* exists in the registry."""
return alias in self._models
@@ -728,7 +860,8 @@ class ModelRegistry:
# refuses on this deployment — say so at every swap chokepoint.
warn_profile_mismatched_aliases(models, app_state)
_validate_registry_args(models, default, fallback, agent_model, task_model)
with self._client_lock:
rerank_close_actions: list[Callable[[], None]] = []
with _deferred_close_actions(rerank_close_actions), self._client_lock:
# FIRST write inside the lock, deliberately BEFORE the map swap
# and the client teardown. The per-send refresh reads the maps
# lock-free and samples the generation AFTER them (see
@@ -761,6 +894,23 @@ class ModelRegistry:
self._admissions[alias] = ModelAdmission(alias, limit)
else:
gate.set_limit(limit)
# Rerank runtimes have their own transport and breaker state. A
# cap-only or unrelated model reload preserves them; any relevant
# endpoint/model/auth/capability change retires the old runtime now.
for alias, entry in list(self._rerank_runtimes.items()):
rerank_cfg = self._models.get(alias)
if (
rerank_cfg is None
or not rerank_cfg.base_url.strip()
or not rerank_cfg.capabilities.get("supports_rerank")
or entry.fingerprint
!= _rerank_runtime_fingerprint(rerank_cfg, entry.instruction)
):
close = entry.runtime.begin_retirement()
if close is not None:
rerank_close_actions.append(close)
del self._rerank_runtimes[alias]
log.info("rerank.runtime_retired alias=%s reason=registry_reload", alias)
# Removed aliases remain as tombstones for this registry's
# lifetime. A stale lane may still hold or queue on that object;
# re-adding the alias must reconfigure the same gate rather than
@@ -800,7 +950,13 @@ class ModelRegistry:
def shutdown(self) -> None:
"""Close all cached client connections."""
with self._client_lock:
rerank_close_actions: list[Callable[[], None]] = []
with _deferred_close_actions(rerank_close_actions), self._client_lock:
for entry in self._rerank_runtimes.values():
close = entry.runtime.begin_retirement()
if close is not None:
rerank_close_actions.append(close)
self._rerank_runtimes.clear()
for client in self._clients.values():
if hasattr(client, "close"):
client.close()
+413 -4
View File
@@ -21,15 +21,22 @@ default and no fall back to a local model.
from __future__ import annotations
import contextlib
import math
import threading
import time
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Protocol
from typing import TYPE_CHECKING, Any, Protocol
import httpx
from turnstone.core.deadline import DeadlineCancelledError
from turnstone.core.log import get_logger
if TYPE_CHECKING:
from turnstone.core.admission import AdmissionLease, ModelAdmission
log = get_logger(__name__)
# A reranker reorders candidate documents by relevance to the query, returning
@@ -49,6 +56,14 @@ class RerankError(RuntimeError):
"""
class RerankCircuitOpenError(RerankError):
"""The endpoint circuit is open, so no rerank request was dispatched."""
class RerankRuntimeRetiredError(RerankError):
"""The resolved runtime retired before this batch could dispatch."""
def _sigmoid(x: float) -> float:
"""Numerically-stable logistic sigmoid (maps a logit to a probability)."""
if x >= 0.0:
@@ -89,12 +104,385 @@ class RerankClient(Protocol):
"""Minimal interface for a rerank backend."""
def rerank(
self, query: str, documents: list[str], *, top_n: int | None = None
self,
query: str,
documents: list[str],
*,
top_n: int | None = None,
timeout: float | None = None,
) -> list[RerankHit]:
"""Score ``documents`` against ``query``; return hits sorted best-first."""
...
@dataclass(frozen=True, slots=True)
class RerankCircuitSnapshot:
"""Non-sensitive instantaneous breaker state for diagnostics and tests."""
state: str
consecutive_failures: int
@dataclass(frozen=True, slots=True)
class RerankRuntimeSnapshot:
"""Non-sensitive instantaneous lifecycle state for diagnostics and tests."""
alias: str
model: str
active_calls: int
retired: bool
closed: bool
circuit: RerankCircuitSnapshot
@dataclass(frozen=True, slots=True)
class _CircuitPermit:
epoch: int
probe: bool
class _RerankCircuit:
"""Small thread-safe consecutive-failure circuit for one rerank runtime."""
def __init__(
self,
alias: str,
model: str,
*,
failure_threshold: int,
cooldown_seconds: float,
clock: Callable[[], float],
) -> None:
self._alias = alias
self._model = model
self._failure_threshold = failure_threshold
self._cooldown_seconds = cooldown_seconds
self._clock = clock
self._lock = threading.Lock()
self._state = "closed"
self._consecutive_failures = 0
self._opened_at: float | None = None
# Invalidates completions from calls admitted before a transition.
self._epoch = 0
def acquire(self) -> _CircuitPermit:
with self._lock:
if self._state == "closed":
return _CircuitPermit(self._epoch, False)
if self._state == "half_open":
raise RerankCircuitOpenError(
f"rerank circuit is half-open for alias {self._alias!r}"
)
opened_at = self._opened_at
if opened_at is None or self._clock() - opened_at < self._cooldown_seconds:
raise RerankCircuitOpenError(f"rerank circuit is open for alias {self._alias!r}")
self._state = "half_open"
log.info(
"rerank.circuit_half_open alias=%s model=%s",
self._alias,
self._model,
)
return _CircuitPermit(self._epoch, True)
def validate(self, permit: _CircuitPermit) -> None:
"""Refuse a permit invalidated while its caller waited for admission."""
with self._lock:
expected_state = "half_open" if permit.probe else "closed"
if permit.epoch != self._epoch or self._state != expected_state:
raise RerankCircuitOpenError(
f"rerank circuit changed while waiting for alias {self._alias!r}"
)
def succeed(self, permit: _CircuitPermit) -> None:
recovered = False
with self._lock:
if permit.epoch != self._epoch:
return
if permit.probe:
if self._state != "half_open":
return
self._state = "closed"
self._opened_at = None
self._consecutive_failures = 0
self._epoch += 1
recovered = True
elif self._state == "closed":
self._consecutive_failures = 0
if recovered:
log.info(
"rerank.circuit_recovered alias=%s model=%s",
self._alias,
self._model,
)
def fail(self, permit: _CircuitPermit) -> None:
opened = False
failures = 0
with self._lock:
if permit.epoch != self._epoch:
return
if permit.probe:
if self._state != "half_open":
return
self._state = "open"
self._opened_at = self._clock()
self._consecutive_failures = self._failure_threshold
self._epoch += 1
failures = self._consecutive_failures
opened = True
elif self._state == "closed":
self._consecutive_failures += 1
failures = self._consecutive_failures
if failures >= self._failure_threshold:
self._state = "open"
self._opened_at = self._clock()
self._epoch += 1
opened = True
if opened:
log.warning(
"rerank.circuit_open alias=%s model=%s failures=%d",
self._alias,
self._model,
failures,
)
def abandon(self, permit: _CircuitPermit) -> None:
"""Release a half-open reservation without judging endpoint health."""
with self._lock:
if permit.epoch == self._epoch and permit.probe and self._state == "half_open":
# Preserve the original opened_at. Its cooldown has already
# elapsed, so the next caller may become the replacement probe.
self._state = "open"
def snapshot(self) -> RerankCircuitSnapshot:
with self._lock:
return RerankCircuitSnapshot(self._state, self._consecutive_failures)
class _RerankRuntimeLease:
"""One idempotently releasable active-call hold on a runtime."""
__slots__ = ("_released", "_runtime")
def __init__(self, runtime: RerankRuntime) -> None:
self._runtime = runtime
self._released = False
def release(self) -> None:
if self._released:
return
self._released = True
self._runtime._release_call()
class RerankRuntime:
"""Lifecycle-owned backend, circuit, and active-call retirement state."""
def __init__(
self,
client: RerankClient,
*,
alias: str,
model: str,
failure_threshold: int = 3,
cooldown_seconds: float = 30.0,
clock: Callable[[], float] = time.monotonic,
) -> None:
if failure_threshold < 1:
raise ValueError("rerank circuit failure threshold must be positive")
if cooldown_seconds < 0:
raise ValueError("rerank circuit cooldown must be non-negative")
self._client = client
self.alias = alias
self.model = model
self._state_lock = threading.Lock()
self._active_calls = 0
self._retired = False
self._closed = False
self._circuit = _RerankCircuit(
alias,
model,
failure_threshold=failure_threshold,
cooldown_seconds=cooldown_seconds,
clock=clock,
)
@property
def client(self) -> RerankClient:
return self._client
def acquire_call(self) -> _RerankRuntimeLease:
with self._state_lock:
if self._retired or self._closed:
raise RerankRuntimeRetiredError(f"rerank runtime retired for alias {self.alias!r}")
self._active_calls += 1
return _RerankRuntimeLease(self)
def _release_call(self) -> None:
close_now = False
with self._state_lock:
if self._active_calls <= 0:
raise RuntimeError("rerank runtime lease released without an active call")
self._active_calls -= 1
if self._retired and self._active_calls == 0 and not self._closed:
self._closed = True
close_now = True
if close_now:
self._close_client()
def begin_retirement(self) -> Callable[[], None] | None:
"""Prevent new calls and return an idle close action for the caller.
The split lets ``ModelRegistry`` mark the runtime while holding its
short-lived registry lock, then execute a potentially blocking client
close after releasing that lock. Active calls close from their final
lease release instead.
"""
with self._state_lock:
self._retired = True
if self._active_calls or self._closed:
return None
self._closed = True
return self._close_client
def retire(self) -> None:
close = self.begin_retirement()
if close is not None:
close()
def _close_client(self) -> None:
close = getattr(self._client, "close", None)
if not callable(close):
return
try:
close()
except Exception:
log.warning(
"rerank.runtime_close_failed alias=%s model=%s",
self.alias,
self.model,
exc_info=True,
)
def circuit_acquire(self) -> _CircuitPermit:
return self._circuit.acquire()
def circuit_succeed(self, permit: _CircuitPermit) -> None:
self._circuit.succeed(permit)
def circuit_validate(self, permit: _CircuitPermit) -> None:
self._circuit.validate(permit)
def circuit_fail(self, permit: _CircuitPermit) -> None:
self._circuit.fail(permit)
def circuit_abandon(self, permit: _CircuitPermit) -> None:
self._circuit.abandon(permit)
def snapshot(self) -> RerankRuntimeSnapshot:
with self._state_lock:
active_calls = self._active_calls
retired = self._retired
closed = self._closed
return RerankRuntimeSnapshot(
alias=self.alias,
model=self.model,
active_calls=active_calls,
retired=retired,
closed=closed,
circuit=self._circuit.snapshot(),
)
@dataclass(frozen=True, slots=True)
class RerankLane:
"""One immutable, per-batch binding to a shared rerank runtime."""
runtime: RerankRuntime
alias: str
model: str
admission: ModelAdmission
registry_generation: int
config_version: int = 0
def _raise_if_aborted(cancel_ref: Any) -> None:
if bool(getattr(cancel_ref, "aborted", False)):
raise DeadlineCancelledError("rerank cancelled")
def rerank(
lane: RerankLane,
query: str,
documents: list[str],
*,
top_n: int | None = None,
timeout: float,
cancel_ref: Any = None,
) -> list[RerankHit]:
"""Dispatch one rerank batch through circuit, admission, and runtime leases."""
if not documents:
return []
_raise_if_aborted(cancel_ref)
permit = lane.runtime.circuit_acquire()
admission: AdmissionLease | None = None
runtime_lease: _RerankRuntimeLease | None = None
dispatched = False
try:
admission = lane.admission.acquire(cancel_ref=cancel_ref)
# The circuit can open while this call is queued behind either rerank
# or model work on the shared alias gate. Do not let an old closed-state
# permit leak one more full-timeout request after that transition.
lane.runtime.circuit_validate(permit)
runtime_lease = lane.runtime.acquire_call()
_raise_if_aborted(cancel_ref)
mark_dispatch = getattr(cancel_ref, "mark_dispatch", None)
if callable(mark_dispatch):
with contextlib.suppress(Exception):
mark_dispatch()
dispatched = True
hits = lane.runtime.client.rerank(
query,
documents,
top_n=top_n,
timeout=timeout,
)
_raise_if_aborted(cancel_ref)
if not hits:
raise RerankError("rerank endpoint returned no scores for non-empty input")
except (DeadlineCancelledError, RerankCircuitOpenError, RerankRuntimeRetiredError):
lane.runtime.circuit_abandon(permit)
raise
except Exception:
try:
# Stop/supersession owns the outcome even when the dispatched
# transport fails while cancellation is landing. Do not charge a
# cancelled request to endpoint health or let retrieval fallback
# absorb it as an ordinary rerank outage.
_raise_if_aborted(cancel_ref)
except DeadlineCancelledError:
lane.runtime.circuit_abandon(permit)
raise
if dispatched:
lane.runtime.circuit_fail(permit)
else:
# Admission and lifecycle failures say nothing about endpoint
# health; only a call that reached the backend can trip its circuit.
lane.runtime.circuit_abandon(permit)
raise
except BaseException:
lane.runtime.circuit_abandon(permit)
raise
else:
lane.runtime.circuit_succeed(permit)
return hits
finally:
if runtime_lease is not None:
runtime_lease.release()
if admission is not None:
admission.release()
class CohereJinaRerankClient:
"""Rerank via a Cohere/Jina-compatible ``POST <url>`` endpoint.
@@ -117,9 +505,17 @@ class CohereJinaRerankClient:
self._api_key = api_key
self._timeout = timeout
self._instruction = instruction
self._client = httpx.Client()
self._close_lock = threading.Lock()
self._closed = False
def rerank(
self, query: str, documents: list[str], *, top_n: int | None = None
self,
query: str,
documents: list[str],
*,
top_n: int | None = None,
timeout: float | None = None,
) -> list[RerankHit]:
if not documents:
return []
@@ -135,10 +531,23 @@ class CohereJinaRerankClient:
if top_n is not None:
payload["top_n"] = top_n
headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
resp = httpx.post(self._url, json=payload, headers=headers, timeout=self._timeout)
resp = self._client.post(
self._url,
json=payload,
headers=headers,
timeout=self._timeout if timeout is None else timeout,
)
resp.raise_for_status()
return _parse_hits(resp.json(), len(documents))
def close(self) -> None:
"""Close the owned HTTP connection pool exactly once."""
with self._close_lock:
if self._closed:
return
self._closed = True
self._client.close()
def _parse_hits(data: Any, n_docs: int) -> list[RerankHit]:
"""Parse a Cohere/Jina ``{"results": [...]}`` or bare-list rerank response.
+6 -1
View File
@@ -298,7 +298,12 @@ def calibrate_model(
)
if client is None:
raise ValueError("no rerank endpoint (base_url is empty)")
return calibrate(client, model=model or base_url)
try:
return calibrate(client, model=model or base_url)
finally:
close = getattr(client, "close", None)
if callable(close):
close()
def _raw_scale(raw: list[float]) -> str:
+62 -41
View File
@@ -1,63 +1,84 @@
"""Resolve the runtime rerank client from configuration.
"""Resolve an immutable runtime rerank lane from configuration.
Sole caller is ``ChatSession._resolve_rerank_client``: the reranker is the model
Sole caller is ``ChatSession._resolve_rerank_lane``: the reranker is the model
definition (capability ``supports_rerank``) selected via the Reranker role
(``tools.reranker_alias``); there is no global endpoint fallback. The calibrate
CLI (``admin.py``) and the calibrate endpoint (``_global_rerank_instruction`` in
console/server.py) resolve the endpoint themselves via ``calibrate_model`` and
only share the instruction precedence below, not this function.
(``tools.reranker_alias``); there is no global endpoint fallback. Resolution
creates a fresh frozen binding around a registry-owned shared runtime, never a
session-owned client. Calibration remains an isolated one-shot path.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, cast
if TYPE_CHECKING:
from turnstone.core.rerank import RerankClient
from turnstone.core.rerank import RerankLane
def resolve_rerank_client_from(
def _effective_snapshot(config_store: Any) -> tuple[int, dict[str, Any]]:
snapshot = getattr(config_store, "effective_snapshot", None)
if callable(snapshot):
version, settings = snapshot()
return int(version), dict(settings)
# Compatibility for lightweight embedders/tests implementing only get().
return int(getattr(config_store, "version", 0)), {
"tools.reranker_alias": config_store.get("tools.reranker_alias"),
"tools.rerank_instruction": config_store.get("tools.rerank_instruction"),
}
def _deactivate(registry: Any) -> None:
deactivate = getattr(registry, "deactivate_rerank_runtime", None)
if callable(deactivate):
deactivate()
def resolve_rerank_lane_from(
config_store: Any | None,
registry: Any | None,
*,
timeout: float,
) -> RerankClient | None:
"""Return a rerank client from the selected Reranker model, or ``None``.
max_attempts: int = 3,
) -> RerankLane | None:
"""Return a coherent lane for the selected Reranker model, or ``None``.
The reranker is the model definition (capability ``supports_rerank``) picked
via the Reranker role (``tools.reranker_alias``) managed like every other
model, its ``base_url`` the full Cohere/Jina-compatible /rerank endpoint.
Returns ``None`` until such a model is selected (there is no bundled rerank
endpoint and no global URL fallback).
model, its ``base_url`` the full Cohere/Jina-compatible /rerank endpoint. A
ConfigStore snapshot and registry generation witness keep multi-key settings
and model binding coherent across concurrent reloads.
"""
if config_store is None or registry is None:
return None
cs = config_store
alias = str(cs.get("tools.reranker_alias") or "").strip()
if not alias:
return None
try:
cfg = registry.get_config(alias)
except Exception:
cfg = None
if cfg is None or not cfg.base_url or not cfg.capabilities.get("supports_rerank"):
return None
from turnstone.core.config import get_rerank_instruction
from turnstone.core.rerank import resolve_rerank_client
# The query instruction is a global task knob (instruction-aware rerankers
# like Qwen3); it applies to whichever reranker model is active. Stored admin
# value wins, else config.toml/env -- the same precedence the calibrate CLI
# (admin.py) and the calibrate endpoint (_global_rerank_instruction) use, so
# the instruction used at calibration time matches the one used at runtime.
instruction = str(cs.get("tools.rerank_instruction") or "").strip() or get_rerank_instruction()
return resolve_rerank_client(
url=cfg.base_url,
model=cfg.model or "",
api_key=cfg.api_key,
timeout=timeout,
instruction=instruction,
)
for _attempt in range(max(1, max_attempts)):
version, settings = _effective_snapshot(config_store)
alias = str(settings.get("tools.reranker_alias") or "").strip()
if not alias:
_deactivate(registry)
return None
instruction = (
str(settings.get("tools.rerank_instruction") or "").strip() or get_rerank_instruction()
)
try:
lane = cast(
"RerankLane",
registry.resolve_rerank_lane(
alias,
instruction=instruction,
config_version=version,
),
)
except (KeyError, ValueError):
_deactivate(registry)
return None
if (
int(getattr(config_store, "version", version)) == version
and int(getattr(registry, "generation", lane.registry_generation))
== lane.registry_generation
):
return lane
# A continuously changing configuration is not a safe binding. Retrieval
# callers preserve native order and retry resolution on their next use.
return None
+136 -42
View File
@@ -278,7 +278,7 @@ if TYPE_CHECKING:
)
from turnstone.core.output_guard import OutputAssessment
from turnstone.core.output_guard_judge import OutputGuardJudge, OutputJudgeVerdict
from turnstone.core.rerank import RerankClient, Reranker
from turnstone.core.rerank import Reranker, RerankHit, RerankLane
from turnstone.core.storage import ForkCloneSnapshot
from turnstone.core.web_search import WebSearchClient
@@ -3505,28 +3505,80 @@ class ChatSession:
timeout=self.tool_timeout,
)
def _resolve_rerank_client(self) -> RerankClient | None:
"""Return a rerank client, or None when reranking is unconfigured.
def _resolve_rerank_lane(self) -> RerankLane | None:
"""Return a fresh rerank lane, or None when reranking is unconfigured.
The reranker is a **model definition** (capability ``supports_rerank``)
selected via the Reranker role (``tools.reranker_alias``); its base_url
is the full /rerank endpoint. There is no bundled rerank endpoint and no
global URL fallback, so reranking stays disabled until such a model is
selected.
global URL fallback. The registry owns the shared runtime; callers must
resolve per batch rather than retain this immutable binding.
"""
from turnstone.core.rerank_config import resolve_rerank_client_from
from turnstone.core.rerank_config import resolve_rerank_lane_from
return resolve_rerank_client_from(
return resolve_rerank_lane_from(
getattr(self, "_config_store", None),
getattr(self, "_registry", None),
timeout=min(self.tool_timeout, _RERANK_TIMEOUT_CAP_S),
)
def _rerank_hits(
self,
query: str,
documents: list[str],
*,
origin_generation: int = 0,
) -> list[RerankHit] | None:
"""Resolve and dispatch one batch, preserving its cancellation owner.
``None`` means no live reranker is configured; callers return identity
order so long-lived BM25/tool-search closures can observe a later role
selection without being rebuilt. Endpoint and circuit failures raise
through the existing retrieval fallback seams.
"""
lane = self._resolve_rerank_lane()
if lane is None:
return None
from turnstone.core.deadline import DeadlineCancelledError
from turnstone.core.rerank import rerank
task_scope = _active_task_agent_cancel_scope.get()
owner_generation = (
origin_generation
or _active_tool_origin_generation.get()
or _active_commit_origin_generation.get()
)
if task_scope is not None:
cancel_ref: Any = task_scope.cancel_ref
elif hasattr(self, "_cancel_event"):
cancel_ref = _CancelRef(self, owner_generation)
else: # lightweight direct seam tests/embedders
cancel_ref = None
try:
return rerank(
lane,
query,
documents,
timeout=min(float(getattr(self, "tool_timeout", 30.0)), _RERANK_TIMEOUT_CAP_S),
cancel_ref=cancel_ref,
)
except DeadlineCancelledError:
# Retrieval fallbacks catch Exception. Translate the deadline seam
# back into GenerationCancelled (a BaseException) before they can
# mistake Stop/supersession for an endpoint outage.
if task_scope is not None:
task_scope.check()
check_cancelled = getattr(self, "_check_cancelled", None)
if callable(check_cancelled):
check_cancelled(owner_generation)
raise
def _rerank_enabled_for(self, tool: str) -> bool:
"""Whether reranking is enabled for ``tool`` (currently: 'web_search', 'bm25').
The per-tool toggles default on; the operative gate is whether an
endpoint is configured (``_resolve_rerank_client`` returns None when
endpoint is configured (``_resolve_rerank_lane`` returns None when
not). A bare CLI without a ConfigStore inherits the on-by-default toggle.
"""
cs = getattr(self, "_config_store", None)
@@ -3534,25 +3586,30 @@ class ChatSession:
return bool(cs.get(f"tools.rerank_{tool}"))
return True
def _web_search_reranker(self) -> Reranker | None:
"""Build a web_search reranker callable, or None when disabled.
def _web_search_reranker(self) -> Reranker:
"""Build a live-resolving web_search reranker callable.
Returns a ``(query, docs) -> ranked indices`` adapter over the configured
rerank endpoint; None when reranking is off or no endpoint is set.
The closure captures no concrete lane/runtime. Config changes are read
at each invocation; disabled/unconfigured state is identity order.
"""
if not self._rerank_enabled_for("web_search"):
return None
rc = self._resolve_rerank_client()
if rc is None:
return None
def _rank(query: str, docs: list[str]) -> list[int]:
return [hit.index for hit in rc.rerank(query, docs)]
if not self._rerank_enabled_for("web_search"):
return list(range(len(docs)))
hits = self._rerank_hits(query, docs)
if hits is None:
return list(range(len(docs)))
return [hit.index for hit in hits]
return _rank
def _bm25_reranker(self, threshold: float = 0.0) -> Reranker | None:
"""Build a BM25 reranker callable, or None when disabled.
def _bm25_reranker(
self,
threshold: float = 0.0,
*,
origin_generation: int = 0,
) -> Reranker:
"""Build a live-resolving BM25 reranker callable.
Mirrors ``_web_search_reranker``. ``threshold`` is a relevance FLOOR
applied in this closure (where scores still exist); the BM25Index seam
@@ -3563,24 +3620,22 @@ class ChatSession:
threshold. Only memory-pointer relevance filtering passes a configured threshold;
reactive surfaces pass 0.
"""
if not self._rerank_enabled_for("bm25"):
return None
rc = self._resolve_rerank_client()
if rc is None:
return None
from turnstone.core.rerank import RerankError, normalize_scores
def _rank(query: str, docs: list[str]) -> list[int]:
hits = rc.rerank(query, docs)
if not self._rerank_enabled_for("bm25"):
return list(range(len(docs)))
hits = self._rerank_hits(
query,
docs,
origin_generation=origin_generation,
)
if hits is None:
return list(range(len(docs)))
if docs and not hits:
# A conforming reranker scores every document; an empty result
# for non-empty input means the endpoint response was
# unparseable (rc.rerank -> _parse_hits returns [] without
# raising). That is an endpoint FAILURE -- a discrete branch
# from the relevance floor below -- so raise and let BM25Index
# fall back to BM25 order in BOTH modes, instead of the
# filter-mode floor honoring it as "nothing relevant".
# The managed dispatcher enforces this too. Keep the adapter
# fail-closed for alternate protocol implementations and
# lightweight test embedders that provide the seam directly.
raise RerankError("rerank endpoint returned no scores for non-empty input")
# Normalise to a 0-1 relevance space (sigmoid for logit endpoints) so
# ``threshold`` means the same on every reranker. Sigmoid is monotonic
@@ -11204,7 +11259,13 @@ class ChatSession:
# below rejects every old mutation. The participant lookup is cached so
# the in-lock apply path performs no storage access.
planned_user_turn = (
self._plan_genuine_user_turn(user_input, turn_principal_id) if not from_wake else None
self._plan_genuine_user_turn(
user_input,
turn_principal_id,
origin_generation=my_generation,
)
if not from_wake
else None
)
shared_state_plan = self._plan_shared_state()
participant_name: str | None = None
@@ -11885,7 +11946,9 @@ class ChatSession:
# must not let this retired worker pop a successor's queue
# or repaint its state as idle.
flushed_out: list[bool] = []
queued_flush = self._prepare_queued_flush()
queued_flush = self._prepare_queued_flush(
origin_generation=my_generation,
)
def _drain_or_idle(
durable: list[Callable[[], None]],
@@ -12396,6 +12459,7 @@ class ChatSession:
queued_flush = self._prepare_queued_flush(
user_feedback or "",
include_queue=False,
origin_generation=my_generation,
)
fold_tool_batch = functools.partial(
_fold_tool_batch,
@@ -12468,7 +12532,10 @@ class ChatSession:
)
raise
except GenerationCancelled:
queued_flush = self._prepare_queued_flush()
queued_flush = self._prepare_queued_flush(
origin_generation=my_generation,
allow_rerank=False,
)
def _finalize_cancelled_generation(
durable: list[Callable[[], None]],
@@ -12575,7 +12642,10 @@ class ChatSession:
# Do NOT re-raise — return normally so server worker thread
# completes cleanly.
except KeyboardInterrupt as exc:
queued_flush = self._prepare_queued_flush()
queued_flush = self._prepare_queued_flush(
origin_generation=my_generation,
allow_rerank=False,
)
def _finalize_interrupted_generation(
error: BaseException,
@@ -12601,7 +12671,10 @@ class ChatSession:
)
raise
except Exception as exc:
queued_flush = self._prepare_queued_flush()
queued_flush = self._prepare_queued_flush(
origin_generation=my_generation,
allow_rerank=False,
)
# Orphan gate: a superseded thread's stream death can escape
# cancel conversion (the successor replaced the cancel event
@@ -17099,6 +17172,8 @@ class ChatSession:
prefix: str = "",
*,
include_queue: bool = True,
origin_generation: int = 0,
allow_rerank: bool = True,
) -> _QueuedFlushPlan:
"""Snapshot and rank one prospective queued USER turn outside locks."""
worker_claim = current_worker_claim(self)
@@ -17133,7 +17208,12 @@ class ChatSession:
return _QueuedFlushPlan(
prefix=prefix,
items=items,
turn=self._plan_genuine_user_turn(content, principal),
turn=self._plan_genuine_user_turn(
content,
principal,
origin_generation=origin_generation,
allow_rerank=allow_rerank,
),
)
def _apply_queued_flush(
@@ -19456,6 +19536,8 @@ class ChatSession:
user_message: str,
*,
access: _MemoryAccess,
origin_generation: int = 0,
allow_rerank: bool = True,
) -> str:
"""Build a live metadata-only pointer for one genuine user turn."""
if not user_message.strip() or not self._nudges_enabled("memory_pointer"):
@@ -19470,7 +19552,14 @@ class ChatSession:
rows,
user_message,
k=self._mem_cfg.relevance_k,
reranker=self._bm25_reranker(threshold),
reranker=(
self._bm25_reranker(
threshold,
origin_generation=origin_generation,
)
if allow_rerank
else None
),
rerank_filters=threshold > 0,
)
except Exception:
@@ -19482,6 +19571,9 @@ class ChatSession:
self,
content: str,
principal_id: str,
*,
origin_generation: int = 0,
allow_rerank: bool = True,
) -> _GenuineUserTurnPlan:
"""Plan all metadata derived from an admitted human turn."""
principal = principal_id.strip()
@@ -19491,6 +19583,8 @@ class ChatSession:
memory_pointer=self._plan_memory_pointer(
content,
access=self._memory_access(principal),
origin_generation=origin_generation,
allow_rerank=allow_rerank,
),
)