tetra-rp 0.17.1__py3-none-any.whl

tetra_rp/execute_class.py

@@ -0,0 +1,319 @@
+"""
+Class execution module for remote class instantiation and method calls.
+
+This module provides functionality to create and execute remote class instances,
+with automatic caching of class serialization data to improve performance and
+prevent memory leaks through LRU eviction.
+"""
+
+import base64
+import hashlib
+import inspect
+import logging
+import textwrap
+import uuid
+from typing import List, Optional, Type
+
+import cloudpickle
+
+from .core.resources import ResourceManager, ServerlessResource
+from .core.utils.constants import HASH_TRUNCATE_LENGTH, UUID_FALLBACK_LENGTH
+from .core.utils.lru_cache import LRUCache
+from .protos.remote_execution import FunctionRequest
+from .stubs import stub_resource
+
+log = logging.getLogger(__name__)
+
+# Global in-memory cache for serialized class data with LRU eviction
+_SERIALIZED_CLASS_CACHE = LRUCache(max_size=1000)
+
+
+def serialize_constructor_args(args, kwargs):
+    """Serialize constructor arguments for caching."""
+    serialized_args = [
+        base64.b64encode(cloudpickle.dumps(arg)).decode("utf-8") for arg in args
+    ]
+    serialized_kwargs = {
+        k: base64.b64encode(cloudpickle.dumps(v)).decode("utf-8")
+        for k, v in kwargs.items()
+    }
+    return serialized_args, serialized_kwargs
+
+
+def get_or_cache_class_data(
+    cls: Type, args: tuple, kwargs: dict, cache_key: str
+) -> str:
+    """Get class code from cache or extract and cache it."""
+    if cache_key not in _SERIALIZED_CLASS_CACHE:
+        # Cache miss - extract and cache class code
+        clean_class_code = extract_class_code_simple(cls)
+
+        try:
+            serialized_args, serialized_kwargs = serialize_constructor_args(
+                args, kwargs
+            )
+
+            # Cache the serialized data
+            _SERIALIZED_CLASS_CACHE.set(
+                cache_key,
+                {
+                    "class_code": clean_class_code,
+                    "constructor_args": serialized_args,
+                    "constructor_kwargs": serialized_kwargs,
+                },
+            )
+
+            log.debug(f"Cached class data for {cls.__name__} with key: {cache_key}")
+
+        except (TypeError, AttributeError, OSError) as e:
+            log.warning(
+                f"Could not serialize constructor arguments for {cls.__name__}: {e}"
+            )
+            log.warning(
+                f"Skipping constructor argument caching for {cls.__name__} due to unserializable arguments"
+            )
+
+            # Store minimal cache entry to avoid repeated attempts
+            _SERIALIZED_CLASS_CACHE.set(
+                cache_key,
+                {
+                    "class_code": clean_class_code,
+                    "constructor_args": None,  # Signal that args couldn't be cached
+                    "constructor_kwargs": None,
+                },
+            )
+
+        return clean_class_code
+    else:
+        # Cache hit - retrieve cached data
+        cached_data = _SERIALIZED_CLASS_CACHE.get(cache_key)
+        log.debug(
+            f"Retrieved cached class data for {cls.__name__} with key: {cache_key}"
+        )
+        return cached_data["class_code"]
+
+
+def extract_class_code_simple(cls: Type) -> str:
+    """Extract clean class code without decorators and proper indentation"""
+    try:
+        # Get source code
+        source = inspect.getsource(cls)
+
+        # Split into lines
+        lines = source.split("\n")
+
+        # Find the class definition line (starts with 'class' and contains ':')
+        class_start_idx = -1
+        for i, line in enumerate(lines):
+            stripped = line.strip()
+            if stripped.startswith("class ") and ":" in stripped:
+                class_start_idx = i
+                break
+
+        if class_start_idx == -1:
+            raise ValueError("Could not find class definition")
+
+        # Take lines from class definition onwards (ignore everything before)
+        class_lines = lines[class_start_idx:]
+
+        # Remove empty lines at the end
+        while class_lines and not class_lines[-1].strip():
+            class_lines.pop()
+
+        # Join back and dedent to remove any leading indentation
+        class_code = "\n".join(class_lines)
+        class_code = textwrap.dedent(class_code)
+
+        # Validate the code by trying to compile it
+        compile(class_code, "<string>", "exec")
+
+        log.debug(f"Successfully extracted class code for {cls.__name__}")
+        return class_code
+
+    except Exception as e:
+        log.warning(f"Could not extract class code for {cls.__name__}: {e}")
+        log.warning("Falling back to basic class structure")
+
+        # Enhanced fallback: try to preserve method signatures
+        fallback_methods = []
+        for name, method in inspect.getmembers(cls, predicate=inspect.isfunction):
+            try:
+                sig = inspect.signature(method)
+                fallback_methods.append(f"    def {name}{sig}:")
+                fallback_methods.append("        pass")
+                fallback_methods.append("")
+            except (TypeError, ValueError, OSError) as e:
+                log.warning(f"Could not extract method signature for {name}: {e}")
+                fallback_methods.append(f"    def {name}(self, *args, **kwargs):")
+                fallback_methods.append("        pass")
+                fallback_methods.append("")
+
+        fallback_code = f"""class {cls.__name__}:
+    def __init__(self, *args, **kwargs):
+        pass
+
+{chr(10).join(fallback_methods)}"""
+
+        return fallback_code
+
+
+def get_class_cache_key(
+    cls: Type, constructor_args: tuple, constructor_kwargs: dict
+) -> str:
+    """Generate a cache key for class serialization based on class source and constructor args.
+
+    Args:
+        cls: The class type to generate a key for
+        constructor_args: Positional arguments passed to class constructor
+        constructor_kwargs: Keyword arguments passed to class constructor
+
+    Returns:
+        A unique cache key string, or a UUID-based fallback if serialization fails
+
+    Note:
+        Falls back to UUID-based key if constructor arguments cannot be serialized,
+        which disables caching benefits but maintains functionality.
+    """
+    try:
+        # Get class source code for hashing
+        class_source = extract_class_code_simple(cls)
+
+        # Create hash of class source
+        class_hash = hashlib.sha256(class_source.encode()).hexdigest()
+
+        # Create hash of constructor arguments
+        args_data = cloudpickle.dumps((constructor_args, constructor_kwargs))
+        args_hash = hashlib.sha256(args_data).hexdigest()
+
+        # Combine hashes for final cache key
+        cache_key = f"{cls.__name__}_{class_hash[:HASH_TRUNCATE_LENGTH]}_{args_hash[:HASH_TRUNCATE_LENGTH]}"
+
+        log.debug(f"Generated cache key for {cls.__name__}: {cache_key}")
+        return cache_key
+
+    except (TypeError, AttributeError, OSError) as e:
+        log.warning(f"Could not generate cache key for {cls.__name__}: {e}")
+        # Fallback to basic key without caching benefits
+        return f"{cls.__name__}_{uuid.uuid4().hex[:UUID_FALLBACK_LENGTH]}"
+
+
+def create_remote_class(
+    cls: Type,
+    resource_config: ServerlessResource,
+    dependencies: Optional[List[str]],
+    system_dependencies: Optional[List[str]],
+    accelerate_downloads: bool,
+    extra: dict,
+):
+    """
+    Create a remote class wrapper.
+    """
+    # Validate inputs
+    if not inspect.isclass(cls):
+        raise TypeError(f"Expected a class, got {type(cls).__name__}")
+    if not hasattr(cls, "__name__"):
+        raise ValueError("Class must have a __name__ attribute")
+
+    class RemoteClassWrapper:
+        def __init__(self, *args, **kwargs):
+            self._class_type = cls
+            self._resource_config = resource_config
+            self._dependencies = dependencies or []
+            self._system_dependencies = system_dependencies or []
+            self._accelerate_downloads = accelerate_downloads
+            self._extra = extra
+            self._constructor_args = args
+            self._constructor_kwargs = kwargs
+            self._instance_id = (
+                f"{cls.__name__}_{uuid.uuid4().hex[:UUID_FALLBACK_LENGTH]}"
+            )
+            self._initialized = False
+
+            # Generate cache key and get class code
+            self._cache_key = get_class_cache_key(cls, args, kwargs)
+            self._clean_class_code = get_or_cache_class_data(
+                cls, args, kwargs, self._cache_key
+            )
+
+            log.debug(f"Created remote class wrapper for {cls.__name__}")
+
+        async def _ensure_initialized(self):
+            """Ensure the remote instance is created."""
+            if self._initialized:
+                return
+
+            # Get remote resource
+            resource_manager = ResourceManager()
+            remote_resource = await resource_manager.get_or_deploy_resource(
+                self._resource_config
+            )
+            self._stub = stub_resource(remote_resource, **self._extra)
+
+            # Create the remote instance by calling a method (which will trigger instance creation)
+            # We'll do this on first method call
+            self._initialized = True
+
+        def __getattr__(self, name):
+            """Dynamically create method proxies for all class methods."""
+            if name.startswith("_"):
+                raise AttributeError(
+                    f"'{self.__class__.__name__}' object has no attribute '{name}'"
+                )
+
+            async def method_proxy(*args, **kwargs):
+                await self._ensure_initialized()
+
+                # Get cached data
+                cached_data = _SERIALIZED_CLASS_CACHE.get(self._cache_key)
+
+                # Serialize method arguments (these change per call, so no caching)
+                method_args = [
+                    base64.b64encode(cloudpickle.dumps(arg)).decode("utf-8")
+                    for arg in args
+                ]
+                method_kwargs = {
+                    k: base64.b64encode(cloudpickle.dumps(v)).decode("utf-8")
+                    for k, v in kwargs.items()
+                }
+
+                # Handle constructor args - use cached if available, else serialize fresh
+                if cached_data["constructor_args"] is not None:
+                    # Use cached constructor args
+                    constructor_args = cached_data["constructor_args"]
+                    constructor_kwargs = cached_data["constructor_kwargs"]
+                else:
+                    # Constructor args couldn't be cached due to serialization issues
+                    # Serialize them fresh for each method call (fallback behavior)
+                    constructor_args = [
+                        base64.b64encode(cloudpickle.dumps(arg)).decode("utf-8")
+                        for arg in self._constructor_args
+                    ]
+                    constructor_kwargs = {
+                        k: base64.b64encode(cloudpickle.dumps(v)).decode("utf-8")
+                        for k, v in self._constructor_kwargs.items()
+                    }
+
+                request = FunctionRequest(
+                    execution_type="class",
+                    class_name=self._class_type.__name__,
+                    class_code=cached_data["class_code"],
+                    method_name=name,
+                    args=method_args,
+                    kwargs=method_kwargs,
+                    constructor_args=constructor_args,
+                    constructor_kwargs=constructor_kwargs,
+                    dependencies=self._dependencies,
+                    system_dependencies=self._system_dependencies,
+                    accelerate_downloads=self._accelerate_downloads,
+                    instance_id=self._instance_id,
+                    create_new_instance=not hasattr(
+                        self, "_stub"
+                    ),  # Create new only on first call
+                )
+
+                # Execute via stub
+                return await self._stub.execute_class_method(request)  # type: ignore
+
+            return method_proxy
+
+    return RemoteClassWrapper

tetra_rp/logger.py

@@ -0,0 +1,34 @@
+import logging
+import os
+import sys
+from typing import Union, Optional
+
+
+def setup_logging(
+    level: Union[int, str] = logging.INFO, stream=sys.stdout, fmt: Optional[str] = None
+):
+    """
+    Sets up the root logger with a stream handler and basic formatting.
+    Does nothing if handlers are already configured.
+    """
+    if isinstance(level, str):
+        level = getattr(logging, level.upper(), logging.INFO)
+
+    if fmt is None:
+        if level == logging.DEBUG:
+            fmt = "%(asctime)s | %(levelname)-5s | %(name)s | %(filename)s:%(lineno)d | %(message)s"
+        else:
+            # Default format for INFO level and above
+            fmt = "%(asctime)s | %(levelname)-5s | %(message)s"
+
+    root_logger = logging.getLogger()
+    if not root_logger.hasHandlers():
+        handler = logging.StreamHandler(stream)
+        handler.setFormatter(logging.Formatter(fmt))
+        root_logger.setLevel(level)
+        root_logger.addHandler(handler)
+
+    # Optionally allow log level override via env var
+    env_level = os.environ.get("LOG_LEVEL")
+    if env_level:
+        root_logger.setLevel(env_level.upper())

tetra_rp/protos/remote_execution.py

@@ -0,0 +1,148 @@
+"""Remote execution protocol definitions using Pydantic models.
+
+This module defines the request/response protocol for remote function and class execution.
+The models align with the protobuf schema for communication with remote workers.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class FunctionRequest(BaseModel):
+    """Request model for remote function or class execution.
+
+    Supports both function-based execution and class instantiation with method calls.
+    All serialized data (args, kwargs, etc.) are base64-encoded cloudpickle strings.
+    """
+
+    # MADE OPTIONAL - can be None for class-only execution
+    function_name: Optional[str] = Field(
+        default=None,
+        description="Name of the function to execute",
+    )
+    function_code: Optional[str] = Field(
+        default=None,
+        description="Source code of the function to execute",
+    )
+    args: List[str] = Field(
+        default_factory=list,
+        description="List of base64-encoded cloudpickle-serialized arguments",
+    )
+    kwargs: Dict[str, str] = Field(
+        default_factory=dict,
+        description="Dictionary of base64-encoded cloudpickle-serialized keyword arguments",
+    )
+    dependencies: Optional[List[str]] = Field(
+        default=None,
+        description="Optional list of pip packages to install before executing the function",
+    )
+    system_dependencies: Optional[List[str]] = Field(
+        default=None,
+        description="Optional list of system dependencies to install before executing the function",
+    )
+
+    # NEW FIELDS FOR CLASS SUPPORT
+    execution_type: str = Field(
+        default="function", description="Type of execution: 'function' or 'class'"
+    )
+    class_name: Optional[str] = Field(
+        default=None,
+        description="Name of the class to instantiate (for class execution)",
+    )
+    class_code: Optional[str] = Field(
+        default=None,
+        description="Source code of the class to instantiate (for class execution)",
+    )
+    constructor_args: List[str] = Field(
+        default_factory=list,
+        description="List of base64-encoded cloudpickle-serialized constructor arguments",
+    )
+    constructor_kwargs: Dict[str, str] = Field(
+        default_factory=dict,
+        description="Dictionary of base64-encoded cloudpickle-serialized constructor keyword arguments",
+    )
+    method_name: str = Field(
+        default="__call__",
+        description="Name of the method to call on the class instance",
+    )
+    instance_id: Optional[str] = Field(
+        default=None,
+        description="Unique identifier for the class instance (for persistence)",
+    )
+    create_new_instance: bool = Field(
+        default=True,
+        description="Whether to create a new instance or reuse existing one",
+    )
+
+    # Download acceleration fields
+    accelerate_downloads: bool = Field(
+        default=True,
+        description="Enable download acceleration for dependencies and models",
+    )
+
+    @model_validator(mode="after")
+    def validate_execution_requirements(self) -> "FunctionRequest":
+        """Validate that required fields are provided based on execution_type"""
+        if self.execution_type == "function":
+            if self.function_name is None:
+                raise ValueError(
+                    'function_name is required when execution_type is "function"'
+                )
+            if self.function_code is None:
+                raise ValueError(
+                    'function_code is required when execution_type is "function"'
+                )
+
+        elif self.execution_type == "class":
+            if self.class_name is None:
+                raise ValueError(
+                    'class_name is required when execution_type is "class"'
+                )
+            if self.class_code is None:
+                raise ValueError(
+                    'class_code is required when execution_type is "class"'
+                )
+
+        return self
+
+
+class FunctionResponse(BaseModel):
+    """Response model for remote function or class execution results.
+
+    Contains execution results, error information, and metadata about class instances
+    when applicable. The result field contains base64-encoded cloudpickle data.
+    """
+
+    success: bool = Field(
+        description="Indicates if the function execution was successful",
+    )
+    result: Optional[str] = Field(
+        default=None,
+        description="Base64-encoded cloudpickle-serialized result of the function",
+    )
+    error: Optional[str] = Field(
+        default=None,
+        description="Error message if the function execution failed",
+    )
+    stdout: Optional[str] = Field(
+        default=None,
+        description="Captured standard output from the function execution",
+    )
+    instance_id: Optional[str] = Field(
+        default=None, description="ID of the class instance that was used/created"
+    )
+    instance_info: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description="Metadata about the class instance (creation time, call count, etc.)",
+    )
+
+
+class RemoteExecutorStub(ABC):
+    """Abstract base class for remote execution."""
+
+    @abstractmethod
+    async def ExecuteFunction(self, request: FunctionRequest) -> FunctionResponse:
+        """Execute a function on the remote resource."""
+        raise NotImplementedError("Subclasses should implement this method.")

tetra_rp/stubs/__init__.py

@@ -0,0 +1,5 @@
+from .registry import stub_resource
+
+__all__ = [
+    "stub_resource",
+]

tetra_rp/stubs/live_serverless.py

@@ -0,0 +1,155 @@
+import ast
+import base64
+import inspect
+import textwrap
+import hashlib
+import traceback
+import threading
+import cloudpickle
+import logging
+from ..core.resources import LiveServerless
+from ..protos.remote_execution import (
+    FunctionRequest,
+    FunctionResponse,
+    RemoteExecutorStub,
+)
+
+log = logging.getLogger(__name__)
+
+
+# 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 (both sync and async)
+    function_def = None
+    for node in ast.walk(module):
+        if (
+            isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
+            and node.name == func.__name__
+        ):
+            function_def = node
+            break
+
+    if not function_def:
+        raise ValueError(f"Could not find function definition for {func.__name__}")
+
+    # Get the line and column offsets
+    lineno = function_def.lineno - 1  # Line numbers are 1-based
+
+    # Split into lines and extract just the function part
+    lines = source.split("\n")
+    function_lines = lines[lineno:]
+
+    # Dedent to remove any extra indentation
+    function_source = textwrap.dedent("\n".join(function_lines))
+
+    # Return the function hash for cache key
+    source_hash = hashlib.sha256(function_source.encode("utf-8")).hexdigest()
+
+    return function_source, source_hash
+
+
+class LiveServerlessStub(RemoteExecutorStub):
+    """Adapter class to make Runpod endpoints look like gRPC stubs."""
+
+    def __init__(self, server: LiveServerless):
+        self.server = server
+
+    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,
+        }
+
+        # 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]
+
+        # Serialize arguments using cloudpickle
+        if args:
+            request["args"] = [
+                base64.b64encode(cloudpickle.dumps(arg)).decode("utf-8") for arg in args
+            ]
+        if kwargs:
+            request["kwargs"] = {
+                k: base64.b64encode(cloudpickle.dumps(v)).decode("utf-8")
+                for k, v in kwargs.items()
+            }
+
+        return FunctionRequest(**request)
+
+    def handle_response(self, response: FunctionResponse):
+        if not (response.success or response.error):
+            raise ValueError("Invalid response from server")
+
+        if response.stdout:
+            for line in response.stdout.splitlines():
+                print(line)
+
+        if response.success:
+            if response.result is None:
+                raise ValueError("Response result is None")
+            return cloudpickle.loads(base64.b64decode(response.result))
+        else:
+            raise Exception(f"Remote execution failed: {response.error}")
+
+    async def ExecuteFunction(
+        self, request: FunctionRequest, sync: bool = False
+    ) -> FunctionResponse:
+        try:
+            # Convert the gRPC request to Runpod format
+            payload = request.model_dump(exclude_none=True)
+
+            if sync:
+                job = await self.server.run_sync(payload)
+            else:
+                job = await self.server.run(payload)
+
+            if job.error:
+                return FunctionResponse(
+                    success=False,
+                    error=job.error,
+                    stdout=job.output.get("stdout", ""),
+                )
+
+            return FunctionResponse(**job.output)
+
+        except Exception as e:
+            error_traceback = traceback.format_exc()
+            return FunctionResponse(
+                success=False,
+                error=f"{str(e)}\n{error_traceback}",
+            )