Repository navigation
fix(kv-router): keep custom selection policies on hard session bindings - #15661
ishandhanani wants to merge 4 commits into
Conversation
…licies In the selection core, a Hard session binding is exclusive only for the default selector; custom worker-selection policies receive it as advisory context. A custom policy that chose another eligible worker had its booking rejected by commit_session. When the core is configured with Hard session affinity, build every policy with exclusive affinity so the scheduler limits candidates to an eligible bound worker, falling back to normal selection when that worker is ineligible. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
Selection narrowed an exclusive selector to an eligible affinity target, but the prefill-busy admission checks still considered every eligible worker. A request bound to a busy worker could be admitted because another worker was free, then booked onto the busy bound worker, bypassing the class prefill threshold. Apply the same narrowing in admission and selection so the request waits for its bound worker, as an explicit pin already does. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
ishandhanani
changed the title
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (5)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review. WalkthroughHard session affinity now makes an eligible session target exclusive across worker-selection policies and queue eligibility checks. Tests cover custom-policy selection and waiting for a busy affinity target until capacity changes. ChangesHard Session Affinity
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to No confirmed Hard-affinity issue remains that would prevent merging. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
![dynamo-review-agent[bot]](https://avatars.githubusercontent.com/u/201626793?s=60&v=4){kind=small}
| self.selector.uses_exclusive_affinity_target(), | ||
| &configs, | ||
| &request, | ||
| request |
There was a problem hiding this comment.
The new admission-side affinity narrowing uses request.eligibility() without the authoritative overloaded-worker exclusion that final selection applies. If the affinity target is overloaded and prefill-busy while another eligible worker is idle, admission incorrectly narrows to the target and parks the request even though selection considers that target ineligible. Apply overload eligibility consistently to all three admission checks.
🤖 AI Fix
Fetch the overloaded worker set and pass eligibility_with_overloaded(...) into with_exclusive_affinity in handle_enqueue, has_dispatchable_ready_head, and drain_ready.
| .request | ||
| .eligibility() | ||
| .with_available_workers(available.as_deref()), | ||
| Self::with_exclusive_affinity( |
There was a problem hiding this comment.
An exclusive-affinity request remains in the shared WorkerPlacement::Any heap, but this new check makes its head non-dispatchable whenever its bound worker is busy. PolicyClassQueue::next_dispatchable examines only that shared head, so later unbound requests in the same policy class cannot use otherwise-idle workers until the affinity target frees or the head expires. Give exclusive-affinity requests a target-specific queue lane, analogous to pinned requests, while preserving worker-wide affinity when no DP rank is specified.
🤖 AI Fix
Queue eligible exclusive-affinity requests in a target-specific lane that can be blocked and rechecked independently of the shared Any lane.
| .request | ||
| .eligibility() | ||
| .with_available_workers(available.as_deref()), | ||
| Self::with_exclusive_affinity( |
There was a problem hiding this comment.
Exclusive-affinity requests remain in the shared WorkerPlacement::Any lane. When the affinity target is busy, this new check marks the shared head nondispatchable; PolicyQueue only examines that head, so later unconstrained requests cannot reach otherwise idle workers until the affinity target frees. Put exclusive-affinity requests in affinity-aware lanes, or allow dispatch to skip a blocked constrained shared head.
🤖 AI Fix
Queue exclusive-affinity requests in affinity-aware worker/rank lanes so a busy bound target cannot block later shared requests for idle workers.
| assert_eq!(bound_worker(&core, "s"), Some(other)); | ||
| } | ||
|
|
||
| /// Picks any candidate other than the session's bound worker. |
There was a problem hiding this comment.
This private test-helper comment merely expands AvoidAffinityPicker into a sentence; the type name and pick implementation already make the behavior explicit.
🤖 AI Fix
Remove this doc comment.
❌ Dynamo PR CI failed — run 37424955074 (attempt 1) on
|
| Other | Jobs |
|---|---|
| dynamo-runtime | ❌ 2 ✅ 6 |
Failure details
2 jobs failed (amd64 and arm64 runtime CPU tests) with the same 2 tests: test_drained_stream_can_require_generated_content and test_drained_stream_accepts_generated_content in tests/fault_tolerance/cancellation/test_utils.py time out after 30s waiting for the HTTP response on both architectures. No cancelled jobs.
❌ dynamo-runtime / test / parallel cuda13.0, amd64: 2 cancellation test_utils tests time out (requests Timeout 30s)
Job: dynamo-runtime / test / parallel cuda13.0, amd64 · Failed step: Run CPU-only tests (parallelized) · Logs: gh run view --job 112154625552 -R ai-dynamo/dynamo --log-failed
FAILED tests/fault_tolerance/cancellation/test_utils.py::test_drained_stream_can_require_generated_content - requests.exceptions.Timeout: Timed out after 30.0s waiting for the HTTP response
FAILED tests/fault_tolerance/cancellation/test_utils.py::test_drained_stream_accepts_generated_content - requests.exceptions.Timeout: Timed out after 30.0s waiting for the HTTP response
tests/fault_tolerance/cancellation/test_utils.py:50: in test_drained_stream_can_require_generated_content
read_streaming_responses(
tests/fault_tolerance/cancellation/utils.py:309: in read_streaming_responses
response_raw = cancellable_req.get_response()
tests/fault_tolerance/cancellation/utils.py:188: in get_response
raise requests.exceptions.Timeout(
E requests.exceptions.Timeout: Timed out after 30.0s waiting for the HTTP response
=== 2 failed, 789 passed, 23 skipped, 5398 deselected in 1129.32s (0:18:49) ===
##[error]The path for one of the files in artifact is not valid: /test_unsafe_or_nonexact_protobuf_pins[protobuf==6.33.*]/test.log.txt. Contains the following character: Asterisk *
Both tests/fault_tolerance/cancellation/test_utils.py drained-stream tests fail in every attempt, including all 3 auto-retries: CancellableRequest.get_response() (utils.py:188) raises requests.exceptions.Timeout after 30s because no HTTP response arrives. The other 789 tests pass. The artifact-upload error after that (the * in a per-test log path) is a follow-on step failure and does not affect the test result.
❌ dynamo-runtime / test / parallel cuda13.0, arm64: 2 cancellation test_utils tests time out (requests Timeout 30s)
Job: dynamo-runtime / test / parallel cuda13.0, arm64 · Failed step: Run CPU-only tests (parallelized) · Logs: gh run view --job 112154625579 -R ai-dynamo/dynamo --log-failed
FAILED tests/fault_tolerance/cancellation/test_utils.py::test_drained_stream_can_require_generated_content - requests.exceptions.Timeout: Timed out after 30.0s waiting for the HTTP response
FAILED tests/fault_tolerance/cancellation/test_utils.py::test_drained_stream_accepts_generated_content - requests.exceptions.Timeout: Timed out after 30.0s waiting for the HTTP response
tests/fault_tolerance/cancellation/test_utils.py:50: in test_drained_stream_can_require_generated_content
read_streaming_responses(
tests/fault_tolerance/cancellation/utils.py:309: in read_streaming_responses
response_raw = cancellable_req.get_response()
tests/fault_tolerance/cancellation/utils.py:188: in get_response
raise requests.exceptions.Timeout(
E requests.exceptions.Timeout: Timed out after 30.0s waiting for the HTTP response
=== 2 failed, 789 passed, 23 skipped, 5398 deselected in 1129.32s (0:18:49) ===
##[error]The path for one of the files in artifact is not valid: /test_unsafe_or_nonexact_protobuf_pins[protobuf==6.33.*]/test.log.txt. Contains the following character: Asterisk *
Both tests/fault_tolerance/cancellation/test_utils.py drained-stream tests fail in every attempt, including all 3 auto-retries: CancellableRequest.get_response() (utils.py:188) raises requests.exceptions.Timeout after 30s because no HTTP response arrives. The other 789 tests pass. The artifact-upload error after that (the * in a per-test log path) is a follow-on step failure and does not affect the test result.
For agents
{"pr": 15661, "run_id": 37424955074, "run_attempt": 1, "head_sha": "437e4a819be8402fbac3ad942188673f3d26dbd4", "failures": [{"job": "dynamo-runtime / test / parallel cuda13.0, amd64", "job_id": 112154625552, "failed_step": "Run CPU-only tests (parallelized)", "signature": "requests.exceptions.Timeout: Timed out after 30.0s waiting for the HTTP response", "tests": ["tests/fault_tolerance/cancellation/test_utils.py::test_drained_stream_can_require_generated_content", "tests/fault_tolerance/cancellation/test_utils.py::test_drained_stream_accepts_generated_content"], "log_cmd": "gh run view --job 112154625552 -R ai-dynamo/dynamo --log-failed"}, {"job": "dynamo-runtime / test / parallel cuda13.0, arm64", "job_id": 112154625579, "failed_step": "Run CPU-only tests (parallelized)", "signature": "requests.exceptions.Timeout: Timed out after 30.0s waiting for the HTTP response", "tests": ["tests/fault_tolerance/cancellation/test_utils.py::test_drained_stream_can_require_generated_content", "tests/fault_tolerance/cancellation/test_utils.py::test_drained_stream_accepts_generated_content"], "log_cmd": "gh run view --job 112154625579 -R ai-dynamo/dynamo --log-failed"}]}Posted automatically by Devin for run 37424955074. Updated on every full-CI run of this PR.
| .request | ||
| .eligibility() | ||
| .with_available_workers(available.as_deref()), | ||
| Self::with_exclusive_affinity( |
There was a problem hiding this comment.
Affinity-bound requests stay in the shared Any lane (placement only looks at pinned_worker, L1245), and next_dispatchable only peeks the shared head (policy_queue.rs:355). So a request waiting on its busy bound worker now stalls later unbound requests in the same class even while other workers are idle. Pinned requests don't do this since they get their own lane, so "matching how an explicit pinned_worker already waits" in the description isn't quite accurate.
This is the same head-of-line limitation the TODO above already acknowledges for allowlist/taint-constrained heads, but this PR widens it to most session traffic on the standalone service when queueing is on. Needs a human call on how to defer it (e.g. a ticket for per-target lanes linked from the TODO, a description fix, an operator note).
There was a problem hiding this comment.
Agreed the admission narrowing was the wrong approach. Reverted it in 437e4a8: a hard-bound request now passes the normal fleet-wide admission check and queues on its bound worker, so it no longer sits at the head of the shared lane waiting for that worker. Description updated.
| } | ||
| // Hard affinity rejects a booking dispatched away from an eligible bound worker, so every | ||
| // policy, not only the default selector, must keep to that worker. | ||
| let worker_selection_policy_factory = match session_affinity { |
There was a problem hiding this comment.
Under Hard, if a custom policy's WorkerFilter rejects the bound worker, narrowing leaves it no candidates and it returns AllEligibleWorkersFiltered (503). That error path in run.rs:465 just drops the hold, which releases the binding and refreshes its TTL rather than invalidating it, so every retry fails the same way and the session never re-binds. Before this change the policy picked another worker, commit rejected once with 400 and invalidated the binding, so the next request recovered.
Suggest invalidating a Hold::Bound on AllEligibleWorkersFiltered under a Hard table (as hold_session already does for unschedulable targets), plus a test: second booking 503, third re-initializes on the other worker.
There was a problem hiding this comment.
Fixed in 303a1f7. Under Hard, AllEligibleWorkersFiltered now invalidates the bound hold, so the session re-binds on its next request. Added hard_mode_rebinds_when_a_filter_rejects_the_bound_worker: second booking fails, binding dropped, third booking re-binds on the other worker.
| request | ||
| .eligibility() | ||
| .with_available_workers(available.as_deref()), | ||
| Self::with_exclusive_affinity( |
There was a problem hiding this comment.
Admission narrows using request.eligibility(), which carries no overloaded set, while selection uses eligibility_with_overloaded and drops the narrowing for an overloaded target. So with an overloaded and prefill-busy bound worker, the request waits for a worker it won't be booked on (or gets QueueRejected in a never-queues class), where before it went to an idle worker. Reachable on frontend Soft + queueing, since only the frontend wires an overload provider. Could fetch the overload set once per drain pass, like active_tokens, and use it for the narrowing decision.
There was a problem hiding this comment.
No longer applies: the admission narrowing is reverted in 437e4a8, so admission is back to the fleet-wide check and does not narrow to an overloaded target.
| /// Narrow `eligibility` to the request's affinity target when the selector treats an eligible | ||
| /// target as exclusive. Admission and selection both apply it, so a request bound to a busy | ||
| /// worker waits for that worker instead of being admitted because another worker is free. | ||
| fn with_exclusive_affinity<'a>( |
There was a problem hiding this comment.
Scope note: this keys on selector exclusivity, not affinity mode, so it also reaches the frontend's Soft mode (routing_host/kv.rs:43 passes Soft bindings as affinity_target, and the default policy is exclusive). With queueing on, Soft requests now wait on a busy target, and in never-queues classes both Soft and standalone Hard requests get QueueRejected where they used to be booked. Worth updating the "frontend/Soft unchanged" lines in the description and adding a never-queues test; the new test only covers the drain-loop site.
There was a problem hiding this comment.
Reverted the admission narrowing in 437e4a8, so frontend Soft mode and never-queues classes are unchanged again. The remaining change keys on Hard only (the core wraps policies only when its session-affinity config is Hard), and the description now says so.
With Hard affinity narrowing selection to the bound worker, a custom policy filter that rejects that worker leaves no candidate and selection fails with AllEligibleWorkersFiltered. The hold was released rather than invalidated, so the binding survived and every retry failed the same way. Invalidate the bound hold on that error so the session re-binds on its next request. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
A hard-bound request should go to its bound worker and queue there, not wait in the router for that worker. Restore fleet-wide admission so the router holds such a request only when every eligible worker is busy, as for any other request. This also removes the head-of-line blocking, overload mismatch, and Soft-mode and never-queues side effects of the admission narrowing. This reverts commit 4c95c34. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
ishandhanani
changed the title
2 participants
Summary
With
--session-affinity-ttl-secs(Hard mode), the standalone selection service rejects a/select_and_reservebooking dispatched away from a session's live binding. The selection core passes that binding to worker selection asaffinity_target, but only the default selector treated it as exclusive. Custom worker-selection policies (dynamo-two-tier-cost-fn,thunderagent, and others linked through the policy catalog) received it as advisory context, so whenever a custom policy picked a different eligible worker,commit_sessionrejected the request.Changes:
SelectionCoreis built with Hard session affinity, it builds every policy withwith_exclusive_affinity(true). Selection is limited to an eligible bound worker and falls back to normal selection when that worker is ineligible, so a worker that departs between the hold and the selection still re-initializes the session throughcommit_session.AllEligibleWorkersFiltered. The binding is now invalidated on that error, so the session re-binds on its next request instead of failing every retry.Scope:
Raised in review of #15600: #15600 (comment)
Validation
hard_mode_keeps_custom_policies_on_the_binding: a custom picker that always avoids the session's bound worker keeps the session on that worker under Hard affinity. Without the fix the second booking is rejected.hard_mode_rebinds_when_a_filter_rejects_the_bound_worker: the second booking fails withAllEligibleWorkersFiltered, the binding is dropped, and the third booking re-binds on the other worker. Without the invalidation the binding survives.cargo test -p dynamo-kv-router --features standalone-selection --lib: 1,266 passed.cargo clippy -p dynamo-kv-router --features standalone-selection --all-targets -- -D warningsandcargo fmt --check: clean.