feat: AE-815 Add class based execution

src/tetra_rp/client.py

@@ -1,10 +1,13 @@
+import inspect
 import logging
 from functools import wraps
 from typing import List, Optional
-from .core.resources import ServerlessResource, ResourceManager, NetworkVolume
-from .stubs import stub_resource
 
 
+from .core.resources import NetworkVolume, ResourceManager, ServerlessResource
+from .stubs import stub_resource
+from .execute_class import create_remote_class
+
 log = logging.getLogger(__name__)
 
 
@@ -46,26 +49,26 @@ async def my_function(data):
     ```
     """
 
-    def decorator(func):
-        @wraps(func)
-        async def wrapper(*args, **kwargs):
-            # Create netowrk volume if mount_volume is provided
-            if mount_volume:
-                try:
-                    network_volume = await mount_volume.deploy()
-                    resource_config.networkVolumeId = network_volume.id
-                except Exception as e:
-                    log.error(f"Failed to create or mount network volume: {e}")
-                    raise
-
-            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

src/tetra_rp/execute_class.py

@@ -0,0 +1,172 @@
+import base64
+import inspect
+import logging
+import textwrap
+import uuid
+from typing import List, Type
+
+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")
+
+        print(f"Successfully extracted class code for {cls.__name__}")
+        return class_code
+
+    except Exception as e:
+        print(f"Warning: Could not extract class code for {cls.__name__}: {e}")
+        print("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:  # noqa: E722
+                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: List[str],
+    system_dependencies: List[str],
+    extra: dict,
+):
+    """
+    Create a remote class wrapper.
+    """
+
+    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)
+
+            return method_proxy
+
+    return RemoteClassWrapper

src/tetra_rp/protos/remote_execution.proto

@@ -4,22 +4,36 @@ package tetra;
 
 // The remote execution service definition
 service RemoteExecutor {
-  // Execute a function remotely
+  // Execute a function or class method remotely
   rpc ExecuteFunction (FunctionRequest) returns (FunctionResponse) {}
 }
 
-// The request message containing function details and arguments
+// The request message containing function/class details and arguments
 message FunctionRequest {
-  string function_name = 1;       // Name of the function to execute
-  string function_code = 2;       // Source code of the function
-  repeated string args = 3;       // Base64-encoded cloudpickle-serialized positional arguments
-  map<string, string> kwargs = 4; // Base64-encoded cloudpickle-serialized keyword arguments
-  repeated string dependencies = 5; // Optional list of pip packages to install before execution
+  optional string function_name = 1;       // Name of the function to execute
+  optional string function_code = 2;       // Source code of the function to execute
+  repeated string args = 3;                // Base64-encoded cloudpickle-serialized positional arguments
+  map<string, string> kwargs = 4;         // Base64-encoded cloudpickle-serialized keyword arguments
+  repeated string dependencies = 5;        // Optional list of pip packages to install before execution
+  repeated string system_dependencies = 6; // Optional list of system dependencies to install before execution
+
+  string execution_type = 7;               // Type of execution: 'function' or 'class'
+  optional string class_name = 8;          // Name of the class to instantiate (for class execution)
+  optional string class_code = 9;          // Source code of the class to instantiate (for class execution)
+  repeated string constructor_args = 10;   // Base64-encoded cloudpickle-serialized constructor arguments
+  map<string, string> constructor_kwargs = 11; // Base64-encoded cloudpickle-serialized constructor keyword arguments
+  string method_name = 12;                 // Name of the method to call on the class instance (default: "__call__")
+  optional string instance_id = 13;        // Unique identifier for the class instance (for persistence)
+  bool create_new_instance = 14;           // Whether to create a new instance or reuse existing one
 }
 
 // The response message containing the execution result or error
 message FunctionResponse {
-  bool success = 1;               // Whether execution was successful
-  string result = 2;              // Base64-encoded cloudpickle-serialized result (if success)
-  string error = 3;               // Error message (if not success)
+  bool success = 1;                        // Whether execution was successful
+  optional string result = 2;              // Base64-encoded cloudpickle-serialized result (if success)
+  optional string error = 3;               // Error message (if not success)
+  optional string stdout = 4;              // Captured standard output from the function execution
+
+  optional string instance_id = 5;         // ID of the class instance that was used/created
+  map<string, string> instance_info = 6;   // Metadata about the class instance (creation time, call count, etc.)
 }

src/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."""