tetra-rp 0.5.5__py3-none-any.whl

tetra_rp/core/utils/singleton.py

@@ -0,0 +1,7 @@
+class SingletonMixin:
+    _instances = {}
+
+    def __new__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            cls._instances[cls] = super().__new__(cls)
+        return cls._instances[cls]

tetra_rp/logger.py

@@ -0,0 +1,34 @@
+import logging
+import os
+import sys
+from typing import Union, Optional
+
+
+def setup_logging(
+    level: Union[int, str] = logging.INFO, stream=sys.stdout, fmt: Optional[str] = None
+):
+    """
+    Sets up the root logger with a stream handler and basic formatting.
+    Does nothing if handlers are already configured.
+    """
+    if isinstance(level, str):
+        level = getattr(logging, level.upper(), logging.INFO)
+
+    if fmt is None:
+        if level == logging.DEBUG:
+            fmt = "%(asctime)s | %(levelname)-5s | %(name)s | %(filename)s:%(lineno)d | %(message)s"
+        else:
+            # Default format for INFO level and above
+            fmt = "%(asctime)s | %(levelname)-5s | %(message)s"
+
+    root_logger = logging.getLogger()
+    if not root_logger.hasHandlers():
+        handler = logging.StreamHandler(stream)
+        handler.setFormatter(logging.Formatter(fmt))
+        root_logger.setLevel(level)
+        root_logger.addHandler(handler)
+
+    # Optionally allow log level override via env var
+    env_level = os.environ.get("LOG_LEVEL")
+    if env_level:
+        root_logger.setLevel(env_level.upper())

tetra_rp/protos/remote_execution.py

@@ -0,0 +1,57 @@
+# TODO: generate using betterproto
+
+from abc import ABC, abstractmethod
+from typing import List, Dict, Optional
+from pydantic import BaseModel, Field
+
+
+class FunctionRequest(BaseModel):
+    function_name: str = Field(
+        description="Name of the function to execute",
+    )
+    function_code: str = Field(
+        description="Source code of the function to execute",
+    )
+    args: List = Field(
+        default_factory=list,
+        description="List of base64-encoded cloudpickle-serialized arguments",
+    )
+    kwargs: Dict = Field(
+        default_factory=dict,
+        description="Dictionary of base64-encoded cloudpickle-serialized keyword arguments",
+    )
+    dependencies: Optional[List] = Field(
+        default=None,
+        description="Optional list of pip packages to install before executing the function",
+    )
+    system_dependencies: Optional[List] = Field(
+        default=None,
+        description="Optional list of system dependencies to install before executing the function",
+    )
+
+
+class FunctionResponse(BaseModel):
+    success: bool = Field(
+        description="Indicates if the function execution was successful",
+    )
+    result: Optional[str] = Field(
+        default=None,
+        description="Base64-encoded cloudpickle-serialized result of the function",
+    )
+    error: Optional[str] = Field(
+        default=None,
+        description="Error message if the function execution failed",
+    )
+    stdout: Optional[str] = Field(
+        default=None,
+        description="Captured standard output from the function execution",
+    )
+
+
+class RemoteExecutorStub(ABC):
+    """Abstract base class for remote execution."""
+
+    @abstractmethod
+    async def ExecuteFunction(self, request: FunctionRequest) -> FunctionResponse:
+        """Execute a function on the remote resource."""
+        raise NotImplementedError("Subclasses should implement this method.")

tetra_rp/stubs/__init__.py

@@ -0,0 +1,5 @@
+from .registry import stub_resource
+
+__all__ = [
+    "stub_resource",
+]

tetra_rp/stubs/live_serverless.py

@@ -0,0 +1,133 @@
+import ast
+import base64
+import inspect
+import textwrap
+import hashlib
+import traceback
+import cloudpickle
+import logging
+from ..core.resources import LiveServerless
+from ..protos.remote_execution import (
+    FunctionRequest,
+    FunctionResponse,
+    RemoteExecutorStub,
+)
+
+log = logging.getLogger(__name__)
+
+
+# global in memory cache, TODO: use a more robust cache in future
+_SERIALIZED_FUNCTION_CACHE = {}
+
+
+def get_function_source(func):
+    """Extract the function source code without the decorator."""
+    # Get the source code of the decorated function
+    source = inspect.getsource(func)
+
+    # Parse the source code
+    module = ast.parse(source)
+
+    # Find the function definition node
+    function_def = None
+    for node in ast.walk(module):
+        if isinstance(node, ast.FunctionDef) and node.name == func.__name__:
+            function_def = node
+            break
+
+    if not function_def:
+        raise ValueError(f"Could not find function definition for {func.__name__}")
+
+    # Get the line and column offsets
+    lineno = function_def.lineno - 1  # Line numbers are 1-based
+
+    # Split into lines and extract just the function part
+    lines = source.split("\n")
+    function_lines = lines[lineno:]
+
+    # Dedent to remove any extra indentation
+    function_source = textwrap.dedent("\n".join(function_lines))
+
+    # Return the function hash for cache key
+    source_hash = hashlib.sha256(function_source.encode("utf-8")).hexdigest()
+
+    return function_source, source_hash
+
+
+class LiveServerlessStub(RemoteExecutorStub):
+    """Adapter class to make Runpod endpoints look like gRPC stubs."""
+
+    def __init__(self, server: LiveServerless):
+        self.server = server
+
+    def prepare_request(self, func, dependencies, system_dependencies, *args, **kwargs):
+        source, src_hash = get_function_source(func)
+
+        request = {
+            "function_name": func.__name__,
+            "dependencies": dependencies,
+            "system_dependencies": system_dependencies,
+        }
+
+        # check if the function is already cached
+        if src_hash not in _SERIALIZED_FUNCTION_CACHE:
+            # Cache the serialized function
+            _SERIALIZED_FUNCTION_CACHE[src_hash] = source
+
+        request["function_code"] = _SERIALIZED_FUNCTION_CACHE[src_hash]
+
+        # Serialize arguments using cloudpickle
+        if args:
+            request["args"] = [
+                base64.b64encode(cloudpickle.dumps(arg)).decode("utf-8") for arg in args
+            ]
+        if kwargs:
+            request["kwargs"] = {
+                k: base64.b64encode(cloudpickle.dumps(v)).decode("utf-8")
+                for k, v in kwargs.items()
+            }
+
+        return FunctionRequest(**request)
+
+    def handle_response(self, response: FunctionResponse):
+        if not (response.success or response.error):
+            raise ValueError("Invalid response from server")
+
+        if response.stdout:
+            for line in response.stdout.splitlines():
+                log.info(f"Remote | {line}")
+
+        if response.success:
+            if response.result is None:
+                raise ValueError("Response result is None")
+            return cloudpickle.loads(base64.b64decode(response.result))
+        else:
+            raise Exception(f"Remote execution failed: {response.error}")
+
+    async def ExecuteFunction(
+        self, request: FunctionRequest, sync: bool = False
+    ) -> FunctionResponse:
+        try:
+            # Convert the gRPC request to Runpod format
+            payload = request.model_dump(exclude_none=True)
+
+            if sync:
+                job = await self.server.run_sync(payload)
+            else:
+                job = await self.server.run(payload)
+
+            if job.error:
+                return FunctionResponse(
+                    success=False,
+                    error=job.error,
+                    stdout=job.output.get("stdout", ""),
+                )
+
+            return FunctionResponse(**job.output)
+
+        except Exception as e:
+            error_traceback = traceback.format_exc()
+            return FunctionResponse(
+                success=False,
+                error=f"{str(e)}\n{error_traceback}",
+            )

tetra_rp/stubs/registry.py

@@ -0,0 +1,85 @@
+import logging
+from functools import singledispatch
+from .live_serverless import LiveServerlessStub
+from .serverless import ServerlessEndpointStub
+from ..core.resources import (
+    CpuServerlessEndpoint,
+    LiveServerless,
+    ServerlessEndpoint,
+)
+
+
+log = logging.getLogger(__name__)
+
+
+@singledispatch
+def stub_resource(resource, **extra):
+    async def fallback(*args, **kwargs):
+        return {"error": f"Cannot stub {resource.__class__.__name__}."}
+
+    return fallback
+
+
+@stub_resource.register(LiveServerless)
+def _(resource, **extra):
+    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)
+
+    return stubbed_resource
+
+
+@stub_resource.register(ServerlessEndpoint)
+def _(resource, **extra):
+    async def stubbed_resource(
+        func, dependencies, system_dependencies, *args, **kwargs
+    ) -> dict:
+        if args == (None,):
+            # cleanup: when the function is called with no args
+            args = []
+
+        if dependencies or system_dependencies:
+            log.warning(
+                "Dependencies are not supported for ServerlessEndpoint. "
+                "They will be ignored."
+            )
+
+        stub = ServerlessEndpointStub(resource)
+        payload = stub.prepare_payload(func, *args, **kwargs)
+        response = await stub.execute(payload, sync=extra.get("sync", False))
+        return stub.handle_response(response)
+
+    return stubbed_resource
+
+
+@stub_resource.register(CpuServerlessEndpoint)
+def _(resource, **extra):
+    async def stubbed_resource(
+        func, dependencies, system_dependencies, *args, **kwargs
+    ) -> dict:
+        if args == (None,):
+            # cleanup: when the function is called with no args
+            args = []
+
+        if dependencies or system_dependencies:
+            log.warning(
+                "Dependencies are not supported for CpuServerlessEndpoint. "
+                "They will be ignored."
+            )
+
+        stub = ServerlessEndpointStub(resource)
+        payload = stub.prepare_payload(func, *args, **kwargs)
+        response = await stub.execute(payload, sync=extra.get("sync", False))
+        return stub.handle_response(response)
+
+    return stubbed_resource

tetra_rp/stubs/serverless.py

@@ -0,0 +1,30 @@
+from ..core.resources import ServerlessEndpoint, JobOutput
+
+
+class ServerlessEndpointStub:
+    """Adapter class to make Runpod endpoints requests."""
+
+    def __init__(self, server: ServerlessEndpoint):
+        self.server = server
+
+    def prepare_payload(self, func, *args, **kwargs) -> dict:
+        return func(*args, **kwargs)
+
+    async def execute(self, payload: dict, sync: bool = False) -> JobOutput:
+        """
+        Executes a serverless endpoint request with the payload.
+        Returns a JobOutput object.
+        """
+        if sync:
+            return await self.server.run_sync(payload)
+        else:
+            return await self.server.run(payload)
+
+    def handle_response(self, response: JobOutput):
+        if response.output:
+            return response.output
+
+        if response.error:
+            raise Exception(f"Remote execution failed: {response.error}")
+
+        raise ValueError("Invalid response from server")