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

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/stubs/live_serverless.py

@@ -4,6 +4,7 @@ import inspect
 import textwrap
 import hashlib
 import traceback
+import threading
 import cloudpickle
 import logging
 from ..core.resources import LiveServerless
@@ -16,8 +17,9 @@ from ..protos.remote_execution import (
 log = logging.getLogger(__name__)
 
 
-# global in memory cache, TODO: use a more robust cache in future
+# Global in-memory cache with thread safety
 _SERIALIZED_FUNCTION_CACHE = {}
+_function_cache_lock = threading.RLock()
 
 
 def get_function_source(func):
@@ -60,21 +62,34 @@ 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
-        if src_hash not in _SERIALIZED_FUNCTION_CACHE:
-            # Cache the serialized function
-            _SERIALIZED_FUNCTION_CACHE[src_hash] = source
+        # Thread-safe cache access
+        with _function_cache_lock:
+            # check if the function is already cached
+            if src_hash not in _SERIALIZED_FUNCTION_CACHE:
+                # Cache the serialized function
+                _SERIALIZED_FUNCTION_CACHE[src_hash] = source
 
-        request["function_code"] = _SERIALIZED_FUNCTION_CACHE[src_hash]
+            request["function_code"] = _SERIALIZED_FUNCTION_CACHE[src_hash]
 
         # Serialize arguments using cloudpickle
         if args:

tetra_rp/stubs/registry.py

@@ -2,6 +2,7 @@ import logging
 from functools import singledispatch
 
 from ..core.resources import (
+    CpuLiveServerless,
     CpuServerlessEndpoint,
     LiveServerless,
     ServerlessEndpoint,
@@ -20,19 +21,31 @@ def stub_resource(resource, **extra):
     return fallback
 
 
-@stub_resource.register(LiveServerless)
-def _(resource, **extra):
+def _create_live_serverless_stub(resource, **extra):
+    """Create a live serverless stub for both LiveServerless and CpuLiveServerless."""
     stub = LiveServerlessStub(resource)
 
     # 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)
@@ -48,15 +61,27 @@ def _(resource, **extra):
     return stubbed_resource
 
 
+@stub_resource.register(LiveServerless)
+def _(resource, **extra):
+    return _create_live_serverless_stub(resource, **extra)
+
+
+@stub_resource.register(CpuLiveServerless)
+def _(resource, **extra):
+    return _create_live_serverless_stub(resource, **extra)
+
+
 @stub_resource.register(ServerlessEndpoint)
 def _(resource, **extra):
     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,):
-            # cleanup: when the function is called with no args
-            args = []
-
         if dependencies or system_dependencies:
             log.warning(
                 "Dependencies are not supported for ServerlessEndpoint. "
@@ -74,12 +99,14 @@ def _(resource, **extra):
 @stub_resource.register(CpuServerlessEndpoint)
 def _(resource, **extra):
     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,):
-            # cleanup: when the function is called with no args
-            args = []
-
         if dependencies or system_dependencies:
             log.warning(
                 "Dependencies are not supported for CpuServerlessEndpoint. "

{tetra_rp-0.10.0.dist-info → tetra_rp-0.12.0.dist-info}/METADATA

@@ -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.12.0.dist-info/RECORD

@@ -0,0 +1,38 @@
+tetra_rp/__init__.py,sha256=_D3Wbtv9tBh_WGS4Si5uQbsefL63GqqjoGTN3R8P6fA,769
+tetra_rp/client.py,sha256=urSVh0j9didd9U8lboPv3TtFYURp2XO6ReOICr9Xrls,3414
+tetra_rp/execute_class.py,sha256=jYNFalqqjKvvCz1zzodRvOkrLQd2FYnLYa4EElEYp8w,12243
+tetra_rp/logger.py,sha256=gk5-PWp3k_GQ5DxndsRkBCX0jarp_3lgZ1oiTFuThQg,1125
+tetra_rp/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tetra_rp/core/api/__init__.py,sha256=oldrEKMwxYoBPLvPfVlaFS3wfUtTTxCN6-HzlpTh6vE,124
+tetra_rp/core/api/runpod.py,sha256=3TTx1fkXMLZ2R5JCrQYPEn8dhdUsBt8i5OEwAfaKQ_k,10451
+tetra_rp/core/resources/__init__.py,sha256=ZvSfceYV4S1Xo3YA03B-sR5VRud8BKfx5Pe44xFfruo,911
+tetra_rp/core/resources/base.py,sha256=UJeDiFN45aO1n5SBcxn56ohLhj-AWHoj0KO7mF4yJ_o,1440
+tetra_rp/core/resources/cloud.py,sha256=XJOWPfzYlDVJGHxgffcfpEaOKrWhGdi7AzTlaGuYj0o,70
+tetra_rp/core/resources/constants.py,sha256=F1gPqFaXcCmfrbUSO9PQtUBv984TxFc3pySgVy-kXk8,158
+tetra_rp/core/resources/cpu.py,sha256=szhkjaJ9OnjzLvaD7Yc2hRCR-C9eyB6tswo1Qe6VuZ0,3962
+tetra_rp/core/resources/environment.py,sha256=FC9kJCa8YLSar75AKUKqJYnNLrUdjZj8ZTOrspBrS00,1267
+tetra_rp/core/resources/gpu.py,sha256=mMOPLhBugFBAMAl3ezhjAxKuvYya5_9A_h7kvaCoAfk,1885
+tetra_rp/core/resources/live_serverless.py,sha256=FLmaQdn5UMczEfkP3qykIfRVfZeyYdvyNHX9Nd13_54,1868
+tetra_rp/core/resources/network_volume.py,sha256=h_1xhrbBm9jJWROOGl5qy9u4_kCKSyV4idzt0567-J8,5193
+tetra_rp/core/resources/resource_manager.py,sha256=K-SgCk2BMNEAnkB87YynxUH-suZcdcOPLMonL7EogIw,4988
+tetra_rp/core/resources/serverless.py,sha256=1T21RkMjGnM1I87AsGQ6qazp7A9cE7LwH_c6yJ5shPQ,13427
+tetra_rp/core/resources/serverless_cpu.py,sha256=OiG1C_5_j7pYPHdGp3lQQanIl2ak81u8-jlZo2OXflA,5567
+tetra_rp/core/resources/template.py,sha256=qQ8Wd7Rzr1_YeAbW1V7_k7AVHzgWR_RPjcaRfKsetAk,3141
+tetra_rp/core/resources/utils.py,sha256=mgXfgz_NuHN_IC7TzMNdH9II-LMjxcDCG7syDTcPiGs,1721
+tetra_rp/core/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tetra_rp/core/utils/backoff.py,sha256=1pfa0smFNpib8nztcIgBbtrVvQeECKh-aNOfL2TztgU,1324
+tetra_rp/core/utils/constants.py,sha256=Dm4XiO5zTzfdqOSeYVfAjaf2LyHnIEVmbOi_s_k1J_E,375
+tetra_rp/core/utils/file_lock.py,sha256=bxtAexD2rbqMhdr94VbmKdNp0gfKRgxDXx1n7LX4Eso,8269
+tetra_rp/core/utils/json.py,sha256=q0r7aEdfh8kKVeHGeh9fBDfuhHYNopSreislAMB6HhM,1163
+tetra_rp/core/utils/lru_cache.py,sha256=drwKg-DfLbeBRGTzuxKqNKMQq0EuZV15LMTZIOyZuVk,2618
+tetra_rp/core/utils/singleton.py,sha256=lSXgEQGX9nzhrc05GMpThn9SHKG45iajBbSEtwCcNyI,632
+tetra_rp/protos/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tetra_rp/protos/remote_execution.py,sha256=flKJG0U4ked84cXyF4Gfs_7fBgLsEVOzBv8ZWB9UlP0,5648
+tetra_rp/stubs/__init__.py,sha256=ozKsHs8q0T7o2qhQEquub9hqomh1Htys53mMraaRu2E,72
+tetra_rp/stubs/live_serverless.py,sha256=2bQ47Hq4D8o3YydkxR6QVuLjYrzYS427r3067beYK2A,4550
+tetra_rp/stubs/registry.py,sha256=akJREObyvTAVp6Rvxw8O49W52eJAjByleLU_phMyaDI,3351
+tetra_rp/stubs/serverless.py,sha256=BM_a5Ml5VADBYu2WRNmo9qnicP8NnXDGl5ywifulbD0,947
+tetra_rp-0.12.0.dist-info/METADATA,sha256=u2OHbJV-8KIiOca3SenK2J6I8ms13P8Jc64KGNhFAyE,28077
+tetra_rp-0.12.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tetra_rp-0.12.0.dist-info/top_level.txt,sha256=bBay7JTDwJXsTYvVjrwno9hnF-j0q272lk65f2AcPjU,9
+tetra_rp-0.12.0.dist-info/RECORD,,

tetra_rp-0.10.0.dist-info/RECORD

@@ -1,36 +0,0 @@
-tetra_rp/__init__.py,sha256=-1S5sYIKtnUV8V1HlSIbX1yZwiUrsO8J5b3ZEIR_phU,687
-tetra_rp/client.py,sha256=rAMMmn4ejAayFXJMZzx7dG_8Y65tCEMI6wSSKgur4zQ,2500
-tetra_rp/execute_class.py,sha256=HoH-qWDA7X6yGvQMwmHn5-MKxbLWHEDEHsuat5dzl2U,11912
-tetra_rp/logger.py,sha256=gk5-PWp3k_GQ5DxndsRkBCX0jarp_3lgZ1oiTFuThQg,1125
-tetra_rp/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tetra_rp/core/api/__init__.py,sha256=oldrEKMwxYoBPLvPfVlaFS3wfUtTTxCN6-HzlpTh6vE,124
-tetra_rp/core/api/runpod.py,sha256=3TTx1fkXMLZ2R5JCrQYPEn8dhdUsBt8i5OEwAfaKQ_k,10451
-tetra_rp/core/resources/__init__.py,sha256=UhIwo1Y6-tw5qsULamR296sQiztuz-oWrSTreqfmFSw,814
-tetra_rp/core/resources/base.py,sha256=UJeDiFN45aO1n5SBcxn56ohLhj-AWHoj0KO7mF4yJ_o,1440
-tetra_rp/core/resources/cloud.py,sha256=XJOWPfzYlDVJGHxgffcfpEaOKrWhGdi7AzTlaGuYj0o,70
-tetra_rp/core/resources/constants.py,sha256=F1gPqFaXcCmfrbUSO9PQtUBv984TxFc3pySgVy-kXk8,158
-tetra_rp/core/resources/cpu.py,sha256=YIE-tKolSU3JJzpPB7ey-PbRdqKWsJZ_Ad4h2OYaaiA,1231
-tetra_rp/core/resources/environment.py,sha256=FC9kJCa8YLSar75AKUKqJYnNLrUdjZj8ZTOrspBrS00,1267
-tetra_rp/core/resources/gpu.py,sha256=2jIIMr8PNnlIAP8ZTKO8Imx-rdxXp2rbdSHJeVfjawk,1858
-tetra_rp/core/resources/live_serverless.py,sha256=A3JRdCYwHR2KN_OlmTLcv-m_ObxNhBhc5CnUzXOpOtc,1177
-tetra_rp/core/resources/network_volume.py,sha256=h11dRlAkkxrqyNvUP9Eb8BHAUSFQyRP4lNgBdKChezw,5391
-tetra_rp/core/resources/resource_manager.py,sha256=kUVZDblfUzaG78S8FwOzu4rN6QSegUgQNK3fJ_X7l0w,2834
-tetra_rp/core/resources/serverless.py,sha256=48mENAPQrR8fMjWFpb7mpGFOMqjXZnRWGULGH7NPa5E,13629
-tetra_rp/core/resources/template.py,sha256=UkflJXZFWIbQkLuUt4oRLAjn-yIpw9_mT2X1cAH69CU,3141
-tetra_rp/core/resources/utils.py,sha256=mgXfgz_NuHN_IC7TzMNdH9II-LMjxcDCG7syDTcPiGs,1721
-tetra_rp/core/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tetra_rp/core/utils/backoff.py,sha256=1pfa0smFNpib8nztcIgBbtrVvQeECKh-aNOfL2TztgU,1324
-tetra_rp/core/utils/constants.py,sha256=Dm4XiO5zTzfdqOSeYVfAjaf2LyHnIEVmbOi_s_k1J_E,375
-tetra_rp/core/utils/json.py,sha256=q0r7aEdfh8kKVeHGeh9fBDfuhHYNopSreislAMB6HhM,1163
-tetra_rp/core/utils/lru_cache.py,sha256=drwKg-DfLbeBRGTzuxKqNKMQq0EuZV15LMTZIOyZuVk,2618
-tetra_rp/core/utils/singleton.py,sha256=JRli0HhBfq4P9mBUOg1TZUUwMvIenRqWdymX3qFMm2k,210
-tetra_rp/protos/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tetra_rp/protos/remote_execution.py,sha256=F4uwobnp5q-lX3lR7NCAB23J6OzlzcsB35cezwuoSnI,4638
-tetra_rp/stubs/__init__.py,sha256=ozKsHs8q0T7o2qhQEquub9hqomh1Htys53mMraaRu2E,72
-tetra_rp/stubs/live_serverless.py,sha256=o1NH5XEwUD-27NXJsEGO0IwnuDp8iXwUiw5nZtaZZOI,4199
-tetra_rp/stubs/registry.py,sha256=dmbyC7uBp04_sXsG2wJCloFfFRzYjYQ-naEBKhTRo-U,2839
-tetra_rp/stubs/serverless.py,sha256=BM_a5Ml5VADBYu2WRNmo9qnicP8NnXDGl5ywifulbD0,947
-tetra_rp-0.10.0.dist-info/METADATA,sha256=Ck626kHGCXM6r5CHIm9P7gcg1q3IGWhB7Wiw7x0yIJs,28046
-tetra_rp-0.10.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tetra_rp-0.10.0.dist-info/top_level.txt,sha256=bBay7JTDwJXsTYvVjrwno9hnF-j0q272lk65f2AcPjU,9
-tetra_rp-0.10.0.dist-info/RECORD,,