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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tetra_rp
-Version: 0.10.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
@@ -13,6 +13,7 @@ Description-Content-Type: text/markdown
 Requires-Dist: cloudpickle>=3.1.1
 Requires-Dist: runpod
 Requires-Dist: python-dotenv>=1.0.0
+Requires-Dist: pydantic>=2.0.0
 
 # Tetra: Serverless computing for AI workloads
 

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

@@ -1,6 +1,6 @@
 [project]
 name = "tetra_rp"
-version = "0.10.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" },
@@ -21,6 +21,7 @@ dependencies = [
     "cloudpickle>=3.1.1",
     "runpod",
     "python-dotenv>=1.0.0",
+    "pydantic>=2.0.0",
 ]
 
 [dependency-groups]

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

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

{tetra_rp-0.10.0 → tetra_rp-0.12.0}/src/tetra_rp/client.py

@@ -14,6 +14,8 @@ def remote(
     resource_config: ServerlessResource,
     dependencies: Optional[List[str]] = None,
     system_dependencies: Optional[List[str]] = None,
+    accelerate_downloads: bool = True,
+    hf_models_to_cache: Optional[List[str]] = None,
     **extra,
 ):
     """
@@ -22,10 +24,17 @@ def remote(
     This decorator allows a function to be executed in a remote serverless environment, with support for
     dynamic resource provisioning and installation of required dependencies.
 
+    Args:
         resource_config (ServerlessResource): Configuration object specifying the serverless resource
             to be provisioned or used.
         dependencies (List[str], optional): A list of pip package names to be installed in the remote
             environment before executing the function. Defaults to None.
+        system_dependencies (List[str], optional): A list of system packages to be installed in the remote
+            environment before executing the function. Defaults to None.
+        accelerate_downloads (bool, optional): Enable download acceleration for dependencies and models.
+            Defaults to True.
+        hf_models_to_cache (List[str], optional): List of HuggingFace model IDs to pre-cache using
+            download acceleration. Defaults to None.
         extra (dict, optional): Additional parameters for the execution of the resource. Defaults to an empty dict.
 
     Returns:
@@ -37,7 +46,8 @@ def remote(
         @remote(
             resource_config=my_resource_config,
             dependencies=["numpy", "pandas"],
-            sync=True  # Optional, to run synchronously
+            accelerate_downloads=True,
+            hf_models_to_cache=["gpt2", "bert-base-uncased"]
         )
         async def my_function(data):
             # Function logic here
@@ -49,7 +59,13 @@ def remote(
         if inspect.isclass(func_or_class):
             # Handle class decoration
             return create_remote_class(
-                func_or_class, resource_config, dependencies, system_dependencies, extra
+                func_or_class,
+                resource_config,
+                dependencies,
+                system_dependencies,
+                accelerate_downloads,
+                hf_models_to_cache,
+                extra,
             )
         else:
             # Handle function decoration (unchanged)
@@ -62,7 +78,13 @@ def remote(
 
                 stub = stub_resource(remote_resource, **extra)
                 return await stub(
-                    func_or_class, dependencies, system_dependencies, *args, **kwargs
+                    func_or_class,
+                    dependencies,
+                    system_dependencies,
+                    accelerate_downloads,
+                    hf_models_to_cache,
+                    *args,
+                    **kwargs,
                 )
 
             return wrapper

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

@@ -3,24 +3,26 @@ 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
+from .network_volume import NetworkVolume, DataCenter
 
 
 __all__ = [
     "runpod",
     "BaseResource",
     "CpuInstanceType",
+    "CpuLiveServerless",
     "CpuServerlessEndpoint",
     "CudaVersion",
+    "DataCenter",
     "DeployableResource",
     "GpuGroup",
     "GpuType",

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.10.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.10.0 → tetra_rp-0.12.0}/src/tetra_rp/core/resources/network_volume.py

@@ -38,8 +38,8 @@ class NetworkVolume(DeployableResource):
     dataCenterId: DataCenter = Field(default=DataCenter.EU_RO_1, frozen=True)
 
     id: Optional[str] = Field(default=None)
-    name: Optional[str] = None
-    size: Optional[int] = Field(default=50, gt=0)  # Size in GB
+    name: str
+    size: Optional[int] = Field(default=100, gt=0)  # Size in GB
 
     def __str__(self) -> str:
         return f"{self.__class__.__name__}:{self.id}"
@@ -47,15 +47,11 @@ class NetworkVolume(DeployableResource):
     @property
     def resource_id(self) -> str:
         """Unique resource ID based on name and datacenter for idempotent behavior."""
-        if self.name:
-            # Use name + datacenter for volumes with names 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()}"
-        else:
-            # Fall back to default behavior for unnamed volumes
-            return super().resource_id
+        # 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]:

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