Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

lib/bindings/python/src/dynamo/_core.pyi (1)

550-568: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Mirror the current preferred_taints behavior in the Python docs.

This stub still describes preferred_taints as actively affecting scoring, but the default selector only consumes required_taints today. Please mark preferred_taints as accepted-but-not-enforced until frontend routing actually uses it.

🤖 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 550 - 568, The
docstring for class RoutingConstraints inaccurately claims preferred_taints
affects scoring; update the RoutingConstraints docstring (and mention in the
__init__ signature if present) to state that preferred_taints is currently
accepted but not enforced/used by the default selector—i.e., keep
required_taints behavior unchanged, but explicitly mark preferred_taints as a
no-op/for-future-use until frontend routing consumes it so readers understand it
does not impact scoring today.

lib/llm/src/kv_router/prefill_router/execution.rs (2)

273-316: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Reject or filter hard-constrained requests in SimpleRouter mode.

routing_constraints is forwarded to the KV-router branch only. In the SimpleRouter branch, required_taints are ignored, so non-KV prefill routing can still choose an incompatible worker.

🤖 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/llm/src/kv_router/prefill_router/execution.rs` around lines 273 - 316,
query_prefill_worker forwards routing_constraints only to the KvRouter branch so
the SimpleRouter branch (InnerPrefillRouter::SimpleRouter) ignores
required_taints and may pick incompatible workers; update the SimpleRouter path
in query_prefill_worker to first inspect routing_constraints.required_taints
(and other hard constraints like allowed_worker_ids) and either filter available
workers returned by r.select_next_worker()/r.peek_next_worker() to those that
satisfy the constraints or reject the request with an error when no compatible
worker exists, ensuring behavior matches the KvRouter branch.

---

24-82: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Validate preselected_worker against required_taints.

This method now threads routing_constraints through the normal selection path, but the preselected_worker branch still bypasses that check entirely. A request with hard constraints can therefore bootstrap against an incompatible prefill worker. Reuse the pinned-worker compatibility check here, or reject preselected_worker when the constraints cannot be validated.

🤖 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/llm/src/kv_router/prefill_router/execution.rs` around lines 24 - 82,
resolve_prefill_worker currently skips validating a provided preselected_worker
against routing constraints/required_taints, allowing incompatible bootstraps;
update the preselected_worker branch in resolve_prefill_worker to perform the
same compatibility check used by query_prefill_worker (reuse the pinned-worker
compatibility logic or call the same validator) using routing_constraints (and
required_taints) derived from req.routing and req.block_mm_routing_info(), and
if validation fails return PrefillResolveDecision::Unavailable (or the same
error path as query_prefill_worker) instead of accepting the worker; reference
resolve_prefill_worker, preselected_worker, routing_constraints,
required_taints, query_prefill_worker, and PrefillResolveDecision in your
change.

lib/llm/src/kv_router.rs (1)

690-716: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Restore the old find_best_match entry point.

This adds RoutingConstraints to a public method, so existing out-of-tree Rust callers will stop compiling. Please keep the old signature as a default-constraints wrapper and move the taints-aware path to a separate method name.

🤖 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/llm/src/kv_router.rs` around lines 690 - 716, Restore the original public
API by keeping the old find_best_match signature (without RoutingConstraints) as
a thin wrapper that calls the existing taints-aware implementation (rename that
implementation to find_best_match_with_constraints or similar) and passes a
default RoutingConstraints value; specifically, keep a public fn
find_best_match(&self, context_id: Option<&str>, tokens: &[u32], block_mm_infos:
Option<&[Option<BlockExtraInfo>]>, router_config_override:
Option<&RouterConfigOverride>, update_states: bool, lora_name: Option<String>,
priority_jump: f64, expected_output_tokens: Option<u32>, allowed_worker_ids:
Option<HashSet<WorkerId>>) -> anyhow::Result<(WorkerWithDpRank, u32)> that
forwards to find_best_match_with_constraints(..., routing_constraints:
RoutingConstraints::default()), and rename the current method that accepts
RoutingConstraints to the new taints-aware name (e.g.,
find_best_match_with_constraints) so callers who need taint behavior can use it
while existing callers keep compiling.

lib/kv-router/src/scheduling/selector.rs (1)

287-296: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Apply preferred-taint bias independently of the base cost.

Multiplying base_score by the preferred-taint multiplier makes soft preferences a no-op whenever the computed cost is 0.0, so fully idle or fully cached candidates can still tie onto an explicitly avoided taint. This needs an additive bias or another tie-breaker that still differentiates zero-cost workers.

🤖 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/selector.rs` around lines 287 - 296, The
get_score closure currently multiplies base_score by preferred_taint_multiplier,
which makes taint preferences ineffective when base_score == 0. Update the
get_score closure (and related logic around worker_logit and
request.routing_constraints.preferred_taint_multiplier) to apply the taint
preference as an additive bias independent of base_score — e.g., compute
base_score = self.worker_logit(...), compute taint_mult =
request.routing_constraints.preferred_taint_multiplier(config.taints()), compute
a small additive bias (constant or epsilon-scaled value) derived from taint_mult
(for example bias = (taint_mult - 1.0) * TAINT_BIAS) and return base_score +
bias so zero-cost workers are still differentiated by taint preference.

lib/kv-router/src/protocols.rs (1)

174-180: ⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Non-finite weight validation still missing.

Serde will deserialize NaN and ±infinity values into preferred_taints without validation. When these flow into preferred_taint_multiplier(), they propagate through tanh and exp, producing all-NaN scores that can panic downstream selection logic (e.g., temperature-0 argmax).

Add a custom Deserialize impl or validation method that rejects or neutralizes non-finite weights at the API boundary.

🛡️ Example validation approach

impl RoutingConstraints {
    pub fn validate(&self) -> Result<(), &'static str> {
        for (taint, weight) in &self.preferred_taints {
            if !weight.is_finite() {
                return Err("preferred_taints must have finite weights");
            }
        }
        Ok(())
    }
}

Or filter at construction/deserialization time.

🤖 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/protocols.rs` around lines 174 - 180, The preferred_taints
map in struct RoutingConstraints can contain non-finite floats (NaN/±inf) from
serde deserialization; add validation to reject or sanitize these values at the
API boundary by either implementing a custom Deserialize for RoutingConstraints
or adding a validate() method called right after deserialization that iterates
over self.preferred_taints and returns an error if any weight is not finite;
ensure callers (or the deserializer) enforce this so
preferred_taint_multiplier() never receives non-finite inputs.

lib/kv-router/src/protocols.rs (1)

174-215: ⚡ Quick win

Consider adding unit tests for RoutingConstraints methods.

The new RoutingConstraints methods (is_empty, has_hard_constraints, is_compatible_with_worker_taints, preferred_taint_matches, preferred_taint_multiplier) have no test coverage. Adding tests would document expected behavior, catch edge cases (empty sets, boundary conditions, formula correctness), and prevent regressions.

📝 Example test structure

#[cfg(test)]
mod routing_constraints_tests {
    use super::*;

    #[test]
    fn test_is_compatible_with_worker_taints() {
        let constraints = RoutingConstraints {
            required_taints: ["gpu", "us-west"].iter().map(|s| s.to_string()).collect(),
            ..Default::default()
        };
        
        let worker_taints_match: HashSet<String> = 
            ["gpu", "us-west", "extra"].iter().map(|s| s.to_string()).collect();
        assert!(constraints.is_compatible_with_worker_taints(&worker_taints_match));
        
        let worker_taints_missing: HashSet<String> = 
            ["gpu"].iter().map(|s| s.to_string()).collect();
        assert!(!constraints.is_compatible_with_worker_taints(&worker_taints_missing));
    }
    
    // Add tests for other methods...
}

🤖 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/protocols.rs` around lines 174 - 215, Add unit tests for
RoutingConstraints covering is_empty, has_hard_constraints,
is_compatible_with_worker_taints, preferred_taint_matches, and
preferred_taint_multiplier: create a #[cfg(test)] mod with tests that verify
empty/default constraints, constraints with required_taints that both match and
don't match worker_taints (use
RoutingConstraints::is_compatible_with_worker_taints and has_hard_constraints),
verify preferred_taint_matches returns correct counts for overlapping keys, and
assert preferred_taint_multiplier produces expected numeric behavior for known
weights (include cases like no preferred_taints -> multiplier ==
(-0.0f64).exp(), single weight, and multiple weights to validate tanh->exp
flow). Reference the RoutingConstraints struct and its methods by name when
locating code to test.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes