podstack 1.2.0__py3-none-any.whl

podstack_gpu/runner.py

@@ -0,0 +1,746 @@
+"""GPU Runner - Main interface for running code on cloud GPUs."""
+
+import os
+import time
+import json
+import threading
+from typing import Optional, Dict, Any, Callable
+from dataclasses import dataclass
+
+import requests
+
+from .exceptions import (
+    PodstackError,
+    AuthenticationError,
+    ValidationError,
+    ExecutionError,
+)
+
+
+@dataclass
+class StatusUpdate:
+    """Real-time status update from execution."""
+    execution_id: str
+    status: str
+    timestamp: str
+    gpu_type: Optional[str] = None
+    gpu_count: Optional[int] = None
+    gpu_fraction: Optional[int] = None
+    queue_position: Optional[int] = None
+    estimated_wait_seconds: Optional[int] = None
+    started_at: Optional[str] = None
+    completed_at: Optional[str] = None
+    gpu_seconds: Optional[float] = None
+    cost_cents: Optional[int] = None
+    error: Optional[str] = None
+    progress: Optional[float] = None
+    message: Optional[str] = None
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "StatusUpdate":
+        """Create StatusUpdate from dictionary."""
+        return cls(
+            execution_id=data.get("execution_id", ""),
+            status=data.get("status", "unknown"),
+            timestamp=data.get("timestamp", ""),
+            gpu_type=data.get("gpu_type"),
+            gpu_count=data.get("gpu_count"),
+            gpu_fraction=data.get("gpu_fraction"),
+            queue_position=data.get("queue_position"),
+            estimated_wait_seconds=data.get("estimated_wait_seconds"),
+            started_at=data.get("started_at"),
+            completed_at=data.get("completed_at"),
+            gpu_seconds=data.get("gpu_seconds"),
+            cost_cents=data.get("cost_cents"),
+            error=data.get("error"),
+            progress=data.get("progress"),
+            message=data.get("message"),
+        )
+
+    @property
+    def is_terminal(self) -> bool:
+        """Check if this is a terminal status."""
+        return self.status in ("completed", "failed", "timeout", "cancelled")
+
+
+@dataclass
+class ExecutionResult:
+    """Result of a GPU execution."""
+    execution_id: str
+    status: str
+    output: str
+    gpu_seconds: float
+    cost_cents: int
+    error: Optional[str] = None
+
+    @property
+    def success(self) -> bool:
+        return self.status == "completed"
+
+    @property
+    def cost_inr(self) -> float:
+        """Cost in INR (paise converted to rupees)."""
+        return self.cost_cents / 100 if self.cost_cents else 0.0
+
+    @property
+    def cost_dollars(self) -> float:
+        """Deprecated: Use cost_inr instead. Returns cost in INR for backwards compatibility."""
+        return self.cost_inr
+
+
+class GPURunner:
+    """
+    Run Python code on cloud GPUs.
+
+    Example:
+        gpu = GPURunner(token, project_id)
+        result = gpu.run('''
+            import torch
+            print(torch.cuda.is_available())
+        ''', gpu="L40S", fraction=25)
+    """
+
+    DEFAULT_API_URL = "https://cloud.podstack.ai/notebooks"
+
+    def __init__(
+        self,
+        token: str,
+        project_id: str,
+        api_url: str = None,
+        timeout: int = 300,
+    ):
+        """
+        Initialize GPU Runner.
+
+        Args:
+            token: API token (psk_xxx for platform tokens, pjt_xxx for project tokens)
+            project_id: Project UUID
+            api_url: API base URL (optional, defaults to PODSTACK_API_URL env var or production URL)
+            timeout: Default execution timeout in seconds
+        """
+        self.token = token
+        self.project_id = project_id
+        # Priority: explicit api_url > PODSTACK_API_URL env var > default URL
+        self.api_url = (api_url or os.environ.get("PODSTACK_API_URL") or self.DEFAULT_API_URL).rstrip("/")
+        self.default_timeout = timeout
+
+        self._session = requests.Session()
+        self._session.headers.update({
+            "Authorization": f"Bearer {token}",
+            "Content-Type": "application/json",
+        })
+
+    def run(
+        self,
+        code: str,
+        gpu: str = "L40S",
+        count: int = 1,
+        fraction: int = 100,
+        timeout: int = None,
+        memory: int = None,
+        env: str = None,
+        pip: list = None,
+        wait: bool = True,
+        poll_interval: float = 2.0,
+    ) -> ExecutionResult:
+        """
+        Run code on a cloud GPU.
+
+        Args:
+            code: Python code to execute
+            gpu: GPU type (A10, L40, L40S, A100-40G, A100-80G, H100)
+            count: Number of GPUs (1-8)
+            fraction: GPU fraction percentage (25, 50, 75, 100)
+            timeout: Max execution time in seconds
+            memory: GPU memory limit in GB (optional)
+            env: Environment preset (ml, cv, nlp, audio, tabular, rl, scientific)
+            pip: List of pip packages to install
+            wait: If True, wait for completion. If False, return immediately.
+            poll_interval: Seconds between status checks when waiting
+
+        Returns:
+            ExecutionResult with output and status
+        """
+        timeout = timeout or self.default_timeout
+
+        # Build annotation
+        annotation = self._build_annotation(
+            gpu=gpu,
+            count=count,
+            fraction=fraction,
+            timeout=timeout,
+            memory=memory,
+            env=env,
+            pip=pip,
+        )
+
+        # Prepend annotation to code
+        full_code = f"{annotation}\n{code}"
+
+        # Submit execution
+        execution_id = self._submit(full_code)
+
+        print(f"Submitted execution: {execution_id}")
+        print(f"  GPU: {gpu} @ {fraction}%")
+
+        if not wait:
+            return ExecutionResult(
+                execution_id=execution_id,
+                status="submitted",
+                output="",
+                gpu_seconds=0,
+                cost_cents=0,
+            )
+
+        # Wait for completion
+        return self._wait_for_completion(execution_id, poll_interval, timeout)
+
+    def _build_annotation(
+        self,
+        gpu: str,
+        count: int,
+        fraction: int,
+        timeout: int,
+        memory: int = None,
+        env: str = None,
+        pip: list = None,
+    ) -> str:
+        """Build the GPU annotation comment."""
+        parts = [f"#@podstack gpu={gpu} count={count} fraction={fraction} timeout={timeout}"]
+
+        if memory:
+            parts.append(f"memory={memory}")
+        if env:
+            parts.append(f"env={env}")
+        if pip:
+            parts.append(f"pip={','.join(pip)}")
+
+        return " ".join(parts)
+
+    def _submit(self, code: str) -> str:
+        """Submit code for execution."""
+        url = f"{self.api_url}/api/v1/executions/submit"
+
+        try:
+            response = self._session.post(
+                url,
+                params={
+                    "code": code,
+                    "project_id": self.project_id,
+                },
+            )
+        except requests.RequestException as e:
+            raise PodstackError(f"Failed to connect to API: {e}")
+
+        if response.status_code == 401:
+            raise AuthenticationError("Invalid or expired token")
+        elif response.status_code == 403:
+            raise AuthenticationError("No access to this project")
+        elif response.status_code == 400:
+            data = response.json()
+            raise ValidationError(data.get("detail", "Invalid request"))
+        elif response.status_code != 202:
+            raise PodstackError(f"API error: {response.status_code} - {response.text}")
+
+        data = response.json()
+        return data["execution_id"]
+
+    def _wait_for_completion(
+        self,
+        execution_id: str,
+        poll_interval: float,
+        timeout: int,
+    ) -> ExecutionResult:
+        """Wait for execution completion using real-time SSE streaming."""
+        # Try SSE streaming first for real-time updates
+        try:
+            return self._wait_with_sse(execution_id, timeout)
+        except Exception as e:
+            print(f"  SSE streaming unavailable ({e}), falling back to polling...")
+            return self._wait_with_polling(execution_id, poll_interval, timeout)
+
+    def _wait_with_sse(
+        self,
+        execution_id: str,
+        timeout: int,
+    ) -> ExecutionResult:
+        """Wait for completion using Server-Sent Events for real-time updates."""
+        url = f"{self.api_url}/api/v1/executions/{execution_id}/status/stream"
+        result_url = f"{self.api_url}/api/v1/executions/{execution_id}/result"
+
+        last_status = None
+        final_update: Optional[StatusUpdate] = None
+
+        try:
+            # Use stream=True for SSE
+            response = self._session.get(url, stream=True, timeout=timeout + 60)
+            if response.status_code != 200:
+                raise PodstackError(f"Failed to connect to status stream: {response.text}")
+
+            for line in response.iter_lines(decode_unicode=True):
+                if line is None:
+                    continue
+
+                line = line.strip()
+
+                # Skip empty lines and comments (heartbeat)
+                if not line or line.startswith(":"):
+                    continue
+
+                # Parse SSE data
+                if line.startswith("data:"):
+                    data_str = line[5:].strip()
+                    try:
+                        data = json.loads(data_str)
+                        update = StatusUpdate.from_dict(data)
+
+                        # Print status updates
+                        if update.status != last_status and update.status != "heartbeat":
+                            self._print_status_update(update)
+                            last_status = update.status
+
+                        # Check for terminal state
+                        if update.is_terminal:
+                            final_update = update
+                            break
+
+                    except json.JSONDecodeError:
+                        continue
+
+        except requests.exceptions.Timeout:
+            raise ExecutionError(
+                f"Timed out waiting for execution after {timeout}s",
+                execution_id=execution_id,
+                status="timeout",
+            )
+        except requests.exceptions.RequestException as e:
+            raise PodstackError(f"Connection error during status streaming: {e}")
+
+        if final_update is None:
+            raise PodstackError("Status stream ended without terminal state")
+
+        # Get final result with output
+        output = ""
+        try:
+            response = self._session.get(result_url)
+            if response.status_code == 200:
+                result_data = response.json()
+                output = "\n".join(
+                    o.get("content", "")
+                    for o in result_data.get("output", [])
+                    if o.get("output_type") in ("stdout", "result")
+                )
+        except Exception:
+            pass
+
+        result = ExecutionResult(
+            execution_id=execution_id,
+            status=final_update.status,
+            output=output,
+            gpu_seconds=final_update.gpu_seconds or 0,
+            cost_cents=final_update.cost_cents or 0,
+            error=final_update.error,
+        )
+
+        # Print final status
+        self._print_result(result)
+
+        if final_update.status == "failed":
+            raise ExecutionError(
+                result.error or "Execution failed",
+                execution_id=execution_id,
+                status=final_update.status,
+            )
+
+        return result
+
+    def _wait_with_polling(
+        self,
+        execution_id: str,
+        poll_interval: float,
+        timeout: int,
+    ) -> ExecutionResult:
+        """Fallback: Poll for execution completion."""
+        url = f"{self.api_url}/api/v1/executions/{execution_id}/status"
+        result_url = f"{self.api_url}/api/v1/executions/{execution_id}/result"
+
+        start_time = time.time()
+        last_status = None
+
+        while True:
+            elapsed = time.time() - start_time
+            if elapsed > timeout + 60:  # Allow extra time for provisioning
+                raise ExecutionError(
+                    f"Timed out waiting for execution after {elapsed:.0f}s",
+                    execution_id=execution_id,
+                    status="timeout",
+                )
+
+            try:
+                response = self._session.get(url)
+                if response.status_code != 200:
+                    raise PodstackError(f"Failed to get status: {response.text}")
+
+                data = response.json()
+                status = data["status"]
+
+                # Print status updates
+                if status != last_status:
+                    self._print_status(status, data)
+                    last_status = status
+
+                # Check terminal states
+                if status in ("completed", "failed", "timeout", "cancelled"):
+                    break
+
+            except requests.RequestException as e:
+                print(f"  Warning: {e}")
+
+            time.sleep(poll_interval)
+
+        # Get final result
+        try:
+            response = self._session.get(result_url)
+            if response.status_code == 200:
+                result_data = response.json()
+                output = "\n".join(
+                    o.get("content", "")
+                    for o in result_data.get("output", [])
+                    if o.get("output_type") in ("stdout", "result")
+                )
+            else:
+                output = ""
+                result_data = data
+        except Exception:
+            output = ""
+            result_data = data
+
+        result = ExecutionResult(
+            execution_id=execution_id,
+            status=status,
+            output=output,
+            gpu_seconds=result_data.get("gpu_seconds", 0) or 0,
+            cost_cents=result_data.get("cost_cents", 0) or 0,
+            error=result_data.get("error"),
+        )
+
+        # Print final status
+        self._print_result(result)
+
+        if status == "failed":
+            raise ExecutionError(
+                result.error or "Execution failed",
+                execution_id=execution_id,
+                status=status,
+            )
+
+        return result
+
+    def _print_status_update(self, update: StatusUpdate):
+        """Print status update from SSE stream."""
+        if update.message:
+            print(f"  {update.message}")
+        else:
+            status_messages = {
+                "pending": "Status: pending",
+                "queued": f"Status: queued (position {update.queue_position or '?'})",
+                "provisioning": "Provisioning GPU...",
+                "running": "Running on GPU...",
+                "completed": "Completed!",
+                "failed": "Failed!",
+                "timeout": "Timed out!",
+                "cancelled": "Cancelled",
+            }
+            print(f"  {status_messages.get(update.status, f'Status: {update.status}')}")
+
+        # Show progress if available
+        if update.progress is not None and update.status == "running":
+            print(f"  Progress: {update.progress:.1f}%")
+
+    def _print_status(self, status: str, data: Dict[str, Any]):
+        """Print status update."""
+        status_messages = {
+            "pending": "  Status: pending",
+            "queued": f"  Status: queued (position {data.get('queue_position', '?')})",
+            "provisioning": "  Provisioning GPU...",
+            "running": "  Running on GPU...",
+            "completed": "  Completed!",
+            "failed": "  Failed!",
+            "timeout": "  Timed out!",
+            "cancelled": "  Cancelled",
+        }
+        print(status_messages.get(status, f"  Status: {status}"))
+
+    def _print_result(self, result: ExecutionResult):
+        """Print execution result summary."""
+        if result.success:
+            print(f"\nExecution completed successfully!")
+            print(f"  GPU time: {result.gpu_seconds:.1f}s")
+            print(f"  Cost: ₹{result.cost_inr:.4f}")
+            if result.output:
+                print(f"\nOutput:\n{result.output}")
+        else:
+            print(f"\nExecution {result.status}")
+            if result.error:
+                print(f"  Error: {result.error}")
+
+    def get_status(self, execution_id: str) -> Dict[str, Any]:
+        """Get current status of an execution."""
+        url = f"{self.api_url}/api/v1/executions/{execution_id}/status"
+
+        response = self._session.get(url)
+        if response.status_code != 200:
+            raise PodstackError(f"Failed to get status: {response.text}")
+
+        return response.json()
+
+    def cancel(self, execution_id: str, reason: str = "Cancelled by user") -> bool:
+        """Cancel a running execution."""
+        url = f"{self.api_url}/api/v1/executions/{execution_id}/cancel"
+
+        response = self._session.post(url, params={"reason": reason})
+        return response.status_code == 200
+
+    def list_gpus(self) -> Dict[str, Any]:
+        """List available GPU types and pricing."""
+        url = f"{self.api_url}/api/v1/gpu-types"
+
+        response = self._session.get(url)
+        if response.status_code != 200:
+            raise PodstackError(f"Failed to get GPU types: {response.text}")
+
+        return response.json()
+
+    def stream_status(
+        self,
+        execution_id: str,
+        callback: Callable[[StatusUpdate], None],
+        timeout: int = 3600,
+    ) -> None:
+        """
+        Stream status updates for an execution with a callback.
+
+        This method blocks until the execution reaches a terminal state.
+        Use this for real-time status tracking in your application.
+
+        Args:
+            execution_id: Execution ID to track
+            callback: Function called with each StatusUpdate
+            timeout: Maximum time to wait in seconds
+
+        Example:
+            def on_status(update):
+                print(f"Status: {update.status}")
+                if update.queue_position:
+                    print(f"Queue position: {update.queue_position}")
+
+            runner.stream_status(execution_id, on_status)
+        """
+        url = f"{self.api_url}/api/v1/executions/{execution_id}/status/stream"
+
+        try:
+            response = self._session.get(url, stream=True, timeout=timeout)
+            if response.status_code != 200:
+                raise PodstackError(f"Failed to connect to status stream: {response.text}")
+
+            for line in response.iter_lines(decode_unicode=True):
+                if line is None:
+                    continue
+
+                line = line.strip()
+
+                # Skip empty lines and comments (heartbeat)
+                if not line or line.startswith(":"):
+                    continue
+
+                # Parse SSE data
+                if line.startswith("data:"):
+                    data_str = line[5:].strip()
+                    try:
+                        data = json.loads(data_str)
+                        update = StatusUpdate.from_dict(data)
+
+                        # Call the callback
+                        if update.status != "heartbeat":
+                            callback(update)
+
+                        # Stop on terminal state
+                        if update.is_terminal:
+                            break
+
+                    except json.JSONDecodeError:
+                        continue
+
+        except requests.exceptions.Timeout:
+            raise ExecutionError(
+                f"Timed out waiting for execution after {timeout}s",
+                execution_id=execution_id,
+                status="timeout",
+            )
+
+    def stream_status_async(
+        self,
+        execution_id: str,
+        callback: Callable[[StatusUpdate], None],
+        timeout: int = 3600,
+    ) -> threading.Thread:
+        """
+        Stream status updates in a background thread.
+
+        Returns immediately with a thread handle. The callback is called
+        from the background thread for each status update.
+
+        Args:
+            execution_id: Execution ID to track
+            callback: Function called with each StatusUpdate
+            timeout: Maximum time to wait in seconds
+
+        Returns:
+            Thread object (already started)
+
+        Example:
+            updates = []
+            def on_status(update):
+                updates.append(update)
+                print(f"Status: {update.status}")
+
+            thread = runner.stream_status_async(execution_id, on_status)
+            # Do other work...
+            thread.join()  # Wait for completion
+        """
+        thread = threading.Thread(
+            target=self.stream_status,
+            args=(execution_id, callback, timeout),
+            daemon=True,
+        )
+        thread.start()
+        return thread
+
+    def stream_logs(
+        self,
+        execution_id: str,
+        timeout: int = 3600,
+    ):
+        """
+        Stream real-time logs in Jupyter notebook with live output display.
+
+        This method displays logs in real-time as they are generated on the GPU.
+        Works with all @podstack annotations - logs will stream automatically.
+
+        Args:
+            execution_id: Execution ID to stream logs for
+            timeout: Maximum time to wait in seconds
+
+        Example (in Jupyter notebook):
+            # Submit execution
+            gpu = GPURunner(token, project_id)
+            result = gpu.run(code, wait=False)
+
+            # Stream logs in real-time
+            gpu.stream_logs(result.execution_id)
+
+        The logs will update live in the notebook output area.
+        """
+        try:
+            # Try to use IPython display for Jupyter notebooks
+            from IPython.display import display, HTML, clear_output
+            import sys
+            from io import StringIO
+
+            # Create output buffer
+            output_area = StringIO()
+
+            print("Streaming logs (press Stop button or wait for completion)...")
+            print("=" * 60)
+
+            url = f"{self.api_url}/api/v1/executions/{execution_id}/status/stream"
+
+            response = self._session.get(url, stream=True, timeout=timeout)
+            if response.status_code != 200:
+                raise PodstackError(f"Failed to connect to log stream: {response.text}")
+
+            for line in response.iter_lines(decode_unicode=True):
+                if line is None:
+                    continue
+
+                line = line.strip()
+
+                # Skip empty lines and comments
+                if not line or line.startswith(":"):
+                    continue
+
+                # Parse SSE data
+                if line.startswith("data:"):
+                    data_str = line[5:].strip()
+                    try:
+                        data = json.loads(data_str)
+
+                        # Check if this is a log message
+                        if data.get("status") == "log":
+                            output_type = data.get("data", {}).get("output_type", "stdout")
+                            content = data.get("data", {}).get("content", "")
+
+                            # Display log content
+                            if output_type == "stderr":
+                                # Red for errors
+                                print(f"\033[91m{content}\033[0m", end="")
+                            else:
+                                # Normal for stdout
+                                print(content, end="")
+
+                            sys.stdout.flush()
+
+                        # Check for terminal state
+                        update = StatusUpdate.from_dict(data)
+                        if update.is_terminal:
+                            print("\n" + "=" * 60)
+                            print(f"Execution {update.status}")
+                            if update.error:
+                                print(f"Error: {update.error}")
+                            if update.gpu_seconds:
+                                print(f"GPU time: {update.gpu_seconds:.1f}s")
+                            if update.cost_cents:
+                                print(f"Cost: ₹{update.cost_cents / 100:.4f}")
+                            break
+
+                    except json.JSONDecodeError:
+                        continue
+
+        except ImportError:
+            # Fallback for non-Jupyter environments
+            print("Warning: IPython not available. Falling back to simple streaming.")
+            self._stream_logs_simple(execution_id, timeout)
+
+    def _stream_logs_simple(self, execution_id: str, timeout: int):
+        """Simple log streaming without Jupyter widgets."""
+        url = f"{self.api_url}/api/v1/executions/{execution_id}/status/stream"
+
+        response = self._session.get(url, stream=True, timeout=timeout)
+        if response.status_code != 200:
+            raise PodstackError(f"Failed to connect to log stream: {response.text}")
+
+        for line in response.iter_lines(decode_unicode=True):
+            if line is None:
+                continue
+
+            line = line.strip()
+
+            if not line or line.startswith(":"):
+                continue
+
+            if line.startswith("data:"):
+                data_str = line[5:].strip()
+                try:
+                    data = json.loads(data_str)
+
+                    if data.get("status") == "log":
+                        content = data.get("data", {}).get("content", "")
+                        print(content, end="")
+
+                    update = StatusUpdate.from_dict(data)
+                    if update.is_terminal:
+                        print(f"\nExecution {update.status}")
+                        break
+
+                except json.JSONDecodeError:
+                    continue