chore: import upstream snapshot with attribution

python/ray/serve/tests/test_controller_recovery.py

@@ -0,0 +1,551 @@
+import asyncio
+import logging
+import os
+import re
+import sys
+import time
+
+import httpx
+import pytest
+
+import ray
+from ray import serve
+from ray._common.test_utils import SignalActor, wait_for_condition
+from ray.exceptions import RayTaskError
+from ray.serve._private.common import DeploymentID, ReplicaState
+from ray.serve._private.constants import (
+    RAY_SERVE_ENABLE_HA_PROXY,
+    SERVE_CONTROLLER_NAME,
+    SERVE_DEFAULT_APP_NAME,
+    SERVE_NAMESPACE,
+    SERVE_PROXY_NAME,
+)
+from ray.serve._private.test_utils import (
+    check_replica_counts,
+    get_application_url,
+    request_with_retries,
+)
+from ray.serve.schema import LoggingConfig, ServeDeploySchema
+from ray.util.state import list_actors
+
+
+@pytest.mark.parametrize(
+    "deployment_options",
+    [
+        {"num_replicas": 2},
+        {"autoscaling_config": {"min_replicas": 2, "max_replicas": 2}},
+    ],
+)
+def test_recover_start_from_replica_actor_names(serve_instance, deployment_options):
+    """Test controller is able to recover starting -> running replicas from
+    actor names.
+    """
+
+    # Test failed to deploy with total of 2 replicas,
+    # but first constructor call fails.
+    @serve.deployment(
+        name="recover_start_from_replica_actor_names", **deployment_options
+    )
+    class TransientConstructorFailureDeployment:
+        def __init__(self):
+            pass
+
+        def __call__(self, *args):
+            return "hii"
+
+    serve.run(TransientConstructorFailureDeployment.bind(), name="app")
+    for _ in range(10):
+        response = request_with_retries(timeout=30, app_name="app")
+        assert response.text == "hii"
+    # Assert 2 replicas are running in deployment deployment after partially
+    # successful deploy() call with transient error
+    deployment_dict = ray.get(serve_instance._controller._all_running_replicas.remote())
+    id = DeploymentID(name="recover_start_from_replica_actor_names", app_name="app")
+    assert len(deployment_dict[id]) == 2
+
+    replica_version_hash = None
+    for replica in deployment_dict[id]:
+        ref = replica.get_actor_handle().initialize_and_get_metadata.remote()
+        _, version, *_ = ray.get(ref)
+        if replica_version_hash is None:
+            replica_version_hash = hash(version)
+        assert replica_version_hash == hash(version), (
+            "Replica version hash should be the same for "
+            "same code version and user config."
+        )
+
+    # Sample: [
+    # 'TransientConstructorFailureDeployment#xlituP',
+    # 'SERVE_CONTROLLER_ACTOR',
+    # 'TransientConstructorFailureDeployment#NosHNA',
+    # 'SERVE_CONTROLLER_ACTOR:SERVE_PROXY_ACTOR-node:192.168.86.165-0']
+    actor_infos = list_actors(filters=[("state", "=", "ALIVE")])
+    replica_names = [
+        actor_info["name"]
+        for actor_info in actor_infos
+        if (
+            SERVE_CONTROLLER_NAME not in actor_info["name"]
+            and SERVE_PROXY_NAME not in actor_info["name"]
+        )
+    ]
+    assert (
+        len(replica_names) == 2
+    ), "Should have two running replicas fetched from ray API."
+
+    # Kill controller and wait for endpoint to be available again
+    ray.kill(serve.context._global_client._controller, no_restart=False)
+    wait_for_condition(
+        lambda: get_application_url("HTTP", "app", use_localhost=True) is not None
+    )
+    for _ in range(10):
+        response = request_with_retries(timeout=30, app_name="app")
+        assert response.text == "hii"
+
+    # Ensure recovered replica names are the same
+    recovered_actor_infos = list_actors(filters=[("state", "=", "ALIVE")])
+    recovered_replica_names = [
+        actor_info["name"]
+        for actor_info in recovered_actor_infos
+        if (
+            SERVE_CONTROLLER_NAME not in actor_info["name"]
+            and SERVE_PROXY_NAME not in actor_info["name"]
+        )
+    ]
+    assert (
+        recovered_replica_names == replica_names
+    ), "Running replica actor names after recovery must match"
+
+    # Ensure recovered replica version has are the same
+    for replica_name in recovered_replica_names:
+        actor_handle = ray.get_actor(replica_name, namespace=SERVE_NAMESPACE)
+        ref = actor_handle.initialize_and_get_metadata.remote()
+        _, version, *_ = ray.get(ref)
+        assert replica_version_hash == hash(
+            version
+        ), "Replica version hash should be the same after recover from actor names"
+
+
+def test_recover_rolling_update_from_replica_actor_names(serve_instance):
+    """Test controller can recover replicas during rolling update.
+
+    Replicas starting -> updating -> running
+    replicas from actor names, with right replica versions during rolling
+    update.
+    """
+
+    signal = SignalActor.remote()
+
+    @serve.deployment(name="test", num_replicas=2)
+    class V1:
+        async def __call__(self):
+            await signal.wait.remote()
+            return "1", os.getpid()
+
+    @serve.deployment(name="test", num_replicas=2)
+    class V2:
+        async def __call__(self):
+            return "2", os.getpid()
+
+    h = serve.run(V1.bind(), name="app")
+
+    # Send requests to get pids of initial 2 replicas
+    signal.send.remote()
+    refs = [h.remote() for _ in range(10)]
+    versions, pids = zip(*[ref.result() for ref in refs])
+    assert versions.count("1") == 10
+    initial_pids = set(pids)
+    assert len(initial_pids) == 2
+
+    # blocked_ref will block a single replica until the signal is sent.
+    signal.send.remote(clear=True)
+    blocked_ref = h.remote()
+
+    # Kill the controller
+    ray.kill(serve.context._global_client._controller, no_restart=False)
+
+    # Redeploy new version.
+    serve._run(V2.bind(), _blocking=False, name="app")
+
+    # One replica of the old version should be stuck in stopping because
+    # of the blocked request. Two replicas of the new version should be
+    # brought up without waiting for the old replica to stop.
+    wait_for_condition(
+        check_replica_counts,
+        controller=serve_instance._controller,
+        deployment_id=DeploymentID(name="test", app_name="app"),
+        total=3,
+        by_state=[
+            (ReplicaState.STOPPING, 1, lambda r: r._actor.pid in initial_pids),
+            (ReplicaState.RUNNING, 2, lambda r: r._actor.pid not in initial_pids),
+        ],
+    )
+
+    # All new requests should be sent to the new running replicas
+    refs = [h.remote() for _ in range(10)]
+    versions, pids = zip(*[ref.result(timeout_s=5) for ref in refs])
+    assert versions.count("2") == 10
+    pids2 = set(pids)
+    assert len(pids2 & initial_pids) == 0
+
+    # Kill the controller
+    ray.kill(serve.context._global_client._controller, no_restart=False)
+
+    # Release the signal so that the old replica can shutdown
+    ray.get(signal.send.remote())
+    val, pid = blocked_ref.result()
+    assert val == "1"
+    assert pid in initial_pids
+
+    # Now the goal and requests to the new version should complete.
+    # We should have two running replicas of the new version.
+    serve_instance._wait_for_application_running("app")
+    check_replica_counts(
+        controller=serve_instance._controller,
+        deployment_id=DeploymentID(name="test", app_name="app"),
+        total=2,
+        by_state=(
+            [(ReplicaState.RUNNING, 2, lambda r: r._actor.pid not in initial_pids)]
+        ),
+    )
+
+
+def test_controller_recover_initializing_actor(serve_instance):
+    """Controller crash while a replica is still in `__init__`.
+
+    The previous controller never finished sending the replica its initial
+    `initialize_and_get_metadata(rank=...)` call, so the live actor has
+    neither a rank nor a fully-initialized user callable. The new
+    controller must detect this via the `was_initialized` probe and
+    replace the half-initialized actor with a fresh one rather than try to
+    recover it (which would silently complete its initialization with
+    `rank=None` and break rank tracking).
+    """
+
+    signal = SignalActor.remote()
+    init_started = SignalActor.remote()
+    client = serve_instance
+
+    @ray.remote
+    def pending_init_indicator():
+        ray.get(init_started.wait.remote())
+        return True
+
+    @serve.deployment
+    class V1:
+        async def __init__(self):
+            ray.get(init_started.send.remote())
+            await signal.wait.remote()
+
+        def __call__(self, request):
+            return f"1|{os.getpid()}"
+
+    serve._run(V1.bind(), _blocking=False, name="app")
+    ray.get(pending_init_indicator.remote())
+
+    def get_actor_info(name: str):
+        all_current_actors = list_actors(filters=[("state", "=", "ALIVE")])
+        for actor in all_current_actors:
+            if SERVE_PROXY_NAME in actor["name"]:
+                continue
+            if name in actor["name"]:
+                return actor["name"], actor["pid"]
+
+    original_actor_tag, _ = get_actor_info(f"app#{V1.name}")
+    _, controller1_pid = get_actor_info(SERVE_CONTROLLER_NAME)
+    ray.kill(serve.context._global_client._controller, no_restart=False)
+
+    # Wait for the controller to be replaced. `list_actors` can briefly
+    # report the killed controller as ALIVE (and any new controller as
+    # PENDING_CREATION) right after `ray.kill`, so wait specifically for
+    # the pid to change.
+    def controller_replaced():
+        info = get_actor_info(SERVE_CONTROLLER_NAME)
+        return info is not None and info[1] != controller1_pid
+
+    wait_for_condition(controller_replaced)
+
+    # The new controller's `was_initialized` probe will report False for
+    # the half-initialized actor, so it is killed and replaced. Wait for
+    # the replacement replica to start and report it has reached its
+    # constructor. Unblock its `__init__` once it has.
+    ray.get(pending_init_indicator.remote())
+    ray.get(signal.send.remote())
+    client._wait_for_application_running("app")
+
+    # The original half-initialized actor should have been replaced with a
+    # fresh one (different replica id baked into the actor name).
+    new_actor_tag, _ = get_actor_info(f"app#{V1.name}")
+    assert new_actor_tag != original_actor_tag
+
+    # And the original actor should actually be dead -- not just hidden by
+    # the ALIVE filter on a different replica id. `list_actors` may take a
+    # moment to reflect the kill in its state.
+    def original_actor_dead():
+        matching = list_actors(filters=[("name", "=", original_actor_tag)])
+        # Either the entry has been pruned entirely, or it is reported DEAD.
+        return not matching or all(a["state"] == "DEAD" for a in matching)
+
+    wait_for_condition(original_actor_dead)
+
+
+def test_replica_deletion_after_controller_recover(serve_instance):
+    """Test that replicas are deleted when controller is recovered"""
+
+    controller = serve.context._global_client._controller
+
+    @serve.deployment(graceful_shutdown_timeout_s=3)
+    class V1:
+        async def __call__(self):
+            while True:
+                await asyncio.sleep(0.1)
+
+    handle = serve.run(V1.bind(), name="app")
+    _ = handle.remote()
+    serve.delete("app", _blocking=False)
+
+    def check_replica(replica_state=None):
+        id = DeploymentID(name="V1", app_name="app")
+        try:
+            replicas = ray.get(controller._dump_replica_states_for_testing.remote(id))
+        except RayTaskError as ex:
+            # Deployment is not existed any more.
+            if isinstance(ex, KeyError):
+                return []
+            # Unexpected exception raised.
+            raise ex
+        if replica_state is None:
+            replica_state = list(ReplicaState)
+        else:
+            replica_state = [replica_state]
+        return replicas.get(replica_state)
+
+    # Make sure the replica is in STOPPING state.
+    wait_for_condition(lambda: len(check_replica(ReplicaState.STOPPING)) > 0)
+    ray.kill(serve.context._global_client._controller, no_restart=False)
+    # Make sure the replica is in STOPPING state.
+    wait_for_condition(lambda: len(check_replica(ReplicaState.STOPPING)) > 0)
+
+    # The graceful shutdown timeout of 3 seconds should be used
+    wait_for_condition(lambda: len(check_replica()) == 0, timeout=20)
+    # Application should be removed soon after
+    wait_for_condition(lambda: "app" not in serve.status().applications, timeout=20)
+
+
+def test_recover_deleting_application(serve_instance):
+    """Test that replicas that are stuck on __del__ when the controller crashes,
+    is properly recovered when the controller is recovered.
+
+    This is similar to the test test_replica_deletion_after_controller_recover,
+    except what's blocking the deployment is __del__ instead of ongoing requests
+    """
+
+    signal = SignalActor.remote()
+
+    @serve.deployment
+    class A:
+        async def __del__(self):
+            await signal.wait.remote()
+
+    id = DeploymentID(name="A")
+    serve.run(A.bind())
+
+    @ray.remote
+    def delete_task():
+        serve.delete(SERVE_DEFAULT_APP_NAME)
+
+    # Delete application and make sure it is stuck on deleting
+    delete_ref = delete_task.remote()
+    print("Started task to delete application `default`")
+
+    def application_deleting():
+        # Confirm application is in deleting state
+        app_status = serve.status().applications[SERVE_DEFAULT_APP_NAME]
+        assert app_status.status == "DELETING"
+
+        # Confirm deployment is in updating state
+        status = serve_instance.get_all_deployment_statuses()[0]
+        assert status.name == "A" and status.status == "UPDATING"
+
+        # Confirm replica is stopping
+        replicas = ray.get(
+            serve_instance._controller._dump_replica_states_for_testing.remote(id)
+        )
+        assert replicas.count(states=[ReplicaState.STOPPING]) == 1
+
+        # Confirm delete task is still blocked
+        finished, pending = ray.wait([delete_ref], timeout=0)
+        assert pending and not finished
+        return True
+
+    def check_deleted():
+        deployment_statuses = serve_instance.get_all_deployment_statuses()
+        if len(deployment_statuses) != 0:
+            return False
+
+        finished, pending = ray.wait([delete_ref], timeout=0)
+        return finished and not pending
+
+    wait_for_condition(application_deleting)
+    for _ in range(10):
+        time.sleep(0.1)
+        assert application_deleting()
+
+    print("Confirmed that application `default` is stuck on deleting.")
+
+    # Kill controller while the application is stuck on deleting
+    ray.kill(serve.context._global_client._controller, no_restart=False)
+    print("Finished killing the controller (with restart).")
+
+    def check_controller_alive():
+        all_current_actors = list_actors(filters=[("state", "=", "ALIVE")])
+        for actor in all_current_actors:
+            if actor["class_name"] == "ServeController":
+                return True
+        return False
+
+    wait_for_condition(check_controller_alive)
+    print("Controller is back alive.")
+
+    wait_for_condition(application_deleting)
+    # Before we send the signal, the application should still be deleting
+    for _ in range(10):
+        time.sleep(0.1)
+        assert application_deleting()
+
+    print("Confirmed that application is still stuck on deleting.")
+
+    # Since we've confirmed the replica is in a stopping state, we can grab
+    # the reference to the in-progress graceful shutdown task
+    replicas = ray.get(
+        serve_instance._controller._dump_replica_states_for_testing.remote(id)
+    )
+    graceful_shutdown_ref = replicas.get()[0]._actor._graceful_shutdown_ref
+
+    signal.send.remote()
+    print("Sent signal to unblock deletion of application")
+    wait_for_condition(check_deleted)
+    print("Confirmed that application finished deleting and delete task has returned.")
+
+    # Make sure graceful shutdown ran successfully
+    ray.get(graceful_shutdown_ref)
+    print("Confirmed that graceful shutdown ran successfully.")
+
+
+def test_controller_crashes_with_logging_config(serve_instance):
+    """Controller persists logging config into kv store, and when controller recover
+    from crash, it will read logging config from kv store and apply to the
+    controller and proxy.
+    """
+
+    @serve.deployment
+    class Model:
+        def __init__(self):
+            self.logger = logging.getLogger("ray.serve")
+
+        def __call__(self):
+            self.logger.debug("this_is_debug_info")
+            return
+
+    serve.run(Model.bind())
+
+    client = serve_instance
+
+    # Update the logging config
+    client.update_global_logging_config(
+        LoggingConfig(encoding="JSON", log_level="DEBUG")
+    )
+
+    def check_log_file(log_file: str, expected_regex: list):
+        with open(log_file, "r") as f:
+            s = f.read()
+            for regex in expected_regex:
+                assert re.findall(regex, s) != []
+        return True
+
+    # Check the controller update
+    def check_log_state():
+        logging_config, _ = ray.get(client._controller._get_logging_config.remote())
+        assert logging_config.encoding == "JSON"
+        assert logging_config.log_level == "DEBUG"
+        return True
+
+    wait_for_condition(check_log_state, timeout=60)
+    _, log_file_path = ray.get(client._controller._get_logging_config.remote())
+    # DEBUG level check & JSON check
+    check_log_file(
+        log_file_path,
+        [".*Configure the serve controller logger.*", '.*"component_name":.*'],
+    )
+
+    ray.kill(client._controller, no_restart=False)
+
+    def check_controller_alive():
+        all_current_actors = list_actors(filters=[("state", "=", "ALIVE")])
+        for actor in all_current_actors:
+            if actor["class_name"] == "ServeController":
+                return True
+        return False
+
+    wait_for_condition(check_controller_alive)
+
+    # Check the controller log config
+    wait_for_condition(check_log_state)
+    _, new_log_file_path = ray.get(client._controller._get_logging_config.remote())
+    assert new_log_file_path != log_file_path
+    # Check again, make sure the logging config is recovered.
+    check_log_file(new_log_file_path, ['.*"component_name":.*'])
+
+    # Check proxy logging
+    def check_proxy_handle_in_controller():
+        proxy_handles = ray.get(client._controller.get_proxies.remote())
+        expected_num_proxies = 1
+        if RAY_SERVE_ENABLE_HA_PROXY:
+            # fallback proxy
+            expected_num_proxies += 1
+        assert len(proxy_handles) == expected_num_proxies
+        return True
+
+    wait_for_condition(check_proxy_handle_in_controller)
+    proxy_handles = ray.get(client._controller.get_proxies.remote())
+    proxy_handle = list(proxy_handles.values())[0]
+    file_path = ray.get(proxy_handle._get_logging_config.remote())
+    # We should see the health check debug log in the proxy logs.
+    wait_for_condition(
+        check_log_file,
+        log_file=file_path,
+        expected_regex=['"message": "Received health check."'],
+        timeout=15,  # The health check period is 10 seconds.
+    )
+
+
+def test_controller_recover_and_deploy(serve_instance):
+    """Ensure that in-progress deploy can finish even after controller dies."""
+    client = serve_instance
+    signal = SignalActor.options(name="signal123").remote()
+
+    config_json = {
+        "applications": [
+            {
+                "name": SERVE_DEFAULT_APP_NAME,
+                "import_path": "ray.serve.tests.test_config_files.hangs.app",
+            }
+        ]
+    }
+    config = ServeDeploySchema.model_validate(config_json)
+    client.deploy_apps(config)
+
+    wait_for_condition(
+        lambda: serve.status().applications["default"].status == "DEPLOYING"
+    )
+    ray.kill(client._controller, no_restart=False)
+
+    signal.send.remote()
+
+    # When controller restarts, it should redeploy config automatically
+    wait_for_condition(
+        lambda: httpx.get(f"{get_application_url()}/").text == "hello world"
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))