tetra-rp 0.7.0__py3-none-any.whl → 0.8.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/execute_class.py

@@ -0,0 +1,178 @@
+import base64
+import inspect
+import logging
+import textwrap
+import uuid
+from typing import List, Type, Optional
+
+import cloudpickle
+
+from .core.resources import ResourceManager, ServerlessResource
+from .protos.remote_execution import FunctionRequest
+from .stubs import stub_resource
+
+log = logging.getLogger(__name__)
+
+
+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 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[:8]}"
+            self._initialized = False
+
+            self._clean_class_code = extract_class_code_simple(cls)
+
+            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()
+
+                # Create class method request
+
+                # class_code = inspect.getsource(self._class_type)
+                class_code = self._clean_class_code
+
+                request = FunctionRequest(
+                    execution_type="class",
+                    class_name=self._class_type.__name__,
+                    class_code=class_code,
+                    method_name=name,
+                    args=[
+                        base64.b64encode(cloudpickle.dumps(arg)).decode("utf-8")
+                        for arg in args
+                    ],
+                    kwargs={
+                        k: base64.b64encode(cloudpickle.dumps(v)).decode("utf-8")
+                        for k, v in kwargs.items()
+                    },
+                    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()
+                    },
+                    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.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tetra_rp
-Version: 0.7.0
+Version: 0.8.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

{tetra_rp-0.7.0.dist-info → tetra_rp-0.8.0.dist-info}/RECORD

@@ -1,5 +1,6 @@
 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=OXP1IkORELNFxOi1WHQOfUepmQGfkKmw85iZccaMEww,6515
 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
@@ -28,12 +29,12 @@ tetra_rp/core/utils/backoff.py,sha256=1pfa0smFNpib8nztcIgBbtrVvQeECKh-aNOfL2Tztg
 tetra_rp/core/utils/json.py,sha256=q0r7aEdfh8kKVeHGeh9fBDfuhHYNopSreislAMB6HhM,1163
 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.8.0.dist-info/METADATA,sha256=M0qEc5SQITXYUy_FQLy_EE24CiTbJtrXjU6g5-q3fws,28055
+tetra_rp-0.8.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tetra_rp-0.8.0.dist-info/top_level.txt,sha256=bBay7JTDwJXsTYvVjrwno9hnF-j0q272lk65f2AcPjU,9
+tetra_rp-0.8.0.dist-info/RECORD,,