chore: import upstream snapshot with attribution

python/ray/serve/tests/test_advanced.py

@@ -0,0 +1,146 @@
+import sys
+import time
+
+import httpx
+import pytest
+from starlette.requests import Request
+
+import ray
+from ray import serve
+from ray._common.test_utils import SignalActor
+from ray.serve._private.constants import SERVE_DEFAULT_APP_NAME
+from ray.serve._private.test_utils import Barrier
+from ray.serve.handle import DeploymentHandle
+from ray.util.state import list_objects
+
+
+def test_serve_forceful_shutdown(serve_instance):
+    @serve.deployment(graceful_shutdown_timeout_s=0.1)
+    def sleeper():
+        while True:
+            time.sleep(1000)
+
+    handle = serve.run(sleeper.bind())
+    response = handle.remote()
+    serve.delete(SERVE_DEFAULT_APP_NAME)
+
+    with pytest.raises(ray.exceptions.RayActorError):
+        response.result()
+
+
+def test_serve_graceful_shutdown(serve_instance):
+    signal = SignalActor.remote()
+
+    @serve.deployment(
+        name="wait",
+        max_ongoing_requests=10,
+        graceful_shutdown_timeout_s=1000,
+        graceful_shutdown_wait_loop_s=0.5,
+    )
+    class Wait:
+        async def __call__(self, signal_actor):
+            await signal_actor.wait.remote()
+
+    handle = serve.run(Wait.bind())
+    responses = [handle.remote(signal) for _ in range(10)]
+
+    # Wait for all the queries to be enqueued
+    with pytest.raises(TimeoutError):
+        responses[0].result(timeout_s=1)
+
+    @ray.remote(num_cpus=0)
+    def do_blocking_delete():
+        serve.delete(SERVE_DEFAULT_APP_NAME)
+
+    # Now delete the deployment. This should trigger the shutdown sequence.
+    delete_ref = do_blocking_delete.remote()
+
+    # The queries should be enqueued but not executed becuase they are blocked
+    # by signal actor.
+    with pytest.raises(TimeoutError):
+        responses[0].result(timeout_s=1)
+
+    signal.send.remote()
+
+    # All the queries should be drained and executed without error.
+    [r.result() for r in responses]
+    # Blocking delete should complete.
+    ray.get(delete_ref)
+
+
+def test_parallel_start(serve_instance):
+    # Test the ability to start multiple replicas in parallel.
+    # In the past, when Serve scale up a deployment, it does so one by one and
+    # wait for each replica to initialize. This test avoid this by preventing
+    # the first replica to finish initialization unless the second replica is
+    # also started.
+    barrier = Barrier.remote(n=2)
+
+    @serve.deployment(num_replicas=2)
+    class LongStartingServable:
+        def __init__(self):
+            ray.get(barrier.wait.remote(), timeout=10)
+
+        def __call__(self):
+            return "Ready"
+
+    handle = serve.run(LongStartingServable.bind())
+    handle.remote().result(timeout_s=10)
+
+
+def test_passing_object_ref_to_deployment_not_pinned_to_memory(serve_instance):
+    """Passing object refs to deployments should not pin the refs in memory.
+
+    We had issue that passing object ref to a deployment will result in memory leak
+    due to _PyObjScanner/ cloudpickler pinning the object to memory. This test will
+    ensure the object ref is released after the request is done.
+
+    See: https://github.com/ray-project/ray/issues/43248
+    """
+
+    def _obj_ref_exists_in_state_api(obj_ref_hex: str) -> bool:
+        return (
+            len(
+                list_objects(
+                    filters=[("object_id", "=", obj_ref_hex)],
+                    raise_on_missing_output=False,
+                )
+            )
+            > 0
+        )
+
+    @serve.deployment
+    class Dep1:
+        def multiply_by_two(self, length: int):
+            return length * 2
+
+    @serve.deployment
+    class Gateway:
+        def __init__(self, dep1: DeploymentHandle):
+            self.dep1: DeploymentHandle = dep1
+
+        async def __call__(self, http_request: Request) -> str:
+            length = int(http_request.query_params.get("length"))
+            length_ref = ray.put(length)
+
+            # Sanity check that the ObjectRef exists in the state API.
+            assert _obj_ref_exists_in_state_api(length_ref.hex())
+            return {
+                "length": length,
+                "result": await self.dep1.multiply_by_two.remote(length_ref),
+                "length_ref_hex": length_ref.hex(),
+            }
+
+    serve.run(Gateway.bind(Dep1.bind()))
+
+    length = 10
+    response = httpx.get(f"http://localhost:8000?length={length}").json()
+    assert response["length"] == length
+    assert response["result"] == length * 2
+
+    # Ensure the object ref is not in the memory anymore.
+    assert not _obj_ref_exists_in_state_api(response["length_ref_hex"])
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))