feat: unify the NetworkVol classes

Signed-off-by: pandyamarut <pandyamarut@gmail.com>

Author: pandyamarut

src/tetra_rp/__init__.py

@@ -19,8 +19,7 @@
     ResourceManager,
     ServerlessEndpoint,
     runpod,
-    NetworkVolumeConfig,
-    NetworkVolumeResource,
+    NetworkVolume,
 )
 
 
@@ -35,6 +34,5 @@
     "ResourceManager",
     "ServerlessEndpoint",
     "runpod",
-    "NetworkVolumeConfig",
-    "NetworkVolumeResource",
+    "NetworkVolume",
 ]

src/tetra_rp/client.py

@@ -1,7 +1,7 @@
 import logging
 from functools import wraps
 from typing import List, Dict, Optional
-from .core.resources import ServerlessResource, ResourceManager, NetworkVolumeConfig, NetworkVolumeResource
+from .core.resources import ServerlessResource, ResourceManager, NetworkVolume
 from .stubs import stub_resource
 
 
@@ -12,7 +12,7 @@ def remote(
     resource_config: ServerlessResource,
     dependencies: List[str] = None,
     system_dependencies: List[str] = None,
-    mount_volume: Optional[NetworkVolumeConfig] = None,
+    mount_volume: Optional[NetworkVolume] = None,
     **extra,
 ):
     """
@@ -52,11 +52,10 @@ 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
+                    network_volume = await mount_volume.deploy()
+                    resource_config.networkVolumeId = network_volume
                     log.info(
-                        f"Updated resource config with network volume: {network_volume.volume_id}"
+                        f"Updated resource config with network volume: {network_volume}"
                     )
                 except Exception as e:
                     log.error(f"Failed to create or mount network volume: {e}")

src/tetra_rp/core/resources/__init__.py

@@ -10,8 +10,7 @@
     ServerlessEndpoint,
     JobOutput,
     CudaVersion,
-    NetworkVolumeConfig,
-    NetworkVolumeResource,
+    NetworkVolume,
 )
 from .template import PodTemplate
 
@@ -32,6 +31,5 @@
     "ServerlessResource",
     "ServerlessEndpoint",
     "PodTemplate",
-    "NetworkVolumeConfig",
-    "NetworkVolumeResource",
+    "NetworkVolume",
 ]

src/tetra_rp/core/resources/serverless.py

@@ -23,6 +23,7 @@
 from .cpu import CpuInstanceType
 from .environment import EnvironmentVars
 
+
 # Environment variables are loaded from the .env file
 def get_env_vars() -> Dict[str, str]:
     """
@@ -232,7 +233,7 @@ async def is_ready_for_requests(self, give_up_threshold=10) -> bool:
 
         Raises:
             ValueError: If the serverless resource is not deployed.
-            RuntimeError: If the health status is THROTTLED, UNHEALTHY, or UNKNOWN 
+            RuntimeError: If the health status is THROTTLED, UNHEALTHY, or UNKNOWN
             after exceeding the give_up_threshold.
         """
         if not self.is_deployed():
@@ -379,7 +380,9 @@ class ServerlessEndpoint(ServerlessResource):
     @model_validator(mode="after")
     def set_serverless_template(self):
         if not any([self.imageName, self.template, self.templateId]):
-            raise ValueError("Either imageName, template, or templateId must be provided")
+            raise ValueError(
+                "Either imageName, template, or templateId must be provided"
+            )
 
         if not self.templateId and not self.template:
             self.template = PodTemplate(
@@ -404,6 +407,7 @@ class CpuServerlessEndpoint(ServerlessEndpoint):
     Represents a CPU-only serverless endpoint distinct from a live serverless.
     Inherits from ServerlessEndpoint.
     """
+
     instanceIds: Optional[List[CpuInstanceType]] = [CpuInstanceType.CPU3G_2_8]
 
 
@@ -471,33 +475,32 @@ class ServerlessHealth(BaseModel):
     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):
+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.
 
     """
-    config: NetworkVolumeConfig
-    volume_id: Optional[str] = Field(default=None)
-    
-        
+
+    dataCenterId: Optional[str] = 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.volume_id:
+        if not self.id:
             raise ValueError("Network volume ID is not set")
         return f"{CONSOLE_BASE_URL}/user/storage"
 
@@ -506,23 +509,21 @@ 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
+            result = await client.create_network_volume(
+                datacenter_id=self.dataCenterId, name=self.name, size=self.size
             )
-            log.info(f"Created network volume: {volume['id']}")
-            return volume["id"]
-    
+        log.info(f"Created network volume: {result['id']}")
+        if volume := self.__class__(**result):
+            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
-    
+        return self.id is not None
+
     async def deploy(self) -> "DeployableResource":
         """
         Deploys the network volume resource using the provided configuration.
@@ -532,17 +533,16 @@ async def deploy(self) -> "DeployableResource":
             # If the resource is already deployed, return it
             if self.is_deployed():
                 log.debug(f"{self} exists")
-                return self
+                return self.id
 
             # Create the network volume
-            self.volume_id = await self.create_network_volume()
+            self.id = await self.create_network_volume()
 
             if self.is_deployed():
-                return self
+                return self.id
 
             raise ValueError("Deployment failed, no volume was created.")
 
         except Exception as e:
             log.error(f"{self} failed to deploy: {e}")
             raise
-