mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
Add per-alias model concurrency admission (#990)
* feat(models): add per-alias concurrency admission Add registry-backed FIFO admission limits with queue-aware deadlines and full-stream leases. Expose max_concurrency through storage, admin configuration, OpenAPI, documentation, and diagrams, with role and live backend count coverage. * fix(api): omit null concurrency schema default Keep max_concurrency optional for presence-keyed updates without advertising a null default for its non-null integer OpenAPI shape.
This commit is contained in:
+18
-2
@@ -920,6 +920,7 @@ configurations so workstreams can use different LLM backends.
|
||||
[models.local]
|
||||
base_url = "http://localhost:8000/v1"
|
||||
model = "qwen3-32b"
|
||||
max_concurrency = 1
|
||||
# provider defaults to "openai"
|
||||
|
||||
[models.claude]
|
||||
@@ -958,6 +959,20 @@ binding on mismatch. Failed reload validation changes neither maps nor
|
||||
generation. If only non-transport fields changed, compatible client pools can
|
||||
remain warm, but the session still receives a new coherent config/lane.
|
||||
|
||||
**Per-alias admission:** Each registry alias owns a stable, hot-resizable
|
||||
`ModelAdmission`. `max_concurrency = 0` is unlimited; a positive value limits
|
||||
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. `model_turn()` materializes attachment fallbacks before admission,
|
||||
then holds one lease across eager stream creation and the complete drain,
|
||||
releasing before retry backoff. Admission 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.
|
||||
Fallback is a lane change, so retry classification and result provenance come
|
||||
@@ -1196,8 +1211,9 @@ Verified quirks of vLLM's Anthropic endpoint:
|
||||
|
||||
**Database model definitions:** On server entry points, models can also be
|
||||
defined in the `model_definitions` table (admin Models tab). DB models support
|
||||
the same per-model sampling overrides. Config.toml models override DB models
|
||||
with the same alias in-memory (the DB rows are never modified).
|
||||
the same per-model sampling overrides and per-alias `max_concurrency`.
|
||||
Config.toml models override DB models with the same alias in-memory (the DB
|
||||
rows are never modified).
|
||||
|
||||
**Lifecycle:**
|
||||
1. `load_model_registry()` loads DB model definitions (if storage available),
|
||||
|
||||
+4
-1
@@ -455,7 +455,10 @@ configuration editor.
|
||||
The **Channels** tab links users to either a Discord or Slack account
|
||||
via a per-row channel-type selector. The **Models** tab is a CRUD
|
||||
editor for `model_definitions`, including static and dynamic backend-auth
|
||||
modes. Model edits rebind existing workstreams at their next send while
|
||||
modes and a per-process **Max concurrent generations** limit for each alias
|
||||
(`0` means unlimited). The limit is shared by every model-backed role using
|
||||
that alias and a streaming generation holds its slot through the full decode.
|
||||
Model edits rebind existing workstreams at their next send while
|
||||
in-flight requests keep their original definition snapshot; see
|
||||
[Settings](settings.md#model-definition-reloads) for the full contract. The **Nodes** tab edits per-node
|
||||
metadata, and the **TLS** tab manages CA and leaf certificates for the
|
||||
|
||||
@@ -134,6 +134,7 @@ class "ModelLane" as ModelLane <<frozen>> {
|
||||
+ capabilities: ModelCapabilities
|
||||
+ extra_params: dict | None
|
||||
+ registry: ModelRegistry | None
|
||||
+ admission: ModelAdmission | None
|
||||
+ backend_auth_config: ModelConfig | None
|
||||
+ backend_auth_resolver: Callable | None
|
||||
}
|
||||
@@ -357,12 +358,13 @@ class "ModelRegistry" as ModelReg {
|
||||
- _models: dict[str, ModelConfig]
|
||||
- _clients: dict[str, Any]
|
||||
- _providers: dict[str, LLMProvider]
|
||||
- _admissions: dict[str, ModelAdmission]
|
||||
- _client_lock: Lock
|
||||
+ default: str
|
||||
+ fallback: list[str]
|
||||
+ agent_model: str | None
|
||||
--
|
||||
+ resolve_binding(alias) → (client, model, config, provider, generation)
|
||||
+ resolve_binding(alias) → (client, model, config, provider, admission, generation)
|
||||
+ get_client(alias) → Any
|
||||
+ get_provider(alias) → LLMProvider
|
||||
+ has_alias(alias) → bool
|
||||
@@ -376,6 +378,22 @@ class "ModelRegistry" as ModelReg {
|
||||
core/model_registry.py
|
||||
}
|
||||
|
||||
class "ModelAdmission" as ModelAdmission {
|
||||
- alias: str
|
||||
- _limit: int
|
||||
- _in_flight: int
|
||||
- _waiters: deque
|
||||
+ acquire(cancel_ref) → AdmissionLease
|
||||
+ set_limit(limit)
|
||||
+ snapshot() → AdmissionSnapshot
|
||||
--
|
||||
Per-process FIFO generation gate.
|
||||
Stable across alias hot reloads;
|
||||
queue time is deadline credit.
|
||||
--
|
||||
core/admission.py
|
||||
}
|
||||
|
||||
class "ModelConfig" as ModelCfg <<frozen>> {
|
||||
+ alias: str
|
||||
+ provider: str
|
||||
@@ -385,6 +403,7 @@ class "ModelConfig" as ModelCfg <<frozen>> {
|
||||
+ temperature: float | None
|
||||
+ max_tokens: int | None
|
||||
+ reasoning_effort: str | None
|
||||
+ max_concurrency: int
|
||||
+ auth_mode: str
|
||||
+ obo_audience: str
|
||||
+ obo_scopes: str
|
||||
@@ -475,12 +494,14 @@ KindAdapter ..> ChatSession : constructs
|
||||
|
||||
ModelReg --> "*" ModelCfg : holds
|
||||
ModelReg --> "*" LLMProvider : caches
|
||||
ModelReg --> "*" ModelAdmission : owns per alias
|
||||
LLMProvider --> ModelCaps : returns
|
||||
ModelReg --> ResolvedBinding : resolves atomically
|
||||
ResolvedBinding --> ModelLane
|
||||
ModelLane --> LLMProvider
|
||||
ModelLane --> ModelCaps
|
||||
ModelLane --> ModelCfg : auth/config snapshot
|
||||
ModelLane --> ModelAdmission : admission lease
|
||||
ModelTurnFn --> ModelLane
|
||||
ModelTurnFn --> ModelTurnResult
|
||||
ModelTurnFn ..> BackendAuth : per-call resolver
|
||||
|
||||
@@ -10,6 +10,7 @@ participant "SessionManager" as Manager
|
||||
participant "ChatSession" as Session
|
||||
participant "SessionUIBase" as UI
|
||||
participant "model_turn()\n+ lowering" as Plant
|
||||
participant "ModelAdmission\n(per alias)" as Admission
|
||||
participant "LLM provider" as Provider
|
||||
participant "Tool workers" as Tools
|
||||
database "StorageBackend\n(SQLite / PostgreSQL)" as Storage
|
||||
@@ -54,6 +55,9 @@ loop until final answer and no queued input
|
||||
Session -> Plant : model_turn(active ModelLane, Turns,\n tools, cancel_ref, on_chunk)
|
||||
activate Plant
|
||||
Plant -> Plant : canonical Turns → provider wire\nrestore ids + repair + lane-specific fold
|
||||
Plant -> Plant : materialize attachment refs\n(nested perception before outer slot)
|
||||
Plant -> Admission : acquire(cancel_ref)
|
||||
activate Admission
|
||||
Plant -> Plant : resolve per-call backend credential\nfrom lane's pinned ModelConfig
|
||||
Plant -> Provider : create_streaming(...)
|
||||
activate Provider
|
||||
@@ -68,6 +72,8 @@ loop until final answer and no queued input
|
||||
Provider --> Plant : finish + usage + native blocks
|
||||
deactivate Provider
|
||||
Plant -> Plant : drain + re-ingest assistant Turn\nwith serving-lane provenance
|
||||
Plant -> Admission : release before retry backoff
|
||||
deactivate Admission
|
||||
Plant --> Session : ModelTurnResult
|
||||
deactivate Plant
|
||||
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:e1431edf3891785b922c52b7897e3af5d39ba9973a815f892d6afdb762c5297b
|
||||
size 612662
|
||||
oid sha256:c74e99c530c3a8af9ab35b1e4d8c4fef0ea35c0c04cc35da7cf3588e71382057
|
||||
size 661175
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:79299b25ccc10484af13684a89ed9457abb9bc1781604bf6a9000ccb84362e55
|
||||
size 311107
|
||||
oid sha256:b237e00c28225c847b3e1a083c5c35535d1bf1549f6e6c7e2b6c4fd0e037a980
|
||||
size 331906
|
||||
|
||||
@@ -54,6 +54,23 @@ When a per-model override is `NULL` (empty in the UI), the global default is
|
||||
used. Switching models via `/model <alias>` re-resolves sampling parameters
|
||||
from the new model's overrides or global defaults.
|
||||
|
||||
### Per-model concurrency
|
||||
|
||||
Each model definition may set `max_concurrency` to limit simultaneous model
|
||||
generations for that alias in one Turnstone process. `0` or an omitted value
|
||||
means unlimited. The gate is shared by every role using the alias—interactive
|
||||
turns, coordinators, task agents, judges, output guards, perception, compaction,
|
||||
and title generation—and a streaming generation holds its slot until the
|
||||
stream is fully drained or closed.
|
||||
|
||||
Admission is strictly per alias. Two aliases remain independent even when they
|
||||
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.
|
||||
|
||||
### Model backend authentication
|
||||
|
||||
Model definitions support four backend credential modes:
|
||||
@@ -445,6 +462,9 @@ send. Endpoint, provider, backend model ID, capabilities, extra parameters, and
|
||||
backend-auth configuration are replaced as one immutable binding. In-flight
|
||||
turns, judges, and task agents finish or cancel against the binding they
|
||||
started with; an admin edit never tears one request across two definitions.
|
||||
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.
|
||||
|
||||
Sampling and other saved workstream configuration remain workstream state. A
|
||||
model-definition edit does not silently rewrite a live workstream's chosen
|
||||
|
||||
@@ -11880,6 +11880,14 @@
|
||||
"title": "Context Window",
|
||||
"type": "integer"
|
||||
},
|
||||
"max_concurrency": {
|
||||
"default": 0,
|
||||
"description": "Maximum concurrent model generations for this alias in one process; zero means unlimited.",
|
||||
"maximum": 2147483647,
|
||||
"minimum": 0,
|
||||
"title": "Max Concurrency",
|
||||
"type": "integer"
|
||||
},
|
||||
"capabilities": {
|
||||
"default": "{}",
|
||||
"title": "Capabilities",
|
||||
@@ -12015,6 +12023,14 @@
|
||||
"title": "Context Window",
|
||||
"type": "integer"
|
||||
},
|
||||
"max_concurrency": {
|
||||
"default": 0,
|
||||
"description": "Maximum concurrent model generations for this alias in one process; zero means unlimited.",
|
||||
"maximum": 2147483647,
|
||||
"minimum": 0,
|
||||
"title": "Max Concurrency",
|
||||
"type": "integer"
|
||||
},
|
||||
"capabilities": {
|
||||
"default": "{}",
|
||||
"title": "Capabilities",
|
||||
@@ -12151,6 +12167,14 @@
|
||||
"title": "Context Window",
|
||||
"type": "integer"
|
||||
},
|
||||
"max_concurrency": {
|
||||
"default": 0,
|
||||
"description": "Maximum concurrent model generations for this alias in one process; zero means unlimited.",
|
||||
"maximum": 2147483647,
|
||||
"minimum": 0,
|
||||
"title": "Max Concurrency",
|
||||
"type": "integer"
|
||||
},
|
||||
"capabilities": {
|
||||
"additionalProperties": true,
|
||||
"title": "Capabilities",
|
||||
@@ -12304,6 +12328,13 @@
|
||||
"default": null,
|
||||
"title": "Context Window"
|
||||
},
|
||||
"max_concurrency": {
|
||||
"description": "Maximum concurrent model generations for this alias in one process; zero means unlimited.",
|
||||
"maximum": 2147483647,
|
||||
"minimum": 0,
|
||||
"title": "Max Concurrency",
|
||||
"type": "integer"
|
||||
},
|
||||
"capabilities": {
|
||||
"additionalProperties": true,
|
||||
"default": null,
|
||||
|
||||
@@ -92,6 +92,7 @@ def _seed_model_def(
|
||||
obo_audience: str = "",
|
||||
obo_scopes: str = "",
|
||||
capabilities: str = "{}",
|
||||
max_concurrency: int = 0,
|
||||
) -> None:
|
||||
"""Insert a model definition row directly via the storage API."""
|
||||
storage.create_model_definition(
|
||||
@@ -108,6 +109,7 @@ def _seed_model_def(
|
||||
auth_mode=auth_mode,
|
||||
obo_audience=obo_audience,
|
||||
obo_scopes=obo_scopes,
|
||||
max_concurrency=max_concurrency,
|
||||
)
|
||||
|
||||
|
||||
@@ -2431,6 +2433,91 @@ def test_update_endpoint_refreshes_registry(storage: SQLiteBackend) -> None:
|
||||
assert registry.get_config("local").model == "new-model"
|
||||
|
||||
|
||||
def test_create_and_update_max_concurrency_refresh_registry(storage: SQLiteBackend) -> None:
|
||||
_seed_model_def(storage, definition_id="m1", alias="local", model="m")
|
||||
registry = _make_registry(alias="local", model="m")
|
||||
client = _make_client(storage, registry)
|
||||
|
||||
created = client.post(
|
||||
"/v1/api/admin/model-definitions",
|
||||
json={"alias": "limited", "model": "m2", "max_concurrency": 3},
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
assert created.json()["max_concurrency"] == 3
|
||||
assert registry.get_config("limited").max_concurrency == 3
|
||||
|
||||
definition_id = created.json()["definition_id"]
|
||||
updated = client.put(
|
||||
f"/v1/api/admin/model-definitions/{definition_id}",
|
||||
json={"max_concurrency": 0},
|
||||
)
|
||||
assert updated.status_code == 200, updated.text
|
||||
assert updated.json()["max_concurrency"] == 0
|
||||
assert registry.get_config("limited").max_concurrency == 0
|
||||
|
||||
|
||||
def test_update_omission_preserves_max_concurrency(storage: SQLiteBackend) -> None:
|
||||
_seed_model_def(
|
||||
storage,
|
||||
definition_id="m1",
|
||||
alias="local",
|
||||
model="m",
|
||||
max_concurrency=2,
|
||||
)
|
||||
registry = _make_registry(alias="local", model="m")
|
||||
client = _make_client(storage, registry)
|
||||
|
||||
resp = client.put(
|
||||
"/v1/api/admin/model-definitions/m1",
|
||||
json={"model": "m2"},
|
||||
)
|
||||
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert resp.json()["max_concurrency"] == 2
|
||||
assert storage.get_model_definition("m1")["max_concurrency"] == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("invalid", [None, True, "1", 1.0, -1, 2_147_483_648])
|
||||
def test_create_rejects_invalid_max_concurrency(
|
||||
storage: SQLiteBackend,
|
||||
invalid: Any,
|
||||
) -> None:
|
||||
client = _make_client(storage, _make_registry())
|
||||
|
||||
resp = client.post(
|
||||
"/v1/api/admin/model-definitions",
|
||||
json={"alias": "invalid", "model": "m", "max_concurrency": invalid},
|
||||
)
|
||||
|
||||
assert resp.status_code == 400, resp.text
|
||||
assert "max_concurrency" in resp.json()["error"]
|
||||
assert storage.get_model_definition_by_alias("invalid") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("invalid", [None, True, "1", 1.0, -1, 2_147_483_648])
|
||||
def test_update_rejects_invalid_max_concurrency(
|
||||
storage: SQLiteBackend,
|
||||
invalid: Any,
|
||||
) -> None:
|
||||
_seed_model_def(
|
||||
storage,
|
||||
definition_id="m1",
|
||||
alias="local",
|
||||
model="m",
|
||||
max_concurrency=2,
|
||||
)
|
||||
client = _make_client(storage, _make_registry(alias="local", model="m"))
|
||||
|
||||
resp = client.put(
|
||||
"/v1/api/admin/model-definitions/m1",
|
||||
json={"max_concurrency": invalid},
|
||||
)
|
||||
|
||||
assert resp.status_code == 400, resp.text
|
||||
assert "max_concurrency" in resp.json()["error"]
|
||||
assert storage.get_model_definition("m1")["max_concurrency"] == 2
|
||||
|
||||
|
||||
def test_update_endpoint_skips_refresh_on_empty_body(
|
||||
storage: SQLiteBackend, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
@@ -3809,6 +3896,7 @@ def test_every_mutable_column_probes_its_classification(
|
||||
"base_url": "https://other.example/v1",
|
||||
"api_key": "sk-new",
|
||||
"context_window": 4096,
|
||||
"max_concurrency": 2,
|
||||
"capabilities": {"note": "probe"},
|
||||
# Arm direction: the loop seeds THIS row disabled, so the probe is the
|
||||
# gated false→true flip rather than the carved-out disarm.
|
||||
|
||||
@@ -909,6 +909,29 @@ def test_model_response_controls_are_capability_driven_and_sparse() -> None:
|
||||
assert 'apiSurfEl.addEventListener("change", _onModelFieldChange)' in admin
|
||||
|
||||
|
||||
def test_model_max_concurrency_form_round_trips_strict_integer() -> None:
|
||||
html = _CONSOLE_INDEX.read_text(encoding="utf-8")
|
||||
admin = _CONSOLE_ADMIN_JS.read_text(encoding="utf-8")
|
||||
|
||||
assert 'id="model-max-concurrency"' in html
|
||||
assert 'max="2147483647"' in html
|
||||
assert "0 = unlimited" in html
|
||||
|
||||
create = _slice_function_body(admin, "showCreateModelModal")
|
||||
edit = _slice_function_body(admin, "showEditModelModal")
|
||||
render = _slice_function_body(admin, "_renderModels")
|
||||
assert create is not None and edit is not None and render is not None
|
||||
assert 'getElementById("model-max-concurrency").value = "0"' in create
|
||||
assert "m.max_concurrency != null ? m.max_concurrency : 0" in edit
|
||||
# submitCreateModel is longer than the balanced-slice helper's bounded
|
||||
# window; these names are unique to that form path, so whole-file pins are
|
||||
# both stable and unambiguous.
|
||||
assert "Number.isInteger(maxConcurrency)" in admin
|
||||
assert "maxConcurrency > 2147483647" in admin
|
||||
assert "form.max_concurrency = maxConcurrency" in admin
|
||||
assert 'overrides.push("limit=" + m.max_concurrency)' in render
|
||||
|
||||
|
||||
def test_shared_utils_defines_set_safe_html_helper() -> None:
|
||||
"""``setSafeHtml`` in ``shared/utils.js`` is the single audited entry
|
||||
point for installing trusted HTML strings into a DOM element outside
|
||||
|
||||
@@ -16,6 +16,8 @@ import pytest
|
||||
from turnstone.core.deadline import (
|
||||
DeadlineCancelledError,
|
||||
DeadlineExceededError,
|
||||
StreamAbortRef,
|
||||
run_abortable_with_deadline,
|
||||
run_with_deadline,
|
||||
)
|
||||
|
||||
@@ -111,6 +113,70 @@ def test_on_abandon_errors_do_not_mask_the_deadline_error() -> None:
|
||||
)
|
||||
|
||||
|
||||
def test_abortable_deadline_credits_only_active_admission_wait() -> None:
|
||||
def _work(ref: StreamAbortRef) -> str:
|
||||
ref.begin_admission_wait()
|
||||
time.sleep(0.12)
|
||||
ref.end_admission_wait()
|
||||
time.sleep(0.02)
|
||||
return "ok"
|
||||
|
||||
start = time.monotonic()
|
||||
assert (
|
||||
run_abortable_with_deadline(
|
||||
_work,
|
||||
timeout=0.05,
|
||||
poll=0.005,
|
||||
thread_name="dl-admission-credit",
|
||||
)
|
||||
== "ok"
|
||||
)
|
||||
assert time.monotonic() - start >= 0.12
|
||||
|
||||
|
||||
def test_provider_time_still_expires_after_admission_credit() -> None:
|
||||
provider_started = threading.Event()
|
||||
|
||||
def _work(ref: StreamAbortRef) -> None:
|
||||
ref.begin_admission_wait()
|
||||
time.sleep(0.08)
|
||||
ref.end_admission_wait()
|
||||
provider_started.set()
|
||||
time.sleep(1.0)
|
||||
|
||||
with pytest.raises(DeadlineExceededError):
|
||||
run_abortable_with_deadline(
|
||||
_work,
|
||||
timeout=0.05,
|
||||
poll=0.005,
|
||||
thread_name="dl-provider-after-admission",
|
||||
)
|
||||
|
||||
# The admission interval was credited (the call reached provider work),
|
||||
# but that work consumed the unchanged logical deadline and timed out.
|
||||
assert provider_started.is_set()
|
||||
|
||||
|
||||
def test_dispatch_marker_is_observability_not_a_clock_reset() -> None:
|
||||
captured: list[StreamAbortRef] = []
|
||||
|
||||
def _work(ref: StreamAbortRef) -> None:
|
||||
captured.append(ref)
|
||||
time.sleep(0.06)
|
||||
ref.mark_dispatch()
|
||||
time.sleep(0.08)
|
||||
|
||||
with pytest.raises(DeadlineExceededError):
|
||||
run_abortable_with_deadline(
|
||||
_work,
|
||||
timeout=0.1,
|
||||
poll=0.005,
|
||||
thread_name="dl-dispatch-marker",
|
||||
)
|
||||
assert captured[0].dispatch_count == 1
|
||||
assert captured[0].last_dispatch_at is not None
|
||||
|
||||
|
||||
class TestStreamAbortRef:
|
||||
def test_abort_closes_captured_stream(self) -> None:
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
+129
-16
@@ -5,18 +5,24 @@ from __future__ import annotations
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import replace
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from tests._session_helpers import as_stream
|
||||
from tests._session_helpers import mock_completion_result as _mock_result
|
||||
from turnstone.core.admission import ModelAdmission
|
||||
from turnstone.core.judge import IntentJudge, IntentVerdict, JudgeConfig, evaluate_heuristic
|
||||
from turnstone.core.model_backend_auth import BackendAuthUnavailableError
|
||||
from turnstone.core.model_registry import ModelConfig
|
||||
from turnstone.core.model_turn import ModelLane, ResolvedModelBinding
|
||||
from turnstone.core.providers._protocol import ModelCapabilities
|
||||
from turnstone.core.providers._protocol import IncompleteStreamError, ModelCapabilities
|
||||
from turnstone.core.trajectory import Role
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pytest
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -412,33 +418,140 @@ class TestCancelEventSemantics:
|
||||
assert all(v.tier == "llm" for v in results)
|
||||
assert provider.create_streaming.call_count == 3
|
||||
|
||||
def test_batch_backend_auth_resolves_once_and_reuses_token(self):
|
||||
"""One judge batch owns one delegated credential snapshot."""
|
||||
def test_queued_backend_auth_resolves_after_admission_per_attempt(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A queued judge never ages and reuses a pre-admission token."""
|
||||
monkeypatch.setattr("turnstone.core.model_turn._DRAIN_RETRY_BASE_DELAY", 0.0)
|
||||
provider = _make_mock_provider(_good_verdict_json())
|
||||
provider.retryable_error_names = frozenset({"IncompleteStreamError"})
|
||||
|
||||
def _incomplete_stream() -> Any:
|
||||
yield from ()
|
||||
raise IncompleteStreamError("retry after admission")
|
||||
|
||||
provider.create_streaming.side_effect = [
|
||||
_incomplete_stream(),
|
||||
as_stream(_mock_result(_good_verdict_json())),
|
||||
]
|
||||
judge = _make_judge(provider)
|
||||
admission = ModelAdmission("judge", 1)
|
||||
auth_config = MagicMock(name="pinned-auth-config")
|
||||
judge._lane = replace(
|
||||
judge._lane,
|
||||
alias="judge",
|
||||
admission=admission,
|
||||
backend_auth_config=auth_config,
|
||||
)
|
||||
|
||||
batch_client = MagicMock()
|
||||
bound_client = object()
|
||||
batch_client.with_options.return_value = bound_client
|
||||
batch_client.with_options.side_effect = lambda *, api_key: f"client:{api_key}"
|
||||
judge._create_client = MagicMock(return_value=batch_client) # type: ignore[method-assign]
|
||||
resolver = MagicMock(return_value="user-a-token")
|
||||
|
||||
token_epoch = "stale"
|
||||
resolutions: list[tuple[str, Any, str]] = []
|
||||
|
||||
def _resolve(alias: str, config: ModelConfig | None) -> str:
|
||||
token = f"{token_epoch}-{len(resolutions) + 1}"
|
||||
resolutions.append((alias, config, token))
|
||||
return token
|
||||
|
||||
results: list[IntentVerdict] = []
|
||||
items = [_make_item(call_id=f"tc_{i}") for i in range(2)]
|
||||
done = threading.Event()
|
||||
holder = admission.acquire()
|
||||
|
||||
judge.evaluate(
|
||||
items,
|
||||
[_make_item()],
|
||||
[{"role": "user", "content": "test"}],
|
||||
results.append,
|
||||
backend_auth_resolver=resolver,
|
||||
done_callback=done.set,
|
||||
backend_auth_resolver=_resolve,
|
||||
)
|
||||
_wait_for(results, 2)
|
||||
|
||||
resolver.assert_called_once_with("", None)
|
||||
deadline = time.monotonic() + 2.0
|
||||
while admission.snapshot().queued != 1 and time.monotonic() < deadline:
|
||||
time.sleep(0.005)
|
||||
queued = admission.snapshot().queued == 1
|
||||
resolutions_while_queued = list(resolutions)
|
||||
token_epoch = "fresh"
|
||||
holder.release()
|
||||
|
||||
assert done.wait(5.0)
|
||||
assert queued
|
||||
assert resolutions_while_queued == []
|
||||
assert resolutions == [
|
||||
("judge", auth_config, "fresh-1"),
|
||||
("judge", auth_config, "fresh-2"),
|
||||
]
|
||||
assert batch_client.with_options.call_count == 2
|
||||
assert all(
|
||||
call.kwargs["client"] is bound_client
|
||||
for call in provider.create_streaming.call_args_list
|
||||
assert [call.kwargs["api_key"] for call in batch_client.with_options.call_args_list] == [
|
||||
"fresh-1",
|
||||
"fresh-2",
|
||||
]
|
||||
assert [call.kwargs["client"] for call in provider.create_streaming.call_args_list] == [
|
||||
"client:fresh-1",
|
||||
"client:fresh-2",
|
||||
]
|
||||
assert len(results) == 1
|
||||
assert results[0].tier == "llm"
|
||||
|
||||
def test_lazy_auth_failure_falls_back_remaining_batch_without_dispatch(self, caplog) -> None:
|
||||
"""A fail-closed mint after admission aborts the whole judge batch."""
|
||||
provider = _make_mock_provider(_good_verdict_json())
|
||||
judge = _make_judge(provider)
|
||||
admission = ModelAdmission("judge", 1)
|
||||
auth_config = MagicMock(name="pinned-auth-config")
|
||||
judge._lane = replace(
|
||||
judge._lane,
|
||||
alias="judge",
|
||||
admission=admission,
|
||||
backend_auth_config=auth_config,
|
||||
)
|
||||
assert [verdict.call_id for verdict in results] == ["tc_0", "tc_1"]
|
||||
|
||||
batch_client = MagicMock()
|
||||
judge._create_client = MagicMock(return_value=batch_client) # type: ignore[method-assign]
|
||||
resolver = MagicMock(side_effect=BackendAuthUnavailableError("mint unavailable"))
|
||||
items = [_make_item(call_id=f"tc_{i}") for i in range(3)]
|
||||
results: list[IntentVerdict] = []
|
||||
done = threading.Event()
|
||||
holder = admission.acquire()
|
||||
|
||||
with caplog.at_level("ERROR", logger="turnstone.core.judge"):
|
||||
judge.evaluate(
|
||||
items,
|
||||
[{"role": "user", "content": "test"}],
|
||||
results.append,
|
||||
done_callback=done.set,
|
||||
backend_auth_resolver=resolver,
|
||||
)
|
||||
|
||||
deadline = time.monotonic() + 2.0
|
||||
while admission.snapshot().queued != 1 and time.monotonic() < deadline:
|
||||
time.sleep(0.005)
|
||||
queued = admission.snapshot().queued == 1
|
||||
resolutions_while_queued = resolver.call_count
|
||||
holder.release()
|
||||
finished = done.wait(5.0)
|
||||
|
||||
assert queued
|
||||
assert resolutions_while_queued == 0
|
||||
assert finished
|
||||
resolver.assert_called_once_with("judge", auth_config)
|
||||
batch_client.with_options.assert_not_called()
|
||||
provider.create_streaming.assert_not_called()
|
||||
assert admission.snapshot().in_flight == 0
|
||||
assert [verdict.call_id for verdict in results] == ["tc_0", "tc_1", "tc_2"]
|
||||
assert all(verdict.tier == "llm_fallback" for verdict in results)
|
||||
assert all(
|
||||
"judge backend authentication failed" in verdict.reasoning for verdict in results
|
||||
)
|
||||
auth_logs = [
|
||||
record
|
||||
for record in caplog.records
|
||||
if "Judge backend authentication failed" in record.getMessage()
|
||||
]
|
||||
assert len(auth_logs) == 1
|
||||
|
||||
def test_cancelled_batch_skips_backend_auth_resolution(self):
|
||||
"""The daemon checks cancellation before doing a credential mint."""
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Migration coverage for per-alias model concurrency."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import command
|
||||
from alembic.config import Config
|
||||
|
||||
_MIGRATIONS_DIR = str(
|
||||
Path(__file__).resolve().parent.parent / "turnstone" / "core" / "storage" / "migrations"
|
||||
)
|
||||
|
||||
|
||||
def _alembic_cfg(db_path: Path) -> Config:
|
||||
cfg = Config()
|
||||
cfg.set_main_option("script_location", _MIGRATIONS_DIR)
|
||||
cfg.set_main_option("sqlalchemy.url", f"sqlite:///{db_path}")
|
||||
return cfg
|
||||
|
||||
|
||||
class TestMigration070:
|
||||
def test_upgrade_defaults_preexisting_rows_to_unlimited(self, tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "070-up.db"
|
||||
cfg = _alembic_cfg(db_path)
|
||||
command.upgrade(cfg, "069")
|
||||
engine = sa.create_engine(f"sqlite:///{db_path}")
|
||||
try:
|
||||
with engine.begin() as conn:
|
||||
conn.execute(
|
||||
sa.text(
|
||||
"INSERT INTO model_definitions "
|
||||
"(definition_id, alias, model, created, updated) "
|
||||
"VALUES ('d1', 'local', 'm', "
|
||||
"'2026-01-01T00:00:00', '2026-01-01T00:00:00')"
|
||||
)
|
||||
)
|
||||
command.upgrade(cfg, "070")
|
||||
with engine.connect() as conn:
|
||||
value = conn.execute(
|
||||
sa.text(
|
||||
"SELECT max_concurrency FROM model_definitions WHERE definition_id = 'd1'"
|
||||
)
|
||||
).scalar_one()
|
||||
assert value == 0
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
def test_downgrade_then_upgrade_round_trip(self, tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "070-roundtrip.db"
|
||||
cfg = _alembic_cfg(db_path)
|
||||
command.upgrade(cfg, "070")
|
||||
command.downgrade(cfg, "069")
|
||||
engine = sa.create_engine(f"sqlite:///{db_path}")
|
||||
try:
|
||||
columns = {c["name"] for c in sa.inspect(engine).get_columns("model_definitions")}
|
||||
assert "max_concurrency" not in columns
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
command.upgrade(cfg, "070")
|
||||
engine = sa.create_engine(f"sqlite:///{db_path}")
|
||||
try:
|
||||
columns = {c["name"] for c in sa.inspect(engine).get_columns("model_definitions")}
|
||||
assert "max_concurrency" in columns
|
||||
finally:
|
||||
engine.dispose()
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Focused contract tests for per-alias model admission."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pytest
|
||||
|
||||
import turnstone.core.admission as admission_mod
|
||||
from turnstone.core.admission import ModelAdmission
|
||||
from turnstone.core.deadline import DeadlineCancelledError, StreamAbortRef
|
||||
from turnstone.core.model_registry import (
|
||||
KEY_GUARD_DEFERRED_TO_LIFESPAN,
|
||||
ModelConfig,
|
||||
ModelRegistry,
|
||||
)
|
||||
|
||||
|
||||
def _wait_until(predicate: Any, *, timeout: float = 2.0) -> None:
|
||||
deadline = time.monotonic() + timeout
|
||||
while not predicate():
|
||||
if time.monotonic() >= deadline:
|
||||
raise AssertionError("condition did not become true")
|
||||
time.sleep(0.005)
|
||||
|
||||
|
||||
def _config(alias: str, limit: int, *, base_url: str | None = None) -> ModelConfig:
|
||||
return ModelConfig(
|
||||
alias=alias,
|
||||
base_url=base_url or f"http://{alias}.example/v1",
|
||||
api_key="test",
|
||||
model="model",
|
||||
max_concurrency=limit,
|
||||
)
|
||||
|
||||
|
||||
def test_unlimited_holders_are_counted_and_live_narrowing_drains() -> None:
|
||||
gate = ModelAdmission("primary", 0)
|
||||
first = gate.acquire()
|
||||
second = gate.acquire()
|
||||
assert gate.snapshot().in_flight == 2
|
||||
|
||||
gate.set_limit(1)
|
||||
acquired = threading.Event()
|
||||
release_waiter = threading.Event()
|
||||
|
||||
def _waiter() -> None:
|
||||
with gate.acquire():
|
||||
acquired.set()
|
||||
release_waiter.wait(2.0)
|
||||
|
||||
thread = threading.Thread(target=_waiter, daemon=True)
|
||||
thread.start()
|
||||
_wait_until(lambda: gate.snapshot().queued == 1)
|
||||
|
||||
first.release()
|
||||
assert not acquired.wait(0.05)
|
||||
second.release()
|
||||
assert acquired.wait(1.0)
|
||||
|
||||
release_waiter.set()
|
||||
thread.join(1.0)
|
||||
assert not thread.is_alive()
|
||||
assert gate.snapshot().in_flight == 0
|
||||
|
||||
|
||||
def test_fifo_waiters_and_hot_widening() -> None:
|
||||
gate = ModelAdmission("primary", 1)
|
||||
original = gate.acquire()
|
||||
acquired_order: list[int] = []
|
||||
acquired = [threading.Event(), threading.Event()]
|
||||
releases = [threading.Event(), threading.Event()]
|
||||
|
||||
def _waiter(index: int) -> None:
|
||||
with gate.acquire():
|
||||
acquired_order.append(index)
|
||||
acquired[index].set()
|
||||
releases[index].wait(2.0)
|
||||
|
||||
threads: list[threading.Thread] = []
|
||||
for index in range(2):
|
||||
thread = threading.Thread(target=_waiter, args=(index,), daemon=True)
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
_wait_until(lambda expected=index + 1: gate.snapshot().queued == expected)
|
||||
|
||||
gate.set_limit(2)
|
||||
assert acquired[0].wait(1.0)
|
||||
assert not acquired[1].wait(0.05)
|
||||
|
||||
original.release()
|
||||
assert acquired[1].wait(1.0)
|
||||
assert acquired_order == [0, 1]
|
||||
|
||||
for event in releases:
|
||||
event.set()
|
||||
for thread in threads:
|
||||
thread.join(1.0)
|
||||
assert not thread.is_alive()
|
||||
|
||||
|
||||
def test_cancelled_waiter_is_removed_and_never_admitted() -> None:
|
||||
gate = ModelAdmission("primary", 1)
|
||||
holder = gate.acquire()
|
||||
cancel_ref = StreamAbortRef()
|
||||
finished = threading.Event()
|
||||
errors: list[BaseException] = []
|
||||
|
||||
def _waiter() -> None:
|
||||
try:
|
||||
gate.acquire(cancel_ref=cancel_ref)
|
||||
except BaseException as exc: # test records the worker's exact exit
|
||||
errors.append(exc)
|
||||
finally:
|
||||
finished.set()
|
||||
|
||||
thread = threading.Thread(target=_waiter, daemon=True)
|
||||
thread.start()
|
||||
_wait_until(lambda: gate.snapshot().queued == 1)
|
||||
cancel_ref.abort()
|
||||
|
||||
assert finished.wait(1.0)
|
||||
assert len(errors) == 1
|
||||
assert isinstance(errors[0], DeadlineCancelledError)
|
||||
assert gate.snapshot().queued == 0
|
||||
holder.release()
|
||||
assert gate.snapshot().in_flight == 0
|
||||
|
||||
|
||||
def test_registry_keeps_one_gate_per_alias_across_resize_remove_and_readd() -> None:
|
||||
registry = ModelRegistry(
|
||||
{
|
||||
"alpha": _config("alpha", 1, base_url="http://shared.example/v1"),
|
||||
"beta": _config("beta", 3, base_url="http://shared.example/v1"),
|
||||
},
|
||||
default="alpha",
|
||||
)
|
||||
alpha = registry.get_admission("alpha")
|
||||
beta = registry.get_admission("beta")
|
||||
assert alpha is not beta
|
||||
|
||||
registry.reload(
|
||||
{"alpha": _config("alpha", 2)},
|
||||
default="alpha",
|
||||
app_state=KEY_GUARD_DEFERRED_TO_LIFESPAN,
|
||||
)
|
||||
assert registry.get_admission("alpha") is alpha
|
||||
assert alpha.limit == 2
|
||||
|
||||
registry.reload(
|
||||
{},
|
||||
default="",
|
||||
app_state=KEY_GUARD_DEFERRED_TO_LIFESPAN,
|
||||
)
|
||||
registry.reload(
|
||||
{"alpha": _config("alpha", 4)},
|
||||
default="alpha",
|
||||
app_state=KEY_GUARD_DEFERRED_TO_LIFESPAN,
|
||||
)
|
||||
assert registry.get_admission("alpha") is alpha
|
||||
assert alpha.limit == 4
|
||||
|
||||
|
||||
def test_wait_stall_and_resize_logs_expose_queue_state(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
events: list[tuple[str, dict[str, Any]]] = []
|
||||
monkeypatch.setattr(admission_mod, "_CANCEL_POLL_SECONDS", 0.005)
|
||||
monkeypatch.setattr(admission_mod, "_STALL_WARNING_SECONDS", 0.01)
|
||||
monkeypatch.setattr(
|
||||
admission_mod,
|
||||
"log",
|
||||
SimpleNamespace(
|
||||
info=lambda event, **fields: events.append((event, fields)),
|
||||
warning=lambda event, **fields: events.append((event, fields)),
|
||||
),
|
||||
)
|
||||
gate = ModelAdmission("alpha", 1)
|
||||
holder = gate.acquire()
|
||||
|
||||
thread = threading.Thread(target=lambda: gate.acquire().release(), daemon=True)
|
||||
thread.start()
|
||||
_wait_until(lambda: gate.snapshot().queued == 1)
|
||||
_wait_until(lambda: any(event == "model.admission_stalled" for event, _ in events))
|
||||
holder.release()
|
||||
thread.join(1.0)
|
||||
gate.set_limit(2)
|
||||
|
||||
by_name = {event: fields for event, fields in events}
|
||||
assert by_name["model.admission_wait"]["alias"] == "alpha"
|
||||
assert by_name["model.admission_wait"]["queued_ahead"] == 0
|
||||
assert by_name["model.admission_stalled"]["in_flight"] == 1
|
||||
assert by_name["model.admission_stalled"]["queued"] == 1
|
||||
assert by_name["model.admission_resized"]["previous_limit"] == 1
|
||||
assert by_name["model.admission_resized"]["limit"] == 2
|
||||
@@ -0,0 +1,462 @@
|
||||
"""Live request-count assertions for per-alias model admission.
|
||||
|
||||
These tests exercise the real OpenAI-compatible backend configured by the
|
||||
same environment contract as :mod:`tests.test_server_live`. A local threaded
|
||||
reverse proxy sits between ``model_turn`` and that backend solely to count
|
||||
HTTP requests. It buffers each complete upstream response, then relays the
|
||||
unchanged SSE payload to Turnstone's real streaming provider path.
|
||||
|
||||
Run explicitly (a backend must be listening at ``TURNSTONE_TEST_BASE_URL``,
|
||||
which defaults to ``http://localhost:8000/v1``)::
|
||||
|
||||
pytest tests/test_model_admission_live.py -m live -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from collections import defaultdict
|
||||
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 openai import OpenAI
|
||||
|
||||
from turnstone.core.model_registry import ModelConfig, ModelRegistry
|
||||
from turnstone.core.model_turn import ModelLane, model_turn, resolve_model_binding
|
||||
from turnstone.core.trajectory import Turn
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from types import TracebackType
|
||||
|
||||
|
||||
_LIVE_BASE_URL = os.environ.get("TURNSTONE_TEST_BASE_URL", "http://localhost:8000/v1")
|
||||
_LIVE_API_KEY = os.environ.get("TURNSTONE_TEST_API_KEY", "not-needed") or "not-needed"
|
||||
_HOP_BY_HOP_HEADERS = frozenset(
|
||||
{
|
||||
"connection",
|
||||
"content-encoding",
|
||||
"content-length",
|
||||
"keep-alive",
|
||||
"proxy-authenticate",
|
||||
"proxy-authorization",
|
||||
"te",
|
||||
"trailer",
|
||||
"transfer-encoding",
|
||||
"upgrade",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ObservedCounts:
|
||||
requests: int
|
||||
forwarded: int
|
||||
completed: int
|
||||
failed: int
|
||||
active: int
|
||||
peak: int
|
||||
forward_active: int
|
||||
forward_peak: int
|
||||
requests_by_alias: dict[str, int]
|
||||
forwarded_by_alias: dict[str, int]
|
||||
completed_by_alias: dict[str, int]
|
||||
peak_by_alias: dict[str, int]
|
||||
forward_peak_by_alias: dict[str, int]
|
||||
rendezvous_timed_out: bool
|
||||
|
||||
|
||||
class _ConcurrencyCounter:
|
||||
"""Thread-safe request counts with one bounded overlap rendezvous.
|
||||
|
||||
The first ``rendezvous_size`` admitted requests wait for each other before
|
||||
any is forwarded upstream. A short settling hold after the rendezvous
|
||||
makes an over-admission regression observable as a peak above the cap,
|
||||
independent of how quickly the live backend generates the tiny response.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
rendezvous_size: int,
|
||||
rendezvous_timeout: float = 3.0,
|
||||
settling_hold: float = 0.5,
|
||||
) -> None:
|
||||
self._lock = threading.Lock()
|
||||
self._rendezvous = threading.Event()
|
||||
self._rendezvous_size = rendezvous_size
|
||||
self._rendezvous_timeout = rendezvous_timeout
|
||||
self._settling_hold = settling_hold
|
||||
self._requests = 0
|
||||
self._forwarded = 0
|
||||
self._completed = 0
|
||||
self._failed = 0
|
||||
self._active = 0
|
||||
self._peak = 0
|
||||
self._forward_active = 0
|
||||
self._forward_peak = 0
|
||||
self._active_by_alias: dict[str, int] = defaultdict(int)
|
||||
self._forward_active_by_alias: dict[str, int] = defaultdict(int)
|
||||
self._requests_by_alias: dict[str, int] = defaultdict(int)
|
||||
self._forwarded_by_alias: dict[str, int] = defaultdict(int)
|
||||
self._completed_by_alias: dict[str, int] = defaultdict(int)
|
||||
self._peak_by_alias: dict[str, int] = defaultdict(int)
|
||||
self._forward_peak_by_alias: dict[str, int] = defaultdict(int)
|
||||
self._rendezvous_timed_out = False
|
||||
|
||||
def enter(self, alias: str) -> None:
|
||||
with self._lock:
|
||||
self._requests += 1
|
||||
self._requests_by_alias[alias] += 1
|
||||
self._active += 1
|
||||
self._active_by_alias[alias] += 1
|
||||
self._peak = max(self._peak, self._active)
|
||||
self._peak_by_alias[alias] = max(
|
||||
self._peak_by_alias[alias], self._active_by_alias[alias]
|
||||
)
|
||||
if self._active >= self._rendezvous_size:
|
||||
self._rendezvous.set()
|
||||
|
||||
if not self._rendezvous.wait(self._rendezvous_timeout):
|
||||
# Release this and every later request so a broken cap fails by
|
||||
# count instead of leaving live-test worker threads parked.
|
||||
with self._lock:
|
||||
self._rendezvous_timed_out = True
|
||||
self._rendezvous.set()
|
||||
time.sleep(self._settling_hold)
|
||||
|
||||
def begin_forward(self, alias: str) -> None:
|
||||
"""Record one HTTP request entering the real upstream backend."""
|
||||
with self._lock:
|
||||
self._forwarded += 1
|
||||
self._forwarded_by_alias[alias] += 1
|
||||
self._forward_active += 1
|
||||
self._forward_active_by_alias[alias] += 1
|
||||
self._forward_peak = max(self._forward_peak, self._forward_active)
|
||||
self._forward_peak_by_alias[alias] = max(
|
||||
self._forward_peak_by_alias[alias],
|
||||
self._forward_active_by_alias[alias],
|
||||
)
|
||||
|
||||
def end_forward(self, alias: str) -> None:
|
||||
with self._lock:
|
||||
self._forward_active -= 1
|
||||
self._forward_active_by_alias[alias] -= 1
|
||||
|
||||
def leave(self, alias: str, *, completed: bool) -> None:
|
||||
with self._lock:
|
||||
self._active -= 1
|
||||
self._active_by_alias[alias] -= 1
|
||||
if completed:
|
||||
self._completed += 1
|
||||
self._completed_by_alias[alias] += 1
|
||||
else:
|
||||
self._failed += 1
|
||||
|
||||
def snapshot(self) -> _ObservedCounts:
|
||||
with self._lock:
|
||||
return _ObservedCounts(
|
||||
requests=self._requests,
|
||||
forwarded=self._forwarded,
|
||||
completed=self._completed,
|
||||
failed=self._failed,
|
||||
active=self._active,
|
||||
peak=self._peak,
|
||||
forward_active=self._forward_active,
|
||||
forward_peak=self._forward_peak,
|
||||
requests_by_alias=dict(self._requests_by_alias),
|
||||
forwarded_by_alias=dict(self._forwarded_by_alias),
|
||||
completed_by_alias=dict(self._completed_by_alias),
|
||||
peak_by_alias=dict(self._peak_by_alias),
|
||||
forward_peak_by_alias=dict(self._forward_peak_by_alias),
|
||||
rendezvous_timed_out=self._rendezvous_timed_out,
|
||||
)
|
||||
|
||||
|
||||
class _CountingProxyServer(ThreadingHTTPServer):
|
||||
daemon_threads = True
|
||||
|
||||
def __init__(self, counter: _ConcurrencyCounter) -> None:
|
||||
super().__init__(("127.0.0.1", 0), _CountingProxyHandler)
|
||||
self.counter = counter
|
||||
self.upstream_base_url = _LIVE_BASE_URL.rstrip("/")
|
||||
self.upstream_authorization = f"Bearer {_LIVE_API_KEY}"
|
||||
self.alias_by_authorization: dict[str, str] = {}
|
||||
|
||||
|
||||
class _CountingProxyHandler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
@property
|
||||
def _proxy(self) -> _CountingProxyServer:
|
||||
if not isinstance(self.server, _CountingProxyServer):
|
||||
raise TypeError("counting handler requires _CountingProxyServer")
|
||||
return self.server
|
||||
|
||||
def do_POST(self) -> None: # noqa: N802 - BaseHTTPRequestHandler API
|
||||
alias = self._proxy.alias_by_authorization.get(self.headers.get("Authorization", ""))
|
||||
if alias is None:
|
||||
self._write_json_error(401, "unknown live-test alias credential")
|
||||
return
|
||||
|
||||
raw_length = self.headers.get("Content-Length", "0")
|
||||
try:
|
||||
content_length = int(raw_length)
|
||||
except ValueError:
|
||||
self._write_json_error(400, "invalid content length")
|
||||
return
|
||||
body = self.rfile.read(content_length)
|
||||
counter = self._proxy.counter
|
||||
counter.enter(alias)
|
||||
completed = False
|
||||
try:
|
||||
headers = {
|
||||
key: value
|
||||
for key, value in self.headers.items()
|
||||
if key.lower() not in _HOP_BY_HOP_HEADERS
|
||||
and key.lower() not in {"authorization", "host"}
|
||||
}
|
||||
# The temporary per-alias credentials are observation tags only;
|
||||
# the real backend sees exactly the configured live credential.
|
||||
headers["Authorization"] = self._proxy.upstream_authorization
|
||||
headers["Accept-Encoding"] = "identity"
|
||||
counter.begin_forward(alias)
|
||||
try:
|
||||
upstream = httpx.post(
|
||||
f"{self._proxy.upstream_base_url}{self.path}",
|
||||
content=body,
|
||||
headers=headers,
|
||||
timeout=httpx.Timeout(60.0, connect=5.0),
|
||||
)
|
||||
finally:
|
||||
counter.end_forward(alias)
|
||||
payload = upstream.content
|
||||
self.send_response(upstream.status_code)
|
||||
for key, value in upstream.headers.multi_items():
|
||||
if key.lower() not in _HOP_BY_HOP_HEADERS:
|
||||
self.send_header(key, value)
|
||||
self.send_header("Content-Length", str(len(payload)))
|
||||
self.send_header("Connection", "close")
|
||||
self.end_headers()
|
||||
self.wfile.write(payload)
|
||||
self.wfile.flush()
|
||||
completed = upstream.status_code < 400
|
||||
except Exception:
|
||||
# The worker observes a normal 502; its exception and the failed
|
||||
# count retain the useful signal without leaking backend details.
|
||||
self._write_json_error(502, "live backend proxy failure")
|
||||
finally:
|
||||
counter.leave(alias, completed=completed)
|
||||
|
||||
def log_message(self, _format: str, *args: Any) -> None:
|
||||
"""Keep expected local proxy traffic out of pytest output."""
|
||||
|
||||
def _write_json_error(self, status: int, detail: str) -> None:
|
||||
payload = json.dumps({"error": {"message": detail}}).encode()
|
||||
try:
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(payload)))
|
||||
self.send_header("Connection", "close")
|
||||
self.end_headers()
|
||||
self.wfile.write(payload)
|
||||
self.wfile.flush()
|
||||
except (BrokenPipeError, ConnectionError):
|
||||
pass
|
||||
|
||||
|
||||
class _CountingProxy:
|
||||
def __init__(self, counter: _ConcurrencyCounter) -> None:
|
||||
self._server = _CountingProxyServer(counter)
|
||||
self._thread = threading.Thread(
|
||||
target=self._server.serve_forever,
|
||||
name="model-admission-live-proxy",
|
||||
daemon=True,
|
||||
)
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
address = self._server.server_address
|
||||
host = address[0]
|
||||
port = address[1]
|
||||
if not isinstance(host, str) or not isinstance(port, int):
|
||||
raise TypeError("counting proxy did not bind an IPv4 TCP address")
|
||||
return f"http://{host}:{port}"
|
||||
|
||||
def credential_for(self, alias: str) -> str:
|
||||
credential = f"turnstone-live-admission-{alias}"
|
||||
self._server.alias_by_authorization[f"Bearer {credential}"] = alias
|
||||
return credential
|
||||
|
||||
def __enter__(self) -> _CountingProxy:
|
||||
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._thread.join(timeout=5.0)
|
||||
if self._thread.is_alive():
|
||||
raise AssertionError("live counting proxy did not stop")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def live_model_id() -> str:
|
||||
"""Auto-detect the model using the established live-backend contract."""
|
||||
client = OpenAI(base_url=_LIVE_BASE_URL, api_key=_LIVE_API_KEY)
|
||||
try:
|
||||
models = client.models.list()
|
||||
finally:
|
||||
client.close()
|
||||
ids = [model.id for model in models.data]
|
||||
assert ids, "No models found on the live backend"
|
||||
return ids[0]
|
||||
|
||||
|
||||
def _live_lane(
|
||||
proxy: _CountingProxy,
|
||||
model: str,
|
||||
*,
|
||||
alias: str,
|
||||
limit: int,
|
||||
) -> tuple[ModelRegistry, ModelLane]:
|
||||
config = ModelConfig(
|
||||
alias=alias,
|
||||
base_url=proxy.base_url,
|
||||
api_key=proxy.credential_for(alias),
|
||||
model=model,
|
||||
provider="openai-compatible",
|
||||
max_concurrency=limit,
|
||||
)
|
||||
registry = ModelRegistry({alias: config}, default=alias)
|
||||
try:
|
||||
binding = resolve_model_binding(registry, alias)
|
||||
# Bound each live call and disable SDK retries: the proxy's exact
|
||||
# request count should equal the number of model_turn invocations,
|
||||
# while model_turn still exercises the real provider and admission
|
||||
# lifecycle.
|
||||
client = binding.lane.client.with_options(timeout=60.0, max_retries=0)
|
||||
return registry, dataclasses.replace(binding.lane, client=client)
|
||||
except BaseException:
|
||||
registry.shutdown()
|
||||
raise
|
||||
|
||||
|
||||
def _run_parallel(lanes: list[ModelLane], model: str) -> None:
|
||||
barrier = threading.Barrier(len(lanes))
|
||||
|
||||
def _one_turn(lane: ModelLane) -> str:
|
||||
barrier.wait(timeout=10.0)
|
||||
result = model_turn(
|
||||
lane,
|
||||
[Turn.user("Reply with OK.")],
|
||||
max_tokens=8,
|
||||
)
|
||||
return result.serving_model
|
||||
|
||||
with ThreadPoolExecutor(max_workers=len(lanes)) as pool:
|
||||
futures = [pool.submit(_one_turn, lane) for lane in lanes]
|
||||
observed_models = [future.result(timeout=75.0) for future in futures]
|
||||
assert observed_models == [model] * len(lanes)
|
||||
|
||||
|
||||
@pytest.mark.live
|
||||
class TestLiveModelAdmissionCounts:
|
||||
"""Count genuine backend requests at the per-alias admission boundary."""
|
||||
|
||||
def test_alias_cap_bounds_exact_peak_and_completes_every_request(
|
||||
self, live_model_id: str
|
||||
) -> None:
|
||||
counter = _ConcurrencyCounter(rendezvous_size=2)
|
||||
with _CountingProxy(counter) as proxy:
|
||||
registry, lane = _live_lane(proxy, live_model_id, alias="limited", limit=2)
|
||||
try:
|
||||
_run_parallel([lane, lane, lane, lane], live_model_id)
|
||||
finally:
|
||||
try:
|
||||
lane.client.close()
|
||||
finally:
|
||||
registry.shutdown()
|
||||
|
||||
counts = counter.snapshot()
|
||||
assert counts.requests == 4
|
||||
assert counts.forwarded == 4
|
||||
assert counts.completed == 4
|
||||
assert counts.failed == 0
|
||||
assert counts.active == 0
|
||||
assert counts.peak == 2
|
||||
assert counts.forward_active == 0
|
||||
assert counts.forward_peak == 2
|
||||
assert counts.requests_by_alias == {"limited": 4}
|
||||
assert counts.forwarded_by_alias == {"limited": 4}
|
||||
assert counts.completed_by_alias == {"limited": 4}
|
||||
assert counts.peak_by_alias == {"limited": 2}
|
||||
assert counts.forward_peak_by_alias == {"limited": 2}
|
||||
assert counts.rendezvous_timed_out is False
|
||||
|
||||
def test_aliases_sharing_endpoint_have_independent_caps(self, live_model_id: str) -> None:
|
||||
counter = _ConcurrencyCounter(rendezvous_size=2)
|
||||
with _CountingProxy(counter) as proxy:
|
||||
configs = {
|
||||
alias: ModelConfig(
|
||||
alias=alias,
|
||||
base_url=proxy.base_url,
|
||||
api_key=proxy.credential_for(alias),
|
||||
model=live_model_id,
|
||||
provider="openai-compatible",
|
||||
max_concurrency=1,
|
||||
)
|
||||
for alias in ("alpha", "beta")
|
||||
}
|
||||
# The two registry aliases deliberately name the exact same URL;
|
||||
# both requests are then forwarded to the same physical backend.
|
||||
assert configs["alpha"].base_url == configs["beta"].base_url
|
||||
registry = ModelRegistry(configs, default="alpha")
|
||||
lanes: list[ModelLane] = []
|
||||
try:
|
||||
alpha = resolve_model_binding(registry, "alpha").lane
|
||||
beta = resolve_model_binding(registry, "beta").lane
|
||||
alpha = dataclasses.replace(
|
||||
alpha, client=alpha.client.with_options(timeout=60.0, max_retries=0)
|
||||
)
|
||||
lanes.append(alpha)
|
||||
beta = dataclasses.replace(
|
||||
beta, client=beta.client.with_options(timeout=60.0, max_retries=0)
|
||||
)
|
||||
lanes.append(beta)
|
||||
_run_parallel(lanes, live_model_id)
|
||||
finally:
|
||||
try:
|
||||
for lane in lanes:
|
||||
lane.client.close()
|
||||
finally:
|
||||
registry.shutdown()
|
||||
|
||||
counts = counter.snapshot()
|
||||
assert counts.requests == 2
|
||||
assert counts.forwarded == 2
|
||||
assert counts.completed == 2
|
||||
assert counts.failed == 0
|
||||
assert counts.active == 0
|
||||
assert counts.peak == 2
|
||||
assert counts.forward_active == 0
|
||||
assert counts.forward_peak == 2
|
||||
assert counts.requests_by_alias == {"alpha": 1, "beta": 1}
|
||||
assert counts.forwarded_by_alias == {"alpha": 1, "beta": 1}
|
||||
assert counts.completed_by_alias == {"alpha": 1, "beta": 1}
|
||||
assert counts.peak_by_alias == {"alpha": 1, "beta": 1}
|
||||
assert counts.forward_peak_by_alias == {"alpha": 1, "beta": 1}
|
||||
assert counts.rendezvous_timed_out is False
|
||||
@@ -33,6 +33,7 @@ class TestModelDefinitionStorage:
|
||||
assert m["base_url"] == "https://api.openai.com/v1"
|
||||
assert m["api_key"] == "sk-test"
|
||||
assert m["context_window"] == 128000
|
||||
assert m["max_concurrency"] == 0
|
||||
assert m["capabilities"] == "{}"
|
||||
assert m["enabled"] is True
|
||||
|
||||
@@ -158,6 +159,7 @@ class TestModelDefinitionStorage:
|
||||
assert m["base_url"] == ""
|
||||
assert m["api_key"] == ""
|
||||
assert m["context_window"] == 32768
|
||||
assert m["max_concurrency"] == 0
|
||||
assert m["capabilities"] == "{}"
|
||||
assert m["enabled"] is True
|
||||
assert m["created_by"] == ""
|
||||
@@ -166,6 +168,23 @@ class TestModelDefinitionStorage:
|
||||
assert m["max_tokens"] is None
|
||||
assert m["reasoning_effort"] is None
|
||||
|
||||
def test_create_update_and_list_max_concurrency(self, db: SQLiteBackend) -> None:
|
||||
did = _make_id()
|
||||
db.create_model_definition(
|
||||
definition_id=did,
|
||||
alias="limited",
|
||||
model="gpt-5",
|
||||
max_concurrency=3,
|
||||
)
|
||||
assert db.get_model_definition(did)["max_concurrency"] == 3
|
||||
|
||||
assert db.update_model_definition(did, max_concurrency=1)
|
||||
assert db.get_model_definition_by_alias("limited")["max_concurrency"] == 1
|
||||
assert db.list_model_definitions()[0]["max_concurrency"] == 1
|
||||
|
||||
assert db.update_model_definition(did, max_concurrency=0)
|
||||
assert db.get_model_definition(did)["max_concurrency"] == 0
|
||||
|
||||
def test_create_with_sampling_params(self, db: SQLiteBackend) -> None:
|
||||
did = _make_id()
|
||||
db.create_model_definition(
|
||||
|
||||
@@ -17,6 +17,7 @@ from turnstone.core import model_registry as mr_module
|
||||
from turnstone.core.model_registry import (
|
||||
KEY_GUARD_DEFERRED_TO_LIFESPAN,
|
||||
DynamicAuthKeyError,
|
||||
ModelConcurrencyConfigError,
|
||||
ModelConfig,
|
||||
ModelRegistry,
|
||||
UnknownModelAliasError,
|
||||
@@ -50,6 +51,7 @@ class TestModelConfig:
|
||||
assert cfg.alias == "local"
|
||||
assert cfg.model == "qwen3-32b"
|
||||
assert cfg.context_window == 32768 # default
|
||||
assert cfg.max_concurrency == 0
|
||||
|
||||
def test_custom_context_window(self) -> None:
|
||||
cfg = ModelConfig(
|
||||
@@ -112,6 +114,16 @@ class TestModelConfig:
|
||||
assert cfg.surface_persisted_reasoning is False
|
||||
assert cfg.replay_reasoning_to_model is True
|
||||
|
||||
def test_max_concurrency_is_strict_and_not_binding_identity(self) -> None:
|
||||
unlimited = ModelConfig(alias="x", base_url="x", api_key="x", model="x")
|
||||
limited = dataclasses.replace(unlimited, max_concurrency=2)
|
||||
assert limited.max_concurrency == 2
|
||||
assert unlimited == limited
|
||||
|
||||
for invalid in (-1, 2_147_483_648, True, 1.0, "1", None):
|
||||
with pytest.raises(ModelConcurrencyConfigError, match="max_concurrency"):
|
||||
dataclasses.replace(unlimited, max_concurrency=invalid) # type: ignore[arg-type]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ModelRegistry
|
||||
@@ -386,6 +398,7 @@ class TestLoadModelRegistry:
|
||||
"api_key": "sk-test",
|
||||
"model": "gpt-4o",
|
||||
"context_window": 128000,
|
||||
"max_concurrency": 2,
|
||||
},
|
||||
},
|
||||
"model": {
|
||||
@@ -404,8 +417,26 @@ class TestLoadModelRegistry:
|
||||
assert reg.has_alias("openai")
|
||||
assert not reg.has_alias("default")
|
||||
assert reg.default == "openai"
|
||||
_, model, _, _ = reg.resolve()
|
||||
_, model, cfg, _ = reg.resolve()
|
||||
assert model == "gpt-4o"
|
||||
assert cfg.max_concurrency == 2
|
||||
|
||||
@pytest.mark.parametrize("invalid", [-1, 2_147_483_648, True, 1.0, "1", None])
|
||||
def test_config_rejects_invalid_max_concurrency(self, invalid: Any) -> None:
|
||||
fake_cfg = {
|
||||
"models": {
|
||||
"local": {
|
||||
"base_url": "http://localhost:8000/v1",
|
||||
"model": "m",
|
||||
"max_concurrency": invalid,
|
||||
}
|
||||
}
|
||||
}
|
||||
with (
|
||||
patch("turnstone.core.model_registry.load_config", return_value=fake_cfg),
|
||||
pytest.raises(ModelConcurrencyConfigError, match="max_concurrency"),
|
||||
):
|
||||
load_model_registry()
|
||||
|
||||
def test_config_context_window_zero_inherits_detected(self) -> None:
|
||||
"""``context_window = 0`` in a [models.*] entry is the auto-detect
|
||||
@@ -794,6 +825,7 @@ class TestLoadModelRegistryWithDB:
|
||||
"temperature": 1.5,
|
||||
"max_tokens": 4096,
|
||||
"reasoning_effort": "high",
|
||||
"max_concurrency": 4,
|
||||
}
|
||||
]
|
||||
)
|
||||
@@ -803,6 +835,7 @@ class TestLoadModelRegistryWithDB:
|
||||
assert cfg.temperature == 1.5
|
||||
assert cfg.max_tokens == 4096
|
||||
assert cfg.reasoning_effort == "high"
|
||||
assert cfg.max_concurrency == 4
|
||||
|
||||
def test_db_sampling_params_null_means_none(self) -> None:
|
||||
"""NULL sampling params in DB map to None (use global default)."""
|
||||
@@ -3638,8 +3671,8 @@ class TestSessionAgentModel:
|
||||
session._run_agent([Turn.user("x")], label="task", agent_alias="fast")
|
||||
assert captured["model"] == "fast-model"
|
||||
|
||||
def test_session_fallback_uses_exact_primary_lane_and_caps(self) -> None:
|
||||
"""A sub-agent inherits the primary lane with auth already consumed."""
|
||||
def test_session_fallback_uses_exact_primary_lane_and_pinned_auth(self) -> None:
|
||||
"""A sub-agent keeps the primary lane and pins its auth resolver."""
|
||||
import turnstone.core.session as session_module
|
||||
|
||||
reg = self._three_model_registry() # no agent_model / plan_model set
|
||||
@@ -3658,8 +3691,11 @@ class TestSessionAgentModel:
|
||||
|
||||
assert model_turn_spy.call_count == 1
|
||||
used_lane = model_turn_spy.call_args.args[0]
|
||||
assert used_lane == dataclasses.replace(primary_lane, backend_auth_resolver=None)
|
||||
assert used_lane.backend_auth_resolver is None
|
||||
assert used_lane == dataclasses.replace(
|
||||
primary_lane,
|
||||
backend_auth_resolver=used_lane.backend_auth_resolver,
|
||||
)
|
||||
assert used_lane.backend_auth_resolver is not None
|
||||
assert used_lane.capabilities is primary_caps
|
||||
assert used_lane.client is _client(session)
|
||||
assert used_lane.alias == "main"
|
||||
|
||||
+227
-1
@@ -12,6 +12,8 @@ import ast
|
||||
import inspect
|
||||
import logging
|
||||
import textwrap
|
||||
import threading
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
@@ -39,7 +41,7 @@ from turnstone.core.providers._protocol import (
|
||||
serialized_tool_chars,
|
||||
)
|
||||
from turnstone.core.session import ChatSession
|
||||
from turnstone.core.trajectory import Role, ToolCall, Turn
|
||||
from turnstone.core.trajectory import AttachmentRef, Role, ToolCall, Turn
|
||||
|
||||
|
||||
class _FakeProvider:
|
||||
@@ -230,6 +232,230 @@ def test_resolve_model_binding_canonicalizes_empty_alias_to_default() -> None:
|
||||
assert binding.registry_generation == 7
|
||||
|
||||
|
||||
def test_model_turn_materializes_before_admission_and_mints_inside_hold() -> None:
|
||||
order: list[str] = []
|
||||
|
||||
class _Gate:
|
||||
held = False
|
||||
|
||||
def acquire(self, *, cancel_ref: Any = None) -> Any:
|
||||
del cancel_ref
|
||||
order.append("acquire")
|
||||
gate = self
|
||||
|
||||
class _Lease:
|
||||
def __enter__(self) -> None:
|
||||
gate.held = True
|
||||
order.append("enter")
|
||||
|
||||
def __exit__(self, *_exc: object) -> None:
|
||||
gate.held = False
|
||||
order.append("release")
|
||||
|
||||
return _Lease()
|
||||
|
||||
gate = _Gate()
|
||||
bound_client = object()
|
||||
base_client = MagicMock()
|
||||
base_client.with_options.return_value = bound_client
|
||||
|
||||
def _resolve(ids: list[str]) -> dict[str, Any]:
|
||||
assert ids == ["image-1"]
|
||||
assert not gate.held
|
||||
order.append("materialize")
|
||||
return {
|
||||
"image-1": {
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "data:image/png;base64,abc"},
|
||||
}
|
||||
}
|
||||
|
||||
def _auth(_alias: str, _cfg: Any) -> str:
|
||||
assert gate.held
|
||||
order.append("auth")
|
||||
return "minted-token"
|
||||
|
||||
class _Provider(_FakeProvider):
|
||||
def create_streaming(self, **kwargs: Any) -> Any:
|
||||
assert gate.held
|
||||
assert kwargs["client"] is bound_client
|
||||
assert kwargs["resolve_attachments"] is None
|
||||
order.append("dispatch")
|
||||
self.calls.append(kwargs)
|
||||
|
||||
def _stream() -> Any:
|
||||
assert gate.held
|
||||
order.append("drain")
|
||||
yield from as_stream(CompletionResult(content="ok"))
|
||||
|
||||
return _stream()
|
||||
|
||||
provider = _Provider([])
|
||||
lane = ModelLane(
|
||||
provider=provider,
|
||||
client=base_client,
|
||||
model="m",
|
||||
alias="primary",
|
||||
backend_auth_resolver=_auth,
|
||||
admission=gate, # type: ignore[arg-type]
|
||||
)
|
||||
turn = Turn(Role.USER, (AttachmentRef(attachment_id="image-1", kind="image"),))
|
||||
|
||||
result = model_turn(lane, [turn], resolve_attachments=_resolve)
|
||||
|
||||
assert result.content == "ok"
|
||||
assert order == ["materialize", "acquire", "enter", "auth", "dispatch", "drain", "release"]
|
||||
assert result.wire_msgs == [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "data:image/png;base64,abc"},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_model_turn_releases_admission_before_retry_backoff(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from turnstone.core.deadline import StreamAbortRef
|
||||
|
||||
order: list[str] = []
|
||||
|
||||
class _Gate:
|
||||
held = False
|
||||
acquire_calls = 0
|
||||
|
||||
def acquire(self, *, cancel_ref: Any = None) -> Any:
|
||||
del cancel_ref
|
||||
self.acquire_calls += 1
|
||||
gate = self
|
||||
|
||||
class _Lease:
|
||||
def __enter__(self) -> None:
|
||||
assert not gate.held
|
||||
gate.held = True
|
||||
order.append("enter")
|
||||
|
||||
def __exit__(self, *_exc: object) -> None:
|
||||
gate.held = False
|
||||
order.append("release")
|
||||
|
||||
return _Lease()
|
||||
|
||||
gate = _Gate()
|
||||
provider = _FlakyProvider([IncompleteStreamError("retry"), CompletionResult(content="ok")])
|
||||
dispatch = provider.create_streaming
|
||||
|
||||
def _dispatch(**kwargs: Any) -> Any:
|
||||
assert gate.held
|
||||
order.append("dispatch")
|
||||
return dispatch(**kwargs)
|
||||
|
||||
provider.create_streaming = _dispatch # type: ignore[method-assign]
|
||||
|
||||
def _sleep(_delay: float) -> None:
|
||||
assert not gate.held
|
||||
order.append("backoff")
|
||||
|
||||
monkeypatch.setattr(model_turn_mod, "time", SimpleNamespace(sleep=_sleep))
|
||||
lane = ModelLane(
|
||||
provider=provider,
|
||||
client=object(),
|
||||
model="m",
|
||||
admission=gate, # type: ignore[arg-type]
|
||||
)
|
||||
ref = StreamAbortRef()
|
||||
|
||||
result = model_turn(lane, [Turn.user("x")], cancel_ref=ref)
|
||||
|
||||
assert result.content == "ok"
|
||||
assert gate.acquire_calls == 2
|
||||
assert ref.dispatch_count == 2
|
||||
assert order == ["enter", "dispatch", "release", "backoff", "enter", "dispatch", "release"]
|
||||
|
||||
|
||||
def test_same_alias_attachment_work_completes_before_outer_admission() -> None:
|
||||
from turnstone.core.admission import ModelAdmission
|
||||
|
||||
gate = ModelAdmission("primary", 1)
|
||||
nested_completed = False
|
||||
|
||||
def _resolve(ids: list[str]) -> dict[str, Any]:
|
||||
nonlocal nested_completed
|
||||
assert ids == ["image-1"]
|
||||
# Models a nested perception call using the same alias. This would
|
||||
# block forever if the outer model_turn had already taken the slot.
|
||||
with gate.acquire():
|
||||
nested_completed = True
|
||||
return {"image-1": {"type": "image_url", "image_url": {"url": "data:x"}}}
|
||||
|
||||
provider = _FakeProvider([CompletionResult(content="ok")])
|
||||
lane = ModelLane(provider=provider, client=object(), model="m", admission=gate)
|
||||
turn = Turn(Role.USER, (AttachmentRef(attachment_id="image-1", kind="image"),))
|
||||
|
||||
assert model_turn(lane, [turn], resolve_attachments=_resolve).content == "ok"
|
||||
assert nested_completed
|
||||
assert gate.snapshot().in_flight == 0
|
||||
|
||||
|
||||
def test_model_turn_releases_admission_when_eager_create_fails() -> None:
|
||||
from turnstone.core.admission import ModelAdmission
|
||||
|
||||
class _CreateFailureProvider(_FakeProvider):
|
||||
def create_streaming(self, **kwargs: Any) -> Any:
|
||||
self.calls.append(kwargs)
|
||||
raise RuntimeError("connect failed")
|
||||
|
||||
gate = ModelAdmission("primary", 1)
|
||||
provider = _CreateFailureProvider([])
|
||||
lane = ModelLane(provider=provider, client=object(), model="m", admission=gate)
|
||||
|
||||
with pytest.raises(RuntimeError, match="connect failed"):
|
||||
model_turn(lane, [Turn.user("x")])
|
||||
|
||||
assert len(provider.calls) == 1
|
||||
assert gate.snapshot().in_flight == 0
|
||||
|
||||
|
||||
def test_model_turn_cancelled_while_queued_never_dispatches() -> None:
|
||||
from turnstone.core.admission import ModelAdmission
|
||||
from turnstone.core.deadline import DeadlineCancelledError, StreamAbortRef
|
||||
|
||||
gate = ModelAdmission("primary", 1)
|
||||
holder = gate.acquire()
|
||||
provider = _FakeProvider([CompletionResult(content="never")])
|
||||
lane = ModelLane(provider=provider, client=object(), model="m", admission=gate)
|
||||
cancel_ref = StreamAbortRef()
|
||||
errors: list[BaseException] = []
|
||||
|
||||
def _call() -> None:
|
||||
try:
|
||||
model_turn(lane, [Turn.user("x")], cancel_ref=cancel_ref)
|
||||
except BaseException as exc: # test records the worker's exact exit
|
||||
errors.append(exc)
|
||||
|
||||
thread = threading.Thread(target=_call, daemon=True)
|
||||
thread.start()
|
||||
deadline = time.monotonic() + 1.0
|
||||
while gate.snapshot().queued != 1:
|
||||
if time.monotonic() >= deadline:
|
||||
holder.release()
|
||||
raise AssertionError("model turn did not queue")
|
||||
time.sleep(0.005)
|
||||
cancel_ref.abort()
|
||||
thread.join(1.0)
|
||||
holder.release()
|
||||
|
||||
assert not thread.is_alive()
|
||||
assert len(errors) == 1
|
||||
assert isinstance(errors[0], DeadlineCancelledError)
|
||||
assert provider.calls == []
|
||||
|
||||
|
||||
class _FlakyProvider:
|
||||
"""Scripted drain-time deaths: each script entry is either a
|
||||
``CompletionResult`` (streamed normally) or an exception instance
|
||||
|
||||
@@ -297,6 +297,26 @@ class TestConsoleSpec:
|
||||
)
|
||||
|
||||
|
||||
def test_model_max_concurrency_schema_is_strict_non_nullable_integer() -> None:
|
||||
from turnstone.api.console_spec import build_console_spec
|
||||
|
||||
schemas = build_console_spec()["components"]["schemas"]
|
||||
for name in ("ModelDefinitionInfo", "CreateModelDefinitionRequest"):
|
||||
prop = schemas[name]["properties"]["max_concurrency"]
|
||||
assert prop["type"] == "integer"
|
||||
assert prop["default"] == 0
|
||||
assert prop["minimum"] == 0
|
||||
assert prop["maximum"] == 2_147_483_647
|
||||
assert "anyOf" not in prop
|
||||
|
||||
update_prop = schemas["UpdateModelDefinitionRequest"]["properties"]["max_concurrency"]
|
||||
assert update_prop["type"] == "integer"
|
||||
assert update_prop["minimum"] == 0
|
||||
assert update_prop["maximum"] == 2_147_483_647
|
||||
assert "anyOf" not in update_prop
|
||||
assert "default" not in update_prop
|
||||
|
||||
|
||||
class TestCheckedInArtifactFreshness:
|
||||
"""The checked-in `sdk/typescript/*.json` specs must match their source.
|
||||
|
||||
|
||||
@@ -19,8 +19,8 @@ class _StubProvider:
|
||||
|
||||
``describe`` routes through ``model_turn``, so the stub carries the lane
|
||||
surface (``provider_name``, ``get_capabilities``) and returns a full
|
||||
``CompletionResult`` shape, and it records the ``resolve_attachments``
|
||||
callback the translator would use to materialize the by-reference parts.
|
||||
``CompletionResult`` shape, and it records that ``model_turn`` already
|
||||
materialized by-reference parts and cleared the provider callback.
|
||||
"""
|
||||
|
||||
provider_name = "openai-compatible"
|
||||
@@ -100,11 +100,10 @@ def test_describe_lowers_prompt_then_by_reference_parts() -> None:
|
||||
assert prov.last_messages is not None
|
||||
content = prov.last_messages[0]["content"]
|
||||
assert content[0]["type"] == "text" # prompt leads
|
||||
# The attachment rides by reference; the translator materializes it via
|
||||
# the threaded resolver, which must return the prebuilt parts verbatim.
|
||||
assert content[1]["attachment_id"] == "perception-input"
|
||||
assert prov.last_resolve is not None
|
||||
assert prov.last_resolve(["perception-input"]) == {"perception-input": _parts()}
|
||||
# model_turn materializes before admission, then hands the provider the
|
||||
# prebuilt inline part with no resolver left to invoke under the gate.
|
||||
assert content[1] == _parts()[0]
|
||||
assert prov.last_resolve is None
|
||||
|
||||
|
||||
def test_describe_empty_parts_skips_backend() -> None:
|
||||
|
||||
@@ -139,6 +139,7 @@ def test_model_status_route_carries_backend_auth_fields():
|
||||
auth_mode="rfc8693_obo",
|
||||
obo_audience="api://gw",
|
||||
obo_scopes="aud-gw openid",
|
||||
max_concurrency=3,
|
||||
)
|
||||
reg = ModelRegistry(models={"gw": cfg}, default="gw")
|
||||
request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(registry=reg)))
|
||||
@@ -147,6 +148,7 @@ def test_model_status_route_carries_backend_auth_fields():
|
||||
assert entry["auth_mode"] == "rfc8693_obo"
|
||||
assert entry["obo_audience"] == "api://gw"
|
||||
assert entry["obo_scopes"] == "aud-gw openid"
|
||||
assert entry["max_concurrency"] == 3
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -326,6 +328,64 @@ def test_model_reload_endpoint_rewrites_models_metadata(monkeypatch, tmp_path):
|
||||
assert {r["alias"] for r in payload} == {"a", "b"}
|
||||
|
||||
|
||||
def test_model_reload_cap_only_change_resizes_the_stable_gate(monkeypatch):
|
||||
"""Operational capacity participates in the endpoint's no-op check even
|
||||
though it deliberately does not participate in ModelConfig identity."""
|
||||
from turnstone.server import internal_model_reload
|
||||
|
||||
old_reg = ModelRegistry(
|
||||
{
|
||||
"a": ModelConfig(
|
||||
alias="a",
|
||||
base_url="http://x",
|
||||
api_key="k",
|
||||
model="a",
|
||||
max_concurrency=1,
|
||||
)
|
||||
},
|
||||
default="a",
|
||||
)
|
||||
new_reg = ModelRegistry(
|
||||
{
|
||||
"a": ModelConfig(
|
||||
alias="a",
|
||||
base_url="http://x",
|
||||
api_key="k",
|
||||
model="a",
|
||||
max_concurrency=3,
|
||||
)
|
||||
},
|
||||
default="a",
|
||||
)
|
||||
original_gate = old_reg.get_admission("a")
|
||||
app_state = SimpleNamespace(
|
||||
registry=old_reg,
|
||||
cli_model_args={
|
||||
"base_url": "",
|
||||
"api_key": "",
|
||||
"model": "",
|
||||
"context_window": 0,
|
||||
"provider": "openai",
|
||||
},
|
||||
config_store=None,
|
||||
node_id="",
|
||||
)
|
||||
request = SimpleNamespace(app=SimpleNamespace(state=app_state))
|
||||
|
||||
monkeypatch.setattr("turnstone.core.model_registry.load_model_registry", lambda **_kw: new_reg)
|
||||
monkeypatch.setattr("turnstone.core.storage._registry.get_storage", MagicMock)
|
||||
monkeypatch.setattr("turnstone.server._broadcast_agent_tool_schema_refresh", lambda _s: None)
|
||||
|
||||
response = internal_model_reload(request) # type: ignore[arg-type]
|
||||
body = json.loads(response.body)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert body.get("noop", False) is False
|
||||
assert old_reg.get_admission("a") is original_gate
|
||||
assert original_gate.limit == 3
|
||||
assert old_reg.generation == 1
|
||||
|
||||
|
||||
def test_model_reload_refuses_dynamic_auth_without_key(monkeypatch, tmp_path, caplog):
|
||||
"""A keyless node cannot acquire a dynamic alias via model-reload: 503,
|
||||
a deployment fault, not the 422 bad-arguments exit."""
|
||||
@@ -420,6 +480,44 @@ def test_model_reload_maps_auth_config_error_to_422(monkeypatch, tmp_path):
|
||||
assert old_reg.has_alias("a")
|
||||
|
||||
|
||||
def test_model_reload_maps_concurrency_config_error_to_422(monkeypatch, tmp_path):
|
||||
"""A corrupt persisted cap uses the same structured config-error exit."""
|
||||
from turnstone.core.model_registry import ModelConcurrencyConfigError
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
from turnstone.server import internal_model_reload
|
||||
|
||||
storage = SQLiteBackend(str(tmp_path / "reload.db"))
|
||||
old_reg = _registry(("a", "http://x"))
|
||||
app_state = SimpleNamespace(
|
||||
registry=old_reg,
|
||||
cli_model_args={
|
||||
"base_url": "",
|
||||
"api_key": "",
|
||||
"model": "",
|
||||
"context_window": 0,
|
||||
"provider": "openai",
|
||||
},
|
||||
config_store=None,
|
||||
node_id="node-a",
|
||||
)
|
||||
request = SimpleNamespace(app=SimpleNamespace(state=app_state))
|
||||
|
||||
def _raise_concurrency_config(**_kw):
|
||||
raise ModelConcurrencyConfigError("Model 'gw' max_concurrency must be an integer")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"turnstone.core.model_registry.load_model_registry",
|
||||
_raise_concurrency_config,
|
||||
)
|
||||
monkeypatch.setattr("turnstone.core.storage._registry.get_storage", lambda: storage)
|
||||
|
||||
response = internal_model_reload(request) # type: ignore[arg-type]
|
||||
|
||||
assert response.status_code == 422
|
||||
assert "max_concurrency" in json.loads(response.body)["reason"]
|
||||
assert old_reg.has_alias("a")
|
||||
|
||||
|
||||
def test_config_reload_maps_dynamic_auth_key_error_to_503(monkeypatch, tmp_path, caplog):
|
||||
"""Every caller of the reload chokepoint handles its refusal the same
|
||||
way: the settings fan-out answers 503, not an unhandled 500."""
|
||||
|
||||
@@ -10057,6 +10057,75 @@ def test_task_agent_static_auth_fallback_never_reresolves_as_successor():
|
||||
lane.client.with_options.assert_not_called()
|
||||
|
||||
|
||||
def test_task_agent_defers_pinned_auth_resolution_until_model_admission():
|
||||
"""A child mints for its initiating user only after admission."""
|
||||
from dataclasses import replace
|
||||
|
||||
from turnstone.core.admission import ModelAdmission
|
||||
from turnstone.core.model_turn import model_turn as real_model_turn
|
||||
|
||||
session = _make_session()
|
||||
session._acting_user_id = "user-b"
|
||||
provider = seam_provider("done", provider_name="openai-compatible")
|
||||
lane = replace_session_lane(session, provider=provider, alias="task-gateway")
|
||||
auth_config = MagicMock(name="pinned-auth-config")
|
||||
gate = ModelAdmission("task-gateway", 1)
|
||||
stale_live_resolver = MagicMock(return_value="token-for-user-b")
|
||||
lane = replace(
|
||||
lane,
|
||||
backend_auth_resolver=stale_live_resolver,
|
||||
backend_auth_config=auth_config,
|
||||
admission=gate,
|
||||
)
|
||||
session._model_binding = replace(session._model_binding, lane=lane)
|
||||
bound_client = object()
|
||||
lane.client.with_options.return_value = bound_client
|
||||
mint_in_flight: list[int] = []
|
||||
|
||||
def _resolve_for_principal(alias, config, *, principal_id):
|
||||
assert alias == "task-gateway"
|
||||
assert config is auth_config
|
||||
assert session._acting_user_id == "user-b"
|
||||
assert principal_id == "user-a"
|
||||
mint_in_flight.append(gate.snapshot().in_flight)
|
||||
return "token-for-user-a"
|
||||
|
||||
pinned_resolver = MagicMock(side_effect=_resolve_for_principal)
|
||||
session._model_backend_auth_token_for_principal = pinned_resolver
|
||||
stream = provider.create_streaming.return_value
|
||||
|
||||
def _dispatch(**kwargs):
|
||||
assert gate.snapshot().in_flight == 1
|
||||
assert kwargs["client"] is bound_client
|
||||
return stream
|
||||
|
||||
provider.create_streaming.side_effect = _dispatch
|
||||
|
||||
with patch("turnstone.core.session.model_turn", wraps=real_model_turn) as plant_call:
|
||||
result = session._run_agent(
|
||||
[Turn.user("finish the task")],
|
||||
tools=[],
|
||||
auto_tools=set(),
|
||||
principal_id="user-a",
|
||||
)
|
||||
|
||||
assert result == "done"
|
||||
assert plant_call.call_count == 1
|
||||
called_lane = plant_call.call_args.args[0]
|
||||
assert called_lane.admission is gate
|
||||
assert called_lane.backend_auth_resolver is not None
|
||||
assert "backend_auth_token" not in plant_call.call_args.kwargs
|
||||
pinned_resolver.assert_called_once_with(
|
||||
lane.alias,
|
||||
auth_config,
|
||||
principal_id="user-a",
|
||||
)
|
||||
stale_live_resolver.assert_not_called()
|
||||
lane.client.with_options.assert_called_once_with(api_key="token-for-user-a")
|
||||
assert mint_in_flight == [1]
|
||||
assert gate.snapshot().in_flight == 0
|
||||
|
||||
|
||||
def test_already_cancelled_task_agent_does_not_resolve_backend_auth():
|
||||
from turnstone.core.session import GenerationCancelled
|
||||
|
||||
|
||||
@@ -782,15 +782,13 @@ class TestPerceptionFallback:
|
||||
)
|
||||
assert part["type"] == "text"
|
||||
assert "DESCRIPTION" in part["text"]
|
||||
# the perception model was handed the rasterized pages, not the raw
|
||||
# PDF: the wire carries the prompt + a by-reference placeholder, and
|
||||
# the threaded resolver materializes the page parts at the translator.
|
||||
# The perception model was handed the rasterized pages, not the raw
|
||||
# PDF: model_turn expands the by-reference placeholder before taking
|
||||
# admission, then the provider receives inline pages and no resolver.
|
||||
sent = prov.create_streaming.call_args.kwargs["messages"][0]["content"]
|
||||
assert sent[0]["type"] == "text"
|
||||
assert sent[1]["attachment_id"] == "perception-input"
|
||||
resolver = prov.create_streaming.call_args.kwargs["resolve_attachments"]
|
||||
pages = resolver(["perception-input"])["perception-input"]
|
||||
assert [p["type"] for p in pages] == ["image_url", "image_url"]
|
||||
assert [p["type"] for p in sent[1:]] == ["image_url", "image_url"]
|
||||
assert prov.create_streaming.call_args.kwargs["resolve_attachments"] is None
|
||||
|
||||
def test_audio_perception_when_omni_and_no_stt(self, tmp_db, mock_openai_client):
|
||||
s = _make_session(mock_openai_client)
|
||||
|
||||
@@ -44,6 +44,8 @@
|
||||
# provider = "openai"
|
||||
# base_url = "http://localhost:8000/v1"
|
||||
# context_window = 8192
|
||||
# max_concurrency = 1 # Per-process generation cap for this alias; 0 = unlimited.
|
||||
# # For llama.cpp, start with its usable -np slot count.
|
||||
#
|
||||
# [models.local.capabilities]
|
||||
# supports_vision = false
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Literal
|
||||
from typing import Annotated, Any, Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
@@ -12,6 +12,7 @@ from pydantic import BaseModel, Field
|
||||
from pydantic.json_schema import SkipJsonSchema # noqa: TC002
|
||||
|
||||
from turnstone.api.server_schemas import CreateWorkstreamRequest, CreateWorkstreamResponse
|
||||
from turnstone.core.model_registry import MAX_MODEL_CONCURRENCY
|
||||
from turnstone.core.skill_kind import SkillKind
|
||||
from turnstone.core.skill_parser import MAX_SKILL_DESCRIPTION_LEN
|
||||
|
||||
@@ -1002,6 +1003,18 @@ class RegistryInstallRequest(BaseModel):
|
||||
# Admin: Model Definitions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
ModelMaxConcurrency: TypeAlias = Annotated[
|
||||
int,
|
||||
Field(
|
||||
strict=True,
|
||||
ge=0,
|
||||
le=MAX_MODEL_CONCURRENCY,
|
||||
description=(
|
||||
"Maximum concurrent model generations for this alias in one process; zero means unlimited."
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class ModelDefinitionInfo(BaseModel):
|
||||
definition_id: str
|
||||
@@ -1011,6 +1024,7 @@ class ModelDefinitionInfo(BaseModel):
|
||||
base_url: str = ""
|
||||
api_key: str = ""
|
||||
context_window: int = 32768
|
||||
max_concurrency: ModelMaxConcurrency = 0
|
||||
capabilities: str = "{}"
|
||||
enabled: bool = True
|
||||
temperature: float | None = None
|
||||
@@ -1059,6 +1073,7 @@ class CreateModelDefinitionRequest(BaseModel):
|
||||
base_url: str = ""
|
||||
api_key: str = ""
|
||||
context_window: int = 32768
|
||||
max_concurrency: ModelMaxConcurrency = 0
|
||||
capabilities: dict[str, Any] = Field(default_factory=dict)
|
||||
enabled: bool = True
|
||||
temperature: float | None = None
|
||||
@@ -1078,6 +1093,12 @@ class UpdateModelDefinitionRequest(BaseModel):
|
||||
base_url: str | None = None
|
||||
api_key: str | None = None
|
||||
context_window: int | None = None
|
||||
# The runtime update handler is presence-keyed. Keep null out of the
|
||||
# advertised union because explicit JSON null is refused; clients clear a
|
||||
# limit by sending the canonical unlimited value, zero.
|
||||
max_concurrency: ModelMaxConcurrency | SkipJsonSchema[None] = Field(
|
||||
default_factory=lambda: None
|
||||
)
|
||||
# SkipJsonSchema drops the null member from the ADVERTISED union while the
|
||||
# Python type still tolerates None: the presence-keyed update handler
|
||||
# refuses an explicit JSON null, so advertising null would let generated
|
||||
|
||||
@@ -64,6 +64,7 @@ from turnstone.core.metacognition import field_str, sanitize_display
|
||||
from turnstone.core.model_registry import (
|
||||
APP_IDENTITY_MODEL_AUTH_MODES,
|
||||
DYNAMIC_MODEL_AUTH_MODES,
|
||||
MAX_MODEL_CONCURRENCY,
|
||||
MODEL_AUTH_MODE_PROFILES,
|
||||
MODEL_AUTH_TEXT_MAX_LEN,
|
||||
SCOPES_MODEL_AUTH_MODES,
|
||||
@@ -11964,10 +11965,12 @@ def _oidc_configured_for_model_auth(request: Request) -> bool:
|
||||
# "changing this column can neither redirect where a minted credential is
|
||||
# sent nor re-arm minting that a disable or revocation stopped": the four
|
||||
# sampling/shaping knobs hit the same endpoint with the same credential,
|
||||
# and the two reasoning toggles only select what history surfaces.
|
||||
# the admission knob only limits callers of that alias, and the two reasoning
|
||||
# toggles only select what history surfaces.
|
||||
MODEL_AUTH_NEUTRAL_FIELDS = frozenset(
|
||||
{
|
||||
"context_window",
|
||||
"max_concurrency",
|
||||
"temperature",
|
||||
"max_tokens",
|
||||
"reasoning_effort",
|
||||
@@ -11976,6 +11979,26 @@ MODEL_AUTH_NEUTRAL_FIELDS = frozenset(
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _parse_model_max_concurrency(raw: Any) -> tuple[int, JSONResponse | None]:
|
||||
"""Parse the strict model-definition concurrency scalar.
|
||||
|
||||
JSON booleans are integer subclasses in Python, so exact type equality is
|
||||
load-bearing. Strings and integral floats are also refused rather than
|
||||
silently normalized into a materially different admission policy.
|
||||
"""
|
||||
if type(raw) is not int or raw < 0 or raw > MAX_MODEL_CONCURRENCY:
|
||||
return 0, JSONResponse(
|
||||
{
|
||||
"error": (
|
||||
f"max_concurrency must be an integer between 0 and {MAX_MODEL_CONCURRENCY}"
|
||||
)
|
||||
},
|
||||
status_code=400,
|
||||
)
|
||||
return raw, None
|
||||
|
||||
|
||||
# The two cross-field refusal messages, shared verbatim by the create and
|
||||
# update twins (their guard CONDITIONS differ — raw body vs post-merge pair —
|
||||
# but the text must not drift). Mode lists come from the frozenset so a
|
||||
@@ -12728,6 +12751,7 @@ async def admin_list_model_definitions(request: Request) -> JSONResponse:
|
||||
"base_url": "",
|
||||
"api_key": "",
|
||||
"context_window": nm.get("context_window", 0),
|
||||
"max_concurrency": nm.get("max_concurrency", 0),
|
||||
"capabilities": "{}",
|
||||
"enabled": True,
|
||||
"temperature": nm.get("temperature"),
|
||||
@@ -12875,6 +12899,9 @@ async def admin_create_model_definition(request: Request) -> JSONResponse:
|
||||
if isinstance(ctx_raw, float) and not math.isfinite(ctx_raw):
|
||||
return JSONResponse({"error": "context_window must be a finite number"}, status_code=400)
|
||||
context_window = max(0, int(ctx_raw)) if isinstance(ctx_raw, (int, float)) else 0
|
||||
max_concurrency, concurrency_err = _parse_model_max_concurrency(body.get("max_concurrency", 0))
|
||||
if concurrency_err is not None:
|
||||
return concurrency_err
|
||||
caps = body.get("capabilities", {})
|
||||
if not isinstance(caps, dict):
|
||||
# Mirror of the update twin's refusal: coercing a null or the STRING
|
||||
@@ -12977,6 +13004,7 @@ async def admin_create_model_definition(request: Request) -> JSONResponse:
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
context_window=context_window,
|
||||
max_concurrency=max_concurrency,
|
||||
capabilities=capabilities,
|
||||
enabled=enabled,
|
||||
created_by=audit_uid,
|
||||
@@ -13120,6 +13148,11 @@ async def admin_update_model_definition(request: Request) -> JSONResponse:
|
||||
{"error": "context_window must be a finite number"}, status_code=400
|
||||
)
|
||||
updates["context_window"] = max(0, int(ctx_raw)) if isinstance(ctx_raw, (int, float)) else 0
|
||||
if "max_concurrency" in body:
|
||||
max_concurrency, concurrency_err = _parse_model_max_concurrency(body["max_concurrency"])
|
||||
if concurrency_err is not None:
|
||||
return concurrency_err
|
||||
updates["max_concurrency"] = max_concurrency
|
||||
if "capabilities" in body:
|
||||
caps = body["capabilities"]
|
||||
if not isinstance(caps, dict):
|
||||
|
||||
@@ -7525,8 +7525,9 @@ function _renderModels(items) {
|
||||
colAlias.appendChild(document.createTextNode(" "));
|
||||
colAlias.appendChild(defBadge);
|
||||
}
|
||||
// Per-model sampling override indicators
|
||||
// Non-default per-model setting indicators
|
||||
const overrides = [];
|
||||
if (m.max_concurrency > 0) overrides.push("limit=" + m.max_concurrency);
|
||||
if (m.temperature != null) overrides.push("temp=" + m.temperature);
|
||||
if (m.max_tokens != null) overrides.push("max_tok=" + m.max_tokens);
|
||||
if (m.reasoning_effort != null)
|
||||
@@ -7569,10 +7570,10 @@ function _renderModels(items) {
|
||||
const ovrSpan = document.createElement("span");
|
||||
ovrSpan.className = "model-overrides-hint";
|
||||
ovrSpan.textContent = overrides.join(", ");
|
||||
ovrSpan.title = "Per-model overrides (override global defaults)";
|
||||
ovrSpan.title = "Per-model settings";
|
||||
ovrSpan.setAttribute(
|
||||
"aria-label",
|
||||
"Per-model overrides: " + overrides.join(", "),
|
||||
"Per-model settings: " + overrides.join(", "),
|
||||
);
|
||||
colAlias.appendChild(document.createElement("br"));
|
||||
colAlias.appendChild(ovrSpan);
|
||||
@@ -7746,6 +7747,7 @@ function showCreateModelModal() {
|
||||
document.getElementById("model-api-key").value = "";
|
||||
document.getElementById("model-api-key").placeholder = "sk-...";
|
||||
document.getElementById("model-ctx-window").value = "0";
|
||||
document.getElementById("model-max-concurrency").value = "0";
|
||||
document.getElementById("model-temperature").value = "";
|
||||
document.getElementById("model-max-tokens").value = "";
|
||||
document.getElementById("model-reasoning-effort").value = "";
|
||||
@@ -7855,6 +7857,8 @@ function showEditModelModal(definitionId) {
|
||||
"\u2022\u2022\u2022 (leave blank to keep existing)";
|
||||
document.getElementById("model-ctx-window").value =
|
||||
m.context_window != null ? m.context_window : 0;
|
||||
document.getElementById("model-max-concurrency").value =
|
||||
m.max_concurrency != null ? m.max_concurrency : 0;
|
||||
document.getElementById("model-temperature").value =
|
||||
m.temperature != null ? m.temperature : "";
|
||||
document.getElementById("model-max-tokens").value =
|
||||
@@ -8142,6 +8146,24 @@ function submitCreateModel() {
|
||||
enabled: document.getElementById("model-enabled").checked,
|
||||
};
|
||||
|
||||
// Per-alias admission. Empty and zero are the canonical unlimited value;
|
||||
// every other spelling must be an exact non-negative integer.
|
||||
const concurrencyText = document
|
||||
.getElementById("model-max-concurrency")
|
||||
.value.trim();
|
||||
const maxConcurrency = concurrencyText === "" ? 0 : Number(concurrencyText);
|
||||
if (
|
||||
!Number.isInteger(maxConcurrency) ||
|
||||
maxConcurrency < 0 ||
|
||||
maxConcurrency > 2147483647
|
||||
) {
|
||||
_showModelError(
|
||||
"Max concurrent generations must be a whole number from 0 to 2147483647",
|
||||
);
|
||||
return;
|
||||
}
|
||||
form.max_concurrency = maxConcurrency;
|
||||
|
||||
// Per-model sampling overrides — null when empty (use global default)
|
||||
const tempVal = document.getElementById("model-temperature").value.trim();
|
||||
if (tempVal !== "") {
|
||||
|
||||
@@ -1695,6 +1695,23 @@
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
<div class="field-pair">
|
||||
<div>
|
||||
<label for="model-max-concurrency"
|
||||
>Max concurrent generations
|
||||
<span class="label-hint">0 = unlimited</span></label
|
||||
>
|
||||
<input
|
||||
type="number"
|
||||
id="model-max-concurrency"
|
||||
value="0"
|
||||
min="0"
|
||||
max="2147483647"
|
||||
step="1"
|
||||
title="Per-process limit for model generations using this alias"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
<!-- Hidden outright on a deployment that cannot mint any dynamic
|
||||
credential and whose row is not already using one; same
|
||||
treatment the Roles sub-tab gets when its scope is missing. -->
|
||||
|
||||
@@ -0,0 +1,210 @@
|
||||
"""Per-alias admission control for model dispatch.
|
||||
|
||||
The registry owns one :class:`ModelAdmission` object per alias and keeps that
|
||||
object stable across hot reloads. A zero limit is unlimited, but calls are
|
||||
still counted so a live ``0 -> N`` resize can drain already-running work before
|
||||
admitting more.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import threading
|
||||
import time
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from turnstone.core.deadline import DeadlineCancelledError
|
||||
from turnstone.core.log import get_logger
|
||||
|
||||
log = get_logger(__name__)
|
||||
|
||||
_CANCEL_POLL_SECONDS = 0.25
|
||||
_STALL_WARNING_SECONDS = 5.0
|
||||
|
||||
|
||||
def _call_hook(target: Any, name: str) -> None:
|
||||
hook = getattr(target, name, None)
|
||||
if callable(hook):
|
||||
with contextlib.suppress(Exception):
|
||||
hook()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AdmissionSnapshot:
|
||||
"""Non-sensitive instantaneous state for logs and tests."""
|
||||
|
||||
alias: str
|
||||
limit: int
|
||||
in_flight: int
|
||||
queued: int
|
||||
|
||||
|
||||
class AdmissionLease:
|
||||
"""One idempotently releasable admission hold."""
|
||||
|
||||
__slots__ = ("_gate", "wait_seconds", "queued_ahead", "_released")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
gate: ModelAdmission,
|
||||
*,
|
||||
wait_seconds: float,
|
||||
queued_ahead: int,
|
||||
) -> None:
|
||||
self._gate = gate
|
||||
self.wait_seconds = wait_seconds
|
||||
self.queued_ahead = queued_ahead
|
||||
self._released = False
|
||||
|
||||
def release(self) -> None:
|
||||
if self._released:
|
||||
return
|
||||
self._released = True
|
||||
self._gate._release()
|
||||
|
||||
def __enter__(self) -> AdmissionLease:
|
||||
return self
|
||||
|
||||
def __exit__(self, *_exc: object) -> None:
|
||||
self.release()
|
||||
|
||||
|
||||
class ModelAdmission:
|
||||
"""FIFO, hot-resizable in-flight gate for one model alias.
|
||||
|
||||
``limit == 0`` preserves unlimited behavior. Unlimited calls still take
|
||||
leases and increment ``in_flight`` so narrowing the gate during a hot reload
|
||||
observes and drains them rather than starting from a false zero.
|
||||
"""
|
||||
|
||||
def __init__(self, alias: str, limit: int = 0) -> None:
|
||||
if type(limit) is not int or limit < 0:
|
||||
raise ValueError("model admission limit must be a non-negative integer")
|
||||
self.alias = alias
|
||||
self._limit = limit
|
||||
self._in_flight = 0
|
||||
self._waiters: deque[object] = deque()
|
||||
self._cv = threading.Condition()
|
||||
|
||||
@property
|
||||
def limit(self) -> int:
|
||||
with self._cv:
|
||||
return self._limit
|
||||
|
||||
def snapshot(self) -> AdmissionSnapshot:
|
||||
with self._cv:
|
||||
return AdmissionSnapshot(
|
||||
alias=self.alias,
|
||||
limit=self._limit,
|
||||
in_flight=self._in_flight,
|
||||
queued=len(self._waiters),
|
||||
)
|
||||
|
||||
def set_limit(self, limit: int) -> None:
|
||||
"""Resize in place; narrowing lets current holders drain naturally."""
|
||||
if type(limit) is not int or limit < 0:
|
||||
raise ValueError("model admission limit must be a non-negative integer")
|
||||
with self._cv:
|
||||
if limit == self._limit:
|
||||
return
|
||||
previous = self._limit
|
||||
self._limit = limit
|
||||
self._cv.notify_all()
|
||||
log.info(
|
||||
"model.admission_resized",
|
||||
alias=self.alias,
|
||||
previous_limit=previous,
|
||||
limit=limit,
|
||||
in_flight=self._in_flight,
|
||||
queued=len(self._waiters),
|
||||
)
|
||||
|
||||
def _available(self) -> bool:
|
||||
return self._limit == 0 or self._in_flight < self._limit
|
||||
|
||||
def acquire(self, *, cancel_ref: Any = None) -> AdmissionLease:
|
||||
"""Wait FIFO for a slot, abandoning promptly when *cancel_ref* aborts."""
|
||||
if bool(getattr(cancel_ref, "aborted", False)):
|
||||
raise DeadlineCancelledError("cancel_ref aborted during model admission")
|
||||
|
||||
started = time.monotonic()
|
||||
ticket = object()
|
||||
queued_ahead = 0
|
||||
waiting = False
|
||||
did_wait = False
|
||||
warned = False
|
||||
try:
|
||||
with self._cv:
|
||||
if self._waiters or not self._available():
|
||||
queued_ahead = len(self._waiters)
|
||||
self._waiters.append(ticket)
|
||||
waiting = True
|
||||
did_wait = True
|
||||
_call_hook(cancel_ref, "begin_admission_wait")
|
||||
|
||||
while waiting:
|
||||
if bool(getattr(cancel_ref, "aborted", False)):
|
||||
with contextlib.suppress(ValueError):
|
||||
self._waiters.remove(ticket)
|
||||
self._cv.notify_all()
|
||||
raise DeadlineCancelledError("cancel_ref aborted during model admission")
|
||||
if self._waiters[0] is ticket and self._available():
|
||||
self._waiters.popleft()
|
||||
waiting = False
|
||||
break
|
||||
waited = time.monotonic() - started
|
||||
if not warned and waited >= _STALL_WARNING_SECONDS:
|
||||
warned = True
|
||||
log.warning(
|
||||
"model.admission_stalled",
|
||||
alias=self.alias,
|
||||
limit=self._limit,
|
||||
in_flight=self._in_flight,
|
||||
queued=len(self._waiters),
|
||||
wait_seconds=round(waited, 3),
|
||||
)
|
||||
self._cv.wait(_CANCEL_POLL_SECONDS)
|
||||
|
||||
self._in_flight += 1
|
||||
# When the limit is wider than one, let the new FIFO head claim
|
||||
# the next already-free slot without waiting for this holder to
|
||||
# release first.
|
||||
self._cv.notify_all()
|
||||
finally:
|
||||
if waiting:
|
||||
# An unexpected hook/condition failure must not strand its FIFO
|
||||
# ticket. The normal cancellation path already removed it.
|
||||
with self._cv:
|
||||
with contextlib.suppress(ValueError):
|
||||
self._waiters.remove(ticket)
|
||||
self._cv.notify_all()
|
||||
if did_wait:
|
||||
_call_hook(cancel_ref, "end_admission_wait")
|
||||
|
||||
waited = time.monotonic() - started if did_wait else 0.0
|
||||
lease = AdmissionLease(
|
||||
self,
|
||||
wait_seconds=waited,
|
||||
queued_ahead=queued_ahead,
|
||||
)
|
||||
if bool(getattr(cancel_ref, "aborted", False)):
|
||||
lease.release()
|
||||
raise DeadlineCancelledError("cancel_ref aborted during model admission")
|
||||
if waited > 0:
|
||||
log.info(
|
||||
"model.admission_wait",
|
||||
alias=self.alias,
|
||||
limit=self.limit,
|
||||
queued_ahead=queued_ahead,
|
||||
wait_seconds=round(waited, 3),
|
||||
)
|
||||
return lease
|
||||
|
||||
def _release(self) -> None:
|
||||
with self._cv:
|
||||
if self._in_flight <= 0:
|
||||
raise RuntimeError("model admission lease released without a holder")
|
||||
self._in_flight -= 1
|
||||
self._cv.notify_all()
|
||||
@@ -57,12 +57,25 @@ class StreamAbortRef(list[Any]):
|
||||
fix here must be mirrored there.
|
||||
"""
|
||||
|
||||
__slots__ = ("_aborted", "_cancel_event")
|
||||
__slots__ = (
|
||||
"_aborted",
|
||||
"_cancel_event",
|
||||
"_timing_lock",
|
||||
"_admission_wait_started",
|
||||
"_admission_wait_credit",
|
||||
"_dispatch_count",
|
||||
"_last_dispatch_at",
|
||||
)
|
||||
|
||||
def __init__(self, cancel_event: threading.Event | None = None) -> None:
|
||||
super().__init__()
|
||||
self._aborted = False
|
||||
self._cancel_event = cancel_event
|
||||
self._timing_lock = threading.Lock()
|
||||
self._admission_wait_started: float | None = None
|
||||
self._admission_wait_credit = 0.0
|
||||
self._dispatch_count = 0
|
||||
self._last_dispatch_at: float | None = None
|
||||
|
||||
def append(self, stream: Any) -> None:
|
||||
super().append(stream)
|
||||
@@ -77,6 +90,49 @@ class StreamAbortRef(list[Any]):
|
||||
with contextlib.suppress(Exception):
|
||||
stream.close()
|
||||
|
||||
def begin_admission_wait(self) -> None:
|
||||
"""Freeze this call's deadline while it waits for model admission."""
|
||||
with self._timing_lock:
|
||||
if self._admission_wait_started is None:
|
||||
self._admission_wait_started = time.monotonic()
|
||||
|
||||
def end_admission_wait(self) -> None:
|
||||
"""Resume the deadline and retain the elapsed admission credit."""
|
||||
now = time.monotonic()
|
||||
with self._timing_lock:
|
||||
started = self._admission_wait_started
|
||||
if started is None:
|
||||
return
|
||||
self._admission_wait_credit += max(0.0, now - started)
|
||||
self._admission_wait_started = None
|
||||
|
||||
def admission_wait_credit(self) -> float:
|
||||
"""Return completed plus currently accruing admission-wait time."""
|
||||
now = time.monotonic()
|
||||
with self._timing_lock:
|
||||
credit = self._admission_wait_credit
|
||||
if self._admission_wait_started is not None:
|
||||
credit += max(0.0, now - self._admission_wait_started)
|
||||
return credit
|
||||
|
||||
def mark_dispatch(self) -> None:
|
||||
"""Record a provider dispatch without changing deadline semantics."""
|
||||
with self._timing_lock:
|
||||
self._dispatch_count += 1
|
||||
self._last_dispatch_at = time.monotonic()
|
||||
|
||||
@property
|
||||
def dispatch_count(self) -> int:
|
||||
"""Number of provider dispatch attempts observed by this ref."""
|
||||
with self._timing_lock:
|
||||
return self._dispatch_count
|
||||
|
||||
@property
|
||||
def last_dispatch_at(self) -> float | None:
|
||||
"""Monotonic timestamp of the latest provider dispatch, if any."""
|
||||
with self._timing_lock:
|
||||
return self._last_dispatch_at
|
||||
|
||||
@property
|
||||
def aborted(self) -> bool:
|
||||
"""Whether :meth:`abort` has fired.
|
||||
@@ -123,6 +179,7 @@ def run_abortable_with_deadline(
|
||||
poll=poll,
|
||||
thread_name=thread_name,
|
||||
on_abandon=abort_ref.abort,
|
||||
deadline_credit=abort_ref.admission_wait_credit,
|
||||
)
|
||||
|
||||
|
||||
@@ -134,6 +191,7 @@ def run_with_deadline(
|
||||
poll: float = 1.0,
|
||||
thread_name: str = "deadline-worker",
|
||||
on_abandon: Callable[[], None] | None = None,
|
||||
deadline_credit: Callable[[], float] | None = None,
|
||||
) -> _T:
|
||||
"""Run ``fn()`` on a daemon thread, bounded by ``timeout``/``cancel_event``.
|
||||
|
||||
@@ -153,6 +211,11 @@ def run_with_deadline(
|
||||
|
||||
``poll`` bounds how often ``cancel_event`` is checked (and thus the worst-
|
||||
case latency from a cancel to this function returning).
|
||||
|
||||
``deadline_credit`` may dynamically extend the original deadline. The
|
||||
abortable wrapper uses it only for time spent queued at model admission;
|
||||
lowering, attachment materialization, credential minting, dispatch,
|
||||
draining, and retry backoff continue to consume the original budget.
|
||||
"""
|
||||
box: queue.Queue[tuple[bool, object]] = queue.Queue(maxsize=1)
|
||||
|
||||
@@ -186,7 +249,11 @@ def run_with_deadline(
|
||||
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
_abandon(DeadlineCancelledError())
|
||||
remaining = deadline - time.monotonic()
|
||||
credit = 0.0
|
||||
if deadline_credit is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
credit = max(0.0, float(deadline_credit()))
|
||||
remaining = deadline + credit - time.monotonic()
|
||||
if remaining <= 0:
|
||||
_abandon(DeadlineExceededError())
|
||||
try:
|
||||
|
||||
+27
-27
@@ -25,6 +25,7 @@ from turnstone.core.deadline import (
|
||||
run_abortable_with_deadline,
|
||||
)
|
||||
from turnstone.core.log import get_logger
|
||||
from turnstone.core.model_backend_auth import BackendAuthUnavailableError
|
||||
from turnstone.core.model_registry import ModelClientConstructionError
|
||||
from turnstone.core.model_turn import (
|
||||
ModelLane,
|
||||
@@ -1318,9 +1319,10 @@ class IntentJudge:
|
||||
live set (parallel task agents each spawn a generation;
|
||||
``close()`` aborts whatever is still live).
|
||||
backend_auth_resolver: Batch-scoped resolver whose closure pins
|
||||
the initiating principal. It is invoked once by the daemon,
|
||||
after its initial cancellation check, and the resulting token
|
||||
is reused for every item and evidence turn in this batch.
|
||||
the initiating principal. The resolver remains on the
|
||||
batch's lane so each plant-call attempt mints after acquiring
|
||||
alias admission. The mint cache keeps this inexpensive while
|
||||
preventing a queued bearer from expiring before dispatch.
|
||||
|
||||
Returns:
|
||||
List of heuristic verdicts (one per item), available immediately.
|
||||
@@ -1398,28 +1400,17 @@ class IntentJudge:
|
||||
)
|
||||
return
|
||||
|
||||
# Resolve delegated credentials exactly once for the batch. The
|
||||
# caller-supplied closure has already captured the initiating
|
||||
# principal, so a later shared-workstream handoff cannot mint a
|
||||
# successor user's token for this payload.
|
||||
backend_auth_token: str | None = None
|
||||
# Preserve the caller-pinned principal on the batch lane, but let
|
||||
# model_turn resolve credentials only after alias admission. An
|
||||
# admission backlog can outlive a bearer token; carrying the
|
||||
# resolver refreshes near-expiry tokens at the dispatch boundary
|
||||
# without allowing a later shared-workstream actor to take over.
|
||||
batch_lane = self._lane
|
||||
if backend_auth_resolver is not None:
|
||||
try:
|
||||
backend_auth_token = backend_auth_resolver(
|
||||
self._lane.alias,
|
||||
self._lane.backend_auth_config,
|
||||
)
|
||||
except Exception:
|
||||
log.exception("Judge backend authentication failed")
|
||||
self._deliver_fallbacks(
|
||||
items,
|
||||
heuristic_verdicts,
|
||||
callback,
|
||||
"judge backend authentication failed",
|
||||
)
|
||||
return
|
||||
batch_lane = replace(batch_lane, backend_auth_resolver=None)
|
||||
batch_lane = replace(
|
||||
batch_lane,
|
||||
backend_auth_resolver=backend_auth_resolver,
|
||||
)
|
||||
|
||||
if cancel_event and cancel_event.is_set():
|
||||
self._deliver_fallbacks(
|
||||
@@ -1434,7 +1425,8 @@ class IntentJudge:
|
||||
# One lane derivative for the whole batch: only the judge-owned
|
||||
# fresh client differs from the immutable constructor binding.
|
||||
# Every item and evidence turn therefore stays on one provider,
|
||||
# model, capability, config, and credential snapshot.
|
||||
# model, capability, auth-config, and initiating-principal
|
||||
# binding while credentials refresh per admitted attempt.
|
||||
batch_lane = replace(batch_lane, client=client)
|
||||
for idx, (item, h_verdict) in enumerate(zip(items, heuristic_verdicts, strict=True)):
|
||||
if cancel_event and cancel_event.is_set():
|
||||
@@ -1453,7 +1445,6 @@ class IntentJudge:
|
||||
cancel_event,
|
||||
client,
|
||||
lane=batch_lane,
|
||||
backend_auth_token=backend_auth_token,
|
||||
)
|
||||
if llm_verdict:
|
||||
log.info(
|
||||
@@ -1497,6 +1488,15 @@ class IntentJudge:
|
||||
"judge cancelled before evaluating this call",
|
||||
)
|
||||
return
|
||||
except BackendAuthUnavailableError:
|
||||
log.exception("Judge backend authentication failed")
|
||||
self._deliver_fallbacks(
|
||||
items[idx:],
|
||||
heuristic_verdicts[idx:],
|
||||
callback,
|
||||
"judge backend authentication failed",
|
||||
)
|
||||
return
|
||||
except Exception:
|
||||
log.exception(
|
||||
"Judge evaluation failed for %s",
|
||||
@@ -1549,7 +1549,6 @@ class IntentJudge:
|
||||
client: Any | None,
|
||||
*,
|
||||
lane: ModelLane | None = None,
|
||||
backend_auth_token: str | None = None,
|
||||
) -> IntentVerdict | None:
|
||||
"""Run LLM judge for a single tool call. Returns verdict or None."""
|
||||
if lane is None:
|
||||
@@ -1654,7 +1653,6 @@ class IntentJudge:
|
||||
tools=_tools,
|
||||
max_tokens=2048,
|
||||
cancel_ref=ref,
|
||||
backend_auth_token=backend_auth_token,
|
||||
)
|
||||
|
||||
result = run_abortable_with_deadline(
|
||||
@@ -1681,6 +1679,8 @@ class IntentJudge:
|
||||
log.info("judge.verdict.from_partial", turn=turn + 1)
|
||||
return verdict
|
||||
return None
|
||||
except BackendAuthUnavailableError:
|
||||
raise
|
||||
except Exception as e:
|
||||
log.info("judge.turn.failed", turn=turn + 1, error=str(e))
|
||||
return None
|
||||
|
||||
@@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Any
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Mapping
|
||||
|
||||
from turnstone.core.admission import ModelAdmission
|
||||
from turnstone.core.config import load_config
|
||||
from turnstone.core.log import get_logger
|
||||
from turnstone.core.providers import LLMProvider, create_client, create_provider
|
||||
@@ -29,6 +30,12 @@ MODEL_AUTH_MODES = frozenset({"static", "entra_obo", "entra_app", "rfc8693_obo"}
|
||||
# drift into "console stores it, registry refuses to load it".
|
||||
MODEL_AUTH_TEXT_MAX_LEN = 2048
|
||||
|
||||
# ``model_definitions.max_concurrency`` is an ``INTEGER`` on both supported
|
||||
# databases. Keep the public/config/API bound aligned with PostgreSQL's
|
||||
# signed 32-bit representation so a value accepted on SQLite cannot fail when
|
||||
# the same definition is moved to PostgreSQL.
|
||||
MAX_MODEL_CONCURRENCY = 2_147_483_647
|
||||
|
||||
# Derived, never hand-listed (fail-safe defaults): any mode later added to
|
||||
# MODEL_AUTH_MODES lands in this set BY CONSTRUCTION unless it is literally
|
||||
# "static", so membership tests fail CLOSED for modes nobody classified.
|
||||
@@ -80,6 +87,10 @@ class ModelAuthConfigError(ValueError):
|
||||
"""A model definition contains unsafe or internally inconsistent auth settings."""
|
||||
|
||||
|
||||
class ModelConcurrencyConfigError(ValueError):
|
||||
"""A model definition has an invalid per-alias concurrency limit."""
|
||||
|
||||
|
||||
class ModelClientConstructionError(ValueError):
|
||||
"""A registry alias exists but its binding could not be constructed.
|
||||
|
||||
@@ -263,6 +274,24 @@ class ModelConfig:
|
||||
auth_mode: str = "static"
|
||||
obo_audience: str = ""
|
||||
obo_scopes: str = ""
|
||||
# Per-process admission limit for this alias. Zero preserves the
|
||||
# historical unlimited behavior. Operational admission changes do not
|
||||
# change binding identity, hence ``compare=False``.
|
||||
max_concurrency: int = field(default=0, compare=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# ``bool`` is an ``int`` subclass, so use exact type equality. A
|
||||
# permissive coercion here could turn ``true`` into a one-request cap
|
||||
# or a typo into unlimited operation.
|
||||
if (
|
||||
type(self.max_concurrency) is not int
|
||||
or self.max_concurrency < 0
|
||||
or self.max_concurrency > MAX_MODEL_CONCURRENCY
|
||||
):
|
||||
raise ModelConcurrencyConfigError(
|
||||
f"Model {self.alias!r} max_concurrency must be an integer "
|
||||
f"between 0 and {MAX_MODEL_CONCURRENCY}"
|
||||
)
|
||||
|
||||
|
||||
def strip_control_characters(value: str) -> str:
|
||||
@@ -453,6 +482,10 @@ class ModelRegistry:
|
||||
self.task_effort = task_effort
|
||||
self._clients: dict[str, Any] = {}
|
||||
self._providers: dict[str, LLMProvider] = {}
|
||||
self._admissions = {
|
||||
alias: ModelAdmission(alias, int(getattr(cfg, "max_concurrency", 0)))
|
||||
for alias, cfg in self._models.items()
|
||||
}
|
||||
self._client_lock = threading.Lock()
|
||||
# Monotone count of completed reload() swaps. A counter, not a field
|
||||
# diff: sessions re-resolve on ANY difference at the next send, so an
|
||||
@@ -536,6 +569,14 @@ class ModelRegistry:
|
||||
raise UnknownModelAliasError(alias)
|
||||
return self._models[alias]
|
||||
|
||||
def get_admission(self, alias: str) -> ModelAdmission:
|
||||
"""Return the stable per-alias admission gate."""
|
||||
with self._client_lock:
|
||||
gate = self._admissions.get(alias)
|
||||
if gate is None:
|
||||
raise UnknownModelAliasError(alias)
|
||||
return gate
|
||||
|
||||
def has_alias(self, alias: str) -> bool:
|
||||
"""Check if *alias* exists in the registry."""
|
||||
return alias in self._models
|
||||
@@ -565,8 +606,8 @@ class ModelRegistry:
|
||||
|
||||
def resolve_binding(
|
||||
self, alias: str | None = None
|
||||
) -> tuple[Any, str, ModelConfig, LLMProvider, int]:
|
||||
"""Resolve *alias* to ``(client, model_name, config, provider, generation)``.
|
||||
) -> tuple[Any, str, ModelConfig, LLMProvider, ModelAdmission, int]:
|
||||
"""Resolve client, model, config, provider, admission, and generation.
|
||||
|
||||
The session bind primitive: everything a rebind commits, read under
|
||||
ONE lock acquisition, so a :meth:`reload` landing between separate
|
||||
@@ -592,7 +633,14 @@ class ModelRegistry:
|
||||
raise
|
||||
except ValueError as exc:
|
||||
raise ModelClientConstructionError(str(exc)) from exc
|
||||
return (client, cfg.model, cfg, provider, self._generation)
|
||||
return (
|
||||
client,
|
||||
cfg.model,
|
||||
cfg,
|
||||
provider,
|
||||
self._admissions[alias],
|
||||
self._generation,
|
||||
)
|
||||
|
||||
def resolve_agent_alias(self, kind: str) -> str | None:
|
||||
"""Return the configured alias for a sub-agent ``kind``.
|
||||
@@ -703,6 +751,20 @@ class ModelRegistry:
|
||||
self.agent_model = agent_model
|
||||
self.task_model = task_model
|
||||
self.task_effort = task_effort
|
||||
# Admission is strictly per alias. Resize surviving gates in
|
||||
# place so live lanes and new resolutions coordinate through the
|
||||
# same FIFO even when the alias moves to a different endpoint.
|
||||
for alias, cfg in self._models.items():
|
||||
limit = int(getattr(cfg, "max_concurrency", 0))
|
||||
gate = self._admissions.get(alias)
|
||||
if gate is None:
|
||||
self._admissions[alias] = ModelAdmission(alias, limit)
|
||||
else:
|
||||
gate.set_limit(limit)
|
||||
# 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
|
||||
# split old and new work across two independent limits.
|
||||
# Selective teardown — close + drop only clients whose
|
||||
# construction/connection target changed (alias removed, or
|
||||
# base_url / api_key / provider / auth_mode differs). Keeps connection
|
||||
@@ -894,8 +956,9 @@ def load_model_registry(
|
||||
auth_mode=row_auth_mode,
|
||||
obo_audience=row_obo_audience,
|
||||
obo_scopes=row_obo_scopes,
|
||||
max_concurrency=row.get("max_concurrency", 0),
|
||||
)
|
||||
except ModelAuthConfigError:
|
||||
except (ModelAuthConfigError, ModelConcurrencyConfigError):
|
||||
# Configuration errors are authoritative row content, not a
|
||||
# transient storage-read failure. Never degrade past them into a
|
||||
# config-only registry or provider SDK environment credentials.
|
||||
@@ -979,6 +1042,7 @@ def load_model_registry(
|
||||
auth_mode=entry_auth_mode,
|
||||
obo_audience=entry_obo_audience,
|
||||
obo_scopes=entry_obo_scopes,
|
||||
max_concurrency=entry.get("max_concurrency", 0),
|
||||
)
|
||||
|
||||
# 3. Back-compat shim: synthesize a "default" alias from CLI/auto-detected
|
||||
|
||||
+119
-73
@@ -38,6 +38,7 @@ Contract, held deliberately narrow:
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import random
|
||||
import time
|
||||
import uuid
|
||||
@@ -50,6 +51,7 @@ if TYPE_CHECKING:
|
||||
|
||||
from turnstone.core.model_registry import ModelConfig, ModelRegistry
|
||||
|
||||
from turnstone.core.admission import ModelAdmission
|
||||
from turnstone.core.deadline import DeadlineCancelledError
|
||||
from turnstone.core.history_decoration import attach_vllm_chat_reasoning_field
|
||||
from turnstone.core.log import get_logger
|
||||
@@ -103,7 +105,13 @@ from turnstone.core.storage._utils import (
|
||||
_CLIENT_TOOL_CALL_BLOCK_TYPES,
|
||||
strip_orphan_client_tool_blocks,
|
||||
)
|
||||
from turnstone.core.trajectory import ProviderNative, ToolCall, Turn, dicts_from_turns
|
||||
from turnstone.core.trajectory import (
|
||||
ProviderNative,
|
||||
ToolCall,
|
||||
Turn,
|
||||
dicts_from_turns,
|
||||
materialize_attachments,
|
||||
)
|
||||
|
||||
log = get_logger(__name__)
|
||||
|
||||
@@ -481,6 +489,9 @@ class ModelLane:
|
||||
# deployment-wide model.auth_fail_closed switch remains a live per-mint
|
||||
# policy read.
|
||||
backend_auth_config: ModelConfig | None = None
|
||||
# Stable per-alias registry gate. The gate object survives hot-resizes,
|
||||
# so old and newly resolved lanes coordinate through one FIFO.
|
||||
admission: ModelAdmission | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -554,6 +565,7 @@ def same_model_lane_binding(left: ModelLane, right: ModelLane) -> bool:
|
||||
and left.provider is right.provider
|
||||
and left.model == right.model
|
||||
and left.alias == right.alias
|
||||
and left.admission is right.admission
|
||||
)
|
||||
|
||||
|
||||
@@ -673,6 +685,7 @@ def resolve_lane(
|
||||
cfg: ModelConfig | None | EllipsisType = ...,
|
||||
config_store: Any | None = None,
|
||||
backend_auth_resolver: Callable[[str, ModelConfig | None], str | None] | None = None,
|
||||
admission: ModelAdmission | None = None,
|
||||
) -> ModelLane:
|
||||
"""Build a :class:`ModelLane`, resolving what the caller didn't supply.
|
||||
|
||||
@@ -705,6 +718,16 @@ def resolve_lane(
|
||||
if extra_params is ...
|
||||
else extra_params
|
||||
)
|
||||
if admission is None and registry is not None and alias:
|
||||
# Direct registry-backed lane rebuilds (notably judge sampling lanes)
|
||||
# still join the alias's stable gate. Duck-typed test registries may
|
||||
# manufacture attributes dynamically, hence the concrete type check.
|
||||
try:
|
||||
candidate = registry.get_admission(alias)
|
||||
except Exception:
|
||||
candidate = None
|
||||
if isinstance(candidate, ModelAdmission):
|
||||
admission = candidate
|
||||
return ModelLane(
|
||||
provider=provider,
|
||||
client=client,
|
||||
@@ -717,6 +740,7 @@ def resolve_lane(
|
||||
reasoning_effort=resolve_effort_setting(resolved_cfg, config_store),
|
||||
backend_auth_resolver=backend_auth_resolver,
|
||||
backend_auth_config=resolved_cfg,
|
||||
admission=admission,
|
||||
)
|
||||
|
||||
|
||||
@@ -733,7 +757,14 @@ def resolve_model_binding(
|
||||
# empty alias on a default binding disables registry-backed live flags and
|
||||
# delegated backend authentication on every later plant call.
|
||||
effective_alias = alias or registry.default
|
||||
client, model, cfg, provider, generation = registry.resolve_binding(effective_alias)
|
||||
resolved: Any = registry.resolve_binding(effective_alias)
|
||||
if len(resolved) == 6:
|
||||
client, model, cfg, provider, admission, generation = resolved
|
||||
else:
|
||||
# Compatibility for lightweight registry fakes that predate admission;
|
||||
# real ModelRegistry snapshots always take the six-value branch.
|
||||
client, model, cfg, provider, generation = resolved
|
||||
admission = None
|
||||
lane = resolve_lane(
|
||||
provider,
|
||||
client,
|
||||
@@ -743,6 +774,7 @@ def resolve_model_binding(
|
||||
cfg=cfg,
|
||||
config_store=config_store,
|
||||
backend_auth_resolver=backend_auth_resolver,
|
||||
admission=admission,
|
||||
)
|
||||
return ResolvedModelBinding(lane=lane, config=cfg, registry_generation=generation)
|
||||
|
||||
@@ -1087,11 +1119,11 @@ def model_turn(
|
||||
operator- or user-resolved knob (the session's own knobs, a CLI flag).
|
||||
|
||||
*resolve_attachments* materializes by-reference ``AttachmentRef``
|
||||
content at the provider translator (``{type: kind, attachment_id}``
|
||||
content immediately before admission (``{type: kind, attachment_id}``
|
||||
placeholders → inline parts; one id may expand to several parts, e.g.
|
||||
a rasterized PDF). Turn IR never carries inline media bytes — a lane
|
||||
with non-text content passes refs plus this resolver, exactly like the
|
||||
main loop's wire path.
|
||||
a rasterized PDF). This ordering keeps any nested perception/audio work
|
||||
outside the outer alias's gate, avoiding self-deadlock at a limit of one.
|
||||
Turn IR itself never carries inline media bytes.
|
||||
|
||||
*mint* rewrites each returned tool call's id (provider-original →
|
||||
caller-scoped) before the Turn is built; the native blocks keep the
|
||||
@@ -1190,7 +1222,8 @@ def model_turn(
|
||||
override its ``x-api-key``, so an injected header is silently dropped.
|
||||
``None`` leaves the lane's static client credential in place. When the
|
||||
explicit argument is absent, ``lane.backend_auth_resolver`` resolves it
|
||||
once before the drain-retry loop.
|
||||
after admission for each transport attempt, so a queued call cannot age a
|
||||
minted credential before it reaches the wire.
|
||||
"""
|
||||
if mint is not None and wire_id_map is None:
|
||||
raise ValueError(
|
||||
@@ -1235,81 +1268,94 @@ def model_turn(
|
||||
or (lane.capabilities.default_reasoning_effort if lane.capabilities else None)
|
||||
or None
|
||||
)
|
||||
# A Stop can land while deterministic wire preparation runs. Re-read the
|
||||
# abort signal before a possibly networked credential mint; the later read
|
||||
# remains necessary for cancellation that lands while the mint itself is
|
||||
# blocked.
|
||||
call_client = lane_call_client(
|
||||
lane,
|
||||
backend_auth_token=backend_auth_token,
|
||||
cancel_ref=cancel_ref,
|
||||
)
|
||||
# Materialization may perform storage reads and nested perception/audio
|
||||
# sampling. Complete it before taking the outer alias's admission slot so
|
||||
# a cap of one cannot deadlock on a nested call that needs the same alias.
|
||||
served_wire = materialize_attachments(wire, resolve_attachments)
|
||||
# A partially-surfaced stream is never silently re-issued — the
|
||||
# streaming caller owns re-issue.
|
||||
drain_retries = 0 if on_chunk is not None else _DRAIN_RETRIES
|
||||
attempt = 0
|
||||
request_metrics: list[ProviderRequestMetrics] = []
|
||||
while True:
|
||||
# Last read before the wire — it covers everything the entry read is
|
||||
# too early to see: the lowering, and on a delegated-auth alias the
|
||||
# credential resolve, which can block.
|
||||
_raise_if_aborted(cancel_ref, lane)
|
||||
# ``create_streaming`` stays OUTSIDE the try: every adapter issues
|
||||
# the HTTP request eagerly in its body (inside the SDK's own
|
||||
# request-level retry), so an exception from it is a request-time
|
||||
# failure that already got its retries; only drain-time failures
|
||||
# are mid-stream deaths the SDK could never see.
|
||||
# Dynamic backends bind the token as the client's api_key so the SDK
|
||||
# emits it as its own auth header; with_options reuses the pool.
|
||||
# (extra_headers can't override the Anthropic SDK's x-api-key.)
|
||||
chunks = lane.provider.create_streaming(
|
||||
client=call_client,
|
||||
model=lane.model,
|
||||
messages=wire,
|
||||
tools=tools,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature if temperature is not None else lane.temperature,
|
||||
reasoning_effort=effective_effort,
|
||||
extra_params=lane.extra_params,
|
||||
deferred_names=deferred_names,
|
||||
cancel_ref=cancel_ref,
|
||||
capabilities=lane.capabilities,
|
||||
replay_reasoning_to_model=resolve_replay_reasoning_to_model(
|
||||
lane.registry, lane.alias, caps=lane.capabilities, cfg=cfg
|
||||
),
|
||||
resolve_attachments=resolve_attachments,
|
||||
request_metrics_ref=request_metrics,
|
||||
)
|
||||
try:
|
||||
result = drain_stream(
|
||||
_tee_chunks(chunks, on_chunk) if on_chunk else chunks,
|
||||
scan_inline_reasoning=lane_scans_inline_reasoning(lane),
|
||||
lease = lane.admission.acquire(cancel_ref=cancel_ref) if lane.admission else None
|
||||
drain_error: Exception | None = None
|
||||
with lease if lease is not None else contextlib.nullcontext():
|
||||
# Admission precedes a dynamic credential mint. This work and the
|
||||
# full create+drain remain inside the hold; the context exits before
|
||||
# any retry backoff below.
|
||||
call_client = lane_call_client(
|
||||
lane,
|
||||
backend_auth_token=backend_auth_token,
|
||||
cancel_ref=cancel_ref,
|
||||
)
|
||||
break
|
||||
except Exception as exc:
|
||||
attempt += 1
|
||||
if (
|
||||
attempt > drain_retries
|
||||
or bool(getattr(cancel_ref, "aborted", False))
|
||||
or type(exc).__name__ not in lane.provider.retryable_error_names
|
||||
):
|
||||
raise
|
||||
delay = _DRAIN_RETRY_BASE_DELAY * (2 ** (attempt - 1)) * (0.5 + random.random())
|
||||
log.warning(
|
||||
"model_turn.drain_retry",
|
||||
error_type=type(exc).__name__,
|
||||
attempt=attempt,
|
||||
_raise_if_aborted(cancel_ref, lane)
|
||||
mark_dispatch = getattr(cancel_ref, "mark_dispatch", None)
|
||||
if callable(mark_dispatch):
|
||||
with contextlib.suppress(Exception):
|
||||
mark_dispatch()
|
||||
# ``create_streaming`` stays OUTSIDE the drain-error catch: every
|
||||
# adapter issues eagerly in its body, so a request-time failure has
|
||||
# already received the SDK's own retries and propagates unchanged.
|
||||
chunks = lane.provider.create_streaming(
|
||||
client=call_client,
|
||||
model=lane.model,
|
||||
retry_in=round(delay, 2),
|
||||
messages=served_wire,
|
||||
tools=tools,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature if temperature is not None else lane.temperature,
|
||||
reasoning_effort=effective_effort,
|
||||
extra_params=lane.extra_params,
|
||||
deferred_names=deferred_names,
|
||||
cancel_ref=cancel_ref,
|
||||
capabilities=lane.capabilities,
|
||||
replay_reasoning_to_model=resolve_replay_reasoning_to_model(
|
||||
lane.registry, lane.alias, caps=lane.capabilities, cfg=cfg
|
||||
),
|
||||
# Already materialized before admission; provider translators
|
||||
# retain their no-op fallback for direct callers.
|
||||
resolve_attachments=None,
|
||||
request_metrics_ref=request_metrics,
|
||||
)
|
||||
if delay > 0:
|
||||
time.sleep(delay)
|
||||
if bool(getattr(cancel_ref, "aborted", False)):
|
||||
# The deadline abandoned this worker while it was backing off.
|
||||
# The loop-top read would stop the re-issue anyway; this arm
|
||||
# exists to die with the ORIGINAL transport failure rather than
|
||||
# the abandonment error, so the cause of the death survives.
|
||||
raise
|
||||
try:
|
||||
result = drain_stream(
|
||||
_tee_chunks(chunks, on_chunk) if on_chunk else chunks,
|
||||
scan_inline_reasoning=lane_scans_inline_reasoning(lane),
|
||||
)
|
||||
except Exception as exc:
|
||||
drain_error = exc
|
||||
finally:
|
||||
close = getattr(chunks, "close", None)
|
||||
if callable(close):
|
||||
with contextlib.suppress(Exception):
|
||||
close()
|
||||
|
||||
if drain_error is None:
|
||||
break
|
||||
attempt += 1
|
||||
if (
|
||||
attempt > drain_retries
|
||||
or bool(getattr(cancel_ref, "aborted", False))
|
||||
or type(drain_error).__name__ not in lane.provider.retryable_error_names
|
||||
):
|
||||
raise drain_error
|
||||
delay = _DRAIN_RETRY_BASE_DELAY * (2 ** (attempt - 1)) * (0.5 + random.random())
|
||||
log.warning(
|
||||
"model_turn.drain_retry",
|
||||
error_type=type(drain_error).__name__,
|
||||
attempt=attempt,
|
||||
model=lane.model,
|
||||
retry_in=round(delay, 2),
|
||||
)
|
||||
if delay > 0:
|
||||
time.sleep(delay)
|
||||
if bool(getattr(cancel_ref, "aborted", False)):
|
||||
# The deadline abandoned this worker while it was backing off.
|
||||
# The loop-top read would stop the re-issue anyway; this arm
|
||||
# exists to die with the ORIGINAL transport failure rather than
|
||||
# the abandonment error, so the cause of the death survives.
|
||||
raise drain_error
|
||||
|
||||
raw_calls: list[dict[str, Any]] = list(result.tool_calls or [])
|
||||
# Record blanks BEFORE the uuid back-fill: a back-filled id exists only
|
||||
@@ -1369,7 +1415,7 @@ def model_turn(
|
||||
finish_reason=result.finish_reason,
|
||||
usage=result.usage,
|
||||
tool_calls=raw_calls,
|
||||
wire_msgs=wire,
|
||||
wire_msgs=served_wire,
|
||||
producer=lane.provider.provider_name,
|
||||
serving_model=lane.model,
|
||||
tool_def_chars=(
|
||||
|
||||
@@ -23,8 +23,9 @@ one alias; a vision-only model covers image/PDF and is simply skipped for audio.
|
||||
The call goes through :func:`turnstone.core.model_turn.model_turn` (the shared
|
||||
plant-call seam, #827), so any provider works: the trajectory carries the
|
||||
attachment by reference and the pre-built OpenAI-shaped parts (``image_url`` /
|
||||
``input_audio``) materialize at the provider translator via the
|
||||
``resolve_attachments`` callback, exactly like the main loop's wire path.
|
||||
``input_audio``) materialize in ``model_turn`` via the ``resolve_attachments``
|
||||
callback before the alias admission slot is acquired, exactly like the main loop's
|
||||
wire path.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -78,9 +79,9 @@ def describe(
|
||||
|
||||
``parts`` are pre-built OpenAI-shaped content parts — ``image_url`` for
|
||||
image/PDF-page perception, ``input_audio`` for audio. The trajectory
|
||||
carries them by reference; ``model_turn`` hands the resolver to the
|
||||
provider translator, which materializes the placeholder into these exact
|
||||
parts (one ref may expand to many, e.g. a rasterized PDF).
|
||||
carries them by reference; ``model_turn`` uses the resolver before model
|
||||
admission to materialize the placeholder into these exact parts (one ref
|
||||
may expand to many, e.g. a rasterized PDF).
|
||||
|
||||
``lane`` is the caller's already-resolved binding snapshot, so the
|
||||
modality gate and the plant call cannot observe different registry
|
||||
|
||||
+17
-26
@@ -979,8 +979,9 @@ def _neutralize_untrusted_fences(text: str) -> str:
|
||||
def _neutralize_attachment_part(part: Any) -> Any:
|
||||
"""Return an attachment part with textual trust-marker forgeries defanged.
|
||||
|
||||
Attachment placeholders are materialized after ordinary message folding,
|
||||
so their text cannot rely on ``fold_system_turns`` for this boundary. Only
|
||||
Attachment placeholders are materialized after ordinary message folding and
|
||||
before model admission, so their text cannot rely on ``fold_system_turns``
|
||||
for this boundary. Only
|
||||
model-visible text and document strings are inspected; binary image/audio
|
||||
data and base64 PDF payloads retain their exact bytes.
|
||||
"""
|
||||
@@ -5921,8 +5922,8 @@ class ChatSession:
|
||||
|
||||
``self.messages`` is the canonical ``Turn`` trajectory; lowering it here
|
||||
emits by-reference attachments as ``{type: kind, attachment_id}``
|
||||
placeholders. The provider translator materializes them to inline bytes
|
||||
via :meth:`_resolve_attachments` — resolution lives at the C layer."""
|
||||
placeholders. ``model_turn`` materializes them to inline wire parts via
|
||||
:meth:`_resolve_attachments` before acquiring the alias admission slot."""
|
||||
return self.system_messages + dicts_from_turns(self.messages)
|
||||
|
||||
def _resolve_attachments(
|
||||
@@ -5936,10 +5937,10 @@ class ChatSession:
|
||||
"""Resolve content-addressed attachment ids to inline wire content parts.
|
||||
|
||||
The send-time materialization of the by-reference content lane: handed to
|
||||
the provider translator, which calls it with the placeholder ids it finds
|
||||
and expands each to the inline part the wire needs. Blobs are
|
||||
batch-fetched from the content-addressed store; a pruned id resolves to
|
||||
nothing and the translator drops its placeholder.
|
||||
``model_turn``, which calls it with the placeholder ids it finds and
|
||||
expands each to the inline part the wire needs before model admission.
|
||||
Blobs are batch-fetched from the content-addressed store; a pruned id
|
||||
resolves to nothing and the shared materializer drops its placeholder.
|
||||
|
||||
Kinds the active model can't ingest natively (pdf without ``supports_pdf``,
|
||||
audio without ``supports_audio_input``) are converted client-side here —
|
||||
@@ -12672,8 +12673,8 @@ class ChatSession:
|
||||
|
||||
# A shared workstream may bind a new actor while this daemon is still
|
||||
# evaluating later items. Capture the initiating identity now; the
|
||||
# daemon resolves one token after its initial cancellation check and
|
||||
# reuses it for the whole batch.
|
||||
# daemon carries this resolver into each admitted plant call so token
|
||||
# refresh cannot switch to a successor principal.
|
||||
judge_principal = (
|
||||
(self._mcp_effective_user_id or "") if principal_id is None else principal_id
|
||||
).strip()
|
||||
@@ -20594,22 +20595,13 @@ class ChatSession:
|
||||
# invocation; the native lane carried here serves the WITHIN-RUN
|
||||
# reasoning continuity of the agent's own tool loop.
|
||||
same_lane = (lane.alias or "") == (primary_lane.alias or "")
|
||||
# Resolve once per sub-agent run, outside its request retry loop. A
|
||||
# fail-open ``None`` is an intentional static-client result, not an
|
||||
# invitation for ``model_turn`` to resolve again through the lane's
|
||||
# mutable session-principal callback after a shared-user handoff.
|
||||
# Keep the initiating principal pinned for the whole sub-agent run,
|
||||
# but defer each credential mint until model_turn has acquired this
|
||||
# alias's admission slot. A long queue can outlive a bearer token;
|
||||
# the principal-bound resolver refreshes at the dispatch boundary
|
||||
# without consulting a later shared-workstream actor.
|
||||
cancel_scope.check()
|
||||
try:
|
||||
agent_backend_auth_token = self._model_backend_auth_token_for_principal(
|
||||
lane.alias,
|
||||
lane.backend_auth_config,
|
||||
principal_id=agent_principal,
|
||||
)
|
||||
except Exception:
|
||||
cancel_scope.check()
|
||||
raise
|
||||
cancel_scope.check()
|
||||
lane = dataclasses.replace(lane, backend_auth_resolver=None)
|
||||
lane = self._lane_for_backend_auth_principal(lane, agent_principal)
|
||||
|
||||
def _api_call(
|
||||
turns: list[Turn],
|
||||
@@ -20645,7 +20637,6 @@ class ChatSession:
|
||||
or (self.reasoning_effort if same_lane else None),
|
||||
mint=mint,
|
||||
wire_id_map=wire_id_map,
|
||||
backend_auth_token=agent_backend_auth_token,
|
||||
cancel_ref=cancel_scope.cancel_ref,
|
||||
prepare_wire=lambda wire, serving_lane: self._prepare_lowered_wire_messages(
|
||||
wire,
|
||||
|
||||
@@ -5537,6 +5537,7 @@ class PostgreSQLBackend:
|
||||
auth_mode: str = "static",
|
||||
obo_audience: str = "",
|
||||
obo_scopes: str = "",
|
||||
max_concurrency: int = 0,
|
||||
) -> None:
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
@@ -5562,6 +5563,7 @@ class PostgreSQLBackend:
|
||||
auth_mode=auth_mode,
|
||||
obo_audience=obo_audience,
|
||||
obo_scopes=obo_scopes,
|
||||
max_concurrency=max_concurrency,
|
||||
created_by=created_by,
|
||||
created=now,
|
||||
updated=now,
|
||||
|
||||
@@ -2640,6 +2640,7 @@ class StorageBackend(Protocol):
|
||||
auth_mode: str = "static",
|
||||
obo_audience: str = "",
|
||||
obo_scopes: str = "",
|
||||
max_concurrency: int = 0,
|
||||
) -> None:
|
||||
"""Create a model definition. No-op if definition_id already exists."""
|
||||
...
|
||||
|
||||
@@ -833,6 +833,7 @@ model_definitions = sa.Table(
|
||||
sa.Column("base_url", sa.Text, nullable=False, server_default=""),
|
||||
sa.Column("api_key", sa.Text, nullable=False, server_default=""),
|
||||
sa.Column("context_window", sa.Integer, nullable=False, server_default="32768"),
|
||||
sa.Column("max_concurrency", sa.Integer, nullable=False, server_default="0"),
|
||||
sa.Column("capabilities", sa.Text, nullable=False, server_default="{}"),
|
||||
sa.Column("enabled", sa.Integer, nullable=False, server_default="1"),
|
||||
sa.Column("temperature", sa.Float, nullable=True),
|
||||
|
||||
@@ -5677,6 +5677,7 @@ class SQLiteBackend:
|
||||
auth_mode: str = "static",
|
||||
obo_audience: str = "",
|
||||
obo_scopes: str = "",
|
||||
max_concurrency: int = 0,
|
||||
) -> None:
|
||||
|
||||
now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S")
|
||||
@@ -5701,6 +5702,7 @@ class SQLiteBackend:
|
||||
"auth_mode": auth_mode,
|
||||
"obo_audience": obo_audience,
|
||||
"obo_scopes": obo_scopes,
|
||||
"max_concurrency": max_concurrency,
|
||||
"created_by": created_by,
|
||||
"created": now,
|
||||
"updated": now,
|
||||
|
||||
@@ -682,6 +682,7 @@ MODEL_DEFINITION_MUTABLE = frozenset(
|
||||
"base_url",
|
||||
"api_key",
|
||||
"context_window",
|
||||
"max_concurrency",
|
||||
"capabilities",
|
||||
"enabled",
|
||||
"temperature",
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
"""Add the per-alias model concurrency limit.
|
||||
|
||||
``max_concurrency`` bounds concurrent model generations for one alias in one
|
||||
process. Zero preserves the pre-feature unlimited behavior for every
|
||||
existing definition.
|
||||
|
||||
Revision ID: 070
|
||||
Revises: 069
|
||||
Create Date: 2026-08-08
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = "070"
|
||||
down_revision = "069"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column(
|
||||
"model_definitions",
|
||||
sa.Column("max_concurrency", sa.Integer, nullable=False, server_default="0"),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("model_definitions", "max_concurrency")
|
||||
@@ -12,7 +12,7 @@ Field set and rationale: ``docs/design/canonical-trajectory-ideal-target.md`` §
|
||||
NOTE: non-text content rides as ``AttachmentRef`` — a reference to a content-addressed
|
||||
blob in ``workstream_attachments``. ``Turn``s never carry bytes, and the dict bridge
|
||||
carries only the ``{type: kind, attachment_id}`` placeholder. Each output boundary (the
|
||||
provider translator, the ``/history`` display, export) materializes the placeholder to an
|
||||
model-dispatch path, the ``/history`` display, export) materializes the placeholder to an
|
||||
inline part by point-lookup against the blob store, via :func:`resolve_attachment_parts`.
|
||||
"""
|
||||
|
||||
@@ -68,9 +68,9 @@ class TextBlock:
|
||||
class AttachmentRef:
|
||||
"""A reference to attachment bytes held in the content-addressed blob store.
|
||||
|
||||
Non-text content is carried *by reference* (never inline bytes): the translator
|
||||
resolves ``attachment_id`` to bytes and expands it to the provider's native
|
||||
format at wire time. ``kind`` is the by-reference placeholder type —
|
||||
Non-text content is carried *by reference* (never inline bytes): the model-dispatch
|
||||
path resolves ``attachment_id`` to bytes and expands it to the provider-neutral
|
||||
inline wire part before provider admission. ``kind`` is the placeholder type —
|
||||
``"image"``, ``"document"`` (text docs), ``"pdf"``, or ``"audio"``. The
|
||||
dict-bridge keys off ``attachment_id`` and is kind-agnostic, so new kinds
|
||||
need no change here.
|
||||
@@ -437,10 +437,10 @@ def resolve_attachment_parts(
|
||||
``{type: kind, attachment_id}`` placeholders in a message's list content;
|
||||
*parts_by_id* maps an id to its inline content part — or a *list* of parts
|
||||
(one placeholder may expand to several, e.g. a PDF rasterized to one image
|
||||
per page for a vision model) — built from the content-addressed blob. This is the materialization the
|
||||
translator — and reconstruct, for display — runs at its output boundary: a
|
||||
placeholder whose blob is missing (pruned) is dropped, so a consumer never
|
||||
sees an unresolved reference. Identity-preserving when no message carries a
|
||||
per page for a vision model) — built from the content-addressed blob. This
|
||||
shared substitution runs at each output boundary (model dispatch or display):
|
||||
a placeholder whose blob is missing (pruned) is dropped, so a consumer never
|
||||
sees an unresolved reference. Identity-preserving when no message carries a
|
||||
placeholder; never mutates the input.
|
||||
"""
|
||||
|
||||
@@ -480,9 +480,9 @@ def materialize_attachments(
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Expand by-reference attachment placeholders to inline parts at the wire.
|
||||
|
||||
The translator's entry point for the by-reference content lane: collect the
|
||||
placeholder ids across *messages*, ask *resolve* (a storage point-lookup the
|
||||
session hands down) for their inline content parts, and substitute via
|
||||
The shared wire-boundary entry point for the by-reference content lane:
|
||||
collect the placeholder ids across *messages*, ask *resolve* (a storage
|
||||
point-lookup the session hands down) for their inline content parts, and substitute via
|
||||
:func:`resolve_attachment_parts`. A ``None`` resolver (no storage — e.g. a
|
||||
unit test or an in-memory sub-agent whose media is already inline) or a
|
||||
placeholder-free trajectory is a no-op, so the common path is allocation-free.
|
||||
|
||||
+10
-4
@@ -4484,6 +4484,7 @@ def internal_model_reload(request: Request) -> JSONResponse:
|
||||
from turnstone.core.model_registry import (
|
||||
DynamicAuthKeyError,
|
||||
ModelAuthConfigError,
|
||||
ModelConcurrencyConfigError,
|
||||
load_model_registry,
|
||||
)
|
||||
from turnstone.core.storage._registry import get_storage
|
||||
@@ -4503,10 +4504,10 @@ def internal_model_reload(request: Request) -> JSONResponse:
|
||||
provider=cli_args["provider"],
|
||||
storage=storage,
|
||||
)
|
||||
except ModelAuthConfigError as exc:
|
||||
# A row whose auth fields violate _normalize_auth_mode, reachable via
|
||||
# direct SQL, a migration mishap, or console version skew (console
|
||||
# writes are validated). The loader deliberately propagates it; this
|
||||
except (ModelAuthConfigError, ModelConcurrencyConfigError) as exc:
|
||||
# A row whose auth or concurrency fields violate registry validation,
|
||||
# reachable via direct SQL, a migration mishap, or console version skew
|
||||
# (console writes are validated). The loader deliberately propagates it; this
|
||||
# arm keeps the node from answering a bare 500 where the console's
|
||||
# refresh path catches the same row. Same structured 422 contract as
|
||||
# the bad-arguments arm below: the reason names the row and field.
|
||||
@@ -4530,8 +4531,12 @@ def internal_model_reload(request: Request) -> JSONResponse:
|
||||
|
||||
# No-op fast path: skip reload when nothing changed (avoids client churn
|
||||
# on broadcast model-reloads where this node has no pending changes).
|
||||
admission_unchanged = {
|
||||
alias: cfg.max_concurrency for alias, cfg in new_registry.models.items()
|
||||
} == {alias: cfg.max_concurrency for alias, cfg in registry.models.items()}
|
||||
unchanged = (
|
||||
new_registry.models == registry.models
|
||||
and admission_unchanged
|
||||
and new_registry.fallback == registry.fallback
|
||||
and new_registry.agent_model == registry.agent_model
|
||||
and eff_default == registry.default
|
||||
@@ -4624,6 +4629,7 @@ def internal_model_status(request: Request) -> JSONResponse:
|
||||
"provider": cfg.provider,
|
||||
"source": cfg.source,
|
||||
"context_window": cfg.context_window,
|
||||
"max_concurrency": cfg.max_concurrency,
|
||||
"enabled": True,
|
||||
"temperature": cfg.temperature,
|
||||
"max_tokens": cfg.max_tokens,
|
||||
|
||||
Reference in New Issue
Block a user