chore: import upstream snapshot with attribution

python/ray/serve/tests/test_handle_1.py

@@ -0,0 +1,735 @@
+import asyncio
+import concurrent.futures
+import sys
+import threading
+from contextlib import AsyncExitStack
+from typing import Any
+
+import pytest
+
+import ray
+from ray import serve
+from ray.serve._private.common import DeploymentHandleSource
+from ray.serve._private.constants import (
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    SERVE_DEFAULT_APP_NAME,
+)
+from ray.serve.api import get_replica_context
+from ray.serve.exceptions import RayServeException
+from ray.serve.handle import DeploymentHandle, DeploymentResponseGenerator
+
+
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode doesn't set handle source",
+)
+def test_replica_handle_source(serve_instance):
+    @serve.deployment
+    def f():
+        return "hi"
+
+    @serve.deployment
+    class Router:
+        def __init__(self, handle):
+            self.handle = handle
+            self.handle._init()
+
+        def check(self):
+            return self.handle.init_options._source == DeploymentHandleSource.REPLICA
+
+    h = serve.run(Router.bind(f.bind()))
+    assert h.check.remote().result()
+
+
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode work with tasks & actors",
+)
+def test_handle_serializable(serve_instance):
+    @serve.deployment
+    def f():
+        return "hello"
+
+    @ray.remote
+    class DelegateActor:
+        async def call_handle(self, handle):
+            return await handle.remote()
+
+    @serve.deployment
+    class Ingress:
+        def __init__(self, handle):
+            self._handle = handle
+
+        async def __call__(self):
+            # Test pickling handle via `actor.method.remote()`.
+            a = DelegateActor.remote()
+            return await a.call_handle.remote(self._handle)
+
+    app_handle = serve.run(Ingress.bind(f.bind()))
+    assert app_handle.remote().result() == "hello"
+
+
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode doesn't support get_app_handle/get_deployment_handle",
+)
+def test_get_and_call_handle_in_thread(serve_instance):
+    @serve.deployment
+    def f():
+        return "hello"
+
+    serve.run(f.bind())
+
+    def get_and_call_app_handle():
+        handle = serve.get_app_handle(SERVE_DEFAULT_APP_NAME)
+        return handle.remote().result()
+
+    def get_and_call_deployment_handle():
+        handle = serve.get_deployment_handle("f", SERVE_DEFAULT_APP_NAME)
+        return handle.remote().result()
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+        fut1 = executor.submit(get_and_call_app_handle)
+        fut2 = executor.submit(get_and_call_deployment_handle)
+        assert fut1.result() == "hello"
+        assert fut2.result() == "hello"
+
+
+def test_handle_option_chaining(serve_instance):
+    # https://github.com/ray-project/ray/issues/12802
+    # https://github.com/ray-project/ray/issues/12798
+
+    @serve.deployment
+    class MultiMethod:
+        def method_a(self):
+            return "method_a"
+
+        def method_b(self):
+            return "method_b"
+
+        def __call__(self):
+            return "__call__"
+
+    handle1 = serve.run(MultiMethod.bind())
+    counter = handle1.request_counter
+    counter_info = counter.info
+    assert handle1.remote().result() == "__call__"
+
+    handle2 = handle1.options(method_name="method_a")
+
+    assert handle2.remote().result() == "method_a"
+    assert handle2.request_counter == counter
+    assert handle2.request_counter.info == counter_info
+
+    handle3 = handle1.options(method_name="method_b")
+
+    assert handle3.remote().result() == "method_b"
+    assert handle3.request_counter == counter
+    assert handle2.request_counter.info == counter_info
+
+
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode doesn't support get_app_handle/get_deployment_handle",
+)
+def test_repeated_get_handle_cached(serve_instance):
+    @serve.deployment
+    def f(_):
+        return ""
+
+    serve.run(f.bind())
+
+    handle_sets = {serve.get_app_handle("default") for _ in range(100)}
+    assert len(handle_sets) == 1
+
+    handle_sets = {serve.get_deployment_handle("f", "default") for _ in range(100)}
+    assert len(handle_sets) == 1
+
+
+def test_nonexistent_method_sync(serve_instance):
+    @serve.deployment
+    class A:
+        def exists(self):
+            pass
+
+    handle = serve.run(A.bind())
+    with pytest.raises(RayServeException) as excinfo:
+        handle.does_not_exist.remote().result()
+
+    exception_string = str(excinfo.value)
+    assert "'does_not_exist'" in exception_string
+    assert "Available methods: ['exists']" in exception_string
+
+
+@pytest.mark.asyncio
+async def test_nonexistent_method_async(serve_instance):
+    @serve.deployment
+    class A:
+        def exists(self):
+            pass
+
+    handle = serve.run(A.bind())
+    with pytest.raises(RayServeException) as excinfo:
+        await handle.does_not_exist.remote()
+
+    exception_string = str(excinfo.value)
+    assert "'does_not_exist'" in exception_string
+    assert "Available methods: ['exists']" in exception_string
+
+
+def _get_asyncio_loop_running_in_thread() -> asyncio.AbstractEventLoop:
+    loop = asyncio.new_event_loop()
+    threading.Thread(
+        daemon=True,
+        target=loop.run_forever,
+    ).start()
+    return loop
+
+
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode doesn't support get_app_handle/get_deployment_handle",
+)
+@pytest.mark.asyncio
+async def test_call_handle_across_asyncio_loops(serve_instance):
+    @serve.deployment
+    class A:
+        def exists(self):
+            return True
+
+    serve.run(A.bind())
+
+    async def refresh_get():
+        handle = serve.get_app_handle("default")
+        assert await handle.exists.remote()
+
+    for _ in range(10):
+        loop = _get_asyncio_loop_running_in_thread()
+        asyncio.run_coroutine_threadsafe(refresh_get(), loop).result()
+
+    handle = serve.get_app_handle("default")
+
+    async def cache_get():
+        assert await handle.exists.remote()
+
+    for _ in range(10):
+        loop = _get_asyncio_loop_running_in_thread()
+        asyncio.run_coroutine_threadsafe(cache_get(), loop).result()
+
+
+def test_handle_typing(serve_instance):
+    @serve.deployment
+    class DeploymentClass:
+        pass
+
+    @serve.deployment
+    def deployment_func():
+        pass
+
+    @serve.deployment
+    class Ingress:
+        def __init__(
+            self, class_downstream: DeploymentHandle, func_downstream: DeploymentHandle
+        ):
+            # serve.run()'ing this deployment fails if these assertions fail.
+            assert isinstance(class_downstream, DeploymentHandle)
+            assert isinstance(func_downstream, DeploymentHandle)
+
+    h = serve.run(Ingress.bind(DeploymentClass.bind(), deployment_func.bind()))
+    assert isinstance(h, DeploymentHandle)
+
+
+def test_call_function_with_argument(serve_instance):
+    @serve.deployment
+    def echo(name: str):
+        return f"Hi {name}"
+
+    @serve.deployment
+    class Ingress:
+        def __init__(self, h: DeploymentHandle):
+            self._h = h
+
+        async def __call__(self, name: str):
+            return await self._h.remote(name)
+
+    h = serve.run(Ingress.bind(echo.bind()))
+    assert h.remote("sned").result() == "Hi sned"
+
+
+def test_handle_options_with_same_router(serve_instance):
+    """Make sure that multiple handles share same router object."""
+
+    @serve.deployment
+    def echo(name: str):
+        return f"Hi {name}"
+
+    handle = serve.run(echo.bind())
+    handle2 = handle.options(multiplexed_model_id="model2")
+    assert handle2._router is handle._router
+
+
+def test_init(serve_instance):
+    @serve.deployment
+    def f():
+        return "hi"
+
+    h = serve.run(f.bind())
+    h._init(_prefer_local_routing=True)
+    for _ in range(10):
+        assert h.remote().result() == "hi"
+
+
+def test_init_twice_fails(serve_instance):
+    @serve.deployment
+    def f():
+        return "hi"
+
+    h = serve.run(f.bind())
+    h._init()
+
+    with pytest.raises(RuntimeError):
+        h._init()
+
+
+def test_init_after_options_fails(serve_instance):
+    @serve.deployment
+    def f():
+        return "hi"
+
+    h = serve.run(f.bind())
+
+    with pytest.raises(RuntimeError):
+        h.options(stream=True)._init(_prefer_local_routing=True)
+
+
+def test_init_after_request_fails(serve_instance):
+    @serve.deployment
+    def f():
+        return "hi"
+
+    h = serve.run(f.bind())
+    assert h.remote().result() == "hi"
+
+    with pytest.raises(RuntimeError):
+        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))")
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode doesn't support choose_replica/dispatch",
+)
+async def test_choose_replica_and_dispatch_single(serve_instance):
+    """Test choose_replica + dispatch for simple single selection pattern."""
+
+    @serve.deployment(num_replicas=2)
+    class Backend:
+        def process(self, msg: str):
+            replica_id = get_replica_context().replica_id.unique_id
+            return {"actual_replica_id": replica_id, "response": msg}
+
+    @serve.deployment
+    class SimpleProxy:
+        def __init__(self, backend: DeploymentHandle):
+            self.backend = backend
+
+        async def handle_request(self, request: str):
+            # Context manager ensures slot is released if dispatch fails or is skipped
+            async with self.backend.process.choose_replica(request) as selection:
+                assert selection.replica_id is not None
+                assert selection.node_ip is not None
+
+                # Dispatch to the selected replica
+                response = await self.backend.process.dispatch(selection, request)
+
+                # Return both the selection and the response for verification
+                return {"selected_replica_id": selection.replica_id, **response}
+
+    h = serve.run(SimpleProxy.bind(Backend.bind()))
+    result = await h.handle_request.remote("test_message")
+
+    # Verify the result contains the message
+    assert result["response"] == "test_message"
+
+    # Verify that dispatch sent the request to the replica we selected
+    assert result["actual_replica_id"] == result["selected_replica_id"], (
+        f"dispatch sent request to wrong replica: "
+        f"selected {result['selected_replica_id']}, but got response from {result['actual_replica_id']}"
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode doesn't support choose_replica/dispatch",
+)
+async def test_choose_replica_early_return_releases_slot(serve_instance):
+    """Early return from choose_replica should release the reserved slot."""
+
+    @serve.deployment(num_replicas=1, max_ongoing_requests=1)
+    class Backend:
+        def process(self, msg: str):
+            return msg
+
+    @serve.deployment
+    class Proxy:
+        def __init__(self, backend: DeploymentHandle):
+            self.backend = backend
+
+        async def handle_request(self, request: str):
+            async with self.backend.process.choose_replica(request):
+                pass
+
+            # Ensure the second choose_replica request work
+            async with self.backend.process.choose_replica(request) as selection:
+                response = await asyncio.wait_for(
+                    self.backend.process.dispatch(selection, request), timeout=2
+                )
+                return response
+
+    h = serve.run(Proxy.bind(Backend.bind()))
+    assert await h.handle_request.remote("test_message") == "test_message"
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode doesn't support choose_replica/dispatch",
+)
+async def test_choose_replica_exception_releases_slot(serve_instance):
+    """Exception in choose_replica context should release the reserved slot."""
+
+    @serve.deployment(num_replicas=1, max_ongoing_requests=1)
+    class Backend:
+        def process(self, msg: str):
+            return msg
+
+    @serve.deployment
+    class Proxy:
+        def __init__(self, backend: DeploymentHandle):
+            self.backend = backend
+
+        async def handle_request(self, request: str):
+            try:
+                async with self.backend.process.choose_replica(request):
+                    raise RuntimeError("test exception")
+            except RuntimeError:
+                pass
+
+            # Ensure the second choose_replica request work
+            async with self.backend.process.choose_replica(request) as selection:
+                response = await asyncio.wait_for(
+                    self.backend.process.dispatch(selection, request), timeout=2
+                )
+                return response
+
+    h = serve.run(Proxy.bind(Backend.bind()))
+    assert await h.handle_request.remote("test_message") == "test_message"
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode doesn't support choose_replica/dispatch",
+)
+async def test_choose_replica_and_dispatch_streaming(serve_instance):
+    """Test choose_replica + dispatch with handle.options(stream=True)."""
+
+    @serve.deployment(num_replicas=2)
+    class Backend:
+        async def stream(self, msg: str):
+            replica_id = get_replica_context().replica_id.unique_id
+            for i in range(3):
+                yield {"actual_replica_id": replica_id, "chunk": f"{msg}-{i}"}
+
+    @serve.deployment
+    class StreamingProxy:
+        def __init__(self, backend: DeploymentHandle):
+            self.backend_stream = backend.stream.options(stream=True)
+
+        async def handle_request(self, request: str):
+            async with self.backend_stream.choose_replica(request) as selection:
+                gen = self.backend_stream.dispatch(selection, request)
+                chunks = [item async for item in gen]
+                return {"selected_replica_id": selection.replica_id, "chunks": chunks}
+
+    h = serve.run(StreamingProxy.bind(Backend.bind()))
+    result = await h.handle_request.remote("stream_test")
+
+    assert [item["chunk"] for item in result["chunks"]] == [
+        "stream_test-0",
+        "stream_test-1",
+        "stream_test-2",
+    ]
+    assert all(
+        item["actual_replica_id"] == result["selected_replica_id"]
+        for item in result["chunks"]
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode doesn't support choose_replica/dispatch",
+)
+async def test_choose_replica_and_dispatch_parallel(serve_instance):
+    """Test parallel selection pattern (e.g., PD proxy) using AsyncExitStack."""
+
+    @serve.deployment(num_replicas=2)
+    class PrefillServer:
+        def chat(self, msg: str):
+            replica_id = get_replica_context().replica_id.unique_id
+            return {"actual_replica_id": replica_id, "response": msg}
+
+    @serve.deployment(num_replicas=2)
+    class DecodeServer:
+        def chat(self, msg: str):
+            replica_id = get_replica_context().replica_id.unique_id
+            return {"actual_replica_id": replica_id, "response": msg}
+
+    @serve.deployment
+    class PDProxy:
+        def __init__(
+            self,
+            prefill_server: DeploymentHandle,
+            decode_server: DeploymentHandle,
+        ):
+            self.prefill = prefill_server
+            self.decode = decode_server
+
+        async def handle_request(self, request: str):
+            # Use AsyncExitStack to manage multiple context managers in parallel
+            async with AsyncExitStack() as stack:
+                #  Select and RESERVE replicas from BOTH deployments in parallel
+                p_selection, d_selection = await asyncio.gather(
+                    stack.enter_async_context(self.prefill.chat.choose_replica()),
+                    stack.enter_async_context(self.decode.chat.choose_replica()),
+                )
+
+                p_msg = f"prefill:{request}"
+                d_msg = f"decode:{request}"
+
+                # Dispatch to both selected replicas
+                p_result, d_result = await asyncio.gather(
+                    self.prefill.chat.dispatch(p_selection, p_msg),
+                    self.decode.chat.dispatch(d_selection, d_msg),
+                )
+                return {
+                    "prefill": {
+                        "selected_replica_id": p_selection.replica_id,
+                        **p_result,
+                    },
+                    "decode": {
+                        "selected_replica_id": d_selection.replica_id,
+                        **d_result,
+                    },
+                }
+
+    h = serve.run(PDProxy.bind(PrefillServer.bind(), DecodeServer.bind()))
+    result = await h.handle_request.remote("test_parallel")
+
+    assert result["prefill"]["response"] == "prefill:test_parallel"
+    assert result["decode"]["response"] == "decode:test_parallel"
+
+    # Verify that dispatch sent the request to the replica we selected
+    assert (
+        result["prefill"]["actual_replica_id"]
+        == result["prefill"]["selected_replica_id"]
+    ), (
+        f"dispatch sent request to wrong replica for prefill: "
+        f"selected {result['prefill']['selected_replica_id']}, but got response from {result['prefill']['actual_replica_id']}"
+    )
+    assert (
+        result["decode"]["actual_replica_id"] == result["decode"]["selected_replica_id"]
+    ), (
+        f"dispatch sent request to wrong replica for decode: "
+        f"selected {result['decode']['selected_replica_id']}, but got response from {result['decode']['actual_replica_id']}"
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode doesn't support choose_replica/dispatch",
+)
+async def test_choose_replica_cancel_releases_slot_across_loop_boundary(
+    serve_instance,
+):
+    """Cancelling a task holding a selection must release the slot, even when
+    the router lives on a different thread (SingletonThreadRouter)."""
+
+    @serve.deployment(num_replicas=1, max_ongoing_requests=1)
+    class Backend:
+        def process(self, msg: str):
+            return msg
+
+    handle = serve.run(Backend.bind())
+    # Warm up so the cancellation isn't racing initialization.
+    assert await handle.process.remote("warmup") == "warmup"
+
+    ready = asyncio.Event()
+    never_set = asyncio.Event()
+
+    async def hold():
+        async with handle.process.choose_replica("first"):
+            ready.set()
+            await never_set.wait()
+
+    task = asyncio.create_task(hold())
+    await asyncio.wait_for(ready.wait(), timeout=5)
+
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    # Hangs if the cancelled task leaked the only slot.
+    async with handle.process.choose_replica("second") as selection:
+        result = await asyncio.wait_for(
+            handle.process.dispatch(selection, "second"),
+            timeout=5,
+        )
+    assert result == "second"
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode doesn't support choose_replica/dispatch",
+)
+async def test_dispatch_wrapper_uses_choose_replica_stream_flag(serve_instance):
+    """Wrapper type follows metadata.is_streaming, not the dispatching handle."""
+
+    @serve.deployment
+    class Backend:
+        async def stream(self, msg: str):
+            yield f"{msg}-0"
+            yield f"{msg}-1"
+
+    @serve.deployment
+    class Proxy:
+        def __init__(self, backend: DeploymentHandle):
+            self.stream_handle = backend.stream.options(stream=True)
+            self.dispatch_handle = backend.stream  # default stream=False
+
+        async def run(self, msg: str) -> bool:
+            async with self.stream_handle.choose_replica(msg) as sel:
+                resp = self.dispatch_handle.dispatch(sel, msg)
+                is_gen = isinstance(resp, DeploymentResponseGenerator)
+                if is_gen:
+                    async for _ in resp:
+                        pass
+                else:
+                    try:
+                        await resp
+                    except Exception:
+                        pass
+            return is_gen
+
+    h = serve.run(Proxy.bind(Backend.bind()))
+    assert await h.run.remote("test") is True
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode doesn't support choose_replica/dispatch",
+)
+async def test_dispatch_rejects_selection_without_deployment_id(serve_instance):
+    """dispatch must reject a selection whose deployment_id is unset."""
+
+    @serve.deployment
+    class Backend:
+        def process(self, msg: str):
+            return msg
+
+    @serve.deployment
+    class Caller:
+        def __init__(self, backend: DeploymentHandle):
+            self.handle = backend.process
+
+        async def run(self) -> str:
+            async with self.handle.choose_replica("msg") as sel:
+                sel._deployment_id = None
+                try:
+                    resp = self.handle.dispatch(sel, "msg")
+                except ValueError:
+                    return "raised"
+                await resp
+                return "no_raise"
+
+    h = serve.run(Caller.bind(Backend.bind()))
+    assert await h.run.remote() == "raised"
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    RAY_SERVE_FORCE_LOCAL_TESTING_MODE,
+    reason="local_testing_mode doesn't support choose_replica/dispatch",
+)
+async def test_request_counter_only_increments_on_dispatch(serve_instance):
+    """request_counter ticks per dispatched request, not per choose_replica."""
+
+    @serve.deployment
+    class Backend:
+        def process(self, msg: str):
+            return msg
+
+    @serve.deployment
+    class Caller:
+        def __init__(self, backend: DeploymentHandle):
+            self.handle = backend.process
+
+        async def run(self, dispatch_too: bool) -> int:
+            inc_calls: list = []
+            original_inc = self.handle.request_counter.inc
+
+            def tracking_inc(*args, **kwargs):
+                inc_calls.append(1)
+                return original_inc(*args, **kwargs)
+
+            self.handle.request_counter.inc = tracking_inc
+            try:
+                if dispatch_too:
+                    async with self.handle.choose_replica("msg") as sel:
+                        await self.handle.dispatch(sel, "msg")
+                else:
+                    async with self.handle.choose_replica("msg"):
+                        pass
+                return len(inc_calls)
+            finally:
+                self.handle.request_counter.inc = original_inc
+
+    h = serve.run(Caller.bind(Backend.bind()))
+    assert await h.run.remote(False) == 0
+    assert await h.run.remote(True) == 1
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))