chore: import upstream snapshot with attribution

python/ray/serve/tests/test_serve_ha.py

@@ -0,0 +1,124 @@
+import sys
+import threading
+from time import sleep
+
+import pytest
+
+from ray._common.constants import HEAD_NODE_RESOURCE_NAME
+from ray._common.test_utils import wait_for_condition
+from ray.tests.conftest_docker import *  # noqa
+
+scripts = """
+import json
+import os
+
+import ray
+
+from ray import serve
+
+@serve.deployment
+class GetPID:
+    def __call__(self, *args):
+        return {{"pid": os.getpid()}}
+
+serve.run(GetPID.options(num_replicas={num_replicas}).bind())
+"""
+
+check_script = """
+import ray
+import requests
+
+@ray.remote
+def get_pid():
+    return requests.get("http://127.0.0.1:8000/").json()["pid"]
+
+pids = {{
+    requests.get("http://127.0.0.1:8000/").json()["pid"]
+    for _ in range(20)
+}}
+print(pids)
+assert len(pids) == {num_replicas}
+"""
+
+check_ray_nodes_script = """
+import ray
+
+ray.init(address="auto")
+print(ray.nodes())
+"""
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="Only works on linux.")
+def test_ray_serve_basic(docker_cluster):
+    # This test covers the basic cases for serve ha
+    # - It starts the serve on worker nodes.
+    # - Check the deployment is OK
+    # - Stop headnode
+    # - Check the serve app is running healthy
+    # - Start a reconfig (2 replicas) and it'll hang
+    # - Start head node. The script will continue once GCS is back
+    # - Make sure two replicas are there
+
+    # TODO(iycheng): Update serve to better integrate with GCS HA:
+    #   - Make sure no task can run in the raylet where GCS is deployed.
+
+    head, worker = docker_cluster
+    output = worker.exec_run(cmd=f"python -c '{scripts.format(num_replicas=1)}'")
+    assert output.exit_code == 0, output.output
+    assert b"Adding 1 replica to Deployment(" in output.output
+
+    output = worker.exec_run(cmd=f"python -c '{check_script.format(num_replicas=1)}'")
+    assert output.exit_code == 0, output.output
+
+    # Kill the head node.
+    head.kill()
+
+    # Make sure the Serve application can still handle traffic.
+    output = worker.exec_run(cmd=f"python -c '{check_script.format(num_replicas=1)}'")
+    assert output.exit_code == 0, output.output
+
+    # Scale up the application in a background thread so we can concurrently kill the
+    # head node.
+    def reconfig():
+        worker.exec_run(cmd=f"python -c '{scripts.format(num_replicas=2)}'")
+
+    t = threading.Thread(target=reconfig)
+    t.start()
+    sleep(1)
+
+    # Updating the application should continue once the head node is back online.
+    head.restart()
+    t.join()
+
+    # Ensure head node is up before calling check_script on the worker again.
+    def check_for_head_node_come_back_up():
+        _output = head.exec_run(cmd=f"python -c '{check_ray_nodes_script}'")
+        return (
+            _output.exit_code == 0
+            and bytes(HEAD_NODE_RESOURCE_NAME, "utf-8") in _output.output
+        )
+
+    wait_for_condition(check_for_head_node_come_back_up)
+
+    # Check that the application has been updated.
+    output = worker.exec_run(cmd=f"python -c '{check_script.format(num_replicas=2)}'")
+    assert output.exit_code == 0, output.output
+
+    # Make sure the serve controller still runs on the head node after restart.
+    check_controller_head_node_script = """
+import ray
+import requests
+from ray.serve.schema import ServeInstanceDetails
+from ray._common.constants import HEAD_NODE_RESOURCE_NAME
+ray.init(address="auto")
+head_node_id = ray.get_runtime_context().get_node_id()
+serve_details = ServeInstanceDetails(
+    **requests.get("http://localhost:8265/api/serve/applications/").json())
+assert serve_details.controller_info.node_id == head_node_id
+"""
+    output = head.exec_run(cmd=f"python -c '{check_controller_head_node_script}'")
+    assert output.exit_code == 0, output.output
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))