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

tetra_rp/runtime/circuit_breaker.py

@@ -0,0 +1,274 @@
+"""Circuit breaker pattern for handling endpoint failures."""
+
+import asyncio
+import logging
+import time
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any, Callable, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class CircuitState(Enum):
+    """Circuit breaker state machine."""
+
+    CLOSED = "closed"
+    OPEN = "open"
+    HALF_OPEN = "half_open"
+
+
+@dataclass
+class CircuitBreakerStats:
+    """Statistics for a circuit breaker instance."""
+
+    state: CircuitState = CircuitState.CLOSED
+    failure_count: int = 0
+    success_count: int = 0
+    last_failure_at: Optional[datetime] = None
+    last_success_at: Optional[datetime] = None
+    state_changed_at: datetime = field(
+        default_factory=lambda: datetime.now(timezone.utc)
+    )
+    total_requests: int = 0
+    total_failures: int = 0
+
+
+class EndpointCircuitBreaker:
+    """Circuit breaker for a single endpoint with sliding window."""
+
+    def __init__(
+        self,
+        endpoint_url: str,
+        failure_threshold: int = 5,
+        success_threshold: int = 2,
+        timeout_seconds: int = 60,
+        window_size: int = 10,
+    ):
+        """Initialize circuit breaker for an endpoint.
+
+        Args:
+            endpoint_url: URL of the endpoint to protect
+            failure_threshold: Failures required to open circuit
+            success_threshold: Successes required to close circuit
+            timeout_seconds: Time before attempting recovery
+            window_size: Size of sliding window for counting failures
+        """
+        self.endpoint_url = endpoint_url
+        self.failure_threshold = failure_threshold
+        self.success_threshold = success_threshold
+        self.timeout_seconds = timeout_seconds
+        self.window_size = window_size
+        self.stats = CircuitBreakerStats()
+        self._lock = asyncio.Lock()
+        self._failure_times: list[float] = []
+
+    async def execute(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+        """Execute function with circuit breaker protection.
+
+        Args:
+            func: Async function to execute
+            *args: Positional arguments for func
+            **kwargs: Keyword arguments for func
+
+        Returns:
+            Result from func
+
+        Raises:
+            CircuitBreakerOpenError: If circuit is open
+            Exception: Any exception raised by func
+        """
+        async with self._lock:
+            state = self.stats.state
+
+            if state == CircuitState.OPEN:
+                # Check if timeout has passed
+                if self._should_attempt_recovery():
+                    self._transition_to_half_open()
+                else:
+                    raise CircuitBreakerOpenError(
+                        f"Circuit OPEN for {self.endpoint_url}. "
+                        f"Retry in {self._seconds_until_recovery()}s"
+                    )
+
+        # Execute function
+        try:
+            result = await func(*args, **kwargs)
+            await self._on_success()
+            return result
+        except Exception as e:
+            await self._on_failure(e)
+            raise
+
+    async def _on_success(self) -> None:
+        """Record successful request."""
+        async with self._lock:
+            self.stats.success_count += 1
+            self.stats.total_requests += 1
+            self.stats.last_success_at = datetime.now(timezone.utc)
+
+            logger.debug(
+                f"Circuit breaker {self.endpoint_url}: "
+                f"success {self.stats.success_count}/{self.success_threshold}"
+            )
+
+            if self.stats.state == CircuitState.HALF_OPEN:
+                # Close circuit after successes
+                if self.stats.success_count >= self.success_threshold:
+                    self._transition_to_closed()
+            elif self.stats.state == CircuitState.CLOSED:
+                # Reset failure count on success
+                self.stats.failure_count = 0
+                self._failure_times.clear()
+
+    async def _on_failure(self, error: Exception) -> None:
+        """Record failed request."""
+        async with self._lock:
+            self.stats.failure_count += 1
+            self.stats.total_failures += 1
+            self.stats.total_requests += 1
+            self.stats.last_failure_at = datetime.now(timezone.utc)
+
+            # Track failure times for sliding window
+            now = time.time()
+            self._failure_times.append(now)
+
+            # Keep only failures within window
+            cutoff = now - self.timeout_seconds
+            self._failure_times = [t for t in self._failure_times if t > cutoff]
+
+            logger.debug(
+                f"Circuit breaker {self.endpoint_url}: "
+                f"failure {self.stats.failure_count}/{self.failure_threshold}, "
+                f"error: {error}"
+            )
+
+            if self.stats.state == CircuitState.HALF_OPEN:
+                # Open circuit on first failure in half-open
+                self._transition_to_open()
+            elif self.stats.state == CircuitState.CLOSED:
+                # Open circuit if threshold reached
+                if len(self._failure_times) >= self.failure_threshold:
+                    self._transition_to_open()
+
+    def _transition_to_open(self) -> None:
+        """Transition circuit to OPEN state."""
+        if self.stats.state == CircuitState.OPEN:
+            return  # Already open
+        self.stats.state = CircuitState.OPEN
+        self.stats.state_changed_at = datetime.now(timezone.utc)
+        self.stats.success_count = 0
+        logger.warning(
+            f"Circuit breaker OPEN for {self.endpoint_url} "
+            f"after {self.stats.failure_count} failures"
+        )
+
+    def _transition_to_half_open(self) -> None:
+        """Transition circuit to HALF_OPEN state."""
+        self.stats.state = CircuitState.HALF_OPEN
+        self.stats.state_changed_at = datetime.now(timezone.utc)
+        self.stats.failure_count = 0
+        self.stats.success_count = 0
+        logger.info(
+            f"Circuit breaker HALF_OPEN for {self.endpoint_url}, testing recovery"
+        )
+
+    def _transition_to_closed(self) -> None:
+        """Transition circuit to CLOSED state."""
+        self.stats.state = CircuitState.CLOSED
+        self.stats.state_changed_at = datetime.now(timezone.utc)
+        self.stats.failure_count = 0
+        self.stats.success_count = 0
+        self._failure_times.clear()
+        logger.info(f"Circuit breaker CLOSED for {self.endpoint_url}, recovered")
+
+    def _should_attempt_recovery(self) -> bool:
+        """Check if enough time has passed to attempt recovery."""
+        if not self.stats.last_failure_at:
+            return False
+        elapsed = datetime.now(timezone.utc) - self.stats.state_changed_at
+        return elapsed.total_seconds() >= self.timeout_seconds
+
+    def _seconds_until_recovery(self) -> int:
+        """Get seconds until recovery can be attempted."""
+        if not self.stats.state_changed_at:
+            return self.timeout_seconds
+        elapsed = datetime.now(timezone.utc) - self.stats.state_changed_at
+        remaining = self.timeout_seconds - int(elapsed.total_seconds())
+        return max(0, remaining)
+
+    def get_state(self) -> CircuitState:
+        """Get current circuit state."""
+        return self.stats.state
+
+    def get_stats(self) -> CircuitBreakerStats:
+        """Get circuit breaker statistics."""
+        return self.stats
+
+
+class CircuitBreakerRegistry:
+    """Manages circuit breakers for multiple endpoints."""
+
+    def __init__(
+        self,
+        failure_threshold: int = 5,
+        success_threshold: int = 2,
+        timeout_seconds: int = 60,
+    ):
+        """Initialize circuit breaker registry.
+
+        Args:
+            failure_threshold: Failures required to open circuit
+            success_threshold: Successes required to close circuit
+            timeout_seconds: Time before attempting recovery
+        """
+        self.failure_threshold = failure_threshold
+        self.success_threshold = success_threshold
+        self.timeout_seconds = timeout_seconds
+        self._breakers: dict[str, EndpointCircuitBreaker] = {}
+        self._lock = asyncio.Lock()
+
+    def get_breaker(self, endpoint_url: str) -> EndpointCircuitBreaker:
+        """Get or create circuit breaker for endpoint.
+
+        Args:
+            endpoint_url: URL of the endpoint
+
+        Returns:
+            EndpointCircuitBreaker instance
+        """
+        if endpoint_url not in self._breakers:
+            self._breakers[endpoint_url] = EndpointCircuitBreaker(
+                endpoint_url,
+                failure_threshold=self.failure_threshold,
+                success_threshold=self.success_threshold,
+                timeout_seconds=self.timeout_seconds,
+            )
+        return self._breakers[endpoint_url]
+
+    def get_state(self, endpoint_url: str) -> CircuitState:
+        """Get state of circuit breaker for endpoint.
+
+        Args:
+            endpoint_url: URL of the endpoint
+
+        Returns:
+            Current circuit state
+        """
+        breaker = self.get_breaker(endpoint_url)
+        return breaker.get_state()
+
+    def get_all_stats(self) -> dict[str, CircuitBreakerStats]:
+        """Get statistics for all circuit breakers.
+
+        Returns:
+            Mapping of endpoint URLs to statistics
+        """
+        return {url: breaker.get_stats() for url, breaker in self._breakers.items()}
+
+
+class CircuitBreakerOpenError(Exception):
+    """Raised when circuit breaker is open."""
+
+    pass

tetra_rp/runtime/config.py

@@ -0,0 +1,12 @@
+"""Configuration constants for runtime module."""
+
+# HTTP client configuration
+DEFAULT_REQUEST_TIMEOUT = 10  # seconds
+DEFAULT_MAX_RETRIES = 3
+DEFAULT_BACKOFF_BASE = 2
+
+# Manifest cache configuration
+DEFAULT_CACHE_TTL = 300  # seconds
+
+# Serialization limits
+MAX_PAYLOAD_SIZE = 10 * 1024 * 1024  # 10MB

tetra_rp/runtime/exceptions.py

@@ -0,0 +1,49 @@
+"""Custom exceptions for cross-endpoint runtime."""
+
+
+class FlashRuntimeError(Exception):
+    """Base exception for runtime errors in cross-endpoint execution."""
+
+    pass
+
+
+class RemoteExecutionError(FlashRuntimeError):
+    """Raised when remote function execution fails."""
+
+    pass
+
+
+class SerializationError(FlashRuntimeError):
+    """Raised when serialization or deserialization of arguments fails."""
+
+    pass
+
+
+class GraphQLError(FlashRuntimeError):
+    """Base exception for GraphQL-related errors."""
+
+    pass
+
+
+class GraphQLMutationError(GraphQLError):
+    """Raised when a GraphQL mutation fails unexpectedly."""
+
+    pass
+
+
+class GraphQLQueryError(GraphQLError):
+    """Raised when a GraphQL query fails unexpectedly."""
+
+    pass
+
+
+class ManifestError(FlashRuntimeError):
+    """Raised when manifest is invalid, missing, or has unexpected structure."""
+
+    pass
+
+
+class ManifestServiceUnavailableError(FlashRuntimeError):
+    """Raised when manifest service is unavailable."""
+
+    pass

tetra_rp/runtime/generic_handler.py

@@ -0,0 +1,206 @@
+"""Generic RunPod serverless handler factory for Flash."""
+
+import json
+import logging
+import traceback
+from pathlib import Path
+from typing import Any, Callable, Dict
+
+from .serialization import deserialize_args, deserialize_kwargs, serialize_arg
+
+logger = logging.getLogger(__name__)
+
+
+def load_manifest(manifest_path: Path | None = None) -> Dict[str, Any]:
+    """Load flash_manifest.json with fallback search.
+
+    Searches multiple locations for manifest:
+    1. Provided path (if given)
+    2. Current working directory
+    3. Module directory
+    4. Three levels up (legacy location)
+
+    Args:
+        manifest_path: Optional explicit path to manifest file
+
+    Returns:
+        Manifest dictionary, or empty dict if not found
+    """
+    if manifest_path and manifest_path.exists():
+        try:
+            with open(manifest_path) as f:
+                return json.load(f)
+        except Exception as e:
+            logger.warning(f"Failed to load manifest from {manifest_path}: {e}")
+            return {"resources": {}, "function_registry": {}}
+
+    # Search multiple locations
+    search_paths = [
+        Path.cwd() / "flash_manifest.json",
+        Path(__file__).parent / "flash_manifest.json",
+        Path(__file__).parent.parent.parent / "flash_manifest.json",
+    ]
+
+    for path in search_paths:
+        if path.exists():
+            try:
+                with open(path) as f:
+                    return json.load(f)
+            except Exception as e:
+                logger.debug(f"Failed to load manifest from {path}: {e}")
+                continue
+
+    logger.warning("flash_manifest.json not found in any expected location")
+    return {"resources": {}, "function_registry": {}}
+
+
+def deserialize_arguments(job_input: Dict[str, Any]) -> tuple[list, dict]:
+    """Deserialize function arguments from job input.
+
+    Args:
+        job_input: Input dict from RunPod job with 'args' and 'kwargs' keys
+
+    Returns:
+        Tuple of (args list, kwargs dict) deserialized from cloudpickle
+    """
+    args = deserialize_args(job_input.get("args", []))
+    kwargs = deserialize_kwargs(job_input.get("kwargs", {}))
+    return args, kwargs
+
+
+def serialize_result(result: Any) -> str:
+    """Serialize function result for response.
+
+    Args:
+        result: Return value from function
+
+    Returns:
+        Base64-encoded cloudpickle of result
+    """
+    return serialize_arg(result)
+
+
+def execute_function(
+    func_or_class: Callable,
+    args: list,
+    kwargs: dict,
+    execution_type: str,
+    job_input: Dict[str, Any],
+) -> Any:
+    """Execute function or class method.
+
+    Args:
+        func_or_class: Function or class to execute
+        args: Positional arguments
+        kwargs: Keyword arguments
+        execution_type: Either "function" or "class"
+        job_input: Full job input for method calls
+
+    Returns:
+        Result of execution
+
+    Raises:
+        Exception: If execution fails
+    """
+    if execution_type == "class":
+        # Instantiate class with constructor args
+        instance = func_or_class(*args, **kwargs)
+        method_name = job_input.get("method_name", "__call__")
+
+        # Call method on instance
+        method = getattr(instance, method_name)
+        method_args, method_kwargs = deserialize_arguments(
+            {
+                "args": job_input.get("method_args", []),
+                "kwargs": job_input.get("method_kwargs", {}),
+            }
+        )
+        return method(*method_args, **method_kwargs)
+    else:
+        # Direct function call
+        return func_or_class(*args, **kwargs)
+
+
+def create_handler(function_registry: Dict[str, Callable]) -> Callable:
+    """Create a RunPod serverless handler with given function registry.
+
+    This factory function creates a handler that:
+    1. Deserializes function arguments from cloudpickle + base64
+    2. Looks up function/class in registry by name
+    3. Executes function or class method
+    4. Serializes result back to cloudpickle + base64
+    5. Returns RunPod-compatible response dict
+
+    Args:
+        function_registry: Dict mapping function names to function/class objects
+
+    Returns:
+        Handler function compatible with runpod.serverless.start()
+
+    Example:
+        ```python
+        from tetra_rp.runtime.generic_handler import create_handler
+        from workers.gpu import process_data, analyze_data
+
+        registry = {
+            "process_data": process_data,
+            "analyze_data": analyze_data,
+        }
+
+        handler = create_handler(registry)
+
+        if __name__ == "__main__":
+            import runpod
+            runpod.serverless.start({"handler": handler})
+        ```
+    """
+
+    def handler(job: Dict[str, Any]) -> Dict[str, Any]:
+        """RunPod serverless handler.
+
+        Args:
+            job: RunPod job dict with 'input' key
+
+        Returns:
+            Response dict with 'success', 'result'/'error' keys
+        """
+        job_input = job.get("input", {})
+        function_name = job_input.get("function_name")
+        execution_type = job_input.get("execution_type", "function")
+
+        if function_name not in function_registry:
+            return {
+                "success": False,
+                "error": f"Function '{function_name}' not found in registry. "
+                f"Available: {list(function_registry.keys())}",
+                "traceback": "",
+            }
+
+        try:
+            # Deserialize arguments
+            args, kwargs = deserialize_arguments(job_input)
+
+            # Get function/class from registry
+            func_or_class = function_registry[function_name]
+
+            # Execute function or class
+            result = execute_function(
+                func_or_class, args, kwargs, execution_type, job_input
+            )
+
+            # Serialize result
+            serialized_result = serialize_result(result)
+
+            return {
+                "success": True,
+                "result": serialized_result,
+            }
+
+        except Exception as e:
+            return {
+                "success": False,
+                "error": str(e),
+                "traceback": traceback.format_exc(),
+            }
+
+    return handler