tetra-rp 0.17.1__py3-none-any.whl

tetra_rp/core/resources/cpu.py

@@ -0,0 +1,146 @@
+from enum import Enum
+from typing import List, Optional
+
+
+class CpuInstanceType(str, Enum):
+    """Valid CPU instance types.
+
+    Format: {generation}{type}-{vcpu}-{memory_gb}
+    Based on Runpod backend validation logic:
+    - memoryInGb = vcpuCount * flavor.ramMultiplier
+
+    RAM Multipliers (DEV environment):
+    - cpu3g: 4.0 (1 vCPU = 4GB, 2 vCPU = 8GB, etc.)
+    - cpu3c: 2.0 (1 vCPU = 2GB, 2 vCPU = 4GB, etc.)
+    - cpu5c: 2.0 (1 vCPU = 2GB, 2 vCPU = 4GB, etc.)
+    - cpu5g: Not available
+    """
+
+    ANY = "any"
+    """Any CPU"""
+
+    # 3rd Generation General Purpose (RAM multiplier: 4.0)
+
+    CPU3G_1_4 = "cpu3g-1-4"
+    """1 vCPU, 4GB RAM, max 10GB container disk"""
+
+    CPU3G_2_8 = "cpu3g-2-8"
+    """2 vCPU, 8GB RAM, max 20GB container disk"""
+
+    CPU3G_4_16 = "cpu3g-4-16"
+    """4 vCPU, 16GB RAM, max 40GB container disk"""
+
+    CPU3G_8_32 = "cpu3g-8-32"
+    """8 vCPU, 32GB RAM, max 80GB container disk"""
+
+    # 3rd Generation Compute-Optimized (RAM multiplier: 2.0)
+
+    CPU3C_1_2 = "cpu3c-1-2"
+    """1 vCPU, 2GB RAM, max 10GB container disk"""
+
+    CPU3C_2_4 = "cpu3c-2-4"
+    """2 vCPU, 4GB RAM, max 20GB container disk"""
+
+    CPU3C_4_8 = "cpu3c-4-8"
+    """4 vCPU, 8GB RAM, max 40GB container disk"""
+
+    CPU3C_8_16 = "cpu3c-8-16"
+    """8 vCPU, 16GB RAM, max 80GB container disk"""
+
+    # 5th Generation Compute-Optimized (RAM multiplier: 2.0)
+
+    CPU5C_1_2 = "cpu5c-1-2"
+    """1 vCPU, 2GB RAM, max 15GB container disk"""
+
+    CPU5C_2_4 = "cpu5c-2-4"
+    """2 vCPU, 4GB RAM, max 30GB container disk"""
+
+    CPU5C_4_8 = "cpu5c-4-8"
+    """4 vCPU, 8GB RAM, max 60GB container disk"""
+
+    CPU5C_8_16 = "cpu5c-8-16"
+    """8 vCPU, 16GB RAM, max 120GB container disk"""
+
+    @classmethod
+    def all(cls) -> List["CpuInstanceType"]:
+        """Returns all CPU Instance Types."""
+        return [c for c in cls if c != cls.ANY]
+
+
+def calculate_max_disk_size(instance_type: CpuInstanceType) -> int:
+    """
+    Calculate the maximum container disk size for a CPU instance type.
+
+    Formula:
+    - CPU3G/CPU3C: vCPU count × 10GB
+    - CPU5C: vCPU count × 15GB
+
+    Args:
+        instance_type: CPU instance type enum
+
+    Returns:
+        Maximum container disk size in GB
+
+    Example:
+        >>> calculate_max_disk_size(CpuInstanceType.CPU3G_1_4)
+        10
+        >>> calculate_max_disk_size(CpuInstanceType.CPU5C_2_4)
+        30
+    """
+    # Parse the instance type string to extract vCPU count
+    # Format: "cpu{generation}{type}-{vcpu}-{memory}"
+    instance_str = instance_type.value
+    parts = instance_str.split("-")
+
+    if len(parts) != 3:
+        raise ValueError(f"Invalid instance type format: {instance_str}")
+
+    vcpu_count = int(parts[1])
+
+    # Determine disk multiplier based on generation
+    if instance_str.startswith("cpu5c"):
+        disk_multiplier = 15  # CPU5C: 15GB per vCPU
+    elif instance_str.startswith(("cpu3g", "cpu3c")):
+        disk_multiplier = 10  # CPU3G/CPU3C: 10GB per vCPU
+    else:
+        raise ValueError(f"Unknown CPU generation/type: {instance_str}")
+
+    return vcpu_count * disk_multiplier
+
+
+# CPU Instance Type Disk Limits (calculated programmatically)
+CPU_INSTANCE_DISK_LIMITS = {
+    instance_type: calculate_max_disk_size(instance_type)
+    for instance_type in CpuInstanceType
+    if instance_type != CpuInstanceType.ANY
+}
+
+
+def get_max_disk_size_for_instances(
+    instance_types: Optional[List[CpuInstanceType]],
+) -> Optional[int]:
+    """
+    Calculate the maximum container disk size for a list of CPU instance types.
+
+    Returns the minimum disk limit across all instance types to ensure compatibility
+    with all specified instances.
+
+    Args:
+        instance_types: List of CPU instance types, or None
+
+    Returns:
+        Maximum allowed disk size in GB, or None if no CPU instances specified
+
+    Example:
+        >>> get_max_disk_size_for_instances([CpuInstanceType.CPU3G_1_4])
+        10
+        >>> get_max_disk_size_for_instances([CpuInstanceType.CPU3G_1_4, CpuInstanceType.CPU3G_2_8])
+        10
+    """
+    if not instance_types:
+        return None
+
+    disk_limits = [
+        CPU_INSTANCE_DISK_LIMITS[instance_type] for instance_type in instance_types
+    ]
+    return min(disk_limits)

tetra_rp/core/resources/environment.py

@@ -0,0 +1,41 @@
+from typing import Dict, Optional
+from dotenv import dotenv_values
+
+
+class EnvironmentVars:
+    def __init__(self):
+        # Store environment variables from .env file
+        self.env = self._load_env()
+
+    def _load_env(self) -> Dict[str, str]:
+        """
+        Loads environment variables specifically from the .env file
+        and returns them as a dictionary.
+
+        Returns:
+            Dict[str, str]: Dictionary containing environment variables from .env file
+        """
+        # Use dotenv_values instead of load_dotenv to get only variables from .env
+        return dict(dotenv_values())
+
+    def get_env(self) -> Dict[str, str]:
+        """
+        Returns the dictionary of environment variables.
+
+        Returns:
+            Dict[str, str]: Dictionary containing environment variables
+        """
+        return self.env
+
+    def get_value(self, key: str, default: str = None) -> Optional[str]:
+        """
+        Gets a specific environment variable by key.
+
+        Args:
+            key (str): The environment variable key
+            default (str, optional): Default value if key doesn't exist
+
+        Returns:
+            Optional[str]: Value of the environment variable or default
+        """
+        return self.env.get(key, default)

tetra_rp/core/resources/gpu.py

@@ -0,0 +1,68 @@
+from typing import Optional, List
+from pydantic import BaseModel
+from enum import Enum
+
+
+class GpuLowestPrice(BaseModel):
+    minimumBidPrice: Optional[float] = None
+    uninterruptablePrice: Optional[float] = None
+
+
+class GpuType(BaseModel):
+    id: str
+    displayName: str
+    memoryInGb: int
+
+
+class GpuTypeDetail(GpuType):
+    communityCloud: Optional[bool] = None
+    communityPrice: Optional[float] = None
+    communitySpotPrice: Optional[float] = None
+    cudaCores: Optional[int] = None
+    lowestPrice: Optional[GpuLowestPrice] = None
+    manufacturer: Optional[str] = None
+    maxGpuCount: Optional[int] = None
+    oneMonthPrice: Optional[float] = None
+    oneWeekPrice: Optional[float] = None
+    secureCloud: Optional[bool] = None
+    securePrice: Optional[float] = None
+    secureSpotPrice: Optional[float] = None
+    threeMonthPrice: Optional[float] = None
+
+
+# TODO: this should be fetched from an API
+class GpuGroup(Enum):
+    ANY = "any"
+    """Any GPU"""
+
+    ADA_24 = "ADA_24"
+    """NVIDIA GeForce RTX 4090"""
+
+    ADA_32_PRO = "ADA_32_PRO"
+    """NVIDIA GeForce RTX 5090"""
+
+    ADA_48_PRO = "ADA_48_PRO"
+    """NVIDIA RTX 6000 Ada Generation, NVIDIA L40, NVIDIA L40S"""
+
+    ADA_80_PRO = "ADA_80_PRO"
+    """NVIDIA H100 PCIe, NVIDIA H100 80GB HBM3, NVIDIA H100 NVL"""
+
+    AMPERE_16 = "AMPERE_16"
+    """NVIDIA RTX A4000, NVIDIA RTX A4500, NVIDIA RTX 4000 Ada Generation, NVIDIA RTX 2000 Ada Generation"""
+
+    AMPERE_24 = "AMPERE_24"
+    """NVIDIA RTX A5000, NVIDIA L4, NVIDIA GeForce RTX 3090"""
+
+    AMPERE_48 = "AMPERE_48"
+    """NVIDIA A40, NVIDIA RTX A6000"""
+
+    AMPERE_80 = "AMPERE_80"
+    """NVIDIA A100 80GB PCIe, NVIDIA A100-SXM4-80GB"""
+
+    HOPPER_141 = "HOPPER_141"
+    """NVIDIA H200"""
+
+    @classmethod
+    def all(cls) -> List["GpuGroup"]:
+        """Returns all GPU groups."""
+        return [cls.AMPERE_48] + [g for g in cls if g != cls.ANY]

tetra_rp/core/resources/live_serverless.py

@@ -0,0 +1,62 @@
+# Ship serverless code as you write it. No builds, no deploys — just run.
+import os
+from pydantic import model_validator
+from .serverless import ServerlessEndpoint
+from .serverless_cpu import CpuServerlessEndpoint
+
+TETRA_IMAGE_TAG = os.environ.get("TETRA_IMAGE_TAG", "latest")
+TETRA_GPU_IMAGE = os.environ.get(
+    "TETRA_GPU_IMAGE", f"runpod/tetra-rp:{TETRA_IMAGE_TAG}"
+)
+TETRA_CPU_IMAGE = os.environ.get(
+    "TETRA_CPU_IMAGE", f"runpod/tetra-rp-cpu:{TETRA_IMAGE_TAG}"
+)
+
+
+class LiveServerlessMixin:
+    """Common mixin for live serverless endpoints that locks the image."""
+
+    @property
+    def _live_image(self) -> str:
+        """Override in subclasses to specify the locked image."""
+        raise NotImplementedError("Subclasses must define _live_image")
+
+    @property
+    def imageName(self):
+        # Lock imageName to specific image
+        return self._live_image
+
+    @imageName.setter
+    def imageName(self, value):
+        # Prevent manual setting of imageName
+        pass
+
+
+class LiveServerless(LiveServerlessMixin, ServerlessEndpoint):
+    """GPU-only live serverless endpoint."""
+
+    @property
+    def _live_image(self) -> str:
+        return TETRA_GPU_IMAGE
+
+    @model_validator(mode="before")
+    @classmethod
+    def set_live_serverless_template(cls, data: dict):
+        """Set default GPU image for Live Serverless."""
+        data["imageName"] = TETRA_GPU_IMAGE
+        return data
+
+
+class CpuLiveServerless(LiveServerlessMixin, CpuServerlessEndpoint):
+    """CPU-only live serverless endpoint with automatic disk sizing."""
+
+    @property
+    def _live_image(self) -> str:
+        return TETRA_CPU_IMAGE
+
+    @model_validator(mode="before")
+    @classmethod
+    def set_live_serverless_template(cls, data: dict):
+        """Set default CPU image for Live Serverless."""
+        data["imageName"] = TETRA_CPU_IMAGE
+        return data

tetra_rp/core/resources/network_volume.py

@@ -0,0 +1,148 @@
+import hashlib
+import logging
+from enum import Enum
+from typing import Optional
+
+from pydantic import (
+    Field,
+    field_serializer,
+)
+
+from ..api.runpod import RunpodRestClient
+from .base import DeployableResource
+from .constants import CONSOLE_BASE_URL
+
+log = logging.getLogger(__name__)
+
+
+class DataCenter(str, Enum):
+    """
+    Enum representing available data centers for network volumes.
+    #TODO: Add more data centers as needed. Lock this to the available data center.
+    """
+
+    EU_RO_1 = "EU-RO-1"
+
+
+class NetworkVolume(DeployableResource):
+    """
+    NetworkVolume resource for creating and managing Runpod network volumes.
+
+    This class handles the creation, deployment, and management of network volumes
+    that can be attached to serverless resources. Supports idempotent deployment
+    where multiple volumes with the same name will reuse existing volumes.
+
+    """
+
+    # Internal fixed value
+    dataCenterId: DataCenter = Field(default=DataCenter.EU_RO_1, frozen=True)
+
+    id: Optional[str] = Field(default=None)
+    name: str
+    size: Optional[int] = Field(default=100, gt=0)  # Size in GB
+
+    def __str__(self) -> str:
+        return f"{self.__class__.__name__}:{self.id}"
+
+    @property
+    def resource_id(self) -> str:
+        """Unique resource ID based on name and datacenter for idempotent behavior."""
+        # Use name + datacenter to ensure idempotence
+        resource_type = self.__class__.__name__
+        config_key = f"{self.name}:{self.dataCenterId.value}"
+        hash_obj = hashlib.md5(f"{resource_type}:{config_key}".encode())
+        return f"{resource_type}_{hash_obj.hexdigest()}"
+
+    @field_serializer("dataCenterId")
+    def serialize_data_center_id(self, value: Optional[DataCenter]) -> Optional[str]:
+        """Convert DataCenter enum to string."""
+        return value.value if value is not None else None
+
+    @property
+    def is_created(self) -> bool:
+        "Returns True if the network volume already exists."
+        return self.id is not None
+
+    @property
+    def url(self) -> str:
+        """
+        Returns the URL for the network volume resource.
+        """
+        if not self.id:
+            raise ValueError("Network volume ID is not set")
+        return f"{CONSOLE_BASE_URL}/user/storage"
+
+    def is_deployed(self) -> bool:
+        """
+        Checks if the network volume resource is deployed and available.
+        """
+        return self.id is not None
+
+    def _normalize_volumes_response(self, volumes_response) -> list:
+        """Normalize API response to list format."""
+        if isinstance(volumes_response, list):
+            return volumes_response
+        return volumes_response.get("networkVolumes", [])
+
+    def _find_matching_volume(self, existing_volumes: list) -> Optional[dict]:
+        """Find existing volume matching name and datacenter."""
+        for volume_data in existing_volumes:
+            if (
+                volume_data.get("name") == self.name
+                and volume_data.get("dataCenterId") == self.dataCenterId.value
+            ):
+                return volume_data
+        return None
+
+    async def _find_existing_volume(self, client) -> Optional["NetworkVolume"]:
+        """Check for existing volume with same name and datacenter."""
+        if not self.name:
+            return None
+
+        log.debug(f"Checking for existing network volume with name: {self.name}")
+        volumes_response = await client.list_network_volumes()
+        existing_volumes = self._normalize_volumes_response(volumes_response)
+
+        if matching_volume := self._find_matching_volume(existing_volumes):
+            log.info(
+                f"Found existing network volume: {matching_volume.get('id')} with name '{self.name}'"
+            )
+            # Update our instance with the existing volume's ID
+            self.id = matching_volume.get("id")
+            return self
+
+        return None
+
+    async def _create_new_volume(self, client) -> "NetworkVolume":
+        """Create a new network volume."""
+        log.debug(f"Creating new network volume: {self.name or 'unnamed'}")
+        payload = self.model_dump(exclude_none=True)
+        result = await client.create_network_volume(payload)
+
+        if volume := self.__class__(**result):
+            return volume
+
+        raise ValueError("Deployment failed, no volume was created.")
+
+    async def deploy(self) -> "DeployableResource":
+        """
+        Deploys the network volume resource using the provided configuration.
+        Returns a DeployableResource object.
+        """
+        try:
+            # If the resource is already deployed, return it
+            if self.is_deployed():
+                log.debug(f"{self} exists")
+                return self
+
+            async with RunpodRestClient() as client:
+                # Check for existing volume first
+                if existing_volume := await self._find_existing_volume(client):
+                    return existing_volume
+
+                # No existing volume found, create a new one
+                return await self._create_new_volume(client)
+
+        except Exception as e:
+            log.error(f"{self} failed to deploy: {e}")
+            raise

tetra_rp/core/resources/resource_manager.py

@@ -0,0 +1,145 @@
+import asyncio
+import cloudpickle
+import logging
+from typing import Dict, Optional
+from pathlib import Path
+
+from ..exceptions import RunpodAPIKeyError
+from ..utils.singleton import SingletonMixin
+from ..utils.file_lock import file_lock, FileLockError
+
+from .base import DeployableResource
+
+
+log = logging.getLogger(__name__)
+
+# File to persist state of resources
+RESOURCE_STATE_FILE = Path(".tetra_resources.pkl")
+
+
+class ResourceManager(SingletonMixin):
+    """Manages dynamic provisioning and tracking of remote resources."""
+
+    # Class variables shared across all instances (singleton)
+    _resources: Dict[str, DeployableResource] = {}
+    _deployment_locks: Dict[str, asyncio.Lock] = {}
+    _global_lock: Optional[asyncio.Lock] = None
+    _lock_initialized = False
+    _resources_initialized = False
+
+    def __init__(self):
+        # Ensure async locks are initialized properly for the singleton instance
+        if not ResourceManager._lock_initialized:
+            ResourceManager._global_lock = asyncio.Lock()
+            ResourceManager._lock_initialized = True
+
+        # Load resources immediately on initialization (only once)
+        if not ResourceManager._resources_initialized:
+            self._load_resources()
+            ResourceManager._resources_initialized = True
+
+    def _load_resources(self) -> Dict[str, DeployableResource]:
+        """Load persisted resource information using cross-platform file locking."""
+        if RESOURCE_STATE_FILE.exists():
+            try:
+                with open(RESOURCE_STATE_FILE, "rb") as f:
+                    # Acquire shared lock for reading (cross-platform)
+                    with file_lock(f, exclusive=False):
+                        self._resources = cloudpickle.load(f)
+                        log.debug(f"Loaded saved resources from {RESOURCE_STATE_FILE}")
+            except (FileLockError, Exception) as e:
+                log.error(f"Failed to load resources from {RESOURCE_STATE_FILE}: {e}")
+        return self._resources
+
+    def _save_resources(self) -> None:
+        """Persist state of resources to disk using cross-platform file locking."""
+        try:
+            with open(RESOURCE_STATE_FILE, "wb") as f:
+                # Acquire exclusive lock for writing (cross-platform)
+                with file_lock(f, exclusive=True):
+                    cloudpickle.dump(self._resources, f)
+                    f.flush()  # Ensure data is written to disk
+                    log.debug(f"Saved resources in {RESOURCE_STATE_FILE}")
+        except (FileLockError, Exception) as e:
+            log.error(f"Failed to save resources to {RESOURCE_STATE_FILE}: {e}")
+            raise
+
+    def add_resource(self, uid: str, resource: DeployableResource):
+        """Add a resource to the manager."""
+        self._resources[uid] = resource
+        self._save_resources()
+
+    # function to check if resource still exists remotely, else remove it
+    def remove_resource(self, uid: str):
+        """Remove a resource from the manager."""
+        if uid not in self._resources:
+            log.warning(f"Resource {uid} not found for removal")
+            return
+
+        del self._resources[uid]
+        log.debug(f"Removed resource {uid}")
+
+        self._save_resources()
+
+    async def _deploy_with_error_context(
+        self, config: DeployableResource
+    ) -> DeployableResource:
+        """Deploy resource with enhanced error context for RunpodAPIKeyError.
+
+        Args:
+            config: Resource configuration to deploy.
+
+        Returns:
+            Deployed resource instance.
+
+        Raises:
+            RunpodAPIKeyError: If deployment fails due to missing API key, with resource context.
+        """
+        try:
+            return await config.deploy()
+        except RunpodAPIKeyError as e:
+            error_msg = f"Cannot deploy resource '{config.name}': {str(e)}"
+            raise RunpodAPIKeyError(error_msg) from e
+
+    async def get_or_deploy_resource(
+        self, config: DeployableResource
+    ) -> DeployableResource:
+        """Get existing or create new resource based on config.
+
+        Thread-safe implementation that prevents concurrent deployments
+        of the same resource configuration.
+        """
+        uid = config.resource_id
+
+        # Ensure global lock is initialized (should be done in __init__)
+        assert ResourceManager._global_lock is not None, "Global lock not initialized"
+
+        # Get or create a per-resource lock
+        async with ResourceManager._global_lock:
+            if uid not in ResourceManager._deployment_locks:
+                ResourceManager._deployment_locks[uid] = asyncio.Lock()
+            resource_lock = ResourceManager._deployment_locks[uid]
+
+        # Acquire per-resource lock for this specific configuration
+        async with resource_lock:
+            # Double-check pattern: check again inside the lock
+            if existing := self._resources.get(uid):
+                if not existing.is_deployed():
+                    log.warning(f"{existing} is no longer valid, redeploying.")
+                    self.remove_resource(uid)
+                    # Don't recursive call - deploy directly within the lock
+                    deployed_resource = await self._deploy_with_error_context(config)
+                    log.info(f"URL: {deployed_resource.url}")
+                    self.add_resource(uid, deployed_resource)
+                    return deployed_resource
+
+                log.debug(f"{existing} exists, reusing.")
+                log.info(f"URL: {existing.url}")
+                return existing
+
+            # No existing resource, deploy new one
+            log.debug(f"Deploying new resource: {uid}")
+            deployed_resource = await self._deploy_with_error_context(config)
+            log.info(f"URL: {deployed_resource.url}")
+            self.add_resource(uid, deployed_resource)
+            return deployed_resource