[dashboard][2/2] Add endpoints to dashboard and dashboard_agent for liveness check of raylet and gcs

dashboard/head.py

@@ -12,7 +12,7 @@
 import ray.experimental.internal_kv as internal_kv
 from ray._private import ray_constants
 from ray._private.gcs_pubsub import GcsAioErrorSubscriber, GcsAioLogSubscriber
-from ray._private.gcs_utils import GcsClient, check_health
+from ray._private.gcs_utils import GcsClient, GcsAioClient, check_health
 from ray.dashboard.datacenter import DataOrganizer
 from ray.dashboard.utils import async_loop_forever
 
@@ -169,9 +169,8 @@ async def run(self):
  # Dashboard will handle connection failure automatically
  self.gcs_client = GcsClient(address=gcs_address, nums_reconnect_retry=0)
  internal_kv._initialize_internal_kv(self.gcs_client)
- self.aiogrpc_gcs_channel = ray._private.utils.init_grpc_channel(
- gcs_address, GRPC_CHANNEL_OPTIONS, asynchronous=True
- )
+ self.gcs_aio_client = GcsAioClient(address=gcs_address)
+ self.aiogrpc_gcs_channel = self.gcs_aio_client.channel.channel()
 
  self.gcs_error_subscriber = GcsAioErrorSubscriber(address=gcs_address)
  self.gcs_log_subscriber = GcsAioLogSubscriber(address=gcs_address)

dashboard/modules/healthz/healthz_agent.py

@@ -0,0 +1,53 @@
+import ray.dashboard.utils as dashboard_utils
+import ray.dashboard.optional_utils as optional_utils
+from ray.dashboard.modules.healthz.utils import HealthChecker
+from aiohttp.web import Request, Response, HTTPServiceUnavailable
+import grpc
+
+routes = optional_utils.ClassMethodRouteTable
+
+
+class HealthzAgent(dashboard_utils.DashboardAgentModule):
+ """Health check in the agent.
+
+ This module adds health check related endpoint to the agent to check
+ local components' health.
+ """
+
+ def __init__(self, dashboard_agent):
+ super().__init__(dashboard_agent)
+ self._health_checker = HealthChecker(
+ dashboard_agent.gcs_aio_client,
+ f"{dashboard_agent.ip}:{dashboard_agent.node_manager_port}",
+ )
+
+ @routes.get("/api/local_raylet_healthz")
+ async def health_check(self, req: Request) -> Response:
+ try:
+ alive = await self._health_checker.check_local_raylet_liveness()
+ if alive is False:
+ return HTTPServiceUnavailable(reason="Local Raylet failed")
+ except grpc.RpcError as e:
+ # We only consider the error other than GCS unreachable as raylet failure
+ # to avoid false positive.
+ # In case of GCS failed, Raylet will crash eventually if GCS is not back
+ # within a given time and the check will fail since agent can't live
+ # without a local raylet.
+ if e.code() not in (
+ grpc.StatusCode.UNAVAILABLE,
+ grpc.StatusCode.UNKNOWN,
+ grpc.StatusCode.DEADLINE_EXCEEDED,
+ ):
+ return HTTPServiceUnavailable(reason=e.message())
+
+ return Response(
+ text="success",
+ content_type="application/text",
+ )
+
+ async def run(self, server):
+ pass
+
+ @staticmethod
+ def is_minimal_module():
+ return True

dashboard/modules/healthz/healthz_head.py

@@ -0,0 +1,40 @@
+import ray.dashboard.utils as dashboard_utils
+import ray.dashboard.optional_utils as optional_utils
+from ray.dashboard.modules.healthz.utils import HealthChecker
+from aiohttp.web import Request, Response, HTTPServiceUnavailable
+
+routes = optional_utils.ClassMethodRouteTable
+
+
+class HealthzHead(dashboard_utils.DashboardHeadModule):
+ """Health check in the head.
+
+ This module adds health check related endpoint to the head to check
+ GCS's heath.
+ """
+
+ def __init__(self, dashboard_head):
+ super().__init__(dashboard_head)
+ self._health_checker = HealthChecker(dashboard_head.gcs_aio_client)
+
+ @routes.get("/api/gcs_healthz")
+ async def health_check(self, req: Request) -> Response:
+ alive = False
+ try:
+ alive = await self._health_checker.check_gcs_liveness()
+ if alive is True:
+ return Response(
+ text="success",
+ content_type="application/text",
+ )
+ except Exception as e:
+ return HTTPServiceUnavailable(reason=f"Health check failed: {e}")
+
+ return HTTPServiceUnavailable(reason="Health check failed")
+
+ async def run(self, server):
+ pass
+
+ @staticmethod
+ def is_minimal_module():
+ return True

dashboard/modules/healthz/tests/test_healthz.py

@@ -0,0 +1,58 @@
+import sys
+import pytest
+import requests
+
+import ray._private.ray_constants as ray_constants
+from ray.tests.conftest import * # noqa: F401 F403
+from ray._private.test_utils import find_free_port, wait_for_condition
+
+
+def test_healthz_head(ray_start_cluster):
+ dashboard_port = find_free_port()
+ h = ray_start_cluster.add_node(dashboard_port=dashboard_port)
+ uri = f"http:https://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(ray_start_cluster):
+ agent_port = find_free_port()
+ h = ray_start_cluster.add_node(dashboard_agent_listen_port=agent_port)
+ uri = f"http:https://localhost:{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_num_heartbeats_timeout", "3")
+
+ agent_port = find_free_port()
+ h = ray_start_cluster.add_node(dashboard_agent_listen_port=agent_port)
+ uri = f"http:https://localhost:{agent_port}/api/local_raylet_healthz"
+
+ wait_for_condition(lambda: requests.get(uri).status_code == 200)
+
+ import signal
+
+ 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)
+
+
+if __name__ == "__main__":
+ sys.exit(pytest.main(["-v", __file__]))

dashboard/modules/healthz/utils.py

@@ -0,0 +1,23 @@
+from typing import Optional
+from ray._private.gcs_utils import GcsAioClient
+
+
+class HealthChecker:
+ def __init__(
+ self, gcs_aio_client: GcsAioClient, local_node_address: Optional[str] = None
+ ):
+ self._gcs_aio_client = gcs_aio_client
+ self._local_node_address = local_node_address
+
+ async def check_local_raylet_liveness(self) -> bool:
+ if self._local_node_address is None:
+ return False
+
+ liveness = await self._gcs_aio_client.check_alive(
+ [self._local_node_address.encode()], 1
+ )
+ return liveness[0]
+
+ async def check_gcs_liveness(self) -> bool:
+ await self._gcs_aio_client.check_alive([], 1)
+ return True