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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tetra_rp
-Version: 0.10.0
+Version: 0.11.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.11.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tetra_rp"
-version = "0.10.0"
+version = "0.11.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.11.0}/src/tetra_rp/__init__.py

@@ -14,6 +14,7 @@ from .core.resources import (  # noqa: E402
     CpuServerlessEndpoint,
     CpuInstanceType,
     CudaVersion,
+    DataCenter,
     GpuGroup,
     LiveServerless,
     PodTemplate,
@@ -29,6 +30,7 @@ __all__ = [
     "CpuServerlessEndpoint",
     "CpuInstanceType",
     "CudaVersion",
+    "DataCenter",
     "GpuGroup",
     "LiveServerless",
     "PodTemplate",

{tetra_rp-0.10.0 → tetra_rp-0.11.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.11.0}/src/tetra_rp/core/resources/__init__.py

@@ -12,7 +12,7 @@ from .serverless import (
     CudaVersion,
 )
 from .template import PodTemplate
-from .network_volume import NetworkVolume
+from .network_volume import NetworkVolume, DataCenter
 
 
 __all__ = [
@@ -21,6 +21,7 @@ __all__ = [
     "CpuInstanceType",
     "CpuServerlessEndpoint",
     "CudaVersion",
+    "DataCenter",
     "DeployableResource",
     "GpuGroup",
     "GpuType",

{tetra_rp-0.10.0 → tetra_rp-0.11.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.10.0 → tetra_rp-0.11.0}/src/tetra_rp/core/resources/serverless.py

@@ -20,7 +20,7 @@ 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 +65,7 @@ class ServerlessResource(DeployableResource):
     _input_only = {
         "id",
         "cudaVersions",
+        "datacenter",
         "env",
         "gpus",
         "flashboot",
@@ -78,8 +79,8 @@ 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
@@ -156,6 +157,17 @@ 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
@@ -197,17 +209,14 @@ class ServerlessResource(DeployableResource):
 
     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:

{tetra_rp-0.10.0 → tetra_rp-0.11.0}/src/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-0.10.0 → tetra_rp-0.11.0}/src/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"

{tetra_rp-0.10.0 → tetra_rp-0.11.0}/src/tetra_rp/protos/remote_execution.py

@@ -1,11 +1,22 @@
-# TODO: generate using betterproto
+"""Remote execution protocol definitions using Pydantic models.
+
+This module defines the request/response protocol for remote function and class execution.
+The models align with the protobuf schema for communication with remote workers.
+"""
+
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 from pydantic import BaseModel, Field, model_validator
 
 
 class FunctionRequest(BaseModel):
+    """Request model for remote function or class execution.
+
+    Supports both function-based execution and class instantiation with method calls.
+    All serialized data (args, kwargs, etc.) are base64-encoded cloudpickle strings.
+    """
+
     # MADE OPTIONAL - can be None for class-only execution
     function_name: Optional[str] = Field(
         default=None,
@@ -15,19 +26,19 @@ class FunctionRequest(BaseModel):
         default=None,
         description="Source code of the function to execute",
     )
-    args: List = Field(
+    args: List[str] = Field(
         default_factory=list,
         description="List of base64-encoded cloudpickle-serialized arguments",
     )
-    kwargs: Dict = Field(
+    kwargs: Dict[str, str] = Field(
         default_factory=dict,
         description="Dictionary of base64-encoded cloudpickle-serialized keyword arguments",
     )
-    dependencies: Optional[List] = Field(
+    dependencies: Optional[List[str]] = Field(
         default=None,
         description="Optional list of pip packages to install before executing the function",
     )
-    system_dependencies: Optional[List] = Field(
+    system_dependencies: Optional[List[str]] = Field(
         default=None,
         description="Optional list of system dependencies to install before executing the function",
     )
@@ -44,11 +55,11 @@ class FunctionRequest(BaseModel):
         default=None,
         description="Source code of the class to instantiate (for class execution)",
     )
-    constructor_args: Optional[List] = Field(
+    constructor_args: List[str] = Field(
         default_factory=list,
         description="List of base64-encoded cloudpickle-serialized constructor arguments",
     )
-    constructor_kwargs: Optional[Dict] = Field(
+    constructor_kwargs: Dict[str, str] = Field(
         default_factory=dict,
         description="Dictionary of base64-encoded cloudpickle-serialized constructor keyword arguments",
     )
@@ -65,6 +76,16 @@ class FunctionRequest(BaseModel):
         description="Whether to create a new instance or reuse existing one",
     )
 
+    # Download acceleration fields
+    accelerate_downloads: bool = Field(
+        default=True,
+        description="Enable download acceleration for dependencies and models",
+    )
+    hf_models_to_cache: Optional[List[str]] = Field(
+        default=None,
+        description="List of HuggingFace model IDs to pre-cache using acceleration",
+    )
+
     @model_validator(mode="after")
     def validate_execution_requirements(self) -> "FunctionRequest":
         """Validate that required fields are provided based on execution_type"""
@@ -92,7 +113,12 @@ class FunctionRequest(BaseModel):
 
 
 class FunctionResponse(BaseModel):
-    # EXISTING FIELDS (unchanged)
+    """Response model for remote function or class execution results.
+
+    Contains execution results, error information, and metadata about class instances
+    when applicable. The result field contains base64-encoded cloudpickle data.
+    """
+
     success: bool = Field(
         description="Indicates if the function execution was successful",
     )
@@ -108,12 +134,10 @@ class FunctionResponse(BaseModel):
         default=None,
         description="Captured standard output from the function execution",
     )
-
-    # NEW FIELDS FOR CLASS SUPPORT
     instance_id: Optional[str] = Field(
         default=None, description="ID of the class instance that was used/created"
     )
-    instance_info: Optional[Dict] = Field(
+    instance_info: Optional[Dict[str, Any]] = Field(
         default=None,
         description="Metadata about the class instance (creation time, call count, etc.)",
     )

{tetra_rp-0.10.0 → tetra_rp-0.11.0}/src/tetra_rp/stubs/live_serverless.py

@@ -60,13 +60,24 @@ class LiveServerlessStub(RemoteExecutorStub):
     def __init__(self, server: LiveServerless):
         self.server = server
 
-    def prepare_request(self, func, dependencies, system_dependencies, *args, **kwargs):
+    def prepare_request(
+        self,
+        func,
+        dependencies,
+        system_dependencies,
+        accelerate_downloads,
+        hf_models_to_cache,
+        *args,
+        **kwargs,
+    ):
         source, src_hash = get_function_source(func)
 
         request = {
             "function_name": func.__name__,
             "dependencies": dependencies,
             "system_dependencies": system_dependencies,
+            "accelerate_downloads": accelerate_downloads,
+            "hf_models_to_cache": hf_models_to_cache,
         }
 
         # check if the function is already cached

{tetra_rp-0.10.0 → tetra_rp-0.11.0}/src/tetra_rp/stubs/registry.py

@@ -26,13 +26,25 @@ def _(resource, **extra):
 
     # Function execution
     async def stubbed_resource(
-        func, dependencies, system_dependencies, *args, **kwargs
+        func,
+        dependencies,
+        system_dependencies,
+        accelerate_downloads,
+        hf_models_to_cache,
+        *args,
+        **kwargs,
     ) -> dict:
         if args == (None,):
             args = []
 
         request = stub.prepare_request(
-            func, dependencies, system_dependencies, *args, **kwargs
+            func,
+            dependencies,
+            system_dependencies,
+            accelerate_downloads,
+            hf_models_to_cache,
+            *args,
+            **kwargs,
         )
         response = await stub.ExecuteFunction(request)
         return stub.handle_response(response)

{tetra_rp-0.10.0 → tetra_rp-0.11.0}/src/tetra_rp.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tetra_rp
-Version: 0.10.0
+Version: 0.11.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.11.0}/src/tetra_rp.egg-info/requires.txt

@@ -1,3 +1,4 @@
 cloudpickle>=3.1.1
 runpod
 python-dotenv>=1.0.0
+pydantic>=2.0.0