tetra-rp 0.17.1__py3-none-any.whl

tetra_rp/cli/utils/skeleton_template/workers/gpu/endpoint.py

@@ -0,0 +1,62 @@
+from tetra_rp import remote, LiveServerless, GpuGroup
+
+
+gpu_config = LiveServerless(
+    name="gpu_worker",
+    gpus=[GpuGroup.ANY],
+    workersMin=0,
+    workersMax=3,
+    idleTimeout=5,
+)
+
+
+@remote(resource_config=gpu_config, dependencies=["torch"])
+async def gpu_hello(input_data: dict) -> dict:
+    """Simple GPU worker example with GPU detection."""
+    import platform
+    from datetime import datetime
+
+    try:
+        import torch
+
+        gpu_available = torch.cuda.is_available()
+        if gpu_available:
+            gpu_name = torch.cuda.get_device_name(0)
+            gpu_count = torch.cuda.device_count()
+            gpu_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3)
+        else:
+            gpu_name = "No GPU detected"
+            gpu_count = 0
+            gpu_memory = 0
+    except Exception as e:
+        gpu_available = False
+        gpu_name = f"Error detecting GPU: {str(e)}"
+        gpu_count = 0
+        gpu_memory = 0
+
+    message = input_data.get("message", "Hello from GPU worker!")
+
+    return {
+        "status": "success",
+        "message": message,
+        "worker_type": "GPU",
+        "gpu_info": {
+            "available": gpu_available,
+            "name": gpu_name,
+            "count": gpu_count,
+            "memory_gb": round(gpu_memory, 2) if gpu_memory else 0,
+        },
+        "timestamp": datetime.now().isoformat(),
+        "platform": platform.system(),
+        "python_version": platform.python_version(),
+    }
+
+
+# Test locally with: python -m workers.gpu.endpoint
+if __name__ == "__main__":
+    import asyncio
+
+    test_payload = {"message": "Testing GPU worker"}
+    print(f"Testing GPU worker with payload: {test_payload}")
+    result = asyncio.run(gpu_hello(test_payload))
+    print(f"Result: {result}")

tetra_rp/client.py

@@ -0,0 +1,128 @@
+import os
+import inspect
+import logging
+from functools import wraps
+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: Optional[List[str]] = None,
+    system_dependencies: Optional[List[str]] = None,
+    accelerate_downloads: bool = True,
+    local: bool = False,
+    **extra,
+):
+    """
+    Decorator to enable dynamic resource provisioning and dependency management for serverless functions.
+
+    This decorator allows a function to be executed in a remote serverless environment, with support for
+    dynamic resource provisioning and installation of required dependencies. It can also bypass remote
+    execution entirely for local testing.
+
+    Supports both sync and async function definitions:
+    - `def my_function(...)` - Regular synchronous function
+    - `async def my_function(...)` - Asynchronous function
+
+    In both cases, the decorated function returns an awaitable that must be called with `await`.
+
+    Args:
+        resource_config (ServerlessResource): Configuration object specifying the serverless resource
+            to be provisioned or used. Not used when local=True.
+        dependencies (List[str], optional): A list of pip package names to be installed in the remote
+            environment before executing the function. Not used when local=True. Defaults to None.
+        system_dependencies (List[str], optional): A list of system packages to be installed in the remote
+            environment before executing the function. Not used when local=True. Defaults to None.
+        accelerate_downloads (bool, optional): Enable download acceleration for dependencies and models.
+            Only applies to remote execution. Defaults to True.
+        local (bool, optional): Execute function/class locally instead of provisioning remote servers.
+            Returns the unwrapped function/class for direct local execution. Users must ensure all required
+            dependencies are already installed in their local environment. Defaults to False.
+        extra (dict, optional): Additional parameters for the execution of the resource. Defaults to an empty dict.
+
+    Returns:
+        Callable: A decorator that wraps the target function, enabling remote execution with the specified
+        resource configuration and dependencies, or returns the unwrapped function/class for local execution.
+
+    Example:
+    ```python
+        # Async function (recommended style)
+        @remote(
+            resource_config=my_resource_config,
+            dependencies=["torch>=2.0.0"],
+        )
+        async def gpu_task(data: dict) -> dict:
+            import torch
+            # GPU processing here
+            return {"result": "processed"}
+
+        # Sync function (also supported)
+        @remote(
+            resource_config=my_resource_config,
+            dependencies=["pandas>=2.0.0"],
+        )
+        def cpu_task(data: dict) -> dict:
+            import pandas as pd
+            # CPU processing here
+            return {"result": "processed"}
+
+        # Local execution (testing/development)
+        @remote(
+            resource_config=my_resource_config,
+            dependencies=["numpy", "pandas"],  # Only used for remote execution
+            local=True,
+        )
+        async def my_test_function(data):
+            # Runs locally - dependencies must be pre-installed
+            pass
+    ```
+    """
+
+    def decorator(func_or_class):
+        if os.getenv("RUNPOD_POD_ID") or os.getenv("RUNPOD_ENDPOINT_ID"):
+            # Worker mode when running on RunPod platform
+            return func_or_class
+
+        # Local execution mode - execute without provisioning remote servers
+        if local:
+            return func_or_class
+
+        # Remote execution mode
+        if inspect.isclass(func_or_class):
+            # Handle class decoration
+            return create_remote_class(
+                func_or_class,
+                resource_config,
+                dependencies,
+                system_dependencies,
+                accelerate_downloads,
+                extra,
+            )
+        else:
+            # Handle function decoration
+            @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_or_class,
+                    dependencies,
+                    system_dependencies,
+                    accelerate_downloads,
+                    *args,
+                    **kwargs,
+                )
+
+            return wrapper
+
+    return decorator

tetra_rp/config.py

@@ -0,0 +1,29 @@
+"""Configuration management for tetra-rp CLI."""
+
+from pathlib import Path
+from typing import NamedTuple
+
+
+class TetraPaths(NamedTuple):
+    """Paths for tetra-rp configuration and data."""
+
+    tetra_dir: Path
+    config_file: Path
+    deployments_file: Path
+
+    def ensure_tetra_dir(self) -> None:
+        """Ensure the .tetra directory exists."""
+        self.tetra_dir.mkdir(exist_ok=True)
+
+
+def get_paths() -> TetraPaths:
+    """Get standardized paths for tetra-rp configuration."""
+    tetra_dir = Path.cwd() / ".tetra"
+    config_file = tetra_dir / "config.json"
+    deployments_file = tetra_dir / "deployments.json"
+
+    return TetraPaths(
+        tetra_dir=tetra_dir,
+        config_file=config_file,
+        deployments_file=deployments_file,
+    )

tetra_rp/core/api/__init__.py

@@ -0,0 +1,6 @@
+from .runpod import RunpodGraphQLClient, RunpodRestClient
+
+__all__ = [
+    "RunpodGraphQLClient",
+    "RunpodRestClient",
+]

tetra_rp/core/api/runpod.py

@@ -0,0 +1,319 @@
+"""
+Direct GraphQL communication with Runpod API.
+Bypasses the outdated runpod-python SDK limitations.
+"""
+
+import json
+import logging
+import os
+from typing import Any, Dict, Optional
+
+import aiohttp
+
+from tetra_rp.core.exceptions import RunpodAPIKeyError
+
+log = logging.getLogger(__name__)
+
+RUNPOD_API_BASE_URL = os.environ.get("RUNPOD_API_BASE_URL", "https://api.runpod.io")
+RUNPOD_REST_API_URL = os.environ.get("RUNPOD_REST_API_URL", "https://rest.runpod.io/v1")
+
+
+class RunpodGraphQLClient:
+    """
+    Runpod GraphQL client for Runpod API.
+    Communicates directly with Runpod's GraphQL endpoint without SDK limitations.
+    """
+
+    GRAPHQL_URL = f"{RUNPOD_API_BASE_URL}/graphql"
+
+    def __init__(self, api_key: Optional[str] = None):
+        self.api_key = api_key or os.getenv("RUNPOD_API_KEY")
+        if not self.api_key:
+            raise RunpodAPIKeyError()
+
+        self.session: Optional[aiohttp.ClientSession] = None
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        """Get or create an aiohttp session."""
+        if self.session is None or self.session.closed:
+            timeout = aiohttp.ClientTimeout(total=300)  # 5 minute timeout
+            self.session = aiohttp.ClientSession(
+                timeout=timeout,
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                },
+            )
+        return self.session
+
+    async def _execute_graphql(
+        self, query: str, variables: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """Execute a GraphQL query/mutation."""
+        session = await self._get_session()
+
+        payload = {"query": query, "variables": variables or {}}
+
+        log.debug(f"GraphQL Query: {query}")
+        log.debug(f"GraphQL Variables: {json.dumps(variables, indent=2)}")
+
+        try:
+            async with session.post(self.GRAPHQL_URL, json=payload) as response:
+                response_data = await response.json()
+
+                log.debug(f"GraphQL Response Status: {response.status}")
+                log.debug(f"GraphQL Response: {json.dumps(response_data, indent=2)}")
+
+                if response.status >= 400:
+                    raise Exception(
+                        f"GraphQL request failed: {response.status} - {response_data}"
+                    )
+
+                if "errors" in response_data:
+                    errors = response_data["errors"]
+                    error_msg = "; ".join([e.get("message", str(e)) for e in errors])
+                    raise Exception(f"GraphQL errors: {error_msg}")
+
+                return response_data.get("data", {})
+
+        except aiohttp.ClientError as e:
+            log.error(f"HTTP client error: {e}")
+            raise Exception(f"HTTP request failed: {e}")
+
+    async def create_endpoint(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Create a serverless endpoint using direct GraphQL mutation.
+        Supports both GPU and CPU endpoints with full field support.
+        """
+        # GraphQL mutation for saveEndpoint (based on actual schema)
+        mutation = """
+        mutation saveEndpoint($input: EndpointInput!) {
+            saveEndpoint(input: $input) {
+                aiKey
+                gpuIds
+                id
+                idleTimeout
+                locations
+                name
+                networkVolumeId
+                scalerType
+                scalerValue
+                templateId
+                type
+                userId
+                version
+                workersMax
+                workersMin
+                workersStandby
+                workersPFBTarget
+                gpuCount
+                allowedCudaVersions
+                executionTimeoutMs
+                instanceIds
+                activeBuildid
+                idePodId
+            }
+        }
+        """
+
+        variables = {"input": input_data}
+
+        log.debug(
+            f"Creating endpoint with GraphQL: {input_data.get('name', 'unnamed')}"
+        )
+
+        result = await self._execute_graphql(mutation, variables)
+
+        if "saveEndpoint" not in result:
+            raise Exception("Unexpected GraphQL response structure")
+
+        endpoint_data = result["saveEndpoint"]
+        log.info(
+            f"Created endpoint: {endpoint_data.get('id', 'unknown')} - {endpoint_data.get('name', 'unnamed')}"
+        )
+
+        return endpoint_data
+
+    async def get_cpu_types(self) -> Dict[str, Any]:
+        """Get available CPU types."""
+        query = """
+        query getCpuTypes {
+            cpuTypes {
+                id
+                displayName
+                manufacturer
+                cores
+                threadsPerCore
+                groupId
+            }
+        }
+        """
+
+        result = await self._execute_graphql(query)
+        return result.get("cpuTypes", [])
+
+    async def get_gpu_types(
+        self, gpu_filter: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """Get available GPU types."""
+        query = """
+        query getGpuTypes($input: GpuTypeFilter) {
+            gpuTypes(input: $input) {
+                id
+                displayName
+                manufacturer
+                memoryInGb
+                cudaCores
+                secureCloud
+                communityCloud
+                securePrice
+                communityPrice
+                communitySpotPrice
+                secureSpotPrice
+                maxGpuCount
+                maxGpuCountCommunityCloud
+                maxGpuCountSecureCloud
+                minPodGpuCount
+                nodeGroupGpuSizes
+                throughput
+            }
+        }
+        """
+
+        variables = {"input": gpu_filter} if gpu_filter else {}
+        result = await self._execute_graphql(query, variables)
+        return result.get("gpuTypes", [])
+
+    async def get_endpoint(self, endpoint_id: str) -> Dict[str, Any]:
+        """Get endpoint details."""
+        # Note: The schema doesn't show a specific endpoint query
+        # This would need to be implemented if such query exists
+        raise NotImplementedError("Get endpoint query not available in current schema")
+
+    async def delete_endpoint(self, endpoint_id: str) -> Dict[str, Any]:
+        """Delete a serverless endpoint."""
+        mutation = """
+        mutation deleteEndpoint($id: String!) {
+            deleteEndpoint(id: $id)
+        }
+        """
+
+        variables = {"id": endpoint_id}
+        log.info(f"Deleting endpoint: {endpoint_id}")
+
+        result = await self._execute_graphql(mutation, variables)
+        return {"success": result.get("deleteEndpoint") is not None}
+
+    async def close(self):
+        """Close the HTTP session."""
+        if self.session and not self.session.closed:
+            await self.session.close()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
+
+class RunpodRestClient:
+    """
+    Runpod REST client for Runpod API.
+    Provides methods to interact with Runpod's REST endpoints.
+    """
+
+    def __init__(self, api_key: Optional[str] = None):
+        self.api_key = api_key or os.getenv("RUNPOD_API_KEY")
+        if not self.api_key:
+            raise RunpodAPIKeyError()
+
+        self.session: Optional[aiohttp.ClientSession] = None
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        """Get or create an aiohttp session."""
+        if self.session is None or self.session.closed:
+            timeout = aiohttp.ClientTimeout(total=300)  # 5 minute timeout
+            self.session = aiohttp.ClientSession(
+                timeout=timeout,
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                },
+            )
+        return self.session
+
+    async def _execute_rest(
+        self, method: str, url: str, data: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """Execute a REST API request."""
+        session = await self._get_session()
+
+        log.debug(f"REST Request: {method} {url}")
+        log.debug(f"REST Data: {json.dumps(data, indent=2) if data else 'None'}")
+
+        try:
+            async with session.request(method, url, json=data) as response:
+                response_data = await response.json()
+
+                log.debug(f"REST Response Status: {response.status}")
+                log.debug(f"REST Response: {json.dumps(response_data, indent=2)}")
+
+                if response.status >= 400:
+                    raise Exception(
+                        f"REST request failed: {response.status} - {response_data}"
+                    )
+
+                return response_data
+
+        except aiohttp.ClientError as e:
+            log.error(f"HTTP client error: {e}")
+            raise Exception(f"HTTP request failed: {e}")
+
+    async def create_network_volume(self, payload: Dict[str, Any]) -> Dict[str, Any]:
+        """Create a network volume in Runpod."""
+        log.debug(f"Creating network volume: {payload.get('name', 'unnamed')}")
+
+        result = await self._execute_rest(
+            "POST", f"{RUNPOD_REST_API_URL}/networkvolumes", payload
+        )
+
+        log.info(
+            f"Created network volume: {result.get('id', 'unknown')} - {result.get('name', 'unnamed')}"
+        )
+
+        return result
+
+    async def list_network_volumes(self) -> Dict[str, Any]:
+        """
+        List all network volumes in Runpod.
+
+        Returns:
+            List of network volume objects or dict containing networkVolumes key.
+            The API may return either format depending on version.
+        """
+        log.debug("Listing network volumes")
+
+        result = await self._execute_rest(
+            "GET", f"{RUNPOD_REST_API_URL}/networkvolumes"
+        )
+
+        # Handle both list and dict responses
+        if isinstance(result, list):
+            volume_count = len(result)
+        else:
+            volume_count = len(result.get("networkVolumes", []))
+
+        log.debug(f"Listed {volume_count} network volumes")
+
+        return result
+
+    async def close(self):
+        """Close the HTTP session."""
+        if self.session and not self.session.closed:
+            await self.session.close()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()

tetra_rp/core/exceptions.py

@@ -0,0 +1,50 @@
+"""Custom exceptions for tetra_rp.
+
+Provides clear, actionable error messages for common failure scenarios.
+"""
+
+
+class RunpodAPIKeyError(Exception):
+    """Raised when RUNPOD_API_KEY environment variable is missing or invalid.
+
+    This exception provides helpful guidance on how to obtain and configure
+    the API key required for remote execution and deployment features.
+    """
+
+    def __init__(self, message: str | None = None):
+        """Initialize with optional custom message.
+
+        Args:
+            message: Optional custom error message. If not provided, uses default.
+        """
+        if message is None:
+            message = self._default_message()
+        super().__init__(message)
+
+    @staticmethod
+    def _default_message() -> str:
+        """Generate default error message with setup instructions.
+
+        Returns:
+            Formatted error message with actionable steps.
+        """
+        return """RUNPOD_API_KEY environment variable is required but not set.
+
+To use Flash remote execution features, you need a Runpod API key.
+
+Get your API key:
+  https://docs.runpod.io/get-started/api-keys
+
+Set your API key using one of these methods:
+
+  1. Environment variable:
+     export RUNPOD_API_KEY=your_api_key_here
+
+  2. In your project's .env file:
+     echo "RUNPOD_API_KEY=your_api_key_here" >> .env
+
+  3. In your shell profile (~/.bashrc, ~/.zshrc):
+     echo 'export RUNPOD_API_KEY=your_api_key_here' >> ~/.bashrc
+
+Note: If you created a .env file, make sure it's in your current directory
+or project root where Flash can find it."""

tetra_rp/core/resources/__init__.py

@@ -0,0 +1,37 @@
+from .base import BaseResource, DeployableResource
+from .cpu import CpuInstanceType
+from .gpu import GpuGroup, GpuType, GpuTypeDetail
+from .resource_manager import ResourceManager
+from .live_serverless import LiveServerless, CpuLiveServerless
+from .serverless import (
+    ServerlessResource,
+    ServerlessEndpoint,
+    JobOutput,
+    CudaVersion,
+    ServerlessType,
+)
+from .serverless_cpu import CpuServerlessEndpoint
+from .template import PodTemplate
+from .network_volume import NetworkVolume, DataCenter
+
+
+__all__ = [
+    "BaseResource",
+    "CpuInstanceType",
+    "CpuLiveServerless",
+    "CpuServerlessEndpoint",
+    "CudaVersion",
+    "DataCenter",
+    "DeployableResource",
+    "GpuGroup",
+    "GpuType",
+    "GpuTypeDetail",
+    "JobOutput",
+    "LiveServerless",
+    "ResourceManager",
+    "ServerlessResource",
+    "ServerlessEndpoint",
+    "ServerlessType",
+    "PodTemplate",
+    "NetworkVolume",
+]

tetra_rp/core/resources/base.py

@@ -0,0 +1,47 @@
+import hashlib
+from abc import ABC, abstractmethod
+from typing import Optional
+from pydantic import BaseModel, ConfigDict
+
+
+class BaseResource(BaseModel):
+    """Base class for all resources."""
+
+    model_config = ConfigDict(
+        validate_by_name=True,
+        validate_default=True,
+        serialize_by_alias=True,
+    )
+
+    id: Optional[str] = None
+
+    @property
+    def resource_id(self) -> str:
+        """Unique resource ID based on configuration."""
+        resource_type = self.__class__.__name__
+        config_str = self.model_dump_json(exclude_none=True)
+        hash_obj = hashlib.md5(f"{resource_type}:{config_str}".encode())
+        return f"{resource_type}_{hash_obj.hexdigest()}"
+
+
+class DeployableResource(BaseResource, ABC):
+    """Base class for deployable resources."""
+
+    def __str__(self) -> str:
+        return f"{self.__class__.__name__}"
+
+    @property
+    @abstractmethod
+    def url(self) -> str:
+        """Public URL of the resource."""
+        raise NotImplementedError("Subclasses should implement this method.")
+
+    @abstractmethod
+    def is_deployed(self) -> bool:
+        """Check the resource if it's still valid or available."""
+        raise NotImplementedError("Subclasses should implement this method.")
+
+    @abstractmethod
+    async def deploy(self) -> "DeployableResource":
+        """Deploy the resource."""
+        raise NotImplementedError("Subclasses should implement this method.")

tetra_rp/core/resources/cloud.py

@@ -0,0 +1,4 @@
+import os
+import runpod
+
+runpod.api_key = os.getenv("RUNPOD_API_KEY")

tetra_rp/core/resources/constants.py

@@ -0,0 +1,4 @@
+import os
+
+CONSOLE_BASE_URL = os.environ.get("CONSOLE_BASE_URL", "https://console.runpod.io")
+CONSOLE_URL = f"{CONSOLE_BASE_URL}/serverless/user/endpoint/%s"