Repository navigation
feat(router): add policy-class queues - #10744
Merged
Merged
Conversation
Signed-off-by: PeaBrane <yanrpei@gmail.com>
Signed-off-by: PeaBrane <yanrpei@gmail.com>
github-actions
Bot
added
feat
documentation
Signed-off-by: PeaBrane <yanrpei@gmail.com>
Contributor
Signed-off-by: PeaBrane <yanrpei@gmail.com>
Signed-off-by: PeaBrane <yanrpei@gmail.com>
Signed-off-by: PeaBrane <yanrpei@gmail.com>
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
Contributor
WalkthroughAdds per-request policy-class queuing to the KV router. A new ChangesPolicy-Class Queuing Feature
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (5)
lib/bindings/python/rust/llm/entrypoint.rs (1)
277-282:⚠️ Potential issue | 🟠 Major | ⚡ Quick winValidate configs loaded through
from_json.
KvRouterConfig.__init__validates the constructed config, butfrom_jsonreturns the parsedRsKvRouterConfigdirectly. That letsrouter_policy_configsupplied via JSON bypass the strict policy YAML validation path until a later startup/use site.Suggested fix
fn from_json(config_json: &str) -> PyResult<Self> { let inner = serde_json::from_str::<RsKvRouterConfig>(config_json).map_err(|e| { PyException::new_err(format!("Failed to parse KvRouterConfig JSON: {e}")) })?; + validate_kv_router_config(&inner)?; Ok(KvRouterConfig { inner }) }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/bindings/python/rust/llm/entrypoint.rs` around lines 277 - 282, The from_json method in KvRouterConfig parses the JSON string into RsKvRouterConfig but returns it directly without applying the same validation that occurs in the __init__ method, allowing invalid router_policy_config to bypass validation. After parsing the JSON with serde_json::from_str and before constructing the KvRouterConfig return value, apply the same validation logic that __init__ uses to ensure the inner RsKvRouterConfig is properly validated, including strict policy YAML validation.lib/kv-router/src/scheduling/types.rs (1)
48-78:⚠️ Potential issue | 🟠 Major | ⚡ Quick winClassify
QueueRejectedas overload.
QueueRejectedis the new policy-queue capacity/backpressure path, butis_overload()still returnsfalsefor it. Any caller using this helper for overload/retry/metrics classification will treat a full policy-class queue as a generic scheduler error.Suggested fix
self, Self::Backpressure { .. } + | Self::QueueRejected(_) | Self::AllEligibleWorkersOverloaded | Self::PinnedWorkerOverloaded { .. } )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/kv-router/src/scheduling/types.rs` around lines 48 - 78, The is_overload() method in the KvSchedulerError implementation does not classify QueueRejected as an overload condition. Update the matches! pattern within is_overload() to include Self::QueueRejected as one of the conditions that should return true, alongside the existing checks for Backpressure, AllEligibleWorkersOverloaded, and PinnedWorkerOverloaded variants. This ensures that callers using is_overload() for classification will correctly treat a full policy-class queue (represented by QueueRejected) as an overload scenario.lib/bindings/python/src/dynamo/_core.pyi (2)
1711-1740:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winKeep
overlap_score_credit_decayin the Python stub API.The Rust binding exposes
overlap_score_credit_decayinKvRouterConfig.__init__, the property accessors, andwith_overrides, but the stub omits it. Valid Python calls are therefore rejected by type checkers.Suggested stub alignment
*, overlap_score_credit: float = 1.0, + overlap_score_credit_decay: float = 0.0, prefill_load_scale: float = 1.0, router_policy_config: Optional[str] = None, @@ `@property` def overlap_score_credit(self) -> float: ... `@overlap_score_credit.setter` def overlap_score_credit(self, value: float) -> None: ... + `@property` + def overlap_score_credit_decay(self) -> float: ... + + `@overlap_score_credit_decay.setter` + def overlap_score_credit_decay(self, value: float) -> None: ... @@ *, overlap_score_credit: Optional[float] = None, + overlap_score_credit_decay: Optional[float] = None, prefill_load_scale: Optional[float] = None, ) -> "KvRouterConfig": ...Also applies to: 1815-1821
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/bindings/python/src/dynamo/_core.pyi` around lines 1711 - 1740, The Python stub file is missing the `overlap_score_credit_decay` parameter that is exposed in the Rust binding's KvRouterConfig. Add the `overlap_score_credit_decay` parameter to the __init__ method in KvRouterConfig at the current location, ensuring it is placed with appropriate type hints and default value alongside the other keyword-only parameters. Additionally, apply the same change to the property accessors and with_overrides method at lines 1815-1821 to maintain consistency across all exposed interfaces in the stub file, ensuring type checkers recognize valid Python calls that match the Rust binding's actual API.
2292-2316:⚠️ Potential issue | 🟠 MajorFix positional parameter order in PlannerReplayBridge signatures.
Both
__init__andcreate_disaggplacemodel_namebeforearrival_speedup_ratioandtrace_block_size, breaking positional compatibility. A call likePlannerReplayBridge(..., router_config, 2.0, 512)now incorrectly binds2.0asmodel_nameinstead ofarrival_speedup_ratio.Move
model_nameafter the existing numeric parameters or make it keyword-only in both the stub (lib/bindings/python/src/dynamo/_core.pyi lines 2299 and 2313) and the Rust PyO3 signature (lib/bindings/python/rust/llm/replay.rs line 1637 for__init__and thecreate_disaggsignature).Suggested stub shape
router_mode: str = "round_robin", router_config: Optional[KvRouterConfig] = None, - model_name: Optional[str] = None, arrival_speedup_ratio: float = 1.0, trace_block_size: int = 512, + *, + model_name: Optional[str] = None, @@ router_mode: str = "round_robin", router_config: Optional[KvRouterConfig] = None, - model_name: Optional[str] = None, arrival_speedup_ratio: float = 1.0, trace_block_size: int = 512, + *, + model_name: Optional[str] = None, ) -> "PlannerReplayBridge": ...🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/bindings/python/src/dynamo/_core.pyi` around lines 2292 - 2316, The PlannerReplayBridge class has parameter order issues in both __init__ and create_disagg methods where model_name is positioned before the numeric parameters arrival_speedup_ratio and trace_block_size, causing positional argument binding to fail. Move model_name to appear after trace_block_size (or make it keyword-only with a leading asterisk) in both method signatures to restore positional compatibility. Apply this fix to both the Python stub file lib/bindings/python/src/dynamo/_core.pyi (for both __init__ at line 2299 and create_disagg at line 2313) and the corresponding Rust PyO3 signatures in lib/bindings/python/rust/llm/replay.rs (the __init__ signature at line 1637 and the create_disagg signature).lib/bindings/python/tests/replay/test_replay_policy_plumbing.py (1)
1-30:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd required pytest category markers to this new test module.
This file currently has no lifecycle/test-type/hardware markers, so marker enforcement can fail for this test.
✅ Suggested fix
+import pytest + import dynamo.replay.api as replay_api +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, +] + def test_replay_api_forwards_policy_model_name(monkeypatch):As per coding guidelines, files under
**/tests/**/*.pymust follow pytest guidance, and this new test should carry required scheduling/marker metadata. Based on learnings, the required marker categories are Lifecycle (pre_merge), Test Type (unit), and Hardware (gpu_0).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/bindings/python/tests/replay/test_replay_policy_plumbing.py` around lines 1 - 30, The test module test_replay_api_forwards_policy_model_name lacks required pytest category markers for lifecycle, test type, and hardware configuration. Import pytest at the top of the file and add three decorator markers to the test_replay_api_forwards_policy_model_name function: `@pytest.mark.pre_merge` for lifecycle, `@pytest.mark.unit` for test type, and `@pytest.mark.gpu_0` for hardware designation. Stack these decorators above the function definition to ensure proper test categorization and scheduling compliance.Sources: Coding guidelines, Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@components/src/dynamo/common/tests/configuration/test_kv_router_args.py`:
- Around line 405-427: The test functions
test_kv_router_kwargs_preserves_policy_config_path and
test_policy_config_cli_overrides_environment use hardcoded temporary paths like
/tmp/router-policy.yaml, which violates hermeticity and platform independence.
Add the tmp_path pytest fixture as a parameter to both test functions, then
replace all hardcoded /tmp paths with paths derived from tmp_path (e.g., by
converting them to string paths based on tmp_path). This ensures tests use
isolated temporary directories and work correctly across different platforms and
test environments.
In `@lib/bindings/python/src/dynamo/llm/exceptions.py`:
- Around line 16-21: The try/except ImportError block around the
RouterQueueLimitExceeded import is too broad and violates Python guidelines.
Instead of catching all ImportErrors, scope the exception handling to only
suppress the ImportError when the specific symbol cannot be imported from
dynamo._core, and re-raise any unexpected ImportErrors that may indicate other
issues. Capture the exception, check if it is related to the missing
RouterQueueLimitExceeded symbol, define the fallback class only in that case,
and re-raise the exception if it represents a different import failure.
In `@lib/data-gen/src/mooncake.rs`:
- Around line 59-60: The new policy_class field has been added to the trace
schema but lacks focused serialize/deserialize round-trip test coverage. Add
wire-format round-trip assertions for the policy_class field on both affected
row structs: the first struct containing the policy_class field (around line
59-60 in mooncake.rs) and the second struct (around line 94-95 in mooncake.rs).
The tests should verify that the field can be serialized and deserialized
correctly, including edge cases like None values which leverage the
serde(default, skip_serializing_if = "Option::is_none") attributes.
Additionally, verify that existing producer/consumer tests still properly cover
the updated schema shape with this new field.
In `@lib/llm/src/kv_router/prefill_router/query.rs`:
- Around line 73-75: The `QueueRejected` match arm in the pattern matching for
`FindBestMatchOutcome` is converting the rejection to a generic error using
`.into()`, which loses the typed rejection details needed by downstream callers
to distinguish throttling/backpressure failures from other error types. Instead
of converting rejection to a generic error, return a structured error variant or
outcome that preserves the QueueRejected semantics so that callers can reliably
classify it as queue throttling rather than treating it as a generic failure.
This may require defining a mapped error type or outcome variant that retains
the rejection information while fitting the function's return type.
In `@lib/mocker/src/replay/offline/components/router.rs`:
- Around line 564-581: The all_workers_busy_with function currently returns
false when workers_with_configs is empty (because checked_any remains false),
which allows drain_pending to pop items and risk losing them if admission fails.
Modify the return statement to handle the empty workers case: instead of
returning checked_any && !any_worker_not_busy, return true (treat as busy) when
checked_any is false, and otherwise return !any_worker_not_busy as before. This
prevents queue items from being popped when there are no workers available to
accept them.
---
Outside diff comments:
In `@lib/bindings/python/rust/llm/entrypoint.rs`:
- Around line 277-282: The from_json method in KvRouterConfig parses the JSON
string into RsKvRouterConfig but returns it directly without applying the same
validation that occurs in the __init__ method, allowing invalid
router_policy_config to bypass validation. After parsing the JSON with
serde_json::from_str and before constructing the KvRouterConfig return value,
apply the same validation logic that __init__ uses to ensure the inner
RsKvRouterConfig is properly validated, including strict policy YAML validation.
In `@lib/bindings/python/src/dynamo/_core.pyi`:
- Around line 1711-1740: The Python stub file is missing the
`overlap_score_credit_decay` parameter that is exposed in the Rust binding's
KvRouterConfig. Add the `overlap_score_credit_decay` parameter to the __init__
method in KvRouterConfig at the current location, ensuring it is placed with
appropriate type hints and default value alongside the other keyword-only
parameters. Additionally, apply the same change to the property accessors and
with_overrides method at lines 1815-1821 to maintain consistency across all
exposed interfaces in the stub file, ensuring type checkers recognize valid
Python calls that match the Rust binding's actual API.
- Around line 2292-2316: The PlannerReplayBridge class has parameter order
issues in both __init__ and create_disagg methods where model_name is positioned
before the numeric parameters arrival_speedup_ratio and trace_block_size,
causing positional argument binding to fail. Move model_name to appear after
trace_block_size (or make it keyword-only with a leading asterisk) in both
method signatures to restore positional compatibility. Apply this fix to both
the Python stub file lib/bindings/python/src/dynamo/_core.pyi (for both __init__
at line 2299 and create_disagg at line 2313) and the corresponding Rust PyO3
signatures in lib/bindings/python/rust/llm/replay.rs (the __init__ signature at
line 1637 and the create_disagg signature).
In `@lib/bindings/python/tests/replay/test_replay_policy_plumbing.py`:
- Around line 1-30: The test module test_replay_api_forwards_policy_model_name
lacks required pytest category markers for lifecycle, test type, and hardware
configuration. Import pytest at the top of the file and add three decorator
markers to the test_replay_api_forwards_policy_model_name function:
`@pytest.mark.pre_merge` for lifecycle, `@pytest.mark.unit` for test type, and
`@pytest.mark.gpu_0` for hardware designation. Stack these decorators above the
function definition to ensure proper test categorization and scheduling
compliance.
In `@lib/kv-router/src/scheduling/types.rs`:
- Around line 48-78: The is_overload() method in the KvSchedulerError
implementation does not classify QueueRejected as an overload condition. Update
the matches! pattern within is_overload() to include Self::QueueRejected as one
of the conditions that should return true, alongside the existing checks for
Backpressure, AllEligibleWorkersOverloaded, and PinnedWorkerOverloaded variants.
This ensures that callers using is_overload() for classification will correctly
treat a full policy-class queue (represented by QueueRejected) as an overload
scenario.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: be725754-b699-483d-9e5b-d096f0da8dd7
⛔ Files ignored due to path filters (3)
Cargo.lockis excluded by!**/*.locklib/bindings/kvbm/Cargo.lockis excluded by!**/*.locklib/bindings/python/Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (56)
components/src/dynamo/common/configuration/groups/kv_router_args.pycomponents/src/dynamo/common/tests/configuration/test_kv_router_args.pycomponents/src/dynamo/mocker/__init__.pycomponents/src/dynamo/replay/api.pycomponents/src/dynamo/replay/main.pydocs/components/frontend/configuration.mddocs/components/router/router-configuration.mddocs/dynosim/runs.mddocs/observability/metrics.mdexamples/router/policy-class-queues.yamllib/bindings/c/src/lib.rslib/bindings/python/rust/errors.rslib/bindings/python/rust/llm/entrypoint.rslib/bindings/python/rust/llm/kv.rslib/bindings/python/rust/llm/replay.rslib/bindings/python/src/dynamo/_core.pyilib/bindings/python/src/dynamo/llm/__init__.pylib/bindings/python/src/dynamo/llm/exceptions.pylib/bindings/python/tests/replay/test_replay_cli.pylib/bindings/python/tests/replay/test_replay_policy_plumbing.pylib/data-gen/src/mooncake.rslib/kv-router/Cargo.tomllib/kv-router/src/protocols.rslib/kv-router/src/scheduling/config.rslib/kv-router/src/scheduling/local.rslib/kv-router/src/scheduling/mod.rslib/kv-router/src/scheduling/policy.rslib/kv-router/src/scheduling/policy_config.rslib/kv-router/src/scheduling/policy_queue.rslib/kv-router/src/scheduling/queue.rslib/kv-router/src/scheduling/selector.rslib/kv-router/src/scheduling/types.rslib/kv-router/src/services/selection/core/mod.rslib/kv-router/src/services/selection/error.rslib/kv-router/src/services/selection/server.rslib/kv-router/src/services/selection/tests.rslib/llm/src/http/service/metadata.rslib/llm/src/http/service/openai.rslib/llm/src/kv_router.rslib/llm/src/kv_router/metrics.rslib/llm/src/kv_router/prefill_router/query.rslib/llm/src/kv_router/push_router.rslib/llm/src/kv_router/push_router/selection.rslib/llm/src/kv_router/scheduler.rslib/mocker/src/common/protocols.rslib/mocker/src/loadgen/driver.rslib/mocker/src/loadgen/tests.rslib/mocker/src/loadgen/trace.rslib/mocker/src/loadgen/types.rslib/mocker/src/replay/offline/components/router.rslib/mocker/src/replay/offline/state.rslib/mocker/src/replay/online/live_runtime.rslib/mocker/src/replay/online/router.rslib/mocker/src/replay/online/tests.rslib/mocker/src/replay/router_shared.rslib/runtime/src/config/environment_names.rs
Signed-off-by: PeaBrane <yanrpei@gmail.com>
Signed-off-by: PeaBrane <yanrpei@gmail.com>
Signed-off-by: PeaBrane <yanrpei@gmail.com>
Signed-off-by: PeaBrane <yanrpei@gmail.com>
Signed-off-by: PeaBrane <yanrpei@gmail.com>
tedzhouhk
approved these changes
Jun 23, 2026
Signed-off-by: PeaBrane <yanrpei@gmail.com>
This comment has been minimized.
This comment has been minimized.
Signed-off-by: PeaBrane <yanrpei@gmail.com>
PeaBrane
enabled auto-merge (squash)
June 23, 2026 21:55
Signed-off-by: PeaBrane <yanrpei@gmail.com>
This branch had an error being deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
--router-policy-config/DYN_ROUTER_POLICY_CONFIG, with strict Rust/serde_yamlvalidation, exact-model profile replacement, root fallback, and the unchanged no-YAML syntheticdefaultqueue.requested policy family × router-observed uncached-ISL bucketto one physical policy-class queue. A class with bothpolicy_familyandcache_bucketis a matrix class; a class with neither is an explicit class that can be selected directly and intentionally bypasses cache-derived classification.x-dynamo-meta-policy-classthrough frontend metadata, prefill cloning, standalone routing, Python bindings, and replay without adding live request-body ornvextfields.Baseten Context
The sample at
examples/router/policy-class-queues.yamlmodelsstandardandlatencypolicy families across router-observedcachedanduncachedbuckets, plus an explicit custom class. Client policy and cache state are independent inputs: ordinary physical-class names cannot bypass cache classification, while a class that omits both binding fields is a trusted explicit override.Compatibility
The policy YAML is a new startup-only surface in this PR. Earlier draft schema names are intentionally not supported, and the replacement for main's missing-ISL admission behavior is intentionally non-equivalent.
defaultqueue using the existingrouter_queue_thresholdandrouter_queue_policy, with no queue limits.default_policy_class,uncached_isl_policy_class_tiers,explicit_only, and absolute queue-limit fields are absent from the strict schema and fail as unknown YAML fields.router_queue_by_incoming_missing_islfield is removed fromKvRouterConfigwithout a compatibility parser and fails normal router-config deserialization as an unknown field.request_queue_limit_per_worker,raw_isl_token_queue_limit_per_worker, andcached_token_queue_limit_per_worker. Effective totals are computed from the current discovered endpoint count at admission; worker removal affects future admission but does not evict queued requests.QueueRejection, HTTP 503, and PythonRouterQueueLimitExceeded, exposing the resolved physical class and computed effective total.Validation
cargo test -p dynamo-kv-router --lib: 608 passed.cargo test -p dynamo-kv-router --lib --features standalone-selection: 654 passed.cargo test -p dynamo-mocker --lib: 371 passed, 1 ignored.cargo test -p dynamo-llm --lib --no-default-features: 1109 passed, 3 ignored..venv/bin/python -m pytest components/src/dynamo/common/tests/configuration/test_kv_router_args.py -q: 31 passed.lib/llm --no-default-features.126c11116fcaf9990f5a765368ac5456da53b5acfa7b3addf74d41ff197334d4; mapped86e8bb1ac65d630856de1418fdb0b3ddd4f3686e8ca4f3bf271a804a80ddac52.4a7b90ed9944d1f867be551e76dfb82dab630a86882634dfa95506097d601e6a.Persistent process-level E2Es remain deferred to a follow-up PR.
Implements a scoped subset of #9755.