diff --git a/handler.py b/handler.py index a518146..a1b94f5 100644 --- a/handler.py +++ b/handler.py @@ -7,7 +7,11 @@ import io import logging import os +import uuid +import sys +from datetime import datetime from contextlib import redirect_stdout, redirect_stderr +from typing import Dict, Any from remote_execution import ( FunctionRequest, FunctionResponse, @@ -15,15 +19,28 @@ ) +logging.basicConfig( + level=logging.DEBUG, # or INFO for less verbose output + stream=sys.stdout, # send logs to stdout (so docker captures it) + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", +) + + class RemoteExecutor(RemoteExecutorStub): """ - RemoteExecutor class for executing functions in a serverless environment. + RemoteExecutor class for executing functions and classes in a serverless environment. Inherits from RemoteExecutorStub. """ + def __init__(self): + super().__init__() + # Instance registry for persistent class instances + self.class_instances: Dict[str, Any] = {} + self.instance_metadata: Dict[str, Dict] = {} + async def ExecuteFunction(self, request: FunctionRequest) -> FunctionResponse: """ - Execute a function on the remote resource. + Execute a function or class method on the remote resource. Args: request: FunctionRequest object containing function details @@ -33,6 +50,9 @@ async def ExecuteFunction(self, request: FunctionRequest) -> FunctionResponse: """ # Install system dependencies first if request.system_dependencies: + sys_installed = self.install_system_dependencies( + request.system_dependencies + ) sys_installed = self.install_system_dependencies( request.system_dependencies ) @@ -47,18 +67,161 @@ async def ExecuteFunction(self, request: FunctionRequest) -> FunctionResponse: return py_installed print(py_installed.stdout) - # Execute the function - return self.execute(request) + # Route to appropriate execution method based on type + execution_type = getattr(request, "execution_type", "function") + if execution_type == "class": + return self.execute_class_method(request) + else: + return self.execute(request) # Your existing function execution - def install_system_dependencies(self, packages) -> FunctionResponse: + # METHOD: Class method execution + def execute_class_method(self, request: FunctionRequest) -> FunctionResponse: """ - Install system packages using apt-get. + Execute a class method with instance management. + """ + stdout_io = io.StringIO() + stderr_io = io.StringIO() + log_io = io.StringIO() - Args: - packages: List of system package names + with redirect_stdout(stdout_io), redirect_stderr(stderr_io): + try: + # Setup logging + log_handler = logging.StreamHandler(log_io) + log_handler.setLevel(logging.DEBUG) + logger = logging.getLogger() + logger.addHandler(log_handler) - Returns: - FunctionResponse: Object indicating success or failure with details + # Get or create class instance + instance, instance_id = self._get_or_create_instance(request) + + # Get the method to call + method_name = getattr(request, "method_name", "__call__") + if not hasattr(instance, method_name): + return FunctionResponse( + success=False, + error=f"Method '{method_name}' not found in class '{request.class_name}'", + ) + + method = getattr(instance, method_name) + + # Deserialize method arguments + args = [ + cloudpickle.loads(base64.b64decode(arg)) for arg in request.args + ] + kwargs = { + k: cloudpickle.loads(base64.b64decode(v)) + for k, v in request.kwargs.items() + } + + # Execute the method + result = method(*args, **kwargs) + + # Update instance metadata + self._update_instance_metadata(instance_id) + + except Exception as e: + # Error handling + combined_output = ( + stdout_io.getvalue() + stderr_io.getvalue() + log_io.getvalue() + ) + traceback_str = traceback.format_exc() + error_message = f"{str(e)}\n{traceback_str}" + + return FunctionResponse( + success=False, + error=error_message, + stdout=combined_output, + ) + + finally: + logger.removeHandler(log_handler) + + # Serialize result + serialized_result = base64.b64encode(cloudpickle.dumps(result)).decode("utf-8") + combined_output = ( + stdout_io.getvalue() + stderr_io.getvalue() + log_io.getvalue() + ) + + return FunctionResponse( + success=True, + result=serialized_result, + stdout=combined_output, + instance_id=instance_id, + instance_info=self.instance_metadata.get(instance_id, {}), + ) + + def _get_or_create_instance(self, request: FunctionRequest) -> tuple[Any, str]: + """ + Get existing instance or create new one. + """ + instance_id = getattr(request, "instance_id", None) + create_new = getattr(request, "create_new_instance", True) + + # Check if we should reuse existing instance + if not create_new and instance_id and instance_id in self.class_instances: + logging.debug(f"Reusing existing instance: {instance_id}") + return self.class_instances[instance_id], instance_id + + # Create new instance + logging.info(f"Creating new instance of class: {request.class_name}") + + # Execute class code + namespace = {} + exec(request.class_code, namespace) + + if request.class_name not in namespace: + raise ValueError( + f"Class '{request.class_name}' not found in the provided code" + ) + + cls = namespace[request.class_name] + + # Deserialize constructor arguments + constructor_args = [] + constructor_kwargs = {} + + if hasattr(request, "constructor_args") and request.constructor_args: + constructor_args = [ + cloudpickle.loads(base64.b64decode(arg)) + for arg in request.constructor_args + ] + + if hasattr(request, "constructor_kwargs") and request.constructor_kwargs: + constructor_kwargs = { + k: cloudpickle.loads(base64.b64decode(v)) + for k, v in request.constructor_kwargs.items() + } + + # Create instance + instance = cls(*constructor_args, **constructor_kwargs) + + # Generate instance ID if not provided + if not instance_id: + instance_id = f"{request.class_name}_{uuid.uuid4().hex[:8]}" + + # Store instance + self.class_instances[instance_id] = instance + self.instance_metadata[instance_id] = { + "class_name": request.class_name, + "created_at": datetime.now().isoformat(), + "method_calls": 0, + "last_used": datetime.now().isoformat(), + } + + logging.info(f"Created instance with ID: {instance_id}") + return instance, instance_id + + def _update_instance_metadata(self, instance_id: str): + """Update metadata for an instance.""" + if instance_id in self.instance_metadata: + self.instance_metadata[instance_id]["method_calls"] += 1 + self.instance_metadata[instance_id]["last_used"] = ( + datetime.now().isoformat() + ) + + def install_system_dependencies(self, packages) -> FunctionResponse: + """ + Install system packages using apt-get. """ if not packages: return FunctionResponse( @@ -84,7 +247,6 @@ def install_system_dependencies(self, packages) -> FunctionResponse: ) # Install the packages - # -y flag for non-interactive, --no-install-recommends to keep it minimal process = subprocess.Popen( ["apt-get", "install", "-y", "--no-install-recommends"] + packages, stdout=subprocess.PIPE, @@ -92,7 +254,7 @@ def install_system_dependencies(self, packages) -> FunctionResponse: env={ **os.environ, "DEBIAN_FRONTEND": "noninteractive", - }, # Prevent prompts + }, # Prevent interactive prompts ) stdout, stderr = process.communicate() @@ -118,10 +280,8 @@ def install_system_dependencies(self, packages) -> FunctionResponse: def install_dependencies(self, packages) -> FunctionResponse: """ Install Python packages using pip with proper process completion handling. - Args: packages: List of package names or package specifications - Returns: FunctionResponse: Object indicating success or failure with details """ @@ -131,21 +291,15 @@ def install_dependencies(self, packages) -> FunctionResponse: print(f"Installing dependencies: {packages}") try: - # Use pip to install the packages - # Note: communicate() already waits for process completion process = subprocess.Popen( ["uv", "pip", "install", "--no-cache-dir"] + packages, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) - # This waits for the process to complete and captures output stdout, stderr = process.communicate() - - # Force reload of installed packages importlib.invalidate_caches() - # Simply rely on pip's return code if process.returncode != 0: return FunctionResponse( success=False, @@ -167,17 +321,14 @@ def install_dependencies(self, packages) -> FunctionResponse: def execute(self, request: FunctionRequest) -> FunctionResponse: """ Execute a function as a remote resource. - Args: request: FunctionRequest object containing function details - Returns: FunctionResponse object with execution result """ stdout_io = io.StringIO() stderr_io = io.StringIO() log_io = io.StringIO() - # Capture all stdout, stderr, and logs into variables and supply them to the FunctionResponse with redirect_stdout(stdout_io), redirect_stderr(stderr_io): try: @@ -198,7 +349,6 @@ def execute(self, request: FunctionRequest) -> FunctionResponse: func = namespace[request.function_name] - # Deserialize arguments using cloudpickle args = [ cloudpickle.loads(base64.b64decode(arg)) for arg in request.args ] @@ -207,15 +357,12 @@ def execute(self, request: FunctionRequest) -> FunctionResponse: for k, v in request.kwargs.items() } - # Execute the function result = func(*args, **kwargs) except Exception as e: - # Combine stdout, stderr, and logs combined_output = ( stdout_io.getvalue() + stderr_io.getvalue() + log_io.getvalue() ) - # Capture full traceback for better debugging traceback_str = traceback.format_exc() error_message = f"{str(e)}\n{traceback_str}" @@ -229,16 +376,13 @@ def execute(self, request: FunctionRequest) -> FunctionResponse: finally: # Remove the log handler to avoid duplicate logs logger.removeHandler(log_handler) - # Serialize result using cloudpickle serialized_result = base64.b64encode(cloudpickle.dumps(result)).decode("utf-8") - # Combine stdout, stderr, and logs combined_output = ( stdout_io.getvalue() + stderr_io.getvalue() + log_io.getvalue() ) - # Return success response return FunctionResponse( success=True, result=serialized_result, @@ -267,6 +411,5 @@ async def handler(event: dict) -> dict: # Start the RunPod serverless handler - if __name__ == "__main__": runpod.serverless.start({"handler": handler}) diff --git a/pyproject.toml b/pyproject.toml index dc35f81..090f25e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,4 +50,4 @@ filterwarnings = [ # Exclude tetra-rp directory since it's a separate repository exclude = [ "tetra-rp/", -] +] \ No newline at end of file diff --git a/test_class_input.json b/test_class_input.json new file mode 100644 index 0000000..e8899a1 --- /dev/null +++ b/test_class_input.json @@ -0,0 +1,15 @@ +{ + "input": { + "execution_type": "class", + "class_name": "TestClass", + "class_code": "class TestClass:\n def __init__(self, value):\n self.value = value\n \n def get_value(self):\n return f'Value is: {self.value}' ", + "method_name": "get_value", + "constructor_args": [ + "gAWVCQAAAAAAAACMBWhlbGxvlC4=" + ], + "constructor_kwargs": {}, + "args": [], + "kwargs": {}, + "create_new_instance": true + } +} \ No newline at end of file diff --git a/tests/integration/test_handler_integration.py b/tests/integration/test_handler_integration.py index f63b13f..aa31096 100644 --- a/tests/integration/test_handler_integration.py +++ b/tests/integration/test_handler_integration.py @@ -1,8 +1,22 @@ import pytest -from handler import handler +import json +import base64 +import cloudpickle +from pathlib import Path + +from handler import handler, RemoteExecutor +from remote_execution import FunctionRequest class TestHandlerIntegration: + """Integration tests using test_input.json and test_class_input.json.""" + + def setup_method(self): + """Setup for each test method.""" + self.test_data_dir = Path(__file__).parent.parent.parent + self.test_input_file = self.test_data_dir / "test_input.json" + self.test_class_input_file = self.test_data_dir / "test_class_input.json" + @pytest.mark.asyncio async def test_handler_end_to_end(self): """Test complete handler workflow with simple function""" @@ -20,3 +34,151 @@ async def test_handler_end_to_end(self): assert result["success"] is True assert result["error"] is None assert result["result"] is not None + + @pytest.mark.asyncio + async def test_handler_with_test_input_json(self): + """Test handler using test_input.json.""" + # Load the test input data + with open(self.test_input_file, "r") as f: + test_data = json.load(f) + + # Execute through the handler + result = await handler(test_data) + + # Verify the response + assert result["success"] is True + assert "result" in result + assert result["error"] is None + + # Decode and verify the actual result + decoded_result = cloudpickle.loads(base64.b64decode(result["result"])) + assert decoded_result == "hello world" + + # Check that stdout was captured + assert "going to say hello" in result["stdout"] + + @pytest.mark.asyncio + async def test_handler_with_test_class_input_json(self): + """Test handler using test_class_input.json.""" + # Load the test class input data + with open(self.test_class_input_file, "r") as f: + test_data = json.load(f) + + # Execute through the handler + result = await handler(test_data) + + # Verify the response + assert result["success"] is True + assert "result" in result + assert result["error"] is None + assert "instance_id" in result + + # Decode and verify the actual result + decoded_result = cloudpickle.loads(base64.b64decode(result["result"])) + assert decoded_result == "Value is: hello" + + # Verify instance information + assert result["instance_id"] is not None + assert "instance_info" in result + assert result["instance_info"]["class_name"] == "TestClass" + assert result["instance_info"]["method_calls"] == 1 + + @pytest.mark.asyncio + async def test_class_instance_reuse(self): + """Test reusing class instances across multiple calls.""" + executor = RemoteExecutor() + + # First call - create instance + request1 = FunctionRequest( + execution_type="class", + class_name="Counter", + class_code="class Counter:\n def __init__(self):\n self.count = 0\n def increment(self):\n self.count += 1\n return self.count", + method_name="increment", + constructor_args=[], + constructor_kwargs={}, + args=[], + kwargs={}, + create_new_instance=True, + ) + + response1 = await executor.ExecuteFunction(request1) + assert response1.success is True + instance_id = response1.instance_id + + result1 = cloudpickle.loads(base64.b64decode(response1.result)) + assert result1 == 1 + + # Second call - reuse instance + request2 = FunctionRequest( + execution_type="class", + class_name="Counter", + class_code="class Counter:\n def __init__(self):\n self.count = 0\n def increment(self):\n self.count += 1\n return self.count", + method_name="increment", + instance_id=instance_id, + create_new_instance=False, + args=[], + kwargs={}, + ) + + response2 = await executor.ExecuteFunction(request2) + assert response2.success is True + assert response2.instance_id == instance_id + + result2 = cloudpickle.loads(base64.b64decode(response2.result)) + assert result2 == 2 # Should increment from previous state + + # Verify metadata was updated + assert response2.instance_info["method_calls"] == 2 + + @pytest.mark.asyncio + async def test_handler_error_scenarios(self): + """Test handler with invalid input scenarios.""" + # Test with completely invalid event structure + invalid_event = {"invalid": "structure"} + result = await handler(invalid_event) + assert result["success"] is False + assert "Error in handler" in result["error"] + + # Test with missing required fields + invalid_event2 = { + "input": { + "execution_type": "function" + # Missing function_name and function_code + } + } + result2 = await handler(invalid_event2) + assert result2["success"] is False + + @pytest.mark.asyncio + async def test_complex_data_serialization(self): + """Test handling complex data types through the full pipeline.""" + test_data = { + "numbers": [1, 2, 3, 4, 5], + "metadata": {"name": "test", "version": 1.0}, + } + + event = { + "input": { + "function_name": "process_data", + "function_code": """ +def process_data(data): + return { + 'sum': sum(data['numbers']), + 'name': data['metadata']['name'], + 'processed': True + } +""", + "args": [ + base64.b64encode(cloudpickle.dumps(test_data)).decode("utf-8") + ], + "kwargs": {}, + } + } + + result = await handler(event) + assert result["success"] is True + + decoded_result = cloudpickle.loads(base64.b64decode(result["result"])) + assert decoded_result["sum"] == 15 + assert decoded_result["name"] == "test" + assert decoded_result["processed"] is True diff --git a/tests/unit/test_remote_executor.py b/tests/unit/test_remote_executor.py new file mode 100644 index 0000000..ba7c767 --- /dev/null +++ b/tests/unit/test_remote_executor.py @@ -0,0 +1,121 @@ +import pytest +import base64 +import cloudpickle +from unittest.mock import Mock, patch +import subprocess + +from handler import RemoteExecutor +from remote_execution import FunctionRequest + + +class TestRemoteExecutor: + """Unit tests for the RemoteExecutor class.""" + + def setup_method(self): + """Setup for each test method.""" + self.executor = RemoteExecutor() + + def encode_args(self, *args): + """Helper to encode arguments.""" + return [ + base64.b64encode(cloudpickle.dumps(arg)).decode("utf-8") for arg in args + ] + + def encode_kwargs(self, **kwargs): + """Helper to encode keyword arguments.""" + return { + k: base64.b64encode(cloudpickle.dumps(v)).decode("utf-8") + for k, v in kwargs.items() + } + + def test_executor_init(self): + """Test RemoteExecutor initialization.""" + assert hasattr(self.executor, "class_instances") + assert hasattr(self.executor, "instance_metadata") + assert len(self.executor.class_instances) == 0 + assert len(self.executor.instance_metadata) == 0 + + @pytest.mark.asyncio + async def test_execute_simple_function(self): + """Test basic function execution.""" + request = FunctionRequest( + function_name="hello", + function_code="def hello():\n return 'hello world'", + args=[], + kwargs={}, + ) + + response = self.executor.execute(request) + + assert response.success is True + result = cloudpickle.loads(base64.b64decode(response.result)) + assert result == "hello world" + + @pytest.mark.asyncio + async def test_execute_function_with_args(self): + """Test function execution with arguments.""" + request = FunctionRequest( + function_name="add", + function_code="def add(a, b):\n return a + b", + args=self.encode_args(5, 3), + kwargs={}, + ) + + response = self.executor.execute(request) + + assert response.success is True + result = cloudpickle.loads(base64.b64decode(response.result)) + assert result == 8 + + @pytest.mark.asyncio + async def test_execute_class_method(self): + """Test class method execution.""" + request = FunctionRequest( + execution_type="class", + class_name="TestClass", + class_code="class TestClass:\n def __init__(self, value):\n self.value = value\n def get_value(self):\n return f'Value: {self.value}'", + method_name="get_value", + constructor_args=self.encode_args("test"), + constructor_kwargs={}, + args=[], + kwargs={}, + ) + + response = self.executor.execute_class_method(request) + + assert response.success is True + assert response.instance_id is not None + result = cloudpickle.loads(base64.b64decode(response.result)) + assert result == "Value: test" + + @pytest.mark.asyncio + async def test_function_error_handling(self): + """Test error handling when function raises exception.""" + request = FunctionRequest( + function_name="error_func", + function_code="def error_func():\n raise ValueError('Test error')", + args=[], + kwargs={}, + ) + + response = self.executor.execute(request) + + assert response.success is False + assert "Test error" in response.error + + @patch("subprocess.Popen") + def test_install_dependencies(self, mock_popen): + """Test dependency installation.""" + mock_process = Mock() + mock_process.returncode = 0 + mock_process.communicate.return_value = (b"Successfully installed", b"") + mock_popen.return_value = mock_process + + response = self.executor.install_dependencies(["numpy"]) + + assert response.success is True + mock_popen.assert_called_once_with( + ["uv", "pip", "install", "--no-cache-dir", "numpy"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + )