dory-processor-sdk 0.0.1__py3-none-any.whl

dory/k8s/__init__.py

@@ -0,0 +1,80 @@
+"""Kubernetes integration utilities."""
+
+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",
+]

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

dory/k8s/client.py

@@ -0,0 +1,251 @@
+"""
+Kubernetes client wrapper.
+
+Provides simplified interface to Kubernetes API.
+"""
+
+import logging
+from typing import Any
+
+from dory.utils.errors import DoryK8sError
+
+logger = logging.getLogger(__name__)
+
+# Optional kubernetes import
+try:
+    from kubernetes import client, config
+    from kubernetes.client.rest import ApiException
+    K8S_AVAILABLE = True
+except ImportError:
+    K8S_AVAILABLE = False
+    client = None
+    config = None
+    ApiException = Exception
+
+
+class K8sClient:
+    """
+    Kubernetes API client wrapper.
+
+    Handles configuration loading and provides
+    simplified access to common operations.
+    """
+
+    def __init__(self, namespace: str | None = None):
+        """
+        Initialize Kubernetes client.
+
+        Args:
+            namespace: Kubernetes namespace (auto-detected if not provided)
+        """
+        self._namespace = namespace
+        self._core_api: Any = None
+        self._initialized = False
+
+    def _ensure_initialized(self) -> None:
+        """Initialize Kubernetes client if not already done."""
+        if self._initialized:
+            return
+
+        if not K8S_AVAILABLE:
+            raise DoryK8sError(
+                "Kubernetes client not available. "
+                "Install with: pip install kubernetes"
+            )
+
+        try:
+            # Try in-cluster config first
+            config.load_incluster_config()
+            logger.debug("Using in-cluster Kubernetes config")
+        except config.ConfigException:
+            try:
+                # Fall back to kubeconfig
+                config.load_kube_config()
+                logger.debug("Using kubeconfig")
+            except config.ConfigException as e:
+                raise DoryK8sError(f"Failed to load Kubernetes config: {e}", cause=e)
+
+        self._core_api = client.CoreV1Api()
+
+        # Auto-detect namespace if not provided
+        if not self._namespace:
+            import os
+            self._namespace = os.environ.get("POD_NAMESPACE", "default")
+
+        self._initialized = True
+
+    @property
+    def namespace(self) -> str:
+        """Get current namespace."""
+        self._ensure_initialized()
+        return self._namespace
+
+    @property
+    def core_api(self):
+        """Get CoreV1Api client."""
+        self._ensure_initialized()
+        return self._core_api
+
+    async def get_pod(self, name: str) -> dict[str, Any]:
+        """
+        Get pod details.
+
+        Args:
+            name: Pod name
+
+        Returns:
+            Pod details as dictionary
+
+        Raises:
+            DoryK8sError: If operation fails
+        """
+        self._ensure_initialized()
+
+        try:
+            pod = self._core_api.read_namespaced_pod(
+                name=name,
+                namespace=self._namespace,
+            )
+            return pod.to_dict()
+
+        except ApiException as e:
+            if e.status == 404:
+                raise DoryK8sError(f"Pod not found: {name}", cause=e)
+            raise DoryK8sError(f"Failed to get pod {name}: {e}", cause=e)
+
+    async def get_pod_annotations(self, name: str) -> dict[str, str]:
+        """
+        Get pod annotations.
+
+        Args:
+            name: Pod name
+
+        Returns:
+            Annotations dictionary
+        """
+        pod = await self.get_pod(name)
+        return pod.get("metadata", {}).get("annotations", {})
+
+    async def patch_pod_annotations(
+        self,
+        name: str,
+        annotations: dict[str, str],
+    ) -> None:
+        """
+        Patch pod annotations.
+
+        Args:
+            name: Pod name
+            annotations: Annotations to add/update
+        """
+        self._ensure_initialized()
+
+        body = {
+            "metadata": {
+                "annotations": annotations,
+            }
+        }
+
+        try:
+            self._core_api.patch_namespaced_pod(
+                name=name,
+                namespace=self._namespace,
+                body=body,
+            )
+            logger.debug(f"Patched annotations on pod {name}")
+
+        except ApiException as e:
+            raise DoryK8sError(f"Failed to patch pod {name}: {e}", cause=e)
+
+    async def get_configmap(self, name: str) -> dict[str, str] | None:
+        """
+        Get ConfigMap data.
+
+        Args:
+            name: ConfigMap name
+
+        Returns:
+            ConfigMap data, or None if not found
+        """
+        self._ensure_initialized()
+
+        try:
+            cm = self._core_api.read_namespaced_config_map(
+                name=name,
+                namespace=self._namespace,
+            )
+            return cm.data or {}
+
+        except ApiException as e:
+            if e.status == 404:
+                return None
+            raise DoryK8sError(f"Failed to get ConfigMap {name}: {e}", cause=e)
+
+    async def create_or_update_configmap(
+        self,
+        name: str,
+        data: dict[str, str],
+        labels: dict[str, str] | None = None,
+    ) -> None:
+        """
+        Create or update a ConfigMap.
+
+        Args:
+            name: ConfigMap name
+            data: ConfigMap data
+            labels: Optional labels
+        """
+        self._ensure_initialized()
+
+        configmap = client.V1ConfigMap(
+            metadata=client.V1ObjectMeta(
+                name=name,
+                namespace=self._namespace,
+                labels=labels,
+            ),
+            data=data,
+        )
+
+        try:
+            # Try create first
+            self._core_api.create_namespaced_config_map(
+                namespace=self._namespace,
+                body=configmap,
+            )
+            logger.debug(f"Created ConfigMap {name}")
+
+        except ApiException as e:
+            if e.status == 409:
+                # Already exists, update
+                self._core_api.replace_namespaced_config_map(
+                    name=name,
+                    namespace=self._namespace,
+                    body=configmap,
+                )
+                logger.debug(f"Updated ConfigMap {name}")
+            else:
+                raise DoryK8sError(f"Failed to create ConfigMap {name}: {e}", cause=e)
+
+    async def delete_configmap(self, name: str) -> bool:
+        """
+        Delete a ConfigMap.
+
+        Args:
+            name: ConfigMap name
+
+        Returns:
+            True if deleted, False if not found
+        """
+        self._ensure_initialized()
+
+        try:
+            self._core_api.delete_namespaced_config_map(
+                name=name,
+                namespace=self._namespace,
+            )
+            return True
+
+        except ApiException as e:
+            if e.status == 404:
+                return False
+            raise DoryK8sError(f"Failed to delete ConfigMap {name}: {e}", cause=e)