feat(router): add policy-class queues

components/src/dynamo/common/configuration/groups/kv_router_args.py

@@ -44,7 +44,7 @@
     "router_reset_states",
     "router_ttl_secs",
     "router_queue_threshold",
-    "router_queue_by_incoming_missing_isl",
+    "router_policy_config",
     "router_event_threads",
     "router_queue_policy",
     "use_remote_indexer",
@@ -127,7 +127,7 @@ class KvRouterConfigBase(ConfigBase):
     router_reset_states: bool
     router_ttl_secs: float
     router_queue_threshold: Optional[float]
-    router_queue_by_incoming_missing_isl: Optional[list[tuple[int, int]]] = None
+    router_policy_config: Optional[str] = None
     router_event_threads: int
     router_queue_policy: str
     use_remote_indexer: bool = False
@@ -403,6 +403,19 @@ def add_arguments(self, parser) -> None:
             ),
             arg_type=nullable_float,
         )
+        add_argument(
+            g,
+            flag_name="--router-policy-config",
+            env_var="DYN_ROUTER_POLICY_CONFIG",
+            default=None,
+            help=(
+                "KV Router: Startup-only YAML policy-family and cache-bucket "
+                "queue configuration. "
+                "When omitted, router_queue_threshold and router_queue_policy define "
+                "the existing single default queue."
+            ),
+            arg_type=str,
+        )
         add_argument(
             g,
             flag_name="--router-event-threads",

components/src/dynamo/common/tests/configuration/test_kv_router_args.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import argparse
+from pathlib import Path
 
 import pytest
 
@@ -145,21 +146,6 @@ def test_deprecated_overlap_score_weight_cli_flows_to_binding_kwargs() -> None:
     assert config.kv_router_kwargs()["overlap_score_weight"] == 2.5
 
 
-def test_deprecated_overlap_score_weight_zero_cli_flows_to_binding_kwargs() -> None:
-    parser = argparse.ArgumentParser()
-    KvRouterArgGroup().add_arguments(parser)
-
-    with pytest.warns(FutureWarning, match="overlap score weight is deprecated"):
-        args = parser.parse_args(["--router-kv-overlap-score-weight", "0"])
-
-    assert args.overlap_score_credit == 1.0
-    assert args.prefill_load_scale == 1.0
-    assert args.overlap_score_weight == 0.0
-
-    config = KvRouterConfigBase.from_cli_args(args)
-    assert config.kv_router_kwargs()["overlap_score_weight"] == 0.0
-
-
 def test_deprecated_overlap_score_weight_env_flows_to_binding_kwargs(
     monkeypatch,
 ) -> None:
@@ -179,132 +165,41 @@ def test_deprecated_overlap_score_weight_env_flows_to_binding_kwargs(
     assert config.kv_router_kwargs()["overlap_score_weight"] == 2.5
 
 
-def test_deprecated_overlap_score_weight_zero_env_flows_to_binding_kwargs(
-    monkeypatch,
-) -> None:
-    monkeypatch.setenv("DYN_ROUTER_KV_OVERLAP_SCORE_WEIGHT", "0")
-
-    with pytest.warns(FutureWarning, match="deprecated"):
-        parser = argparse.ArgumentParser()
-        KvRouterArgGroup().add_arguments(parser)
-
-    args = parser.parse_args([])
-
-    assert args.overlap_score_credit == 1.0
-    assert args.prefill_load_scale == 1.0
-    assert args.overlap_score_weight == 0.0
-
-    config = KvRouterConfigBase.from_cli_args(args)
-    assert config.kv_router_kwargs()["overlap_score_weight"] == 0.0
-
-
-def test_deprecated_overlap_score_weight_env_overrides_new_scale_env(
-    monkeypatch,
-) -> None:
-    monkeypatch.setenv("DYN_ROUTER_KV_OVERLAP_SCORE_WEIGHT", "0")
-    monkeypatch.setenv("DYN_ROUTER_PREFILL_LOAD_SCALE", "2.5")
-
-    with pytest.warns(FutureWarning, match="deprecated"):
-        parser = argparse.ArgumentParser()
-        KvRouterArgGroup().add_arguments(parser)
-
-    args = parser.parse_args([])
-
-    assert args.overlap_score_credit == 1.0
-    assert args.prefill_load_scale == 2.5
-    assert args.overlap_score_weight == 0.0
-
-    config = KvRouterConfigBase.from_cli_args(args)
-    assert config.kv_router_kwargs()["overlap_score_weight"] == 0.0
-
-
-def test_deprecated_overlap_score_weight_env_overrides_new_credit_env(
-    monkeypatch,
-) -> None:
-    monkeypatch.setenv("DYN_ROUTER_KV_OVERLAP_SCORE_WEIGHT", "0")
-    monkeypatch.setenv("DYN_ROUTER_KV_OVERLAP_SCORE_CREDIT", "0.5")
-
-    with pytest.warns(FutureWarning, match="deprecated"):
-        parser = argparse.ArgumentParser()
-        KvRouterArgGroup().add_arguments(parser)
-
-    args = parser.parse_args([])
-
-    assert args.overlap_score_credit == 0.5
-    assert args.prefill_load_scale == 1.0
-    assert args.overlap_score_weight == 0.0
-
-    config = KvRouterConfigBase.from_cli_args(args)
-    assert config.kv_router_kwargs()["overlap_score_weight"] == 0.0
-
-
-def test_deprecated_overlap_score_weight_env_overrides_new_scale_cli(
-    monkeypatch,
-) -> None:
-    monkeypatch.setenv("DYN_ROUTER_KV_OVERLAP_SCORE_WEIGHT", "0")
-
-    with pytest.warns(FutureWarning, match="deprecated"):
-        parser = argparse.ArgumentParser()
-        KvRouterArgGroup().add_arguments(parser)
-
-    args = parser.parse_args(["--router-prefill-load-scale", "3"])
-
-    assert args.overlap_score_credit == 1.0
-    assert args.prefill_load_scale == 3.0
-    assert args.overlap_score_weight == 0.0
-
-    config = KvRouterConfigBase.from_cli_args(args)
-    assert config.kv_router_kwargs()["overlap_score_weight"] == 0.0
-
-
-def test_deprecated_overlap_score_weight_env_overrides_new_credit_cli(
+@pytest.mark.parametrize(
+    ("canonical_env", "cli_args", "expected_credit", "expected_scale"),
+    [
+        (("DYN_ROUTER_PREFILL_LOAD_SCALE", "2.5"), [], 1.0, 2.5),
+        (("DYN_ROUTER_KV_OVERLAP_SCORE_CREDIT", "0.5"), [], 0.5, 1.0),
+        (None, ["--router-prefill-load-scale", "3"], 1.0, 3.0),
+        (None, ["--router-kv-overlap-score-credit", "0.5"], 0.5, 1.0),
+    ],
+    ids=["scale-env", "credit-env", "scale-cli", "credit-cli"],
+)
+def test_deprecated_overlap_score_weight_env_coexists_with_canonical_settings(
     monkeypatch,
+    canonical_env,
+    cli_args,
+    expected_credit,
+    expected_scale,
 ) -> None:
     monkeypatch.setenv("DYN_ROUTER_KV_OVERLAP_SCORE_WEIGHT", "0")
+    if canonical_env is not None:
+        monkeypatch.setenv(*canonical_env)
 
     with pytest.warns(FutureWarning, match="deprecated"):
         parser = argparse.ArgumentParser()
         KvRouterArgGroup().add_arguments(parser)
 
-    args = parser.parse_args(["--router-kv-overlap-score-credit", "0.5"])
+    args = parser.parse_args(cli_args)
 
-    assert args.overlap_score_credit == 0.5
-    assert args.prefill_load_scale == 1.0
+    assert args.overlap_score_credit == expected_credit
+    assert args.prefill_load_scale == expected_scale
     assert args.overlap_score_weight == 0.0
 
     config = KvRouterConfigBase.from_cli_args(args)
     assert config.kv_router_kwargs()["overlap_score_weight"] == 0.0
 
 
-def test_deprecated_overlap_score_weight_cli_order_does_not_change_presence() -> None:
-    parser = argparse.ArgumentParser()
-    KvRouterArgGroup().add_arguments(parser)
-
-    with pytest.warns(FutureWarning, match="overlap score weight is deprecated"):
-        old_then_new = parser.parse_args(
-            [
-                "--router-kv-overlap-score-weight",
-                "0",
-                "--router-prefill-load-scale",
-                "3",
-            ]
-        )
-    with pytest.warns(FutureWarning, match="overlap score weight is deprecated"):
-        new_then_old = parser.parse_args(
-            [
-                "--router-prefill-load-scale",
-                "3",
-                "--router-kv-overlap-score-weight",
-                "0",
-            ]
-        )
-
-    assert old_then_new.overlap_score_weight == 0.0
-    assert old_then_new.prefill_load_scale == 3.0
-    assert new_then_old.overlap_score_weight == 0.0
-    assert new_then_old.prefill_load_scale == 3.0
-
-
 def test_prefill_load_scale_cli_uses_kv_router_config_field() -> None:
     parser = argparse.ArgumentParser()
     KvRouterArgGroup().add_arguments(parser)
@@ -402,27 +297,19 @@ def test_load_aware_preserves_cache_hit_weights() -> None:
     assert kwargs["disk_cache_hit_weight"] == 0.1
 
 
-def test_kv_router_kwargs_preserves_explicit_queue_tiers() -> None:
-    parser = argparse.ArgumentParser()
-    KvRouterArgGroup().add_arguments(parser)
-
-    config = KvRouterConfigBase.from_cli_args(parser.parse_args([]))
-    config.router_queue_by_incoming_missing_isl = [(0, 2048), (2048, 512)]
-
-    kwargs = config.kv_router_kwargs()
-
-    assert kwargs["router_queue_by_incoming_missing_isl"] == [(0, 2048), (2048, 512)]
-
-
-def test_kv_router_kwargs_uses_unbounded_queue_cap_by_default() -> None:
+def test_policy_config_cli_overrides_environment(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+    env_policy_path = str(tmp_path / "env-policy.yaml")
+    explicit_policy_path = str(tmp_path / "explicit-policy.yaml")
+    monkeypatch.setenv("DYN_ROUTER_POLICY_CONFIG", env_policy_path)
     parser = argparse.ArgumentParser()
     KvRouterArgGroup().add_arguments(parser)
 
-    config = KvRouterConfigBase.from_cli_args(parser.parse_args([]))
-
-    kwargs = config.kv_router_kwargs()
+    args = parser.parse_args(["--router-policy-config", explicit_policy_path])
+    config = KvRouterConfigBase.from_cli_args(args)
 
-    assert kwargs["router_queue_by_incoming_missing_isl"] is None
+    assert config.kv_router_kwargs()["router_policy_config"] == explicit_policy_path
 
 
 def test_load_aware_clears_predicted_ttl() -> None:
@@ -646,29 +533,6 @@ def test_admission_control_default_none_with_explicit_threshold_auto_switches(
     assert config.active_prefill_tokens_threshold_frac == 64.0
 
 
-def test_admission_control_default_none_with_no_thresholds_stays_none(
-    monkeypatch,
-) -> None:
-    """When no threshold flag is explicitly set, the default 'none' is
-    preserved (no auto-switch). Already covered by
-    test_frontend_admission_control_defaults_to_none — this is the
-    explicit symmetry check next to the auto-switch tests above."""
-    _clear_admission_control_env(monkeypatch)
-    parser = argparse.ArgumentParser()
-    FrontendArgGroup().add_arguments(parser)
-
-    args = parser.parse_args(["--router-queue-threshold", "32.0"])
-
-    config = FrontendConfig.from_cli_args(args)
-    config.validate()
-
-    assert config.admission_control == "none"
-    assert config.active_decode_blocks_threshold is None
-    assert config.active_prefill_tokens_threshold is None
-    assert config.active_prefill_tokens_threshold_frac is None
-    assert config.router_queue_threshold == 32.0
-
-
 def test_admission_control_apply_is_idempotent(monkeypatch) -> None:
     """``apply_admission_control()`` runs once in ``validate()`` and again
     when ``router_kwargs()`` builds the worker config. The second call

components/src/dynamo/mocker/__init__.py

@@ -50,6 +50,7 @@ def run_mocker_trace_replay(
         trace_format="mooncake",
         trace_shared_prefix_ratio=0.0,
         trace_num_prefix_groups=0,
+        model_name=None,
     ):
         return _run_mocker_trace_replay(
             trace_file,
@@ -64,6 +65,7 @@ def run_mocker_trace_replay(
             trace_format=trace_format,
             trace_shared_prefix_ratio=trace_shared_prefix_ratio,
             trace_num_prefix_groups=trace_num_prefix_groups,
+            model_name=model_name,
         )
 
 

components/src/dynamo/replay/api.py

@@ -28,6 +28,7 @@ def run_trace_replay(
     trace_num_prefix_groups=0,
     report_jsonl_path=None,
     max_sim_time_ms=None,
+    model_name=None,
 ):
     return _run_mocker_trace_replay(
         trace_file,
@@ -49,6 +50,7 @@ def run_trace_replay(
         trace_num_prefix_groups=trace_num_prefix_groups,
         report_jsonl_path=report_jsonl_path,
         max_sim_time_ms=max_sim_time_ms,
+        model_name=model_name,
     )
 
 
@@ -74,6 +76,7 @@ def run_synthetic_trace_replay(
     shared_prefix_ratio=0.0,
     num_prefix_groups=0,
     inter_turn_delay_ms=0.0,
+    model_name=None,
 ):
     return _run_mocker_synthetic_trace_replay(
         input_tokens,
@@ -96,4 +99,5 @@ def run_synthetic_trace_replay(
         shared_prefix_ratio=shared_prefix_ratio,
         num_prefix_groups=num_prefix_groups,
         inter_turn_delay_ms=inter_turn_delay_ms,
+        model_name=model_name,
     )