dory-sdk 2.1.0__py3-none-any.whl → 2.1.4__py3-none-any.whl

dory/health/server.py

@@ -5,17 +5,32 @@ Provides endpoints for:
 - /healthz - Liveness probe
 - /ready - Readiness probe (matches Kubernetes convention)
 - /metrics - Prometheus metrics
-- /state - State transfer (GET/POST) for pod migration
+- /state - State transfer (GET/POST) for pod migration (authenticated, with timeout)
 - /prestop - PreStop hook handler for graceful shutdown
 """
 
+import json
 import logging
+import os
+import secrets
+import time
+from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
 from typing import TYPE_CHECKING, Callable, Awaitable
 
 from aiohttp import web
 
 from dory.health.probes import LivenessProbe, ReadinessProbe
 from dory.utils.errors import DoryHealthError
+from dory.migration.transfer import (
+    TransferConfig,
+    validate_state_size,
+    StateSizeExceeded,
+    StateTransferTimeout,
+    log_transfer_summary,
+    TransferMetrics,
+    ORCHESTRATOR_STATE_TIMEOUT_SEC,
+    ORCHESTRATOR_MAX_STATE_SIZE,
+)
 
 if TYPE_CHECKING:
     from dory.metrics.collector import MetricsCollector
@@ -50,6 +65,8 @@ class HealthServer:
         state_getter: StateGetter | None = None,
         state_restorer: StateRestorer | None = None,
         prestop_handler: PreStopHandler | None = None,
+        state_token: str | None = None,
+        transfer_config: TransferConfig | None = None,
     ):
         """
         Initialize health server.
@@ -63,6 +80,11 @@ class HealthServer:
             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
+            state_token: Authentication token for /state endpoints. If not provided,
+                reads from DORY_STATE_TOKEN environment variable. If neither is set,
+                state endpoints are unauthenticated (not recommended in production).
+            transfer_config: Configuration for state transfer timeouts and size limits.
+                If not provided, uses defaults aligned with Orchestrator (25s timeout, 8MB max).
         """
         self._port = port
         self._health_path = health_path
@@ -73,6 +95,26 @@ class HealthServer:
         self._state_restorer = state_restorer
         self._prestop_handler = prestop_handler
 
+        # State endpoint authentication token (matches Orchestrator's DORY_STATE_TOKEN)
+        self._state_token = state_token or os.environ.get("DORY_STATE_TOKEN")
+        if not self._state_token:
+            logger.warning(
+                "DORY_STATE_TOKEN not configured - state endpoints are unauthenticated. "
+                "Set DORY_STATE_TOKEN environment variable for production deployments."
+            )
+
+        # State transfer configuration (aligned with Orchestrator limits)
+        self._transfer_config = transfer_config or TransferConfig()
+        logger.info(
+            f"State transfer configured: capture_timeout={self._transfer_config.capture_timeout_sec}s "
+            f"(Orchestrator: {ORCHESTRATOR_STATE_TIMEOUT_SEC}s), "
+            f"max_size={self._transfer_config.max_size_bytes:,}B "
+            f"(Orchestrator: {ORCHESTRATOR_MAX_STATE_SIZE:,}B)"
+        )
+
+        # Thread pool for running synchronous state getter with timeout
+        self._executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="state-capture")
+
         self._liveness = LivenessProbe()
         self._readiness = ReadinessProbe()
 
@@ -110,6 +152,44 @@ class HealthServer:
         """Set the callback for PreStop hook."""
         self._prestop_handler = handler
 
+    def set_state_token(self, token: str) -> None:
+        """Set the authentication token for state endpoints."""
+        self._state_token = token
+
+    def _verify_state_auth(self, request: web.Request) -> tuple[bool, str]:
+        """
+        Verify authentication for state endpoints.
+
+        Uses Bearer token authentication matching Orchestrator's implementation.
+        Token comparison uses constant-time comparison to prevent timing attacks.
+
+        Args:
+            request: The incoming HTTP request
+
+        Returns:
+            Tuple of (is_authenticated, error_message)
+        """
+        # If no token configured, allow unauthenticated access (with warning logged at startup)
+        if not self._state_token:
+            return True, ""
+
+        auth_header = request.headers.get("Authorization", "")
+
+        if not auth_header:
+            return False, "Missing Authorization header"
+
+        # Expected format: "Bearer <token>"
+        if not auth_header.startswith("Bearer "):
+            return False, "Invalid Authorization header format (expected 'Bearer <token>')"
+
+        provided_token = auth_header[7:]  # Strip "Bearer " prefix
+
+        # Use constant-time comparison to prevent timing attacks
+        if not secrets.compare_digest(provided_token, self._state_token):
+            return False, "Invalid authentication token"
+
+        return True, ""
+
     async def start(self) -> None:
         """
         Start the health server.
@@ -140,7 +220,11 @@ class HealthServer:
         """Stop the health server."""
         if self._runner:
             await self._runner.cleanup()
-            logger.info("Health server stopped")
+
+        # Shutdown executor
+        self._executor.shutdown(wait=False)
+
+        logger.info("Health server stopped")
 
     def _setup_routes(self) -> None:
         """Configure HTTP routes."""
@@ -227,9 +311,21 @@ class HealthServer:
         Called by Dory Orchestrator to capture state from old pod
         before transferring to new pod.
 
+        Requires authentication via Bearer token (DORY_STATE_TOKEN).
+        Enforces timeout and size limits to prevent Orchestrator timeouts.
+
         Returns:
             JSON response with processor state
         """
+        # Verify authentication
+        is_authenticated, auth_error = self._verify_state_auth(request)
+        if not is_authenticated:
+            logger.warning(f"State GET authentication failed: {auth_error}")
+            return web.json_response(
+                {"error": auth_error},
+                status=401,
+            )
+
         if self._state_getter is None:
             logger.warning("State getter not configured, returning empty state")
             return web.json_response({
@@ -237,10 +333,105 @@ class HealthServer:
                 "data": {},
             }, status=503)
 
+        start_time = time.monotonic()
+        metrics: TransferMetrics | None = None
+
         try:
-            state = self._state_getter()
-            logger.info("State captured for transfer", extra={"state_keys": list(state.keys())})
+            # Run state getter with timeout in thread pool
+            # (state_getter is synchronous, but we need timeout control)
+            loop = __import__("asyncio").get_event_loop()
+            state = await loop.run_in_executor(
+                self._executor,
+                self._state_getter,
+            )
+
+            # Check if we're already over timeout (executor doesn't enforce timeout)
+            capture_duration = time.monotonic() - start_time
+            if capture_duration > self._transfer_config.capture_timeout_sec:
+                metrics = TransferMetrics(
+                    duration_sec=capture_duration,
+                    size_bytes=0,
+                    size_ratio=0,
+                    timed_out=True,
+                    size_exceeded=False,
+                )
+                log_transfer_summary("capture", metrics, self._transfer_config)
+                return web.json_response(
+                    {
+                        "error": f"State capture timed out after {capture_duration:.1f}s "
+                        f"(limit: {self._transfer_config.capture_timeout_sec}s)",
+                        "timeout": True,
+                    },
+                    status=504,  # Gateway Timeout
+                )
+
+            # Serialize and validate size
+            state_json = json.dumps(state)
+            size_bytes = len(state_json.encode("utf-8"))
+
+            # Validate size
+            try:
+                validate_state_size(
+                    state_json,
+                    max_size=self._transfer_config.max_size_bytes,
+                    warn_threshold=self._transfer_config.size_warn_threshold,
+                )
+            except StateSizeExceeded as e:
+                metrics = TransferMetrics(
+                    duration_sec=time.monotonic() - start_time,
+                    size_bytes=size_bytes,
+                    size_ratio=size_bytes / self._transfer_config.max_size_bytes,
+                    timed_out=False,
+                    size_exceeded=True,
+                )
+                log_transfer_summary("capture", metrics, self._transfer_config)
+                return web.json_response(
+                    {
+                        "error": str(e),
+                        "size_exceeded": True,
+                        "size_bytes": size_bytes,
+                        "max_bytes": self._transfer_config.max_size_bytes,
+                    },
+                    status=413,  # Payload Too Large
+                )
+
+            # Success - log metrics
+            metrics = TransferMetrics(
+                duration_sec=time.monotonic() - start_time,
+                size_bytes=size_bytes,
+                size_ratio=size_bytes / self._transfer_config.max_size_bytes,
+                timed_out=False,
+                size_exceeded=False,
+            )
+            log_transfer_summary("capture", metrics, self._transfer_config)
+
+            logger.info(
+                "State captured for transfer",
+                extra={
+                    "state_keys": list(state.keys()),
+                    "size_bytes": size_bytes,
+                    "duration_sec": metrics.duration_sec,
+                },
+            )
             return web.json_response(state)
+
+        except FuturesTimeoutError:
+            metrics = TransferMetrics(
+                duration_sec=time.monotonic() - start_time,
+                size_bytes=0,
+                size_ratio=0,
+                timed_out=True,
+                size_exceeded=False,
+            )
+            log_transfer_summary("capture", metrics, self._transfer_config)
+            return web.json_response(
+                {
+                    "error": f"State capture timed out after {self._transfer_config.capture_timeout_sec}s",
+                    "timeout": True,
+                },
+                status=504,
+            )
+
         except Exception as e:
             logger.error(f"Failed to capture state: {e}")
             return web.json_response(
@@ -255,21 +446,104 @@ class HealthServer:
         Called by Dory Orchestrator to restore state to new pod
         after capturing from old pod.
 
+        Requires authentication via Bearer token (DORY_STATE_TOKEN).
+        Enforces timeout to prevent Orchestrator timeouts.
+
         Returns:
             JSON response confirming state restoration
         """
+        import asyncio
+
+        # Verify authentication
+        is_authenticated, auth_error = self._verify_state_auth(request)
+        if not is_authenticated:
+            logger.warning(f"State POST authentication failed: {auth_error}")
+            return web.json_response(
+                {"error": auth_error},
+                status=401,
+            )
+
         if self._state_restorer is None:
             logger.warning("State restorer not configured")
             return web.json_response({
                 "error": "state_restorer not configured",
             }, status=503)
 
+        start_time = time.monotonic()
+
         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"})
+            # Read request body
+            body = await request.read()
+            size_bytes = len(body)
+
+            # Validate incoming state size
+            if size_bytes > self._transfer_config.max_size_bytes:
+                logger.error(
+                    f"Incoming state size ({size_bytes:,}B) exceeds maximum "
+                    f"({self._transfer_config.max_size_bytes:,}B)"
+                )
+                return web.json_response(
+                    {
+                        "error": f"State size ({size_bytes:,} bytes) exceeds maximum "
+                        f"({self._transfer_config.max_size_bytes:,} bytes)",
+                        "size_exceeded": True,
+                    },
+                    status=413,
+                )
+
+            state = json.loads(body.decode("utf-8"))
+            logger.info(
+                "Restoring state from transfer",
+                extra={"state_keys": list(state.keys()), "size_bytes": size_bytes},
+            )
+
+            # Execute restore with timeout
+            try:
+                await asyncio.wait_for(
+                    self._state_restorer(state),
+                    timeout=self._transfer_config.restore_timeout_sec,
+                )
+            except asyncio.TimeoutError:
+                duration = time.monotonic() - start_time
+                logger.error(
+                    f"State restore timed out after {duration:.1f}s "
+                    f"(limit: {self._transfer_config.restore_timeout_sec}s)"
+                )
+                return web.json_response(
+                    {
+                        "error": f"State restore timed out after {self._transfer_config.restore_timeout_sec}s",
+                        "timeout": True,
+                    },
+                    status=504,
+                )
+
+            duration = time.monotonic() - start_time
+            logger.info(
+                f"State restored successfully in {duration:.2f}s",
+                extra={"size_bytes": size_bytes, "duration_sec": duration},
+            )
+
+            # Warn if restore took significant time
+            if duration > self._transfer_config.restore_timeout_sec * 0.5:
+                logger.warning(
+                    f"State restore took {duration:.1f}s "
+                    f"({duration/self._transfer_config.restore_timeout_sec:.0%} of "
+                    f"{self._transfer_config.restore_timeout_sec}s timeout)"
+                )
+
+            return web.json_response({
+                "status": "ok",
+                "message": "State restored",
+                "duration_sec": round(duration, 3),
+                "size_bytes": size_bytes,
+            })
+
+        except json.JSONDecodeError as e:
+            logger.error(f"Invalid JSON in state restore request: {e}")
+            return web.json_response(
+                {"error": f"Invalid JSON: {e}"},
+                status=400,
+            )
         except Exception as e:
             logger.error(f"Failed to restore state: {e}")
             return web.json_response(

dory/k8s/__init__.py

@@ -3,9 +3,78 @@
 from dory.k8s.client import K8sClient
 from dory.k8s.pod_metadata import PodMetadata
 from dory.k8s.annotation_watcher import AnnotationWatcher
+from dory.k8s.labels import (
+    # Label keys
+    LABEL_MANAGED_BY,
+    LABEL_APP_NAME,
+    LABEL_PROCESSOR_ID,
+    LABEL_WORKLOAD_TYPE,
+    LABEL_WORKLOAD_LOCATION,
+    LABEL_NODE_TYPE,
+    LABEL_MIGRATED_FROM_EDGE,
+    LABEL_ORIGINAL_NODE,
+    # Label values
+    VALUE_ORCHESTRATOR_NAME,
+    VALUE_EDGE,
+    VALUE_MANAGED,
+    # Enums
+    WorkloadLocation,
+    NodeType,
+    # Builder
+    DoryLabels,
+    # Utilities
+    get_label,
+    get_app_name,
+    get_processor_id,
+    get_workload_location,
+    is_managed_by_dory,
+    is_edge_workload,
+    is_migrated_from_edge,
+    get_original_node,
+    detect_workload_context,
+    edge_node_selector,
+    managed_node_selector,
+    edge_toleration,
+    get_contract_documentation,
+    CONTRACT_VERSION,
+)
 
 __all__ = [
+    # Core
     "K8sClient",
     "PodMetadata",
     "AnnotationWatcher",
+    # Label keys
+    "LABEL_MANAGED_BY",
+    "LABEL_APP_NAME",
+    "LABEL_PROCESSOR_ID",
+    "LABEL_WORKLOAD_TYPE",
+    "LABEL_WORKLOAD_LOCATION",
+    "LABEL_NODE_TYPE",
+    "LABEL_MIGRATED_FROM_EDGE",
+    "LABEL_ORIGINAL_NODE",
+    # Label values
+    "VALUE_ORCHESTRATOR_NAME",
+    "VALUE_EDGE",
+    "VALUE_MANAGED",
+    # Enums
+    "WorkloadLocation",
+    "NodeType",
+    # Builder
+    "DoryLabels",
+    # Utilities
+    "get_label",
+    "get_app_name",
+    "get_processor_id",
+    "get_workload_location",
+    "is_managed_by_dory",
+    "is_edge_workload",
+    "is_migrated_from_edge",
+    "get_original_node",
+    "detect_workload_context",
+    "edge_node_selector",
+    "managed_node_selector",
+    "edge_toleration",
+    "get_contract_documentation",
+    "CONTRACT_VERSION",
 ]