[Serve] Refactor RequestMetadata and Query objects (#10483)

python/ray/serve/router.py

@@ -2,40 +2,31 @@ import asyncio
 import copy
 from collections import defaultdict, deque
 import time
-from typing import DefaultDict, List
+from typing import DefaultDict, List, Dict, Any, Optional
 import pickle
+from dataclasses import dataclass
 
 from ray.exceptions import RayTaskError
 
 import ray
 from ray import serve
 from ray.experimental import metrics
+from ray.serve.context import TaskContext
 from ray.serve.endpoint_policy import RandomEndpointPolicy
+from ray.serve.request_params import RequestMetadata
 from ray.serve.utils import logger, chain_future
 
 REPORT_QUEUE_LENGTH_PERIOD_S = 1.0
 
 
+@dataclass
 class Query:
-    def __init__(
-            self,
-            request_args,
-            request_kwargs,
-            request_context,
-            call_method="__call__",
-            shard_key=None,
-            async_future=None,
-            is_shadow_query=False,
-    ):
-        self.request_args = request_args
-        self.request_kwargs = request_kwargs
-        self.request_context = request_context
+    args: List[Any]
+    kwargs: Dict[Any, Any]
+    context: TaskContext
 
-        self.async_future = async_future
-
-        self.call_method = call_method
-        self.shard_key = shard_key
-        self.is_shadow_query = is_shadow_query
+    metadata: RequestMetadata
+    async_future: Optional[asyncio.Future] = None
 
     def __reduce__(self):
         return type(self).ray_deserialize, (self.ray_serialize(), )
@@ -56,27 +47,6 @@ class Query:
         return Query(**kwargs)
 
 
-def _make_future_unwrapper(client_futures: List[asyncio.Future],
-                           host_future: asyncio.Future):
-    """Distribute the result of host_future to each of client_future"""
-    for client_future in client_futures:
-        # Keep a reference to host future so the host future won't get
-        # garbage collected.
-        client_future.host_ref = host_future
-
-    def unwrap_future(_):
-        result = host_future.result()
-
-        if isinstance(result, list):
-            for client_future, result_item in zip(client_futures, result):
-                client_future.set_result(result_item)
-        else:  # Result is an exception.
-            for client_future in client_futures:
-                client_future.set_result(result)
-
-    return unwrap_future
-
-
 class Router:
     """A router that routes request to available workers."""
 
@@ -175,8 +145,7 @@ class Router:
             request_args,
             request_kwargs,
             request_context,
-            call_method=request_meta.call_method,
-            shard_key=request_meta.shard_key,
+            metadata=request_meta,
             async_future=asyncio.get_event_loop().create_future())
         async with self.flush_lock:
             self.endpoint_queues[endpoint].appendleft(query)
@@ -301,7 +270,7 @@ class Router:
         worker = self.replicas[backend_replica_tag]
         try:
             object_ref = worker.handle_request.remote(req.ray_serialize())
-            if req.is_shadow_query:
+            if req.metadata.is_shadow_query:
                 # No need to actually get the result, but we do need to wait
                 # until the call completes to mark the worker idle.
                 await asyncio.wait([object_ref])
@@ -351,7 +320,7 @@ class Router:
                 self._do_query(backend, backend_replica_tag, request))
 
             # For shadow queries, just ignore the result.
-            if not request.is_shadow_query:
+            if not request.metadata.is_shadow_query:
                 chain_future(future, request.async_future)
 
             worker_queue.appendleft(backend_replica_tag)