chore: import upstream snapshot with attribution

python/ray/serve/tests/test_shutdown.py

@@ -0,0 +1,566 @@
+"""Tests for Serve shutdown behavior (isolated from test_standalone)."""
+
+import sys
+
+import pytest
+
+import ray
+from ray import serve
+from ray._common.test_utils import wait_for_condition
+from ray._raylet import GcsClient
+from ray.serve._private.constants import (
+    SERVE_CONTROLLER_NAME,
+    SERVE_DEPLOYMENT_ACTOR_PREFIX,
+    SERVE_NAMESPACE,
+    SERVE_PROXY_NAME,
+)
+from ray.serve._private.default_impl import create_cluster_node_info_cache
+from ray.serve._private.test_utils import expected_proxy_actors
+from ray.serve._private.utils import format_actor_name
+from ray.serve.config import DeploymentActorConfig
+from ray.util.state import list_actors
+
+
+def test_shutdown(ray_shutdown):
+    ray.init(num_cpus=8)
+    serve.start(http_options=dict(port=8003))
+    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
+    cluster_node_info_cache = create_cluster_node_info_cache(gcs_client)
+    cluster_node_info_cache.update()
+
+    @serve.deployment
+    def f():
+        pass
+
+    serve.run(f.bind())
+
+    actor_names = [
+        SERVE_CONTROLLER_NAME,
+        format_actor_name(
+            SERVE_PROXY_NAME,
+            cluster_node_info_cache.get_alive_nodes()[0][0],
+        ),
+    ]
+
+    def check_alive():
+        alive = True
+        for actor_name in actor_names:
+            try:
+                ray.get_actor(actor_name, namespace=SERVE_NAMESPACE)
+            except ValueError:
+                alive = False
+        return alive
+
+    wait_for_condition(check_alive)
+
+    serve.shutdown()
+
+    def check_dead():
+        for actor_name in actor_names:
+            try:
+                ray.get_actor(actor_name, namespace=SERVE_NAMESPACE)
+                return False
+            except ValueError:
+                pass
+        return True
+
+    wait_for_condition(check_dead)
+
+
+@pytest.mark.asyncio
+async def test_shutdown_async(ray_shutdown):
+    ray.init(num_cpus=8)
+    serve.start(http_options=dict(port=8003))
+    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
+    cluster_node_info_cache = create_cluster_node_info_cache(gcs_client)
+    cluster_node_info_cache.update()
+
+    @serve.deployment
+    def f():
+        pass
+
+    serve.run(f.bind())
+
+    actor_names = [
+        SERVE_CONTROLLER_NAME,
+        format_actor_name(
+            SERVE_PROXY_NAME,
+            cluster_node_info_cache.get_alive_nodes()[0][0],
+        ),
+    ]
+
+    def check_alive():
+        alive = True
+        for actor_name in actor_names:
+            try:
+                ray.get_actor(actor_name, namespace=SERVE_NAMESPACE)
+            except ValueError:
+                alive = False
+        return alive
+
+    wait_for_condition(check_alive)
+
+    await serve.shutdown_async()
+
+    def check_dead():
+        for actor_name in actor_names:
+            try:
+                ray.get_actor(actor_name, namespace=SERVE_NAMESPACE)
+                return False
+            except ValueError:
+                pass
+        return True
+
+    wait_for_condition(check_dead)
+
+
+def test_single_app_shutdown_actors(ray_shutdown):
+    """Tests serve.shutdown() works correctly in single-app case
+
+    Ensures that after deploying a (nameless) app using serve.run(), serve.shutdown()
+    deletes all actors (controller, http proxy, all replicas) in the "serve" namespace.
+    """
+    address = ray.init(num_cpus=8)["address"]
+    serve.start(http_options=dict(port=8003))
+
+    @serve.deployment
+    def f():
+        pass
+
+    serve.run(f.bind(), name="app")
+
+    actor_names = {
+        "ServeController",
+        *expected_proxy_actors(),
+        "ServeReplica:app:f",
+    }
+
+    def check_alive():
+        actors = list_actors(
+            address=address,
+            filters=[("ray_namespace", "=", SERVE_NAMESPACE), ("state", "=", "ALIVE")],
+        )
+        return {actor["class_name"] for actor in actors} == actor_names
+
+    def check_dead():
+        actors = list_actors(
+            address=address,
+            filters=[("ray_namespace", "=", SERVE_NAMESPACE), ("state", "=", "ALIVE")],
+        )
+        return len(actors) == 0
+
+    wait_for_condition(check_alive)
+    serve.shutdown()
+    wait_for_condition(check_dead)
+
+
+@pytest.mark.asyncio
+async def test_single_app_shutdown_actors_async(ray_shutdown):
+    """Tests serve.shutdown_async() works correctly in single-app case
+
+    Ensures that after deploying a (nameless) app using serve.run(), serve.shutdown_async()
+    deletes all actors (controller, http proxy, all replicas) in the "serve" namespace.
+    """
+    address = ray.init(num_cpus=8)["address"]
+    serve.start(http_options=dict(port=8003))
+
+    @serve.deployment
+    def f():
+        pass
+
+    serve.run(f.bind(), name="app")
+
+    actor_names = {
+        "ServeController",
+        *expected_proxy_actors(),
+        "ServeReplica:app:f",
+    }
+
+    def check_alive():
+        actors = list_actors(
+            address=address,
+            filters=[("ray_namespace", "=", SERVE_NAMESPACE), ("state", "=", "ALIVE")],
+        )
+        return {actor["class_name"] for actor in actors} == actor_names
+
+    def check_dead():
+        actors = list_actors(
+            address=address,
+            filters=[("ray_namespace", "=", SERVE_NAMESPACE), ("state", "=", "ALIVE")],
+        )
+        return len(actors) == 0
+
+    wait_for_condition(check_alive)
+    await serve.shutdown_async()
+    wait_for_condition(check_dead)
+
+
+def test_multi_app_shutdown_actors(ray_shutdown):
+    """Tests serve.shutdown() works correctly in multi-app case.
+
+    Ensures that after deploying multiple distinct applications, serve.shutdown()
+    deletes all actors (controller, http proxy, all replicas) in the "serve" namespace.
+    """
+    address = ray.init(num_cpus=8)["address"]
+    serve.start(http_options=dict(port=8003))
+
+    @serve.deployment
+    def f():
+        pass
+
+    serve.run(f.bind(), name="app1", route_prefix="/app1")
+    serve.run(f.bind(), name="app2", route_prefix="/app2")
+
+    actor_names = {
+        "ServeController",
+        *expected_proxy_actors(),
+        "ServeReplica:app1:f",
+        "ServeReplica:app2:f",
+    }
+
+    def check_alive():
+        actors = list_actors(
+            address=address,
+            filters=[("ray_namespace", "=", SERVE_NAMESPACE), ("state", "=", "ALIVE")],
+        )
+        return {actor["class_name"] for actor in actors} == actor_names
+
+    def check_dead():
+        actors = list_actors(
+            address=address,
+            filters=[("ray_namespace", "=", SERVE_NAMESPACE), ("state", "=", "ALIVE")],
+        )
+        return len(actors) == 0
+
+    wait_for_condition(check_alive)
+    serve.shutdown()
+    wait_for_condition(check_dead)
+
+
+@pytest.mark.asyncio
+async def test_multi_app_shutdown_actors_async(ray_shutdown):
+    """Tests serve.shutdown_async() works correctly in multi-app case.
+
+    Ensures that after deploying multiple distinct applications, serve.shutdown_async()
+    deletes all actors (controller, http proxy, all replicas) in the "serve" namespace.
+    """
+    address = ray.init(num_cpus=8)["address"]
+    serve.start(http_options=dict(port=8003))
+
+    @serve.deployment
+    def f():
+        pass
+
+    serve.run(f.bind(), name="app1", route_prefix="/app1")
+    serve.run(f.bind(), name="app2", route_prefix="/app2")
+
+    actor_names = {
+        "ServeController",
+        *expected_proxy_actors(),
+        "ServeReplica:app1:f",
+        "ServeReplica:app2:f",
+    }
+
+    def check_alive():
+        actors = list_actors(
+            address=address,
+            filters=[("ray_namespace", "=", SERVE_NAMESPACE), ("state", "=", "ALIVE")],
+        )
+        return {actor["class_name"] for actor in actors} == actor_names
+
+    def check_dead():
+        actors = list_actors(
+            address=address,
+            filters=[("ray_namespace", "=", SERVE_NAMESPACE), ("state", "=", "ALIVE")],
+        )
+        return len(actors) == 0
+
+    wait_for_condition(check_alive)
+    await serve.shutdown_async()
+    wait_for_condition(check_dead)
+
+
+def test_registered_cleanup_actors_killed_on_shutdown(ray_shutdown):
+    """Test that actors registered via _register_shutdown_cleanup_actor are killed.
+
+    This tests the internal actor registration API that allows deployments to register
+    auxiliary actors (like caches, coordinators, etc.) for cleanup on serve.shutdown().
+    """
+    ray.init(num_cpus=4)
+    serve.start()
+
+    # Create a detached actor that we'll register for cleanup
+    @ray.remote
+    class DummyActor:
+        def ping(self):
+            return "pong"
+
+    dummy_actor_name = "test_registered_cleanup_dummy"
+    dummy = DummyActor.options(
+        name=dummy_actor_name, namespace=SERVE_NAMESPACE, lifetime="detached"
+    ).remote()
+
+    # Verify actor is alive
+    assert ray.get(dummy.ping.remote()) == "pong"
+
+    # Register the actor with the controller for cleanup
+    controller = ray.get_actor(SERVE_CONTROLLER_NAME, namespace=SERVE_NAMESPACE)
+    ray.get(controller._register_shutdown_cleanup_actor.remote(dummy))
+
+    # Shutdown serve
+    serve.shutdown()
+
+    # Verify the registered actor is killed
+    def check_actor_dead():
+        try:
+            ray.get_actor(dummy_actor_name, namespace=SERVE_NAMESPACE)
+            return False
+        except ValueError:
+            return True
+
+    wait_for_condition(check_actor_dead)
+
+
+def test_serve_shutdown(ray_shutdown):
+    ray.init(namespace="serve")
+    serve.start()
+
+    @serve.deployment
+    class A:
+        def __call__(self, *args):
+            return "hi"
+
+    serve.run(A.bind())
+
+    assert len(serve.status().applications) == 1
+
+    serve.shutdown()
+    serve.start()
+
+    assert len(serve.status().applications) == 0
+
+    serve.run(A.bind())
+
+    assert len(serve.status().applications) == 1
+
+
+@pytest.mark.asyncio
+async def test_serve_shutdown_async(ray_shutdown):
+    ray.init(namespace="serve")
+    serve.start()
+
+    @serve.deployment
+    class A:
+        def __call__(self, *args):
+            return "hi"
+
+    serve.run(A.bind())
+
+    assert len(serve.status().applications) == 1
+
+    await serve.shutdown_async()
+    serve.start()
+
+    assert len(serve.status().applications) == 0
+
+    serve.run(A.bind())
+
+    assert len(serve.status().applications) == 1
+
+
+def test_shutdown_after_driver_reconnect(ray_shutdown):
+    """serve.shutdown() should work after the driver disconnects and reconnects.
+
+    Regression test for https://github.com/ray-project/ray/issues/61608.
+    When using `with ray.init()` as a context manager, ray.shutdown() is called
+    on exit but the module-level _global_client in ray.serve.context is not
+    cleared. On the next ray.init(), serve.shutdown() must health-check the
+    cached controller handle so it can reconnect instead of using a stale
+    handle from the previous session.
+    """
+
+    @serve.deployment
+    class Greeter:
+        def __call__(self, *args):
+            return "hi"
+
+    # Session 1: start Serve and deploy. The context manager calls
+    # ray.shutdown() on exit, but the detached controller survives.
+    with ray.init(namespace="serve"):
+        serve.start()
+        serve.run(Greeter.bind())
+        assert len(serve.status().applications) == 1
+
+    # Session 2: reconnect and shut down Serve. This must not fail due to
+    # a stale _global_client from session 1.
+    with ray.init(namespace="serve"):
+        serve.shutdown()
+
+    # Session 3: verify Serve is actually shut down — starting fresh should
+    # show zero applications.
+    with ray.init(namespace="serve"):
+        serve.start()
+        assert len(serve.status().applications) == 0
+        serve.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_shutdown_async_after_driver_reconnect(ray_shutdown):
+    """Async variant: serve.shutdown_async() should work after driver reconnect.
+
+    Regression test for https://github.com/ray-project/ray/issues/61608.
+    """
+
+    @serve.deployment
+    class Greeter:
+        def __call__(self, *args):
+            return "hi"
+
+    with ray.init(namespace="serve"):
+        serve.start()
+        serve.run(Greeter.bind())
+        assert len(serve.status().applications) == 1
+
+    with ray.init(namespace="serve"):
+        await serve.shutdown_async()
+
+    with ray.init(namespace="serve"):
+        serve.start()
+        assert len(serve.status().applications) == 0
+        await serve.shutdown_async()
+
+
+def test_shutdown_after_reconnect_shuts_down_live_serve(ray_cluster):
+    """serve.shutdown() must shut down a still-running Serve after the driver
+    reconnects to a long-lived cluster.
+
+    Regression test for https://github.com/ray-project/ray/issues/64647.
+    Unlike test_shutdown_after_driver_reconnect (which starts a fresh in-process
+    cluster on each ray.init(), so no controller ever survives), this uses a
+    persistent cluster. The detached controller survives the driver's
+    ray.shutdown(), so a stale cached _global_client left over from the first
+    session must not make serve.shutdown() skip shutting the live controller down.
+    """
+    cluster = ray_cluster
+    head_node = cluster.add_node(num_cpus=8)
+
+    @serve.deployment
+    class First:
+        def __call__(self):
+            return "first"
+
+    # Session 1: start Serve. The detached controller survives disconnect.
+    ray.init(head_node.address, namespace=SERVE_NAMESPACE)
+    serve.start()
+    serve.run(First.bind())
+    ray.get_actor(SERVE_CONTROLLER_NAME, namespace=SERVE_NAMESPACE)  # sanity: exists
+    # Intentionally do NOT clear serve.context._global_client: the stale cached
+    # client left behind here is exactly what triggers the bug.
+    ray.shutdown()
+
+    # Session 2: reconnect to the same cluster and shut Serve down.
+    ray.init(head_node.address, namespace=SERVE_NAMESPACE)
+    serve.shutdown()
+
+    # The live controller must actually be gone now.
+    def controller_gone():
+        try:
+            ray.get_actor(SERVE_CONTROLLER_NAME, namespace=SERVE_NAMESPACE)
+            return False
+        except ValueError:
+            return True
+
+    wait_for_condition(controller_gone, timeout=20)
+
+
+def test_serve_shutdown_cleans_up_deployment_actors(ray_shutdown):
+    """serve.shutdown() kills all deployment actors.
+
+    Deployment actors are detached, so they must be explicitly killed during
+    shutdown.
+    """
+
+    ray_context = ray.init(num_cpus=4, namespace="default_test_namespace")
+    address = ray_context.address_info["address"]
+    serve.start()
+
+    @ray.remote
+    class SharedCounter:
+        def __init__(self, start: int = 0):
+            self.count = start
+
+        def get(self):
+            return self.count
+
+    @serve.deployment(
+        deployment_actors=[
+            DeploymentActorConfig(
+                name="counter",
+                actor_class=SharedCounter,
+                init_kwargs={"start": 0},
+            ),
+        ],
+    )
+    class MyDeployment:
+        def __call__(self):
+            counter = serve.get_deployment_actor("counter")
+            return str(ray.get(counter.get.remote()))
+
+    serve.run(MyDeployment.bind())
+
+    def _check_deployment_actor_count(expected: int):
+        actors = list_actors(address=address, filters=[("state", "=", "ALIVE")])
+        names = [
+            a["name"]
+            for a in actors
+            if a["name"] and a["name"].startswith(SERVE_DEPLOYMENT_ACTOR_PREFIX)
+        ]
+        return len(names) == expected
+
+    wait_for_condition(lambda: _check_deployment_actor_count(1))
+
+    serve.shutdown()
+
+    wait_for_condition(lambda: _check_deployment_actor_count(0), timeout=15)
+
+
+def test_crash_during_shutdown_no_orphaned_actors(ray_shutdown):
+    """Controller crash mid-shutdown should not orphan replica actors."""
+    address = ray.init(num_cpus=4)["address"]
+    serve.start()
+
+    @serve.deployment(num_replicas=2)
+    class MyModel:
+        def __call__(self, request):
+            return "hello"
+
+    serve.run(MyModel.bind(), name="my_app")
+
+    def replicas_running():
+        actors = list_actors(
+            address=address,
+            filters=[("ray_namespace", "=", SERVE_NAMESPACE), ("state", "=", "ALIVE")],
+        )
+        replicas = [a for a in actors if "ServeReplica" in a.get("class_name", "")]
+        return len(replicas) == 2
+
+    wait_for_condition(replicas_running, timeout=30)
+
+    # Start shutdown, then immediately kill controller to simulate crash.
+    controller = ray.get_actor(SERVE_CONTROLLER_NAME, namespace=SERVE_NAMESPACE)
+    ray.get(controller.graceful_shutdown.remote(wait=False), timeout=10)
+    ray.kill(controller, no_restart=False)
+
+    # Restarted controller should recover and clean up all actors.
+    def check_dead():
+        actors = list_actors(
+            address=address,
+            filters=[("ray_namespace", "=", SERVE_NAMESPACE), ("state", "=", "ALIVE")],
+        )
+        return len(actors) == 0
+
+    wait_for_condition(check_dead, timeout=60)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))