podstack 1.2.0__py3-none-any.whl

podstack/models.py

@@ -0,0 +1,274 @@
+"""
+Podstack Data Models
+
+Data classes representing various resources in the Podstack API.
+"""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Optional, List, Dict, Any
+from enum import Enum
+
+
+class GPUType(str, Enum):
+    """Available GPU types"""
+    A10 = "A10"
+    A100_40GB = "A100_40GB"
+    A100_80GB = "A100_80GB"
+    H100 = "H100"
+    H100_80GB = "H100_80GB"
+    L4 = "L4"
+    L40S = "L40S"
+    T4 = "T4"
+    RTX_4090 = "RTX_4090"
+    RTX_6000_ADA = "RTX_6000_ADA"
+
+    @property
+    def memory_gb(self) -> int:
+        """Get GPU memory in GB"""
+        memory_map = {
+            "A10": 24,
+            "A100_40GB": 40,
+            "A100_80GB": 80,
+            "H100": 80,
+            "H100_80GB": 80,
+            "L4": 24,
+            "L40S": 48,
+            "T4": 16,
+            "RTX_4090": 24,
+            "RTX_6000_ADA": 48,
+        }
+        return memory_map.get(self.value, 0)
+
+    @property
+    def price_per_hour_cents(self) -> int:
+        """Get price per hour in cents"""
+        price_map = {
+            "A10": 90,
+            "A100_40GB": 250,
+            "A100_80GB": 350,
+            "H100": 500,
+            "H100_80GB": 550,
+            "L4": 70,
+            "L40S": 120,
+            "T4": 50,
+            "RTX_4090": 80,
+            "RTX_6000_ADA": 130,
+        }
+        return price_map.get(self.value, 0)
+
+
+class Environment(str, Enum):
+    """Pre-configured environments"""
+    PYTORCH = "pytorch"
+    TENSORFLOW = "tensorflow"
+    JAX = "jax"
+    RAPIDS = "rapids"
+    HUGGINGFACE = "huggingface"
+    CUSTOM = "custom"
+
+
+@dataclass
+class GPUInfo:
+    """Information about an available GPU type"""
+    type: str
+    name: str
+    memory_gb: int
+    price_per_hour_cents: int
+    available: int
+    total: int
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "GPUInfo":
+        return cls(
+            type=data["type"],
+            name=data["name"],
+            memory_gb=data["memory_gb"],
+            price_per_hour_cents=data["price_per_hour_cents"],
+            available=data["available"],
+            total=data["total"]
+        )
+
+
+@dataclass
+class Project:
+    """A project for organizing notebooks"""
+    id: str
+    name: str
+    description: Optional[str]
+    created_at: datetime
+    updated_at: datetime
+    notebook_count: int = 0
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Project":
+        return cls(
+            id=data["id"],
+            name=data["name"],
+            description=data.get("description"),
+            created_at=datetime.fromisoformat(data["created_at"].replace("Z", "+00:00")),
+            updated_at=datetime.fromisoformat(data["updated_at"].replace("Z", "+00:00")),
+            notebook_count=data.get("notebook_count", 0),
+            metadata=data.get("metadata", {})
+        )
+
+
+@dataclass
+class Version:
+    """A notebook version (checkpoint)"""
+    id: str
+    notebook_id: str
+    version_number: int
+    message: Optional[str]
+    created_at: datetime
+    cells_hash: str
+    parent_version_id: Optional[str] = None
+    branch_name: str = "main"
+    is_auto_checkpoint: bool = False
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Version":
+        return cls(
+            id=data["id"],
+            notebook_id=data["notebook_id"],
+            version_number=data["version_number"],
+            message=data.get("message"),
+            created_at=datetime.fromisoformat(data["created_at"].replace("Z", "+00:00")),
+            cells_hash=data["cells_hash"],
+            parent_version_id=data.get("parent_version_id"),
+            branch_name=data.get("branch_name", "main"),
+            is_auto_checkpoint=data.get("is_auto_checkpoint", False)
+        )
+
+
+@dataclass
+class WalletBalance:
+    """User's wallet balance information"""
+    balance_cents: int
+    currency: str
+    last_updated: datetime
+    pending_charges_cents: int = 0
+
+    @property
+    def balance(self) -> float:
+        """Get balance in dollars"""
+        return self.balance_cents / 100
+
+    @property
+    def available_balance(self) -> float:
+        """Get available balance (minus pending charges)"""
+        return (self.balance_cents - self.pending_charges_cents) / 100
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "WalletBalance":
+        return cls(
+            balance_cents=data["balance_cents"],
+            currency=data.get("currency", "USD"),
+            last_updated=datetime.fromisoformat(data["last_updated"].replace("Z", "+00:00")),
+            pending_charges_cents=data.get("pending_charges_cents", 0)
+        )
+
+
+@dataclass
+class UsageRecord:
+    """Usage record for billing"""
+    date: str
+    gpu_type: str
+    gpu_seconds: int
+    cost_cents: int
+    notebook_id: Optional[str] = None
+
+    @property
+    def cost(self) -> float:
+        """Get cost in dollars"""
+        return self.cost_cents / 100
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "UsageRecord":
+        return cls(
+            date=data["date"],
+            gpu_type=data["gpu_type"],
+            gpu_seconds=data["gpu_seconds"],
+            cost_cents=data["cost_cents"],
+            notebook_id=data.get("notebook_id")
+        )
+
+
+@dataclass
+class UsageSummary:
+    """Summary of usage for a period"""
+    total_cost_cents: int
+    total_gpu_seconds: int
+    breakdown: List[UsageRecord]
+    start_date: Optional[str] = None
+    end_date: Optional[str] = None
+
+    @property
+    def total_cost(self) -> float:
+        """Get total cost in dollars"""
+        return self.total_cost_cents / 100
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "UsageSummary":
+        return cls(
+            total_cost_cents=data["total_cost_cents"],
+            total_gpu_seconds=data["total_gpu_seconds"],
+            breakdown=[UsageRecord.from_dict(r) for r in data.get("breakdown", [])],
+            start_date=data.get("start_date"),
+            end_date=data.get("end_date")
+        )
+
+
+@dataclass
+class Webhook:
+    """Webhook configuration"""
+    id: str
+    url: str
+    events: List[str]
+    created_at: datetime
+    is_active: bool = True
+    secret: Optional[str] = None
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Webhook":
+        return cls(
+            id=data["id"],
+            url=data["url"],
+            events=data["events"],
+            created_at=datetime.fromisoformat(data["created_at"].replace("Z", "+00:00")),
+            is_active=data.get("is_active", True),
+            secret=data.get("secret")
+        )
+
+
+@dataclass
+class Cell:
+    """A notebook cell"""
+    id: str
+    cell_type: str  # "code" or "markdown"
+    source: str
+    outputs: List[Dict[str, Any]] = field(default_factory=list)
+    execution_count: Optional[int] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Cell":
+        return cls(
+            id=data["id"],
+            cell_type=data["cell_type"],
+            source=data["source"],
+            outputs=data.get("outputs", []),
+            execution_count=data.get("execution_count"),
+            metadata=data.get("metadata", {})
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "id": self.id,
+            "cell_type": self.cell_type,
+            "source": self.source,
+            "outputs": self.outputs,
+            "execution_count": self.execution_count,
+            "metadata": self.metadata
+        }

podstack/notebook.py

@@ -0,0 +1,410 @@
+"""
+Podstack Notebook Module
+
+Handles notebook operations including creation, execution, and versioning.
+"""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Optional, Dict, Any, List, TYPE_CHECKING
+from enum import Enum
+import asyncio
+
+from .models import Cell, Version
+
+if TYPE_CHECKING:
+    from .client import Client
+
+
+class NotebookStatus(str, Enum):
+    """Status of a notebook"""
+    CREATING = "creating"
+    RUNNING = "running"
+    IDLE = "idle"
+    STOPPING = "stopping"
+    STOPPED = "stopped"
+    FAILED = "failed"
+
+
+@dataclass
+class Notebook:
+    """
+    Represents a GPU notebook.
+
+    Attributes:
+        id: Unique notebook ID
+        name: Notebook name
+        status: Current status
+        gpu_type: Type of GPU allocated
+        environment: Environment preset
+        created_at: Creation timestamp
+        startup_time_ms: Time to start in milliseconds
+        endpoint: WebSocket endpoint for real-time communication
+        jupyter_url: URL to access JupyterLab
+        project_id: Associated project ID
+        idle_timeout_minutes: Auto-shutdown after idle minutes
+        auto_shutdown_enabled: Whether auto-shutdown is enabled
+        metadata: Custom metadata
+    """
+    id: str
+    name: str
+    status: NotebookStatus
+    gpu_type: str
+    environment: str
+    created_at: Optional[datetime] = None
+    startup_time_ms: Optional[float] = None
+    endpoint: Optional[str] = None
+    jupyter_url: Optional[str] = None
+    project_id: Optional[str] = None
+    idle_timeout_minutes: int = 30
+    auto_shutdown_enabled: bool = True
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    cells: List[Cell] = field(default_factory=list)
+
+    _client: Optional["Client"] = field(default=None, repr=False)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any], client: "Client" = None) -> "Notebook":
+        created_at = None
+        if data.get("created_at"):
+            created_at = datetime.fromisoformat(data["created_at"].replace("Z", "+00:00"))
+
+        cells = []
+        if data.get("cells"):
+            cells = [Cell.from_dict(c) for c in data["cells"]]
+
+        return cls(
+            id=data["id"],
+            name=data["name"],
+            status=NotebookStatus(data["status"]),
+            gpu_type=data["gpu_type"],
+            environment=data["environment"],
+            created_at=created_at,
+            startup_time_ms=data.get("startup_time_ms"),
+            endpoint=data.get("endpoint"),
+            jupyter_url=data.get("jupyter_url"),
+            project_id=data.get("project_id"),
+            idle_timeout_minutes=data.get("idle_timeout_minutes", 30),
+            auto_shutdown_enabled=data.get("auto_shutdown_enabled", True),
+            metadata=data.get("metadata", {}),
+            cells=cells,
+            _client=client
+        )
+
+    @property
+    def is_running(self) -> bool:
+        """Check if notebook is running"""
+        return self.status in (NotebookStatus.RUNNING, NotebookStatus.IDLE)
+
+    @property
+    def is_ready(self) -> bool:
+        """Check if notebook is ready for execution"""
+        return self.status == NotebookStatus.RUNNING
+
+    async def refresh(self) -> "Notebook":
+        """Refresh notebook status from API"""
+        if not self._client:
+            raise RuntimeError("Notebook not bound to client")
+
+        data = await self._client._request("GET", f"/notebooks/{self.id}")
+        updated = Notebook.from_dict(data, self._client)
+
+        # Update self with new data
+        self.status = updated.status
+        self.endpoint = updated.endpoint
+        self.jupyter_url = updated.jupyter_url
+        self.cells = updated.cells
+
+        return self
+
+    async def wait_ready(self, timeout: float = 60, poll_interval: float = 0.5) -> "Notebook":
+        """
+        Wait for notebook to be ready.
+
+        Args:
+            timeout: Maximum seconds to wait
+            poll_interval: Seconds between status checks
+
+        Returns:
+            Updated Notebook object
+        """
+        from .exceptions import NotebookStateError
+
+        start_time = asyncio.get_event_loop().time()
+
+        while self.status == NotebookStatus.CREATING:
+            if (asyncio.get_event_loop().time() - start_time) > timeout:
+                raise NotebookStateError(self.id, self.status.value, "running")
+
+            await asyncio.sleep(poll_interval)
+            await self.refresh()
+
+        if self.status == NotebookStatus.FAILED:
+            raise NotebookStateError(self.id, "failed", "running")
+
+        return self
+
+    async def execute(
+        self,
+        code: str,
+        timeout_seconds: int = 300,
+        wait: bool = True
+    ) -> "Execution":
+        """
+        Execute code in this notebook.
+
+        Args:
+            code: Python code to execute
+            timeout_seconds: Maximum execution time
+            wait: Whether to wait for completion
+
+        Returns:
+            Execution object
+        """
+        from .execution import Execution
+
+        if not self._client:
+            raise RuntimeError("Notebook not bound to client")
+
+        data = await self._client._request(
+            "POST",
+            f"/notebooks/{self.id}/execute",
+            json={"code": code, "timeout_seconds": timeout_seconds}
+        )
+
+        execution = Execution.from_dict(data, self._client)
+        execution.notebook_id = self.id
+
+        if wait:
+            await execution.wait(timeout=timeout_seconds)
+
+        return execution
+
+    async def stop(self) -> "Notebook":
+        """Stop the notebook"""
+        if not self._client:
+            raise RuntimeError("Notebook not bound to client")
+
+        await self._client._request("POST", f"/notebooks/{self.id}/stop")
+        await self.refresh()
+        return self
+
+    async def start(self) -> "Notebook":
+        """Start a stopped notebook"""
+        if not self._client:
+            raise RuntimeError("Notebook not bound to client")
+
+        await self._client._request("POST", f"/notebooks/{self.id}/start")
+        await self.refresh()
+        return self
+
+    async def delete(self):
+        """Delete the notebook"""
+        if not self._client:
+            raise RuntimeError("Notebook not bound to client")
+
+        await self._client._request("DELETE", f"/notebooks/{self.id}")
+
+    # Versioning methods
+    async def save(self, message: str = None) -> Version:
+        """
+        Save notebook state (create a version).
+
+        Args:
+            message: Optional commit message
+
+        Returns:
+            Version object
+        """
+        if not self._client:
+            raise RuntimeError("Notebook not bound to client")
+
+        data = await self._client._request(
+            "POST",
+            f"/notebooks/{self.id}/versions",
+            json={
+                "message": message,
+                "cells": [c.to_dict() for c in self.cells] if self.cells else None
+            }
+        )
+        return Version.from_dict(data)
+
+    async def list_versions(self, branch: str = None, limit: int = 20) -> List[Version]:
+        """
+        List notebook versions.
+
+        Args:
+            branch: Filter by branch name
+            limit: Maximum results
+
+        Returns:
+            List of Version objects
+        """
+        if not self._client:
+            raise RuntimeError("Notebook not bound to client")
+
+        params = {"limit": limit}
+        if branch:
+            params["branch"] = branch
+
+        data = await self._client._request(
+            "GET",
+            f"/notebooks/{self.id}/versions",
+            params=params
+        )
+        return [Version.from_dict(v) for v in data.get("versions", [])]
+
+    async def get_version(self, version_id: str) -> Version:
+        """Get a specific version"""
+        if not self._client:
+            raise RuntimeError("Notebook not bound to client")
+
+        data = await self._client._request(
+            "GET",
+            f"/notebooks/{self.id}/versions/{version_id}"
+        )
+        return Version.from_dict(data)
+
+    async def restore_version(self, version_id: str) -> "Notebook":
+        """
+        Restore notebook to a previous version.
+
+        Args:
+            version_id: Version to restore
+
+        Returns:
+            Updated Notebook object
+        """
+        if not self._client:
+            raise RuntimeError("Notebook not bound to client")
+
+        await self._client._request(
+            "POST",
+            f"/notebooks/{self.id}/versions/{version_id}/restore"
+        )
+        await self.refresh()
+        return self
+
+    async def create_branch(self, name: str, from_version_id: str = None) -> Dict[str, Any]:
+        """
+        Create a branch from a version.
+
+        Args:
+            name: Branch name
+            from_version_id: Version to branch from (default: latest)
+
+        Returns:
+            Branch info
+        """
+        if not self._client:
+            raise RuntimeError("Notebook not bound to client")
+
+        data = await self._client._request(
+            "POST",
+            f"/notebooks/{self.id}/branches",
+            json={"name": name, "from_version_id": from_version_id}
+        )
+        return data
+
+    async def list_branches(self) -> List[Dict[str, Any]]:
+        """List notebook branches"""
+        if not self._client:
+            raise RuntimeError("Notebook not bound to client")
+
+        data = await self._client._request("GET", f"/notebooks/{self.id}/branches")
+        return data.get("branches", [])
+
+    def __str__(self) -> str:
+        return f"Notebook({self.id}, name={self.name!r}, status={self.status.value})"
+
+
+class NotebooksAPI:
+    """API for managing notebooks"""
+
+    def __init__(self, client: "Client"):
+        self._client = client
+
+    async def create(
+        self,
+        name: str,
+        gpu_type: str = "A10",
+        environment: str = "pytorch",
+        project_id: str = None,
+        idle_timeout_minutes: int = 30,
+        auto_shutdown_enabled: bool = True,
+        metadata: Dict[str, Any] = None,
+        wait_ready: bool = True
+    ) -> Notebook:
+        """
+        Create a new notebook.
+
+        Args:
+            name: Notebook name
+            gpu_type: GPU type (A10, A100, H100, etc.)
+            environment: Environment preset (pytorch, tensorflow, jax, etc.)
+            project_id: Optional project ID
+            idle_timeout_minutes: Auto-shutdown after idle minutes
+            auto_shutdown_enabled: Whether to enable auto-shutdown
+            metadata: Custom metadata
+            wait_ready: Whether to wait for notebook to be ready
+
+        Returns:
+            Notebook object
+        """
+        payload = {
+            "name": name,
+            "gpu_type": gpu_type,
+            "environment": environment,
+            "idle_timeout_minutes": idle_timeout_minutes,
+            "auto_shutdown_enabled": auto_shutdown_enabled
+        }
+
+        if project_id:
+            payload["project_id"] = project_id
+        if metadata:
+            payload["metadata"] = metadata
+
+        data = await self._client._request("POST", "/notebooks", json=payload)
+        notebook = Notebook.from_dict(data, self._client)
+
+        if wait_ready:
+            await notebook.wait_ready()
+
+        return notebook
+
+    async def get(self, notebook_id: str) -> Notebook:
+        """Get a notebook by ID"""
+        data = await self._client._request("GET", f"/notebooks/{notebook_id}")
+        return Notebook.from_dict(data, self._client)
+
+    async def list(
+        self,
+        status: NotebookStatus = None,
+        project_id: str = None,
+        limit: int = 20,
+        offset: int = 0
+    ) -> List[Notebook]:
+        """
+        List notebooks.
+
+        Args:
+            status: Filter by status
+            project_id: Filter by project
+            limit: Maximum results
+            offset: Pagination offset
+
+        Returns:
+            List of Notebook objects
+        """
+        params = {"limit": limit, "offset": offset}
+        if status:
+            params["status"] = status.value
+        if project_id:
+            params["project_id"] = project_id
+
+        data = await self._client._request("GET", "/notebooks", params=params)
+        return [Notebook.from_dict(n, self._client) for n in data.get("notebooks", [])]
+
+    async def delete(self, notebook_id: str):
+        """Delete a notebook"""
+        await self._client._request("DELETE", f"/notebooks/{notebook_id}")