[serve][3/3] Expose choose_replica/dispatch on deployment handles

[jeffreywang-anyscale]

Description

Plumb choose_replica and dispatch through DeploymentHandle, not just the underlying AsyncioRouter which #63254 ends at.

Deployment handle

• DeploymentHandle.choose_replica() / dispatch(selection, ...) validates that the selection belongs to this deployment.

SingletonThreadRouter / CurrentLoopRouter

• SingletonThreadRouter.choose_replica: bridges the async context manager onto the router thread via run_coroutine_threadsafe, so all state mutations happen on the singleton router loop.
• SingletonThreadRouter.dispatch and CurrentLoopRouter.dispatch: both call selection._mark_dispatched() synchronously before scheduling the dispatch coroutine, so the choose_replica __aexit__ cleanup can't race with the dispatch task. (The race would otherwise double-decrement the queue-length cache on CurrentLoopRouter, where __aexit__ runs synchronously without yielding.)
• CurrentLoopRouter.choose_replica: thin async-with passthrough since it shares an event loop with the caller.

Test-only router

• LocalRouter placeholder in local_testing_mode: choose_replica / dispatch raise NotImplementedError because the local-testing path invokes user callables directly and doesn't exercise router selection.

Benchmark

The overhead of choose_replica() + dispatch() over remote() is constant. This is consistent with the one extra run_coroutine_threadsafe round-trip: remote does one thread crossing, choose_dispatch does two (one for the __aenter__ bridge, one for the _dispatch_to_marked_selection task). Replica count doesn't change the delta, as the overhead is in the bridge, not routing.

Related issues

RFC: #59792
Original PR: #60865
Next PR: Metric changes #62356

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

[gemini-code-assist]

Code Review

This pull request introduces a two-step routing mechanism in Ray Serve, adding choose_replica and dispatch methods to the Router interface and DeploymentHandle. This allows for explicit replica selection and slot reservation prior to request execution, supporting both standard and streaming workflows. The implementation includes logic to bridge these calls across event loops in SingletonThreadRouter and provides validation to ensure selections are dispatched to the correct deployment. Feedback identifies a critical resource leak vulnerability in SingletonThreadRouter.choose_replica where cancellation during context entry could result in leaked request slots. It is recommended to wrap the context entry in a try...finally block to ensure proper cleanup even if the task is cancelled early.

[cursor]

Race causes double-fire of on_request_completed callback

Medium Severity

CurrentLoopRouter.dispatch and SingletonThreadRouter.dispatch schedule _dispatch_to_marked_selection as a task/future that hasn't executed yet when they return. If the choose_replica context exits before that task completes, the AsyncioRouter.choose_replica finally block sees _dispatched=True but _completion_callback_registered=False, and fires on_request_completed. Later, when the dispatch task actually runs and registers its done-callback, on_request_completed fires a second time upon request completion — violating the "exactly once" invariant stated in the comment.

Additional Locations (2)

• python/ray/serve/_private/router.py#L1761-L1766
• python/ray/serve/_private/router.py#L1307-L1315

 

^{Reviewed by Cursor Bugbot for commit 874103a. Configure here.}

[kouroshHakha]

Overall the design is sound — the slot-reservation contract, the _mark_dispatched synchronous flag to prevent the choose_replica finally from racing the dispatch task, and the _completion_callback_registered exactly-once invariant are all well-thought-out. The CurrentLoopRouter wrapper is clean. One critical resource safety issue in the SingletonThreadRouter bridge needs fixing before merge; everything else is minor.

The open Gemini bot comment (on SingletonThreadRouter.choose_replica) is still unaddressed — I've investigated it and believe it's valid; details in the inline comment.

Note

This review was co-written with AI assistance (Claude Code).

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 3a4056c. Configure here.}

[kouroshHakha]

Before we merge, I'd like to understand the overhead of using choose_replica() + dispatch() as a drop-in for .remote(). Your other LLM benchmarks suggest parity, but it would be good to confirm that here with a small Serve microbenchmark.

The interesting case is SingletonThreadRouter (handle called from outside a deployment): .remote() does one run_coroutine_threadsafe round-trip, whereas choose_replica + dispatch does two — one for __aenter__ and one for _dispatch_to_marked_selection. If those thread crossings dominate routing latency you could see a meaningful delta.

Concretely, could you add a choose_dispatch mode to python/ray/serve/_private/benchmarks/handle_noop_latency.py — something like a second path in Benchmarker that does:

async with self._handle.choose_replica() as sel:
    await self._handle.dispatch(sel)

and run both modes (remote vs choose_dispatch) at e.g. num_replicas=1 and num_replicas=4? Even a single table of p50/p99 latency numbers pasted in the PR description would be enough — the goal is just to establish a documented baseline so that if someone later refactors assign_request to use these primitives, there's a clear regression target to compare against.

[jeffreywang-anyscale]

Going to follow up nightly benchmarks with #63293.

[kouroshHakha]

ok so we cannot just replace .remote() with choose + dispatch.