tetra-rp 0.7.0__tar.gz → 0.9.0__tar.gz

{tetra_rp-0.7.0 → tetra_rp-0.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tetra_rp
-Version: 0.7.0
+Version: 0.9.0
 Summary: A Python library for distributed inference and serving of machine learning models
 Author-email: Marut Pandya <pandyamarut@gmail.com>, Patrick Rachford <prachford@icloud.com>, Dean Quinanola <dean.quinanola@runpod.io>
 License: MIT
@@ -11,7 +11,7 @@ Classifier: Operating System :: OS Independent
 Requires-Python: <3.14,>=3.9
 Description-Content-Type: text/markdown
 Requires-Dist: cloudpickle>=3.1.1
-Requires-Dist: runpod~=1.7.9
+Requires-Dist: runpod
 Requires-Dist: python-dotenv>=1.0.0
 
 # Tetra: Serverless computing for AI workloads
@@ -801,6 +801,6 @@ def fetch_data(url):
 This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
 
 <p align="center">
-  <a href="https://github.com/yourusername/tetra">Tetra</a> •
+  <a href="https://github.com/runpod/tetra-rp">Tetra</a> •
   <a href="https://runpod.io">Runpod</a>
 </p>

{tetra_rp-0.7.0 → tetra_rp-0.9.0}/README.md

@@ -785,6 +785,6 @@ def fetch_data(url):
 This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
 
 <p align="center">
-  <a href="https://github.com/yourusername/tetra">Tetra</a> •
+  <a href="https://github.com/runpod/tetra-rp">Tetra</a> •
   <a href="https://runpod.io">Runpod</a>
 </p>

tetra_rp-0.9.0/pyproject.toml

@@ -0,0 +1,105 @@
+[project]
+name = "tetra_rp"
+version = "0.9.0"
+description = "A Python library for distributed inference and serving of machine learning models"
+authors = [
+    { name = "Marut Pandya", email = "pandyamarut@gmail.com" },
+    { name = "Patrick Rachford", email = "prachford@icloud.com" },
+    { name = "Dean Quinanola", email = "dean.quinanola@runpod.io" },
+]
+readme = "README.md"
+license = { text = "MIT" }
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+requires-python = ">=3.9,<3.14"
+
+dependencies = [
+    "cloudpickle>=3.1.1",
+    "runpod",
+    "python-dotenv>=1.0.0",
+]
+
+[dependency-groups]
+dev = [
+    "mypy>=1.16.1",
+    "ruff>=0.11.9",
+]
+test = [
+    "pytest>=8.4.1",
+    "pytest-mock>=3.14.0",
+    "pytest-asyncio>=1.0.0",
+    "pytest-cov>=6.2.1",
+    "twine>=6.1.0",
+]
+
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]
+addopts = [
+    "-v",
+    "--tb=short",
+    "--cov=tetra_rp",
+    "--cov-report=term-missing",
+    "--cov-fail-under=35"
+]
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"
+markers = [
+    "unit: Unit tests",
+    "integration: Integration tests",
+    "slow: Slow tests"
+]
+filterwarnings = [
+    "ignore::DeprecationWarning",
+    "ignore::PendingDeprecationWarning",
+    "ignore::pytest.PytestDeprecationWarning",
+    "ignore::pytest.PytestUnknownMarkWarning"
+]
+
+[tool.ruff]
+# Exclude tetra-examples directory since it's a separate repository
+exclude = [
+    "tetra-examples/",
+]
+
+[tool.mypy]
+# Basic configuration
+python_version = "3.9"
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = false  # Start lenient, can be stricter later
+disallow_incomplete_defs = false
+check_untyped_defs = true
+
+# Import discovery
+mypy_path = "."
+namespace_packages = true
+
+# Error output
+show_error_codes = true
+show_column_numbers = true
+pretty = true
+
+# Exclude directories
+exclude = [
+    "tetra-examples/",
+    "tests/",  # Start by excluding tests, can add later
+]
+
+# Per-module options
+[[tool.mypy.overrides]]
+module = [
+    "runpod.*",
+    "cloudpickle.*",
+]
+ignore_missing_imports = true

{tetra_rp-0.7.0 → tetra_rp-0.9.0}/src/tetra_rp/client.py

@@ -1,17 +1,19 @@
+import inspect
 import logging
 from functools import wraps
-from typing import List
-from .core.resources import ServerlessResource, ResourceManager
-from .stubs import stub_resource
+from typing import List, Optional
 
+from .core.resources import ResourceManager, ServerlessResource
+from .execute_class import create_remote_class
+from .stubs import stub_resource
 
 log = logging.getLogger(__name__)
 
 
 def remote(
     resource_config: ServerlessResource,
-    dependencies: List[str] = None,
-    system_dependencies: List[str] = None,
+    dependencies: Optional[List[str]] = None,
+    system_dependencies: Optional[List[str]] = None,
     **extra,
 ):
     """
@@ -24,8 +26,6 @@ def remote(
             to be provisioned or used.
         dependencies (List[str], optional): A list of pip package names to be installed in the remote
             environment before executing the function. Defaults to None.
-        mount_volume (NetworkVolume, optional): Configuration for creating and mounting a network volume.
-            Should contain 'size', 'datacenter_id', and 'name' keys. Defaults to None.
         extra (dict, optional): Additional parameters for the execution of the resource. Defaults to an empty dict.
 
     Returns:
@@ -45,17 +45,26 @@ def remote(
     ```
     """
 
-    def decorator(func):
-        @wraps(func)
-        async def wrapper(*args, **kwargs):
-            resource_manager = ResourceManager()
-            remote_resource = await resource_manager.get_or_deploy_resource(
-                resource_config
+    def decorator(func_or_class):
+        if inspect.isclass(func_or_class):
+            # Handle class decoration
+            return create_remote_class(
+                func_or_class, resource_config, dependencies, system_dependencies, extra
             )
+        else:
+            # Handle function decoration (unchanged)
+            @wraps(func_or_class)
+            async def wrapper(*args, **kwargs):
+                resource_manager = ResourceManager()
+                remote_resource = await resource_manager.get_or_deploy_resource(
+                    resource_config
+                )
 
-            stub = stub_resource(remote_resource, **extra)
-            return await stub(func, dependencies, system_dependencies, *args, **kwargs)
+                stub = stub_resource(remote_resource, **extra)
+                return await stub(
+                    func_or_class, dependencies, system_dependencies, *args, **kwargs
+                )
 
-        return wrapper
+            return wrapper
 
     return decorator

{tetra_rp-0.7.0 → tetra_rp-0.9.0}/src/tetra_rp/core/resources/live_serverless.py

@@ -3,9 +3,13 @@ import os
 from pydantic import model_validator
 from .serverless import ServerlessEndpoint
 
-
-TETRA_GPU_IMAGE = os.environ.get("TETRA_GPU_IMAGE", "runpod/tetra-rp:dev")
-TETRA_CPU_IMAGE = os.environ.get("TETRA_CPU_IMAGE", "runpod/tetra-rp-cpu:dev")
+TETRA_IMAGE_TAG = os.environ.get("TETRA_IMAGE_TAG", "latest")
+TETRA_GPU_IMAGE = os.environ.get(
+    "TETRA_GPU_IMAGE", f"runpod/tetra-rp:{TETRA_IMAGE_TAG}"
+)
+TETRA_CPU_IMAGE = os.environ.get(
+    "TETRA_CPU_IMAGE", f"runpod/tetra-rp-cpu:{TETRA_IMAGE_TAG}"
+)
 
 
 class LiveServerless(ServerlessEndpoint):

tetra_rp-0.9.0/src/tetra_rp/core/utils/constants.py

@@ -0,0 +1,10 @@
+"""
+Constants for utility modules and caching configurations.
+
+This module contains configurable constants used across the tetra-rp codebase
+to ensure consistency and easy maintenance.
+"""
+
+# Cache key generation constants
+HASH_TRUNCATE_LENGTH = 16  # Length to truncate hash values for cache keys
+UUID_FALLBACK_LENGTH = 8  # Length to truncate UUID values for fallback keys

tetra_rp-0.9.0/src/tetra_rp/core/utils/lru_cache.py

@@ -0,0 +1,75 @@
+"""
+LRU Cache implementation using OrderedDict for memory-efficient caching with automatic eviction.
+
+This module provides a Least Recently Used (LRU) cache implementation that automatically
+manages memory by evicting the least recently used items when the cache exceeds its
+maximum size limit. It maintains O(1) access time and provides a dict-like interface.
+Thread-safe for concurrent access.
+"""
+
+import threading
+from collections import OrderedDict
+from typing import Any, Dict, Optional
+
+
+class LRUCache:
+    """
+    A Least Recently Used (LRU) cache implementation using OrderedDict.
+
+    Automatically evicts the least recently used items when the cache exceeds
+    the maximum size limit. Provides dict-like interface with O(1) operations.
+    Thread-safe for concurrent access using RLock.
+
+    Args:
+        max_size: Maximum number of items to store in cache (default: 1000)
+    """
+
+    def __init__(self, max_size: int = 1000):
+        self.max_size = max_size
+        self.cache = OrderedDict()
+        self._lock = threading.RLock()
+
+    def get(self, key: str) -> Optional[Dict[str, Any]]:
+        """Get item from cache, moving it to end (most recent) if found."""
+        with self._lock:
+            if key in self.cache:
+                self.cache.move_to_end(key)
+                return self.cache[key]
+            return None
+
+    def set(self, key: str, value: Dict[str, Any]) -> None:
+        """Set item in cache, evicting oldest if at capacity."""
+        with self._lock:
+            if key in self.cache:
+                self.cache.move_to_end(key)
+            else:
+                if len(self.cache) >= self.max_size:
+                    self.cache.popitem(last=False)  # Remove oldest
+            self.cache[key] = value
+
+    def clear(self) -> None:
+        """Clear all items from cache."""
+        with self._lock:
+            self.cache.clear()
+
+    def __contains__(self, key: str) -> bool:
+        """Check if key exists in cache."""
+        with self._lock:
+            return key in self.cache
+
+    def __len__(self) -> int:
+        """Return number of items in cache."""
+        with self._lock:
+            return len(self.cache)
+
+    def __getitem__(self, key: str) -> Dict[str, Any]:
+        """Get item using bracket notation, moving to end if found."""
+        with self._lock:
+            if key in self.cache:
+                self.cache.move_to_end(key)
+                return self.cache[key]
+            raise KeyError(key)
+
+    def __setitem__(self, key: str, value: Dict[str, Any]) -> None:
+        """Set item using bracket notation."""
+        self.set(key, value)

tetra_rp-0.9.0/src/tetra_rp/execute_class.py

@@ -0,0 +1,316 @@
+"""
+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]],
+    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._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,
+                    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