tetra-rp 0.17.1__py3-none-any.whl

tetra_rp/core/resources/serverless.py

@@ -0,0 +1,463 @@
+import asyncio
+import logging
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+from pydantic import (
+    BaseModel,
+    Field,
+    field_serializer,
+    field_validator,
+    model_validator,
+)
+from runpod.endpoint.runner import Job
+
+from ..api.runpod import RunpodGraphQLClient
+from ..utils.backoff import get_backoff_delay
+from .base import DeployableResource
+from .cloud import runpod
+from .constants import CONSOLE_URL
+from .environment import EnvironmentVars
+from .gpu import GpuGroup
+from .network_volume import NetworkVolume, DataCenter
+from .template import KeyValuePair, PodTemplate
+
+
+# Environment variables are loaded from the .env file
+def get_env_vars() -> Dict[str, str]:
+    """
+    Returns the environment variables from the .env file.
+    {
+        "KEY": "VALUE",
+    }
+    """
+    env_vars = EnvironmentVars()
+    return env_vars.get_env()
+
+
+log = logging.getLogger(__name__)
+
+
+class ServerlessScalerType(Enum):
+    QUEUE_DELAY = "QUEUE_DELAY"
+    REQUEST_COUNT = "REQUEST_COUNT"
+
+
+class ServerlessType(Enum):
+    """
+    Serverless endpoint execution model.
+
+    QB (Queue-based): Traditional queue processing with automatic retries.
+                      Requests are placed in queue and processed sequentially.
+                      JSON input/output only. Higher latency but built-in error recovery.
+
+    LB (Load-balancer): Direct HTTP routing to healthy workers.
+                        Supports custom HTTP endpoints and any data format.
+                        Lower latency but no automatic retries.
+    """
+
+    QB = "QB"
+    LB = "LB"
+
+
+class CudaVersion(Enum):
+    V11_8 = "11.8"
+    V12_0 = "12.0"
+    V12_1 = "12.1"
+    V12_2 = "12.2"
+    V12_3 = "12.3"
+    V12_4 = "12.4"
+    V12_5 = "12.5"
+    V12_6 = "12.6"
+    V12_7 = "12.7"
+    V12_8 = "12.8"
+
+
+class ServerlessResource(DeployableResource):
+    """
+    Base class for GPU serverless resource
+    """
+
+    _input_only = {
+        "id",
+        "cudaVersions",
+        "datacenter",
+        "env",
+        "gpus",
+        "flashboot",
+        "imageName",
+        "networkVolume",
+    }
+
+    # === Input-only Fields ===
+    cudaVersions: Optional[List[CudaVersion]] = []  # for allowedCudaVersions
+    env: Optional[Dict[str, str]] = Field(default_factory=get_env_vars)
+    flashboot: Optional[bool] = True
+    gpus: Optional[List[GpuGroup]] = [GpuGroup.ANY]  # for gpuIds
+    imageName: Optional[str] = ""  # for template.imageName
+    networkVolume: Optional[NetworkVolume] = None
+    datacenter: DataCenter = Field(default=DataCenter.EU_RO_1)
+
+    # === Input Fields ===
+    executionTimeoutMs: Optional[int] = None
+    gpuCount: Optional[int] = 1
+    idleTimeout: Optional[int] = 5
+    locations: Optional[str] = None
+    name: str
+    networkVolumeId: Optional[str] = None
+    scalerType: Optional[ServerlessScalerType] = ServerlessScalerType.QUEUE_DELAY
+    scalerValue: Optional[int] = 4
+    templateId: Optional[str] = None
+    type: Optional[ServerlessType] = None
+    workersMax: Optional[int] = 3
+    workersMin: Optional[int] = 0
+    workersPFBTarget: Optional[int] = None
+
+    # === Runtime Fields ===
+    activeBuildid: Optional[str] = None
+    aiKey: Optional[str] = None
+    allowedCudaVersions: Optional[str] = None
+    computeType: Optional[str] = None
+    createdAt: Optional[str] = None  # TODO: use datetime
+    gpuIds: Optional[str] = ""
+    hubRelease: Optional[str] = None
+    repo: Optional[str] = None
+    template: Optional[PodTemplate] = None
+    userId: Optional[str] = None
+
+    def __str__(self) -> str:
+        return f"{self.__class__.__name__}:{self.id}"
+
+    @property
+    def url(self) -> str:
+        if not self.id:
+            raise ValueError("Missing self.id")
+        return CONSOLE_URL % self.id
+
+    @property
+    def endpoint(self) -> runpod.Endpoint:
+        """
+        Returns the Runpod endpoint object for this serverless resource.
+        """
+        if not self.id:
+            raise ValueError("Missing self.id")
+        return runpod.Endpoint(self.id)
+
+    @property
+    def endpoint_url(self) -> str:
+        base_url = self.endpoint.rp_client.endpoint_url_base
+        return f"{base_url}/{self.id}"
+
+    @field_serializer("scalerType")
+    def serialize_scaler_type(
+        self, value: Optional[ServerlessScalerType]
+    ) -> Optional[str]:
+        """Convert ServerlessScalerType enum to string."""
+        return value.value if value is not None else None
+
+    @field_serializer("type")
+    def serialize_type(self, value: Optional[ServerlessType]) -> Optional[str]:
+        """Convert ServerlessType enum to string."""
+        return value.value if value is not None else None
+
+    @field_validator("gpus")
+    @classmethod
+    def validate_gpus(cls, value: List[GpuGroup]) -> List[GpuGroup]:
+        """Expand ANY to all GPU groups"""
+        if value == [GpuGroup.ANY]:
+            return GpuGroup.all()
+        return value
+
+    @model_validator(mode="after")
+    def sync_input_fields(self):
+        """Sync between temporary inputs and exported fields"""
+        if self.flashboot:
+            self.name += "-fb"
+
+        # Sync datacenter to locations field for API
+        if not self.locations:
+            self.locations = self.datacenter.value
+
+        # Validate datacenter consistency between endpoint and network volume
+        if self.networkVolume and self.networkVolume.dataCenterId != self.datacenter:
+            raise ValueError(
+                f"Network volume datacenter ({self.networkVolume.dataCenterId.value}) "
+                f"must match endpoint datacenter ({self.datacenter.value})"
+            )
+
+        if self.networkVolume and self.networkVolume.is_created:
+            # Volume already exists, use its ID
+            self.networkVolumeId = self.networkVolume.id
+
+        self._sync_input_fields_gpu()
+
+        return self
+
+    def _sync_input_fields_gpu(self):
+        # GPU-specific fields
+        if self.gpus:
+            # Convert gpus list to gpuIds string
+            self.gpuIds = ",".join(gpu.value for gpu in self.gpus)
+        elif self.gpuIds:
+            # Convert gpuIds string to gpus list (from backend responses)
+            gpu_values = [v.strip() for v in self.gpuIds.split(",") if v.strip()]
+            self.gpus = [GpuGroup(value) for value in gpu_values]
+
+        if self.cudaVersions:
+            # Convert cudaVersions list to allowedCudaVersions string
+            self.allowedCudaVersions = ",".join(v.value for v in self.cudaVersions)
+        elif self.allowedCudaVersions:
+            # Convert allowedCudaVersions string to cudaVersions list (from backend responses)
+            version_values = [
+                v.strip() for v in self.allowedCudaVersions.split(",") if v.strip()
+            ]
+            self.cudaVersions = [CudaVersion(value) for value in version_values]
+
+        return self
+
+    async def _ensure_network_volume_deployed(self) -> None:
+        """
+        Ensures network volume is deployed and ready if one is specified.
+        Updates networkVolumeId with the deployed volume ID.
+        """
+        if self.networkVolumeId:
+            return
+
+        if self.networkVolume:
+            deployedNetworkVolume = await self.networkVolume.deploy()
+            self.networkVolumeId = deployedNetworkVolume.id
+
+    def is_deployed(self) -> bool:
+        """
+        Checks if the serverless resource is deployed and available.
+        """
+        try:
+            if not self.id:
+                return False
+
+            response = self.endpoint.health()
+            return response is not None
+        except Exception as e:
+            log.error(f"Error checking {self}: {e}")
+            return False
+
+    async def deploy(self) -> "DeployableResource":
+        """
+        Deploys the serverless 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
+
+            # NEW: Ensure network volume is deployed first
+            await self._ensure_network_volume_deployed()
+
+            async with RunpodGraphQLClient() as client:
+                payload = self.model_dump(exclude=self._input_only, exclude_none=True)
+                result = await client.create_endpoint(payload)
+
+            if endpoint := self.__class__(**result):
+                return endpoint
+
+            raise ValueError("Deployment failed, no endpoint was returned.")
+
+        except Exception as e:
+            log.error(f"{self} failed to deploy: {e}")
+            raise
+
+    async def run_sync(self, payload: Dict[str, Any]) -> "JobOutput":
+        """
+        Executes a serverless endpoint request with the payload.
+        Returns a JobOutput object.
+        """
+        if not self.id:
+            raise ValueError("Serverless is not deployed")
+
+        def _fetch_job():
+            return self.endpoint.rp_client.post(
+                f"{self.id}/runsync", payload, timeout=60
+            )
+
+        try:
+            # log.debug(f"[{self}] Payload: {payload}")
+
+            log.info(f"{self} | API /run_sync")
+            response = await asyncio.to_thread(_fetch_job)
+            return JobOutput(**response)
+
+        except Exception as e:
+            health = await asyncio.to_thread(self.endpoint.health)
+            health = ServerlessHealth(**health)
+            log.info(f"{self} | Health {health.workers.status}")
+            log.error(f"{self} | Exception: {e}")
+            raise
+
+    async def run(self, payload: Dict[str, Any]) -> "JobOutput":
+        """
+        Executes a serverless endpoint async request with the payload.
+        Returns a JobOutput object.
+        """
+        if not self.id:
+            raise ValueError("Serverless is not deployed")
+
+        job: Optional[Job] = None
+
+        try:
+            # log.debug(f"[{self}] Payload: {payload}")
+
+            # Create a job using the endpoint
+            log.info(f"{self} | API /run")
+            job = await asyncio.to_thread(self.endpoint.run, request_input=payload)
+
+            log_subgroup = f"Job:{job.job_id}"
+
+            log.info(f"{self} | Started {log_subgroup}")
+
+            current_pace = 0
+            attempt = 0
+            job_status = Status.UNKNOWN
+            last_status = job_status
+
+            # Poll for job status
+            while True:
+                await asyncio.sleep(current_pace)
+
+                # Check job status
+                job_status = await asyncio.to_thread(job.status)
+
+                if last_status == job_status:
+                    # nothing changed, increase the gap
+                    attempt += 1
+                    indicator = "." * (attempt // 2) if attempt % 2 == 0 else ""
+                    if indicator:
+                        log.info(f"{log_subgroup} | {indicator}")
+                else:
+                    # status changed, reset the gap
+                    log.info(f"{log_subgroup} | Status: {job_status}")
+                    attempt = 0
+
+                last_status = job_status
+
+                # Adjust polling pace appropriately
+                current_pace = get_backoff_delay(attempt)
+
+                if job_status in ("COMPLETED", "FAILED", "CANCELLED"):
+                    response = await asyncio.to_thread(job._fetch_job)
+                    return JobOutput(**response)
+
+        except Exception as e:
+            if job and job.job_id:
+                log.info(f"{self} | Cancelling job {job.job_id}")
+                await asyncio.to_thread(job.cancel)
+
+            log.error(f"{self} | Exception: {e}")
+            raise
+
+
+class ServerlessEndpoint(ServerlessResource):
+    """
+    Represents a serverless endpoint distinct from a live serverless.
+    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]):
+            raise ValueError(
+                "Either imageName, template, or templateId must be provided"
+            )
+
+        if not self.templateId and not self.template:
+            self.template = self._create_new_template()
+        elif self.template:
+            self._configure_existing_template()
+
+        return self
+
+
+class JobOutput(BaseModel):
+    id: str
+    workerId: str
+    status: str
+    delayTime: int
+    executionTime: int
+    output: Optional[Any] = None
+    error: Optional[str] = ""
+
+    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")
+
+
+class Status(str, Enum):
+    READY = "READY"
+    INITIALIZING = "INITIALIZING"
+    THROTTLED = "THROTTLED"
+    UNHEALTHY = "UNHEALTHY"
+    UNKNOWN = "UNKNOWN"
+
+
+class WorkersHealth(BaseModel):
+    idle: int
+    initializing: int
+    ready: int
+    running: int
+    throttled: int
+    unhealthy: int
+
+    @property
+    def status(self) -> Status:
+        if self.ready or self.idle or self.running:
+            return Status.READY
+
+        if self.initializing:
+            return Status.INITIALIZING
+
+        if self.throttled:
+            return Status.THROTTLED
+
+        if self.unhealthy:
+            return Status.UNHEALTHY
+
+        return Status.UNKNOWN
+
+
+class JobsHealth(BaseModel):
+    completed: int
+    failed: int
+    inProgress: int
+    inQueue: int
+    retried: int
+
+
+class ServerlessHealth(BaseModel):
+    workers: WorkersHealth
+    jobs: JobsHealth
+
+    @property
+    def is_ready(self) -> bool:
+        return self.workers.status == Status.READY

tetra_rp/core/resources/serverless_cpu.py

@@ -0,0 +1,162 @@
+"""
+CPU-specific serverless endpoint classes.
+
+This module contains all CPU-related serverless functionality, separate from GPU serverless.
+"""
+
+from typing import List, Optional
+
+from pydantic import field_serializer, model_validator, field_validator
+
+from .cpu import (
+    CpuInstanceType,
+    CPU_INSTANCE_DISK_LIMITS,
+    get_max_disk_size_for_instances,
+)
+from .serverless import ServerlessEndpoint, get_env_vars
+from .template import KeyValuePair, PodTemplate
+
+
+class CpuEndpointMixin:
+    """Mixin class that provides CPU-specific functionality for serverless endpoints."""
+
+    instanceIds: Optional[List[CpuInstanceType]]
+
+    def _is_cpu_endpoint(self) -> bool:
+        """Check if this is a CPU endpoint (has instanceIds)."""
+        return (
+            hasattr(self, "instanceIds")
+            and self.instanceIds is not None
+            and len(self.instanceIds) > 0
+        )
+
+    def _get_cpu_container_disk_size(self) -> Optional[int]:
+        """Get the appropriate container disk size for CPU instances."""
+        if not self._is_cpu_endpoint():
+            return None
+        return get_max_disk_size_for_instances(self.instanceIds)
+
+    def _apply_cpu_disk_sizing(self, template: PodTemplate) -> None:
+        """Apply CPU disk sizing to a template if it's using the default size."""
+        if not self._is_cpu_endpoint():
+            return
+
+        # Only auto-size if template is using the default value
+        default_disk_size = PodTemplate.model_fields["containerDiskInGb"].default
+        if template.containerDiskInGb == default_disk_size:
+            cpu_disk_size = self._get_cpu_container_disk_size()
+            if cpu_disk_size is not None:
+                template.containerDiskInGb = cpu_disk_size
+
+    def validate_cpu_container_disk_size(self) -> None:
+        """
+        Validate that container disk size doesn't exceed limits for CPU instances.
+
+        Raises:
+            ValueError: If container disk size exceeds the limit for any CPU instance
+        """
+        if (
+            not self._is_cpu_endpoint()
+            or not hasattr(self, "template")
+            or not self.template
+            or not self.template.containerDiskInGb
+        ):
+            return
+
+        max_allowed_disk_size = self._get_cpu_container_disk_size()
+        if max_allowed_disk_size is None:
+            return
+
+        if self.template.containerDiskInGb > max_allowed_disk_size:
+            instance_limits = []
+            for instance_type in self.instanceIds:
+                limit = CPU_INSTANCE_DISK_LIMITS[instance_type]
+                instance_limits.append(f"{instance_type.value}: max {limit}GB")
+
+            raise ValueError(
+                f"Container disk size {self.template.containerDiskInGb}GB exceeds the maximum "
+                f"allowed for CPU instances. Instance limits: {', '.join(instance_limits)}. "
+                f"Maximum allowed: {max_allowed_disk_size}GB"
+            )
+
+    def _sync_cpu_fields(self):
+        """Sync CPU-specific fields, overriding GPU defaults."""
+        # Override GPU-specific fields for CPU
+        if hasattr(self, "gpuCount"):
+            self.gpuCount = 0
+        if hasattr(self, "allowedCudaVersions"):
+            self.allowedCudaVersions = ""
+        if hasattr(self, "gpuIds"):
+            self.gpuIds = ""
+
+    @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]
+
+
+class CpuServerlessEndpoint(CpuEndpointMixin, ServerlessEndpoint):
+    """
+    CPU-only serverless endpoint with automatic disk sizing and validation.
+    Represents a CPU-only serverless endpoint distinct from a live serverless.
+    """
+
+    instanceIds: Optional[List[CpuInstanceType]] = [CpuInstanceType.ANY]
+
+    def _create_new_template(self) -> PodTemplate:
+        """Create a new PodTemplate with CPU-appropriate disk sizing."""
+        template = PodTemplate(
+            name=self.resource_id,
+            imageName=self.imageName,
+            env=KeyValuePair.from_dict(self.env or get_env_vars()),
+        )
+        # Apply CPU-specific disk sizing
+        self._apply_cpu_disk_sizing(template)
+        return template
+
+    def _configure_existing_template(self) -> None:
+        """Configure an existing template with necessary overrides and CPU sizing."""
+        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)
+
+        # Apply CPU-specific disk sizing
+        self._apply_cpu_disk_sizing(self.template)
+
+    @field_validator("instanceIds")
+    @classmethod
+    def validate_cpus(cls, value: List[CpuInstanceType]) -> List[CpuInstanceType]:
+        """Expand ANY to all GPU groups"""
+        if value == [CpuInstanceType.ANY]:
+            return CpuInstanceType.all()
+        return value
+
+    @model_validator(mode="after")
+    def set_serverless_template(self):
+        # Sync CPU-specific fields first
+        self._sync_cpu_fields()
+
+        if not any([self.imageName, self.template, self.templateId]):
+            raise ValueError(
+                "Either imageName, template, or templateId must be provided"
+            )
+
+        if not self.templateId and not self.template:
+            self.template = self._create_new_template()
+        elif self.template:
+            self._configure_existing_template()
+
+        # Validate container disk size for CPU instances
+        self.validate_cpu_container_disk_size()
+
+        return self