tetra-rp 0.5.5__py3-none-any.whl → 0.6.0__py3-none-any.whl

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/client.py

@@ -1,7 +1,7 @@
 import logging
 from functools import wraps
-from typing import List
-from .core.resources import ServerlessResource, ResourceManager
+from typing import List, Optional
+from .core.resources import ServerlessResource, ResourceManager, NetworkVolume
 from .stubs import stub_resource
 
 
@@ -12,6 +12,7 @@ def remote(
     resource_config: ServerlessResource,
     dependencies: List[str] = None,
     system_dependencies: List[str] = None,
+    mount_volume: Optional[NetworkVolume] = None,
     **extra,
 ):
     """
@@ -24,6 +25,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:
@@ -46,6 +49,15 @@ def remote(
     def decorator(func):
         @wraps(func)
         async def wrapper(*args, **kwargs):
+            # Create netowrk volume if mount_volume is provided
+            if mount_volume:
+                try:
+                    network_volume = await mount_volume.deploy()
+                    resource_config.networkVolumeId = network_volume.id
+                except Exception as e:
+                    log.error(f"Failed to create or mount network volume: {e}")
+                    raise
+
             resource_manager = ResourceManager()
             remote_resource = await resource_manager.get_or_deploy_resource(
                 resource_config

tetra_rp/core/api/__init__.py

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

tetra_rp/core/api/runpod.py

@@ -12,6 +12,7 @@ import logging
 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 +211,101 @@ 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.
+
+        Args:
+            datacenter_id (str): The ID of the datacenter where the volume will be created.
+            name (str): The name of the network volume.
+            size_gb (int): The size of the volume in GB.
+
+        Returns:
+            Dict[str, Any]: The created network volume details.
+        """
+        datacenter_id = payload.get("dataCenterId")
+        if hasattr(datacenter_id, "value"):
+            # If datacenter_id is an enum, get its value
+            datacenter_id = datacenter_id.value
+        data = {
+            "dataCenterId": datacenter_id,
+            "name": payload.get("name"),
+            "size": payload.get("size"),
+        }
+        url = f"{RUNPOD_REST_API_URL}/networkvolumes"
+
+        log.debug(f"Creating network volume: {data.get('name', 'unnamed')}")
+
+        result = await self._execute_rest("POST", url, data)
+
+        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/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/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/core/resources/network_volume.py

@@ -0,0 +1,98 @@
+import logging
+from enum import Enum
+from typing import Optional
+
+from pydantic import (
+    Field,
+)
+
+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"
+    US_WA_1 = "US-WA-1"
+    US_CA_1 = "US-CA-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.
+
+    """
+
+    dataCenterId: Optional[DataCenter] = None
+    id: Optional[str] = Field(default=None)
+    name: Optional[str] = None
+    size: Optional[int] = None  # Size in GB
+
+    @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"Network volume {self.id} is already deployed. Mounting existing volume."
+                )
+                log.info(f"Mounted existing network volume: {self.id}")
+                return self
+
+            # Create the network volume
+            self = await self.create_network_volume()
+
+            if self.is_deployed():
+                return self
+
+            raise ValueError("Deployment failed, no volume was created.")
+
+        except Exception as e:
+            log.error(f"{self} failed to deploy: {e}")
+            raise

tetra_rp/core/resources/serverless.py

@@ -1,6 +1,5 @@
 import asyncio
 import logging
-import os
 from typing import Any, Dict, List, Optional
 from enum import Enum
 from pydantic import (
@@ -22,6 +21,7 @@ from .template import PodTemplate, KeyValuePair
 from .gpu import GpuGroup
 from .cpu import CpuInstanceType
 from .environment import EnvironmentVars
+from .constants import CONSOLE_URL
 
 
 # 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"

{tetra_rp-0.5.5.dist-info → tetra_rp-0.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tetra_rp
-Version: 0.5.5
+Version: 0.6.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.dist-info → tetra_rp-0.6.0.dist-info}/RECORD

@@ -1,24 +1,26 @@
-tetra_rp/__init__.py,sha256=_F4CMO3eS_GkPow864tZbBz81zYCfYQ-J3KLikHTli8,647
-tetra_rp/client.py,sha256=sE6GyGai9gpuBBQUzd2Lo7LKUTVbgY606KwlZgyhtGM,2016
+tetra_rp/__init__.py,sha256=-1S5sYIKtnUV8V1HlSIbX1yZwiUrsO8J5b3ZEIR_phU,687
+tetra_rp/client.py,sha256=sJk3kneFJo40OAtbxucbaJ6NxMs-kD2jkpXYwG0vJhM,2681
 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=ErrKSdi9cn9BzT4cJ25X2fbS--GwI3Jwko5NgM-6Otc,82
-tetra_rp/core/api/runpod.py,sha256=v7Jo8vPjH1robl1-25V2PVDefCF9rpCKK758w5oUpG4,6869
+tetra_rp/core/api/__init__.py,sha256=oldrEKMwxYoBPLvPfVlaFS3wfUtTTxCN6-HzlpTh6vE,124
+tetra_rp/core/api/runpod.py,sha256=n1pfvzDceU_yO9iOn4B9Kst6YL2KD0m-migm3bvnfJM,10374
 tetra_rp/core/pool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tetra_rp/core/pool/cluster_manager.py,sha256=KJxEp_044HjnbOhfIdiXZbks_bFDYE1KgKeR5W9VvbY,6007
 tetra_rp/core/pool/dataclass.py,sha256=YngS328_NTewY8Etitj4k7MmdM5GWqqE_OMbytrVNlw,338
 tetra_rp/core/pool/ex.py,sha256=AZOrn9t_X5ycMl-tDg7-jcIURj_9kVmzn9_da8h1TFI,1273
 tetra_rp/core/pool/job.py,sha256=4bisW_ZwiQ2-qD5l0y9SbHcO4EQvSKimmBBU1fpI_YE,567
 tetra_rp/core/pool/worker.py,sha256=N4cOnf8MiDcPFH2XSMmSnnWMACZYUNnKWVhOx2aSxvM,478
-tetra_rp/core/resources/__init__.py,sha256=HKUNv3wktSHGMjNyZxolrYye4uNeWdKf8dF3eQJ51a8,751
+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=6r4I4TEx9AmZ0-OJvE86qrY0S7BEx9t_P2zwHVdtbew,1074
+tetra_rp/core/resources/network_volume.py,sha256=mzYA7Bf4JTFiDyH0sLrUbPjq-26MWgbbRWAMbX2h6iA,2881
 tetra_rp/core/resources/resource_manager.py,sha256=kUVZDblfUzaG78S8FwOzu4rN6QSegUgQNK3fJ_X7l0w,2834
-tetra_rp/core/resources/serverless.py,sha256=HTln6FzjhINYSzEH_iGabdEY9u5BA6o6Q4sUlePe6Es,14850
+tetra_rp/core/resources/serverless.py,sha256=lZMITI1cwOBVNX3dbFyR01AeRZ0eatNC4CRE7vwUrqs,14726
 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
@@ -31,7 +33,7 @@ tetra_rp/stubs/__init__.py,sha256=ozKsHs8q0T7o2qhQEquub9hqomh1Htys53mMraaRu2E,72
 tetra_rp/stubs/live_serverless.py,sha256=o1NH5XEwUD-27NXJsEGO0IwnuDp8iXwUiw5nZtaZZOI,4199
 tetra_rp/stubs/registry.py,sha256=V4m3CeXl8j1pguHuuflxqpWeBgVDQ93YkhxJbElyP7Q,2599
 tetra_rp/stubs/serverless.py,sha256=BM_a5Ml5VADBYu2WRNmo9qnicP8NnXDGl5ywifulbD0,947
-tetra_rp-0.5.5.dist-info/METADATA,sha256=q_m_qvvJylQiHwbX79OyLAc9oVfLLeFqWcEZc5oMrO8,28055
-tetra_rp-0.5.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tetra_rp-0.5.5.dist-info/top_level.txt,sha256=bBay7JTDwJXsTYvVjrwno9hnF-j0q272lk65f2AcPjU,9
-tetra_rp-0.5.5.dist-info/RECORD,,
+tetra_rp-0.6.0.dist-info/METADATA,sha256=CPOdAGUy5vZOSLazBQubYChiQkuwDldW3G0K9KsyBlw,28055
+tetra_rp-0.6.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tetra_rp-0.6.0.dist-info/top_level.txt,sha256=bBay7JTDwJXsTYvVjrwno9hnF-j0q272lk65f2AcPjU,9
+tetra_rp-0.6.0.dist-info/RECORD,,