feat: AE-811 Add a class to create network volume using REST

src/tetra_rp/__init__.py

@@ -20,6 +20,7 @@
     ResourceManager,
     ServerlessEndpoint,
     runpod,
+    NetworkVolume,
 )
 
 
@@ -34,4 +35,5 @@
     "ResourceManager",
     "ServerlessEndpoint",
     "runpod",
+    "NetworkVolume",
 ]

src/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 @@ async def my_function(data):
     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

src/tetra_rp/core/api/__init__.py

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

src/tetra_rp/core/api/runpod.py

@@ -12,6 +12,7 @@
 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 @@ async def __aenter__(self):
 
     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()

src/tetra_rp/core/resources/__init__.py

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

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"

src/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

src/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 .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"