podstack 1.2.0__py3-none-any.whl

podstack/annotations.py

@@ -0,0 +1,725 @@
+"""
+Podstack Annotations
+
+Decorators and context managers for GPU provisioning, experiment tracking,
+and model registry in Podstack.
+
+Usage:
+    import podstack
+
+    # GPU execution - actually provisions and runs on remote GPU
+    @podstack.gpu(type="L40S", fraction=100)
+    def train():
+        import torch
+        print(f"Running on: {torch.cuda.get_device_name(0)}")
+        return {"status": "done"}
+
+    result = train()  # Executes on remote GPU!
+
+    # Experiment tracking
+    @podstack.experiment(name="my-experiment")
+    @podstack.run(name="training-v1")
+    def training_loop():
+        podstack.registry.log_metrics({"loss": 0.5})
+
+    # Model registration
+    @podstack.model.register(name="my-model")
+    def save_model():
+        ...
+"""
+
+import functools
+import os
+import time
+import inspect
+import textwrap
+from typing import Optional, Dict, Any, Callable, Union
+from contextlib import contextmanager
+
+from . import registry
+from .gpu_runner import (
+    GPURunner,
+    GPUExecutionResult,
+    get_runner,
+    init as init_runner,
+    PodstackError,
+    PodstackTimeoutError,
+    PodstackExecutionError,
+    PodstackProvisioningError,
+)
+
+
+# Global state for configuration
+_current_gpu_config: Dict[str, Any] = {}
+_current_environment: str = "pytorch"
+_auto_shutdown_minutes: int = 60
+_remote_execution_enabled: bool = True
+
+
+def enable_remote_execution(enabled: bool = True):
+    """Enable or disable remote GPU execution for decorators."""
+    global _remote_execution_enabled
+    _remote_execution_enabled = enabled
+
+
+def is_remote_execution_enabled() -> bool:
+    """Check if remote execution is enabled."""
+    return _remote_execution_enabled
+
+
+class GPUConfig:
+    """
+    GPU configuration decorator and context manager.
+
+    When used as a decorator, executes the function on a remote GPU.
+    When used as a context manager, sets the GPU config for nested operations.
+    """
+
+    def __init__(
+        self,
+        type: str = "L40S",
+        count: int = 1,
+        allocation: int = 100,
+        fraction: int = None,
+        memory_gb: Optional[int] = None,
+        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,
+        remote: bool = None
+    ):
+        """
+        Configure GPU for execution.
+
+        Args:
+            type: GPU type (L40S, A100-40G, A100-80G, H100, A10, T4)
+            count: Number of GPUs (1-8)
+            allocation: Alias for fraction (for backward compatibility)
+            fraction: GPU time-slice percentage (25, 50, 75, 100)
+            memory_gb: Optional memory override (not used in remote execution)
+            timeout: Maximum 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)
+            remote: Whether to execute remotely (default: True if configured)
+        """
+        self.type = type
+        self.count = count
+        self.fraction = fraction if fraction is not None else allocation
+        self.memory_gb = memory_gb
+        self.timeout = timeout
+        self.env = env
+        self.pip = pip
+        self.uv = uv
+        self.conda = conda
+        self.requirements = requirements
+        self.use_uv = use_uv
+        self.remote = remote if remote is not None else _remote_execution_enabled
+
+        # Store in global config
+        global _current_gpu_config
+        _current_gpu_config = {
+            "type": type,
+            "count": count,
+            "fraction": self.fraction,
+            "allocation": self.fraction,
+            "memory_gb": memory_gb,
+            "timeout": timeout,
+            "env": env,
+            "pip": pip,
+            "uv": uv,
+            "conda": conda,
+            "requirements": requirements,
+            "use_uv": use_uv,
+        }
+
+    def __call__(self, func: Callable) -> Callable:
+        """Decorator usage - executes function on remote GPU."""
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            if not self.remote:
+                # Local execution - just run the function
+                print(f"[Podstack] GPU Config (local): {self.type} x{self.count} @ {self.fraction}%")
+                return func(*args, **kwargs)
+
+            # Remote execution on GPU
+            print(f"[Podstack] Provisioning GPU: {self.type} x{self.count} @ {self.fraction}%")
+
+            try:
+                runner = get_runner()
+            except ValueError as e:
+                print(f"[Podstack] Warning: {e}")
+                print(f"[Podstack] Falling back to local execution")
+                return func(*args, **kwargs)
+
+            # Get function source
+            try:
+                source = inspect.getsource(func)
+                source = textwrap.dedent(source)
+            except OSError:
+                print("[Podstack] Warning: Cannot get function source, running locally")
+                return func(*args, **kwargs)
+
+            # Remove the decorator from source (to avoid recursion)
+            # Handle multi-line decorators by tracking parentheses balance
+            lines = source.split('\n')
+            clean_lines = []
+            skip_decorator = False
+            paren_depth = 0
+            for line in lines:
+                stripped = line.strip()
+                if stripped.startswith('@podstack.gpu') or stripped.startswith('@gpu'):
+                    skip_decorator = True
+                    # Count parentheses in this line
+                    paren_depth += line.count('(') - line.count(')')
+                    continue
+                if skip_decorator:
+                    # Continue skipping until parentheses are balanced
+                    paren_depth += line.count('(') - line.count(')')
+                    if paren_depth <= 0:
+                        skip_decorator = False
+                    continue
+                # Skip other decorators too
+                if stripped.startswith('@'):
+                    continue
+                clean_lines.append(line)
+            source = '\n'.join(clean_lines)
+
+            # Build kwargs string
+            kwargs_parts = []
+            for k, v in kwargs.items():
+                kwargs_parts.append(f"{k}={repr(v)}")
+            for i, arg in enumerate(args):
+                kwargs_parts.append(repr(arg))
+            kwargs_str = ", ".join(kwargs_parts)
+
+            # Build execution code
+            func_name = func.__name__
+            code = f"""{source}
+
+# Execute the function
+__podstack_result__ = {func_name}({kwargs_str})
+if __podstack_result__ is not None:
+    print("__PODSTACK_RESULT__")
+    print(repr(__podstack_result__))
+"""
+
+            # Execute on remote GPU with improved error handling
+            # Auto-enable streaming in Jupyter notebooks for real-time output
+            try:
+                result = runner.run(
+                    code=code,
+                    gpu=self.type,
+                    count=self.count,
+                    fraction=self.fraction,
+                    timeout=self.timeout,
+                    env=self.env,
+                    pip=self.pip,
+                    uv=self.uv,
+                    conda=self.conda,
+                    requirements=self.requirements,
+                    use_uv=self.use_uv,
+                    wait=True,
+                    stream=None  # Auto-detect: True in Jupyter, False otherwise
+                )
+            except PodstackTimeoutError as e:
+                print(f"[Podstack] ⚠️ Timeout Error:")
+                print(f"  Execution ID: {e.execution_id}")
+                print(f"  Timeout: {e.timeout}s")
+                print(f"  Last status: {e.last_status}")
+                print(f"\n  To debug, check the execution result:")
+                print(f"    from podstack import get_runner")
+                print(f"    result = get_runner().result('{e.execution_id}')")
+                raise
+            except PodstackProvisioningError as e:
+                print(f"[Podstack] ⚠️ Provisioning Error:")
+                print(f"  {e.message}")
+                raise
+            except PodstackExecutionError as e:
+                print(f"[Podstack] ⚠️ Execution Error:")
+                print(f"  Execution ID: {e.execution_id}")
+                print(f"  Error: {e.error}")
+                if e.output:
+                    print(f"\n  Output (last 500 chars):\n{e.output[-500:]}")
+                raise
+
+            if not result.success:
+                error_details = result.error or 'Unknown error'
+                error_msg = f"GPU execution failed: {error_details}"
+
+                # Include helpful debugging info
+                if result.output:
+                    # Look for Python tracebacks in output
+                    if "Traceback" in result.output or "Error" in result.output:
+                        error_msg += f"\n\nExecution output (last 1000 chars):\n{result.output[-1000:]}"
+                    else:
+                        error_msg += f"\n\nOutput preview:\n{result.output[-500:]}"
+
+                error_msg += f"\n\nExecution ID: {result.execution_id}"
+                error_msg += f"\nTo get full output: get_runner().result('{result.execution_id}')"
+
+                raise PodstackExecutionError(
+                    execution_id=result.execution_id,
+                    error=error_details,
+                    output=result.output
+                )
+
+            # Parse result from output
+            output = result.output
+            if "__PODSTACK_RESULT__" in output:
+                result_part = output.split("__PODSTACK_RESULT__")[1].strip()
+                if result_part:
+                    lines = result_part.split('\n')
+                    if lines:
+                        try:
+                            return eval(lines[0])
+                        except Exception as e:
+                            # If eval fails, return the raw string
+                            print(f"[Podstack] Warning: Could not parse result ({e}), returning raw output")
+                            return lines[0]
+
+            # Print any other output
+            if output and "__PODSTACK_RESULT__" not in output:
+                print(output)
+
+            return None
+
+        return wrapper
+
+    def __enter__(self):
+        """Context manager entry."""
+        print(f"[Podstack] GPU Config: {self.type} x{self.count} @ {self.fraction}%")
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        pass
+
+
+def gpu(
+    type: str = "L40S",
+    count: int = 1,
+    allocation: int = 100,
+    fraction: int = None,
+    memory_gb: Optional[int] = None,
+    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,
+    remote: bool = None
+) -> GPUConfig:
+    """
+    Configure and provision GPU for execution.
+
+    Can be used as a decorator (executes on remote GPU) or context manager.
+
+    Examples:
+        # As decorator - runs on remote GPU!
+        @podstack.gpu(type="L40S", fraction=100)
+        def train():
+            import torch
+            print(f"GPU: {torch.cuda.get_device_name(0)}")
+            return {"trained": True}
+
+        result = train()  # Executes remotely
+
+        # With pip packages
+        @podstack.gpu(type="L40S", pip=["transformers", "datasets"])
+        def train_llm():
+            from transformers import AutoModel
+            ...
+
+        # With uv packages (faster)
+        @podstack.gpu(type="L40S", uv=["torch", "transformers"])
+        def train_fast():
+            ...
+
+        # Use uv for all pip installs (faster)
+        @podstack.gpu(type="L40S", pip=["torch"], use_uv=True)
+        def train_uv():
+            ...
+
+        # With requirements.txt
+        @podstack.gpu(type="L40S", requirements="requirements.txt")
+        def train_with_deps():
+            ...
+
+        # With requirements.txt using uv (faster)
+        @podstack.gpu(type="L40S", requirements="requirements.txt", use_uv=True)
+        def train_fast_deps():
+            ...
+
+        # With conda packages
+        @podstack.gpu(type="L40S", conda="cudatoolkit=11.8")
+        def train_cuda():
+            ...
+
+        # As context manager (sets config only)
+        with podstack.gpu(type="A100-80G", count=2):
+            # GPU config is set for tracking
+            ...
+
+    Args:
+        type: GPU type (L40S, A100-40G, A100-80G, H100, A10, T4)
+        count: Number of GPUs (1-8, default: 1)
+        allocation: Alias for fraction (backward compatibility)
+        fraction: GPU time-slice percentage (25, 50, 75, 100, default: 100)
+        memory_gb: Optional memory limit in GB
+        timeout: Max execution time in seconds (default: 3600)
+        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)
+        remote: Execute remotely (default: True if configured)
+    """
+    return GPUConfig(
+        type=type,
+        count=count,
+        allocation=allocation,
+        fraction=fraction,
+        memory_gb=memory_gb,
+        timeout=timeout,
+        env=env,
+        pip=pip,
+        uv=uv,
+        conda=conda,
+        requirements=requirements,
+        use_uv=use_uv,
+        remote=remote
+    )
+
+
+def environment(env: str = "pytorch") -> str:
+    """
+    Set the runtime environment.
+
+    Args:
+        env: Environment name (pytorch, tensorflow, jax, huggingface, rapids, custom)
+
+    Returns:
+        The environment name
+    """
+    global _current_environment
+    _current_environment = env
+    print(f"[Podstack] Environment: {env}")
+    return env
+
+
+def auto_shutdown(minutes: int = 60) -> int:
+    """
+    Configure auto-shutdown timer.
+
+    Args:
+        minutes: Minutes of idle time before auto-shutdown
+
+    Returns:
+        The configured minutes
+    """
+    global _auto_shutdown_minutes
+    _auto_shutdown_minutes = minutes
+    print(f"[Podstack] Auto-shutdown: {minutes} minutes")
+    return minutes
+
+
+class ExperimentDecorator:
+    """Experiment tracking decorator."""
+
+    def __init__(self, name: str, description: str = None):
+        """
+        Create or set an experiment.
+
+        Args:
+            name: Experiment name
+            description: Optional description
+        """
+        self.name = name
+        self.description = description
+        self._experiment = None
+
+    def __call__(self, func: Callable) -> Callable:
+        """Decorator usage."""
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            self._experiment = registry.set_experiment(self.name, self.description)
+            print(f"[Podstack] Experiment: {self.name} (ID: {self._experiment.id})")
+            return func(*args, **kwargs)
+        return wrapper
+
+    def __enter__(self):
+        """Context manager entry."""
+        self._experiment = registry.set_experiment(self.name, self.description)
+        print(f"[Podstack] Experiment: {self.name} (ID: {self._experiment.id})")
+        return self._experiment
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        pass
+
+
+def experiment(name: str, description: str = None) -> ExperimentDecorator:
+    """
+    Create or set an experiment for tracking.
+
+    Can be used as decorator or context manager.
+
+    Examples:
+        @podstack.experiment(name="my-experiment")
+        def train():
+            ...
+
+        with podstack.experiment(name="my-experiment") as exp:
+            print(exp.id)
+
+    Args:
+        name: Experiment name
+        description: Optional description
+
+    Returns:
+        ExperimentDecorator instance
+    """
+    return ExperimentDecorator(name=name, description=description)
+
+
+class RunDecorator:
+    """Run tracking decorator with automatic metrics logging."""
+
+    def __init__(
+        self,
+        name: str = None,
+        track_gpu: bool = True,
+        track_time: bool = True,
+        tags: Dict[str, str] = None
+    ):
+        """
+        Start a tracked run.
+
+        Args:
+            name: Run name
+            track_gpu: Auto-track GPU metrics
+            track_time: Auto-track execution time
+            tags: Optional tags
+        """
+        self.name = name
+        self.track_gpu = track_gpu
+        self.track_time = track_time
+        self.tags = tags or {}
+        self._run = None
+        self._start_time = None
+
+    def __call__(self, func: Callable) -> Callable:
+        """Decorator usage."""
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            with self:
+                return func(*args, **kwargs)
+        return wrapper
+
+    def __enter__(self):
+        """Context manager entry."""
+        self._start_time = time.time()
+
+        # Add GPU config to tags if available
+        if _current_gpu_config and self.track_gpu:
+            self.tags["gpu_type"] = _current_gpu_config.get("type", "unknown")
+            self.tags["gpu_fraction"] = str(_current_gpu_config.get("fraction", 100))
+
+        self._run = registry.start_run(name=self.name, tags=self.tags)
+        print(f"[Podstack] Run started: {self.name} (ID: {self._run.id})")
+
+        # Log GPU config as params
+        if _current_gpu_config and self.track_gpu:
+            registry.log_params({
+                "gpu_type": _current_gpu_config.get("type"),
+                "gpu_count": _current_gpu_config.get("count", 1),
+                "gpu_fraction_percent": _current_gpu_config.get("fraction", 100),
+            })
+
+        return self._run
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        if self.track_time and self._start_time:
+            elapsed = time.time() - self._start_time
+            registry.log_metrics({"total_execution_time_seconds": elapsed})
+
+        if exc_type is not None:
+            registry.end_run(status="failed")
+            print(f"[Podstack] Run failed: {self.name}")
+        else:
+            registry.end_run(status="completed")
+            print(f"[Podstack] Run completed: {self.name}")
+
+
+def run(
+    name: str = None,
+    track_gpu: bool = True,
+    track_time: bool = True,
+    tags: Dict[str, str] = None
+) -> RunDecorator:
+    """
+    Start a tracked run.
+
+    Can be used as decorator or context manager.
+
+    Examples:
+        @podstack.run(name="training-v1", track_gpu=True)
+        def train():
+            ...
+
+        with podstack.run(name="training-v1") as r:
+            podstack.registry.log_metrics({"loss": 0.5})
+
+    Args:
+        name: Run name
+        track_gpu: Auto-track GPU metrics (default: True)
+        track_time: Auto-track execution time (default: True)
+        tags: Optional tags dict
+
+    Returns:
+        RunDecorator instance
+    """
+    return RunDecorator(name=name, track_gpu=track_gpu, track_time=track_time, tags=tags)
+
+
+class ModelRegistry:
+    """Model registry operations."""
+
+    @staticmethod
+    def register(
+        name: str,
+        run_id: str = None,
+        description: str = None,
+        tags: Dict[str, str] = None
+    ) -> Callable:
+        """
+        Decorator to register a model after function execution.
+
+        Examples:
+            @podstack.model.register(name="my-model")
+            def save_model():
+                torch.save(model, "model.pt")
+                podstack.registry.log_artifact("model.pt", "model")
+
+        Args:
+            name: Model name
+            run_id: Optional run ID (uses current run if not specified)
+            description: Optional description
+            tags: Optional tags
+        """
+        def decorator(func: Callable) -> Callable:
+            @functools.wraps(func)
+            def wrapper(*args, **kwargs):
+                result = func(*args, **kwargs)
+
+                # Get run_id from current run if not specified
+                actual_run_id = run_id
+                if actual_run_id is None:
+                    client = registry._get_client()
+                    if client._active_run:
+                        actual_run_id = client._active_run.id
+
+                if actual_run_id:
+                    model = registry.register_model(
+                        name=name,
+                        run_id=actual_run_id,
+                        description=description,
+                        tags=tags
+                    )
+                    print(f"[Podstack] Model registered: {name} (ID: {model.id})")
+                else:
+                    print(f"[Podstack] Warning: No active run, model not registered")
+
+                return result
+            return wrapper
+        return decorator
+
+    @staticmethod
+    def promote(
+        name: str = None,
+        version: int = 1,
+        stage: str = "staging",
+        comment: str = None
+    ) -> Callable:
+        """
+        Decorator to promote a model after function execution.
+
+        Examples:
+            @podstack.model.promote(name="my-model", version=1, stage="production")
+            def validate_model():
+                # validation code
+                pass
+
+        Args:
+            name: Model name
+            version: Version number to promote
+            stage: Target stage (development, staging, production, archived)
+            comment: Optional comment
+        """
+        def decorator(func: Callable) -> Callable:
+            @functools.wraps(func)
+            def wrapper(*args, **kwargs):
+                result = func(*args, **kwargs)
+
+                if name:
+                    version_obj = registry.set_model_stage(
+                        model_name=name,
+                        version=version,
+                        stage=stage,
+                        comment=comment
+                    )
+                    print(f"[Podstack] Model promoted: {name} v{version} -> {stage}")
+
+                return result
+            return wrapper
+        return decorator
+
+
+# Singleton instance
+model = ModelRegistry()
+
+
+def get_gpu_config() -> Dict[str, Any]:
+    """Get current GPU configuration."""
+    return _current_gpu_config.copy()
+
+
+def get_environment() -> str:
+    """Get current environment."""
+    return _current_environment
+
+
+def get_auto_shutdown_minutes() -> int:
+    """Get auto-shutdown configuration."""
+    return _auto_shutdown_minutes
+
+
+__all__ = [
+    "gpu",
+    "environment",
+    "auto_shutdown",
+    "experiment",
+    "run",
+    "model",
+    "get_gpu_config",
+    "get_environment",
+    "get_auto_shutdown_minutes",
+    "enable_remote_execution",
+    "is_remote_execution_enabled",
+    "GPUConfig",
+    # Exceptions
+    "PodstackError",
+    "PodstackTimeoutError",
+    "PodstackExecutionError",
+    "PodstackProvisioningError",
+]