tetra-rp 0.7.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

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/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/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/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/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

tetra_rp/protos/remote_execution.py

@@ -1,15 +1,18 @@
 # TODO: generate using betterproto
-
 from abc import ABC, abstractmethod
-from typing import List, Dict, Optional
-from pydantic import BaseModel, Field
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, Field, model_validator
 
 
 class FunctionRequest(BaseModel):
-    function_name: str = Field(
+    # 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: str = Field(
+    function_code: Optional[str] = Field(
+        default=None,
         description="Source code of the function to execute",
     )
     args: List = Field(
@@ -29,8 +32,67 @@ class FunctionRequest(BaseModel):
         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: Optional[List] = Field(
+        default_factory=list,
+        description="List of base64-encoded cloudpickle-serialized constructor arguments",
+    )
+    constructor_kwargs: Optional[Dict] = 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",
+    )
+
+    @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):
+    # EXISTING FIELDS (unchanged)
     success: bool = Field(
         description="Indicates if the function execution was successful",
     )
@@ -47,6 +109,15 @@ class FunctionResponse(BaseModel):
         description="Captured standard output from the function execution",
     )
 
+    # NEW FIELDS FOR CLASS SUPPORT
+    instance_id: Optional[str] = Field(
+        default=None, description="ID of the class instance that was used/created"
+    )
+    instance_info: Optional[Dict] = Field(
+        default=None,
+        description="Metadata about the class instance (creation time, call count, etc.)",
+    )
+
 
 class RemoteExecutorStub(ABC):
     """Abstract base class for remote execution."""

tetra_rp/stubs/registry.py

@@ -1,13 +1,13 @@
 import logging
 from functools import singledispatch
-from .live_serverless import LiveServerlessStub
-from .serverless import ServerlessEndpointStub
+
 from ..core.resources import (
     CpuServerlessEndpoint,
     LiveServerless,
     ServerlessEndpoint,
 )
-
+from .live_serverless import LiveServerlessStub
+from .serverless import ServerlessEndpointStub
 
 log = logging.getLogger(__name__)
 
@@ -22,20 +22,29 @@ def stub_resource(resource, **extra):
 
 @stub_resource.register(LiveServerless)
 def _(resource, **extra):
+    stub = LiveServerlessStub(resource)
+
+    # Function execution
     async def stubbed_resource(
         func, dependencies, system_dependencies, *args, **kwargs
     ) -> dict:
         if args == (None,):
-            # cleanup: when the function is called with no args
             args = []
 
-        stub = LiveServerlessStub(resource)
         request = stub.prepare_request(
             func, dependencies, system_dependencies, *args, **kwargs
         )
         response = await stub.ExecuteFunction(request)
         return stub.handle_response(response)
 
+    # Class method execution
+    async def execute_class_method(request):
+        response = await stub.ExecuteFunction(request)
+        return stub.handle_response(response)
+
+    # Attach the method to the function
+    stubbed_resource.execute_class_method = execute_class_method
+
     return stubbed_resource
 
 

{tetra_rp-0.7.0.dist-info → tetra_rp-0.9.0.dist-info}/METADATA

@@ -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.dist-info → tetra_rp-0.9.0.dist-info}/RECORD

@@ -1,15 +1,10 @@
 tetra_rp/__init__.py,sha256=-1S5sYIKtnUV8V1HlSIbX1yZwiUrsO8J5b3ZEIR_phU,687
-tetra_rp/client.py,sha256=5zerW5tTUnTSe75cRGgTBhqsKNXoCVWgb3Kzh9tJvPA,2209
+tetra_rp/client.py,sha256=rAMMmn4ejAayFXJMZzx7dG_8Y65tCEMI6wSSKgur4zQ,2500
+tetra_rp/execute_class.py,sha256=HoH-qWDA7X6yGvQMwmHn5-MKxbLWHEDEHsuat5dzl2U,11912
 tetra_rp/logger.py,sha256=gk5-PWp3k_GQ5DxndsRkBCX0jarp_3lgZ1oiTFuThQg,1125
 tetra_rp/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tetra_rp/core/api/__init__.py,sha256=oldrEKMwxYoBPLvPfVlaFS3wfUtTTxCN6-HzlpTh6vE,124
 tetra_rp/core/api/runpod.py,sha256=sux4q6xg2PDRKJI5kLkcW4i8UISZUOmQxsdf0g6wgpw,9711
-tetra_rp/core/pool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tetra_rp/core/pool/cluster_manager.py,sha256=KJxEp_044HjnbOhfIdiXZbks_bFDYE1KgKeR5W9VvbY,6007
-tetra_rp/core/pool/dataclass.py,sha256=YngS328_NTewY8Etitj4k7MmdM5GWqqE_OMbytrVNlw,338
-tetra_rp/core/pool/ex.py,sha256=AZOrn9t_X5ycMl-tDg7-jcIURj_9kVmzn9_da8h1TFI,1273
-tetra_rp/core/pool/job.py,sha256=4bisW_ZwiQ2-qD5l0y9SbHcO4EQvSKimmBBU1fpI_YE,567
-tetra_rp/core/pool/worker.py,sha256=N4cOnf8MiDcPFH2XSMmSnnWMACZYUNnKWVhOx2aSxvM,478
 tetra_rp/core/resources/__init__.py,sha256=UhIwo1Y6-tw5qsULamR296sQiztuz-oWrSTreqfmFSw,814
 tetra_rp/core/resources/base.py,sha256=UJeDiFN45aO1n5SBcxn56ohLhj-AWHoj0KO7mF4yJ_o,1440
 tetra_rp/core/resources/cloud.py,sha256=XJOWPfzYlDVJGHxgffcfpEaOKrWhGdi7AzTlaGuYj0o,70
@@ -17,7 +12,7 @@ tetra_rp/core/resources/constants.py,sha256=F1gPqFaXcCmfrbUSO9PQtUBv984TxFc3pySg
 tetra_rp/core/resources/cpu.py,sha256=YIE-tKolSU3JJzpPB7ey-PbRdqKWsJZ_Ad4h2OYaaiA,1231
 tetra_rp/core/resources/environment.py,sha256=FC9kJCa8YLSar75AKUKqJYnNLrUdjZj8ZTOrspBrS00,1267
 tetra_rp/core/resources/gpu.py,sha256=2jIIMr8PNnlIAP8ZTKO8Imx-rdxXp2rbdSHJeVfjawk,1858
-tetra_rp/core/resources/live_serverless.py,sha256=6r4I4TEx9AmZ0-OJvE86qrY0S7BEx9t_P2zwHVdtbew,1074
+tetra_rp/core/resources/live_serverless.py,sha256=A3JRdCYwHR2KN_OlmTLcv-m_ObxNhBhc5CnUzXOpOtc,1177
 tetra_rp/core/resources/network_volume.py,sha256=5_gwJlxt77VHs7T0d41l3IMZR0LhdoyQhroXCYfFF7w,3274
 tetra_rp/core/resources/resource_manager.py,sha256=kUVZDblfUzaG78S8FwOzu4rN6QSegUgQNK3fJ_X7l0w,2834
 tetra_rp/core/resources/serverless.py,sha256=RYH-gl_edEguGOlxR669Hfi_rXII4OEaYzlB2PhzOhI,15753
@@ -25,15 +20,17 @@ tetra_rp/core/resources/template.py,sha256=UkflJXZFWIbQkLuUt4oRLAjn-yIpw9_mT2X1c
 tetra_rp/core/resources/utils.py,sha256=mgXfgz_NuHN_IC7TzMNdH9II-LMjxcDCG7syDTcPiGs,1721
 tetra_rp/core/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tetra_rp/core/utils/backoff.py,sha256=1pfa0smFNpib8nztcIgBbtrVvQeECKh-aNOfL2TztgU,1324
+tetra_rp/core/utils/constants.py,sha256=Dm4XiO5zTzfdqOSeYVfAjaf2LyHnIEVmbOi_s_k1J_E,375
 tetra_rp/core/utils/json.py,sha256=q0r7aEdfh8kKVeHGeh9fBDfuhHYNopSreislAMB6HhM,1163
+tetra_rp/core/utils/lru_cache.py,sha256=drwKg-DfLbeBRGTzuxKqNKMQq0EuZV15LMTZIOyZuVk,2618
 tetra_rp/core/utils/singleton.py,sha256=JRli0HhBfq4P9mBUOg1TZUUwMvIenRqWdymX3qFMm2k,210
 tetra_rp/protos/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tetra_rp/protos/remote_execution.py,sha256=I4641-Dlzj4O7AuZbVei8O9aV2VNrRTcE8r7Fm0e-V8,1901
+tetra_rp/protos/remote_execution.py,sha256=F4uwobnp5q-lX3lR7NCAB23J6OzlzcsB35cezwuoSnI,4638
 tetra_rp/stubs/__init__.py,sha256=ozKsHs8q0T7o2qhQEquub9hqomh1Htys53mMraaRu2E,72
 tetra_rp/stubs/live_serverless.py,sha256=o1NH5XEwUD-27NXJsEGO0IwnuDp8iXwUiw5nZtaZZOI,4199
-tetra_rp/stubs/registry.py,sha256=V4m3CeXl8j1pguHuuflxqpWeBgVDQ93YkhxJbElyP7Q,2599
+tetra_rp/stubs/registry.py,sha256=dmbyC7uBp04_sXsG2wJCloFfFRzYjYQ-naEBKhTRo-U,2839
 tetra_rp/stubs/serverless.py,sha256=BM_a5Ml5VADBYu2WRNmo9qnicP8NnXDGl5ywifulbD0,947
-tetra_rp-0.7.0.dist-info/METADATA,sha256=jcXAGoiAFJVTYosJV2SGMw8L14UkUw2PHNKKH5mXkR8,28055
-tetra_rp-0.7.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tetra_rp-0.7.0.dist-info/top_level.txt,sha256=bBay7JTDwJXsTYvVjrwno9hnF-j0q272lk65f2AcPjU,9
-tetra_rp-0.7.0.dist-info/RECORD,,
+tetra_rp-0.9.0.dist-info/METADATA,sha256=qkRvg25koaP7AHCTpdd9mbZU1GGpBYebK5Gu3aHxics,28045
+tetra_rp-0.9.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tetra_rp-0.9.0.dist-info/top_level.txt,sha256=bBay7JTDwJXsTYvVjrwno9hnF-j0q272lk65f2AcPjU,9
+tetra_rp-0.9.0.dist-info/RECORD,,

tetra_rp/core/pool/cluster_manager.py

@@ -1,177 +0,0 @@
-import time
-from worker import Worker
-from job import Job
-
-from dataclass import WorkerStatus, JobStatus
-
-import logging
-import inspect
-
-
-def setup_logging(level=logging.INFO, fmt=None):
-    if fmt is None:
-        fmt = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-    logging.basicConfig(level=level, format=fmt)
-
-
-def get_logger(name=None):
-    """
-    Returns a logger. If no name is provided, it infers the caller's module name.
-    """
-    if name is None:
-        # Get the caller's module name.
-        frame = inspect.stack()[1]
-        module = inspect.getmodule(frame[0])
-        name = module.__name__ if module else "__main__"
-    return logging.getLogger(name)
-
-
-logger = get_logger(__name__)
-
-
-class ClusterManager:
-    """
-    Manages workers and Jobs currently in Memory:
-    - Runpod for provisioning
-    - Real remote execution
-    - Data base for the
-    """
-
-    def __init__(self):
-        self.workers = {}  # Worker ID -> Worker
-        self.jobs = {}  # Job ID -> Job
-
-    # ----------------- Worker Management -----------------
-    # ------------------------------------------------------
-    def add_worker(self, resource_config: dict):
-        """
-        Add a new worker to the cluster
-        """
-        # here will go the logic to create a worker and add it to the cluster: RUNPOD LOGIC will be added here.
-        worker = Worker(resource_config)
-        self.workers[worker.worker_id] = worker
-
-        logger.info(f"Added worker {worker.worker_id} to the cluster")
-        return worker.worker_id
-
-    def remove_worker(self, worker_id):
-        """
-        Remove a worker from the cluster
-        """
-        worker = self.workers.get(worker_id)
-        if not worker:
-            logger.error(f"Worker {worker_id} not found")
-            return False
-        if worker.status == WorkerStatus.RUNNING:
-            logger.error(f"Worker {worker_id} is still running")
-            return False
-        del self.workers[worker_id]
-        logger.info(f"Removed worker {worker_id} from the cluster")
-        return True
-
-    def list_workers(self):
-        """
-        List all workers in the cluster
-        """
-        return list(self.workers.values())
-
-    # ----------------- Job Management -----------------
-    # ---------------------------------------------------
-
-    def submit_job(self, resource_config: dict):
-        """
-        Submit a new job to the cluster (Queueud). Then attempt to scheduel it.
-        """
-        job = Job(resource_config)
-        self.jobs[job.job_id] = job
-        logger.info(f"Submitted job {job.job_id} to the cluster")
-        # attempt to schedule the job
-        self.schedule_job(job)
-        return job.job_id
-
-    def schedule_job(self, job: Job):
-        """
-        find a suitable worker for the job. It none, Job remains queued.
-        If we want to a auto provision we can actually add a logic here to add a worker if none is available.
-        """
-        if job.status != JobStatus.QUEUED:
-            logger.error(f"Job {job.job_id} is not pending")
-            return False
-
-        # Find worker candidate
-        candidate = self.find_idle_worker(job.resource_config)
-        if candidate:
-            self.assign_job_to_worker(job, candidate)
-        else:
-            logger.info(f"No worker available for job {job.job_id}")
-            # we cn either provision new worker from here and then scehediule the job from here.
-
-    def find_idle_worker(self, resource_config: dict):
-        """
-        Find an idle worker that can run the job
-        """
-        for w in self.workers.values():
-            if w.status == WorkerStatus.IDLE:
-                # check the resource config
-                if w.resource_config == resource_config:
-                    continue
-                return w
-        return None
-
-    def assign_job_to_worker(self, job: Job, worker: Worker):
-        """
-        Mark the job as running and the worker as Running and 'execute' the job.
-        In a real system, we would send a remote command to the worker (eg: gRPC) to execute the job.
-        """
-        job.worker_id = worker.worker_id
-        job.status = JobStatus.RUNNING
-        worker.status = WorkerStatus.RUNNING
-        worker.current_job_id = job.job_id
-        logger.info(f"Assigned job {job.job_id} to worker {worker.worker_id}")
-        self._execute_job(job, worker)
-
-    def _execute_job(self, job: Job, worker: Worker):
-        """
-        Simulate the remote execution. right now, we jsut sleep for 1s.
-        In production, what we we can do is:
-        - Open a gRPC connection to the worker
-        - pass the job details
-        - wait for the compeltion call back
-        """
-        try:
-            logger.info(f"Executing job {job.job_id} on worker {worker.worker_id}")
-            time.sleep(
-                1
-            )  # Here we can add the actual execution logic, currently it mimics the execution.
-
-            # mark the job as completed
-            job.status = JobStatus.COMPLETED
-            job.result = "Job completed successfully"
-            logger.info(f"[Cluster Manager] Job {job.job_id} completed successfully")
-        except Exception as e:
-            job.status = JobStatus.FAILED
-            job.result = f"Job failed: {str(e)}"
-            logger.error(f"[Cluster Manager] Job {job.job_id} failed: {str(e)}")
-        finally:
-            worker.status = WorkerStatus.IDLE
-            worker.current_job_id = None
-
-    def get_job_status(self, job_id):
-        """
-        Get the job details
-        """
-        job = self.jobs.get(job_id)
-        if not job:
-            logger.error(f"Job {job_id} not found")
-            return None
-        return job
-
-    # this function has retry logic but it's currently fuzzy, we might have to change it.
-
-    def retry_queued_jobs(self):
-        """
-        Retry all queued jobs
-        """
-        for job in self.jobs.values():
-            if job.status == JobStatus.QUEUED:
-                self.schedule_job(job)

tetra_rp/core/pool/dataclass.py

@@ -1,18 +0,0 @@
-from enum import Enum
-
-
-class WorkerStatus(Enum):
-    """Enum representing the status of a worker"""
-
-    IDLE = "idle"
-    RUNNING = "running"
-    OFFLINE = "offline"
-
-
-class JobStatus(Enum):
-    """Enum representing the status of a job"""
-
-    QUEUED = "queued"
-    RUNNING = "running"
-    COMPLETED = "completed"
-    FAILED = "failed"

tetra_rp/core/pool/ex.py

@@ -1,38 +0,0 @@
-from cluster_manager import ClusterManager
-
-
-if __name__ == "__main__":
-    cm = ClusterManager()
-
-    # 1) Submit a job with no existing workers (use resource_config dict)
-    job_id = cm.submit_job(
-        resource_config={"gpu": "H100", "memory": 16, "network_volume": 50}
-    )
-    print(
-        "Job status:", cm.get_job_status(job_id)
-    )  # should be QUEUED, no suitable worker
-
-    # 2) Add a worker that doesn't match the GPU
-    w1 = cm.add_worker(
-        resource_config={"gpu": "H100", "memory": 16, "network_volume": 50}
-    )
-    # Re-try scheduling
-    cm.retry_queued_jobs()
-    print("Job status (still queued):", cm.get_job_status(job_id))
-
-    # 3) Add a matching worker
-    w2 = cm.add_worker(
-        resource_config={"gpu": "H100", "memory": 16, "network_volume": 50}
-    )
-    # Re-try scheduling
-    cm.retry_queued_jobs()
-    print("Job status (should complete):", cm.get_job_status(job_id))
-
-    # 4) Submit another job that requires less resources
-    job_id2 = cm.submit_job(resource_config={"memory": 8, "network_volume": 10})
-    # Should be assigned to w1 if it's idle
-    print("Job2 final status:", cm.get_job_status(job_id2))
-
-    # 5) Show final state of workers
-    for worker in cm.list_workers():
-        print("Worker:", worker)

tetra_rp/core/pool/job.py

@@ -1,22 +0,0 @@
-import uuid
-from dataclass import JobStatus
-
-
-class Job:
-    """Represents a 'job' in the system
-
-    In a real system, this might contain the function to run,
-    arguments, and reference to data or code.
-    """
-
-    def __init__(self, resource_config: dict):
-        self.job_id = str(uuid.uuid4())[:8]
-        self.resource_config = resource_config
-        self.status = JobStatus.QUEUED
-
-        self.worker_id = None
-        self.result = None
-        self.error = None
-
-    def __repr__(self):
-        return f"Job(job_id={self.job_id}, status={self.status})"

tetra_rp/core/pool/worker.py

@@ -1,19 +0,0 @@
-import uuid
-from dataclass import WorkerStatus
-
-
-class Worker:
-    """Represents a single worker in the pool
-
-    For Now we store ressources in memory
-    """
-
-    def __init__(self, resource_config: dict):
-        self.worker_id = str(uuid.uuid4())[:8]
-        self.resource_config = resource_config
-        self.status = WorkerStatus.IDLE
-
-        self.current_job_id = None
-
-    def __repr__(self):
-        return f"Worker(worker_id={self.worker_id}, status={self.status})"