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

src/tetra_rp/__init__.py

@@ -19,6 +19,8 @@
     ResourceManager,
     ServerlessEndpoint,
     runpod,
+    NetworkVolumeConfig,
+    NetworkVolumeResource,
 )
 
 
@@ -33,4 +35,6 @@
     "ResourceManager",
     "ServerlessEndpoint",
     "runpod",
+    "NetworkVolumeConfig",
+    "NetworkVolumeResource",
 ]

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, Dict, Optional
+from .core.resources import ServerlessResource, ResourceManager, NetworkVolumeConfig, NetworkVolumeResource
 from .stubs import stub_resource
 
 
@@ -12,7 +12,8 @@ def remote(
     resource_config: ServerlessResource,
     dependencies: List[str] = None,
     system_dependencies: List[str] = None,
-    **extra
+    mount_volume: Optional[NetworkVolumeConfig] = None,
+    **extra,
 ):
     """
     Decorator to enable dynamic resource provisioning and dependency management for serverless functions.
@@ -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 (NetworkVolumeConfig, 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,19 @@ 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:
+                    nv  = NetworkVolumeResource(config=mount_volume)
+                    network_volume = await nv.deploy()
+                    resource_config.networkVolumeId = network_volume.volume_id
+                    log.info(
+                        f"Updated resource config with network volume: {network_volume.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:
@@ -114,7 +115,9 @@ async def create_endpoint(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
 
         variables = {"input": input_data}
 
-        log.debug(f"Creating endpoint with GraphQL: {input_data.get('name', 'unnamed')}")
+        log.debug(
+            f"Creating endpoint with GraphQL: {input_data.get('name', 'unnamed')}"
+        )
 
         result = await self._execute_graphql(mutation, variables)
 
@@ -142,11 +145,13 @@ async def get_cpu_types(self) -> Dict[str, Any]:
             }
         }
         """
-        
+
         result = await self._execute_graphql(query)
         return result.get("cpuTypes", [])
 
-    async def get_gpu_types(self, gpu_filter: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+    async def get_gpu_types(
+        self, gpu_filter: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
         """Get available GPU types."""
         query = """
         query getGpuTypes($input: GpuTypeFilter) {
@@ -171,7 +176,7 @@ async def get_gpu_types(self, gpu_filter: Optional[Dict[str, Any]] = None) -> Di
             }
         }
         """
-        
+
         variables = {"input": gpu_filter} if gpu_filter else {}
         result = await self._execute_graphql(query, variables)
         return result.get("gpuTypes", [])
@@ -206,3 +211,87 @@ 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, datacenter_id: str, name: str, size: int
+    ) -> 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.
+        """
+        url = f"{RUNPOD_REST_API_URL}/networkvolumes"
+        data = {"dataCenterId": datacenter_id, "name": name, "size": size}
+
+        return await self._execute_rest("POST", url, data)
+
+    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

@@ -10,6 +10,8 @@
     ServerlessEndpoint,
     JobOutput,
     CudaVersion,
+    NetworkVolumeConfig,
+    NetworkVolumeResource,
 )
 from .template import PodTemplate
 
@@ -30,4 +32,6 @@
     "ServerlessResource",
     "ServerlessEndpoint",
     "PodTemplate",
+    "NetworkVolumeConfig",
+    "NetworkVolumeResource",
 ]

src/tetra_rp/core/resources/serverless.py

@@ -13,7 +13,7 @@
 
 from runpod.endpoint.runner import Job
 
-from ..api.runpod import RunpodGraphQLClient
+from ..api.runpod import RunpodGraphQLClient, RunpodRestClient
 from ..utils.backoff import get_backoff_delay
 
 from .cloud import runpod
@@ -470,3 +470,79 @@ class ServerlessHealth(BaseModel):
     @property
     def is_ready(self) -> bool:
         return self.workers.status == Status.READY
+
+class NetworkVolumeConfig(BaseModel):
+    """
+    NetworkvolumeConfig Represents a network volume configuration for serverless resources.
+    """
+    size: int  # Size in GB
+    datacenter_id: str
+    name: str
+
+
+class NetworkVolumeResource(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.
+
+    """
+    config: NetworkVolumeConfig
+    volume_id: Optional[str] = Field(default=None)
+
+
+    @property
+    def url(self) -> str:
+        """
+        Returns the URL for the network volume resource.
+        """
+        if not self.volume_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.
+        """
+        print(self.config)
+        async with RunpodRestClient() as client:
+            # Create the network volume
+            volume = await client.create_network_volume(
+                datacenter_id=self.config.datacenter_id,
+                name=self.config.name,
+                size=self.config.size
+            )
+            log.info(f"Created network volume: {volume['id']}")
+            return volume["id"]
+
+    def is_deployed(self) -> bool:
+        """
+        Checks if the network volume resource is deployed and available.
+        """
+        return self.volume_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
+            self.volume_id = 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
+