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

tetra_rp/runtime/service_registry.py

@@ -0,0 +1,346 @@
+"""Runtime service registry for cross-endpoint function routing."""
+
+import asyncio
+import json
+import logging
+import os
+import time
+from pathlib import Path
+from typing import Dict, Optional
+from urllib.parse import urlparse
+
+from tetra_rp.core.resources.serverless import ServerlessResource
+
+from .config import DEFAULT_CACHE_TTL
+from .state_manager_client import StateManagerClient, ManifestServiceUnavailableError
+from .models import Manifest
+
+logger = logging.getLogger(__name__)
+
+
+class ServiceRegistry:
+    """Service discovery and routing for cross-endpoint function calls.
+
+    Loads manifest to map functions to resource configs, queries mothership
+    manifest for endpoint URLs, and determines if function calls are local
+    or remote.
+    """
+
+    def __init__(
+        self,
+        manifest_path: Optional[Path] = None,
+        cache_ttl: int = DEFAULT_CACHE_TTL,
+    ):
+        """Initialize service registry with peer-to-peer State Manager access.
+
+        All endpoints query State Manager directly for manifest updates.
+        No Mothership dependency - all endpoints are equal peers.
+
+        Args:
+            manifest_path: Path to flash_manifest.json. Defaults to
+                FLASH_MANIFEST_PATH env var or auto-detection.
+            cache_ttl: Manifest cache lifetime in seconds (default: 300).
+
+        Environment Variables:
+            FLASH_RESOURCE_NAME: Resource config name for this endpoint.
+                Identifies which resource config this endpoint represents.
+            RUNPOD_ENDPOINT_ID: Endpoint ID (used for State Manager queries and fallback).
+            RUNPOD_API_KEY: API key for State Manager GraphQL access.
+
+        Raises:
+            FileNotFoundError: If manifest_path doesn't exist.
+        """
+        self.cache_ttl = cache_ttl
+        self._endpoint_registry: Dict[str, str] = {}
+        self._endpoint_registry_loaded_at = 0.0
+        self._manifest: Manifest = Manifest(
+            version="1.0",
+            generated_at="",
+            project_name="",
+            function_registry={},
+            resources={},
+        )
+        self._endpoint_registry_lock = asyncio.Lock()
+
+        # Load manifest
+        self._load_manifest(manifest_path)
+
+        # Peer-to-peer: All endpoints use StateManagerClient directly
+        try:
+            self._manifest_client = StateManagerClient()
+        except Exception as e:
+            logger.warning(f"Failed to initialize State Manager client: {e}")
+            self._manifest_client = None
+
+        # Current endpoint identification for local vs remote detection
+        self._current_endpoint = os.getenv("FLASH_RESOURCE_NAME") or os.getenv(
+            "RUNPOD_ENDPOINT_ID"
+        )
+
+    def _load_manifest(self, manifest_path: Optional[Path]) -> None:
+        """Load flash_manifest.json.
+
+        Args:
+            manifest_path: Explicit path to manifest. Tries env var and
+                auto-detection if not provided.
+
+        Raises:
+            FileNotFoundError: If manifest not found.
+        """
+        paths_to_try = []
+
+        # Explicit path
+        if manifest_path:
+            paths_to_try.append(manifest_path)
+
+        # Environment variable
+        env_path = os.getenv("FLASH_MANIFEST_PATH")
+        if env_path:
+            paths_to_try.append(Path(env_path))
+
+        # Auto-detection: same directory as this file, or cwd
+        paths_to_try.extend(
+            [
+                Path(__file__).parent.parent.parent / "flash_manifest.json",
+                Path.cwd() / "flash_manifest.json",
+            ]
+        )
+
+        # Try each path
+        for path in paths_to_try:
+            if path and path.exists():
+                try:
+                    with open(path) as f:
+                        manifest_dict = json.load(f)
+                    self._manifest = Manifest.from_dict(manifest_dict)
+                    logger.debug(f"Manifest loaded from {path}")
+                    return
+                except Exception as e:
+                    logger.warning(f"Failed to load manifest from {path}: {e}")
+                    continue
+
+        # No manifest found - log warning but don't fail
+        logger.warning(
+            "flash_manifest.json not found. Cross-endpoint routing disabled. "
+            "Manifest is required for routing functions between endpoints."
+        )
+        self._manifest = Manifest(
+            version="1.0",
+            generated_at="",
+            project_name="",
+            function_registry={},
+            resources={},
+        )
+
+    async def _ensure_manifest_loaded(self) -> None:
+        """Load manifest from State Manager if cache expired or not loaded.
+
+        Peer-to-Peer Architecture:
+            Each endpoint queries State Manager independently using its own
+            RUNPOD_ENDPOINT_ID. No mothership dependency - all endpoints
+            are equal peers discovering each other through the manifest.
+
+        Query Flow:
+            1. get_flash_environment(RUNPOD_ENDPOINT_ID) → activeBuildId
+            2. get_flash_build(activeBuildId) → manifest
+            3. Extract manifest["resources_endpoints"] mapping
+            4. Cache for 300s (DEFAULT_CACHE_TTL)
+
+        State Manager Consistency:
+            - CLI updates manifest after provisioning all endpoints
+            - Endpoints cache manifest to reduce API calls
+            - TTL ensures eventual consistency (300s by default)
+
+        Returns:
+            None. Updates self._endpoint_registry internally.
+        """
+        async with self._endpoint_registry_lock:
+            now = time.time()
+            cache_age = now - self._endpoint_registry_loaded_at
+
+            if cache_age > self.cache_ttl:
+                if self._manifest_client is None:
+                    logger.debug("State Manager client not available, skipping refresh")
+                    return
+
+                try:
+                    mothership_id = os.getenv("RUNPOD_ENDPOINT_ID")
+                    if not mothership_id:
+                        logger.warning(
+                            "RUNPOD_ENDPOINT_ID not set, cannot query State Manager"
+                        )
+                        return
+
+                    # Query State Manager directly for full manifest
+                    full_manifest = await self._manifest_client.get_persisted_manifest(
+                        mothership_id
+                    )
+
+                    # Extract resources_endpoints mapping
+                    resources_endpoints = full_manifest.get("resources_endpoints", {})
+
+                    self._endpoint_registry = resources_endpoints
+                    self._endpoint_registry_loaded_at = now
+                    logger.debug(
+                        f"Manifest loaded from State Manager: {len(self._endpoint_registry)} endpoints, "
+                        f"cache TTL {self.cache_ttl}s"
+                    )
+                except ManifestServiceUnavailableError as e:
+                    logger.warning(
+                        f"Failed to load manifest from State Manager: {e}. "
+                        f"Cross-endpoint routing unavailable."
+                    )
+                    self._endpoint_registry = {}
+
+    async def get_endpoint_for_function(self, function_name: str) -> Optional[str]:
+        """Get endpoint URL for a function.
+
+        Determines if function is local (same endpoint) or remote (different
+        endpoint), returning None for local and URL for remote.
+
+        Queries State Manager if endpoint registry cache is expired.
+
+        Args:
+            function_name: Name of the function to route.
+
+        Returns:
+            Endpoint URL if function is remote, None if local.
+
+        Raises:
+            ValueError: If function not in manifest.
+        """
+        # Ensure manifest is loaded from State Manager (with caching)
+        await self._ensure_manifest_loaded()
+
+        function_registry = self._manifest.function_registry
+
+        if function_name not in function_registry:
+            raise ValueError(
+                f"Function '{function_name}' not found in manifest. "
+                f"Available functions: {list(function_registry.keys())}"
+            )
+
+        resource_config_name = function_registry[function_name]
+
+        # Check if this is the current endpoint (local)
+        if resource_config_name == self._current_endpoint:
+            return None
+
+        # Check manifest for remote endpoint URL
+        endpoint_url = self._endpoint_registry.get(resource_config_name)
+        if not endpoint_url:
+            logger.debug(
+                f"Endpoint URL for '{resource_config_name}' not in manifest. "
+                f"Manifest has: {list(self._endpoint_registry.keys())}"
+            )
+
+        return endpoint_url
+
+    async def get_resource_for_function(
+        self, function_name: str
+    ) -> Optional[ServerlessResource]:
+        """Get ServerlessResource for a function.
+
+        Creates a ServerlessResource with the correct endpoint ID if the function
+        is remote, returns None if local.
+
+        Args:
+            function_name: Name of the function to route.
+
+        Returns:
+            ServerlessResource with ID set if function is remote
+            None if function runs on current endpoint
+
+        Raises:
+            ValueError: If function not in manifest.
+        """
+        endpoint_url = await self.get_endpoint_for_function(function_name)
+
+        if endpoint_url is None:
+            return None  # Local function
+
+        # Extract endpoint ID from URL (format: https://{endpoint_base_url}/v2/{endpoint_id})
+        try:
+            parsed = urlparse(endpoint_url)
+            # Get the last path component (the endpoint ID)
+            path_parts = parsed.path.rstrip("/").split("/")
+            endpoint_id = path_parts[-1] if path_parts else ""
+
+            if not endpoint_id:
+                raise ValueError(
+                    f"Invalid endpoint URL format: {endpoint_url} - no endpoint ID found"
+                )
+        except Exception as e:
+            raise ValueError(
+                f"Failed to parse endpoint URL '{endpoint_url}': {e}"
+            ) from e
+
+        # Create and return ServerlessResource
+        resource = ServerlessResource(name=f"remote_{function_name}")
+        resource.id = endpoint_id
+
+        return resource
+
+    async def is_local_function(self, function_name: str) -> bool:
+        """Check if function executes on current endpoint.
+
+        Args:
+            function_name: Name of the function.
+
+        Returns:
+            True if function is local, False if remote or not found.
+        """
+        try:
+            endpoint_url = await self.get_endpoint_for_function(function_name)
+            return endpoint_url is None
+        except ValueError:
+            # Function not in manifest, assume local (will execute and fail)
+            return True
+
+    def get_current_endpoint_id(self) -> Optional[str]:
+        """Get ID of current endpoint from environment.
+
+        Returns:
+            Endpoint ID from FLASH_RESOURCE_NAME or RUNPOD_ENDPOINT_ID, or None if not set.
+        """
+        return self._current_endpoint
+
+    def refresh_manifest(self) -> None:
+        """Force refresh manifest from mothership on next access."""
+        self._endpoint_registry_loaded_at = 0
+
+    def get_manifest(self) -> Manifest:
+        """Get loaded manifest.
+
+        Returns:
+            Loaded Manifest object.
+        """
+        return self._manifest
+
+    def get_all_resources(self) -> Dict[str, Dict]:
+        """Get all resource configs from manifest.
+
+        Returns:
+            Dictionary of resource configs as dictionaries.
+        """
+        from dataclasses import asdict
+
+        return {
+            name: asdict(config) for name, config in self._manifest.resources.items()
+        }
+
+    def get_resource_functions(self, resource_name: str) -> list:
+        """Get list of functions for a resource.
+
+        Args:
+            resource_name: Name of the resource config.
+
+        Returns:
+            List of function metadata dictionaries.
+        """
+        resource = self._manifest.resources.get(resource_name)
+        if not resource:
+            return []
+        from dataclasses import asdict
+
+        return [asdict(func) for func in resource.functions]

tetra_rp/runtime/state_manager_client.py

@@ -0,0 +1,248 @@
+"""GraphQL client for State Manager API to persist and reconcile manifests."""
+
+import asyncio
+import logging
+from typing import Any, Dict, Optional
+
+from tetra_rp.core.api.runpod import RunpodGraphQLClient
+
+from .config import DEFAULT_MAX_RETRIES
+from .exceptions import GraphQLError, ManifestServiceUnavailableError
+
+logger = logging.getLogger(__name__)
+
+
+class StateManagerClient:
+    """GraphQL client for State Manager manifest persistence.
+
+    The State Manager persists manifest state via RunPod GraphQL API,
+    providing reconciliation capabilities for the mothership to track
+    deployed resources across boots.
+
+    Thread Safety:
+        Uses asyncio.Lock to serialize read-modify-write operations,
+        preventing race conditions during concurrent resource updates.
+
+    Architecture:
+        Manifest updates follow a read-modify-write pattern:
+        1. Fetch environment -> activeBuildId
+        2. Fetch build -> manifest
+        3. Merge changes into manifest
+        4. Call updateFlashBuildManifest mutation
+
+    Performance:
+        Each update requires 3 GraphQL roundtrips. Consider batching
+        updates when provisioning multiple resources.
+    """
+
+    def __init__(
+        self,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+    ):
+        """Initialize State Manager client.
+
+        Args:
+            max_retries: Maximum retry attempts for operations.
+
+        Raises:
+            RunpodAPIKeyError: If RUNPOD_API_KEY environment variable is not set (raised by RunpodGraphQLClient).
+        """
+        self.max_retries = max_retries
+        self._manifest_lock = asyncio.Lock()
+
+    async def get_persisted_manifest(
+        self, mothership_id: str
+    ) -> Optional[Dict[str, Any]]:
+        """Fetch persisted manifest from State Manager.
+
+        Args:
+            mothership_id: ID of the mothership endpoint.
+
+        Returns:
+            Manifest dict.
+
+        Raises:
+            ManifestServiceUnavailableError: If State Manager unavailable after retries.
+        """
+        last_exception: Optional[Exception] = None
+
+        for attempt in range(self.max_retries):
+            try:
+                async with RunpodGraphQLClient() as client:
+                    _, manifest = await self._fetch_build_and_manifest(
+                        client, mothership_id
+                    )
+
+                logger.debug(f"Persisted manifest loaded for {mothership_id}")
+                return manifest
+
+            except (
+                asyncio.TimeoutError,
+                ManifestServiceUnavailableError,
+                GraphQLError,
+                ConnectionError,
+            ) as e:
+                last_exception = e
+                if attempt < self.max_retries - 1:
+                    backoff = 2**attempt
+                    logger.warning(
+                        f"State Manager request failed (attempt {attempt + 1}): {e}, "
+                        f"retrying in {backoff}s..."
+                    )
+                    await asyncio.sleep(backoff)
+                    continue
+
+        raise ManifestServiceUnavailableError(
+            f"Failed to fetch persisted manifest after {self.max_retries} attempts: "
+            f"{last_exception}"
+        )
+
+    async def update_resource_state(
+        self,
+        mothership_id: str,
+        resource_name: str,
+        resource_data: Dict[str, Any],
+    ) -> None:
+        """Update single resource entry in State Manager.
+
+        Uses locking to prevent race conditions when multiple resources
+        are deployed concurrently.
+
+        Args:
+            mothership_id: ID of the mothership endpoint.
+            resource_name: Name of the resource.
+            resource_data: Resource metadata (config_hash, endpoint_url, status, etc).
+
+        Raises:
+            ManifestServiceUnavailableError: If State Manager unavailable.
+        """
+        last_exception: Optional[Exception] = None
+
+        for attempt in range(self.max_retries):
+            try:
+                async with self._manifest_lock:
+                    async with RunpodGraphQLClient() as client:
+                        build_id, manifest = await self._fetch_build_and_manifest(
+                            client, mothership_id
+                        )
+                        resources = manifest.setdefault("resources", {})
+                        existing = resources.get(resource_name)
+                        if not isinstance(existing, dict):
+                            existing = {}
+                        resources[resource_name] = {**existing, **resource_data}
+                        await client.update_build_manifest(build_id, manifest)
+
+                logger.debug(
+                    f"Updated resource state in State Manager: {mothership_id}/{resource_name}"
+                )
+                return
+
+            except (
+                asyncio.TimeoutError,
+                ManifestServiceUnavailableError,
+                GraphQLError,
+                ConnectionError,
+            ) as e:
+                last_exception = e
+                if attempt < self.max_retries - 1:
+                    backoff = 2**attempt
+                    logger.warning(
+                        f"State Manager request failed (attempt {attempt + 1}): {e}, "
+                        f"retrying in {backoff}s..."
+                    )
+                    await asyncio.sleep(backoff)
+                    continue
+
+        raise ManifestServiceUnavailableError(
+            f"Failed to update resource state after {self.max_retries} attempts: "
+            f"{last_exception}"
+        )
+
+    async def remove_resource_state(
+        self, mothership_id: str, resource_name: str
+    ) -> None:
+        """Remove resource entry from State Manager.
+
+        Uses locking to prevent race conditions when multiple resources
+        are deployed concurrently.
+
+        Args:
+            mothership_id: ID of the mothership endpoint.
+            resource_name: Name of the resource.
+
+        Raises:
+            ManifestServiceUnavailableError: If State Manager unavailable.
+        """
+        last_exception: Optional[Exception] = None
+
+        for attempt in range(self.max_retries):
+            try:
+                async with self._manifest_lock:
+                    async with RunpodGraphQLClient() as client:
+                        build_id, manifest = await self._fetch_build_and_manifest(
+                            client, mothership_id
+                        )
+                        resources = manifest.setdefault("resources", {})
+                        resources.pop(resource_name, None)
+                        await client.update_build_manifest(build_id, manifest)
+
+                logger.debug(
+                    f"Removed resource state from State Manager: {mothership_id}/{resource_name}"
+                )
+                return
+
+            except (
+                asyncio.TimeoutError,
+                ManifestServiceUnavailableError,
+                GraphQLError,
+                ConnectionError,
+            ) as e:
+                last_exception = e
+                if attempt < self.max_retries - 1:
+                    backoff = 2**attempt
+                    logger.warning(
+                        f"State Manager request failed (attempt {attempt + 1}): {e}, "
+                        f"retrying in {backoff}s..."
+                    )
+                    await asyncio.sleep(backoff)
+                    continue
+
+        raise ManifestServiceUnavailableError(
+            f"Failed to remove resource state after {self.max_retries} attempts: "
+            f"{last_exception}"
+        )
+
+    async def _fetch_build_and_manifest(
+        self, client: RunpodGraphQLClient, mothership_id: str
+    ) -> tuple[str, Dict[str, Any]]:
+        """Fetch active build ID and manifest for an environment.
+
+        Args:
+            client: Authenticated GraphQL client.
+            mothership_id: Flash environment ID.
+
+        Returns:
+            Tuple of (build_id, manifest_dict).
+
+        Raises:
+            ManifestServiceUnavailableError: If environment, build, or manifest not found.
+        """
+        environment = await client.get_flash_environment(
+            {"flashEnvironmentId": mothership_id}
+        )
+        build_id = environment.get("activeBuildId")
+        if not build_id:
+            raise ManifestServiceUnavailableError(
+                f"Active build not found for environment {mothership_id}. "
+                f"Environment may not be fully initialized or has no deployed build."
+            )
+
+        build = await client.get_flash_build(build_id)
+        manifest = build.get("manifest")
+        if not manifest:
+            raise ManifestServiceUnavailableError(
+                f"Manifest not found for build {build.get('id', build_id)}. "
+                f"Build may be corrupted, not yet published, or manifest was not generated."
+            )
+
+        return build_id, manifest

tetra_rp/stubs/live_serverless.py

@@ -4,6 +4,7 @@ import inspect
 import textwrap
 import hashlib
 import traceback
+import threading
 import cloudpickle
 import logging
 from ..core.resources import LiveServerless
@@ -12,26 +13,37 @@ from ..protos.remote_execution import (
     FunctionResponse,
     RemoteExecutorStub,
 )
+from ..runtime.serialization import serialize_args, serialize_kwargs
 
 log = logging.getLogger(__name__)
 
 
-# global in memory cache, TODO: use a more robust cache in future
+# Global in-memory cache with thread safety
 _SERIALIZED_FUNCTION_CACHE = {}
+_function_cache_lock = threading.RLock()
 
 
 def get_function_source(func):
     """Extract the function source code without the decorator."""
+    # Unwrap any decorators to get the original function
+    func = inspect.unwrap(func)
+
     # Get the source code of the decorated function
     source = inspect.getsource(func)
 
+    # Dedent the source to handle functions defined in classes or indented contexts
+    source = textwrap.dedent(source)
+
     # Parse the source code
     module = ast.parse(source)
 
-    # Find the function definition node
+    # Find the function definition node (both sync and async)
     function_def = None
     for node in ast.walk(module):
-        if isinstance(node, ast.FunctionDef) and node.name == func.__name__:
+        if (
+            isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
+            and node.name == func.__name__
+        ):
             function_def = node
             break
 
@@ -60,32 +72,38 @@ class LiveServerlessStub(RemoteExecutorStub):
     def __init__(self, server: LiveServerless):
         self.server = server
 
-    def prepare_request(self, func, dependencies, system_dependencies, *args, **kwargs):
+    def prepare_request(
+        self,
+        func,
+        dependencies,
+        system_dependencies,
+        accelerate_downloads,
+        *args,
+        **kwargs,
+    ):
         source, src_hash = get_function_source(func)
 
         request = {
             "function_name": func.__name__,
             "dependencies": dependencies,
             "system_dependencies": system_dependencies,
+            "accelerate_downloads": accelerate_downloads,
         }
 
-        # check if the function is already cached
-        if src_hash not in _SERIALIZED_FUNCTION_CACHE:
-            # Cache the serialized function
-            _SERIALIZED_FUNCTION_CACHE[src_hash] = source
+        # Thread-safe cache access
+        with _function_cache_lock:
+            # check if the function is already cached
+            if src_hash not in _SERIALIZED_FUNCTION_CACHE:
+                # Cache the serialized function
+                _SERIALIZED_FUNCTION_CACHE[src_hash] = source
 
-        request["function_code"] = _SERIALIZED_FUNCTION_CACHE[src_hash]
+            request["function_code"] = _SERIALIZED_FUNCTION_CACHE[src_hash]
 
         # Serialize arguments using cloudpickle
         if args:
-            request["args"] = [
-                base64.b64encode(cloudpickle.dumps(arg)).decode("utf-8") for arg in args
-            ]
+            request["args"] = serialize_args(args)
         if kwargs:
-            request["kwargs"] = {
-                k: base64.b64encode(cloudpickle.dumps(v)).decode("utf-8")
-                for k, v in kwargs.items()
-            }
+            request["kwargs"] = serialize_kwargs(kwargs)
 
         return FunctionRequest(**request)
 
@@ -95,7 +113,7 @@ class LiveServerlessStub(RemoteExecutorStub):
 
         if response.stdout:
             for line in response.stdout.splitlines():
-                log.info(f"Remote | {line}")
+                print(line)
 
         if response.success:
             if response.result is None: