tetra-rp 0.11.0__tar.gz → 0.12.0__tar.gz

{tetra_rp-0.11.0 → tetra_rp-0.12.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tetra_rp
-Version: 0.11.0
+Version: 0.12.0
 Summary: A Python library for distributed inference and serving of machine learning models
 Author-email: Marut Pandya <pandyamarut@gmail.com>, Patrick Rachford <prachford@icloud.com>, Dean Quinanola <dean.quinanola@runpod.io>
 License: MIT

{tetra_rp-0.11.0 → tetra_rp-0.12.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tetra_rp"
-version = "0.11.0"
+version = "0.12.0"
 description = "A Python library for distributed inference and serving of machine learning models"
 authors = [
     { name = "Marut Pandya", email = "pandyamarut@gmail.com" },

{tetra_rp-0.11.0 → tetra_rp-0.12.0}/src/tetra_rp/__init__.py

@@ -13,6 +13,7 @@ from .client import remote  # noqa: E402
 from .core.resources import (  # noqa: E402
     CpuServerlessEndpoint,
     CpuInstanceType,
+    CpuLiveServerless,
     CudaVersion,
     DataCenter,
     GpuGroup,
@@ -29,6 +30,7 @@ __all__ = [
     "remote",
     "CpuServerlessEndpoint",
     "CpuInstanceType",
+    "CpuLiveServerless",
     "CudaVersion",
     "DataCenter",
     "GpuGroup",

{tetra_rp-0.11.0 → tetra_rp-0.12.0}/src/tetra_rp/core/resources/__init__.py

@@ -3,14 +3,14 @@ from .cloud import runpod
 from .cpu import CpuInstanceType
 from .gpu import GpuGroup, GpuType, GpuTypeDetail
 from .resource_manager import ResourceManager
-from .live_serverless import LiveServerless
+from .live_serverless import LiveServerless, CpuLiveServerless
 from .serverless import (
-    CpuServerlessEndpoint,
     ServerlessResource,
     ServerlessEndpoint,
     JobOutput,
     CudaVersion,
 )
+from .serverless_cpu import CpuServerlessEndpoint
 from .template import PodTemplate
 from .network_volume import NetworkVolume, DataCenter
 
@@ -19,6 +19,7 @@ __all__ = [
     "runpod",
     "BaseResource",
     "CpuInstanceType",
+    "CpuLiveServerless",
     "CpuServerlessEndpoint",
     "CudaVersion",
     "DataCenter",

tetra_rp-0.12.0/src/tetra_rp/core/resources/cpu.py

@@ -0,0 +1,137 @@
+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
+    """
+
+    # 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"""
+
+
+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
+}
+
+
+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-0.11.0 → tetra_rp-0.12.0}/src/tetra_rp/core/resources/gpu.py

@@ -32,20 +32,35 @@ class GpuTypeDetail(GpuType):
 
 # 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"
+    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"]:

tetra_rp-0.12.0/src/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-0.12.0/src/tetra_rp/core/resources/resource_manager.py

@@ -0,0 +1,121 @@
+import asyncio
+import cloudpickle
+import logging
+from typing import Dict, Optional
+from pathlib import Path
+
+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  # Will be initialized lazily
+    _lock_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
+
+        if not self._resources:
+            self._load_resources()
+
+    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 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 config.deploy()
+                    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 config.deploy()
+            log.info(f"URL: {deployed_resource.url}")
+            self.add_resource(uid, deployed_resource)
+            return deployed_resource

{tetra_rp-0.11.0 → tetra_rp-0.12.0}/src/tetra_rp/core/resources/serverless.py

@@ -17,7 +17,6 @@ from ..utils.backoff import get_backoff_delay
 from .base import DeployableResource
 from .cloud import runpod
 from .constants import CONSOLE_URL
-from .cpu import CpuInstanceType
 from .environment import EnvironmentVars
 from .gpu import GpuGroup
 from .network_volume import NetworkVolume, DataCenter
@@ -86,7 +85,6 @@ class ServerlessResource(DeployableResource):
     executionTimeoutMs: Optional[int] = None
     gpuCount: Optional[int] = 1
     idleTimeout: Optional[int] = 5
-    instanceIds: Optional[List[CpuInstanceType]] = None
     locations: Optional[str] = None
     name: str
     networkVolumeId: Optional[str] = None
@@ -134,15 +132,6 @@ class ServerlessResource(DeployableResource):
         """Convert ServerlessScalerType enum to string."""
         return value.value if value is not None else None
 
-    @field_serializer("instanceIds")
-    def serialize_instance_ids(
-        self, value: Optional[List[CpuInstanceType]]
-    ) -> Optional[List[str]]:
-        """Convert CpuInstanceType enums to strings."""
-        if value is None:
-            return None
-        return [item.value if hasattr(item, "value") else str(item) for item in value]
-
     @field_validator("gpus")
     @classmethod
     def validate_gpus(cls, value: List[GpuGroup]) -> List[GpuGroup]:
@@ -172,10 +161,9 @@ class ServerlessResource(DeployableResource):
             # Volume already exists, use its ID
             self.networkVolumeId = self.networkVolume.id
 
-        if self.instanceIds:
-            return self._sync_input_fields_cpu()
-        else:
-            return self._sync_input_fields_gpu()
+        self._sync_input_fields_gpu()
+
+        return self
 
     def _sync_input_fields_gpu(self):
         # GPU-specific fields
@@ -199,14 +187,6 @@ class ServerlessResource(DeployableResource):
 
         return self
 
-    def _sync_input_fields_cpu(self):
-        # Override GPU-specific fields for CPU
-        self.gpuCount = 0
-        self.allowedCudaVersions = ""
-        self.gpuIds = ""
-
-        return self
-
     async def _ensure_network_volume_deployed(self) -> None:
         """
         Ensures network volume is deployed and ready if one is specified.
@@ -274,7 +254,7 @@ class ServerlessResource(DeployableResource):
             )
 
         try:
-            # log.debug(f"[{log_group}] Payload: {payload}")
+            # log.debug(f"[{self}] Payload: {payload}")
 
             log.info(f"{self} | API /run_sync")
             response = await asyncio.to_thread(_fetch_job)
@@ -355,6 +335,26 @@ class ServerlessEndpoint(ServerlessResource):
     Inherits from ServerlessResource.
     """
 
+    def _create_new_template(self) -> PodTemplate:
+        """Create a new PodTemplate with standard configuration."""
+        return PodTemplate(
+            name=self.resource_id,
+            imageName=self.imageName,
+            env=KeyValuePair.from_dict(self.env or get_env_vars()),
+        )
+
+    def _configure_existing_template(self) -> None:
+        """Configure an existing template with necessary overrides."""
+        if self.template is None:
+            return
+
+        self.template.name = f"{self.resource_id}__{self.template.resource_id}"
+
+        if self.imageName:
+            self.template.imageName = self.imageName
+        if self.env:
+            self.template.env = KeyValuePair.from_dict(self.env)
+
     @model_validator(mode="after")
     def set_serverless_template(self):
         if not any([self.imageName, self.template, self.templateId]):
@@ -363,32 +363,13 @@ class ServerlessEndpoint(ServerlessResource):
             )
 
         if not self.templateId and not self.template:
-            self.template = PodTemplate(
-                name=self.resource_id,
-                imageName=self.imageName,
-                env=KeyValuePair.from_dict(self.env or get_env_vars()),
-            )
-
+            self.template = self._create_new_template()
         elif self.template:
-            self.template.name = f"{self.resource_id}__{self.template.resource_id}"
-            if self.imageName:
-                self.template.imageName = self.imageName
-            if self.env:
-                self.template.env = KeyValuePair.from_dict(self.env)
+            self._configure_existing_template()
 
         return self
 
 
-class CpuServerlessEndpoint(ServerlessEndpoint):
-    """
-    Convenience class for CPU serverless endpoint.
-    Represents a CPU-only serverless endpoint distinct from a live serverless.
-    Inherits from ServerlessEndpoint.
-    """
-
-    instanceIds: Optional[List[CpuInstanceType]] = [CpuInstanceType.CPU3G_2_8]
-
-
 class JobOutput(BaseModel):
     id: str
     workerId: str
@@ -398,7 +379,7 @@ class JobOutput(BaseModel):
     output: Optional[Any] = None
     error: Optional[str] = ""
 
-    def model_post_init(self, __context):
+    def model_post_init(self, _: Any) -> None:
         log_group = f"Worker:{self.workerId}"
         log.info(f"{log_group} | Delay Time: {self.delayTime} ms")
         log.info(f"{log_group} | Execution Time: {self.executionTime} ms")