[Serve] decouple routing primitives

[machichima]

Description

Based on the RFC #59792, to better support Parallel PD pattern for prefill-decode disaggregation, we can add two new primitives to Ray Serve's DeploymentHandle API: choose_replica() (an async context manager) and dispatch()

Main changes

Public API

python/ray/serve/handle.py

• Implemented DeploymentHandle.choose_replica(): Returns async context manager for replica selection
• Implemented DeploymentHandle.dispatch(): Dispatch request to selected replica
  • Added deployment ID validation to prevent dispatching to wrong deployment

python/ray/serve/exceptions.py

• Added ReplicaUnavailableError: Raised when selected replica is no longer available

Core Router Logic

python/ray/serve/_private/router.py

• Added AsyncioRouter.choose_replica(): Context manager to select and reserve a replica slot
• Added AsyncioRouter.dispatch(): Send request to a previously selected replica
• Add choose_replica() and dispatch() to SingletonThreadRouter and CurrentLoopRouter

python/ray/serve/_private/request_router/request_router.py

• Added RequestRouter.on_replica_result_finished(): Decrement queue length cache when reserved slot is released
• Updated replica selection logic to account for reserved slots when checking availability:

python/ray/serve/_private/request_router/replica_wrapper.py

• Added slot reservation tracking: _reserved_slots set to track active reservations
• Added reserve_slot() and release_slot() to reserve and release a slot and return unique token

Tests

python/ray/serve/tests/test_handle_1.py (integration test based on the examples showed in #59792)

• Added test_choose_replica_and_dispatch_single(): Tests basic single selection pattern
• Added test_choose_replica_and_dispatch_parallel(): Tests parallel selection pattern (PD proxy use case) using AsyncExitStack

python/ray/serve/tests/unit/test_router.py (unit test)

• Add tests under TestChooseReplica for checking AsyncioRouter's choose_replica() and dispatch() behavior

Related issues

Related to #59792

Additional information

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

[gemini-code-assist]

Code Review

This pull request introduces new primitives to decouple replica selection from request dispatching in Ray Serve's routing logic. It adds a new ReplicaSelection data class and choose_replica and dispatch methods to the AsyncioRouter. This change enables a two-phase process where a replica can be chosen and a slot reserved before the request is actually sent.

While this is a good step towards more flexible routing, I've found a few critical issues in the current implementation that will cause runtime errors. Specifically, there are calls to attributes and methods that don't exist (_deployment_handle on AsyncioRouter, and reserve_slot, release_slot, send_request_with_slot on RunningReplica). Additionally, there's a logic error in checking for replica availability that would cause all dispatches to fail. Please see my detailed comments for suggestions on how to fix these issues.

[jeffreywang-anyscale]

Going to break this PR into 2-3 PRs.

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

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

[cursor]

Slot leak when CancelledError interrupts dispatch cleanup

Medium Severity

When _dispatch_to_marked_selection catches asyncio.CancelledError via except BaseException, the subsequent await selection._release_slot(force=True) can itself raise CancelledError (since the task is still cancelled). This prevents the slot from being released. The choose_replica finally block won't help because _dispatched is already True, causing _release_slot() without force to return None. The reserved semaphore slot on the replica leaks until the actor restarts.

Additional Locations (1)

• python/ray/serve/_private/router.py#L1353-L1364

 

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

[cursor]

Unnecessary remote call on every dispatch completion

Medium Severity

The _release_slot_if_still_reserved done callback fires on every completed dispatch, unconditionally calling _release_slot(force=True) which always issues a remote actor call to release_slot. In the normal happy path (the vast majority of requests), _start_request has already consumed the token, so this remote call is a guaranteed no-op that adds latency and actor mailbox pressure. For the PD disaggregation use case (latency-sensitive), this adds overhead proportional to request throughput.

Additional Locations (1)

• python/ray/serve/_private/request_router/replica_wrapper.py#L447-L455

 

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

[jeffreywang-anyscale]

Closing this as the following PRs are merged:

• [serve][1/3] Add replica-side slot reservation primitive #63252
• [serve][2/3] Add choose_replica/dispatch to AsyncioRouter #63254
• [serve][3/3] Expose choose_replica/dispatch on deployment handles #63255

Thanks again @machichima!