Feat: Worker Health Checks (hatchet-dev#289)

* feat: add prom client

* feat: naive healthcheck

* feat: add health check prom support + custom port

* fix: lint

* fix: prom gauge

* feat: minor version

* feat: conditionally run health check

* fix: use -1 for prom gauge

* fix: revert back to 0

Author: hatchet-temporary

hatchet_sdk/loader.py

@@ -43,6 +43,8 @@ def __init__(
         otel_service_name: str | None = None,
         otel_exporter_oltp_headers: dict[str, str] | None = None,
         otel_exporter_oltp_protocol: str | None = None,
+        worker_healthcheck_port: int | None = None,
+        worker_healthcheck_enabled: bool | None = None,
     ):
         self.tenant_id = tenant_id
         self.tls_config = tls_config
@@ -57,6 +59,8 @@ def __init__(
         self.otel_service_name = otel_service_name
         self.otel_exporter_oltp_headers = otel_exporter_oltp_headers
         self.otel_exporter_oltp_protocol = otel_exporter_oltp_protocol
+        self.worker_healthcheck_port = worker_healthcheck_port
+        self.worker_healthcheck_enabled = worker_healthcheck_enabled
 
         if not self.logInterceptor:
             self.logInterceptor = getLogger()
@@ -163,6 +167,23 @@ def get_config_value(key, env_var):
             "otel_exporter_oltp_protocol", "HATCHET_CLIENT_OTEL_EXPORTER_OTLP_PROTOCOL"
         )
 
+        worker_healthcheck_port = int(
+            get_config_value(
+                "worker_healthcheck_port", "HATCHET_CLIENT_WORKER_HEALTHCHECK_PORT"
+            )
+            or 8001
+        )
+
+        worker_healthcheck_enabled = (
+            str(
+                get_config_value(
+                    "worker_healthcheck_port",
+                    "HATCHET_CLIENT_WORKER_HEALTHCHECK_ENABLED",
+                )
+            )
+            == "True"
+        )
+
         return ClientConfig(
             tenant_id=tenant_id,
             tls_config=tls_config,
@@ -178,6 +199,8 @@ def get_config_value(key, env_var):
             otel_service_name=otel_service_name,
             otel_exporter_oltp_headers=otel_exporter_oltp_headers,
             otel_exporter_oltp_protocol=otel_exporter_oltp_protocol,
+            worker_healthcheck_port=worker_healthcheck_port,
+            worker_healthcheck_enabled=worker_healthcheck_enabled,
         )
 
     def _load_tls_config(self, tls_data: Dict, host_port) -> ClientTLSConfig:

hatchet_sdk/worker/worker.py

@@ -12,7 +12,10 @@
 from types import FrameType
 from typing import Any, Callable, TypeVar, get_type_hints
 
-from pydantic import BaseModel
+from aiohttp import web
+from aiohttp.web_request import Request
+from aiohttp.web_response import Response
+from prometheus_client import CONTENT_TYPE_LATEST, Gauge, generate_latest
 
 from hatchet_sdk import Context
 from hatchet_sdk.client import Client, new_client_raw
@@ -88,6 +91,10 @@ def __init__(
 
         self._setup_signal_handlers()
 
+        self.worker_status_gauge = Gauge(
+            "hatchet_worker_status", "Current status of the Hatchet worker"
+        )
+
     def register_function(self, action: str, func: Callable[[Context], Any]) -> None:
         self.action_registry[action] = func
 
@@ -155,6 +162,39 @@ def setup_loop(self, loop: asyncio.AbstractEventLoop | None = None) -> bool:
             created_loop = True
             return created_loop
 
+    async def health_check_handler(self, request: Request) -> Response:
+        status = self.status()
+
+        return web.json_response({"status": status.name})
+
+    async def metrics_handler(self, request: Request) -> Response:
+        self.worker_status_gauge.set(1 if self.status() == WorkerStatus.HEALTHY else 0)
+
+        return web.Response(body=generate_latest(), content_type="text/plain")
+
+    async def start_health_server(self) -> None:
+        port = self.config.worker_healthcheck_port or 8001
+
+        app = web.Application()
+        app.add_routes(
+            [
+                web.get("/health", self.health_check_handler),
+                web.get("/metrics", self.metrics_handler),
+            ]
+        )
+
+        runner = web.AppRunner(app)
+
+        try:
+            await runner.setup()
+            await web.TCPSite(runner, "0.0.0.0", port).start()
+        except Exception as e:
+            logger.error("failed to start healthcheck server")
+            logger.error(str(e))
+            return
+
+        logger.info(f"healthcheck server running on port {port}")
+
     def start(
         self, options: WorkerStartOptions = WorkerStartOptions()
     ) -> Future[asyncio.Task[Any] | None]:
@@ -196,6 +236,9 @@ async def async_start(
         if not _from_start:
             self.setup_loop(options.loop)
 
+        if self.config.worker_healthcheck_enabled:
+            await self.start_health_server()
+
         self.action_listener_process = self._start_listener()
 
         self.action_runner = self._run_action_runner()

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "hatchet-sdk"
-version = "0.42.5"
+version = "0.43.0"
 description = ""
 authors = ["Alexander Belanger <[email protected]>"]
 readme = "README.md"
@@ -28,6 +28,7 @@ opentelemetry-instrumentation = "^0.48b0"
 opentelemetry-distro = "^0.48b0"
 opentelemetry-exporter-otlp = "^1.27.0"
 opentelemetry-exporter-otlp-proto-http = "^1.27.0"
+prometheus-client = "^0.21.1"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.2.2"