tetra-rp 0.6.0__py3-none-any.whl → 0.24.0__py3-none-any.whl

tetra_rp/runtime/production_wrapper.py

@@ -0,0 +1,266 @@
+"""Production wrapper for cross-endpoint function routing."""
+
+import logging
+from typing import Any, Callable, Dict, Optional
+
+from tetra_rp.core.resources.serverless import ServerlessResource
+
+from .exceptions import RemoteExecutionError
+from .serialization import serialize_args, serialize_kwargs
+from .service_registry import ServiceRegistry
+
+logger = logging.getLogger(__name__)
+
+
+class ProductionWrapper:
+    """Wrapper that routes function execution between endpoints.
+
+    Intercepts stub execution and determines if the call is local (execute
+    directly) or remote (call via HTTP to another endpoint).
+    """
+
+    def __init__(self, service_registry: ServiceRegistry):
+        """Initialize production wrapper.
+
+        Args:
+            service_registry: Service registry for routing decisions.
+        """
+        self.service_registry = service_registry
+
+    async def wrap_function_execution(
+        self,
+        original_stub_func: Callable,
+        func: Callable,
+        dependencies: Optional[list],
+        system_dependencies: Optional[list],
+        accelerate_downloads: bool,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        """Route function execution to local or remote endpoint.
+
+        Args:
+            original_stub_func: The original stubbed_resource function.
+            func: The decorated function being called.
+            dependencies: Pip dependencies (for local execution).
+            system_dependencies: System dependencies (for local execution).
+            accelerate_downloads: Download acceleration flag (for local).
+            *args: Function positional arguments.
+            **kwargs: Function keyword arguments.
+
+        Returns:
+            Function execution result.
+
+        Raises:
+            Exception: If execution fails.
+        """
+        function_name = func.__name__
+
+        # Ensure manifest is loaded
+        await self.service_registry._ensure_manifest_loaded()
+
+        # Determine routing
+        try:
+            resource = await self.service_registry.get_resource_for_function(
+                function_name
+            )
+        except ValueError as e:
+            # Function not in manifest, execute locally
+            logger.debug(
+                f"Function {function_name} not in manifest: {e}, executing locally"
+            )
+            return await original_stub_func(
+                func,
+                dependencies,
+                system_dependencies,
+                accelerate_downloads,
+                *args,
+                **kwargs,
+            )
+
+        # Local execution
+        if resource is None:
+            logger.debug(f"Executing local function: {function_name}")
+            return await original_stub_func(
+                func,
+                dependencies,
+                system_dependencies,
+                accelerate_downloads,
+                *args,
+                **kwargs,
+            )
+
+        # Remote execution
+        logger.debug(f"Routing function {function_name} to remote endpoint")
+        return await self._execute_remote(
+            resource,
+            function_name,
+            args,
+            kwargs,
+            execution_type="function",
+        )
+
+    async def wrap_class_method_execution(
+        self,
+        original_method_func: Callable,
+        request: Any,
+    ) -> Any:
+        """Route class method execution to local or remote endpoint.
+
+        Args:
+            original_method_func: The original execute_class_method function.
+            request: FunctionRequest containing class and method info.
+
+        Returns:
+            Method execution result.
+
+        Raises:
+            Exception: If execution fails.
+        """
+        # Ensure manifest is loaded
+        await self.service_registry._ensure_manifest_loaded()
+
+        class_name = getattr(request, "class_name", None)
+
+        if not class_name:
+            # No class name, execute locally
+            return await original_method_func(request)
+
+        # Determine routing
+        try:
+            resource = await self.service_registry.get_resource_for_function(class_name)
+        except ValueError:
+            # Class not in manifest, execute locally
+            logger.debug(f"Class {class_name} not in manifest, executing locally")
+            return await original_method_func(request)
+
+        # Local execution
+        if resource is None:
+            logger.debug(f"Executing local class method: {class_name}")
+            return await original_method_func(request)
+
+        # Remote execution
+        logger.debug(f"Routing class {class_name} to remote endpoint")
+
+        # Convert FunctionRequest to dict payload
+        payload = self._build_class_payload(request)
+        return await self._execute_remote(
+            resource,
+            class_name,
+            (),
+            payload.get("input", {}),
+            execution_type="class",
+        )
+
+    async def _execute_remote(
+        self,
+        resource: ServerlessResource,
+        function_name: str,
+        args: tuple,
+        kwargs: dict,
+        execution_type: str = "function",
+    ) -> Any:
+        """Execute function on remote endpoint.
+
+        Args:
+            resource: ServerlessResource with endpoint ID set.
+            function_name: Name of function/class to execute.
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+            execution_type: "function" or "class".
+
+        Returns:
+            Execution result.
+
+        Raises:
+            RemoteExecutionError: If remote execution fails.
+        """
+        # Serialize arguments
+        serialized_args = serialize_args(args)
+        serialized_kwargs = serialize_kwargs(kwargs)
+
+        # Build payload matching RunPod format
+        payload = {
+            "input": {
+                "function_name": function_name,
+                "execution_type": execution_type,
+                "args": serialized_args,
+                "kwargs": serialized_kwargs,
+            }
+        }
+
+        # Execute via ServerlessResource
+        result = await resource.run_sync(payload)
+
+        # Handle response
+        if result.error:
+            raise RemoteExecutionError(
+                f"Remote execution of {function_name} failed: {result.error}"
+            )
+
+        return result.output
+
+    def _build_class_payload(self, request: Any) -> Dict[str, Any]:
+        """Build payload from FunctionRequest for class execution.
+
+        Args:
+            request: FunctionRequest object.
+
+        Returns:
+            RunPod-format payload dict.
+        """
+        # Extract request data - handle both dict and object access patterns
+        if isinstance(request, dict):
+            data = request
+        else:
+            data = (
+                request.model_dump(exclude_none=True)
+                if hasattr(request, "model_dump")
+                else {}
+            )
+
+        # Extract class execution data
+        payload = {
+            "input": {
+                "function_name": data.get("class_name"),
+                "execution_type": "class",
+                "args": data.get("args", []),
+                "kwargs": data.get("kwargs", {}),
+                "method_name": data.get("method_name"),
+            }
+        }
+
+        return payload
+
+
+# Singleton instance management
+_wrapper_instance: Optional[ProductionWrapper] = None
+
+
+def create_production_wrapper(
+    service_registry: Optional[ServiceRegistry] = None,
+) -> ProductionWrapper:
+    """Create or get singleton ProductionWrapper instance.
+
+    Args:
+        service_registry: Service registry. Creates if not provided.
+
+    Returns:
+        ProductionWrapper instance.
+    """
+    global _wrapper_instance
+
+    if _wrapper_instance is None:
+        # Create components if not provided
+        if service_registry is None:
+            service_registry = ServiceRegistry()
+
+        _wrapper_instance = ProductionWrapper(service_registry)
+
+    return _wrapper_instance
+
+
+def reset_wrapper() -> None:
+    """Reset singleton wrapper (mainly for testing)."""
+    global _wrapper_instance
+    _wrapper_instance = None

tetra_rp/runtime/reliability_config.py

@@ -0,0 +1,149 @@
+"""Centralized configuration for reliability features."""
+
+import os
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Optional
+
+
+class LoadBalancerStrategy(Enum):
+    """Load balancing strategies for endpoint selection."""
+
+    ROUND_ROBIN = "round_robin"
+    LEAST_CONNECTIONS = "least_connections"
+    RANDOM = "random"
+
+
+@dataclass
+class CircuitBreakerConfig:
+    """Configuration for circuit breaker behavior."""
+
+    enabled: bool = True
+    failure_threshold: int = 5
+    success_threshold: int = 2
+    timeout_seconds: int = 60
+    window_size: int = 10
+
+
+@dataclass
+class LoadBalancerConfig:
+    """Configuration for load balancer behavior."""
+
+    enabled: bool = False
+    strategy: LoadBalancerStrategy = LoadBalancerStrategy.ROUND_ROBIN
+
+
+@dataclass
+class RetryConfig:
+    """Configuration for retry behavior with exponential backoff."""
+
+    enabled: bool = True
+    max_attempts: int = 3
+    base_delay: float = 0.5
+    max_delay: float = 10.0
+    jitter: float = 0.2
+    retryable_exceptions: tuple = field(
+        default_factory=lambda: (TimeoutError, ConnectionError)
+    )
+    retryable_status_codes: set = field(
+        default_factory=lambda: {408, 429, 500, 502, 503, 504}
+    )
+
+
+@dataclass
+class MetricsConfig:
+    """Configuration for metrics collection."""
+
+    enabled: bool = True
+    namespace: str = "tetra.metrics"
+
+
+@dataclass
+class ReliabilityConfig:
+    """Centralized reliability features configuration."""
+
+    circuit_breaker: CircuitBreakerConfig = field(default_factory=CircuitBreakerConfig)
+    load_balancer: LoadBalancerConfig = field(default_factory=LoadBalancerConfig)
+    retry: RetryConfig = field(default_factory=RetryConfig)
+    metrics: MetricsConfig = field(default_factory=MetricsConfig)
+
+    @classmethod
+    def from_env(cls) -> "ReliabilityConfig":
+        """Load configuration from environment variables.
+
+        Environment variables:
+        - TETRA_CIRCUIT_BREAKER_ENABLED: Enable circuit breaker (default: true)
+        - TETRA_CB_FAILURE_THRESHOLD: Failures before opening (default: 5)
+        - TETRA_CB_SUCCESS_THRESHOLD: Successes to close (default: 2)
+        - TETRA_CB_TIMEOUT_SECONDS: Time before half-open (default: 60)
+        - TETRA_LOAD_BALANCER_ENABLED: Enable load balancer (default: false)
+        - TETRA_LB_STRATEGY: Load balancer strategy (default: round_robin)
+        - TETRA_RETRY_ENABLED: Enable retry (default: true)
+        - TETRA_RETRY_MAX_ATTEMPTS: Max retry attempts (default: 3)
+        - TETRA_RETRY_BASE_DELAY: Base delay for backoff (default: 0.5)
+        - TETRA_METRICS_ENABLED: Enable metrics (default: true)
+
+        Returns:
+            ReliabilityConfig initialized from environment variables.
+        """
+        circuit_breaker = CircuitBreakerConfig(
+            enabled=os.getenv("TETRA_CIRCUIT_BREAKER_ENABLED", "true").lower()
+            == "true",
+            failure_threshold=int(os.getenv("TETRA_CB_FAILURE_THRESHOLD", "5")),
+            success_threshold=int(os.getenv("TETRA_CB_SUCCESS_THRESHOLD", "2")),
+            timeout_seconds=int(os.getenv("TETRA_CB_TIMEOUT_SECONDS", "60")),
+        )
+
+        strategy_str = os.getenv("TETRA_LB_STRATEGY", "round_robin").lower()
+        try:
+            strategy = LoadBalancerStrategy(strategy_str)
+        except ValueError:
+            strategy = LoadBalancerStrategy.ROUND_ROBIN
+
+        load_balancer = LoadBalancerConfig(
+            enabled=os.getenv("TETRA_LOAD_BALANCER_ENABLED", "false").lower() == "true",
+            strategy=strategy,
+        )
+
+        retry = RetryConfig(
+            enabled=os.getenv("TETRA_RETRY_ENABLED", "true").lower() == "true",
+            max_attempts=int(os.getenv("TETRA_RETRY_MAX_ATTEMPTS", "3")),
+            base_delay=float(os.getenv("TETRA_RETRY_BASE_DELAY", "0.5")),
+        )
+
+        metrics = MetricsConfig(
+            enabled=os.getenv("TETRA_METRICS_ENABLED", "true").lower() == "true",
+        )
+
+        return cls(
+            circuit_breaker=circuit_breaker,
+            load_balancer=load_balancer,
+            retry=retry,
+            metrics=metrics,
+        )
+
+
+# Global default configuration
+_config: Optional[ReliabilityConfig] = None
+
+
+def get_reliability_config() -> ReliabilityConfig:
+    """Get global reliability configuration (lazy-loaded).
+
+    Returns:
+        ReliabilityConfig instance initialized from environment.
+    """
+    global _config
+    if _config is None:
+        _config = ReliabilityConfig.from_env()
+    return _config
+
+
+def set_reliability_config(config: ReliabilityConfig) -> None:
+    """Set global reliability configuration (for testing).
+
+    Args:
+        config: ReliabilityConfig to set as global.
+    """
+    global _config
+    _config = config

tetra_rp/runtime/retry_manager.py

@@ -0,0 +1,118 @@
+"""Retry logic with exponential backoff for failed remote calls."""
+
+import asyncio
+import logging
+from typing import Any, Callable, Optional, Set, Tuple, Type
+
+from tetra_rp.core.utils.backoff import get_backoff_delay
+
+logger = logging.getLogger(__name__)
+
+
+class RetryExhaustedError(Exception):
+    """Raised when max retry attempts are exceeded."""
+
+    pass
+
+
+async def retry_with_backoff(
+    func: Callable[..., Any],
+    max_attempts: int = 3,
+    base_delay: float = 0.5,
+    max_delay: float = 10.0,
+    jitter: float = 0.2,
+    retryable_exceptions: Optional[Tuple[Type[Exception], ...]] = None,
+    retryable_status_codes: Optional[Set[int]] = None,
+    circuit_breaker: Optional[Any] = None,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    """Execute async function with retry and exponential backoff.
+
+    Args:
+        func: Async function to execute
+        max_attempts: Maximum number of attempts (default: 3)
+        base_delay: Base delay between retries in seconds (default: 0.5)
+        max_delay: Maximum delay between retries (default: 10.0)
+        jitter: Jitter factor (0.0-1.0) to add randomness (default: 0.2)
+        retryable_exceptions: Tuple of exception types to retry on
+            (default: (asyncio.TimeoutError, ConnectionError))
+        retryable_status_codes: Set of HTTP status codes to retry on
+            (default: {408, 429, 500, 502, 503, 504})
+        circuit_breaker: Optional circuit breaker to check before retry
+        *args: Positional arguments for func
+        **kwargs: Keyword arguments for func
+
+    Returns:
+        Result from successful function call
+
+    Raises:
+        RetryExhaustedError: If max attempts exceeded
+        Exception: If non-retryable exception occurs
+    """
+    if retryable_exceptions is None:
+        retryable_exceptions = (asyncio.TimeoutError, ConnectionError)
+
+    if retryable_status_codes is None:
+        retryable_status_codes = {408, 429, 500, 502, 503, 504}
+
+    last_exception: Optional[Exception] = None
+
+    for attempt in range(max_attempts):
+        try:
+            # Check circuit breaker before attempting
+            if circuit_breaker is not None:
+                from tetra_rp.runtime.circuit_breaker import CircuitState
+
+                if circuit_breaker.get_state() == CircuitState.OPEN:
+                    raise RuntimeError(
+                        f"Circuit breaker OPEN, skipping retry attempt {attempt + 1}"
+                    )
+
+            result = await func(*args, **kwargs)
+
+            # Log success on retry
+            if attempt > 0:
+                logger.info(f"Retry succeeded on attempt {attempt + 1}/{max_attempts}")
+
+            return result
+
+        except Exception as e:
+            last_exception = e
+
+            # Check if exception is retryable
+            if not isinstance(e, retryable_exceptions):
+                logger.debug(
+                    f"Non-retryable exception in {func.__name__}: {type(e).__name__}"
+                )
+                raise
+
+            # Check for retryable status codes (if exception has status_code)
+            if hasattr(e, "status_code"):
+                if e.status_code not in retryable_status_codes:  # type: ignore
+                    logger.debug(
+                        f"Non-retryable status code {e.status_code} in {func.__name__}"
+                    )
+                    raise
+
+            # If this is the last attempt, don't retry
+            if attempt >= max_attempts - 1:
+                logger.warning(
+                    f"Max retries ({max_attempts}) exhausted for {func.__name__}"
+                )
+                raise RetryExhaustedError(
+                    f"Failed after {max_attempts} attempts: {e}"
+                ) from e
+
+            # Calculate delay with exponential backoff and jitter
+            delay = get_backoff_delay(attempt, base_delay, max_delay, jitter=jitter)
+            logger.debug(
+                f"Retry {attempt + 1}/{max_attempts} for {func.__name__} "
+                f"after {delay:.2f}s"
+            )
+            await asyncio.sleep(delay)
+
+    # Should never reach here, but handle edge case
+    if last_exception:
+        raise last_exception
+    raise RetryExhaustedError(f"Failed after {max_attempts} attempts")

tetra_rp/runtime/serialization.py

@@ -0,0 +1,124 @@
+"""Shared serialization utilities for cloudpickle + base64 encoding."""
+
+import base64
+from typing import Any, Dict, List
+
+import cloudpickle
+
+from .exceptions import SerializationError
+
+
+def serialize_arg(arg: Any) -> str:
+    """Serialize single argument with cloudpickle + base64.
+
+    Args:
+        arg: Argument to serialize.
+
+    Returns:
+        Base64-encoded cloudpickle serialized string.
+
+    Raises:
+        SerializationError: If serialization fails.
+    """
+    try:
+        return base64.b64encode(cloudpickle.dumps(arg)).decode("utf-8")
+    except Exception as e:
+        raise SerializationError(f"Failed to serialize argument: {e}") from e
+
+
+def serialize_args(args: tuple) -> List[str]:
+    """Serialize positional arguments.
+
+    Args:
+        args: Tuple of arguments to serialize.
+
+    Returns:
+        List of base64-encoded serialized arguments.
+
+    Raises:
+        SerializationError: If serialization fails.
+    """
+    try:
+        return [serialize_arg(arg) for arg in args]
+    except SerializationError:
+        raise
+    except Exception as e:
+        raise SerializationError(f"Failed to serialize args: {e}") from e
+
+
+def serialize_kwargs(kwargs: dict) -> Dict[str, str]:
+    """Serialize keyword arguments.
+
+    Args:
+        kwargs: Dictionary of keyword arguments.
+
+    Returns:
+        Dictionary with base64-encoded serialized values.
+
+    Raises:
+        SerializationError: If serialization fails.
+    """
+    try:
+        return {k: serialize_arg(v) for k, v in kwargs.items()}
+    except SerializationError:
+        raise
+    except Exception as e:
+        raise SerializationError(f"Failed to serialize kwargs: {e}") from e
+
+
+def deserialize_arg(arg_b64: str) -> Any:
+    """Deserialize single base64-encoded cloudpickle argument.
+
+    Args:
+        arg_b64: Base64-encoded serialized argument.
+
+    Returns:
+        Deserialized argument.
+
+    Raises:
+        SerializationError: If deserialization fails.
+    """
+    try:
+        return cloudpickle.loads(base64.b64decode(arg_b64))
+    except Exception as e:
+        raise SerializationError(f"Failed to deserialize argument: {e}") from e
+
+
+def deserialize_args(args_b64: List[str]) -> List[Any]:
+    """Deserialize list of base64-encoded arguments.
+
+    Args:
+        args_b64: List of base64-encoded serialized arguments.
+
+    Returns:
+        List of deserialized arguments.
+
+    Raises:
+        SerializationError: If deserialization fails.
+    """
+    try:
+        return [deserialize_arg(arg) for arg in args_b64]
+    except SerializationError:
+        raise
+    except Exception as e:
+        raise SerializationError(f"Failed to deserialize args: {e}") from e
+
+
+def deserialize_kwargs(kwargs_b64: Dict[str, str]) -> Dict[str, Any]:
+    """Deserialize dict of base64-encoded keyword arguments.
+
+    Args:
+        kwargs_b64: Dictionary with base64-encoded serialized values.
+
+    Returns:
+        Dictionary with deserialized values.
+
+    Raises:
+        SerializationError: If deserialization fails.
+    """
+    try:
+        return {k: deserialize_arg(v) for k, v in kwargs_b64.items()}
+    except SerializationError:
+        raise
+    except Exception as e:
+        raise SerializationError(f"Failed to deserialize kwargs: {e}") from e