feat: rule routing v0 for model selection (shadow/active, fail-closed)
Pre-dispatch routing only fills in tasks with no explicit model, classified by structural signals only (plan > goal > default). Explicit task model and worker CLI --model win over policy; policy via CLOUD_ROUTING_POLICY_JSON/FILE + CLOUD_ROUTING_MODE. Shadow mode records a suggestion without changing dispatch. Only policy/shadow decisions emit a model_routing event, so the existing events contract (fallback tests) is unchanged. Fallback rules, effort/approval/sandbox/quota and provider are untouched. Tests: test_models_routing.py (20 tests).
This commit is contained in:
1 parent
e6e1e4b4f3
commit
e364970b46
3 files changed
+654
-3
No files matched your search
+280
-1
@@ -1,12 +1,24 @@
|
||||
"""Model fallback routing.
|
||||
"""Model fallback routing and pre-dispatch rule routing (v0).
|
||||
|
||||
Zero third-party dependencies; standard library only.
|
||||
|
||||
This module decides *whether* a failed model attempt may fall back to
|
||||
the next model in a priority list, and *which* model to try next. It
|
||||
never runs a model turn itself and never talks to any provider.
|
||||
|
||||
Rule routing (v0) is deliberately narrower than that: it only fills
|
||||
in the model for tasks that did not specify one, using structural
|
||||
signals already stored on the task (``plan``, ``goal_mode`` /
|
||||
``verify_command``). It never changes ``effort``, ``approval``,
|
||||
``plan``, ``executor`` or any quota, never switches provider/gateway,
|
||||
and never invents a quality ranking between models. A policy that
|
||||
cannot be parsed, or a tier with no model, falls back to the existing
|
||||
worker-default chain instead of failing the task.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
FALLBACKABLE = {404, 429, 500, 502, 503, 504, "timeout"}
|
||||
|
||||
|
||||
@@ -102,6 +114,273 @@ class ModelRouter:
|
||||
return self.next_after(last_model, attempts[-1][1])
|
||||
|
||||
|
||||
ROUTING_MODES = ("off", "shadow", "active")
|
||||
ROUTING_TIERS = ("plan", "goal", "default")
|
||||
|
||||
|
||||
def prompt_len_bucket(prompt) -> str:
|
||||
"""Bucket a prompt by length for logging only, never for routing.
|
||||
|
||||
Length is recorded so shadow reports can test whether it correlates
|
||||
with outcomes; v0 deliberately does not route on it (a long prompt
|
||||
may just be a pasted log, not a hard task).
|
||||
"""
|
||||
length = len(prompt) if isinstance(prompt, str) else 0
|
||||
if length < 500:
|
||||
return "<500"
|
||||
if length <= 3000:
|
||||
return "500-3000"
|
||||
return ">3000"
|
||||
|
||||
|
||||
def classify_task(task) -> str:
|
||||
"""Classify a task into one routing tier, first match wins.
|
||||
|
||||
Only structural signals are used, in this order: ``plan`` (read-only,
|
||||
never mutates files) before ``goal`` (multi-round goal loop, possibly
|
||||
unbounded rounds), everything else is ``default``. Keyword, prompt
|
||||
length, priority, ``depends_on`` count and executor are intentionally
|
||||
excluded: they are weak or unrelated proxies for difficulty.
|
||||
"""
|
||||
if not isinstance(task, dict):
|
||||
return "default"
|
||||
if task.get("plan"):
|
||||
return "plan"
|
||||
if task.get("goal_mode") or task.get("verify_command"):
|
||||
return "goal"
|
||||
return "default"
|
||||
|
||||
|
||||
def _clean_model_name(value):
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value.strip()
|
||||
return None
|
||||
|
||||
|
||||
def _clean_tiers(raw):
|
||||
"""Return {tier: model} keeping only valid non-empty entries."""
|
||||
tiers = {}
|
||||
if not isinstance(raw, dict):
|
||||
return tiers
|
||||
for tier in ROUTING_TIERS:
|
||||
model = _clean_model_name(raw.get(tier))
|
||||
if isinstance(model, str):
|
||||
# A tier entry may itself be a dict {"model": ...}; accept
|
||||
# that shape too so policy authors are not punished for it.
|
||||
tiers[tier] = model
|
||||
elif isinstance(raw.get(tier), dict):
|
||||
model = _clean_model_name(raw[tier].get("model"))
|
||||
if model:
|
||||
tiers[tier] = model
|
||||
return tiers
|
||||
|
||||
|
||||
class RoutingPolicy:
|
||||
"""A parsed, versioned routing policy (or an empty one).
|
||||
|
||||
``mode`` is one of ``off`` / ``shadow`` / ``active``. ``tiers`` maps
|
||||
a tier name to a model name; a tier absent from the map falls back
|
||||
to the worker default. ``fallback_models`` is parsed and validated
|
||||
but deliberately NOT applied by v0: the worker's explicit
|
||||
``--fallback-models`` remains the only source that can enable a real
|
||||
fallback retry, so a policy file alone can never start burning a
|
||||
second attempt's quota.
|
||||
"""
|
||||
|
||||
def __init__(self, version=None, mode="shadow", tiers=None,
|
||||
fallback_models=None):
|
||||
self.version = version if isinstance(version, str) \
|
||||
and version.strip() else None
|
||||
self.mode = mode if mode in ROUTING_MODES else "off"
|
||||
self.tiers = dict(tiers or {})
|
||||
self.fallback_models = []
|
||||
for item in fallback_models or []:
|
||||
model = _clean_model_name(item)
|
||||
if model and model not in self.fallback_models:
|
||||
self.fallback_models.append(model)
|
||||
|
||||
|
||||
def parse_routing_policy(data):
|
||||
"""Parse one policy document, or return None when it is unusable.
|
||||
|
||||
Unusable means: not a dict, no usable ``tiers`` map, or a ``mode``
|
||||
outside ``off``/``shadow``/``active``. Returning None (instead of
|
||||
raising) is the fail-closed path: the caller falls back to the
|
||||
worker-default chain and the task is never rejected for a bad
|
||||
policy file.
|
||||
"""
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
mode = data.get("mode", "shadow")
|
||||
if mode not in ROUTING_MODES:
|
||||
return None
|
||||
tiers = _clean_tiers(data.get("tiers"))
|
||||
if not tiers:
|
||||
return None
|
||||
fallback_models = data.get("fallback_models")
|
||||
if fallback_models is not None and not isinstance(fallback_models,
|
||||
list):
|
||||
return None
|
||||
return RoutingPolicy(version=data.get("version"), mode=mode,
|
||||
tiers=tiers, fallback_models=fallback_models)
|
||||
|
||||
|
||||
def load_routing_policy(env=None):
|
||||
"""Load a policy from env, following the cost-table convention.
|
||||
|
||||
``CLOUD_ROUTING_POLICY_FILE`` is read first, then
|
||||
``CLOUD_ROUTING_POLICY_JSON`` overrides its fields per tier;
|
||||
``CLOUD_ROUTING_MODE`` overrides the mode when valid. Any source
|
||||
that fails to parse is ignored, never fatal. Returns None when no
|
||||
usable policy was found, so the caller keeps the old behaviour.
|
||||
"""
|
||||
env = env if env is not None else os.environ
|
||||
merged = None
|
||||
path = env.get("CLOUD_ROUTING_POLICY_FILE") or ""
|
||||
if path:
|
||||
try:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
merged = json.load(f)
|
||||
except (OSError, ValueError):
|
||||
merged = None
|
||||
raw = env.get("CLOUD_ROUTING_POLICY_JSON") or ""
|
||||
if raw:
|
||||
try:
|
||||
override = json.loads(raw)
|
||||
except ValueError:
|
||||
override = None
|
||||
if isinstance(override, dict):
|
||||
if not isinstance(merged, dict):
|
||||
merged = override
|
||||
else:
|
||||
merged = dict(merged)
|
||||
merged.update(override)
|
||||
if isinstance(merged.get("tiers"), dict) and \
|
||||
isinstance(override.get("tiers"), dict):
|
||||
tiers = dict(merged["tiers"])
|
||||
tiers.update(override["tiers"])
|
||||
merged["tiers"] = tiers
|
||||
policy = parse_routing_policy(merged)
|
||||
if policy is None:
|
||||
return None
|
||||
mode_override = env.get("CLOUD_ROUTING_MODE") or ""
|
||||
if mode_override in ROUTING_MODES:
|
||||
policy.mode = mode_override
|
||||
return policy
|
||||
|
||||
|
||||
class RoutingDecision:
|
||||
"""One pre-dispatch model decision, kept auditable by construction."""
|
||||
|
||||
def __init__(self, task_id, tier, chosen_model, suggested_model,
|
||||
source, policy_version, applied, reason, signals):
|
||||
self.task_id = task_id
|
||||
self.tier = tier
|
||||
self.chosen_model = chosen_model
|
||||
self.suggested_model = suggested_model
|
||||
self.source = source
|
||||
self.policy_version = policy_version
|
||||
self.applied = applied
|
||||
self.reason = reason
|
||||
self.signals = dict(signals or {})
|
||||
|
||||
def to_event(self) -> dict:
|
||||
return {
|
||||
"type": "model_routing",
|
||||
"task_id": self.task_id,
|
||||
"tier": self.tier,
|
||||
"chosen_model": self.chosen_model,
|
||||
"suggested_model": self.suggested_model,
|
||||
"source": self.source,
|
||||
"policy_version": self.policy_version,
|
||||
"applied": self.applied,
|
||||
"reason": self.reason,
|
||||
}
|
||||
|
||||
|
||||
def decide_routing(task, worker_model, policy=None,
|
||||
worker_model_explicit=False):
|
||||
"""Decide the model for one task, in the documented priority order.
|
||||
|
||||
1. task-level ``model`` (explicit) -- never overridden;
|
||||
2. worker CLI ``--model`` (explicit) -- a pinned deployment is never
|
||||
silently changed by a policy file;
|
||||
3. active policy tier model -- only when the task named no model;
|
||||
4. worker default chain (config file / ``DEFAULT_MODEL``).
|
||||
|
||||
Shadow mode returns the worker default as ``chosen_model`` and the
|
||||
policy tier model as ``suggested_model`` with ``applied=False``:
|
||||
it observes without changing behaviour. This function is pure: it
|
||||
never reads env, files, or a provider, so it is unit-testable with
|
||||
plain dicts.
|
||||
"""
|
||||
task = task if isinstance(task, dict) else {}
|
||||
tier = classify_task(task)
|
||||
signals = {
|
||||
"plan": bool(task.get("plan")),
|
||||
"goal_mode": bool(task.get("goal_mode")),
|
||||
"prompt_len_bucket": prompt_len_bucket(task.get("prompt")),
|
||||
"effort": task.get("effort"),
|
||||
}
|
||||
task_model = _clean_model_name(task.get("model"))
|
||||
if task_model:
|
||||
return RoutingDecision(
|
||||
task_id=task.get("id"), tier=tier, chosen_model=task_model,
|
||||
suggested_model=None, source="explicit", policy_version=None,
|
||||
applied=True, reason="task model explicit", signals=signals)
|
||||
worker_model = _clean_model_name(worker_model)
|
||||
if worker_model_explicit and worker_model:
|
||||
return RoutingDecision(
|
||||
task_id=task.get("id"), tier=tier, chosen_model=worker_model,
|
||||
suggested_model=None, source="worker_explicit",
|
||||
policy_version=None, applied=True,
|
||||
reason="worker CLI model explicit", signals=signals)
|
||||
if policy is None or policy.mode == "off":
|
||||
return RoutingDecision(
|
||||
task_id=task.get("id"), tier=tier, chosen_model=worker_model,
|
||||
suggested_model=None, source="worker_default",
|
||||
policy_version=getattr(policy, "version", None),
|
||||
applied=True, reason="routing off or no policy",
|
||||
signals=signals)
|
||||
tier_model = policy.tiers.get(tier)
|
||||
if not tier_model:
|
||||
return RoutingDecision(
|
||||
task_id=task.get("id"), tier=tier, chosen_model=worker_model,
|
||||
suggested_model=None, source="worker_default",
|
||||
policy_version=policy.version, applied=True,
|
||||
reason="%s tier has no policy model" % tier, signals=signals)
|
||||
if policy.mode == "shadow":
|
||||
return RoutingDecision(
|
||||
task_id=task.get("id"), tier=tier, chosen_model=worker_model,
|
||||
suggested_model=tier_model, source="shadow",
|
||||
policy_version=policy.version, applied=False,
|
||||
reason="%s tier suggests %s (shadow, not applied)"
|
||||
% (tier, tier_model), signals=signals)
|
||||
# mode == "active"
|
||||
return RoutingDecision(
|
||||
task_id=task.get("id"), tier=tier, chosen_model=tier_model,
|
||||
suggested_model=None, source="policy",
|
||||
policy_version=policy.version, applied=True,
|
||||
reason="%s tier -> %s" % (tier, tier_model), signals=signals)
|
||||
|
||||
|
||||
def record_routing(task_record, decision) -> dict:
|
||||
"""Append a ``model_routing`` event to a task record, like fallback.
|
||||
|
||||
Mirrors ``record_fallback``: the event is appended to
|
||||
``task_record["events"]`` (created if absent) so a routed model is
|
||||
never silent. Returns the event.
|
||||
"""
|
||||
event = decision.to_event() if isinstance(decision, RoutingDecision) \
|
||||
else dict(decision)
|
||||
events = task_record.get("events")
|
||||
if not isinstance(events, list):
|
||||
events = []
|
||||
task_record["events"] = events
|
||||
events.append(event)
|
||||
return event
|
||||
|
||||
|
||||
def record_fallback(task_record, from_model, to_model, reason) -> dict:
|
||||
"""Record a model fallback on ``task_record`` and return the event.
|
||||
|
||||
|
||||
@@ -0,0 +1,326 @@
|
||||
"""cloud 规则路由 v0 单测:分类、决策、策略解析、worker 接入。
|
||||
|
||||
运行(在仓库根目录下):
|
||||
python3 -m unittest discover -s cloud/tests -p "test_models_routing.py" -v
|
||||
|
||||
说明:只用纯函数与注入的假 runner,不跑真实模型回合、不联网、
|
||||
不登录任何服务;env 映射也是注入的普通 dict,里面的值只是
|
||||
占位用的假名,绝不是真实密钥。
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
try:
|
||||
from cloud.models import (
|
||||
ModelRouter,
|
||||
classify_task,
|
||||
decide_routing,
|
||||
is_fallbackable,
|
||||
load_routing_policy,
|
||||
parse_routing_policy,
|
||||
prompt_len_bucket,
|
||||
record_routing,
|
||||
)
|
||||
from cloud.server import init_db
|
||||
from cloud.worker import Worker
|
||||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.models import (
|
||||
ModelRouter,
|
||||
classify_task,
|
||||
decide_routing,
|
||||
is_fallbackable,
|
||||
load_routing_policy,
|
||||
parse_routing_policy,
|
||||
prompt_len_bucket,
|
||||
record_routing,
|
||||
)
|
||||
from cloud.server import init_db
|
||||
from cloud.worker import Worker
|
||||
|
||||
|
||||
POLICY = {
|
||||
"version": "routing-v0.1",
|
||||
"mode": "active",
|
||||
"tiers": {
|
||||
"plan": "gpt-6.1-sol",
|
||||
"goal": "gpt-6.1-sol",
|
||||
"default": "gpt-5.6-terra",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _policy(**overrides):
|
||||
policy = json.loads(json.dumps(POLICY))
|
||||
policy.update(overrides)
|
||||
return parse_routing_policy(policy)
|
||||
|
||||
|
||||
def _insert_task(db_path, task_id, prompt="do work", plan=0, model=None):
|
||||
now = time.time()
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||||
" idempotency_key, depends_on, priority, created_by,"
|
||||
" created_at, updated_at, last_heartbeat)"
|
||||
" VALUES(?, 'p1', ?, 'queued', NULL, '[]', 0, 'tester',"
|
||||
" ?, ?, NULL)",
|
||||
(task_id, prompt, now, now),
|
||||
)
|
||||
conn.execute(
|
||||
"UPDATE tasks SET plan = ?, model = ? WHERE id = ?",
|
||||
(plan, model, task_id),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _completed_runner(cmd):
|
||||
if "dispatch" in cmd:
|
||||
name = cmd[cmd.index("--name") + 1]
|
||||
else:
|
||||
name = cmd[cmd.index("--task") + 1]
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout=json.dumps({"task": name, "status": "completed",
|
||||
"text": "done"}),
|
||||
stderr="")
|
||||
|
||||
|
||||
class ClassifyTestCase(unittest.TestCase):
|
||||
def test_plan_beats_goal_and_default(self):
|
||||
self.assertEqual(classify_task({"plan": 1}), "plan")
|
||||
self.assertEqual(classify_task({"plan": 1, "goal_mode": 1}),
|
||||
"plan")
|
||||
self.assertEqual(classify_task({"goal_mode": 1}), "goal")
|
||||
self.assertEqual(
|
||||
classify_task({"verify_command": "make test"}), "goal")
|
||||
self.assertEqual(classify_task({}), "default")
|
||||
self.assertEqual(classify_task({"priority": 9}), "default")
|
||||
|
||||
def test_prompt_len_bucket_is_logging_only(self):
|
||||
self.assertEqual(prompt_len_bucket("x" * 10), "<500")
|
||||
self.assertEqual(prompt_len_bucket("x" * 500), "500-3000")
|
||||
self.assertEqual(prompt_len_bucket("x" * 3001), ">3000")
|
||||
# 长度再长也不改变分层:分层只看结构信号。
|
||||
self.assertEqual(classify_task({"prompt": "x" * 5000}),
|
||||
"default")
|
||||
|
||||
|
||||
class ParsePolicyTestCase(unittest.TestCase):
|
||||
def test_valid_policy(self):
|
||||
policy = _policy()
|
||||
self.assertEqual(policy.mode, "active")
|
||||
self.assertEqual(policy.version, "routing-v0.1")
|
||||
self.assertEqual(policy.tiers["plan"], "gpt-6.1-sol")
|
||||
|
||||
def test_invalid_mode_or_missing_tiers_fails_closed(self):
|
||||
self.assertIsNone(parse_routing_policy(
|
||||
{"mode": "bogus", "tiers": {"default": "x"}}))
|
||||
self.assertIsNone(parse_routing_policy({"mode": "active"}))
|
||||
self.assertIsNone(parse_routing_policy("not-a-dict"))
|
||||
self.assertIsNone(parse_routing_policy(None))
|
||||
|
||||
def test_empty_tier_model_is_dropped_not_fatal(self):
|
||||
policy = parse_routing_policy(
|
||||
{"mode": "active",
|
||||
"tiers": {"plan": " ", "default": "gpt-5.6-terra"}})
|
||||
self.assertNotIn("plan", policy.tiers)
|
||||
self.assertEqual(policy.tiers["default"], "gpt-5.6-terra")
|
||||
|
||||
def test_load_from_json_and_mode_override(self):
|
||||
policy = load_routing_policy({
|
||||
"CLOUD_ROUTING_POLICY_JSON": json.dumps(
|
||||
{"mode": "shadow", "tiers": {"default": "gpt-6.1-sol"}}),
|
||||
"CLOUD_ROUTING_MODE": "active",
|
||||
})
|
||||
self.assertEqual(policy.mode, "active")
|
||||
self.assertIsNone(load_routing_policy(
|
||||
{"CLOUD_ROUTING_POLICY_JSON": "bad json"}))
|
||||
|
||||
def test_load_from_file(self):
|
||||
with tempfile.NamedTemporaryFile(
|
||||
"w", suffix=".json", delete=False) as f:
|
||||
json.dump({"mode": "shadow",
|
||||
"tiers": {"plan": "gpt-6.1-sol"}}, f)
|
||||
path = f.name
|
||||
try:
|
||||
policy = load_routing_policy(
|
||||
{"CLOUD_ROUTING_POLICY_FILE": path})
|
||||
self.assertEqual(policy.mode, "shadow")
|
||||
self.assertEqual(policy.tiers["plan"], "gpt-6.1-sol")
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
class DecideTestCase(unittest.TestCase):
|
||||
def test_explicit_task_model_wins_over_policy(self):
|
||||
decision = decide_routing(
|
||||
{"id": "t", "model": "gpt-6-sol", "plan": 1},
|
||||
"gpt-5.6-terra", _policy())
|
||||
self.assertEqual(decision.chosen_model, "gpt-6-sol")
|
||||
self.assertEqual(decision.source, "explicit")
|
||||
|
||||
def test_worker_cli_model_wins_over_policy(self):
|
||||
decision = decide_routing(
|
||||
{"id": "t", "plan": 1}, "gpt-6.1-sol", _policy(),
|
||||
worker_model_explicit=True)
|
||||
self.assertEqual(decision.chosen_model, "gpt-6.1-sol")
|
||||
self.assertEqual(decision.source, "worker_explicit")
|
||||
|
||||
def test_shadow_does_not_apply(self):
|
||||
decision = decide_routing(
|
||||
{"id": "t", "plan": 1}, "gpt-5.6-terra",
|
||||
_policy(mode="shadow"))
|
||||
self.assertFalse(decision.applied)
|
||||
self.assertEqual(decision.chosen_model, "gpt-5.6-terra")
|
||||
self.assertEqual(decision.suggested_model, "gpt-6.1-sol")
|
||||
self.assertEqual(decision.source, "shadow")
|
||||
|
||||
def test_active_applies_tier_model(self):
|
||||
decision = decide_routing(
|
||||
{"id": "t", "plan": 1}, "gpt-5.6-terra", _policy())
|
||||
self.assertTrue(decision.applied)
|
||||
self.assertEqual(decision.chosen_model, "gpt-6.1-sol")
|
||||
self.assertEqual(decision.source, "policy")
|
||||
self.assertEqual(decision.policy_version, "routing-v0.1")
|
||||
|
||||
def test_off_or_no_policy_or_missing_tier_falls_back(self):
|
||||
for policy in (None, _policy(mode="off"),
|
||||
_policy(tiers={"plan": "gpt-6.1-sol"})):
|
||||
decision = decide_routing(
|
||||
{"id": "t", "goal_mode": 1}, "gpt-5.6-terra", policy)
|
||||
self.assertEqual(decision.chosen_model, "gpt-5.6-terra")
|
||||
self.assertEqual(decision.source, "worker_default")
|
||||
|
||||
def test_decide_does_not_mutate_task_or_other_fields(self):
|
||||
task = {"id": "t", "plan": 1, "effort": "low",
|
||||
"approval": "ask", "executor": "docker"}
|
||||
before = dict(task)
|
||||
decide_routing(task, "gpt-5.6-terra", _policy())
|
||||
self.assertEqual(task, before)
|
||||
|
||||
def test_decide_is_deterministic(self):
|
||||
task = {"id": "t", "plan": 1}
|
||||
first = decide_routing(task, "gpt-5.6-terra", _policy()).to_event()
|
||||
second = decide_routing(task, "gpt-5.6-terra", _policy()).to_event()
|
||||
self.assertEqual(first, second)
|
||||
|
||||
def test_record_routing_appends_event(self):
|
||||
record = {"id": "t"}
|
||||
decision = decide_routing(
|
||||
{"id": "t", "plan": 1}, "gpt-5.6-terra", _policy())
|
||||
event = record_routing(record, decision)
|
||||
self.assertEqual(event["type"], "model_routing")
|
||||
self.assertEqual(record["events"][0], event)
|
||||
|
||||
def test_fallback_rules_unchanged(self):
|
||||
# 路由 v0 不改事后 fallback:401/400 不触发,404 触发。
|
||||
self.assertFalse(is_fallbackable(401))
|
||||
self.assertFalse(is_fallbackable(400))
|
||||
self.assertTrue(is_fallbackable(404))
|
||||
router = ModelRouter(["gpt-6.1-sol", "gpt-6-sol"])
|
||||
self.assertEqual(router.next_after("gpt-6.1-sol", 503),
|
||||
"gpt-6-sol")
|
||||
self.assertIsNone(router.next_after("gpt-6.1-sol", 401))
|
||||
|
||||
|
||||
class WorkerRoutingTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self._old_push_allow = os.environ.get(
|
||||
"CODEX_PUSH_ALLOW_PRIVATE_TARGETS")
|
||||
os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = "1"
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||||
init_db(self.db_path)
|
||||
self.env = {"XCM_API_KEY": "test-placeholder"}
|
||||
self.calls = []
|
||||
|
||||
def tearDown(self):
|
||||
if getattr(self, "_old_push_allow", None) is None:
|
||||
os.environ.pop("CODEX_PUSH_ALLOW_PRIVATE_TARGETS", None)
|
||||
else:
|
||||
os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = (
|
||||
self._old_push_allow)
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _recording_runner(self):
|
||||
def fake(cmd):
|
||||
self.calls.append(list(cmd))
|
||||
return _completed_runner(cmd)
|
||||
return fake
|
||||
|
||||
def _worker(self, **kwargs):
|
||||
kwargs.setdefault("env", self.env)
|
||||
kwargs.setdefault("model_explicit", False)
|
||||
return Worker(self.db_path, "/fake/host.py",
|
||||
runner=self._recording_runner(), **kwargs)
|
||||
|
||||
def _dispatch_model(self):
|
||||
for cmd in self.calls:
|
||||
if "dispatch" in cmd and "--model" in cmd:
|
||||
return cmd[cmd.index("--model") + 1]
|
||||
return None
|
||||
|
||||
def test_worker_init_explicit_flag_and_bad_policy(self):
|
||||
# _worker 默认 model_explicit=False(模拟未钉死的 worker);
|
||||
# 这里显式传 True,验证显式标志与策略解析。
|
||||
worker = self._worker(model="gpt-6.1-sol", model_explicit=True,
|
||||
routing_policy=POLICY)
|
||||
self.assertTrue(worker.model_explicit)
|
||||
self.assertEqual(worker.routing_policy.mode, "active")
|
||||
worker2 = self._worker(
|
||||
routing_policy={"mode": "bogus", "tiers": {"default": "x"}})
|
||||
self.assertIsNone(worker2.routing_policy)
|
||||
# 不传 model_explicit 时按是否传 model 推断:传了即显式。
|
||||
worker3 = Worker(self.db_path, "/fake/host.py",
|
||||
runner=self._recording_runner(), env=self.env,
|
||||
model="gpt-6.1-sol")
|
||||
self.assertTrue(worker3.model_explicit)
|
||||
|
||||
def test_run_all_active_plan_uses_policy_model(self):
|
||||
_insert_task(self.db_path, "T1", plan=1)
|
||||
worker = self._worker(routing_policy=POLICY)
|
||||
results = worker.run_all()
|
||||
self.assertEqual(results[0]["status"], "completed")
|
||||
self.assertEqual(results[0]["actual_model"], "gpt-6.1-sol")
|
||||
self.assertEqual(self._dispatch_model(), "gpt-6.1-sol")
|
||||
event = results[0]["events"][0]
|
||||
self.assertEqual(event["type"], "model_routing")
|
||||
self.assertEqual(event["tier"], "plan")
|
||||
self.assertEqual(event["source"], "policy")
|
||||
self.assertTrue(event["applied"])
|
||||
|
||||
def test_run_all_shadow_does_not_change_dispatch(self):
|
||||
_insert_task(self.db_path, "T1", plan=1)
|
||||
worker = self._worker(routing_policy=_policy(mode="shadow"))
|
||||
results = worker.run_all()
|
||||
self.assertEqual(results[0]["actual_model"], "gpt-5.6-terra")
|
||||
self.assertEqual(self._dispatch_model(), "gpt-5.6-terra")
|
||||
event = results[0]["events"][0]
|
||||
self.assertFalse(event["applied"])
|
||||
self.assertEqual(event["suggested_model"], "gpt-6.1-sol")
|
||||
|
||||
def test_run_all_explicit_task_model_has_no_routing_event(self):
|
||||
_insert_task(self.db_path, "T1", plan=1, model="gpt-6-sol")
|
||||
worker = self._worker(routing_policy=POLICY)
|
||||
results = worker.run_all()
|
||||
self.assertEqual(results[0]["actual_model"], "gpt-6-sol")
|
||||
self.assertEqual(self._dispatch_model(), "gpt-6-sol")
|
||||
self.assertEqual(
|
||||
[e for e in results[0]["events"]
|
||||
if e["type"] == "model_routing"], [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+48
-2
@@ -146,7 +146,9 @@ from cloud.cost import CostStore, estimate_tokens, price_usd_micros, \
|
||||
from cloud.db import connect as _db_connect
|
||||
from cloud.goal_command import parse_goal_command as _parse_goal_command
|
||||
from cloud.logging_config import get_logger, log_event, setup_logging
|
||||
from cloud.models import ModelRouter, is_fallbackable, record_fallback
|
||||
from cloud.models import ModelRouter, RoutingPolicy, decide_routing, \
|
||||
is_fallbackable, load_routing_policy, parse_routing_policy, \
|
||||
record_fallback, record_routing
|
||||
from cloud.auto_memory import AutoMemoryStore
|
||||
from cloud.auto_memory import build_injection_prefix as _auto_prefix
|
||||
from cloud.auto_memory import process_due_tasks as _auto_process_due
|
||||
@@ -424,6 +426,7 @@ class Worker:
|
||||
def __init__(self, db_path, host_py_path, config_path=None, repo=None,
|
||||
runner=None, env=None, env_key=None,
|
||||
codex_config_path=None, model=None, router=None,
|
||||
model_explicit=None, routing_policy=None,
|
||||
max_turns=None, max_cost_usd=None, executor=None,
|
||||
cost_store=None, process_factory=None,
|
||||
task_timeout=1800, max_retries=1,
|
||||
@@ -472,9 +475,32 @@ class Worker:
|
||||
self._env = env if env is not None else os.environ
|
||||
self.env_key = resolve_env_key(env_key, codex_config_path)
|
||||
self.model = _resolve_model(model, config_path)
|
||||
# CLI --model 是否显式:未传参时按“是否传了 model 参数”推断,
|
||||
# 传了即显式(路由绝不覆盖已钉死的部署模型);调用方也可
|
||||
# 显式传 model_explicit=False,把 model 当作普通默认值。
|
||||
if model_explicit is None:
|
||||
self.model_explicit = model is not None
|
||||
else:
|
||||
self.model_explicit = bool(model_explicit)
|
||||
if router is not None and not isinstance(router, ModelRouter):
|
||||
raise ValueError("router 必须是 cloud.models.ModelRouter")
|
||||
self.router = router
|
||||
# 规则路由策略 v0:显式 RoutingPolicy > dict(解析)> 未传时
|
||||
# 从环境读(CLOUD_ROUTING_POLICY_*,见 models.py)。
|
||||
# 解析失败 / 未配时为 None,派发回落旧链,绝不因此拒绝任务。
|
||||
# 策略里的 fallback_models 在 v0 不启用真实重试(见 models.py),
|
||||
# 只有 CLI --fallback-models 能启用,策略文件 alone 不能烧钱。
|
||||
if isinstance(routing_policy, RoutingPolicy):
|
||||
self.routing_policy = routing_policy
|
||||
elif isinstance(routing_policy, dict):
|
||||
self.routing_policy = parse_routing_policy(routing_policy)
|
||||
elif routing_policy is None:
|
||||
try:
|
||||
self.routing_policy = load_routing_policy(self._env)
|
||||
except Exception:
|
||||
self.routing_policy = None
|
||||
else:
|
||||
raise ValueError("routing_policy 必须是 RoutingPolicy 或 dict")
|
||||
# 真实 fallback 换模型列表:默认空 = 关闭,只记事件不换模型
|
||||
# (老行为,生产 worker 未配该参数时绝不突然多花模型钱)。
|
||||
# 归一化:去空白、去重保序、去掉与主模型相同的项(换成同一
|
||||
@@ -5817,7 +5843,27 @@ class Worker:
|
||||
|
||||
task = dict(task)
|
||||
if not task.get("model"):
|
||||
task["model"] = self.model
|
||||
# 规则路由 v0(只在模型仍为空时跑一次,整个派发只算一次,
|
||||
# 后续重试/fallback/续跑都复用这里定下的 task["model"],
|
||||
# 不再重算;网关步已填的 sub2api 默认模型在此之前已生效,
|
||||
# 路由不覆盖它)。显式任务模型与 worker CLI 显式模型
|
||||
# 在 decide_routing 内同样优先,见 models.py 的优先级。
|
||||
decision = decide_routing(
|
||||
task, self.model, policy=self.routing_policy,
|
||||
worker_model_explicit=self.model_explicit)
|
||||
# 只在策略真正介入时记事件与审计:worker_default /
|
||||
# worker_explicit 就是旧行为,记一条 model_routing 会
|
||||
# 打破现有调用方对 events 的假设(例如 fallback 单测
|
||||
# 断言 events[0] 就是 fallback 事件、无事件时为空),
|
||||
# 而旧行为本身已由 task["model"] / actual_model 体现,
|
||||
# 不需要再用一条事件重复说明。
|
||||
if decision.source in ("policy", "shadow"):
|
||||
record_routing({"id": task_id, "events": events}, decision)
|
||||
self._audit(task["project_id"], "worker_routing",
|
||||
"%s tier=%s source=%s applied=%s" % (
|
||||
task_id, decision.tier, decision.source,
|
||||
decision.applied))
|
||||
task["model"] = decision.chosen_model or self.model
|
||||
actual_model = task["model"]
|
||||
# 有限重试与真实 fallback 的优先级:取消不重试;一次尝试失败
|
||||
# 且可 fallback、且显式配了 fallback_models 时,优先真换模型
|
||||
|
||||
Reference in new issue
Block a user