podstack 1.2.0__py3-none-any.whl

podstack_gpu/app.py

@@ -0,0 +1,675 @@
+"""Podstack App - Application class with function decorators."""
+
+from __future__ import annotations
+import os
+import inspect
+import textwrap
+from typing import Optional, List, Dict, Any, Callable, Union, Generator, TYPE_CHECKING
+from functools import wraps
+
+from .exceptions import PodstackError, ValidationError
+
+if TYPE_CHECKING:
+    from .image import Image
+    from .volume import Volume
+    from .secret import Secret
+
+
+# Valid GPU types
+VALID_GPU_TYPES = ["A10", "L40", "L40S", "A100-40G", "A100-80G", "H100"]
+VALID_FRACTIONS = [25, 50, 75, 100]
+
+
+class Function:
+    """
+    A GPU-accelerated function that can be invoked remotely.
+
+    Example:
+        @app.function(gpu="H100")
+        def train(epochs: int):
+            import torch
+            ...
+
+        # Call remotely
+        result = train.remote(epochs=10)
+
+        # Call locally
+        result = train.local(epochs=10)
+    """
+
+    def __init__(
+        self,
+        func: Callable,
+        app: "App",
+        gpu: str = "L40S",
+        count: int = 1,
+        fraction: int = 100,
+        timeout: int = 3600,
+        memory: Optional[int] = None,
+        image: Optional["Image"] = None,
+        volumes: Optional[Dict[str, "Volume"]] = None,
+        secrets: Optional[List["Secret"]] = None,
+        retries: int = 0,
+        concurrency_limit: Optional[int] = None,
+    ):
+        self._func = func
+        self._app = app
+        self._gpu = gpu.upper()
+        self._count = count
+        self._fraction = fraction
+        self._timeout = timeout
+        self._memory = memory
+        self._image = image
+        self._volumes = volumes or {}
+        self._secrets = secrets or []
+        self._retries = retries
+        self._concurrency_limit = concurrency_limit
+
+        # Validate
+        self._validate()
+
+        # Preserve function metadata
+        wraps(func)(self)
+
+    def _validate(self):
+        """Validate function configuration."""
+        if self._gpu not in VALID_GPU_TYPES:
+            raise ValidationError(
+                f"Invalid GPU type: {self._gpu}. "
+                f"Valid types: {', '.join(VALID_GPU_TYPES)}"
+            )
+
+        if self._fraction not in VALID_FRACTIONS:
+            raise ValidationError(
+                f"Invalid GPU fraction: {self._fraction}. "
+                f"Valid fractions: {VALID_FRACTIONS}"
+            )
+
+        if self._count < 1 or self._count > 8:
+            raise ValidationError("GPU count must be between 1 and 8")
+
+        if self._timeout < 60:
+            raise ValidationError("Timeout must be at least 60 seconds")
+
+        if self._timeout > 86400:
+            raise ValidationError("Timeout cannot exceed 86400 seconds (24 hours)")
+
+    def __call__(self, *args, **kwargs):
+        """Call the function locally (default behavior)."""
+        return self.local(*args, **kwargs)
+
+    def local(self, *args, **kwargs):
+        """Execute the function locally without GPU."""
+        return self._func(*args, **kwargs)
+
+    def remote(self, *args, **kwargs) -> "FunctionCall":
+        """
+        Execute the function on a remote GPU.
+
+        Returns:
+            FunctionCall object to get results
+        """
+        return self._app._execute_remote(self, args, kwargs)
+
+    def spawn(self, *args, **kwargs) -> "FunctionCall":
+        """
+        Spawn the function on a remote GPU without waiting.
+
+        Returns:
+            FunctionCall object to get results later
+        """
+        return self._app._execute_remote(self, args, kwargs, wait=False)
+
+    def map(self, *iterables, order_outputs: bool = True, return_exceptions: bool = False):
+        """
+        Map the function over iterables in parallel.
+
+        Args:
+            *iterables: Input iterables to map over
+            order_outputs: If True, return results in input order
+            return_exceptions: If True, return exceptions instead of raising
+
+        Returns:
+            Generator of results
+        """
+        return self._app._execute_map(
+            self, iterables,
+            order_outputs=order_outputs,
+            return_exceptions=return_exceptions
+        )
+
+    def starmap(self, args_list: List[tuple], kwargs_list: List[dict] = None, **options):
+        """
+        Map the function over a list of argument tuples.
+
+        Args:
+            args_list: List of positional argument tuples
+            kwargs_list: List of keyword argument dicts (optional)
+        """
+        kwargs_list = kwargs_list or [{}] * len(args_list)
+        return self._app._execute_starmap(self, args_list, kwargs_list, **options)
+
+    def get_source(self) -> str:
+        """Get the source code of the function without decorators."""
+        source = inspect.getsource(self._func)
+        source = textwrap.dedent(source)
+
+        # Remove decorator lines (lines starting with @)
+        lines = source.split('\n')
+        cleaned_lines = []
+        in_decorator = False
+
+        for line in lines:
+            stripped = line.strip()
+            # Skip decorator lines
+            if stripped.startswith('@'):
+                in_decorator = True
+                # Check if decorator continues (ends with backslash or has unclosed parens)
+                if stripped.endswith('\\') or (stripped.count('(') > stripped.count(')')):
+                    continue
+                else:
+                    in_decorator = False
+                    continue
+            # Skip continuation of multi-line decorators
+            if in_decorator:
+                if stripped.endswith('\\') or (stripped.count('(') > stripped.count(')')):
+                    continue
+                else:
+                    in_decorator = False
+                    continue
+            cleaned_lines.append(line)
+
+        return '\n'.join(cleaned_lines)
+
+    def _build_annotation(self) -> str:
+        """Build the GPU annotation string."""
+        parts = [f"#@podstack gpu={self._gpu} count={self._count} fraction={self._fraction} timeout={self._timeout}"]
+
+        if self._memory:
+            parts.append(f"memory={self._memory}")
+
+        if self._image:
+            annotation_part = self._image.definition.to_annotation()
+            if annotation_part:
+                parts.append(annotation_part)
+
+        return " ".join(parts)
+
+    @property
+    def info(self) -> dict:
+        """Get function metadata."""
+        return {
+            "name": self._func.__name__,
+            "gpu": self._gpu,
+            "count": self._count,
+            "fraction": self._fraction,
+            "timeout": self._timeout,
+            "memory": self._memory,
+            "image": self._image.to_dict() if self._image else None,
+        }
+
+
+class FunctionCall:
+    """
+    Represents an in-progress or completed function call.
+
+    Example:
+        call = train.spawn(epochs=10)
+        # Do other work...
+        result = call.get()  # Wait for result
+        print(f"Duration: {call.gpu_seconds}s, Cost: ₹{call.cost_inr:.4f}")
+
+        # Or stream status updates in real-time
+        for update in call.stream_status():
+            print(f"Status: {update.status}")
+    """
+
+    def __init__(self, execution_id: str, app: "App"):
+        self._execution_id = execution_id
+        self._app = app
+        self._result = None
+        self._done = False
+        self._error = None
+        self._gpu_seconds = 0.0
+        self._cost_cents = 0
+        self._status = "submitted"
+        self._queue_position = None
+        self._estimated_wait = None
+
+    @property
+    def object_id(self) -> str:
+        """Get the execution ID."""
+        return self._execution_id
+
+    @property
+    def status(self) -> str:
+        """Get current execution status."""
+        return self._status
+
+    @property
+    def queue_position(self) -> Optional[int]:
+        """Get queue position (if queued)."""
+        return self._queue_position
+
+    @property
+    def estimated_wait_seconds(self) -> Optional[int]:
+        """Get estimated wait time in seconds (if queued)."""
+        return self._estimated_wait
+
+    @property
+    def gpu_seconds(self) -> float:
+        """Get GPU execution time in seconds."""
+        return self._gpu_seconds
+
+    @property
+    def cost_cents(self) -> int:
+        """Get execution cost in cents."""
+        return self._cost_cents
+
+    @property
+    def cost_inr(self) -> float:
+        """Get execution cost in INR (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
+
+    def get(self, timeout: float = None):
+        """
+        Wait for the function to complete and return the result.
+
+        Args:
+            timeout: Maximum time to wait in seconds
+
+        Returns:
+            The function's return value
+
+        Raises:
+            TimeoutError: If timeout is exceeded
+            ExecutionError: If the function failed
+        """
+        if self._done:
+            if self._error:
+                raise self._error
+            return self._result
+
+        result = self._app._wait_for_result(self._execution_id, timeout)
+        self._result = result
+        self._done = True
+        return result
+
+    def stream_status(self, callback: Callable = None):
+        """
+        Stream status updates in real-time.
+
+        Can be used as an iterator or with a callback.
+
+        Args:
+            callback: Optional callback function for each update
+
+        Yields:
+            StatusUpdate objects (if no callback provided)
+
+        Example:
+            # As iterator
+            for update in call.stream_status():
+                print(f"Status: {update.status}")
+
+            # With callback
+            call.stream_status(callback=lambda u: print(u.status))
+        """
+        runner = self._app._get_runner()
+
+        if callback:
+            def _callback(update):
+                self._update_from_status(update)
+                callback(update)
+            runner.stream_status(self._execution_id, _callback)
+        else:
+            # Return a generator
+            updates = []
+            def _collect(update):
+                self._update_from_status(update)
+                updates.append(update)
+
+            runner.stream_status(self._execution_id, _collect)
+            yield from updates
+
+    def _update_from_status(self, update):
+        """Update internal state from status update."""
+        from .runner import StatusUpdate
+        self._status = update.status
+        self._queue_position = update.queue_position
+        self._estimated_wait = update.estimated_wait_seconds
+        if update.gpu_seconds:
+            self._gpu_seconds = update.gpu_seconds
+        if update.cost_cents:
+            self._cost_cents = update.cost_cents
+        if update.is_terminal:
+            self._done = True
+            if update.error:
+                from .exceptions import ExecutionError
+                self._error = ExecutionError(update.error, self._execution_id, update.status)
+
+    def cancel(self, reason: str = "Cancelled by user") -> bool:
+        """Cancel the function call."""
+        return self._app._cancel_execution(self._execution_id, reason)
+
+
+class App:
+    """
+    Podstack Application - container for GPU functions.
+
+    Example:
+        import podstack
+
+        app = podstack.App("my-training-app")
+
+        @app.function(gpu="H100", image=podstack.Image.ml())
+        def train(epochs: int):
+            import torch
+            ...
+
+        if __name__ == "__main__":
+            result = train.remote(epochs=10)
+    """
+
+    # Default API URL for the Podstack Notebooks API
+    DEFAULT_API_URL = "https://cloud.podstack.ai/notebooks"
+
+    def __init__(
+        self,
+        name: str = None,
+        token: str = None,
+        project_id: str = None,
+        api_url: str = None,
+    ):
+        """
+        Initialize a Podstack App.
+
+        Args:
+            name: Application name (for organization)
+            token: API token (or set PODSTACK_TOKEN env var) - supports psk_xxx platform tokens
+            project_id: Project ID (or set PODSTACK_PROJECT_ID env var)
+            api_url: API URL (optional, defaults to PODSTACK_API_URL env var or production URL)
+        """
+        self.name = name or "default"
+        self._token = token or os.environ.get("PODSTACK_TOKEN")
+        self._project_id = project_id or os.environ.get("PODSTACK_PROJECT_ID")
+        self._api_url = (api_url or os.environ.get("PODSTACK_API_URL") or self.DEFAULT_API_URL).rstrip("/")
+
+        self._functions: Dict[str, Function] = {}
+        self._runner = None  # Lazy initialized
+
+    def _get_runner(self):
+        """Get or create the GPU runner."""
+        if self._runner is None:
+            from .runner import GPURunner
+
+            if not self._token:
+                raise PodstackError(
+                    "No API token provided. Set PODSTACK_TOKEN environment variable "
+                    "or pass token to App()"
+                )
+            if not self._project_id:
+                raise PodstackError(
+                    "No project ID provided. Set PODSTACK_PROJECT_ID environment variable "
+                    "or pass project_id to App()"
+                )
+
+            self._runner = GPURunner(
+                token=self._token,
+                project_id=self._project_id,
+                api_url=self._api_url,
+            )
+        return self._runner
+
+    def function(
+        self,
+        gpu: str = "L40S",
+        count: int = 1,
+        fraction: int = 100,
+        timeout: int = 3600,
+        memory: int = None,
+        image: "Image" = None,
+        volumes: Dict[str, "Volume"] = None,
+        secrets: List["Secret"] = None,
+        retries: int = 0,
+        concurrency_limit: int = None,
+    ) -> Callable[[Callable], Function]:
+        """
+        Decorator to create a GPU-accelerated function.
+
+        Args:
+            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) - use lower fractions for cost savings
+            timeout: Maximum execution time in seconds
+            memory: GPU memory limit in GB
+            image: Container image specification
+            volumes: Mount volumes {"/path": Volume}
+            secrets: List of secrets to inject
+            retries: Number of retries on failure
+            concurrency_limit: Max concurrent executions
+
+        Example:
+            # Full GPU for training
+            @app.function(gpu="H100", image=podstack.Image.ml())
+            def train(epochs: int):
+                import torch
+                ...
+
+            # Fractional GPU for inference (cost-effective)
+            @app.function(gpu="L40S", fraction=25, image=podstack.Image.ml())
+            def inference(data):
+                return model.predict(data)
+
+            # Multiple GPUs for distributed training
+            @app.function(gpu="A100-80G", count=4, image=podstack.Image.ml())
+            def distributed_train():
+                ...
+        """
+        def decorator(func: Callable) -> Function:
+            fn = Function(
+                func=func,
+                app=self,
+                gpu=gpu,
+                count=count,
+                fraction=fraction,
+                timeout=timeout,
+                memory=memory,
+                image=image,
+                volumes=volumes,
+                secrets=secrets,
+                retries=retries,
+                concurrency_limit=concurrency_limit,
+            )
+            self._functions[func.__name__] = fn
+            return fn
+
+        return decorator
+
+    def cls(
+        self,
+        gpu: str = "L40S",
+        count: int = 1,
+        fraction: int = 100,
+        timeout: int = 3600,
+        memory: int = None,
+        image: "Image" = None,
+        volumes: Dict[str, "Volume"] = None,
+        secrets: List["Secret"] = None,
+    ):
+        """
+        Decorator to create a GPU-accelerated class.
+
+        Methods decorated with @method will run on the GPU.
+
+        Example:
+            @app.cls(gpu="H100", image=podstack.Image.ml())
+            class ModelServer:
+                def __init__(self):
+                    self.model = load_model()
+
+                @podstack.method()
+                def predict(self, data):
+                    return self.model(data)
+        """
+        def decorator(cls):
+            # Store GPU config on the class
+            cls._podstack_config = {
+                "gpu": gpu,
+                "count": count,
+                "fraction": fraction,
+                "timeout": timeout,
+                "memory": memory,
+                "image": image,
+                "volumes": volumes,
+                "secrets": secrets,
+            }
+            cls._podstack_app = self
+            return cls
+
+        return decorator
+
+    def _execute_remote(self, fn: Function, args: tuple, kwargs: dict, wait: bool = True) -> FunctionCall:
+        """Execute a function remotely."""
+        runner = self._get_runner()
+
+        # Build the code to execute
+        source = fn.get_source()
+
+        # Build argument serialization
+        import json
+        args_json = json.dumps(args)
+        kwargs_json = json.dumps(kwargs)
+
+        code = f'''
+{fn._build_annotation()}
+{source}
+
+import json
+_args = json.loads({args_json!r})
+_kwargs = json.loads({kwargs_json!r})
+_result = {fn._func.__name__}(*_args, **_kwargs)
+print("__RESULT__:", json.dumps(_result))
+'''
+
+        # Submit
+        result = runner.run(
+            code,
+            gpu=fn._gpu,
+            count=fn._count,
+            fraction=fn._fraction,
+            timeout=fn._timeout,
+            memory=fn._memory,
+            env=fn._image.definition.env_preset if fn._image else None,
+            pip=fn._image.definition.pip_packages if fn._image else None,
+            wait=wait,
+        )
+
+        call = FunctionCall(result.execution_id, self)
+        call._gpu_seconds = result.gpu_seconds
+        call._cost_cents = result.cost_cents
+
+        if wait and result.success:
+            # Parse result from output
+            for line in result.output.split('\n'):
+                if line.startswith("__RESULT__:"):
+                    import json
+                    call._result = json.loads(line[11:].strip())
+                    call._done = True
+                    break
+
+            # If no __RESULT__ found, return the raw output
+            if not call._done:
+                call._result = result.output
+                call._done = True
+
+        return call
+
+    def _execute_map(self, fn: Function, iterables, order_outputs: bool, return_exceptions: bool):
+        """Execute map over iterables."""
+        # Simple implementation - submit all and collect
+        calls = []
+        for args in zip(*iterables):
+            call = self._execute_remote(fn, args, {}, wait=False)
+            calls.append(call)
+
+        # Collect results
+        results = []
+        for call in calls:
+            try:
+                results.append(call.get())
+            except Exception as e:
+                if return_exceptions:
+                    results.append(e)
+                else:
+                    raise
+
+        return results if order_outputs else iter(results)
+
+    def _execute_starmap(self, fn: Function, args_list, kwargs_list, **options):
+        """Execute starmap over argument lists."""
+        calls = []
+        for args, kwargs in zip(args_list, kwargs_list):
+            call = self._execute_remote(fn, args, kwargs, wait=False)
+            calls.append(call)
+
+        # Collect results
+        return [call.get() for call in calls]
+
+    def _wait_for_result(self, execution_id: str, timeout: float = None):
+        """Wait for an execution to complete."""
+        runner = self._get_runner()
+        import time
+
+        start = time.time()
+        while True:
+            if timeout and (time.time() - start) > timeout:
+                from .exceptions import TimeoutError
+                raise TimeoutError(f"Timed out waiting for {execution_id}")
+
+            status = runner.get_status(execution_id)
+            if status["status"] in ("completed", "failed", "timeout", "cancelled"):
+                break
+            time.sleep(2)
+
+        if status["status"] != "completed":
+            from .exceptions import ExecutionError
+            raise ExecutionError(
+                status.get("error", f"Execution {status['status']}"),
+                execution_id=execution_id,
+                status=status["status"],
+            )
+
+        return status.get("result")
+
+    def _cancel_execution(self, execution_id: str, reason: str) -> bool:
+        """Cancel an execution."""
+        runner = self._get_runner()
+        return runner.cancel(execution_id, reason)
+
+    @property
+    def registered_functions(self) -> Dict[str, Function]:
+        """Get all registered functions."""
+        return self._functions.copy()
+
+
+def method():
+    """
+    Decorator to mark a method as GPU-executable within a @app.cls class.
+
+    Example:
+        @app.cls(gpu="H100")
+        class Model:
+            @podstack.method()
+            def predict(self, x):
+                ...
+    """
+    def decorator(func):
+        func._is_podstack_method = True
+        return func
+    return decorator

podstack_gpu/exceptions.py

@@ -0,0 +1,35 @@
+"""Podstack SDK exceptions."""
+
+
+class PodstackError(Exception):
+    """Base exception for Podstack SDK."""
+    pass
+
+
+class AuthenticationError(PodstackError):
+    """Authentication failed - invalid or expired token."""
+    pass
+
+
+class ValidationError(PodstackError):
+    """Invalid parameters or annotation."""
+    pass
+
+
+class ExecutionError(PodstackError):
+    """GPU execution failed."""
+
+    def __init__(self, message: str, execution_id: str = None, status: str = None):
+        super().__init__(message)
+        self.execution_id = execution_id
+        self.status = status
+
+
+class TimeoutError(PodstackError):
+    """Execution timed out."""
+    pass
+
+
+class InsufficientBalanceError(PodstackError):
+    """Insufficient wallet balance."""
+    pass