[Serve] update input arg for choose_replicas to be a list of candidate replicas (#53456)

Signed-off-by: Gene Su <[email protected]>

Author: GeneDer

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

@@ -51,7 +51,7 @@ class PowerOfTwoChoicesRequestRouter(
 
     async def choose_replicas(
         self,
-        replicas_ranks: List[List[RunningReplica]],
+        candidate_replicas: List[RunningReplica],
         pending_request: Optional[PendingRequest] = None,
     ) -> List[List[RunningReplica]]:
         """One iteration of the power of two choices procedure that chooses
@@ -84,6 +84,6 @@ async def choose_replicas(
             k=min(2, len(candidate_replica_ids)),
         )
         replica_id_to_replica_map = {
-            replica.replica_id: replica for replica in replicas_ranks[0]
+            replica.replica_id: replica for replica in candidate_replicas
         }
         return [[replica_id_to_replica_map[chosen_id] for chosen_id in chosen_ids]]

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

@@ -837,11 +837,11 @@ async def choose_replicas_with_backoff(
                         extra={"log_to_stderr": False},
                     )
 
-                replica_ranks = [list(self._replicas.values())]
+                replica_ranks = list(self._replicas.values())
                 chosen_replicas: List[
                     List[RunningReplica]
                 ] = await self.choose_replicas(
-                    replicas_ranks=replica_ranks,
+                    candidate_replicas=replica_ranks,
                     pending_request=pending_request,
                 )
                 for replicas in chosen_replicas:
@@ -1026,7 +1026,7 @@ def select_available_replicas(
     @abstractmethod
     async def choose_replicas(
         self,
-        replicas_ranks: List[List[RunningReplica]],
+        candidate_replicas: List[RunningReplica],
         pending_request: Optional[PendingRequest] = None,
     ) -> List[List[RunningReplica]]:
         """Chooses a subset of candidate replicas from available replicas.
@@ -1036,9 +1036,8 @@ async def choose_replicas(
         replica selection.
 
         Args:
-            replicas_ranks: A list of lists of replicas, where each inner list
-                represents a rank of replicas. The first rank is the most
-                preferred and the last rank is the least preferred.
+            candidate_replicas: A list of candidate replicas to be considered in the
+                policy.
             pending_request: The request to be routed. This is used to
                 determine which replicas are eligible for routing.
 

python/ray/serve/tests/test_api.py

@@ -74,10 +74,10 @@ async def __call__(self):
 class FakeRequestRouter(RequestRouter):
     async def choose_replicas(
         self,
-        replicas_ranks: List[List[RunningReplica]],
+        candidate_replicas: List[RunningReplica],
         pending_request: Optional[PendingRequest] = None,
     ) -> List[List[RunningReplica]]:
-        return replicas_ranks
+        return [candidate_replicas]
 
 
 @serve.deployment(request_router_class=FakeRequestRouter)

python/ray/serve/tests/unit/test_router.py

@@ -218,7 +218,7 @@ async def choose_replica_for_request(
 
     async def choose_replicas(
         self,
-        replicas_ranks: List[List[RunningReplica]],
+        candidate_replicas: List[RunningReplica],
         pending_request: Optional[PendingRequest] = None,
     ) -> List[List[RunningReplica]]:
         pass