[Serve] Remove register_custom_serializer

python/ray/serve/router.py

@@ -37,13 +37,16 @@ def __init__(
         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)
 

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