dory-sdk 2.1.0__py3-none-any.whl

dory/health/server.py

@@ -0,0 +1,306 @@
+"""
+Health and metrics HTTP server.
+
+Provides endpoints for:
+- /healthz - Liveness probe
+- /ready - Readiness probe (matches Kubernetes convention)
+- /metrics - Prometheus metrics
+- /state - State transfer (GET/POST) for pod migration
+- /prestop - PreStop hook handler for graceful shutdown
+"""
+
+import logging
+from typing import TYPE_CHECKING, Callable, Awaitable
+
+from aiohttp import web
+
+from dory.health.probes import LivenessProbe, ReadinessProbe
+from dory.utils.errors import DoryHealthError
+
+if TYPE_CHECKING:
+    from dory.metrics.collector import MetricsCollector
+
+logger = logging.getLogger(__name__)
+
+# Type aliases for callbacks
+StateGetter = Callable[[], dict]
+StateRestorer = Callable[[dict], Awaitable[None]]
+PreStopHandler = Callable[[], Awaitable[None]]
+
+
+class HealthServer:
+    """
+    HTTP server for health probes, metrics, and state transfer.
+
+    Runs on a separate port from the main application.
+    Provides endpoints required by Dory Orchestrator for:
+    - Health probes (liveness/readiness)
+    - Prometheus metrics
+    - State transfer during pod migration
+    - PreStop hook for graceful shutdown
+    """
+
+    def __init__(
+        self,
+        port: int = 8080,
+        health_path: str = "/healthz",
+        ready_path: str = "/ready",  # Changed from /readyz to match Orchestrator
+        metrics_path: str = "/metrics",
+        metrics_collector: "MetricsCollector | None" = None,
+        state_getter: StateGetter | None = None,
+        state_restorer: StateRestorer | None = None,
+        prestop_handler: PreStopHandler | None = None,
+    ):
+        """
+        Initialize health server.
+
+        Args:
+            port: Port to listen on
+            health_path: Path for liveness probe
+            ready_path: Path for readiness probe
+            metrics_path: Path for Prometheus metrics
+            metrics_collector: Optional metrics collector for /metrics endpoint
+            state_getter: Callback to get processor state for /state GET
+            state_restorer: Callback to restore processor state for /state POST
+            prestop_handler: Callback for /prestop PreStop hook
+        """
+        self._port = port
+        self._health_path = health_path
+        self._ready_path = ready_path
+        self._metrics_path = metrics_path
+        self._metrics_collector = metrics_collector
+        self._state_getter = state_getter
+        self._state_restorer = state_restorer
+        self._prestop_handler = prestop_handler
+
+        self._liveness = LivenessProbe()
+        self._readiness = ReadinessProbe()
+
+        self._app: web.Application | None = None
+        self._runner: web.AppRunner | None = None
+        self._site: web.TCPSite | None = None
+
+    @property
+    def liveness_probe(self) -> LivenessProbe:
+        """Get liveness probe for adding custom checks."""
+        return self._liveness
+
+    @property
+    def readiness_probe(self) -> ReadinessProbe:
+        """Get readiness probe for adding custom checks."""
+        return self._readiness
+
+    def mark_ready(self) -> None:
+        """Mark the application as ready to receive traffic."""
+        self._readiness.mark_ready()
+
+    def mark_not_ready(self) -> None:
+        """Mark the application as not ready."""
+        self._readiness.mark_not_ready()
+
+    def set_state_getter(self, getter: StateGetter) -> None:
+        """Set the callback for getting processor state."""
+        self._state_getter = getter
+
+    def set_state_restorer(self, restorer: StateRestorer) -> None:
+        """Set the callback for restoring processor state."""
+        self._state_restorer = restorer
+
+    def set_prestop_handler(self, handler: PreStopHandler) -> None:
+        """Set the callback for PreStop hook."""
+        self._prestop_handler = handler
+
+    async def start(self) -> None:
+        """
+        Start the health server.
+
+        Raises:
+            DoryHealthError: If server fails to start
+        """
+        try:
+            self._app = web.Application()
+            self._setup_routes()
+
+            self._runner = web.AppRunner(self._app)
+            await self._runner.setup()
+
+            self._site = web.TCPSite(
+                self._runner,
+                host="0.0.0.0",
+                port=self._port,
+            )
+            await self._site.start()
+
+            logger.info(f"Health server started on port {self._port}")
+
+        except Exception as e:
+            raise DoryHealthError(f"Failed to start health server: {e}", cause=e)
+
+    async def stop(self) -> None:
+        """Stop the health server."""
+        if self._runner:
+            await self._runner.cleanup()
+            logger.info("Health server stopped")
+
+    def _setup_routes(self) -> None:
+        """Configure HTTP routes."""
+        self._app.router.add_get(self._health_path, self._handle_health)
+        self._app.router.add_get(self._ready_path, self._handle_ready)
+        self._app.router.add_get(self._metrics_path, self._handle_metrics)
+
+        # State transfer endpoints (required by Dory Orchestrator)
+        self._app.router.add_get("/state", self._handle_state_get)
+        self._app.router.add_post("/state", self._handle_state_post)
+
+        # PreStop hook endpoint (required by Dory Orchestrator)
+        self._app.router.add_get("/prestop", self._handle_prestop)
+
+        # Root endpoint for basic info
+        self._app.router.add_get("/", self._handle_root)
+
+    async def _handle_root(self, request: web.Request) -> web.Response:
+        """Handle root endpoint."""
+        return web.json_response({
+            "service": "dory-processor",
+            "endpoints": [
+                self._health_path,
+                self._ready_path,
+                self._metrics_path,
+                "/state",
+                "/prestop",
+            ],
+        })
+
+    async def _handle_health(self, request: web.Request) -> web.Response:
+        """
+        Handle liveness probe.
+
+        Returns 200 if alive, 503 if unhealthy.
+        """
+        result = await self._liveness.check()
+
+        status = 200 if result.healthy else 503
+        return web.json_response(result.to_dict(), status=status)
+
+    async def _handle_ready(self, request: web.Request) -> web.Response:
+        """
+        Handle readiness probe.
+
+        Returns 200 if ready, 503 if not ready.
+        """
+        result = await self._readiness.check()
+
+        status = 200 if result.healthy else 503
+        return web.json_response(result.to_dict(), status=status)
+
+    async def _handle_metrics(self, request: web.Request) -> web.Response:
+        """
+        Handle Prometheus metrics endpoint.
+
+        Returns metrics in Prometheus text format.
+        """
+        if self._metrics_collector is None:
+            return web.Response(
+                text="# No metrics collector configured\n",
+                content_type="text/plain",
+            )
+
+        try:
+            metrics_text = self._metrics_collector.export_prometheus()
+            return web.Response(
+                text=metrics_text,
+                content_type="text/plain; version=0.0.4",
+                charset="utf-8",
+            )
+        except Exception as e:
+            logger.error(f"Error exporting metrics: {e}")
+            return web.Response(
+                text=f"# Error exporting metrics: {e}\n",
+                content_type="text/plain",
+                status=500,
+            )
+
+    async def _handle_state_get(self, request: web.Request) -> web.Response:
+        """
+        Handle GET /state for state capture during migration.
+
+        Called by Dory Orchestrator to capture state from old pod
+        before transferring to new pod.
+
+        Returns:
+            JSON response with processor state
+        """
+        if self._state_getter is None:
+            logger.warning("State getter not configured, returning empty state")
+            return web.json_response({
+                "error": "state_getter not configured",
+                "data": {},
+            }, status=503)
+
+        try:
+            state = self._state_getter()
+            logger.info("State captured for transfer", extra={"state_keys": list(state.keys())})
+            return web.json_response(state)
+        except Exception as e:
+            logger.error(f"Failed to capture state: {e}")
+            return web.json_response(
+                {"error": f"Failed to capture state: {e}"},
+                status=500,
+            )
+
+    async def _handle_state_post(self, request: web.Request) -> web.Response:
+        """
+        Handle POST /state for state restoration during migration.
+
+        Called by Dory Orchestrator to restore state to new pod
+        after capturing from old pod.
+
+        Returns:
+            JSON response confirming state restoration
+        """
+        if self._state_restorer is None:
+            logger.warning("State restorer not configured")
+            return web.json_response({
+                "error": "state_restorer not configured",
+            }, status=503)
+
+        try:
+            state = await request.json()
+            logger.info("Restoring state from transfer", extra={"state_keys": list(state.keys())})
+            await self._state_restorer(state)
+            logger.info("State restored successfully")
+            return web.json_response({"status": "ok", "message": "State restored"})
+        except Exception as e:
+            logger.error(f"Failed to restore state: {e}")
+            return web.json_response(
+                {"error": f"Failed to restore state: {e}"},
+                status=500,
+            )
+
+    async def _handle_prestop(self, request: web.Request) -> web.Response:
+        """
+        Handle GET /prestop for PreStop hook.
+
+        Called by Kubernetes PreStop hook before pod termination.
+        Allows the application to prepare for graceful shutdown.
+
+        Returns:
+            JSON response confirming prestop handling
+        """
+        logger.info("PreStop hook invoked - preparing for shutdown")
+
+        # Mark as not ready to stop receiving new traffic
+        self._readiness.mark_not_ready()
+
+        if self._prestop_handler:
+            try:
+                await self._prestop_handler()
+                logger.info("PreStop handler completed")
+            except Exception as e:
+                logger.error(f"PreStop handler error: {e}")
+                # Continue anyway - don't block shutdown
+
+        return web.json_response({
+            "status": "ok",
+            "message": "PreStop hook processed, ready for termination",
+        })

dory/k8s/__init__.py

@@ -0,0 +1,11 @@
+"""Kubernetes integration utilities."""
+
+from dory.k8s.client import K8sClient
+from dory.k8s.pod_metadata import PodMetadata
+from dory.k8s.annotation_watcher import AnnotationWatcher
+
+__all__ = [
+    "K8sClient",
+    "PodMetadata",
+    "AnnotationWatcher",
+]

dory/k8s/annotation_watcher.py

@@ -0,0 +1,184 @@
+"""
+Annotation watcher for migration signals.
+
+Watches pod annotations for migration-related signals
+from the orchestrator.
+"""
+
+import asyncio
+import logging
+from typing import Callable, Any
+
+from dory.k8s.client import K8sClient
+from dory.utils.errors import DoryK8sError
+
+logger = logging.getLogger(__name__)
+
+
+class AnnotationWatcher:
+    """
+    Watches pod annotations for orchestrator signals.
+
+    Monitors annotations:
+    - dory.io/migration: "true" when migration imminent
+    - dory.io/shutdown: "true" when shutdown requested
+    - dory.io/snapshot: "true" when snapshot requested
+    """
+
+    MIGRATION_ANNOTATION = "dory.io/migration"
+    SHUTDOWN_ANNOTATION = "dory.io/shutdown"
+    SNAPSHOT_ANNOTATION = "dory.io/snapshot"
+    DEADLINE_ANNOTATION = "dory.io/migration-deadline"
+
+    def __init__(
+        self,
+        k8s_client: K8sClient,
+        pod_name: str,
+        poll_interval: float = 5.0,
+    ):
+        """
+        Initialize annotation watcher.
+
+        Args:
+            k8s_client: Kubernetes client
+            pod_name: Name of pod to watch
+            poll_interval: Seconds between polls
+        """
+        self._k8s_client = k8s_client
+        self._pod_name = pod_name
+        self._poll_interval = poll_interval
+
+        self._running = False
+        self._watch_task: asyncio.Task | None = None
+
+        # Callbacks
+        self._on_migration: Callable[[], Any] | None = None
+        self._on_shutdown: Callable[[], Any] | None = None
+        self._on_snapshot: Callable[[], Any] | None = None
+
+        # State tracking
+        self._last_annotations: dict[str, str] = {}
+
+    def on_migration(self, callback: Callable[[], Any]) -> None:
+        """Set callback for migration signal."""
+        self._on_migration = callback
+
+    def on_shutdown(self, callback: Callable[[], Any]) -> None:
+        """Set callback for shutdown signal."""
+        self._on_shutdown = callback
+
+    def on_snapshot(self, callback: Callable[[], Any]) -> None:
+        """Set callback for snapshot signal."""
+        self._on_snapshot = callback
+
+    async def start(self) -> None:
+        """Start watching annotations."""
+        if self._running:
+            return
+
+        self._running = True
+        self._watch_task = asyncio.create_task(self._watch_loop())
+        logger.info(f"Started annotation watcher for pod {self._pod_name}")
+
+    async def stop(self) -> None:
+        """Stop watching annotations."""
+        self._running = False
+
+        if self._watch_task:
+            self._watch_task.cancel()
+            try:
+                await self._watch_task
+            except asyncio.CancelledError:
+                pass
+            self._watch_task = None
+
+        logger.info("Annotation watcher stopped")
+
+    async def _watch_loop(self) -> None:
+        """Main watch loop."""
+        while self._running:
+            try:
+                await self._check_annotations()
+            except DoryK8sError as e:
+                logger.warning(f"Failed to check annotations: {e}")
+            except Exception as e:
+                logger.error(f"Unexpected error in annotation watcher: {e}")
+
+            await asyncio.sleep(self._poll_interval)
+
+    async def _check_annotations(self) -> None:
+        """Check annotations for changes."""
+        try:
+            annotations = await self._k8s_client.get_pod_annotations(self._pod_name)
+        except DoryK8sError:
+            # Pod might not exist yet or API unavailable
+            return
+
+        # Check migration annotation
+        if self._annotation_changed(self.MIGRATION_ANNOTATION, annotations, "true"):
+            logger.info("Migration signal detected")
+            if self._on_migration:
+                await self._invoke_callback(self._on_migration)
+
+        # Check shutdown annotation
+        if self._annotation_changed(self.SHUTDOWN_ANNOTATION, annotations, "true"):
+            logger.info("Shutdown signal detected")
+            if self._on_shutdown:
+                await self._invoke_callback(self._on_shutdown)
+
+        # Check snapshot annotation
+        if self._annotation_changed(self.SNAPSHOT_ANNOTATION, annotations, "true"):
+            logger.info("Snapshot signal detected")
+            if self._on_snapshot:
+                await self._invoke_callback(self._on_snapshot)
+            # Clear snapshot annotation after processing
+            await self._clear_annotation(self.SNAPSHOT_ANNOTATION)
+
+        self._last_annotations = annotations
+
+    def _annotation_changed(
+        self,
+        key: str,
+        new_annotations: dict[str, str],
+        trigger_value: str,
+    ) -> bool:
+        """Check if annotation changed to trigger value."""
+        old_value = self._last_annotations.get(key)
+        new_value = new_annotations.get(key)
+
+        return old_value != new_value and new_value == trigger_value
+
+    async def _invoke_callback(self, callback: Callable[[], Any]) -> None:
+        """Invoke callback, handling async/sync."""
+        try:
+            if asyncio.iscoroutinefunction(callback):
+                await callback()
+            else:
+                callback()
+        except Exception as e:
+            logger.error(f"Callback error: {e}")
+
+    async def _clear_annotation(self, key: str) -> None:
+        """Clear an annotation after processing."""
+        try:
+            await self._k8s_client.patch_pod_annotations(
+                self._pod_name,
+                {key: None},  # Setting to None removes the annotation
+            )
+        except DoryK8sError as e:
+            logger.warning(f"Failed to clear annotation {key}: {e}")
+
+    def get_migration_deadline(self) -> float | None:
+        """
+        Get migration deadline from annotations.
+
+        Returns:
+            Unix timestamp of deadline, or None
+        """
+        deadline_str = self._last_annotations.get(self.DEADLINE_ANNOTATION)
+        if deadline_str:
+            try:
+                return float(deadline_str)
+            except ValueError:
+                pass
+        return None