[serve][2/3] Add choose_replica/dispatch to AsyncioRouter (#63254)

Signed-off-by: Jeffrey Wang <jeffreywang@anyscale.com>
Co-authored-by: machichima <nary12321@gmail.com>

Author: jeffreywang88

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

@@ -348,7 +348,12 @@ async def reserve_slot(
             accepted, num_ongoing_requests = await obj_ref
         except asyncio.CancelledError:
             ray.cancel(obj_ref)
-            # TODO (#63254): Add error handling for actor-unavailable cleanup.
+            self._actor_handle.release_slot.remote(slot_token)
+            raise
+        except Exception:
+            # The actor may have reserved the slot before the reply was lost
+            # (e.g. ActorUnavailableError). `release_slot` is idempotent for unknown
+            # tokens, so this is safe even when the reservation never actually happened.
             self._actor_handle.release_slot.remote(slot_token)
             raise
 
@@ -408,6 +413,11 @@ class ReplicaSelection:
         default=False, init=False
     )  # Tracks if dispatch was called
 
+    # Set by dispatch once the result's done-callback is wired up. Read by
+    # choose_replica's finally to decide whether to fire on_request_completed
+    # manually (only one of the two paths should fire it).
+    _completion_callback_registered: bool = field(default=False, init=False)
+
     @property
     def address(self) -> str:
         """Returns the replica address in host:port format."""

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

@@ -24,7 +24,6 @@
     DeploymentHandleSource,
     DeploymentID,
     ReplicaID,
-    ReplicaQueueLengthInfo,
     RequestMetadata,
     RunningReplicaInfo,
 )
@@ -830,17 +829,11 @@ def _insert_pending_request_sorted(
             index += 1
         queue.insert(index, pending_request)
 
-    def on_new_queue_len_info(
-        self, replica_id: ReplicaID, queue_len_info: ReplicaQueueLengthInfo
-    ):
+    def on_new_queue_len_info(self, replica_id: ReplicaID, num_ongoing_requests: int):
         """Update queue length cache with new info received from replica."""
         if self._use_replica_queue_len_cache:
-            self._replica_queue_len_cache.update(
-                replica_id, queue_len_info.num_ongoing_requests
-            )
-            self._update_router_queue_len_gauge(
-                replica_id, queue_len_info.num_ongoing_requests
-            )
+            self._replica_queue_len_cache.update(replica_id, num_ongoing_requests)
+            self._update_router_queue_len_gauge(replica_id, num_ongoing_requests)
 
     def on_send_request(self, replica_id: ReplicaID):
         """Increment queue length cache when a request is sent to a replica."""