diff --git a/python/ray/serve/router.py b/python/ray/serve/router.py index 501e04453..825182d7d 100644 --- a/python/ray/serve/router.py +++ b/python/ray/serve/router.py @@ -37,13 +37,16 @@ class Query: self.shard_key = shard_key self.is_shadow_query = is_shadow_query + def __reduce__(self): + return type(self).ray_deserialize, (self.ray_serialize(), ) + def ray_serialize(self): # NOTE: this method is needed because Query need to be serialized and # sent to the replica worker. However, after we send the query to # replica worker the async_future is still needed to retrieve the final # result. Therefore we need a way to pass the information to replica # worker without removing async_future. - clone = copy.copy(self).__dict__ + clone = copy.copy(self.__dict__) clone.pop("async_future") return pickle.dumps(clone) diff --git a/python/ray/serve/tests/test_router.py b/python/ray/serve/tests/test_router.py index 3e4514126..3634b11e7 100644 --- a/python/ray/serve/tests/test_router.py +++ b/python/ray/serve/tests/test_router.py @@ -162,12 +162,6 @@ async def test_shard_key(serve_instance, task_runner_mock_actor): async def test_router_use_max_concurrency(serve_instance): - # The VisibleRouter::get_queues method needs to pickle queries - # so we register serializer here. In regular code path, query - # serialization is done by Serve manually for performance. - ray.register_custom_serializer(Query, Query.ray_serialize, - Query.ray_deserialize) - signal = SignalActor.remote() @ray.remote