podstack 1.2.0__py3-none-any.whl

podstack/gpu_runner.py

@@ -0,0 +1,1141 @@
+"""
+Podstack GPU Runner
+
+Handles actual GPU provisioning and remote code execution.
+"""
+
+import os
+import sys
+import time
+import inspect
+import textwrap
+import logging
+import threading
+from typing import Optional, Dict, Any, Callable, Union, List, Iterator
+from dataclasses import dataclass, field
+import httpx
+
+# Configure logging
+logger = logging.getLogger("podstack.gpu_runner")
+
+
+def is_jupyter() -> bool:
+    """Check if running in a Jupyter notebook."""
+    try:
+        from IPython import get_ipython
+        shell = get_ipython()
+        if shell is None:
+            return False
+        if shell.__class__.__name__ == 'ZMQInteractiveShell':
+            return True  # Jupyter notebook or qtconsole
+        if shell.__class__.__name__ == 'TerminalInteractiveShell':
+            return False  # Terminal IPython
+    except (ImportError, NameError):
+        pass
+    return False
+
+
+class OutputStreamer:
+    """Handles real-time output streaming from GPU executions."""
+
+    def __init__(self, execution_id: str, api_url: str, headers: Dict[str, str], timeout: float = 30.0):
+        self.execution_id = execution_id
+        self.api_url = api_url
+        self.headers = headers
+        self.timeout = timeout
+        self._output_buffer: List[str] = []
+        self._status = "running"
+        self._final_result: Dict[str, Any] = {}
+        self._jupyter_widget = None
+
+    def stream(self, show_output: bool = True) -> Iterator[Dict[str, Any]]:
+        """
+        Stream output events from the execution.
+
+        Yields:
+            Dict with event data (type, content, timestamp, etc.)
+        """
+        url = f"{self.api_url}/api/v1/executions/{self.execution_id}/stream"
+
+        try:
+            with httpx.Client(timeout=httpx.Timeout(self.timeout, read=None)) as client:
+                with client.stream("GET", url, headers=self.headers) as response:
+                    if response.status_code != 200:
+                        raise RuntimeError(f"Failed to connect to stream: HTTP {response.status_code}")
+
+                    for line in response.iter_lines():
+                        if not line:
+                            continue
+
+                        # SSE format: "data: {...}"
+                        if line.startswith("data: "):
+                            data_str = line[6:]  # Remove "data: " prefix
+                            try:
+                                event = __import__('json').loads(data_str)
+                                yield event
+
+                                # Track status
+                                if "status" in event:
+                                    self._status = event["status"]
+                                    self._final_result = event
+
+                                # Display output
+                                if show_output and event.get("type") in ("stdout", "stderr", "output"):
+                                    content = event.get("content", "")
+                                    if content:
+                                        self._output_buffer.append(content)
+                                        self._display_output(content, event.get("type", "stdout"))
+
+                            except __import__('json').JSONDecodeError as e:
+                                logger.warning(f"Failed to parse SSE event: {e}")
+                                continue
+
+        except httpx.ConnectError as e:
+            raise ConnectionError(f"Failed to connect to stream: {e}")
+        except httpx.TimeoutException:
+            raise TimeoutError(f"Stream connection timed out")
+
+    def _display_output(self, content: str, output_type: str = "stdout"):
+        """Display output, with special handling for Jupyter."""
+        if is_jupyter():
+            # In Jupyter, print directly (output appears in cell)
+            if output_type == "stderr":
+                sys.stderr.write(content)
+                sys.stderr.flush()
+            else:
+                sys.stdout.write(content)
+                sys.stdout.flush()
+        else:
+            # Terminal output
+            if output_type == "stderr":
+                sys.stderr.write(content)
+                sys.stderr.flush()
+            else:
+                sys.stdout.write(content)
+                sys.stdout.flush()
+
+    def get_full_output(self) -> str:
+        """Get the complete buffered output."""
+        return "".join(self._output_buffer)
+
+    def get_final_result(self) -> Dict[str, Any]:
+        """Get the final result after streaming completes."""
+        return self._final_result
+
+
+class PodstackError(Exception):
+    """Base exception for Podstack SDK errors."""
+    pass
+
+
+class PodstackTimeoutError(PodstackError):
+    """Raised when GPU execution times out."""
+    def __init__(self, execution_id: str, timeout: int, last_status: str, message: str = None):
+        self.execution_id = execution_id
+        self.timeout = timeout
+        self.last_status = last_status
+        self.message = message or f"Execution {execution_id} timed out after {timeout}s (last status: {last_status})"
+        super().__init__(self.message)
+
+
+class PodstackExecutionError(PodstackError):
+    """Raised when GPU execution fails."""
+    def __init__(self, execution_id: str, error: str, output: str = None):
+        self.execution_id = execution_id
+        self.error = error
+        self.output = output
+        message = f"Execution {execution_id} failed: {error}"
+        if output:
+            message += f"\n\nOutput:\n{output[:1000]}"
+        super().__init__(message)
+
+
+class PodstackProvisioningError(PodstackError):
+    """Raised when GPU provisioning fails or takes too long."""
+    def __init__(self, execution_id: str, message: str):
+        self.execution_id = execution_id
+        self.message = message
+        super().__init__(message)
+
+
+@dataclass
+class GPUExecutionResult:
+    """Result of a GPU execution."""
+    execution_id: str
+    status: str
+    output: str = ""
+    error: Optional[str] = None
+    gpu_seconds: float = 0.0
+    cost_paise: int = 0
+    gpu_type: str = ""
+    gpu_count: int = 1
+    success: bool = False
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "GPUExecutionResult":
+        return cls(
+            execution_id=data.get("execution_id", ""),
+            status=data.get("status", "unknown"),
+            output=data.get("output", ""),
+            error=data.get("error"),
+            gpu_seconds=data.get("gpu_seconds", 0.0),
+            cost_paise=data.get("actual_cost_paise", 0),
+            gpu_type=data.get("gpu_type", ""),
+            gpu_count=data.get("gpu_count", 1),
+            success=data.get("status") == "completed"
+        )
+
+
+class GPURunner:
+    """
+    GPU Runner for executing code on remote GPU instances.
+
+    Usage:
+        runner = GPURunner(api_key="...", project_id="...")
+
+        # Execute code on GPU
+        result = runner.run('''
+            import torch
+            print(f"GPU: {torch.cuda.get_device_name(0)}")
+        ''', gpu="L40S", fraction=100)
+
+        print(result.output)
+    """
+
+    def __init__(
+        self,
+        api_key: str = None,
+        project_id: str = None,
+        api_url: str = None,
+        timeout: float = 30.0
+    ):
+        """
+        Initialize the GPU Runner.
+
+        Args:
+            api_key: API key or JWT token for authentication
+            project_id: Project ID for billing and tracking
+            api_url: Notebook service API URL
+            timeout: HTTP request timeout in seconds
+        """
+        self.api_key = api_key or os.environ.get("PODSTACK_API_KEY")
+        self.project_id = project_id or os.environ.get("PODSTACK_PROJECT_ID")
+
+        # Determine API URL - prefer internal cluster URL if running inside K8s
+        if api_url:
+            self.api_url = api_url
+        elif os.environ.get("PODSTACK_NOTEBOOK_SERVICE_URL"):
+            self.api_url = os.environ.get("PODSTACK_NOTEBOOK_SERVICE_URL")
+        elif os.environ.get("KUBERNETES_SERVICE_HOST"):
+            # Running inside K8s cluster - use internal service URL
+            self.api_url = "http://podstack-svc-podstack-services-notebook.podstack.svc.cluster.local:8084"
+        else:
+            # External access
+            self.api_url = "https://cloud.podstack.ai/notebooks"
+        self.timeout = timeout
+
+        if not self.api_key:
+            raise ValueError("API key is required. Set PODSTACK_API_KEY or pass api_key.")
+        if not self.project_id:
+            raise ValueError("Project ID is required. Set PODSTACK_PROJECT_ID or pass project_id.")
+
+        # Debug info
+        if os.environ.get("PODSTACK_DEBUG"):
+            print(f"[Podstack Debug] API URL: {self.api_url}")
+            print(f"[Podstack Debug] Project ID: {self.project_id}")
+
+    def _get_headers(self) -> Dict[str, str]:
+        """Get request headers."""
+        return {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+            "User-Agent": "podstack-python-sdk/1.0.3"
+        }
+
+    def _build_annotation(
+        self,
+        gpu: str = "L40S",
+        count: int = 1,
+        fraction: int = 100,
+        timeout: int = 3600,
+        env: str = None,
+        pip: Union[str, list] = None,
+        uv: Union[str, list] = None,
+        conda: Union[str, list] = None,
+        requirements: str = None,
+        use_uv: bool = False
+    ) -> str:
+        """Build the @podstack annotation string."""
+        parts = [f"#@podstack gpu={gpu}"]
+        if count > 1:
+            parts.append(f"count={count}")
+        if fraction != 100:
+            parts.append(f"fraction={fraction}")
+        if timeout != 3600:
+            parts.append(f"timeout={timeout}")
+        if env:
+            parts.append(f"env={env}")
+
+        # Handle pip packages
+        if pip:
+            if isinstance(pip, list):
+                pip = ",".join(pip)
+            parts.append(f"pip={pip}")
+
+        # Handle uv packages
+        if uv:
+            if isinstance(uv, list):
+                uv = ",".join(uv)
+            parts.append(f"uv={uv}")
+
+        # Handle conda packages
+        if conda:
+            if isinstance(conda, list):
+                conda = ",".join(conda)
+            parts.append(f"conda={conda}")
+
+        # Handle requirements.txt
+        if requirements:
+            parts.append(f"requirements={requirements}")
+
+        # Use uv for pip installs
+        if use_uv:
+            parts.append("use_uv=true")
+
+        return " ".join(parts)
+
+    def submit(
+        self,
+        code: str,
+        gpu: str = "L40S",
+        count: int = 1,
+        fraction: int = 100,
+        timeout: int = 3600,
+        env: str = None,
+        pip: Union[str, list] = None,
+        uv: Union[str, list] = None,
+        conda: Union[str, list] = None,
+        requirements: str = None,
+        use_uv: bool = False,
+        add_annotation: bool = True
+    ) -> Dict[str, Any]:
+        """
+        Submit code for GPU execution (non-blocking).
+
+        Args:
+            code: Python code to execute
+            gpu: GPU type (L40S, A100-40G, A100-80G, H100, A10, T4)
+            count: Number of GPUs (1-8)
+            fraction: GPU time-slice percentage (25, 50, 75, 100)
+            timeout: Maximum execution time in seconds
+            env: Environment preset (ml, nlp, cv, audio, tabular, rl, scientific)
+            pip: Pip packages to install (string "pkg1,pkg2" or list ["pkg1", "pkg2"])
+            uv: UV packages to install (faster than pip)
+            conda: Conda packages to install (string or list)
+            requirements: Path to requirements.txt file to install
+            use_uv: Use uv instead of pip for installation (faster)
+            add_annotation: Whether to add @podstack annotation to code
+
+        Returns:
+            Dict with execution_id and status
+        """
+        # Read requirements.txt if provided
+        requirements_content = None
+        if requirements and os.path.exists(requirements):
+            with open(requirements, 'r') as f:
+                requirements_content = f.read()
+
+        # Add annotation if not present
+        if add_annotation and not code.strip().startswith("#@podstack"):
+            annotation = self._build_annotation(gpu, count, fraction, timeout, env, pip, uv, conda, requirements, use_uv)
+            code = f"{annotation}\n\n{code}"
+
+        # Build installation code for packages
+        install_parts = []
+
+        # Helper function for streaming subprocess output
+        install_helper = '''
+# Helper function to stream install output with progress
+import subprocess
+import sys
+
+def _stream_install(cmd, description):
+    """Run command and stream output in real-time."""
+    print(f"\\n{'='*60}")
+    print(f"[Podstack] {description}")
+    print(f"{'='*60}")
+    print(f"$ {' '.join(cmd)}\\n")
+    sys.stdout.flush()
+
+    process = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1
+    )
+
+    for line in process.stdout:
+        print(line, end='')
+        sys.stdout.flush()
+
+    process.wait()
+    if process.returncode != 0:
+        raise subprocess.CalledProcessError(process.returncode, cmd)
+
+    print(f"\\n[Podstack] ✓ {description} - Complete!")
+    print(f"{'='*60}\\n")
+    sys.stdout.flush()
+'''
+
+        # Handle pip packages (use uv if use_uv=True)
+        if pip:
+            if isinstance(pip, list):
+                pip_list = pip
+            else:
+                pip_list = [p.strip() for p in pip.replace(",", " ").split() if p.strip()]
+            if use_uv:
+                install_parts.append(f'''
+# Install pip packages using uv (faster)
+_pip_packages = {repr(pip_list)}
+_stream_install(
+    ["uv", "pip", "install", "--system", "--progress"] + _pip_packages,
+    f"Installing pip packages with uv: {{' '.join(_pip_packages)}}"
+)
+''')
+            else:
+                install_parts.append(f'''
+# Install pip packages
+_pip_packages = {repr(pip_list)}
+_stream_install(
+    [sys.executable, "-m", "pip", "install", "--progress-bar", "on"] + _pip_packages,
+    f"Installing pip packages: {{' '.join(_pip_packages)}}"
+)
+''')
+
+        # Handle uv packages (always use uv)
+        if uv:
+            if isinstance(uv, list):
+                uv_list = uv
+            else:
+                uv_list = [u.strip() for u in uv.replace(",", " ").split() if u.strip()]
+            install_parts.append(f'''
+# Install packages using uv (fast package manager)
+_uv_packages = {repr(uv_list)}
+_stream_install(
+    ["uv", "pip", "install", "--system", "--progress"] + _uv_packages,
+    f"Installing packages with uv: {{' '.join(_uv_packages)}}"
+)
+''')
+
+        # Handle conda/mamba packages (using micromamba)
+        if conda:
+            if isinstance(conda, list):
+                conda_list = conda
+            else:
+                conda_list = [c.strip() for c in conda.replace(",", " ").split() if c.strip()]
+            install_parts.append(f'''
+# Install conda packages using micromamba (fast conda alternative)
+import shutil
+_conda_packages = {repr(conda_list)}
+# Use micromamba if available, fall back to mamba, then conda
+if shutil.which("micromamba"):
+    _pkg_mgr = "micromamba"
+elif shutil.which("mamba"):
+    _pkg_mgr = "mamba"
+else:
+    _pkg_mgr = "conda"
+_stream_install(
+    [_pkg_mgr, "install", "-y"] + _conda_packages,
+    f"Installing conda packages: {{' '.join(_conda_packages)}}"
+)
+''')
+
+        # Handle requirements.txt (use uv if use_uv=True)
+        if requirements_content:
+            if use_uv:
+                install_parts.append(f'''
+# Install from requirements.txt using uv (faster)
+_requirements = """{requirements_content}"""
+
+with open("/tmp/requirements.txt", "w") as f:
+    f.write(_requirements)
+
+print("\\n[Podstack] requirements.txt contents:")
+print(_requirements)
+
+_stream_install(
+    ["uv", "pip", "install", "--system", "--progress", "-r", "/tmp/requirements.txt"],
+    "Installing from requirements.txt with uv"
+)
+''')
+            else:
+                install_parts.append(f'''
+# Install from requirements.txt
+_requirements = """{requirements_content}"""
+
+with open("/tmp/requirements.txt", "w") as f:
+    f.write(_requirements)
+
+print("\\n[Podstack] requirements.txt contents:")
+print(_requirements)
+
+_stream_install(
+    [sys.executable, "-m", "pip", "install", "--progress-bar", "on", "-r", "/tmp/requirements.txt"],
+    "Installing from requirements.txt"
+)
+''')
+
+        # Prepend all installation code with helper function
+        if install_parts:
+            install_code = install_helper + "\n".join(install_parts) + "\nprint('\\n[Podstack] ✓ All packages installed successfully!\\n')\n\n"
+            code = install_code + code
+
+        url = f"{self.api_url}/api/v1/executions/submit"
+
+        with httpx.Client(timeout=self.timeout) as client:
+            try:
+                response = client.post(
+                    url,
+                    headers=self._get_headers(),
+                    json={
+                        "code": code,
+                        "project_id": self.project_id
+                    }
+                )
+            except httpx.ConnectError as e:
+                raise ConnectionError(f"Failed to connect to {url}: {e}")
+            except httpx.TimeoutException:
+                raise TimeoutError(f"Request to {url} timed out")
+
+            if response.status_code == 401:
+                raise PermissionError("Authentication failed. Check your API key.")
+            elif response.status_code == 402:
+                raise ValueError("Insufficient balance. Please add funds to your wallet.")
+            elif response.status_code >= 400:
+                # Try to parse JSON error, fallback to text
+                try:
+                    error_data = response.json()
+                    error_msg = error_data.get("error", error_data.get("message", response.text))
+                except:
+                    error_msg = response.text[:500] if response.text else f"HTTP {response.status_code}"
+                raise RuntimeError(f"Submission failed ({response.status_code}): {error_msg}")
+
+            # Parse successful response
+            try:
+                return response.json()
+            except:
+                raise RuntimeError(f"Invalid JSON response from {url}: {response.text[:200]}")
+
+    def status(self, execution_id: str) -> Dict[str, Any]:
+        """
+        Get execution status.
+
+        Args:
+            execution_id: The execution ID
+
+        Returns:
+            Dict with status, queue_position, etc.
+        """
+        url = f"{self.api_url}/api/v1/executions/{execution_id}/status"
+
+        with httpx.Client(timeout=self.timeout) as client:
+            try:
+                response = client.get(url, headers=self._get_headers())
+            except httpx.ConnectError as e:
+                raise ConnectionError(f"Failed to connect to {url}: {e}")
+
+            if response.status_code >= 400:
+                try:
+                    error_msg = response.json().get("error", response.text)
+                except:
+                    error_msg = response.text[:500] if response.text else f"HTTP {response.status_code}"
+                raise RuntimeError(f"Failed to get status: {error_msg}")
+
+            try:
+                return response.json()
+            except:
+                raise RuntimeError(f"Invalid JSON response: {response.text[:200]}")
+
+    def result(self, execution_id: str, timeout: float = None) -> GPUExecutionResult:
+        """
+        Get execution result (blocks until complete).
+
+        Args:
+            execution_id: The execution ID
+            timeout: Max time to wait in seconds
+
+        Returns:
+            GPUExecutionResult object
+        """
+        url = f"{self.api_url}/api/v1/executions/{execution_id}/result"
+        request_timeout = timeout or 3600
+
+        with httpx.Client(timeout=request_timeout) as client:
+            response = client.get(url, headers=self._get_headers())
+
+            if response.status_code >= 400:
+                raise RuntimeError(f"Failed to get result: {response.text}")
+
+            return GPUExecutionResult.from_dict(response.json())
+
+    def stream_output(self, execution_id: str, show_output: bool = True) -> Iterator[Dict[str, Any]]:
+        """
+        Stream real-time output from a running execution.
+
+        This connects to the Server-Sent Events (SSE) endpoint and yields
+        output events as they arrive. Ideal for Jupyter notebooks.
+
+        Args:
+            execution_id: The execution ID to stream
+            show_output: Whether to print output to stdout/stderr (default: True)
+
+        Yields:
+            Dict with event data:
+                - type: "stdout", "stderr", "output", or "status"
+                - content: The output content (for output events)
+                - status: Execution status (for status events)
+                - timestamp: Event timestamp
+                - execution_id: The execution ID
+
+        Example:
+            for event in runner.stream_output("exec_123"):
+                if event.get("status") == "completed":
+                    print(f"Done! Cost: {event.get('cost_cents', 0)/100:.2f}")
+        """
+        streamer = OutputStreamer(
+            execution_id=execution_id,
+            api_url=self.api_url,
+            headers=self._get_headers(),
+            timeout=self.timeout
+        )
+        yield from streamer.stream(show_output=show_output)
+
+    def run(
+        self,
+        code: str,
+        gpu: str = "L40S",
+        count: int = 1,
+        fraction: int = 100,
+        timeout: int = 3600,
+        env: str = None,
+        pip: Union[str, list] = None,
+        uv: Union[str, list] = None,
+        conda: Union[str, list] = None,
+        requirements: str = None,
+        use_uv: bool = False,
+        wait: bool = True,
+        poll_interval: float = 2.0,
+        max_retries: int = 3,
+        provisioning_timeout: int = 300,
+        cancel_on_timeout: bool = True,
+        stream: bool = None
+    ) -> GPUExecutionResult:
+        """
+        Execute code on GPU and optionally wait for completion.
+
+        Args:
+            code: Python code to execute
+            gpu: GPU type (L40S, A100-40G, A100-80G, H100, A10, T4)
+            count: Number of GPUs (1-8)
+            fraction: GPU time-slice percentage (25, 50, 75, 100)
+            timeout: Maximum execution time in seconds
+            env: Environment preset (ml, nlp, cv, audio, tabular, rl, scientific)
+            pip: Pip packages to install (string "pkg1,pkg2" or list ["pkg1", "pkg2"])
+            uv: UV packages to install (faster than pip)
+            conda: Conda packages to install (string or list)
+            requirements: Path to requirements.txt file
+            use_uv: Use uv instead of pip for installation (faster)
+            wait: Whether to wait for completion
+            poll_interval: Seconds between status checks when waiting
+            max_retries: Max retries for transient network errors (default: 3)
+            provisioning_timeout: Max time to wait for GPU provisioning (default: 300s)
+            cancel_on_timeout: Whether to cancel execution on timeout (default: True)
+            stream: Stream output in real-time (default: True in Jupyter, False otherwise)
+
+        Returns:
+            GPUExecutionResult object
+
+        Raises:
+            PodstackTimeoutError: If execution times out
+            PodstackProvisioningError: If GPU provisioning fails or times out
+            PodstackExecutionError: If execution fails
+            ConnectionError: If unable to connect to API
+            ValueError: If parameters are invalid
+        """
+        # Submit the code
+        submission = self.submit(code, gpu, count, fraction, timeout, env, pip, uv, conda, requirements, use_uv)
+        execution_id = submission.get("execution_id")
+
+        if not execution_id:
+            raise RuntimeError(f"No execution_id in response: {submission}")
+
+        print(f"[Podstack] Execution submitted: {execution_id}")
+
+        if not wait:
+            return GPUExecutionResult(
+                execution_id=execution_id,
+                status="submitted",
+                gpu_type=gpu,
+                gpu_count=count
+            )
+
+        # Determine if we should stream output
+        should_stream = stream if stream is not None else is_jupyter()
+
+        if should_stream:
+            return self._run_with_streaming(execution_id, gpu, count, timeout, max_retries, cancel_on_timeout)
+        else:
+            return self._run_with_polling(execution_id, gpu, count, timeout, poll_interval, max_retries, provisioning_timeout, cancel_on_timeout)
+
+    def _run_with_streaming(
+        self,
+        execution_id: str,
+        gpu: str,
+        count: int,
+        timeout: int,
+        max_retries: int,
+        cancel_on_timeout: bool
+    ) -> GPUExecutionResult:
+        """Run execution with real-time output streaming."""
+        print(f"[Podstack] Waiting for GPU runner ({gpu} x{count})...")
+
+        start_time = time.time()
+        output_buffer = []
+        final_status = "unknown"
+        final_event = {}
+
+        try:
+            for event in self.stream_output(execution_id, show_output=True):
+                elapsed = time.time() - start_time
+                if elapsed > timeout:
+                    if cancel_on_timeout:
+                        try:
+                            self.cancel(execution_id)
+                            print(f"\n[Podstack] Execution cancelled due to timeout")
+                        except Exception as e:
+                            logger.warning(f"Failed to cancel execution: {e}")
+
+                    raise PodstackTimeoutError(
+                        execution_id=execution_id,
+                        timeout=timeout,
+                        last_status=final_status,
+                        message=(
+                            f"Execution {execution_id} timed out after {timeout}s.\n"
+                            f"Last status: {final_status}\n"
+                            f"Suggestions:\n"
+                            f"  - Increase timeout parameter (current: {timeout}s)\n"
+                            f"  - Check execution logs: runner.result('{execution_id}')"
+                        )
+                    )
+
+                # Track output
+                if event.get("type") in ("stdout", "stderr", "output"):
+                    content = event.get("content", "")
+                    if content:
+                        output_buffer.append(content)
+
+                # Track status
+                if "status" in event:
+                    new_status = event["status"]
+                    if new_status != final_status:
+                        final_status = new_status
+                        if final_status == "provisioning":
+                            print(f"\n[Podstack] Provisioning GPU runner...")
+                        elif final_status == "running":
+                            print(f"\n[Podstack] Running on GPU...")
+
+                    # Check for terminal status
+                    if final_status in ("completed", "failed", "timeout", "cancelled"):
+                        final_event = event
+                        break
+
+        except (ConnectionError, httpx.ConnectError) as e:
+            # Try to recover and get the result
+            logger.warning(f"Stream connection lost: {e}")
+            print(f"\n[Podstack] Stream connection lost, fetching final result...")
+
+        # Get final result
+        result = None
+        for attempt in range(max_retries):
+            try:
+                result = self.result(execution_id)
+                break
+            except (ConnectionError, httpx.ConnectError, httpx.TimeoutException) as e:
+                if attempt == max_retries - 1:
+                    raise ConnectionError(f"Failed to get result after {max_retries} attempts: {e}")
+                time.sleep(2 * (attempt + 1))
+
+        if result.success:
+            print(f"\n[Podstack] Completed in {result.gpu_seconds:.1f}s (cost: ₹{result.cost_paise/100:.2f})")
+        else:
+            error_msg = result.error or 'Unknown error'
+            print(f"\n[Podstack] Failed: {error_msg}")
+
+        return result
+
+    def _run_with_polling(
+        self,
+        execution_id: str,
+        gpu: str,
+        count: int,
+        timeout: int,
+        poll_interval: float,
+        max_retries: int,
+        provisioning_timeout: int,
+        cancel_on_timeout: bool
+    ) -> GPUExecutionResult:
+        """Run execution with polling (non-streaming mode)."""
+        print(f"[Podstack] Waiting for GPU runner ({gpu} x{count})...")
+        start_time = time.time()
+        provisioning_start = None
+        last_status = ""
+        consecutive_errors = 0
+        last_error = None
+
+        while True:
+            elapsed = time.time() - start_time
+
+            # Check for overall timeout
+            if elapsed > timeout:
+                if cancel_on_timeout:
+                    try:
+                        self.cancel(execution_id)
+                        print(f"[Podstack] Execution cancelled due to timeout")
+                    except Exception as e:
+                        logger.warning(f"Failed to cancel execution: {e}")
+
+                raise PodstackTimeoutError(
+                    execution_id=execution_id,
+                    timeout=timeout,
+                    last_status=last_status,
+                    message=(
+                        f"Execution {execution_id} timed out after {timeout}s.\n"
+                        f"Last status: {last_status}\n"
+                        f"Suggestions:\n"
+                        f"  - Increase timeout parameter (current: {timeout}s)\n"
+                        f"  - Check execution logs: runner.result('{execution_id}')\n"
+                        f"  - Cancel and retry: runner.cancel('{execution_id}')"
+                    )
+                )
+
+            # Try to get status with retry on transient errors
+            try:
+                status_data = self.status(execution_id)
+                consecutive_errors = 0  # Reset on success
+            except (ConnectionError, httpx.ConnectError, httpx.TimeoutException) as e:
+                consecutive_errors += 1
+                last_error = str(e)
+                if consecutive_errors >= max_retries:
+                    raise ConnectionError(
+                        f"Lost connection to API after {consecutive_errors} retries. "
+                        f"Last error: {last_error}\n"
+                        f"Execution ID: {execution_id}"
+                    )
+                logger.warning(f"Transient error (attempt {consecutive_errors}/{max_retries}): {e}")
+                time.sleep(poll_interval * consecutive_errors)  # Exponential backoff
+                continue
+            except Exception as e:
+                # For other errors, fail immediately
+                raise RuntimeError(f"Failed to get status for {execution_id}: {e}")
+
+            current_status = status_data.get("status", "unknown")
+
+            if current_status != last_status:
+                last_status = current_status
+                if current_status == "pending":
+                    print(f"[Podstack] Pending...")
+                elif current_status == "queued":
+                    pos = status_data.get("queue_position", "?")
+                    print(f"[Podstack] Queued (position: {pos})")
+                elif current_status == "provisioning":
+                    provisioning_start = time.time()
+                    print(f"[Podstack] Provisioning GPU runner...")
+                elif current_status == "running":
+                    print(f"[Podstack] Running on GPU...")
+                elif current_status == "streaming":
+                    print(f"[Podstack] Streaming output...")
+                elif current_status in ("completed", "failed", "timeout", "cancelled"):
+                    pass  # Terminal states handled below
+                else:
+                    # Unknown status - log but continue
+                    logger.warning(f"Unknown status: {current_status}")
+                    print(f"[Podstack] Status: {current_status}")
+
+            # Check for provisioning timeout
+            if provisioning_start and current_status == "provisioning":
+                provisioning_elapsed = time.time() - provisioning_start
+                if provisioning_elapsed > provisioning_timeout:
+                    if cancel_on_timeout:
+                        try:
+                            self.cancel(execution_id)
+                        except Exception:
+                            pass
+                    raise PodstackProvisioningError(
+                        execution_id=execution_id,
+                        message=(
+                            f"GPU provisioning timed out after {provisioning_timeout}s for {execution_id}.\n"
+                            f"This may indicate:\n"
+                            f"  - No available {gpu} GPUs at this time\n"
+                            f"  - Cluster resource constraints\n"
+                            f"  - Network issues in the cluster\n"
+                            f"Suggestions:\n"
+                            f"  - Try a different GPU type\n"
+                            f"  - Retry after a few minutes\n"
+                            f"  - Contact support if the issue persists"
+                        )
+                    )
+
+            # Check for terminal states
+            if current_status in ("completed", "failed", "timeout", "cancelled"):
+                break
+
+            time.sleep(poll_interval)
+
+        # Get final result with retry
+        result = None
+        for attempt in range(max_retries):
+            try:
+                result = self.result(execution_id)
+                break
+            except (ConnectionError, httpx.ConnectError, httpx.TimeoutException) as e:
+                if attempt == max_retries - 1:
+                    raise ConnectionError(f"Failed to get result after {max_retries} attempts: {e}")
+                time.sleep(poll_interval * (attempt + 1))
+
+        if result.success:
+            print(f"[Podstack] Completed in {result.gpu_seconds:.1f}s (cost: ₹{result.cost_paise/100:.2f})")
+        else:
+            error_msg = result.error or 'Unknown error'
+            print(f"[Podstack] Failed: {error_msg}")
+            # Include partial output in the error for debugging
+            if result.output:
+                print(f"[Podstack] Output (last 500 chars):\n{result.output[-500:]}")
+
+        return result
+
+    def cancel(self, execution_id: str) -> Dict[str, Any]:
+        """
+        Cancel a running execution.
+
+        Args:
+            execution_id: The execution ID to cancel
+
+        Returns:
+            Dict with cancellation status
+        """
+        url = f"{self.api_url}/api/v1/executions/{execution_id}/cancel"
+
+        with httpx.Client(timeout=self.timeout) as client:
+            try:
+                response = client.post(url, headers=self._get_headers())
+            except httpx.ConnectError as e:
+                raise ConnectionError(f"Failed to connect to {url}: {e}")
+            except httpx.TimeoutException:
+                raise TimeoutError(f"Request to {url} timed out")
+
+            if response.status_code >= 400:
+                try:
+                    error_msg = response.json().get("error", response.text)
+                except:
+                    error_msg = response.text[:500] if response.text else f"HTTP {response.status_code}"
+                raise RuntimeError(f"Failed to cancel execution: {error_msg}")
+
+            try:
+                return response.json()
+            except:
+                return {"status": "cancelled"}
+
+    def run_function(
+        self,
+        func: Callable,
+        gpu: str = "L40S",
+        count: int = 1,
+        fraction: int = 100,
+        timeout: int = 3600,
+        env: str = None,
+        pip: str = None,
+        **kwargs
+    ) -> GPUExecutionResult:
+        """
+        Execute a function on remote GPU.
+
+        Args:
+            func: Function to execute
+            gpu: GPU type
+            count: Number of GPUs
+            fraction: GPU fraction percentage
+            timeout: Max execution time
+            env: Environment preset
+            pip: Additional pip packages
+            **kwargs: Arguments to pass to the function
+
+        Returns:
+            GPUExecutionResult object
+        """
+        # Get function source code
+        try:
+            source = inspect.getsource(func)
+        except OSError:
+            raise ValueError("Cannot get source code for function. Define it in a file, not interactively.")
+
+        # Dedent if needed
+        source = textwrap.dedent(source)
+
+        # Build the execution code
+        func_name = func.__name__
+
+        # Serialize kwargs
+        kwargs_str = ", ".join(f"{k}={repr(v)}" for k, v in kwargs.items())
+
+        code = f"""
+{source}
+
+# Execute the function
+__result__ = {func_name}({kwargs_str})
+if __result__ is not None:
+    print(__result__)
+"""
+
+        return self.run(
+            code=code,
+            gpu=gpu,
+            count=count,
+            fraction=fraction,
+            timeout=timeout,
+            env=env,
+            pip=pip
+        )
+
+
+# Global runner instance
+_runner: Optional[GPURunner] = None
+
+
+def init(
+    api_key: str = None,
+    project_id: str = None,
+    api_url: str = None
+):
+    """
+    Initialize the global GPU runner.
+
+    Args:
+        api_key: API key for authentication
+        project_id: Project ID for billing
+        api_url: Notebook service URL
+    """
+    global _runner
+    _runner = GPURunner(
+        api_key=api_key,
+        project_id=project_id,
+        api_url=api_url
+    )
+
+
+def get_runner() -> GPURunner:
+    """Get the global GPU runner instance."""
+    global _runner
+    if _runner is None:
+        _runner = GPURunner()
+    return _runner
+
+
+def run(
+    code: str,
+    gpu: str = "L40S",
+    count: int = 1,
+    fraction: int = 100,
+    timeout: int = 3600,
+    env: str = None,
+    pip: Union[str, list] = None,
+    uv: Union[str, list] = None,
+    conda: Union[str, list] = None,
+    requirements: str = None,
+    use_uv: bool = False,
+    wait: bool = True,
+    stream: bool = None
+) -> GPUExecutionResult:
+    """
+    Execute code on remote GPU.
+
+    Args:
+        code: Python code to execute
+        gpu: GPU type (L40S, A100-40G, A100-80G, H100, A10, T4)
+        count: Number of GPUs (1-8)
+        fraction: GPU fraction percentage (25, 50, 75, 100)
+        timeout: Max execution time in seconds
+        env: Environment preset (ml, nlp, cv, audio, tabular, rl, scientific)
+        pip: Pip packages - string "pkg1,pkg2" or list ["pkg1", "pkg2"]
+        uv: UV packages (faster than pip) - string or list
+        conda: Conda packages - string or list
+        requirements: Path to requirements.txt file
+        use_uv: Use uv instead of pip for all installations (faster)
+        wait: Whether to wait for completion
+        stream: Stream output in real-time (default: True in Jupyter, False otherwise)
+
+    Returns:
+        GPUExecutionResult object
+
+    Examples:
+        # Install single package
+        podstack.run_on_gpu(code, pip="transformers")
+
+        # Install multiple packages
+        podstack.run_on_gpu(code, pip=["torch", "transformers", "datasets"])
+
+        # Install with uv (faster)
+        podstack.run_on_gpu(code, uv=["torch", "transformers"])
+
+        # Use uv for all pip installs (faster)
+        podstack.run_on_gpu(code, pip=["torch"], use_uv=True)
+
+        # Install from requirements.txt
+        podstack.run_on_gpu(code, requirements="requirements.txt")
+
+        # Install from requirements.txt using uv (faster)
+        podstack.run_on_gpu(code, requirements="requirements.txt", use_uv=True)
+
+        # Combine pip and conda
+        podstack.run_on_gpu(code, pip="transformers", conda="cudatoolkit=11.8")
+
+        # Force streaming in non-Jupyter environment
+        podstack.run(code, stream=True)
+    """
+    return get_runner().run(
+        code=code,
+        gpu=gpu,
+        count=count,
+        fraction=fraction,
+        timeout=timeout,
+        env=env,
+        pip=pip,
+        uv=uv,
+        conda=conda,
+        requirements=requirements,
+        use_uv=use_uv,
+        wait=wait,
+        stream=stream
+    )
+
+
+def stream_output(execution_id: str, show_output: bool = True) -> Iterator[Dict[str, Any]]:
+    """
+    Stream real-time output from a running execution.
+
+    This connects to the Server-Sent Events (SSE) endpoint and yields
+    output events as they arrive. Ideal for Jupyter notebooks.
+
+    Args:
+        execution_id: The execution ID to stream
+        show_output: Whether to print output to stdout/stderr (default: True)
+
+    Yields:
+        Dict with event data
+
+    Example:
+        for event in podstack.stream_output("exec_123"):
+            if event.get("status") == "completed":
+                print(f"Done!")
+    """
+    return get_runner().stream_output(execution_id, show_output=show_output)