tetra-rp 0.10.0__py3-none-any.whl → 0.12.0__py3-none-any.whl

tetra_rp/core/resources/serverless.py

@@ -17,10 +17,9 @@ 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
+from .network_volume import NetworkVolume, DataCenter
 from .template import KeyValuePair, PodTemplate
 
 
@@ -65,6 +64,7 @@ class ServerlessResource(DeployableResource):
     _input_only = {
         "id",
         "cudaVersions",
+        "datacenter",
         "env",
         "gpus",
         "flashboot",
@@ -78,14 +78,13 @@ class ServerlessResource(DeployableResource):
     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
-    instanceIds: Optional[List[CpuInstanceType]] = None
     locations: Optional[str] = None
     name: str
     networkVolumeId: Optional[str] = None
@@ -133,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]:
@@ -156,14 +146,24 @@ class ServerlessResource(DeployableResource):
         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
 
-        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
@@ -187,27 +187,16 @@ 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.
+        Ensures network volume is deployed and ready if one is specified.
         Updates networkVolumeId with the deployed volume ID.
         """
         if self.networkVolumeId:
             return
 
-        if not self.networkVolume:
-            log.info(f"{self.name} requires a default network volume")
-            self.networkVolume = NetworkVolume(name=f"{self.name}-volume")
-
-        if deployedNetworkVolume := await self.networkVolume.deploy():
+        if self.networkVolume:
+            deployedNetworkVolume = await self.networkVolume.deploy()
             self.networkVolumeId = deployedNetworkVolume.id
 
     def is_deployed(self) -> bool:
@@ -265,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)
@@ -346,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]):
@@ -354,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
@@ -389,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")

tetra_rp/core/resources/serverless_cpu.py

@@ -0,0 +1,154 @@
+"""
+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
+
+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.CPU3G_2_8]
+
+    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)
+
+    @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

tetra_rp/core/resources/template.py

@@ -22,7 +22,7 @@ class KeyValuePair(BaseModel):
 class PodTemplate(BaseResource):
     advancedStart: Optional[bool] = False
     config: Optional[Dict[str, Any]] = {}
-    containerDiskInGb: Optional[int] = 10
+    containerDiskInGb: Optional[int] = 64
     containerRegistryAuthId: Optional[str] = ""
     dockerArgs: Optional[str] = ""
     env: Optional[List[KeyValuePair]] = []

tetra_rp/core/utils/file_lock.py

@@ -0,0 +1,260 @@
+"""
+Cross-platform file locking utilities.
+
+Provides unified file locking interface that works across Windows, macOS, and Linux.
+Uses platform-appropriate locking mechanisms:
+- Windows: msvcrt.locking()
+- Unix/Linux/macOS: fcntl.flock()
+- Fallback: Basic file existence checking (limited protection)
+"""
+
+import contextlib
+import logging
+import platform
+import time
+from pathlib import Path
+from typing import BinaryIO, Optional
+
+log = logging.getLogger(__name__)
+
+# Platform detection
+_IS_WINDOWS = platform.system() == "Windows"
+_IS_UNIX = platform.system() in ("Linux", "Darwin")
+
+# Initialize availability flags
+_WINDOWS_LOCKING_AVAILABLE = False
+_UNIX_LOCKING_AVAILABLE = False
+
+# Import platform-specific modules
+if _IS_WINDOWS:
+    try:
+        import msvcrt
+
+        _WINDOWS_LOCKING_AVAILABLE = True
+    except ImportError:
+        msvcrt = None
+        log.warning("msvcrt not available on Windows platform")
+
+if _IS_UNIX:
+    try:
+        import fcntl
+
+        _UNIX_LOCKING_AVAILABLE = True
+    except ImportError:
+        fcntl = None
+        log.warning("fcntl not available on Unix platform")
+
+
+class FileLockError(Exception):
+    """Exception raised when file locking operations fail."""
+
+    pass
+
+
+class FileLockTimeout(FileLockError):
+    """Exception raised when file locking times out."""
+
+    pass
+
+
+@contextlib.contextmanager
+def file_lock(
+    file_handle: BinaryIO,
+    exclusive: bool = True,
+    timeout: Optional[float] = 10.0,
+    retry_interval: float = 0.1,
+):
+    """
+    Cross-platform file locking context manager.
+
+    Args:
+        file_handle: Open file handle to lock
+        exclusive: True for exclusive lock, False for shared lock
+        timeout: Maximum seconds to wait for lock (None = no timeout)
+        retry_interval: Seconds to wait between lock attempts
+
+    Raises:
+        FileLockTimeout: If lock cannot be acquired within timeout
+        FileLockError: If locking operation fails
+
+    Usage:
+        with open("file.dat", "rb") as f:
+            with file_lock(f, exclusive=False):  # Shared read lock
+                data = f.read()
+
+        with open("file.dat", "wb") as f:
+            with file_lock(f, exclusive=True):   # Exclusive write lock
+                f.write(data)
+    """
+    lock_acquired = False
+    start_time = time.time()
+
+    try:
+        # Platform-specific locking
+        while not lock_acquired:
+            try:
+                if _IS_WINDOWS and _WINDOWS_LOCKING_AVAILABLE:
+                    _acquire_windows_lock(file_handle, exclusive)
+                elif _IS_UNIX and _UNIX_LOCKING_AVAILABLE:
+                    _acquire_unix_lock(file_handle, exclusive)
+                else:
+                    # Fallback - limited protection via file existence
+                    _acquire_fallback_lock(file_handle, exclusive, timeout)
+
+                lock_acquired = True
+                log.debug(f"File lock acquired (exclusive={exclusive})")
+
+            except (OSError, IOError, FileLockError) as e:
+                # Check timeout
+                if timeout is not None and (time.time() - start_time) >= timeout:
+                    raise FileLockTimeout(
+                        f"Could not acquire file lock within {timeout} seconds: {e}"
+                    ) from e
+
+                # Retry after interval
+                time.sleep(retry_interval)
+
+        # Lock acquired successfully
+        yield
+
+    finally:
+        # Release lock
+        if lock_acquired:
+            try:
+                if _IS_WINDOWS and _WINDOWS_LOCKING_AVAILABLE:
+                    _release_windows_lock(file_handle)
+                elif _IS_UNIX and _UNIX_LOCKING_AVAILABLE:
+                    _release_unix_lock(file_handle)
+                else:
+                    _release_fallback_lock(file_handle)
+
+                log.debug("File lock released")
+
+            except Exception as e:
+                log.error(f"Error releasing file lock: {e}")
+                # Don't raise - we're in cleanup
+
+
+def _acquire_windows_lock(file_handle: BinaryIO, exclusive: bool) -> None:
+    """Acquire Windows file lock using msvcrt.locking()."""
+    if not _WINDOWS_LOCKING_AVAILABLE:
+        raise FileLockError("Windows file locking not available (msvcrt missing)")
+
+    # Windows locking modes
+    if exclusive:
+        lock_mode = msvcrt.LK_NBLCK  # Non-blocking exclusive lock
+    else:
+        # Windows doesn't have shared locks in msvcrt
+        # Fall back to exclusive for compatibility
+        lock_mode = msvcrt.LK_NBLCK
+        log.debug("Windows: Using exclusive lock instead of shared (msvcrt limitation)")
+
+    try:
+        # Lock the entire file (position 0, length 1)
+        file_handle.seek(0)
+        msvcrt.locking(file_handle.fileno(), lock_mode, 1)
+    except OSError as e:
+        raise FileLockError(f"Failed to acquire Windows file lock: {e}") from e
+
+
+def _release_windows_lock(file_handle: BinaryIO) -> None:
+    """Release Windows file lock."""
+    if not _WINDOWS_LOCKING_AVAILABLE:
+        return
+
+    try:
+        file_handle.seek(0)
+        msvcrt.locking(file_handle.fileno(), msvcrt.LK_UNLCK, 1)
+    except OSError as e:
+        raise FileLockError(f"Failed to release Windows file lock: {e}") from e
+
+
+def _acquire_unix_lock(file_handle: BinaryIO, exclusive: bool) -> None:
+    """Acquire Unix file lock using fcntl.flock()."""
+    if not _UNIX_LOCKING_AVAILABLE:
+        raise FileLockError("Unix file locking not available (fcntl missing)")
+
+    # Unix locking modes
+    if exclusive:
+        lock_mode = fcntl.LOCK_EX | fcntl.LOCK_NB  # Non-blocking exclusive
+    else:
+        lock_mode = fcntl.LOCK_SH | fcntl.LOCK_NB  # Non-blocking shared
+
+    try:
+        fcntl.flock(file_handle.fileno(), lock_mode)
+    except (OSError, IOError) as e:
+        raise FileLockError(f"Failed to acquire Unix file lock: {e}") from e
+
+
+def _release_unix_lock(file_handle: BinaryIO) -> None:
+    """Release Unix file lock."""
+    if not _UNIX_LOCKING_AVAILABLE:
+        return
+
+    try:
+        fcntl.flock(file_handle.fileno(), fcntl.LOCK_UN)
+    except (OSError, IOError) as e:
+        raise FileLockError(f"Failed to release Unix file lock: {e}") from e
+
+
+def _acquire_fallback_lock(
+    file_handle: BinaryIO, exclusive: bool, timeout: Optional[float]
+) -> None:
+    """
+    Fallback locking using lock files.
+
+    This provides minimal protection but doesn't prevent all race conditions.
+    It's better than no locking but not as robust as OS-level file locks.
+    """
+    log.warning(
+        "Using fallback file locking - limited protection against race conditions"
+    )
+
+    # Create lock file based on the original file
+    file_path = (
+        Path(file_handle.name) if hasattr(file_handle, "name") else Path("unknown")
+    )
+    lock_file = file_path.with_suffix(file_path.suffix + ".lock")
+
+    start_time = time.time()
+
+    while True:
+        try:
+            # Try to create lock file atomically
+            lock_file.touch(mode=0o600, exist_ok=False)
+            log.debug(f"Fallback lock file created: {lock_file}")
+            return
+
+        except FileExistsError:
+            # Lock file exists, check timeout
+            if timeout is not None and (time.time() - start_time) >= timeout:
+                raise FileLockError(f"Fallback lock timeout: {lock_file} exists")
+
+            # Wait and retry
+            time.sleep(0.1)
+
+
+def _release_fallback_lock(file_handle: BinaryIO) -> None:
+    """Release fallback lock by removing lock file."""
+    try:
+        file_path = (
+            Path(file_handle.name) if hasattr(file_handle, "name") else Path("unknown")
+        )
+        lock_file = file_path.with_suffix(file_path.suffix + ".lock")
+
+        if lock_file.exists():
+            lock_file.unlink()
+            log.debug(f"Fallback lock file removed: {lock_file}")
+
+    except Exception as e:
+        log.error(f"Failed to remove fallback lock file: {e}")
+
+
+def get_platform_info() -> dict:
+    """Get information about current platform and available locking mechanisms."""
+    return {
+        "platform": platform.system(),
+        "windows_locking": _IS_WINDOWS and _WINDOWS_LOCKING_AVAILABLE,
+        "unix_locking": _IS_UNIX and _UNIX_LOCKING_AVAILABLE,
+        "fallback_only": not (_WINDOWS_LOCKING_AVAILABLE or _UNIX_LOCKING_AVAILABLE),
+    }

tetra_rp/core/utils/singleton.py

@@ -1,7 +1,21 @@
+import threading
+
+
 class SingletonMixin:
+    """Thread-safe singleton mixin class.
+
+    Uses threading.Lock to ensure only one instance is created
+    per class, even under concurrent access.
+    """
+
     _instances = {}
+    _lock = threading.Lock()
 
     def __new__(cls, *args, **kwargs):
+        # Use double-checked locking pattern for performance
         if cls not in cls._instances:
-            cls._instances[cls] = super().__new__(cls)
+            with cls._lock:
+                # Check again inside the lock (double-checked locking)
+                if cls not in cls._instances:
+                    cls._instances[cls] = super().__new__(cls)
         return cls._instances[cls]

tetra_rp/execute_class.py

@@ -202,6 +202,8 @@ def create_remote_class(
     resource_config: ServerlessResource,
     dependencies: Optional[List[str]],
     system_dependencies: Optional[List[str]],
+    accelerate_downloads: bool,
+    hf_models_to_cache: Optional[List[str]],
     extra: dict,
 ):
     """
@@ -219,6 +221,8 @@ def create_remote_class(
             self._resource_config = resource_config
             self._dependencies = dependencies or []
             self._system_dependencies = system_dependencies or []
+            self._accelerate_downloads = accelerate_downloads
+            self._hf_models_to_cache = hf_models_to_cache
             self._extra = extra
             self._constructor_args = args
             self._constructor_kwargs = kwargs
@@ -302,6 +306,8 @@ def create_remote_class(
                     constructor_kwargs=constructor_kwargs,
                     dependencies=self._dependencies,
                     system_dependencies=self._system_dependencies,
+                    accelerate_downloads=self._accelerate_downloads,
+                    hf_models_to_cache=self._hf_models_to_cache,
                     instance_id=self._instance_id,
                     create_new_instance=not hasattr(
                         self, "_stub"