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

dory/migration/state_manager.py

@@ -14,6 +14,7 @@ from typing import Any, TYPE_CHECKING
 from dory.types import StateBackend
 from dory.migration.serialization import StateSerializer
 from dory.migration.configmap import ConfigMapStore
+from dory.migration.s3_store import S3Store, S3Config
 from dory.utils.errors import DoryStateError
 
 if TYPE_CHECKING:
@@ -55,6 +56,7 @@ class StateManager:
         self._config = config
         self._serializer = StateSerializer()
         self._configmap_store: ConfigMapStore | None = None
+        self._s3_store: S3Store | None = None
 
         # Get namespace from environment
         self._namespace = os.environ.get("POD_NAMESPACE", "default")
@@ -285,17 +287,73 @@ class StateManager:
             raise DoryStateError(f"Failed to delete state from PVC {path}: {e}", cause=e)
 
     # =========================================================================
-    # S3 Backend (placeholder - would need boto3)
+    # S3 Backend
     # =========================================================================
 
+    def _get_s3_store(self) -> S3Store:
+        """Get or create S3 store instance."""
+        if self._s3_store is None:
+            # Try to get S3 config from DoryConfig if available
+            s3_config = None
+            if self._config and hasattr(self._config, "s3_config"):
+                s3_config = self._config.s3_config
+
+            self._s3_store = S3Store(config=s3_config)
+
+        return self._s3_store
+
     async def _save_to_s3(self, processor_id: str, state_json: str) -> None:
-        """Save state to S3."""
-        raise DoryStateError("S3 backend not yet implemented")
+        """Save state to S3 with offline buffering support."""
+        store = self._get_s3_store()
+        await store.save(
+            processor_id,
+            state_json,
+            metadata={
+                "pod-name": self._pod_name,
+                "namespace": self._namespace,
+            },
+        )
 
     async def _load_from_s3(self, processor_id: str) -> str | None:
-        """Load state from S3."""
-        raise DoryStateError("S3 backend not yet implemented")
+        """Load state from S3 (falls back to local buffer if unavailable)."""
+        store = self._get_s3_store()
+        return await store.load(processor_id)
 
     async def _delete_from_s3(self, processor_id: str) -> bool:
         """Delete state from S3."""
-        raise DoryStateError("S3 backend not yet implemented")
+        store = self._get_s3_store()
+        return await store.delete(processor_id)
+
+    async def sync_s3_buffer(self) -> int:
+        """
+        Sync locally buffered states to S3.
+
+        Call this periodically on edge nodes to upload states
+        that were buffered during connectivity issues.
+
+        Returns:
+            Number of states synced
+        """
+        if self._backend != StateBackend.S3:
+            return 0
+
+        store = self._get_s3_store()
+        return await store.sync_buffer()
+
+    async def start_s3_background_sync(self, interval_seconds: float = 60.0) -> None:
+        """
+        Start background S3 sync for edge nodes.
+
+        Args:
+            interval_seconds: Interval between sync attempts
+        """
+        if self._backend != StateBackend.S3:
+            return
+
+        store = self._get_s3_store()
+        await store.start_background_sync(interval_seconds)
+
+    async def stop_s3_background_sync(self) -> None:
+        """Stop background S3 sync."""
+        if self._s3_store:
+            await self._s3_store.stop_background_sync()

dory/migration/transfer.py

@@ -0,0 +1,382 @@
+"""
+State transfer utilities with timeout and size validation.
+
+Provides utilities for safe state capture and restore operations
+that align with Orchestrator timeout expectations.
+"""
+
+import asyncio
+import functools
+import logging
+import time
+from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
+from dataclasses import dataclass
+from typing import Any, Callable, TypeVar
+
+from dory.utils.errors import DoryStateError
+
+logger = logging.getLogger(__name__)
+
+# Type variable for generic return type
+T = TypeVar("T")
+
+# Orchestrator constants (from transfer.go)
+ORCHESTRATOR_STATE_TIMEOUT_SEC = 30  # DefaultHTTPTimeout in transfer.go
+ORCHESTRATOR_MAX_STATE_SIZE = 10 * 1024 * 1024  # MaxResponseBodySize in transfer.go (10MB)
+
+# Default SDK limits (with safety margin)
+DEFAULT_CAPTURE_TIMEOUT_SEC = 25  # 5s buffer before Orchestrator timeout
+DEFAULT_RESTORE_TIMEOUT_SEC = 25
+DEFAULT_MAX_STATE_SIZE = 8 * 1024 * 1024  # 8MB, 2MB buffer before Orchestrator limit
+DEFAULT_SIZE_WARN_THRESHOLD = 0.75  # Warn at 75% of max
+
+
+@dataclass
+class TransferConfig:
+    """Configuration for state transfer operations."""
+
+    capture_timeout_sec: float = DEFAULT_CAPTURE_TIMEOUT_SEC
+    restore_timeout_sec: float = DEFAULT_RESTORE_TIMEOUT_SEC
+    max_size_bytes: int = DEFAULT_MAX_STATE_SIZE
+    size_warn_threshold: float = DEFAULT_SIZE_WARN_THRESHOLD
+
+    def __post_init__(self):
+        """Validate configuration against Orchestrator limits."""
+        if self.capture_timeout_sec >= ORCHESTRATOR_STATE_TIMEOUT_SEC:
+            logger.warning(
+                f"state_capture_timeout_sec ({self.capture_timeout_sec}s) >= "
+                f"Orchestrator timeout ({ORCHESTRATOR_STATE_TIMEOUT_SEC}s). "
+                "Reducing to {ORCHESTRATOR_STATE_TIMEOUT_SEC - 5}s."
+            )
+            self.capture_timeout_sec = ORCHESTRATOR_STATE_TIMEOUT_SEC - 5
+
+        if self.restore_timeout_sec >= ORCHESTRATOR_STATE_TIMEOUT_SEC:
+            logger.warning(
+                f"state_restore_timeout_sec ({self.restore_timeout_sec}s) >= "
+                f"Orchestrator timeout ({ORCHESTRATOR_STATE_TIMEOUT_SEC}s). "
+                f"Reducing to {ORCHESTRATOR_STATE_TIMEOUT_SEC - 5}s."
+            )
+            self.restore_timeout_sec = ORCHESTRATOR_STATE_TIMEOUT_SEC - 5
+
+        if self.max_size_bytes > ORCHESTRATOR_MAX_STATE_SIZE:
+            logger.warning(
+                f"state_max_size_bytes ({self.max_size_bytes}) > "
+                f"Orchestrator limit ({ORCHESTRATOR_MAX_STATE_SIZE}). "
+                f"Reducing to {ORCHESTRATOR_MAX_STATE_SIZE}."
+            )
+            self.max_size_bytes = ORCHESTRATOR_MAX_STATE_SIZE
+
+
+@dataclass
+class TransferMetrics:
+    """Metrics from a state transfer operation."""
+
+    duration_sec: float
+    size_bytes: int
+    size_ratio: float  # size / max_size
+    timed_out: bool
+    size_exceeded: bool
+
+
+class StateTransferError(DoryStateError):
+    """Error during state transfer operation."""
+
+    def __init__(
+        self,
+        message: str,
+        metrics: TransferMetrics | None = None,
+        cause: Exception | None = None,
+    ):
+        super().__init__(message, cause=cause)
+        self.metrics = metrics
+
+
+class StateTransferTimeout(StateTransferError):
+    """State transfer operation timed out."""
+    pass
+
+
+class StateSizeExceeded(StateTransferError):
+    """State size exceeds configured maximum."""
+    pass
+
+
+def validate_state_size(
+    state_json: str,
+    max_size: int = DEFAULT_MAX_STATE_SIZE,
+    warn_threshold: float = DEFAULT_SIZE_WARN_THRESHOLD,
+) -> TransferMetrics:
+    """
+    Validate state size against limits.
+
+    Args:
+        state_json: Serialized state JSON string
+        max_size: Maximum allowed size in bytes
+        warn_threshold: Fraction of max_size to trigger warning
+
+    Returns:
+        TransferMetrics with size information
+
+    Raises:
+        StateSizeExceeded: If state exceeds max_size
+    """
+    size_bytes = len(state_json.encode("utf-8"))
+    size_ratio = size_bytes / max_size if max_size > 0 else 0
+
+    metrics = TransferMetrics(
+        duration_sec=0,
+        size_bytes=size_bytes,
+        size_ratio=size_ratio,
+        timed_out=False,
+        size_exceeded=size_bytes > max_size,
+    )
+
+    if size_bytes > max_size:
+        raise StateSizeExceeded(
+            f"State size ({size_bytes:,} bytes) exceeds maximum "
+            f"({max_size:,} bytes). Orchestrator will reject this state.",
+            metrics=metrics,
+        )
+
+    if size_ratio >= warn_threshold:
+        logger.warning(
+            f"State size ({size_bytes:,} bytes) is {size_ratio:.1%} of maximum "
+            f"({max_size:,} bytes). Consider reducing state size to avoid "
+            "transfer failures."
+        )
+
+    return metrics
+
+
+def with_timeout(
+    timeout_sec: float,
+    operation_name: str = "operation",
+) -> Callable[[Callable[..., T]], Callable[..., T]]:
+    """
+    Decorator to add timeout to synchronous functions.
+
+    Runs the function in a thread pool executor with timeout.
+
+    Args:
+        timeout_sec: Timeout in seconds
+        operation_name: Name for error messages
+
+    Returns:
+        Decorated function with timeout
+    """
+    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> T:
+            start_time = time.monotonic()
+
+            with ThreadPoolExecutor(max_workers=1) as executor:
+                future = executor.submit(func, *args, **kwargs)
+
+                try:
+                    result = future.result(timeout=timeout_sec)
+                    duration = time.monotonic() - start_time
+
+                    # Log if operation took significant time
+                    if duration > timeout_sec * 0.5:
+                        logger.warning(
+                            f"{operation_name} took {duration:.2f}s "
+                            f"({duration/timeout_sec:.1%} of {timeout_sec}s timeout)"
+                        )
+
+                    return result
+
+                except FuturesTimeoutError:
+                    duration = time.monotonic() - start_time
+                    metrics = TransferMetrics(
+                        duration_sec=duration,
+                        size_bytes=0,
+                        size_ratio=0,
+                        timed_out=True,
+                        size_exceeded=False,
+                    )
+                    raise StateTransferTimeout(
+                        f"{operation_name} timed out after {timeout_sec}s. "
+                        "Consider reducing state size or optimizing get_state().",
+                        metrics=metrics,
+                    )
+
+        return wrapper
+    return decorator
+
+
+async def async_with_timeout(
+    coro: Any,
+    timeout_sec: float,
+    operation_name: str = "operation",
+) -> Any:
+    """
+    Execute coroutine with timeout.
+
+    Args:
+        coro: Coroutine to execute
+        timeout_sec: Timeout in seconds
+        operation_name: Name for error messages
+
+    Returns:
+        Result of the coroutine
+
+    Raises:
+        StateTransferTimeout: If operation times out
+    """
+    start_time = time.monotonic()
+
+    try:
+        result = await asyncio.wait_for(coro, timeout=timeout_sec)
+        duration = time.monotonic() - start_time
+
+        # Log if operation took significant time
+        if duration > timeout_sec * 0.5:
+            logger.warning(
+                f"{operation_name} took {duration:.2f}s "
+                f"({duration/timeout_sec:.1%} of {timeout_sec}s timeout)"
+            )
+
+        return result
+
+    except asyncio.TimeoutError:
+        duration = time.monotonic() - start_time
+        metrics = TransferMetrics(
+            duration_sec=duration,
+            size_bytes=0,
+            size_ratio=0,
+            timed_out=True,
+            size_exceeded=False,
+        )
+        raise StateTransferTimeout(
+            f"{operation_name} timed out after {timeout_sec}s. "
+            "Consider reducing state size or optimizing the operation.",
+            metrics=metrics,
+        )
+
+
+class StateCaptureGuard:
+    """
+    Context manager for safe state capture with timeout and size validation.
+
+    Usage:
+        config = TransferConfig(capture_timeout_sec=25, max_size_bytes=8*1024*1024)
+
+        with StateCaptureGuard(config) as guard:
+            state = processor.get_state()
+            state_json = json.dumps(state)
+            guard.validate(state_json)
+    """
+
+    def __init__(self, config: TransferConfig | None = None):
+        """
+        Initialize capture guard.
+
+        Args:
+            config: Transfer configuration
+        """
+        self._config = config or TransferConfig()
+        self._start_time: float = 0
+        self._metrics: TransferMetrics | None = None
+
+    def __enter__(self) -> "StateCaptureGuard":
+        self._start_time = time.monotonic()
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
+        duration = time.monotonic() - self._start_time
+
+        if duration > self._config.capture_timeout_sec:
+            logger.error(
+                f"State capture took {duration:.2f}s, exceeding "
+                f"{self._config.capture_timeout_sec}s timeout"
+            )
+
+        return False  # Don't suppress exceptions
+
+    def validate(self, state_json: str) -> TransferMetrics:
+        """
+        Validate captured state.
+
+        Args:
+            state_json: Serialized state JSON
+
+        Returns:
+            TransferMetrics with capture information
+
+        Raises:
+            StateSizeExceeded: If state exceeds max size
+            StateTransferTimeout: If capture exceeded timeout
+        """
+        duration = time.monotonic() - self._start_time
+
+        # Check timeout
+        if duration > self._config.capture_timeout_sec:
+            self._metrics = TransferMetrics(
+                duration_sec=duration,
+                size_bytes=len(state_json.encode("utf-8")),
+                size_ratio=0,
+                timed_out=True,
+                size_exceeded=False,
+            )
+            raise StateTransferTimeout(
+                f"State capture took {duration:.2f}s, exceeding "
+                f"{self._config.capture_timeout_sec}s timeout",
+                metrics=self._metrics,
+            )
+
+        # Check size
+        size_metrics = validate_state_size(
+            state_json,
+            max_size=self._config.max_size_bytes,
+            warn_threshold=self._config.size_warn_threshold,
+        )
+
+        self._metrics = TransferMetrics(
+            duration_sec=duration,
+            size_bytes=size_metrics.size_bytes,
+            size_ratio=size_metrics.size_ratio,
+            timed_out=False,
+            size_exceeded=size_metrics.size_exceeded,
+        )
+
+        return self._metrics
+
+    @property
+    def metrics(self) -> TransferMetrics | None:
+        """Get capture metrics."""
+        return self._metrics
+
+
+def log_transfer_summary(
+    operation: str,
+    metrics: TransferMetrics,
+    config: TransferConfig,
+) -> None:
+    """
+    Log a summary of the transfer operation.
+
+    Args:
+        operation: Operation name (e.g., "capture", "restore")
+        metrics: Transfer metrics
+        config: Transfer configuration
+    """
+    level = logging.INFO
+    status = "completed"
+
+    if metrics.timed_out:
+        level = logging.ERROR
+        status = "TIMED OUT"
+    elif metrics.size_exceeded:
+        level = logging.ERROR
+        status = "SIZE EXCEEDED"
+    elif metrics.size_ratio >= config.size_warn_threshold:
+        level = logging.WARNING
+        status = "completed (size warning)"
+
+    logger.log(
+        level,
+        f"State {operation} {status}: "
+        f"duration={metrics.duration_sec:.2f}s/{config.capture_timeout_sec}s, "
+        f"size={metrics.size_bytes:,}B/{config.max_size_bytes:,}B "
+        f"({metrics.size_ratio:.1%})",
+    )