[serve][P0] fix await deployment response hanging issue (#52561)

## Why are these changes needed?

If the user passes a `DeploymentResponse` object directly into a
deployment handle, we will resolve it to the object ref under the hood
before scheduling the request. But we were incorrectly using a
threading.Lock when getting the object ref so this can cause deadlock if
the same `DeploymentResponse` object is used for multiple downstream
calls. Instead we should be using an asyncio lock.

## Related issue number

closes #51190

---------

Signed-off-by: Cindy Zhang <cindyzyx9@gmail.com>

Author: zcin

python/ray/serve/_private/replica_result.py

@@ -1,3 +1,4 @@
+import asyncio
 import inspect
 import threading
 import time
@@ -62,6 +63,7 @@ def __init__(
         self._is_streaming: bool = metadata.is_streaming
         self._request_id: str = metadata.request_id
         self._object_ref_or_gen_sync_lock = threading.Lock()
+        self._lazy_object_ref_or_gen_asyncio_lock = None
 
         if isinstance(obj_ref_or_gen, ray.ObjectRefGenerator):
             self._obj_ref_gen = obj_ref_or_gen
@@ -73,6 +75,14 @@ def __init__(
                 self._obj_ref_gen is not None
             ), "An ObjectRefGenerator must be passed for streaming requests."
 
+    @property
+    def _object_ref_or_gen_asyncio_lock(self) -> asyncio.Lock:
+        """Lazy `asyncio.Lock` object."""
+        if self._lazy_object_ref_or_gen_asyncio_lock is None:
+            self._lazy_object_ref_or_gen_asyncio_lock = asyncio.Lock()
+
+        return self._lazy_object_ref_or_gen_asyncio_lock
+
     def _process_response(f: Union[Callable, Coroutine]):
         @wraps(f)
         def wrapper(self, *args, **kwargs):
@@ -174,7 +184,7 @@ async def to_object_ref_async(self) -> ray.ObjectRef:
         # object ref cached in order to avoid calling `__anext__()` to
         # resolve to the underlying object ref more than once.
         # See: https://github.com/ray-project/ray/issues/43879.
-        with self._object_ref_or_gen_sync_lock:
+        async with self._object_ref_or_gen_asyncio_lock:
             if self._obj_ref is None:
                 self._obj_ref = await self._obj_ref_gen.__anext__()
 

python/ray/serve/tests/test_handle_1.py

@@ -2,6 +2,7 @@
 import concurrent.futures
 import sys
 import threading
+from typing import Any
 
 import pytest
 
@@ -311,5 +312,32 @@ def f():
         h._init(_prefer_local_routing=True)
 
 
+def test_response_used_in_multiple_calls(serve_instance):
+    @serve.deployment(graceful_shutdown_timeout_s=0)
+    class F:
+        async def __call__(self, sleep_amt: int, x: Any):
+            await asyncio.sleep(sleep_amt)
+            return f"({x})"
+
+    @serve.deployment(graceful_shutdown_timeout_s=0)
+    class Ingress:
+        async def __init__(self, h):
+            self.h = h
+
+        async def __call__(self):
+            # r1 will take 5 seconds to finish. This makes sure when h.remote() is
+            # started for r2 and r3 (and both rely on r1), r1 is still executing.
+            r1 = self.h.remote(5, "r1")
+
+            # Neither of these should get stuck.
+            r2 = self.h.remote(0, r1)
+            r3 = self.h.remote(0, r1)
+
+            return await r2, await r3
+
+    h = serve.run(Ingress.bind(F.bind()))
+    assert h.remote().result(timeout_s=10) == ("((r1))", "((r1))")
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", "-s", __file__]))