podstack 1.2.0__py3-none-any.whl

podstack/execution.py

@@ -0,0 +1,291 @@
+"""
+Podstack Execution Module
+
+Handles code execution operations.
+"""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Optional, Dict, Any, List, TYPE_CHECKING
+from enum import Enum
+import asyncio
+
+if TYPE_CHECKING:
+    from .client import Client
+
+
+class ExecutionStatus(str, Enum):
+    """Status of an execution"""
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+    TIMEOUT = "timeout"
+
+
+@dataclass
+class ExecutionOutput:
+    """Output from an execution"""
+    output_type: str  # "stdout", "stderr", "display_data", "execute_result", "error"
+    data: Any
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ExecutionOutput":
+        return cls(
+            output_type=data["output_type"],
+            data=data["data"],
+            metadata=data.get("metadata", {})
+        )
+
+
+@dataclass
+class Execution:
+    """
+    Represents a code execution.
+
+    Attributes:
+        id: Unique execution ID
+        status: Current status
+        code: The code that was executed
+        output: Execution output (stdout)
+        error: Error message if failed
+        execution_time_ms: Execution time in milliseconds
+        gpu_memory_used_mb: GPU memory used in MB
+        created_at: When the execution was created
+        completed_at: When the execution completed
+    """
+    id: str
+    status: ExecutionStatus
+    code: str
+    output: Optional[str] = None
+    error: Optional[str] = None
+    execution_time_ms: Optional[float] = None
+    gpu_memory_used_mb: Optional[int] = None
+    created_at: Optional[datetime] = None
+    completed_at: Optional[datetime] = None
+    notebook_id: Optional[str] = None
+    outputs: List[ExecutionOutput] = field(default_factory=list)
+
+    _client: Optional["Client"] = field(default=None, repr=False)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any], client: "Client" = None) -> "Execution":
+        created_at = None
+        if data.get("created_at"):
+            created_at = datetime.fromisoformat(data["created_at"].replace("Z", "+00:00"))
+
+        completed_at = None
+        if data.get("completed_at"):
+            completed_at = datetime.fromisoformat(data["completed_at"].replace("Z", "+00:00"))
+
+        outputs = []
+        if data.get("outputs"):
+            outputs = [ExecutionOutput.from_dict(o) for o in data["outputs"]]
+
+        return cls(
+            id=data["id"] if "id" in data else data.get("execution_id", ""),
+            status=ExecutionStatus(data["status"]),
+            code=data.get("code", ""),
+            output=data.get("output"),
+            error=data.get("error"),
+            execution_time_ms=data.get("execution_time_ms"),
+            gpu_memory_used_mb=data.get("gpu_memory_used_mb"),
+            created_at=created_at,
+            completed_at=completed_at,
+            notebook_id=data.get("notebook_id"),
+            outputs=outputs,
+            _client=client
+        )
+
+    @property
+    def is_complete(self) -> bool:
+        """Check if execution is complete"""
+        return self.status in (
+            ExecutionStatus.COMPLETED,
+            ExecutionStatus.FAILED,
+            ExecutionStatus.CANCELLED,
+            ExecutionStatus.TIMEOUT
+        )
+
+    @property
+    def is_success(self) -> bool:
+        """Check if execution completed successfully"""
+        return self.status == ExecutionStatus.COMPLETED
+
+    async def refresh(self) -> "Execution":
+        """Refresh execution status from API"""
+        if not self._client:
+            raise RuntimeError("Execution not bound to client")
+
+        data = await self._client._request("GET", f"/executions/{self.id}")
+        updated = Execution.from_dict(data, self._client)
+
+        # Update self with new data
+        self.status = updated.status
+        self.output = updated.output
+        self.error = updated.error
+        self.execution_time_ms = updated.execution_time_ms
+        self.gpu_memory_used_mb = updated.gpu_memory_used_mb
+        self.completed_at = updated.completed_at
+        self.outputs = updated.outputs
+
+        return self
+
+    async def wait(self, poll_interval: float = 0.5, timeout: float = None) -> "Execution":
+        """
+        Wait for execution to complete.
+
+        Args:
+            poll_interval: Seconds between status checks
+            timeout: Maximum seconds to wait (None for no limit)
+
+        Returns:
+            Updated Execution object
+        """
+        start_time = asyncio.get_event_loop().time()
+
+        while not self.is_complete:
+            if timeout and (asyncio.get_event_loop().time() - start_time) > timeout:
+                from .exceptions import ExecutionTimeoutError
+                raise ExecutionTimeoutError(self.id, int(timeout))
+
+            await asyncio.sleep(poll_interval)
+            await self.refresh()
+
+        return self
+
+    async def cancel(self) -> "Execution":
+        """Cancel the execution"""
+        if not self._client:
+            raise RuntimeError("Execution not bound to client")
+
+        await self._client._request("POST", f"/executions/{self.id}/cancel")
+        await self.refresh()
+        return self
+
+    def get_stdout(self) -> str:
+        """Get stdout output"""
+        if self.output:
+            return self.output
+
+        stdout_parts = []
+        for out in self.outputs:
+            if out.output_type == "stdout":
+                stdout_parts.append(str(out.data))
+        return "".join(stdout_parts)
+
+    def get_stderr(self) -> str:
+        """Get stderr output"""
+        stderr_parts = []
+        for out in self.outputs:
+            if out.output_type == "stderr":
+                stderr_parts.append(str(out.data))
+        return "".join(stderr_parts)
+
+    def get_display_data(self) -> List[Dict[str, Any]]:
+        """Get display data outputs (images, HTML, etc.)"""
+        return [
+            {"data": out.data, "metadata": out.metadata}
+            for out in self.outputs
+            if out.output_type == "display_data"
+        ]
+
+    def __str__(self) -> str:
+        return f"Execution({self.id}, status={self.status.value})"
+
+
+class ExecutionsAPI:
+    """API for managing executions"""
+
+    def __init__(self, client: "Client"):
+        self._client = client
+
+    async def create(
+        self,
+        code: str,
+        gpu_type: str = "A10",
+        environment: str = "pytorch",
+        timeout_seconds: int = 300
+    ) -> Execution:
+        """
+        Create a serverless execution (run code without managing notebooks).
+
+        Args:
+            code: Python code to execute
+            gpu_type: GPU type to use
+            environment: Environment preset
+            timeout_seconds: Maximum execution time
+
+        Returns:
+            Execution object
+        """
+        data = await self._client._request("POST", "/executions", json={
+            "code": code,
+            "gpu_type": gpu_type,
+            "environment": environment,
+            "timeout_seconds": timeout_seconds
+        })
+        return Execution.from_dict(data, self._client)
+
+    async def get(self, execution_id: str) -> Execution:
+        """Get an execution by ID"""
+        data = await self._client._request("GET", f"/executions/{execution_id}")
+        return Execution.from_dict(data, self._client)
+
+    async def list(
+        self,
+        notebook_id: str = None,
+        status: ExecutionStatus = None,
+        limit: int = 20,
+        offset: int = 0
+    ) -> List[Execution]:
+        """
+        List executions.
+
+        Args:
+            notebook_id: Filter by notebook
+            status: Filter by status
+            limit: Maximum results
+            offset: Pagination offset
+
+        Returns:
+            List of Execution objects
+        """
+        params = {"limit": limit, "offset": offset}
+        if notebook_id:
+            params["notebook_id"] = notebook_id
+        if status:
+            params["status"] = status.value
+
+        data = await self._client._request("GET", "/executions", params=params)
+        return [Execution.from_dict(e, self._client) for e in data.get("executions", [])]
+
+    async def run(
+        self,
+        code: str,
+        gpu_type: str = "A10",
+        environment: str = "pytorch",
+        timeout_seconds: int = 300,
+        wait: bool = True
+    ) -> Execution:
+        """
+        Run code and optionally wait for completion.
+
+        Args:
+            code: Python code to execute
+            gpu_type: GPU type to use
+            environment: Environment preset
+            timeout_seconds: Maximum execution time
+            wait: Whether to wait for completion
+
+        Returns:
+            Execution object
+        """
+        execution = await self.create(code, gpu_type, environment, timeout_seconds)
+
+        if wait:
+            await execution.wait(timeout=timeout_seconds)
+
+        return execution