podstack 1.2.0__py3-none-any.whl

podstack/client.py

@@ -0,0 +1,322 @@
+"""
+Podstack Client
+
+Main client for interacting with the Podstack API.
+"""
+
+import os
+import asyncio
+from typing import Optional, Dict, Any, List
+from urllib.parse import urljoin
+import httpx
+
+from .exceptions import (
+    PodstackError,
+    AuthenticationError,
+    NotFoundError,
+    RateLimitError,
+    ValidationError,
+    ConnectionError
+)
+from .notebook import NotebooksAPI, Notebook
+from .execution import ExecutionsAPI, Execution
+from .models import (
+    GPUInfo,
+    GPUType,
+    Project,
+    WalletBalance,
+    UsageSummary,
+    Webhook
+)
+
+
+class Client:
+    """
+    Podstack API Client.
+
+    Usage:
+        # Async context manager (recommended)
+        async with Client(api_key="your-api-key") as client:
+            notebook = await client.notebooks.create(name="experiment", gpu_type="A100")
+            result = await notebook.execute("print('Hello GPU!')")
+
+        # Manual lifecycle
+        client = Client(api_key="your-api-key")
+        await client.connect()
+        # ... use client ...
+        await client.close()
+
+        # Sync wrapper for simple scripts
+        client = Client(api_key="your-api-key")
+        notebook = client.sync_create_notebook(name="experiment", gpu_type="A100")
+    """
+
+    DEFAULT_BASE_URL = "https://api.podstack.io/v1"
+    DEFAULT_TIMEOUT = 30.0
+
+    def __init__(
+        self,
+        api_key: str = None,
+        base_url: str = None,
+        timeout: float = None,
+        max_retries: int = 3
+    ):
+        """
+        Initialize the Podstack client.
+
+        Args:
+            api_key: API key for authentication. If not provided, uses PODSTACK_API_KEY env var.
+            base_url: Base URL for the API. Defaults to https://api.podstack.io/v1
+            timeout: Request timeout in seconds. Defaults to 30.
+            max_retries: Maximum number of retries for failed requests.
+        """
+        self.api_key = api_key or os.environ.get("PODSTACK_API_KEY")
+        if not self.api_key:
+            raise AuthenticationError("API key is required. Provide api_key or set PODSTACK_API_KEY env var.")
+
+        self.base_url = base_url or os.environ.get("PODSTACK_BASE_URL", self.DEFAULT_BASE_URL)
+        self.timeout = timeout or self.DEFAULT_TIMEOUT
+        self.max_retries = max_retries
+
+        self._http_client: Optional[httpx.AsyncClient] = None
+
+        # API resources
+        self.notebooks = NotebooksAPI(self)
+        self.executions = ExecutionsAPI(self)
+
+    async def connect(self):
+        """Initialize HTTP client"""
+        if self._http_client is None:
+            self._http_client = httpx.AsyncClient(
+                base_url=self.base_url,
+                timeout=self.timeout,
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                    "User-Agent": "podstack-python/1.2.0"
+                }
+            )
+
+    async def close(self):
+        """Close HTTP client"""
+        if self._http_client:
+            await self._http_client.aclose()
+            self._http_client = None
+
+    async def __aenter__(self) -> "Client":
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
+    async def _request(
+        self,
+        method: str,
+        path: str,
+        params: Dict[str, Any] = None,
+        json: Dict[str, Any] = None,
+        retry_count: int = 0
+    ) -> Dict[str, Any]:
+        """
+        Make an API request.
+
+        Args:
+            method: HTTP method
+            path: API path
+            params: Query parameters
+            json: JSON body
+            retry_count: Current retry count
+
+        Returns:
+            Response data as dict
+        """
+        if self._http_client is None:
+            await self.connect()
+
+        try:
+            response = await self._http_client.request(
+                method=method,
+                url=path,
+                params=params,
+                json=json
+            )
+
+            # Handle errors
+            if response.status_code == 401:
+                raise AuthenticationError()
+            elif response.status_code == 404:
+                raise NotFoundError("Resource", path)
+            elif response.status_code == 429:
+                retry_after = response.headers.get("Retry-After")
+                if retry_count < self.max_retries:
+                    await asyncio.sleep(int(retry_after or 1))
+                    return await self._request(method, path, params, json, retry_count + 1)
+                raise RateLimitError(int(retry_after) if retry_after else None)
+            elif response.status_code == 400:
+                error_data = response.json().get("error", {})
+                raise ValidationError(
+                    error_data.get("message", "Invalid request"),
+                    error_data.get("field")
+                )
+            elif response.status_code >= 500:
+                if retry_count < self.max_retries:
+                    await asyncio.sleep(2 ** retry_count)
+                    return await self._request(method, path, params, json, retry_count + 1)
+                raise PodstackError(
+                    f"Server error: {response.status_code}",
+                    code="server_error"
+                )
+            elif response.status_code >= 400:
+                error_data = response.json().get("error", {})
+                raise PodstackError(
+                    error_data.get("message", f"Request failed: {response.status_code}"),
+                    code=error_data.get("code")
+                )
+
+            # Return JSON response
+            if response.status_code == 204:
+                return {}
+            return response.json()
+
+        except httpx.ConnectError as e:
+            if retry_count < self.max_retries:
+                await asyncio.sleep(2 ** retry_count)
+                return await self._request(method, path, params, json, retry_count + 1)
+            raise ConnectionError(f"Failed to connect: {e}")
+        except httpx.TimeoutException:
+            if retry_count < self.max_retries:
+                await asyncio.sleep(2 ** retry_count)
+                return await self._request(method, path, params, json, retry_count + 1)
+            raise ConnectionError("Request timed out")
+
+    # GPU methods
+    async def list_gpus(self) -> List[GPUInfo]:
+        """List available GPU types"""
+        data = await self._request("GET", "/gpus")
+        return [GPUInfo.from_dict(g) for g in data.get("gpus", [])]
+
+    async def get_gpu_availability(self, gpu_type: str) -> Dict[str, Any]:
+        """Get availability for a specific GPU type"""
+        data = await self._request("GET", f"/gpus/{gpu_type}")
+        return data
+
+    # Project methods
+    async def create_project(self, name: str, description: str = None) -> Project:
+        """Create a new project"""
+        data = await self._request("POST", "/projects", json={
+            "name": name,
+            "description": description
+        })
+        return Project.from_dict(data)
+
+    async def list_projects(self, limit: int = 20, offset: int = 0) -> List[Project]:
+        """List projects"""
+        data = await self._request("GET", "/projects", params={
+            "limit": limit,
+            "offset": offset
+        })
+        return [Project.from_dict(p) for p in data.get("projects", [])]
+
+    async def get_project(self, project_id: str) -> Project:
+        """Get a project by ID"""
+        data = await self._request("GET", f"/projects/{project_id}")
+        return Project.from_dict(data)
+
+    async def delete_project(self, project_id: str):
+        """Delete a project"""
+        await self._request("DELETE", f"/projects/{project_id}")
+
+    # Billing methods
+    async def get_wallet_balance(self) -> WalletBalance:
+        """Get wallet balance"""
+        data = await self._request("GET", "/billing/wallet")
+        return WalletBalance.from_dict(data)
+
+    async def get_usage(
+        self,
+        start_date: str = None,
+        end_date: str = None,
+        group_by: str = "day"
+    ) -> UsageSummary:
+        """
+        Get usage summary.
+
+        Args:
+            start_date: Start date (ISO 8601)
+            end_date: End date (ISO 8601)
+            group_by: Grouping (day, week, month)
+
+        Returns:
+            UsageSummary object
+        """
+        params = {"group_by": group_by}
+        if start_date:
+            params["start_date"] = start_date
+        if end_date:
+            params["end_date"] = end_date
+
+        data = await self._request("GET", "/billing/usage", params=params)
+        return UsageSummary.from_dict(data)
+
+    # Webhook methods
+    async def create_webhook(self, url: str, events: List[str]) -> Webhook:
+        """
+        Create a webhook.
+
+        Args:
+            url: Webhook URL
+            events: List of events to subscribe to
+
+        Returns:
+            Webhook object
+        """
+        data = await self._request("POST", "/webhooks", json={
+            "url": url,
+            "events": events
+        })
+        return Webhook.from_dict(data)
+
+    async def list_webhooks(self) -> List[Webhook]:
+        """List webhooks"""
+        data = await self._request("GET", "/webhooks")
+        return [Webhook.from_dict(w) for w in data.get("webhooks", [])]
+
+    async def delete_webhook(self, webhook_id: str):
+        """Delete a webhook"""
+        await self._request("DELETE", f"/webhooks/{webhook_id}")
+
+    # Sync wrappers for simple scripts
+    def _run_async(self, coro):
+        """Run async coroutine in sync context"""
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        return loop.run_until_complete(coro)
+
+    def sync_create_notebook(self, **kwargs) -> Notebook:
+        """Sync wrapper for notebooks.create()"""
+        return self._run_async(self.notebooks.create(**kwargs))
+
+    def sync_get_notebook(self, notebook_id: str) -> Notebook:
+        """Sync wrapper for notebooks.get()"""
+        return self._run_async(self.notebooks.get(notebook_id))
+
+    def sync_list_notebooks(self, **kwargs) -> List[Notebook]:
+        """Sync wrapper for notebooks.list()"""
+        return self._run_async(self.notebooks.list(**kwargs))
+
+    def sync_run(self, code: str, **kwargs) -> Execution:
+        """Sync wrapper for executions.run()"""
+        return self._run_async(self.executions.run(code, **kwargs))
+
+    def sync_list_gpus(self) -> List[GPUInfo]:
+        """Sync wrapper for list_gpus()"""
+        return self._run_async(self.list_gpus())
+
+    def sync_get_wallet_balance(self) -> WalletBalance:
+        """Sync wrapper for get_wallet_balance()"""
+        return self._run_async(self.get_wallet_balance())

podstack/exceptions.py

@@ -0,0 +1,125 @@
+"""
+Podstack Exception Classes
+
+Custom exceptions for handling various error conditions in the Podstack SDK.
+"""
+
+
+class PodstackError(Exception):
+    """Base exception for all Podstack errors"""
+
+    def __init__(self, message: str, code: str = None, details: dict = None):
+        super().__init__(message)
+        self.message = message
+        self.code = code
+        self.details = details or {}
+
+    def __str__(self):
+        if self.code:
+            return f"[{self.code}] {self.message}"
+        return self.message
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(message={self.message!r}, code={self.code!r})"
+
+
+class AuthenticationError(PodstackError):
+    """Raised when API authentication fails"""
+
+    def __init__(self, message: str = "Invalid or missing API key"):
+        super().__init__(message, code="authentication_error")
+
+
+class NotFoundError(PodstackError):
+    """Raised when a requested resource is not found"""
+
+    def __init__(self, resource_type: str, resource_id: str):
+        message = f"{resource_type} '{resource_id}' not found"
+        super().__init__(message, code="not_found")
+        self.resource_type = resource_type
+        self.resource_id = resource_id
+
+
+class RateLimitError(PodstackError):
+    """Raised when API rate limit is exceeded"""
+
+    def __init__(self, retry_after: int = None):
+        message = "Rate limit exceeded"
+        if retry_after:
+            message += f". Retry after {retry_after} seconds"
+        super().__init__(message, code="rate_limit_exceeded")
+        self.retry_after = retry_after
+
+
+class GPUNotAvailableError(PodstackError):
+    """Raised when the requested GPU type is not available"""
+
+    def __init__(self, gpu_type: str, available_types: list = None):
+        message = f"GPU type '{gpu_type}' is not available"
+        super().__init__(message, code="gpu_not_available")
+        self.gpu_type = gpu_type
+        self.available_types = available_types or []
+
+
+class ExecutionTimeoutError(PodstackError):
+    """Raised when code execution times out"""
+
+    def __init__(self, execution_id: str, timeout_seconds: int):
+        message = f"Execution '{execution_id}' timed out after {timeout_seconds} seconds"
+        super().__init__(message, code="execution_timeout")
+        self.execution_id = execution_id
+        self.timeout_seconds = timeout_seconds
+
+
+class ValidationError(PodstackError):
+    """Raised when request validation fails"""
+
+    def __init__(self, message: str, field: str = None):
+        super().__init__(message, code="validation_error")
+        self.field = field
+
+
+class QuotaExceededError(PodstackError):
+    """Raised when account quota is exceeded"""
+
+    def __init__(self, quota_type: str, limit: int, current: int):
+        message = f"{quota_type} quota exceeded: {current}/{limit}"
+        super().__init__(message, code="quota_exceeded")
+        self.quota_type = quota_type
+        self.limit = limit
+        self.current = current
+
+
+class InsufficientBalanceError(PodstackError):
+    """Raised when wallet balance is insufficient"""
+
+    def __init__(self, required: float, available: float):
+        message = f"Insufficient balance: required ${required:.2f}, available ${available:.2f}"
+        super().__init__(message, code="insufficient_balance")
+        self.required = required
+        self.available = available
+
+
+class NotebookStateError(PodstackError):
+    """Raised when notebook is in an invalid state for the operation"""
+
+    def __init__(self, notebook_id: str, current_state: str, required_state: str):
+        message = f"Notebook '{notebook_id}' is {current_state}, must be {required_state}"
+        super().__init__(message, code="invalid_notebook_state")
+        self.notebook_id = notebook_id
+        self.current_state = current_state
+        self.required_state = required_state
+
+
+class ConnectionError(PodstackError):
+    """Raised when connection to the API fails"""
+
+    def __init__(self, message: str = "Failed to connect to Podstack API"):
+        super().__init__(message, code="connection_error")
+
+
+class WebSocketError(PodstackError):
+    """Raised when WebSocket connection fails"""
+
+    def __init__(self, message: str = "WebSocket connection error"):
+        super().__init__(message, code="websocket_error")