podstack 1.2.0__py3-none-any.whl

podstack/__init__.py

@@ -0,0 +1,222 @@
+"""
+Podstack Python SDK
+
+High-performance GPU platform SDK for ML workloads with experiment tracking
+and model registry.
+
+Usage:
+    import podstack
+
+    # Initialize (required for GPU execution and registry)
+    podstack.init(
+        api_key="your-api-key",
+        project_id="your-project-id"
+    )
+
+    # GPU Execution - Actually 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 {"trained": True}
+
+    result = train()  # Executes on remote GPU
+
+    # Or run code directly
+    result = podstack.run_on_gpu('''
+        import torch
+        print(f"GPU: {torch.cuda.get_device_name(0)}")
+    ''', gpu="L40S")
+
+Registry (Experiment Tracking & Model Management):
+    from podstack import registry
+
+    registry.init(api_key="your-api-key", project_id="your-project-id")
+    registry.set_experiment("my-experiment")
+
+    with registry.start_run(name="training") as run:
+        registry.log_params({"lr": 0.001})
+        registry.log_metrics({"loss": 0.5}, step=1)
+
+    registry.register_model(name="my-model", run_id=run.id)
+
+Decorators:
+    import podstack
+
+    @podstack.gpu(type="L40S", fraction=100)
+    @podstack.experiment(name="my-experiment")
+    @podstack.run(name="training-v1", track_gpu=True)
+    def train():
+        ...
+
+    @podstack.model.register(name="my-model")
+    def save_model():
+        ...
+"""
+
+__version__ = "1.2.0"
+
+from .client import Client
+from .notebook import Notebook, NotebookStatus
+from .execution import Execution, ExecutionStatus
+from .exceptions import (
+    PodstackError,
+    AuthenticationError,
+    NotFoundError,
+    RateLimitError,
+    GPUNotAvailableError,
+    ExecutionTimeoutError,
+)
+from .models import (
+    GPUType,
+    Environment,
+    Project,
+    Version,
+    WalletBalance,
+)
+
+# Registry module import
+from . import registry
+
+# GPU Runner module import
+from . import gpu_runner
+from .gpu_runner import (
+    GPURunner,
+    GPUExecutionResult,
+    run as run_on_gpu,
+)
+
+# Annotations module import
+from . import annotations
+from .annotations import (
+    gpu,
+    environment,
+    auto_shutdown,
+    experiment,
+    run,
+    model,
+    get_gpu_config,
+    get_environment,
+    get_auto_shutdown_minutes,
+    enable_remote_execution,
+)
+
+
+def init(
+    api_key: str = None,
+    project_id: str = None,
+    api_url: str = None,
+    registry_url: str = None
+):
+    """
+    Initialize the Podstack SDK.
+
+    This sets up both GPU execution and registry services.
+
+    Args:
+        api_key: API key for authentication (or set PODSTACK_API_KEY env var)
+        project_id: Project ID (or set PODSTACK_PROJECT_ID env var)
+        api_url: Notebook service URL (optional)
+        registry_url: Registry service URL (optional)
+
+    Example:
+        import podstack
+
+        podstack.init(
+            api_key="your-api-key",
+            project_id="your-project-id"
+        )
+
+        # Now you can use GPU execution
+        @podstack.gpu(type="L40S")
+        def train():
+            ...
+    """
+    # Initialize GPU runner
+    gpu_runner.init(
+        api_key=api_key,
+        project_id=project_id,
+        api_url=api_url
+    )
+
+    # Initialize registry
+    registry.init(
+        api_key=api_key,
+        project_id=project_id,
+        api_url=registry_url
+    )
+
+
+def auto_init():
+    """
+    Auto-initialize the SDK from environment variables.
+
+    Checks for PODSTACK_API_KEY and PODSTACK_PROJECT_ID env vars.
+    Called automatically at import time when PODSTACK_AUTO_INIT=1 or
+    when running inside a Podstack notebook (PODSTACK_NOTEBOOK_ID is set).
+    """
+    import os
+    api_key = os.getenv("PODSTACK_API_KEY")
+    project_id = os.getenv("PODSTACK_PROJECT_ID")
+
+    if api_key and project_id:
+        init(
+            api_key=api_key,
+            project_id=project_id,
+            api_url=os.getenv("PODSTACK_API_URL"),
+            registry_url=os.getenv("PODSTACK_REGISTRY_URL"),
+        )
+
+
+# Auto-initialize when running in a Podstack notebook or when explicitly requested
+import os as _os
+if _os.getenv("PODSTACK_AUTO_INIT") == "1" or _os.getenv("PODSTACK_NOTEBOOK_ID"):
+    auto_init()
+del _os
+
+
+__all__ = [
+    # Version
+    "__version__",
+    # Initialization
+    "init",
+    "auto_init",
+    # Core
+    "Client",
+    "Notebook",
+    "NotebookStatus",
+    "Execution",
+    "ExecutionStatus",
+    # Exceptions
+    "PodstackError",
+    "AuthenticationError",
+    "NotFoundError",
+    "RateLimitError",
+    "GPUNotAvailableError",
+    "ExecutionTimeoutError",
+    # Models
+    "GPUType",
+    "Environment",
+    "Project",
+    "Version",
+    "WalletBalance",
+    # Registry
+    "registry",
+    # GPU Runner
+    "gpu_runner",
+    "GPURunner",
+    "GPUExecutionResult",
+    "run_on_gpu",
+    # Annotations
+    "annotations",
+    "gpu",
+    "environment",
+    "auto_shutdown",
+    "experiment",
+    "run",
+    "model",
+    "get_gpu_config",
+    "get_environment",
+    "get_auto_shutdown_minutes",
+    "enable_remote_execution",
+]