diff --git a/src/tetra_rp/__init__.py b/src/tetra_rp/__init__.py index 7f81dcb4..a8bca517 100644 --- a/src/tetra_rp/__init__.py +++ b/src/tetra_rp/__init__.py @@ -20,6 +20,7 @@ ResourceManager, ServerlessEndpoint, runpod, + NetworkVolume, ) @@ -34,4 +35,5 @@ "ResourceManager", "ServerlessEndpoint", "runpod", + "NetworkVolume", ] diff --git a/src/tetra_rp/client.py b/src/tetra_rp/client.py index 9f68874e..12c5d14c 100644 --- a/src/tetra_rp/client.py +++ b/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 diff --git a/src/tetra_rp/core/api/__init__.py b/src/tetra_rp/core/api/__init__.py index 5a268e76..478b2dbe 100644 --- a/src/tetra_rp/core/api/__init__.py +++ b/src/tetra_rp/core/api/__init__.py @@ -1,5 +1,6 @@ -from .runpod import RunpodGraphQLClient +from .runpod import RunpodGraphQLClient, RunpodRestClient __all__ = [ "RunpodGraphQLClient", + "RunpodRestClient", ] diff --git a/src/tetra_rp/core/api/runpod.py b/src/tetra_rp/core/api/runpod.py index 91c69c59..c5d94347 100644 --- a/src/tetra_rp/core/api/runpod.py +++ b/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() diff --git a/src/tetra_rp/core/resources/__init__.py b/src/tetra_rp/core/resources/__init__.py index 6fc8ef31..4f149747 100644 --- a/src/tetra_rp/core/resources/__init__.py +++ b/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", ] diff --git a/src/tetra_rp/core/resources/constants.py b/src/tetra_rp/core/resources/constants.py new file mode 100644 index 00000000..28a0b2ae --- /dev/null +++ b/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" diff --git a/src/tetra_rp/core/resources/network_volume.py b/src/tetra_rp/core/resources/network_volume.py new file mode 100644 index 00000000..c4f6e095 --- /dev/null +++ b/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 diff --git a/src/tetra_rp/core/resources/serverless.py b/src/tetra_rp/core/resources/serverless.py index c7773308..51267fab 100644 --- a/src/tetra_rp/core/resources/serverless.py +++ b/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"