chore: import upstream snapshot with attribution

python/ray/dashboard/modules/reporter/tests/test_healthz.py

@@ -0,0 +1,121 @@
+import signal
+import sys
+
+import pytest
+import requests
+
+import ray._private.ray_constants as ray_constants
+from ray._common.network_utils import find_free_port
+from ray._common.test_utils import wait_for_condition
+from ray.tests.conftest import *  # noqa: F401 F403
+
+
+def test_healthz_head(monkeypatch, ray_start_cluster):
+    dashboard_port = find_free_port()
+    h = ray_start_cluster.add_node(dashboard_port=dashboard_port)
+    uri = f"http://localhost:{dashboard_port}/api/gcs_healthz"
+    wait_for_condition(lambda: requests.get(uri).status_code == 200)
+    h.all_processes[ray_constants.PROCESS_TYPE_GCS_SERVER][0].process.kill()
+    # It'll either timeout or just return an error
+    try:
+        wait_for_condition(lambda: requests.get(uri, timeout=1) != 200, timeout=4)
+    except RuntimeError as e:
+        assert "Read timed out" in str(e)
+
+
+def test_healthz_agent_1(monkeypatch, ray_start_cluster):
+    agent_port = find_free_port()
+    h = ray_start_cluster.add_node(dashboard_agent_listen_port=agent_port)
+    uri = f"http://{h.node_ip_address}:{agent_port}/api/local_raylet_healthz"
+
+    wait_for_condition(lambda: requests.get(uri).status_code == 200)
+
+    h.all_processes[ray_constants.PROCESS_TYPE_GCS_SERVER][0].process.kill()
+    # GCS's failure will not lead to healthz failure
+    assert requests.get(uri).status_code == 200
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="SIGSTOP only on posix")
+def test_healthz_agent_2(monkeypatch, ray_start_cluster):
+    monkeypatch.setenv("RAY_health_check_failure_threshold", "3")
+    monkeypatch.setenv("RAY_health_check_timeout_ms", "100")
+    monkeypatch.setenv("RAY_health_check_period_ms", "1000")
+    monkeypatch.setenv("RAY_health_check_initial_delay_ms", "0")
+
+    agent_port = find_free_port()
+    h = ray_start_cluster.add_node(dashboard_agent_listen_port=agent_port)
+    uri = f"http://{h.node_ip_address}:{agent_port}/api/local_raylet_healthz"
+
+    wait_for_condition(lambda: requests.get(uri).status_code == 200)
+
+    h.all_processes[ray_constants.PROCESS_TYPE_RAYLET][0].process.send_signal(
+        signal.SIGSTOP
+    )
+
+    # GCS still think raylet is alive.
+    assert requests.get(uri).status_code == 200
+    # But after heartbeat timeout, it'll think the raylet is down.
+    wait_for_condition(lambda: requests.get(uri).status_code != 200)
+
+
+def test_unified_healthz_head(monkeypatch, ray_start_cluster):
+    agent_port = find_free_port()
+    h = ray_start_cluster.add_node(dashboard_agent_listen_port=agent_port)
+    uri = f"http://{h.node_ip_address}:{agent_port}/api/healthz"
+
+    wait_for_condition(lambda: requests.get(uri).status_code == 200)
+    resp = requests.get(uri)
+    assert "raylet: success" in resp.text
+    assert "gcs: success" in resp.text
+
+    h.all_processes[ray_constants.PROCESS_TYPE_GCS_SERVER][0].process.kill()
+    wait_for_condition(lambda: requests.get(uri).status_code == 503)
+    resp = requests.get(uri)
+    assert "gcs: " in resp.text
+    assert "gcs: success" not in resp.text
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="SIGSTOP only on posix")
+def test_unified_healthz_worker(monkeypatch, ray_start_cluster):
+    monkeypatch.setenv("RAY_health_check_failure_threshold", "3")
+    monkeypatch.setenv("RAY_health_check_timeout_ms", "100")
+    monkeypatch.setenv("RAY_health_check_period_ms", "1000")
+    monkeypatch.setenv("RAY_health_check_initial_delay_ms", "0")
+
+    ray_start_cluster.add_node()
+    agent_port = find_free_port()
+    h = ray_start_cluster.add_node(dashboard_agent_listen_port=agent_port)
+    uri = f"http://{h.node_ip_address}:{agent_port}/api/healthz"
+
+    wait_for_condition(lambda: requests.get(uri).status_code == 200)
+    resp = requests.get(uri)
+    assert "gcs: success (no local gcs)" in resp.text
+
+    # Stop local raylet and verify this makes /healthz fail.
+    h.all_processes[ray_constants.PROCESS_TYPE_RAYLET][0].process.send_signal(
+        signal.SIGSTOP
+    )
+    wait_for_condition(lambda: requests.get(uri).status_code == 503)
+    resp = requests.get(uri)
+    assert "raylet: Local Raylet failed" in resp.text
+
+
+def test_unified_healthz_worker_gcs_down(monkeypatch, ray_start_cluster):
+    h_head = ray_start_cluster.add_node()
+    agent_port = find_free_port()
+    h_worker = ray_start_cluster.add_node(dashboard_agent_listen_port=agent_port)
+    uri = f"http://{h_worker.node_ip_address}:{agent_port}/api/healthz"
+
+    wait_for_condition(lambda: requests.get(uri).status_code == 200)
+    resp = requests.get(uri)
+    assert "gcs: success (no local gcs)" in resp.text
+
+    # Stop the head GCS server.
+    h_head.all_processes[ray_constants.PROCESS_TYPE_GCS_SERVER][0].process.kill()
+
+    # Worker health check should still succeed.
+    assert requests.get(uri).status_code == 200
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))