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

dory/k8s/labels.py

@@ -0,0 +1,505 @@
+"""Kubernetes label contract between SDK and Orchestrator.
+
+This module defines the label contract that enables coordination between
+the Dory SDK and the Dory Orchestrator. Both components must use identical
+label keys and values for proper operation.
+
+Label Categories:
+    1. Identity Labels - Identify the processor and its managing controller
+    2. Workload Labels - Control scheduling and workload placement
+    3. Migration Labels - Track failover and migration state
+
+Contract Version: 1.0
+Orchestrator Compatibility: v1.0+
+"""
+
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+# =============================================================================
+# Label Keys (must match orchestrator/pkg/utils/k8s.go)
+# =============================================================================
+
+# Identity Labels
+LABEL_MANAGED_BY = "managed-by"
+LABEL_APP_NAME = "app"
+LABEL_PROCESSOR_ID = "processor-id"
+
+# Workload Labels
+LABEL_WORKLOAD_TYPE = "workload-type"
+LABEL_WORKLOAD_LOCATION = "workload-location"
+LABEL_NODE_TYPE = "node-type"
+LABEL_INSTANCE_FAMILY = "instance-family"
+
+# Migration Labels
+LABEL_MIGRATED_FROM_EDGE = "migrated-from-edge"
+LABEL_ORIGINAL_WORKLOAD_TYPE = "original-workload-type"
+LABEL_ORIGINAL_NODE = "original-edge-node"
+
+
+# =============================================================================
+# Label Values (must match orchestrator/pkg/utils/k8s.go)
+# =============================================================================
+
+VALUE_ORCHESTRATOR_NAME = "dory-orchestrator"
+VALUE_EDGE = "edge"
+VALUE_MANAGED = "managed"
+VALUE_APPLICATION = "application"
+
+
+# =============================================================================
+# Label Selectors (must match orchestrator/pkg/utils/k8s.go)
+# =============================================================================
+
+SELECTOR_EDGE_NODES = "node-type=edge"
+SELECTOR_MANAGED_NODES = "workload-type=application"
+SELECTOR_MANAGED_PODS = "managed-by=dory-orchestrator"
+
+
+# =============================================================================
+# Enums for Type Safety
+# =============================================================================
+
+class WorkloadLocation(Enum):
+    """Workload location for pods.
+
+    Determines where the pod should run:
+    - EDGE: Runs on externally-managed edge nodes (node-type=edge)
+    - MANAGED: Runs on Karpenter-managed cloud nodes (workload-type=application)
+    """
+    EDGE = VALUE_EDGE
+    MANAGED = VALUE_MANAGED
+
+
+class NodeType(Enum):
+    """Node type classification.
+
+    Used for node selection:
+    - EDGE: Edge nodes with intermittent connectivity
+    - SYSTEM: System nodes (control plane, etc.)
+    - APPLICATION: Karpenter-managed application nodes
+    """
+    EDGE = "edge"
+    SYSTEM = "system"
+    APPLICATION = VALUE_APPLICATION
+
+
+# =============================================================================
+# Label Builder
+# =============================================================================
+
+@dataclass
+class DoryLabels:
+    """Builder for Dory-compatible Kubernetes labels.
+
+    Creates labels that are compatible with the Dory Orchestrator's
+    label contract. Use this to ensure your pods are properly managed.
+
+    Usage:
+        # Create labels for an edge processor
+        labels = DoryLabels(
+            app_name="my-processor",
+            processor_id="proc-123",
+            workload_location=WorkloadLocation.EDGE,
+        )
+
+        # Get as dictionary for Kubernetes manifest
+        pod_labels = labels.to_dict()
+
+        # Create labels for a migrated pod
+        migrated_labels = DoryLabels(
+            app_name="my-processor",
+            processor_id="proc-123",
+            workload_location=WorkloadLocation.EDGE,
+            migrated_from_edge=True,
+            original_node="edge-node-1",
+        )
+    """
+
+    app_name: str
+    processor_id: str | None = None
+    workload_location: WorkloadLocation = WorkloadLocation.MANAGED
+    migrated_from_edge: bool = False
+    original_workload_type: str | None = None
+    original_node: str | None = None
+    custom_labels: dict[str, str] | None = None
+
+    def to_dict(self) -> dict[str, str]:
+        """Convert to Kubernetes labels dictionary.
+
+        Returns:
+            Dictionary of label key-value pairs
+        """
+        labels = {
+            LABEL_MANAGED_BY: VALUE_ORCHESTRATOR_NAME,
+            LABEL_APP_NAME: self.app_name,
+            LABEL_WORKLOAD_LOCATION: self.workload_location.value,
+        }
+
+        if self.processor_id:
+            labels[LABEL_PROCESSOR_ID] = self.processor_id
+
+        if self.migrated_from_edge:
+            labels[LABEL_MIGRATED_FROM_EDGE] = "true"
+
+        if self.original_workload_type:
+            labels[LABEL_ORIGINAL_WORKLOAD_TYPE] = self.original_workload_type
+
+        if self.original_node:
+            labels[LABEL_ORIGINAL_NODE] = self.original_node
+
+        if self.custom_labels:
+            labels.update(self.custom_labels)
+
+        return labels
+
+    @classmethod
+    def from_dict(cls, labels: dict[str, str]) -> "DoryLabels":
+        """Create from existing labels dictionary.
+
+        Args:
+            labels: Kubernetes labels dictionary
+
+        Returns:
+            DoryLabels instance
+        """
+        workload_location = WorkloadLocation.MANAGED
+        if labels.get(LABEL_WORKLOAD_LOCATION) == VALUE_EDGE:
+            workload_location = WorkloadLocation.EDGE
+
+        return cls(
+            app_name=labels.get(LABEL_APP_NAME, ""),
+            processor_id=labels.get(LABEL_PROCESSOR_ID),
+            workload_location=workload_location,
+            migrated_from_edge=labels.get(LABEL_MIGRATED_FROM_EDGE) == "true",
+            original_workload_type=labels.get(LABEL_ORIGINAL_WORKLOAD_TYPE),
+            original_node=labels.get(LABEL_ORIGINAL_NODE),
+        )
+
+    def is_edge_workload(self) -> bool:
+        """Check if this is an edge workload."""
+        return self.workload_location == WorkloadLocation.EDGE
+
+    def is_migrated(self) -> bool:
+        """Check if this pod was migrated from edge."""
+        return self.migrated_from_edge
+
+
+# =============================================================================
+# Label Reader Utilities
+# =============================================================================
+
+def get_label(labels: dict[str, str] | None, key: str, default: str = "") -> str:
+    """Safely get a label value.
+
+    Args:
+        labels: Labels dictionary (may be None)
+        key: Label key
+        default: Default value if not found
+
+    Returns:
+        Label value or default
+    """
+    if labels is None:
+        return default
+    return labels.get(key, default)
+
+
+def get_app_name(labels: dict[str, str] | None) -> str:
+    """Get app name from labels."""
+    return get_label(labels, LABEL_APP_NAME)
+
+
+def get_processor_id(labels: dict[str, str] | None) -> str:
+    """Get processor ID from labels."""
+    return get_label(labels, LABEL_PROCESSOR_ID)
+
+
+def get_workload_location(labels: dict[str, str] | None) -> WorkloadLocation | None:
+    """Get workload location from labels.
+
+    Returns:
+        WorkloadLocation enum or None if not set
+    """
+    value = get_label(labels, LABEL_WORKLOAD_LOCATION)
+    if value == VALUE_EDGE:
+        return WorkloadLocation.EDGE
+    elif value == VALUE_MANAGED:
+        return WorkloadLocation.MANAGED
+    return None
+
+
+def is_managed_by_dory(labels: dict[str, str] | None) -> bool:
+    """Check if pod is managed by Dory Orchestrator."""
+    return get_label(labels, LABEL_MANAGED_BY) == VALUE_ORCHESTRATOR_NAME
+
+
+def is_edge_workload(labels: dict[str, str] | None) -> bool:
+    """Check if pod is an edge workload."""
+    return get_label(labels, LABEL_WORKLOAD_LOCATION) == VALUE_EDGE
+
+
+def is_migrated_from_edge(labels: dict[str, str] | None) -> bool:
+    """Check if pod was migrated from edge."""
+    return get_label(labels, LABEL_MIGRATED_FROM_EDGE) == "true"
+
+
+def get_original_node(labels: dict[str, str] | None) -> str:
+    """Get original edge node for migrated pod."""
+    return get_label(labels, LABEL_ORIGINAL_NODE)
+
+
+# =============================================================================
+# Environment-Based Label Detection
+# =============================================================================
+
+def get_pod_labels_from_env() -> dict[str, str]:
+    """Get pod labels from environment variables.
+
+    The Kubernetes Downward API can expose labels as environment variables.
+    This function reads common label values from the environment.
+
+    Expected environment variables:
+        - POD_NAME: Pod name
+        - POD_NAMESPACE: Pod namespace
+        - POD_IP: Pod IP address
+        - NODE_NAME: Node name
+        - DORY_APP_NAME: App name (from label)
+        - DORY_PROCESSOR_ID: Processor ID (from label)
+        - DORY_WORKLOAD_LOCATION: Workload location (from label)
+
+    Returns:
+        Dictionary of detected labels
+    """
+    labels = {}
+
+    if app_name := os.environ.get("DORY_APP_NAME"):
+        labels[LABEL_APP_NAME] = app_name
+
+    if processor_id := os.environ.get("DORY_PROCESSOR_ID"):
+        labels[LABEL_PROCESSOR_ID] = processor_id
+
+    if workload_location := os.environ.get("DORY_WORKLOAD_LOCATION"):
+        labels[LABEL_WORKLOAD_LOCATION] = workload_location
+
+    # Managed by is always the orchestrator if we're in Dory context
+    if labels:
+        labels[LABEL_MANAGED_BY] = VALUE_ORCHESTRATOR_NAME
+
+    return labels
+
+
+def detect_workload_context() -> dict[str, Any]:
+    """Detect workload context from environment.
+
+    Reads environment variables set by Kubernetes Downward API and
+    Dory-specific variables to determine the workload context.
+
+    Returns:
+        Dictionary with context information:
+        - pod_name: str | None
+        - pod_namespace: str | None
+        - pod_ip: str | None
+        - node_name: str | None
+        - app_name: str | None
+        - processor_id: str | None
+        - workload_location: WorkloadLocation | None
+        - is_edge: bool
+        - is_kubernetes: bool
+    """
+    pod_name = os.environ.get("POD_NAME")
+    pod_namespace = os.environ.get("POD_NAMESPACE")
+    pod_ip = os.environ.get("POD_IP")
+    node_name = os.environ.get("NODE_NAME")
+    app_name = os.environ.get("DORY_APP_NAME")
+    processor_id = os.environ.get("DORY_PROCESSOR_ID")
+    workload_location_str = os.environ.get("DORY_WORKLOAD_LOCATION", "")
+
+    workload_location = None
+    if workload_location_str == VALUE_EDGE:
+        workload_location = WorkloadLocation.EDGE
+    elif workload_location_str == VALUE_MANAGED:
+        workload_location = WorkloadLocation.MANAGED
+
+    is_kubernetes = pod_name is not None and pod_namespace is not None
+    is_edge = workload_location == WorkloadLocation.EDGE
+
+    return {
+        "pod_name": pod_name,
+        "pod_namespace": pod_namespace,
+        "pod_ip": pod_ip,
+        "node_name": node_name,
+        "app_name": app_name,
+        "processor_id": processor_id,
+        "workload_location": workload_location,
+        "is_edge": is_edge,
+        "is_kubernetes": is_kubernetes,
+    }
+
+
+# =============================================================================
+# Node Selector Builders
+# =============================================================================
+
+def edge_node_selector() -> dict[str, str]:
+    """Get node selector for edge nodes.
+
+    Returns:
+        Node selector dictionary for edge node placement
+    """
+    return {LABEL_NODE_TYPE: NodeType.EDGE.value}
+
+
+def managed_node_selector() -> dict[str, str]:
+    """Get node selector for managed (Karpenter) nodes.
+
+    Returns:
+        Node selector dictionary for managed node placement
+    """
+    return {LABEL_WORKLOAD_TYPE: NodeType.APPLICATION.value}
+
+
+def edge_toleration() -> dict[str, Any]:
+    """Get toleration for edge nodes.
+
+    Edge nodes typically have a taint to prevent non-edge workloads
+    from scheduling on them.
+
+    Returns:
+        Toleration dictionary for edge nodes
+    """
+    return {
+        "key": "edge-node",
+        "operator": "Equal",
+        "value": "true",
+        "effect": "NoSchedule",
+    }
+
+
+# =============================================================================
+# Contract Documentation
+# =============================================================================
+
+CONTRACT_VERSION = "1.0"
+
+LABEL_CONTRACT = """
+Dory SDK/Orchestrator Label Contract v1.0
+=========================================
+
+This document defines the label contract between the Dory SDK and
+the Dory Orchestrator. Both components MUST use these exact labels
+for proper coordination.
+
+IDENTITY LABELS
+---------------
+These labels identify the processor and its managing controller.
+
+| Label Key      | Description                    | Required | Example Value      |
+|----------------|--------------------------------|----------|--------------------|
+| managed-by     | Managing controller            | Yes      | dory-orchestrator  |
+| app            | Application name               | Yes      | my-processor       |
+| processor-id   | Unique processor identifier    | No       | proc-123           |
+
+WORKLOAD LABELS
+---------------
+These labels control scheduling and workload placement.
+
+| Label Key          | Description                     | Values        | Default  |
+|--------------------|---------------------------------|---------------|----------|
+| workload-location  | Where the pod runs              | edge, managed | managed  |
+| workload-type      | Node workload type (Karpenter)  | application   | -        |
+| node-type          | Node classification             | edge, system  | -        |
+
+MIGRATION LABELS
+----------------
+These labels track failover and migration state.
+
+| Label Key              | Description                    | Values     | Set By        |
+|------------------------|--------------------------------|------------|---------------|
+| migrated-from-edge     | Pod was migrated from edge     | true       | Orchestrator  |
+| original-workload-type | Original workload type         | edge       | Orchestrator  |
+| original-edge-node     | Original edge node name        | <node>     | Orchestrator  |
+
+NODE SELECTORS
+--------------
+Edge workloads:
+    node-type: edge
+
+Managed workloads (Karpenter):
+    workload-type: application
+
+TOLERATIONS
+-----------
+Edge nodes have a taint that requires this toleration:
+    key: edge-node
+    operator: Equal
+    value: true
+    effect: NoSchedule
+
+ENVIRONMENT VARIABLES
+---------------------
+The SDK expects these environment variables (set via Downward API):
+
+| Variable               | Source                        | Description           |
+|------------------------|-------------------------------|-----------------------|
+| POD_NAME               | metadata.name                 | Pod name              |
+| POD_NAMESPACE          | metadata.namespace            | Pod namespace         |
+| POD_IP                 | status.podIP                  | Pod IP address        |
+| NODE_NAME              | spec.nodeName                 | Node name             |
+| DORY_APP_NAME          | metadata.labels['app']        | App label value       |
+| DORY_PROCESSOR_ID      | metadata.labels['processor-id']| Processor ID label   |
+| DORY_WORKLOAD_LOCATION | metadata.labels['workload-location'] | Location label |
+
+FAILOVER FLOW
+-------------
+1. Edge pod running with:
+   - workload-location: edge
+   - node-type selector: edge
+
+2. Edge node fails, Orchestrator:
+   - Deletes edge pod
+   - Creates new pod with:
+     - workload-location: edge (preserved)
+     - migrated-from-edge: true (added)
+     - original-edge-node: <old-node> (added)
+     - workload-type selector: application (changed)
+
+3. Edge node recovers, Orchestrator:
+   - Deletes migrated pod
+   - Creates new pod on edge with original labels
+
+USAGE EXAMPLE
+-------------
+```python
+from dory.k8s.labels import DoryLabels, WorkloadLocation
+
+# Create labels for edge processor
+labels = DoryLabels(
+    app_name="my-processor",
+    processor_id="proc-123",
+    workload_location=WorkloadLocation.EDGE,
+)
+
+# Use in Kubernetes manifest
+pod_spec = {
+    "metadata": {
+        "labels": labels.to_dict(),
+    },
+    "spec": {
+        "nodeSelector": edge_node_selector(),
+        "tolerations": [edge_toleration()],
+    },
+}
+```
+"""
+
+
+def get_contract_documentation() -> str:
+    """Get the label contract documentation.
+
+    Returns:
+        Complete label contract documentation string
+    """
+    return LABEL_CONTRACT

dory/migration/__init__.py

@@ -3,9 +3,58 @@
 from dory.migration.state_manager import StateManager
 from dory.migration.serialization import StateSerializer
 from dory.migration.configmap import ConfigMapStore
+from dory.migration.s3_store import S3Store, S3Config, OfflineBuffer
+from dory.migration.transfer import (
+    TransferConfig,
+    TransferMetrics,
+    StateTransferError,
+    StateTransferTimeout,
+    StateSizeExceeded,
+    validate_state_size,
+    ORCHESTRATOR_STATE_TIMEOUT_SEC,
+    ORCHESTRATOR_MAX_STATE_SIZE,
+)
+from dory.migration.versioning import (
+    StateFormatVersion,
+    VersionedState,
+    StateMetadata,
+    VersionNegotiator,
+    VersionNegotiationResult,
+    get_sdk_version,
+    get_supported_versions,
+    is_version_supported,
+    serialize_state,
+    deserialize_state,
+    get_version_info,
+    SDK_STATE_VERSION,
+)
 
 __all__ = [
     "StateManager",
     "StateSerializer",
     "ConfigMapStore",
+    "S3Store",
+    "S3Config",
+    "OfflineBuffer",
+    "TransferConfig",
+    "TransferMetrics",
+    "StateTransferError",
+    "StateTransferTimeout",
+    "StateSizeExceeded",
+    "validate_state_size",
+    "ORCHESTRATOR_STATE_TIMEOUT_SEC",
+    "ORCHESTRATOR_MAX_STATE_SIZE",
+    # Versioning
+    "StateFormatVersion",
+    "VersionedState",
+    "StateMetadata",
+    "VersionNegotiator",
+    "VersionNegotiationResult",
+    "get_sdk_version",
+    "get_supported_versions",
+    "is_version_supported",
+    "serialize_state",
+    "deserialize_state",
+    "get_version_info",
+    "SDK_STATE_VERSION",
 ]