feat: AE-892 - Lock the network volume region

src/tetra_rp/client.py

@@ -1,7 +1,7 @@
 import logging
 from functools import wraps
-from typing import List, Optional
-from .core.resources import ServerlessResource, ResourceManager, NetworkVolume
+from typing import List
+from .core.resources import ServerlessResource, ResourceManager
 from .stubs import stub_resource
 
 
@@ -12,7 +12,6 @@ def remote(
     resource_config: ServerlessResource,
     dependencies: List[str] = None,
     system_dependencies: List[str] = None,
-    mount_volume: Optional[NetworkVolume] = None,
     **extra,
 ):
     """
@@ -49,15 +48,6 @@ 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/resources/network_volume.py

@@ -20,8 +20,6 @@ class DataCenter(str, Enum):
     """
 
     EU_RO_1 = "EU-RO-1"
-    US_WA_1 = "US-WA-1"
-    US_CA_1 = "US-CA-1"
 
 
 class NetworkVolume(DeployableResource):
@@ -33,10 +31,12 @@ class NetworkVolume(DeployableResource):
 
     """
 
-    dataCenterId: Optional[DataCenter] = None
+    # 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] = None  # Size in GB
+    size: Optional[int] = Field(default=10, gt=0)  # Size in GB
 
     @property
     def is_created(self) -> bool:

src/tetra_rp/core/resources/serverless.py

@@ -1,6 +1,6 @@
 import asyncio
 import logging
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 from enum import Enum
 from pydantic import (
     field_serializer,
@@ -22,6 +22,7 @@
 from .cpu import CpuInstanceType
 from .environment import EnvironmentVars
 from .constants import CONSOLE_URL
+from .network_volume import NetworkVolume
 
 
 # Environment variables are loaded from the .env file
@@ -62,7 +63,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
@@ -71,14 +80,19 @@ class ServerlessResource(DeployableResource):
     gpus: Optional[List[GpuGroup]] = [GpuGroup.ANY]  # for gpuIds
     imageName: Optional[str] = ""  # for template.imageName
 
+    #  Input-only field that accepts NetworkVolume object or string ID
+    networkVolume: Optional[Union[NetworkVolume, str]] = Field(
+        default=None, exclude=True
+    )
+
     # === Input Fields ===
     executionTimeoutMs: Optional[int] = None
     gpuCount: Optional[int] = 1
     idleTimeout: Optional[int] = 5
     instanceIds: Optional[List[CpuInstanceType]] = None
     locations: Optional[str] = None
     name: str
-    networkVolumeId: Optional[str] = None
+    networkVolumeId: Optional[str] = None  # This gets set from networkVolume
     scalerType: Optional[ServerlessScalerType] = ServerlessScalerType.QUEUE_DELAY
     scalerValue: Optional[int] = 4
     templateId: Optional[str] = None
@@ -116,6 +130,26 @@ def endpoint(self) -> runpod.Endpoint:
             raise ValueError("Missing self.id")
         return runpod.Endpoint(self.id)
 
+    @field_validator("networkVolume")
+    @classmethod
+    def validate_network_volume(
+        cls, value: Optional[Union[NetworkVolume, str]]
+    ) -> Optional[Union[NetworkVolume, str]]:
+        """Validate networkVolume input"""
+        if value is None:
+            return None
+
+        if isinstance(value, str):
+            # If it's a string, assume it's a volume ID
+            return value
+        elif isinstance(value, NetworkVolume):
+            # If it's a NetworkVolume object, validate it
+            return value
+        else:
+            raise ValueError(
+                "networkVolume must be either a NetworkVolume object or a string ID"
+            )
+
     @field_serializer("scalerType")
     def serialize_scaler_type(
         self, value: Optional[ServerlessScalerType]
@@ -142,6 +176,16 @@ def sync_input_fields(self):
         if self.flashboot:
             self.name += "-fb"
 
+        if self.networkVolume:
+            if isinstance(self.networkVolume, str):
+                # It's already an ID
+                self.networkVolumeId = self.networkVolume
+            elif isinstance(self.networkVolume, NetworkVolume):
+                # It's a NetworkVolume object
+                if 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:
@@ -177,6 +221,37 @@ def _sync_input_fields_cpu(self):
 
         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 not self.networkVolume:
+            log.info(
+                f"No network volume provided for {self.name}, creating default network volume"
+            )
+            default_volume = NetworkVolume(
+                name=f"{self.name}-volume",
+            )
+            self.networkVolume = default_volume
+
+        if isinstance(self.networkVolume, str):
+            # It's already an ID, set it
+            self.networkVolumeId = self.networkVolume
+            return
+
+        if isinstance(self.networkVolume, NetworkVolume):
+            if not self.networkVolume.is_created:
+                # Deploy the network volume
+                log.info(f"Deploying network volume for {self.name}")
+                deployed_volume = await self.networkVolume.deploy()
+                self.networkVolume = deployed_volume
+                self.networkVolumeId = deployed_volume.id
+                log.info(f"Network volume deployed with ID: {deployed_volume.id}")
+            else:
+                # Already deployed, just set the ID
+                self.networkVolumeId = self.networkVolume.id
+
     def is_deployed(self) -> bool:
         """
         Checks if the serverless resource is deployed and available.
@@ -202,6 +277,9 @@ async def deploy(self) -> "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)