tetra-rp 0.5.5__tar.gz → 0.7.0__tar.gz

{tetra_rp-0.5.5 → tetra_rp-0.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tetra_rp
-Version: 0.5.5
+Version: 0.7.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

{tetra_rp-0.5.5 → tetra_rp-0.7.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tetra_rp"
-version = "0.5.5"
+version = "0.7.0"
 description = "A Python library for distributed inference and serving of machine learning models"
 authors = [
     { name = "Marut Pandya", email = "pandyamarut@gmail.com" },

{tetra_rp-0.5.5 → tetra_rp-0.7.0}/src/tetra_rp/__init__.py

@@ -20,6 +20,7 @@ from .core.resources import (  # noqa: E402
     ResourceManager,
     ServerlessEndpoint,
     runpod,
+    NetworkVolume,
 )
 
 
@@ -34,4 +35,5 @@ __all__ = [
     "ResourceManager",
     "ServerlessEndpoint",
     "runpod",
+    "NetworkVolume",
 ]

{tetra_rp-0.5.5 → tetra_rp-0.7.0}/src/tetra_rp/client.py

@@ -24,6 +24,8 @@ def remote(
             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.
+        mount_volume (NetworkVolume, optional): Configuration for creating and mounting a network volume.
+            Should contain 'size', 'datacenter_id', and 'name' keys. Defaults to None.
         extra (dict, optional): Additional parameters for the execution of the resource. Defaults to an empty dict.
 
     Returns:

tetra_rp-0.7.0/src/tetra_rp/core/api/__init__.py

@@ -0,0 +1,6 @@
+from .runpod import RunpodGraphQLClient, RunpodRestClient
+
+__all__ = [
+    "RunpodGraphQLClient",
+    "RunpodRestClient",
+]

{tetra_rp-0.5.5 → tetra_rp-0.7.0}/src/tetra_rp/core/api/runpod.py

@@ -3,15 +3,17 @@ Direct GraphQL communication with Runpod API.
 Bypasses the outdated runpod-python SDK limitations.
 """
 
-import os
 import json
-import aiohttp
-from typing import Dict, Any, Optional
 import logging
+import os
+from typing import Any, Dict, Optional
+
+import aiohttp
 
 log = logging.getLogger(__name__)
 
 RUNPOD_API_BASE_URL = os.environ.get("RUNPOD_API_BASE_URL", "https://api.runpod.io")
+RUNPOD_REST_API_URL = os.environ.get("RUNPOD_REST_API_URL", "https://rest.runpod.io/v1")
 
 
 class RunpodGraphQLClient:
@@ -210,3 +212,82 @@ class RunpodGraphQLClient:
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         await self.close()
+
+
+class RunpodRestClient:
+    """
+    Runpod REST client for Runpod API.
+    Provides methods to interact with Runpod's REST endpoints.
+    """
+
+    def __init__(self, api_key: Optional[str] = None):
+        self.api_key = api_key or os.getenv("RUNPOD_API_KEY")
+        if not self.api_key:
+            raise ValueError("Runpod API key is required")
+
+        self.session: Optional[aiohttp.ClientSession] = None
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        """Get or create an aiohttp session."""
+        if self.session is None or self.session.closed:
+            timeout = aiohttp.ClientTimeout(total=300)  # 5 minute timeout
+            self.session = aiohttp.ClientSession(
+                timeout=timeout,
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                },
+            )
+        return self.session
+
+    async def _execute_rest(
+        self, method: str, url: str, data: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """Execute a REST API request."""
+        session = await self._get_session()
+
+        log.debug(f"REST Request: {method} {url}")
+        log.debug(f"REST Data: {json.dumps(data, indent=2) if data else 'None'}")
+
+        try:
+            async with session.request(method, url, json=data) as response:
+                response_data = await response.json()
+
+                log.debug(f"REST Response Status: {response.status}")
+                log.debug(f"REST Response: {json.dumps(response_data, indent=2)}")
+
+                if response.status >= 400:
+                    raise Exception(
+                        f"REST request failed: {response.status} - {response_data}"
+                    )
+
+                return response_data
+
+        except aiohttp.ClientError as e:
+            log.error(f"HTTP client error: {e}")
+            raise Exception(f"HTTP request failed: {e}")
+
+    async def create_network_volume(self, payload: Dict[str, Any]) -> Dict[str, Any]:
+        """Create a network volume in Runpod."""
+        log.debug(f"Creating network volume: {payload.get('name', 'unnamed')}")
+
+        result = await self._execute_rest(
+            "POST", f"{RUNPOD_REST_API_URL}/networkvolumes", payload
+        )
+
+        log.info(
+            f"Created network volume: {result.get('id', 'unknown')} - {result.get('name', 'unnamed')}"
+        )
+
+        return result
+
+    async def close(self):
+        """Close the HTTP session."""
+        if self.session and not self.session.closed:
+            await self.session.close()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()

{tetra_rp-0.5.5 → tetra_rp-0.7.0}/src/tetra_rp/core/resources/__init__.py

@@ -12,6 +12,7 @@ from .serverless import (
     CudaVersion,
 )
 from .template import PodTemplate
+from .network_volume import NetworkVolume
 
 
 __all__ = [
@@ -30,4 +31,5 @@ __all__ = [
     "ServerlessResource",
     "ServerlessEndpoint",
     "PodTemplate",
+    "NetworkVolume",
 ]

tetra_rp-0.7.0/src/tetra_rp/core/resources/constants.py

@@ -0,0 +1,4 @@
+import os
+
+CONSOLE_BASE_URL = os.environ.get("CONSOLE_BASE_URL", "https://console.runpod.io")
+CONSOLE_URL = f"{CONSOLE_BASE_URL}/serverless/user/endpoint/%s"

tetra_rp-0.7.0/src/tetra_rp/core/resources/network_volume.py

@@ -0,0 +1,107 @@
+import logging
+from enum import Enum
+from typing import Optional
+
+from pydantic import (
+    Field,
+    field_serializer,
+)
+
+from ..api.runpod import RunpodRestClient
+from .base import DeployableResource
+from .constants import CONSOLE_BASE_URL
+
+log = logging.getLogger(__name__)
+
+
+class DataCenter(str, Enum):
+    """
+    Enum representing available data centers for network volumes.
+    #TODO: Add more data centers as needed. Lock this to the available data center.
+    """
+
+    EU_RO_1 = "EU-RO-1"
+
+
+class NetworkVolume(DeployableResource):
+    """
+    NetworkVolume resource for creating and managing Runpod netowrk volumes.
+
+    This class handles the creation, deployment, and management of network volumes
+    that can be attached to serverless resources.
+
+    """
+
+    # Internal fixed value
+    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=10, gt=0)  # Size in GB
+
+    def __str__(self) -> str:
+        return f"{self.__class__.__name__}:{self.id}"
+
+    @field_serializer("dataCenterId")
+    def serialize_data_center_id(self, value: Optional[DataCenter]) -> Optional[str]:
+        """Convert DataCenter enum to string."""
+        return value.value if value is not None else None
+
+    @property
+    def is_created(self) -> bool:
+        "Returns True if the network volume already exists."
+        return self.id is not None
+
+    @property
+    def url(self) -> str:
+        """
+        Returns the URL for the network volume resource.
+        """
+        if not self.id:
+            raise ValueError("Network volume ID is not set")
+        return f"{CONSOLE_BASE_URL}/user/storage"
+
+    async def create_network_volume(self) -> str:
+        """
+        Creates a network volume using the provided configuration.
+        Returns the volume ID.
+        """
+        async with RunpodRestClient() as client:
+            # Create the network volume
+            payload = self.model_dump(exclude_none=True)
+            result = await client.create_network_volume(payload)
+
+        if volume := self.__class__(**result):
+            return volume
+
+    def is_deployed(self) -> bool:
+        """
+        Checks if the network volume resource is deployed and available.
+        """
+        return self.id is not None
+
+    async def deploy(self) -> "DeployableResource":
+        """
+        Deploys the network volume 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
+
+            # Create the network volume
+            async with RunpodRestClient() as client:
+                # Create the network volume
+                payload = self.model_dump(exclude_none=True)
+                result = await client.create_network_volume(payload)
+
+            if volume := self.__class__(**result):
+                return volume
+
+            raise ValueError("Deployment failed, no volume was created.")
+
+        except Exception as e:
+            log.error(f"{self} failed to deploy: {e}")
+            raise

{tetra_rp-0.5.5 → tetra_rp-0.7.0}/src/tetra_rp/core/resources/serverless.py

@@ -1,27 +1,27 @@
 import asyncio
 import logging
-import os
-from typing import Any, Dict, List, Optional
 from enum import Enum
+from typing import Any, Dict, List, Optional
+
 from pydantic import (
+    BaseModel,
+    Field,
     field_serializer,
     field_validator,
     model_validator,
-    BaseModel,
-    Field,
 )
-
 from runpod.endpoint.runner import Job
 
 from ..api.runpod import RunpodGraphQLClient
 from ..utils.backoff import get_backoff_delay
-
-from .cloud import runpod
 from .base import DeployableResource
-from .template import PodTemplate, KeyValuePair
-from .gpu import GpuGroup
+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 .template import KeyValuePair, PodTemplate
 
 
 # Environment variables are loaded from the .env file
@@ -39,10 +39,6 @@ def get_env_vars() -> Dict[str, str]:
 log = logging.getLogger(__name__)
 
 
-CONSOLE_BASE_URL = os.environ.get("CONSOLE_BASE_URL", "https://console.runpod.io")
-CONSOLE_URL = f"{CONSOLE_BASE_URL}/serverless/user/endpoint/%s"
-
-
 class ServerlessScalerType(Enum):
     QUEUE_DELAY = "QUEUE_DELAY"
     REQUEST_COUNT = "REQUEST_COUNT"
@@ -66,7 +62,15 @@ class ServerlessResource(DeployableResource):
     Base class for GPU serverless resource
     """
 
-    _input_only = {"id", "cudaVersions", "env", "gpus", "flashboot", "imageName"}
+    _input_only = {
+        "id",
+        "cudaVersions",
+        "env",
+        "gpus",
+        "flashboot",
+        "imageName",
+        "networkVolume",
+    }
 
     # === Input-only Fields ===
     cudaVersions: Optional[List[CudaVersion]] = []  # for allowedCudaVersions
@@ -75,6 +79,8 @@ class ServerlessResource(DeployableResource):
     gpus: Optional[List[GpuGroup]] = [GpuGroup.ANY]  # for gpuIds
     imageName: Optional[str] = ""  # for template.imageName
 
+    networkVolume: Optional[NetworkVolume] = None
+
     # === Input Fields ===
     executionTimeoutMs: Optional[int] = None
     gpuCount: Optional[int] = 1
@@ -146,6 +152,10 @@ class ServerlessResource(DeployableResource):
         if self.flashboot:
             self.name += "-fb"
 
+        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:
@@ -181,6 +191,21 @@ class ServerlessResource(DeployableResource):
 
         return self
 
+    async def _ensure_network_volume_deployed(self) -> None:
+        """
+        Ensures network volume is deployed and ready.
+        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():
+            self.networkVolumeId = deployedNetworkVolume.id
+
     def is_deployed(self) -> bool:
         """
         Checks if the serverless resource is deployed and available.
@@ -206,6 +231,9 @@ class ServerlessResource(DeployableResource):
                 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)

{tetra_rp-0.5.5 → tetra_rp-0.7.0}/src/tetra_rp.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tetra_rp
-Version: 0.5.5
+Version: 0.7.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

{tetra_rp-0.5.5 → tetra_rp-0.7.0}/src/tetra_rp.egg-info/SOURCES.txt

@@ -20,10 +20,12 @@ src/tetra_rp/core/pool/worker.py
 src/tetra_rp/core/resources/__init__.py
 src/tetra_rp/core/resources/base.py
 src/tetra_rp/core/resources/cloud.py
+src/tetra_rp/core/resources/constants.py
 src/tetra_rp/core/resources/cpu.py
 src/tetra_rp/core/resources/environment.py
 src/tetra_rp/core/resources/gpu.py
 src/tetra_rp/core/resources/live_serverless.py
+src/tetra_rp/core/resources/network_volume.py
 src/tetra_rp/core/resources/resource_manager.py
 src/tetra_rp/core/resources/serverless.py
 src/tetra_rp/core/resources/template.py

tetra_rp-0.5.5/src/tetra_rp/core/api/__init__.py

@@ -1,5 +0,0 @@
-from .runpod import RunpodGraphQLClient
-
-__all__ = [
-    "RunpodGraphQLClient",
-]