tetra-rp 0.6.0__py3-none-any.whl → 0.24.0__py3-none-any.whl

tetra_rp/core/resources/base.py

@@ -1,6 +1,6 @@
 import hashlib
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Optional, Dict, Any
 from pydantic import BaseModel, ConfigDict
 
 
@@ -17,11 +17,137 @@ class BaseResource(BaseModel):
 
     @property
     def resource_id(self) -> str:
-        """Unique resource ID based on configuration."""
+        """Unique resource ID based on configuration.
+
+        Computed once and cached to ensure stability across the object's lifetime.
+        This prevents hash changes if validators mutate the object after first access.
+
+        The hash excludes the 'id' field since it's assigned by the provider after
+        deployment and should not affect resource identity.
+
+        If the resource defines _hashed_fields as a class variable, only those fields
+        are included in the hash.
+        """
+        # Use a private attribute in __dict__ for caching (handles pickle correctly)
+        cache_key = "_cached_resource_id"
+        if cache_key not in self.__dict__:
+            resource_type = self.__class__.__name__
+            # Check if resource defines _hashed_fields as a set/frozenset or ModelPrivateAttr
+            hashed_fields_attr = getattr(self.__class__, "_hashed_fields", None)
+            include_fields = None
+
+            if isinstance(hashed_fields_attr, (set, frozenset)):
+                # Direct set/frozenset
+                include_fields = hashed_fields_attr - {"id"}
+            elif hasattr(hashed_fields_attr, "default") and isinstance(
+                hashed_fields_attr.default, (set, frozenset)
+            ):
+                # Pydantic ModelPrivateAttr with set/frozenset default
+                include_fields = hashed_fields_attr.default - {"id"}
+
+            if include_fields:
+                config_str = self.model_dump_json(
+                    exclude_none=True, include=include_fields
+                )
+            else:
+                # Fallback: Exclude only id field
+                config_str = self.model_dump_json(exclude_none=True, exclude={"id"})
+
+            hash_obj = hashlib.md5(f"{resource_type}:{config_str}".encode())
+            self.__dict__[cache_key] = f"{resource_type}_{hash_obj.hexdigest()}"
+        return self.__dict__[cache_key]
+
+    @property
+    def config_hash(self) -> str:
+        """Get hash of current configuration (excluding id and server-assigned fields).
+
+        Unlike resource_id which is cached, this always computes fresh hash.
+        Used for drift detection.
+
+        For resources with _input_only set, only those fields are included in the hash
+        to avoid drift from server-assigned fields.
+        """
+        import json
+        import logging
+
         resource_type = self.__class__.__name__
-        config_str = self.model_dump_json(exclude_none=True)
+
+        # If resource defines input_only fields, use only those for hash
+        if hasattr(self, "_input_only"):
+            # Include only user-provided input fields, not server-assigned ones
+            include_fields = self._input_only - {"id"}  # Exclude id from input fields
+            config_dict = self.model_dump(
+                exclude_none=True, include=include_fields, mode="json"
+            )
+        else:
+            # Fallback: exclude only id field
+            config_dict = self.model_dump(
+                exclude_none=True, exclude={"id"}, mode="json"
+            )
+
+        # Convert to JSON string for hashing
+        config_str = json.dumps(config_dict, sort_keys=True)
         hash_obj = hashlib.md5(f"{resource_type}:{config_str}".encode())
-        return f"{resource_type}_{hash_obj.hexdigest()}"
+        hash_value = hash_obj.hexdigest()
+
+        # Debug logging to see what's being hashed
+        log = logging.getLogger(__name__)
+        if hasattr(self, "name"):
+            log.debug(
+                f"CONFIG HASH for {self.name} ({resource_type}):\n"
+                f"  Fields included: {sorted(config_dict.keys())}\n"
+                f"  Config dict: {config_str}\n"
+                f"  Hash: {hash_value}"
+            )
+
+        return hash_value
+
+    def get_resource_key(self) -> str:
+        """Get stable resource key for tracking.
+
+        Format: {ResourceType}:{name}
+        This provides stable identity even when config changes.
+        """
+        resource_type = self.__class__.__name__
+        name = getattr(self, "name", None)
+        if name:
+            return f"{resource_type}:{name}"
+        # Fallback to resource_id for resources without names
+        return self.resource_id
+
+    def __getstate__(self) -> Dict[str, Any]:
+        """Get state for pickling, excluding non-pickleable items."""
+        import weakref as weakref_module
+
+        state = self.__dict__.copy()
+
+        # Remove any weakrefs from the state dict
+        # This handles cases where threading.Lock or similar objects leak weakrefs
+        keys_to_remove = []
+        for key, value in state.items():
+            # Direct weakref
+            if isinstance(value, weakref_module.ref):
+                keys_to_remove.append(key)
+                continue
+
+            # Check if value holds weakrefs in its __dict__
+            if hasattr(value, "__dict__"):
+                try:
+                    for sub_value in value.__dict__.values():
+                        if isinstance(sub_value, weakref_module.ref):
+                            keys_to_remove.append(key)
+                            break
+                except Exception:
+                    pass
+
+        for key in keys_to_remove:
+            del state[key]
+
+        return state
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        """Restore state from pickling."""
+        self.__dict__.update(state)
 
 
 class DeployableResource(BaseResource, ABC):
@@ -45,3 +171,12 @@ class DeployableResource(BaseResource, ABC):
     async def deploy(self) -> "DeployableResource":
         """Deploy the resource."""
         raise NotImplementedError("Subclasses should implement this method.")
+
+    @abstractmethod
+    async def undeploy(self) -> bool:
+        """Undeploy/delete the resource.
+
+        Returns:
+            True if successfully undeployed, False otherwise
+        """
+        raise NotImplementedError("Subclasses should implement this method.")

tetra_rp/core/resources/constants.py

@@ -1,4 +1,25 @@
 import os
+from urllib.parse import urlparse
+
+import runpod
 
 CONSOLE_BASE_URL = os.environ.get("CONSOLE_BASE_URL", "https://console.runpod.io")
 CONSOLE_URL = f"{CONSOLE_BASE_URL}/serverless/user/endpoint/%s"
+
+
+def _endpoint_domain_from_base_url(base_url: str) -> str:
+    if not base_url:
+        return "api.runpod.ai"
+    if "://" not in base_url:
+        base_url = f"https://{base_url}"
+    parsed = urlparse(base_url)
+    return parsed.netloc or "api.runpod.ai"
+
+
+ENDPOINT_DOMAIN = _endpoint_domain_from_base_url(runpod.endpoint_url_base)
+
+# Flash app artifact upload constants
+TARBALL_CONTENT_TYPE = "application/gzip"
+MAX_TARBALL_SIZE_MB = 500  # Maximum tarball size in megabytes
+VALID_TARBALL_EXTENSIONS = (".tar.gz", ".tgz")  # Valid tarball file extensions
+GZIP_MAGIC_BYTES = (0x1F, 0x8B)  # Magic bytes for gzip files

tetra_rp/core/resources/cpu.py

@@ -1,4 +1,5 @@
 from enum import Enum
+from typing import List, Optional
 
 
 class CpuInstanceType(str, Enum):
@@ -12,23 +13,124 @@ class CpuInstanceType(str, Enum):
     - 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
-    CPU3G_2_8 = "cpu3g-2-8"  # 2 vCPU, 8GB RAM
-    CPU3G_4_16 = "cpu3g-4-16"  # 4 vCPU, 16GB RAM
-    CPU3G_8_32 = "cpu3g-8-32"  # 8 vCPU, 32GB RAM
+
+    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
-    CPU3C_2_4 = "cpu3c-2-4"  # 2 vCPU, 4GB RAM
-    CPU3C_4_8 = "cpu3c-4-8"  # 4 vCPU, 8GB RAM
-    CPU3C_8_16 = "cpu3c-8-16"  # 8 vCPU, 16GB RAM
+
+    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
-    CPU5C_2_4 = "cpu5c-2-4"  # 2 vCPU, 4GB RAM
-    CPU5C_4_8 = "cpu5c-4-8"  # 4 vCPU, 8GB RAM
-    CPU5C_8_16 = "cpu5c-8-16"  # 8 vCPU, 16GB RAM
+
+    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/core/resources/gpu.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Optional, List
 from pydantic import BaseModel
 from enum import Enum
@@ -8,13 +10,13 @@ class GpuLowestPrice(BaseModel):
     uninterruptablePrice: Optional[float] = None
 
 
-class GpuType(BaseModel):
+class GpuTypeModel(BaseModel):
     id: str
     displayName: str
     memoryInGb: int
 
 
-class GpuTypeDetail(GpuType):
+class GpuTypeDetail(GpuTypeModel):
     communityCloud: Optional[bool] = None
     communityPrice: Optional[float] = None
     communitySpotPrice: Optional[float] = None
@@ -32,22 +34,186 @@ 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"]:
         """Returns all GPU groups."""
         return [cls.AMPERE_48] + [g for g in cls if g != cls.ANY]
+
+    @classmethod
+    def to_gpu_ids_str(cls, gpu_types: List[GpuType | GpuGroup]) -> str:
+        """
+        The API expects a comma-separated list of pool IDs, with GPU ID negations (-4090, -3090, etc.). So to convert a list of GPU types to a string of pool IDs, we need to:
+        1. Convert the GPU types to pool IDs
+        2. Add a negation for each GPU type that is not in the list
+        3. Join the pool IDs with a comma
+        """
+        pool_ids = set()
+        pool_ids_from_groups = set()
+        explicit_gpu_types = set()
+
+        for gpu_type in gpu_types:
+            if isinstance(gpu_type, GpuGroup):
+                pool_id = gpu_type
+                pool_ids_from_groups.add(pool_id)
+            else:
+                pool_id = _pool_from_gpu_type(gpu_type)
+                explicit_gpu_types.add(gpu_type)
+
+            if pool_id:
+                pool_ids.add(pool_id)
+
+        # only add negations for pools selected via explicit gpu types
+        if explicit_gpu_types:
+            # iterate over a snapshot because we add negations into the same set
+            for pool_id in list(pool_ids):
+                if pool_id in pool_ids_from_groups:
+                    continue
+                for gpu_type in POOLS_TO_TYPES.get(pool_id, []):
+                    if gpu_type not in explicit_gpu_types:
+                        pool_ids.add(f"-{gpu_type.value}")
+
+        # normalize to strings for the api
+        out = []
+        for pool_id in pool_ids:
+            if isinstance(pool_id, GpuGroup):
+                out.append(pool_id.value)
+            else:
+                out.append(str(pool_id))
+        return ",".join(out)
+
+    @classmethod
+    def from_gpu_ids_str(cls, gpu_ids_str: str) -> List[GpuGroup | GpuType]:
+        """
+        Convert a comma-separated list of pool IDs to a list of GPU types.
+        """
+        ids = gpu_ids_str.split(",")
+        pool_ids = []
+        gpu_types = []
+        negated_gpu_types = []
+        for id in ids:
+            if id.startswith("-") and GpuType.is_gpu_type(id[1:]):
+                negated_gpu_types.append(GpuType(id[1:]))
+            elif GpuType.is_gpu_type(id):
+                gpu_types.append(GpuType(id))
+            else:
+                pool_ids.append(id)
+
+        ids = []
+
+        for pool_id in pool_ids:
+            try:
+                pool = GpuGroup(pool_id)
+            except ValueError:
+                # ignore unknown pool ids from backend
+                continue
+
+            pool_gpus = POOLS_TO_TYPES.get(pool, [])
+            # check if there are any negated gpu types in the pool
+            if any(gpu_type in negated_gpu_types for gpu_type in pool_gpus):
+                # add the gpu types that are not in the negated gpu types
+                ids.extend(
+                    [
+                        gpu_type
+                        for gpu_type in pool_gpus
+                        if gpu_type not in negated_gpu_types
+                    ]
+                )
+            else:
+                ids.append(pool)
+
+        ids.extend(gpu_types)
+        return ids
+
+
+# TODO: fetch from central registry at some point
+class GpuType(Enum):
+    ANY = "any"
+    """Any GPU"""
+
+    NVIDIA_GEFORCE_RTX_4090 = "NVIDIA GeForce RTX 4090"
+    NVIDIA_GEFORCE_RTX_5090 = "NVIDIA GeForce RTX 5090"
+    NVIDIA_RTX_6000_ADA_GENERATION = "NVIDIA RTX 6000 Ada Generation"
+    NVIDIA_H100_80GB_HBM3 = "NVIDIA H100 80GB HBM3"
+    NVIDIA_RTX_A4000 = "NVIDIA RTX A4000"
+    NVIDIA_RTX_A4500 = "NVIDIA RTX A4500"
+    NVIDIA_RTX_4000_ADA_GENERATION = "NVIDIA RTX 4000 Ada Generation"
+    NVIDIA_RTX_2000_ADA_GENERATION = "NVIDIA RTX 2000 Ada Generation"
+    NVIDIA_RTX_A5000 = "NVIDIA RTX A5000"
+    NVIDIA_L4 = "NVIDIA L4"
+    NVIDIA_GEFORCE_RTX_3090 = "NVIDIA GeForce RTX 3090"
+    NVIDIA_A40 = "NVIDIA A40"
+    NVIDIA_RTX_A6000 = "NVIDIA RTX A6000"
+    NVIDIA_A100_80GB_PCIe = "NVIDIA A100 80GB PCIe"
+    NVIDIA_A100_SXM4_80GB = "NVIDIA A100-SXM4-80GB"
+    NVIDIA_H200 = "NVIDIA H200"
+
+    @classmethod
+    def all(cls) -> List["GpuType"]:
+        """Returns all GPU types."""
+        return [g for g in cls if g != cls.ANY]
+
+    @classmethod
+    def is_gpu_type(cls, gpu_type: str) -> bool:
+        """
+        Check if a string is a valid GPU type.
+        """
+        return gpu_type in {m.value for m in cls}
+
+
+POOLS_TO_TYPES = {
+    GpuGroup.ADA_24: [GpuType.NVIDIA_GEFORCE_RTX_4090],
+    GpuGroup.ADA_32_PRO: [GpuType.NVIDIA_GEFORCE_RTX_5090],
+    GpuGroup.ADA_48_PRO: [GpuType.NVIDIA_RTX_6000_ADA_GENERATION],
+    GpuGroup.ADA_80_PRO: [GpuType.NVIDIA_H100_80GB_HBM3],
+    GpuGroup.AMPERE_16: [
+        GpuType.NVIDIA_RTX_A4000,
+        GpuType.NVIDIA_RTX_A4500,
+        GpuType.NVIDIA_RTX_4000_ADA_GENERATION,
+        GpuType.NVIDIA_RTX_2000_ADA_GENERATION,
+    ],
+    GpuGroup.AMPERE_24: [
+        GpuType.NVIDIA_RTX_A5000,
+        GpuType.NVIDIA_L4,
+        GpuType.NVIDIA_GEFORCE_RTX_3090,
+    ],
+    GpuGroup.AMPERE_48: [GpuType.NVIDIA_A40, GpuType.NVIDIA_RTX_A6000],
+    GpuGroup.AMPERE_80: [GpuType.NVIDIA_A100_80GB_PCIe, GpuType.NVIDIA_A100_SXM4_80GB],
+    GpuGroup.HOPPER_141: [GpuType.NVIDIA_H200],
+}
+
+
+def _pool_from_gpu_type(gpu_type: GpuType) -> str:
+    for group, types in POOLS_TO_TYPES.items():
+        if gpu_type in types:
+            return group
+    return None