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

tetra_rp/runtime/lb_handler.py

@@ -0,0 +1,189 @@
+"""Factory for creating FastAPI load-balanced handlers.
+
+This module provides the factory function for generating FastAPI applications
+that handle load-balanced serverless endpoints. It supports:
+- User-defined HTTP routes
+- /execute endpoint for @remote function execution (LiveLoadBalancer only)
+
+Security Model:
+    The /execute endpoint accepts and executes serialized function code. This is
+    secure because:
+    1. The function code originates from the client's @remote decorator
+    2. The client (user) controls what function gets sent
+    3. This mirrors the trusted client model of LiveServerlessStub
+    4. In production, API authentication should protect the /execute endpoint
+
+    Users should NOT expose the /execute endpoint to untrusted clients.
+"""
+
+import inspect
+import logging
+from typing import Any, Callable, Dict
+
+from fastapi import FastAPI, Request
+
+from .serialization import (
+    deserialize_args,
+    deserialize_kwargs,
+    serialize_arg,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def create_lb_handler(
+    route_registry: Dict[tuple[str, str], Callable],
+    include_execute: bool = False,
+    lifespan: Callable = None,
+) -> FastAPI:
+    """Create FastAPI app with routes from registry.
+
+    Args:
+        route_registry: Mapping of (HTTP_METHOD, path) -> handler_function
+                       Example: {("GET", "/api/health"): health_check}
+        include_execute: Whether to register /execute endpoint for @remote execution.
+                        Only used for LiveLoadBalancer (local development).
+                        Deployed endpoints should not expose /execute for security.
+        lifespan: Optional lifespan context manager for startup/shutdown hooks.
+
+    Returns:
+        Configured FastAPI application with routes registered.
+    """
+    app = FastAPI(title="Flash Load-Balanced Handler", lifespan=lifespan)
+
+    # Register /execute endpoint for @remote stub execution (if enabled)
+    if include_execute:
+
+        @app.post("/execute")
+        async def execute_remote_function(request: Request) -> Dict[str, Any]:
+            """Framework endpoint for @remote decorator execution.
+
+            WARNING: This endpoint is INTERNAL to the Flash framework. It should only be
+            called by the @remote stub from tetra_rp.stubs.load_balancer_sls. Exposing
+            this endpoint to untrusted clients could allow arbitrary code execution.
+
+            Accepts serialized function code and arguments, executes them,
+            and returns serialized result.
+
+            Request body:
+                {
+                    "function_name": "process_data",
+                    "function_code": "def process_data(x, y): return x + y",
+                    "args": [base64_encoded_arg1, base64_encoded_arg2],
+                    "kwargs": {"key": base64_encoded_value}
+                }
+
+            Returns:
+                {
+                    "success": true,
+                    "result": base64_encoded_result
+                }
+                or
+                {
+                    "success": false,
+                    "error": "error message"
+                }
+            """
+            try:
+                body = await request.json()
+            except Exception as e:
+                logger.error(f"Failed to parse request body: {e}")
+                return {"success": False, "error": f"Invalid request body: {e}"}
+
+            try:
+                # Extract function metadata
+                function_name = body.get("function_name")
+                function_code = body.get("function_code")
+
+                if not function_name or not function_code:
+                    return {
+                        "success": False,
+                        "error": "Missing function_name or function_code in request",
+                    }
+
+                # Deserialize arguments
+                try:
+                    args = deserialize_args(body.get("args", []))
+                    kwargs = deserialize_kwargs(body.get("kwargs", {}))
+                except Exception as e:
+                    logger.error(f"Failed to deserialize arguments: {e}")
+                    return {
+                        "success": False,
+                        "error": f"Failed to deserialize arguments: {e}",
+                    }
+
+                # Execute function in isolated namespace
+                namespace: Dict[str, Any] = {}
+                try:
+                    exec(function_code, namespace)
+                except SyntaxError as e:
+                    logger.error(f"Syntax error in function code: {e}")
+                    return {
+                        "success": False,
+                        "error": f"Syntax error in function code: {e}",
+                    }
+                except Exception as e:
+                    logger.error(f"Error executing function code: {e}")
+                    return {
+                        "success": False,
+                        "error": f"Error executing function code: {e}",
+                    }
+
+                # Get function from namespace
+                if function_name not in namespace:
+                    return {
+                        "success": False,
+                        "error": f"Function '{function_name}' not found in executed code",
+                    }
+
+                func = namespace[function_name]
+
+                # Execute function
+                try:
+                    result = func(*args, **kwargs)
+
+                    # Handle async functions
+                    if inspect.iscoroutine(result):
+                        result = await result
+                except Exception as e:
+                    logger.error(f"Function execution failed: {e}")
+                    return {
+                        "success": False,
+                        "error": f"Function execution failed: {e}",
+                    }
+
+                # Serialize result
+                try:
+                    result_b64 = serialize_arg(result)
+                    return {"success": True, "result": result_b64}
+                except Exception as e:
+                    logger.error(f"Failed to serialize result: {e}")
+                    return {
+                        "success": False,
+                        "error": f"Failed to serialize result: {e}",
+                    }
+
+            except Exception as e:
+                logger.error(f"Unexpected error in /execute endpoint: {e}")
+                return {"success": False, "error": f"Unexpected error: {e}"}
+
+    # Register user-defined routes from registry
+    for (method, path), handler in route_registry.items():
+        method_upper = method.upper()
+
+        if method_upper == "GET":
+            app.get(path)(handler)
+        elif method_upper == "POST":
+            app.post(path)(handler)
+        elif method_upper == "PUT":
+            app.put(path)(handler)
+        elif method_upper == "DELETE":
+            app.delete(path)(handler)
+        elif method_upper == "PATCH":
+            app.patch(path)(handler)
+        else:
+            logger.warning(
+                f"Unsupported HTTP method '{method}' for path '{path}'. Skipping."
+            )
+
+    return app

tetra_rp/runtime/load_balancer.py

@@ -0,0 +1,160 @@
+"""Load balancing strategies for distributed endpoint routing."""
+
+import asyncio
+import logging
+import random
+from typing import TYPE_CHECKING, List, Optional
+
+from tetra_rp.runtime.reliability_config import LoadBalancerStrategy
+
+if TYPE_CHECKING:
+    from tetra_rp.runtime.circuit_breaker import CircuitBreakerRegistry
+
+logger = logging.getLogger(__name__)
+
+
+class LoadBalancer:
+    """Load balancer for selecting endpoints using various strategies."""
+
+    def __init__(
+        self, strategy: LoadBalancerStrategy = LoadBalancerStrategy.ROUND_ROBIN
+    ):
+        """Initialize load balancer.
+
+        Args:
+            strategy: Load balancing strategy to use
+        """
+        self.strategy = strategy
+        self._round_robin_index = 0
+        self._lock = asyncio.Lock()
+        self._in_flight_requests: dict[str, int] = {}
+
+    async def select_endpoint(
+        self,
+        endpoints: List[str],
+        circuit_breaker_registry: Optional["CircuitBreakerRegistry"] = None,
+    ) -> Optional[str]:
+        """Select an endpoint using configured strategy.
+
+        Args:
+            endpoints: List of available endpoint URLs
+            circuit_breaker_registry: Optional circuit breaker registry to check health
+
+        Returns:
+            Selected endpoint URL or None if all endpoints are unhealthy
+        """
+        if not endpoints:
+            return None
+
+        # Filter out unhealthy endpoints if circuit breaker available
+        healthy_endpoints = endpoints
+        if circuit_breaker_registry is not None:
+            from tetra_rp.runtime.circuit_breaker import CircuitState
+
+            healthy_endpoints = [
+                url
+                for url in endpoints
+                if circuit_breaker_registry.get_state(url) != CircuitState.OPEN
+            ]
+
+        if not healthy_endpoints:
+            logger.warning(
+                f"All {len(endpoints)} endpoints are unhealthy (circuit open)"
+            )
+            return None
+
+        if self.strategy == LoadBalancerStrategy.ROUND_ROBIN:
+            return await self._round_robin_select(healthy_endpoints)
+        elif self.strategy == LoadBalancerStrategy.LEAST_CONNECTIONS:
+            return await self._least_connections_select(healthy_endpoints)
+        elif self.strategy == LoadBalancerStrategy.RANDOM:
+            return await self._random_select(healthy_endpoints)
+        else:
+            # Default to round-robin
+            return await self._round_robin_select(healthy_endpoints)
+
+    async def _round_robin_select(self, endpoints: List[str]) -> str:
+        """Select endpoint using round-robin strategy.
+
+        Args:
+            endpoints: List of available endpoints
+
+        Returns:
+            Selected endpoint URL
+        """
+        async with self._lock:
+            selected = endpoints[self._round_robin_index % len(endpoints)]
+            self._round_robin_index += 1
+        logger.debug(
+            f"Load balancer: ROUND_ROBIN selected {selected} "
+            f"(index {self._round_robin_index - 1})"
+        )
+        return selected
+
+    async def _least_connections_select(self, endpoints: List[str]) -> str:
+        """Select endpoint with fewest in-flight requests.
+
+        Args:
+            endpoints: List of available endpoints
+
+        Returns:
+            Selected endpoint URL
+        """
+        async with self._lock:
+            # Initialize counts for endpoints
+            for endpoint in endpoints:
+                if endpoint not in self._in_flight_requests:
+                    self._in_flight_requests[endpoint] = 0
+
+            # Find endpoint with minimum connections
+            selected = min(endpoints, key=lambda e: self._in_flight_requests.get(e, 0))
+
+        logger.debug(
+            f"Load balancer: LEAST_CONNECTIONS selected {selected} "
+            f"({self._in_flight_requests.get(selected, 0)} in-flight)"
+        )
+        return selected
+
+    async def _random_select(self, endpoints: List[str]) -> str:
+        """Select endpoint using random strategy.
+
+        Args:
+            endpoints: List of available endpoints
+
+        Returns:
+            Selected endpoint URL
+        """
+        selected = random.choice(endpoints)
+        logger.debug(f"Load balancer: RANDOM selected {selected}")
+        return selected
+
+    async def record_request(self, endpoint: str) -> None:
+        """Record that a request is starting on endpoint.
+
+        Args:
+            endpoint: Endpoint URL
+        """
+        async with self._lock:
+            self._in_flight_requests[endpoint] = (
+                self._in_flight_requests.get(endpoint, 0) + 1
+            )
+
+    async def record_request_complete(self, endpoint: str) -> None:
+        """Record that a request completed on endpoint.
+
+        Args:
+            endpoint: Endpoint URL
+        """
+        async with self._lock:
+            if endpoint in self._in_flight_requests:
+                self._in_flight_requests[endpoint] = max(
+                    0, self._in_flight_requests[endpoint] - 1
+                )
+
+    def get_stats(self) -> dict[str, int]:
+        """Get current in-flight request counts.
+
+        Returns:
+            Mapping of endpoint URLs to in-flight request counts
+        """
+        return dict(self._in_flight_requests)

tetra_rp/runtime/manifest_fetcher.py

@@ -0,0 +1,192 @@
+"""Manifest fetcher with RunPod GQL integration and caching.
+
+This module provides manifest fetching from RunPod GraphQL API (source of truth)
+with local file caching and fallback.
+"""
+
+import asyncio
+import json
+import logging
+import time
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from .config import DEFAULT_CACHE_TTL
+from .generic_handler import load_manifest
+
+logger = logging.getLogger(__name__)
+
+
+class ManifestFetcher:
+    """Fetches and caches manifest from RunPod GraphQL API.
+
+    RunPod's GraphQL API is the source of truth for manifest data. This
+    fetcher pulls from it using RunpodGraphQLClient, caches locally, and
+    falls back to local file if RunPod API is unavailable.
+    """
+
+    def __init__(
+        self,
+        cache_ttl: int = DEFAULT_CACHE_TTL,
+        manifest_path: Optional[Path] = None,
+    ):
+        """Initialize manifest fetcher.
+
+        Args:
+            cache_ttl: Cache time-to-live in seconds (default: 300)
+            manifest_path: Optional path to local manifest file
+        """
+        self.cache_ttl = cache_ttl
+        self.manifest_path = manifest_path
+
+        # Cache state
+        self._cached_manifest: Optional[Dict[str, Any]] = None
+        self._cache_loaded_at: float = 0
+        self._cache_lock = asyncio.Lock()
+
+    async def get_manifest(
+        self,
+        mothership_id: Optional[str] = None,
+    ) -> Optional[Dict[str, Any]]:
+        """Get manifest from cache or fetch from RunPod GraphQL API.
+
+        Flow:
+        1. Check if cached and not expired → return cached
+        2. If expired/not cached → fetch from RunPod GraphQL API
+        3. Update local flash_manifest.json with fetched data
+        4. Cache the result
+        5. Return manifest
+
+        If RunPod GQL fetch fails, falls back to local file.
+
+        Args:
+            mothership_id: Optional mothership endpoint ID for tracking
+
+        Returns:
+            Manifest dictionary or None if unavailable
+        """
+        async with self._cache_lock:
+            now = time.time()
+            cache_age = now - self._cache_loaded_at
+
+            # Return cached if still valid
+            if self._cached_manifest and cache_age < self.cache_ttl:
+                logger.debug(
+                    f"Serving cached manifest (age: {cache_age:.1f}s, "
+                    f"TTL: {self.cache_ttl}s)"
+                )
+                return self._cached_manifest
+
+            # Cache expired or not loaded - fetch from RunPod GQL
+            logger.debug("Cache expired or empty, fetching from RunPod GraphQL API")
+
+            try:
+                # Fetch from RunPod GraphQL API (placeholder)
+                manifest = await self._fetch_from_gql(mothership_id)
+
+                # Update local flash_manifest.json
+                if manifest:
+                    self._update_local_file(manifest)
+
+                    # Update cache
+                    self._cached_manifest = manifest
+                    self._cache_loaded_at = now
+
+                    logger.info(
+                        f"Manifest fetched from RunPod GQL and cached "
+                        f"({len(manifest.get('resources', {}))} resources)"
+                    )
+                    return manifest
+
+            except NotImplementedError:
+                logger.debug(
+                    "RunPod GQL fetch not implemented, falling back to local file"
+                )
+            except Exception as e:
+                logger.warning(
+                    f"RunPod GQL fetch failed: {e}, falling back to local file"
+                )
+
+            # Fallback: load from local file
+            manifest = load_manifest(self.manifest_path)
+            if manifest:
+                # Cache the fallback manifest
+                self._cached_manifest = manifest
+                self._cache_loaded_at = now
+                logger.debug("Loaded and cached manifest from local file")
+
+            return manifest
+
+    async def _fetch_from_gql(
+        self,
+        mothership_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Fetch manifest from RunPod GraphQL API.
+
+        TBD: Future implementation will query RunPod's GraphQL API
+        to retrieve the manifest configuration.
+
+        Args:
+            mothership_id: Optional mothership endpoint ID
+
+        Returns:
+            Manifest dictionary from RunPod GQL
+
+        Raises:
+            NotImplementedError: Placeholder for future implementation
+
+        Note:
+            Future implementation will use RunpodGraphQLClient:
+
+            ```python
+            async with RunpodGraphQLClient() as client:
+                query = '''
+                query GetManifest($mothershipId: ID!) {
+                    getManifest(mothershipId: $mothershipId) {
+                        version
+                        projectName
+                        generatedAt
+                        resources
+                        functionRegistry
+                    }
+                }
+                '''
+                result = await client.execute(query, {"mothershipId": mothership_id})
+                return result["data"]["getManifest"]
+            ```
+        """
+        raise NotImplementedError(
+            "RunPod manifest query not yet implemented. "
+            "Falling back to local flash_manifest.json file."
+        )
+
+    def _update_local_file(self, manifest: Dict[str, Any]) -> None:
+        """Update local flash_manifest.json with fetched data.
+
+        Args:
+            manifest: Manifest dictionary to write
+        """
+        try:
+            # Determine file path
+            if self.manifest_path:
+                file_path = self.manifest_path
+            else:
+                file_path = Path.cwd() / "flash_manifest.json"
+
+            # Write manifest to file
+            with open(file_path, "w") as f:
+                json.dump(manifest, f, indent=2)
+
+            logger.debug(f"Updated local manifest file: {file_path}")
+
+        except Exception as e:
+            logger.warning(f"Failed to update local manifest file: {e}")
+            # Non-critical error - cached manifest still valid
+
+    def invalidate_cache(self) -> None:
+        """Manually invalidate the cache.
+
+        Next get_manifest() call will fetch from GQL.
+        """
+        self._cache_loaded_at = 0
+        logger.debug("Manifest cache invalidated")