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

dory/edge/detector.py

@@ -0,0 +1,546 @@
+"""Edge/Cloud workload detection.
+
+Provides multiple methods to detect whether the SDK is running on an
+edge node or a cloud node, with fallback strategies.
+
+Detection Methods (in priority order):
+    1. Environment variables (DORY_WORKLOAD_LOCATION, NODE_TYPE)
+    2. Pod labels via Kubernetes API
+    3. Node labels via Kubernetes API
+    4. Node name pattern matching (fallback)
+"""
+
+import logging
+import os
+import re
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+from dory.k8s.labels import (
+    LABEL_WORKLOAD_LOCATION,
+    LABEL_NODE_TYPE,
+    LABEL_MIGRATED_FROM_EDGE,
+    LABEL_ORIGINAL_NODE,
+    VALUE_EDGE,
+    VALUE_MANAGED,
+    WorkloadLocation,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class NodeType(Enum):
+    """Type of node the workload is running on."""
+
+    EDGE = "edge"
+    CLOUD = "cloud"
+    UNKNOWN = "unknown"
+
+
+@dataclass
+class WorkloadContext:
+    """Complete context about where the workload is running.
+
+    Attributes:
+        node_type: Whether running on edge or cloud node
+        workload_location: Intended workload location (edge/managed)
+        is_edge: True if running on edge node
+        is_migrated: True if this is a migrated edge workload on cloud
+        original_node: Original edge node name if migrated
+        node_name: Current node name
+        pod_name: Current pod name
+        detection_method: How the context was detected
+        confidence: Confidence level (high/medium/low)
+    """
+
+    node_type: NodeType
+    workload_location: WorkloadLocation | None
+    is_edge: bool
+    is_migrated: bool
+    original_node: str | None
+    node_name: str | None
+    pod_name: str | None
+    detection_method: str
+    confidence: str  # "high", "medium", "low"
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "node_type": self.node_type.value,
+            "workload_location": self.workload_location.value if self.workload_location else None,
+            "is_edge": self.is_edge,
+            "is_migrated": self.is_migrated,
+            "original_node": self.original_node,
+            "node_name": self.node_name,
+            "pod_name": self.pod_name,
+            "detection_method": self.detection_method,
+            "confidence": self.confidence,
+        }
+
+
+class WorkloadDetector:
+    """Detects whether the workload is running on edge or cloud.
+
+    Uses multiple detection strategies with fallbacks:
+
+    1. **Environment Variables** (highest confidence)
+       - DORY_WORKLOAD_LOCATION: "edge" or "managed"
+       - DORY_NODE_TYPE: "edge" or "cloud"
+       - NODE_TYPE: Node type from Downward API
+
+    2. **Pod Labels** (high confidence)
+       - workload-location label on the pod
+       - migrated-from-edge label for failover pods
+
+    3. **Node Labels** (high confidence)
+       - node-type label on the node (requires K8s API access)
+
+    4. **Node Name Pattern** (low confidence)
+       - Pattern matching on node names (e.g., "edge-*", "cloud-*")
+
+    Usage:
+        detector = WorkloadDetector()
+        context = detector.detect()
+
+        if context.is_edge:
+            # Running on edge node
+            enable_offline_mode()
+        else:
+            # Running on cloud node
+            enable_full_connectivity()
+
+        if context.is_migrated:
+            # This is a failover pod
+            await restore_state_from_edge(context.original_node)
+    """
+
+    # Default patterns for node name detection
+    EDGE_NODE_PATTERNS = [
+        r"^edge-",
+        r"-edge-",
+        r"-edge$",
+        r"^edge\d+",
+        r"^outpost-",
+        r"^local-",
+        r"^onprem-",
+    ]
+
+    CLOUD_NODE_PATTERNS = [
+        r"^ip-\d+",  # AWS node naming
+        r"^gke-",    # GKE node naming
+        r"^aks-",    # AKS node naming
+        r"^cloud-",
+        r"^managed-",
+    ]
+
+    def __init__(
+        self,
+        k8s_client: Any = None,
+        edge_patterns: list[str] | None = None,
+        cloud_patterns: list[str] | None = None,
+    ):
+        """Initialize detector.
+
+        Args:
+            k8s_client: Optional Kubernetes client for API-based detection
+            edge_patterns: Custom regex patterns for edge node names
+            cloud_patterns: Custom regex patterns for cloud node names
+        """
+        self._k8s_client = k8s_client
+        self._edge_patterns = edge_patterns or self.EDGE_NODE_PATTERNS
+        self._cloud_patterns = cloud_patterns or self.CLOUD_NODE_PATTERNS
+        self._cached_context: WorkloadContext | None = None
+
+    def detect(self, use_cache: bool = True) -> WorkloadContext:
+        """Detect workload context.
+
+        Tries multiple detection methods in order of reliability.
+
+        Args:
+            use_cache: If True, return cached result if available
+
+        Returns:
+            WorkloadContext with detection results
+        """
+        if use_cache and self._cached_context:
+            return self._cached_context
+
+        # Try detection methods in order
+        context = (
+            self._detect_from_env()
+            or self._detect_from_pod_labels()
+            or self._detect_from_node_labels()
+            or self._detect_from_node_name()
+            or self._default_context()
+        )
+
+        self._cached_context = context
+        logger.info(
+            f"Workload detection: node_type={context.node_type.value}, "
+            f"is_edge={context.is_edge}, method={context.detection_method}, "
+            f"confidence={context.confidence}"
+        )
+
+        return context
+
+    def _detect_from_env(self) -> WorkloadContext | None:
+        """Detect from environment variables."""
+        # Check DORY_WORKLOAD_LOCATION
+        workload_loc = os.environ.get("DORY_WORKLOAD_LOCATION", "").lower()
+        node_type_env = os.environ.get("DORY_NODE_TYPE", "").lower()
+        node_type_k8s = os.environ.get("NODE_TYPE", "").lower()  # From Downward API
+
+        node_name = os.environ.get("NODE_NAME")
+        pod_name = os.environ.get("POD_NAME")
+
+        # Check for migrated pod
+        is_migrated = os.environ.get("DORY_MIGRATED_FROM_EDGE", "").lower() == "true"
+        original_node = os.environ.get("DORY_ORIGINAL_NODE")
+
+        # Determine node type
+        node_type = NodeType.UNKNOWN
+        if node_type_env == "edge" or node_type_k8s == "edge":
+            node_type = NodeType.EDGE
+        elif node_type_env == "cloud" or node_type_k8s in ("cloud", "managed", "application"):
+            node_type = NodeType.CLOUD
+
+        # Determine workload location
+        workload_location = None
+        if workload_loc == "edge":
+            workload_location = WorkloadLocation.EDGE
+        elif workload_loc == "managed":
+            workload_location = WorkloadLocation.MANAGED
+
+        # If we have clear indicators, return result
+        if node_type != NodeType.UNKNOWN or workload_location is not None:
+            is_edge = node_type == NodeType.EDGE
+            if node_type == NodeType.UNKNOWN and workload_location:
+                # Infer node type from workload location (less reliable)
+                is_edge = workload_location == WorkloadLocation.EDGE and not is_migrated
+
+            return WorkloadContext(
+                node_type=node_type if node_type != NodeType.UNKNOWN else (
+                    NodeType.EDGE if is_edge else NodeType.CLOUD
+                ),
+                workload_location=workload_location,
+                is_edge=is_edge,
+                is_migrated=is_migrated,
+                original_node=original_node,
+                node_name=node_name,
+                pod_name=pod_name,
+                detection_method="environment",
+                confidence="high",
+            )
+
+        return None
+
+    def _detect_from_pod_labels(self) -> WorkloadContext | None:
+        """Detect from pod labels via Kubernetes API."""
+        if not self._k8s_client:
+            return None
+
+        try:
+            pod_name = os.environ.get("POD_NAME")
+            namespace = os.environ.get("POD_NAMESPACE", "default")
+
+            if not pod_name:
+                return None
+
+            # Get pod labels
+            pod = self._k8s_client.get_pod(namespace, pod_name)
+            if not pod or not pod.metadata or not pod.metadata.labels:
+                return None
+
+            labels = pod.metadata.labels
+            node_name = pod.spec.node_name if pod.spec else None
+
+            workload_loc = labels.get(LABEL_WORKLOAD_LOCATION, "")
+            is_migrated = labels.get(LABEL_MIGRATED_FROM_EDGE) == "true"
+            original_node = labels.get(LABEL_ORIGINAL_NODE)
+
+            workload_location = None
+            if workload_loc == VALUE_EDGE:
+                workload_location = WorkloadLocation.EDGE
+            elif workload_loc == VALUE_MANAGED:
+                workload_location = WorkloadLocation.MANAGED
+
+            if workload_location:
+                # For migrated pods, they're edge workloads but on cloud nodes
+                is_edge = workload_location == WorkloadLocation.EDGE and not is_migrated
+
+                return WorkloadContext(
+                    node_type=NodeType.EDGE if is_edge else NodeType.CLOUD,
+                    workload_location=workload_location,
+                    is_edge=is_edge,
+                    is_migrated=is_migrated,
+                    original_node=original_node,
+                    node_name=node_name,
+                    pod_name=pod_name,
+                    detection_method="pod_labels",
+                    confidence="high",
+                )
+
+        except Exception as e:
+            logger.debug(f"Pod label detection failed: {e}")
+
+        return None
+
+    def _detect_from_node_labels(self) -> WorkloadContext | None:
+        """Detect from node labels via Kubernetes API."""
+        if not self._k8s_client:
+            return None
+
+        try:
+            node_name = os.environ.get("NODE_NAME")
+            if not node_name:
+                return None
+
+            # Get node labels
+            node = self._k8s_client.get_node(node_name)
+            if not node or not node.metadata or not node.metadata.labels:
+                return None
+
+            labels = node.metadata.labels
+            node_type_label = labels.get(LABEL_NODE_TYPE, "")
+
+            if node_type_label == "edge":
+                return WorkloadContext(
+                    node_type=NodeType.EDGE,
+                    workload_location=WorkloadLocation.EDGE,
+                    is_edge=True,
+                    is_migrated=False,
+                    original_node=None,
+                    node_name=node_name,
+                    pod_name=os.environ.get("POD_NAME"),
+                    detection_method="node_labels",
+                    confidence="high",
+                )
+            elif node_type_label in ("cloud", "managed", "application"):
+                return WorkloadContext(
+                    node_type=NodeType.CLOUD,
+                    workload_location=WorkloadLocation.MANAGED,
+                    is_edge=False,
+                    is_migrated=False,
+                    original_node=None,
+                    node_name=node_name,
+                    pod_name=os.environ.get("POD_NAME"),
+                    detection_method="node_labels",
+                    confidence="high",
+                )
+
+        except Exception as e:
+            logger.debug(f"Node label detection failed: {e}")
+
+        return None
+
+    def _detect_from_node_name(self) -> WorkloadContext | None:
+        """Detect from node name patterns (fallback)."""
+        node_name = os.environ.get("NODE_NAME", "")
+        if not node_name:
+            return None
+
+        # Check edge patterns
+        for pattern in self._edge_patterns:
+            if re.search(pattern, node_name, re.IGNORECASE):
+                return WorkloadContext(
+                    node_type=NodeType.EDGE,
+                    workload_location=WorkloadLocation.EDGE,
+                    is_edge=True,
+                    is_migrated=False,
+                    original_node=None,
+                    node_name=node_name,
+                    pod_name=os.environ.get("POD_NAME"),
+                    detection_method="node_name_pattern",
+                    confidence="low",
+                )
+
+        # Check cloud patterns
+        for pattern in self._cloud_patterns:
+            if re.search(pattern, node_name, re.IGNORECASE):
+                return WorkloadContext(
+                    node_type=NodeType.CLOUD,
+                    workload_location=WorkloadLocation.MANAGED,
+                    is_edge=False,
+                    is_migrated=False,
+                    original_node=None,
+                    node_name=node_name,
+                    pod_name=os.environ.get("POD_NAME"),
+                    detection_method="node_name_pattern",
+                    confidence="low",
+                )
+
+        return None
+
+    def _default_context(self) -> WorkloadContext:
+        """Return default context when detection fails."""
+        return WorkloadContext(
+            node_type=NodeType.UNKNOWN,
+            workload_location=None,
+            is_edge=False,  # Default to cloud behavior (safer)
+            is_migrated=False,
+            original_node=None,
+            node_name=os.environ.get("NODE_NAME"),
+            pod_name=os.environ.get("POD_NAME"),
+            detection_method="default",
+            confidence="none",
+        )
+
+    def clear_cache(self) -> None:
+        """Clear cached detection result."""
+        self._cached_context = None
+
+
+# =============================================================================
+# Convenience Functions
+# =============================================================================
+
+# Global detector instance
+_default_detector: WorkloadDetector | None = None
+
+
+def get_detector() -> WorkloadDetector:
+    """Get or create default detector instance."""
+    global _default_detector
+    if _default_detector is None:
+        _default_detector = WorkloadDetector()
+    return _default_detector
+
+
+def is_edge_node() -> bool:
+    """Check if currently running on an edge node.
+
+    Returns:
+        True if running on edge node
+    """
+    return get_detector().detect().is_edge
+
+
+def is_cloud_node() -> bool:
+    """Check if currently running on a cloud node.
+
+    Returns:
+        True if running on cloud node
+    """
+    return not get_detector().detect().is_edge
+
+
+def is_migrated_workload() -> bool:
+    """Check if this is a migrated edge workload.
+
+    Returns:
+        True if this is a failover pod from edge
+    """
+    return get_detector().detect().is_migrated
+
+
+def get_workload_context() -> WorkloadContext:
+    """Get full workload context.
+
+    Returns:
+        WorkloadContext with all detection information
+    """
+    return get_detector().detect()
+
+
+def get_node_type() -> NodeType:
+    """Get current node type.
+
+    Returns:
+        NodeType enum (EDGE, CLOUD, or UNKNOWN)
+    """
+    return get_detector().detect().node_type
+
+
+# =============================================================================
+# Kubernetes Pod Spec Helper
+# =============================================================================
+
+def get_recommended_env_vars() -> list[dict[str, Any]]:
+    """Get recommended environment variables for pod spec.
+
+    Returns a list of environment variable definitions that should be
+    added to the pod spec to enable proper workload detection.
+
+    Returns:
+        List of env var definitions for Kubernetes pod spec
+    """
+    return [
+        # Standard Downward API fields
+        {
+            "name": "POD_NAME",
+            "valueFrom": {"fieldRef": {"fieldPath": "metadata.name"}},
+        },
+        {
+            "name": "POD_NAMESPACE",
+            "valueFrom": {"fieldRef": {"fieldPath": "metadata.namespace"}},
+        },
+        {
+            "name": "POD_IP",
+            "valueFrom": {"fieldRef": {"fieldPath": "status.podIP"}},
+        },
+        {
+            "name": "NODE_NAME",
+            "valueFrom": {"fieldRef": {"fieldPath": "spec.nodeName"}},
+        },
+        # Dory-specific labels exposed as env vars
+        {
+            "name": "DORY_APP_NAME",
+            "valueFrom": {"fieldRef": {"fieldPath": "metadata.labels['app']"}},
+        },
+        {
+            "name": "DORY_WORKLOAD_LOCATION",
+            "valueFrom": {"fieldRef": {"fieldPath": "metadata.labels['workload-location']"}},
+        },
+        # Node type (if available)
+        {
+            "name": "DORY_NODE_TYPE",
+            "value": "$(NODE_TYPE)",  # Set by node or externally
+        },
+    ]
+
+
+def get_pod_spec_yaml_snippet() -> str:
+    """Get YAML snippet for pod spec environment variables.
+
+    Returns:
+        YAML string to add to pod spec
+    """
+    return """
+env:
+  # Standard Kubernetes Downward API
+  - name: POD_NAME
+    valueFrom:
+      fieldRef:
+        fieldPath: metadata.name
+  - name: POD_NAMESPACE
+    valueFrom:
+      fieldRef:
+        fieldPath: metadata.namespace
+  - name: POD_IP
+    valueFrom:
+      fieldRef:
+        fieldPath: status.podIP
+  - name: NODE_NAME
+    valueFrom:
+      fieldRef:
+        fieldPath: spec.nodeName
+
+  # Dory workload detection
+  - name: DORY_APP_NAME
+    valueFrom:
+      fieldRef:
+        fieldPath: metadata.labels['app']
+  - name: DORY_WORKLOAD_LOCATION
+    valueFrom:
+      fieldRef:
+        fieldPath: metadata.labels['workload-location']
+  - name: DORY_MIGRATED_FROM_EDGE
+    valueFrom:
+      fieldRef:
+        fieldPath: metadata.labels['migrated-from-edge']
+  - name: DORY_ORIGINAL_NODE
+    valueFrom:
+      fieldRef:
+        fieldPath: metadata.labels['original-edge-node']
+"""