podstack 1.2.1

Official Python SDK for Podstack GPU Notebook Platform

Podstack Python SDK

Official Python SDK for the Podstack GPU Platform. Run ML workloads on remote GPUs with simple decorators, track experiments, and manage models.

Installation

pip install podstack

With optional dependencies:

pip install podstack[torch]        # PyTorch support
pip install podstack[huggingface]  # HuggingFace Transformers
pip install podstack[all]          # All ML frameworks

Quick Start

import podstack

# Initialize the SDK
podstack.init(
    api_key="your-api-key",
    project_id="your-project-id"
)

# Run a function on a remote GPU with a single decorator
@podstack.gpu(type="L40S", fraction=100)
def train():
    import torch
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    return {"status": "done"}

result = train()  # Executes on remote GPU!

Decorators & Annotations

Podstack provides decorators that turn any Python function into a remote GPU workload with built-in experiment tracking.

@podstack.gpu - Remote GPU Execution

import podstack

# Basic GPU execution
@podstack.gpu(type="L40S")
def train_model():
    import torch
    model = torch.nn.Linear(768, 10).cuda()
    return {"params": sum(p.numel() for p in model.parameters())}

result = train_model()

# Specify GPU type, count, and fraction
@podstack.gpu(type="A100-80G", count=2, fraction=100)
def train_large_model():
    import torch
    print(f"GPUs available: {torch.cuda.device_count()}")

# Install pip packages on the fly
@podstack.gpu(type="L40S", pip=["transformers", "datasets", "accelerate"])
def finetune_llm():
    from transformers import AutoModelForCausalLM, AutoTokenizer
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
    ...

# Use uv for faster package installation
@podstack.gpu(type="L40S", uv=["torch", "transformers"])
def fast_setup():
    ...

# Install from requirements.txt
@podstack.gpu(type="L40S", requirements="requirements.txt", use_uv=True)
def train_with_deps():
    ...

# Use conda packages
@podstack.gpu(type="L40S", conda="cudatoolkit=11.8")
def train_with_conda():
    ...

# Use a pre-built environment
@podstack.gpu(type="L40S", env="nlp")
def nlp_task():
    ...

# Set execution timeout (default: 3600s)
@podstack.gpu(type="L40S", timeout=7200)
def long_training():
    ...

# Disable remote execution (run locally for debugging)
@podstack.gpu(type="L40S", remote=False)
def debug_locally():
    print("This runs on your local machine")

# Use as a context manager
with podstack.gpu(type="A100-80G", count=2) as cfg:
    print(f"GPU config set: {cfg.type}")

Available GPU types: T4, L4, A10, L40S, A100-40G, A100-80G, H100

Available environments: ml, nlp, cv, audio, tabular, rl, scientific

@podstack.experiment - Experiment Tracking

import podstack

# As a decorator
@podstack.experiment(name="transformer-experiments")
def run_experiment():
    ...

# As a context manager
with podstack.experiment(name="transformer-experiments") as exp:
    print(f"Experiment ID: {exp.id}")

@podstack.run - Run Tracking

Automatically tracks execution time and GPU configuration.

import podstack

# As a decorator
@podstack.experiment(name="my-experiment")
@podstack.run(name="training-v1", track_gpu=True)
def train():
    podstack.registry.log_params({"lr": 0.001, "batch_size": 32})
    for epoch in range(10):
        loss = 1.0 / (epoch + 1)
        podstack.registry.log_metrics({"loss": loss}, step=epoch)

# As a context manager
with podstack.run(name="training-v1") as run:
    podstack.registry.log_params({"lr": 0.001})
    podstack.registry.log_metrics({"loss": 0.5}, step=1)
    print(f"Run ID: {run.id}")

# With tags
@podstack.run(name="ablation-study", tags={"variant": "no-dropout"})
def ablation():
    ...

@podstack.model - Model Registration

import podstack

# Register model after function completes
@podstack.experiment(name="my-experiment")
@podstack.run(name="training-v1")
@podstack.model.register(name="my-classifier")
def train_and_save():
    import torch
    model = torch.nn.Linear(768, 10)
    torch.save(model.state_dict(), "model.pt")
    podstack.registry.log_artifact("model.pt", "model")

# Promote model to production after validation
@podstack.model.promote(name="my-classifier", version=1, stage="production")
def validate_and_promote():
    # Run validation checks
    accuracy = 0.95
    assert accuracy > 0.90, "Model doesn't meet threshold"

Combining Decorators

Stack decorators for a complete ML workflow:

import podstack

podstack.init(api_key="your-api-key", project_id="your-project-id")

@podstack.gpu(type="L40S", pip=["transformers", "datasets"])
@podstack.experiment(name="sentiment-analysis")
@podstack.run(name="bert-finetune-v1", track_gpu=True)
@podstack.model.register(name="sentiment-bert")
def full_pipeline():
    from transformers import AutoModelForSequenceClassification, Trainer

    model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")

    # Log hyperparameters
    podstack.registry.log_params({
        "model": "bert-base-uncased",
        "learning_rate": 2e-5,
        "epochs": 3
    })

    # Train...
    podstack.registry.log_metrics({"accuracy": 0.92, "f1": 0.89})

    return {"accuracy": 0.92}

result = full_pipeline()  # Runs on remote L40S GPU with full tracking

Registry - Experiment Tracking & Model Management

Initialize

from podstack import registry

registry.init(
    api_key="your-api-key",
    project_id="your-project-id"
)

Track Experiments and Runs

from podstack import registry

# Set experiment
registry.set_experiment("my-experiment")

# Start a tracked run
with registry.start_run(name="training-v1") as run:
    # Log hyperparameters
    registry.log_params({
        "learning_rate": 0.001,
        "batch_size": 32,
        "epochs": 10,
        "optimizer": "adam"
    })

    # Log metrics at each step
    for epoch in range(10):
        loss = train_epoch()
        accuracy = evaluate()
        registry.log_metrics({"loss": loss, "accuracy": accuracy}, step=epoch)

    # Set tags
    registry.set_tag("framework", "pytorch")

    # Log artifacts
    registry.log_artifact("model.pt", "model")
    registry.log_artifact("training_curves.png", "plots")

    # Log dataset metadata
    registry.log_dataset(
        name="imdb-reviews",
        path="s3://datasets/imdb",
        num_rows=50000,
        num_features=2
    )

Log and Load Models

from podstack import registry

# Log a model object (auto-detects framework)
registry.log_model(model, artifact_path="model", framework="pytorch")

# Register in model registry
registry.register_model(
    name="my-classifier",
    run_id=run.id,
    description="BERT sentiment classifier"
)

# Promote to production
registry.set_model_stage("my-classifier", version=1, stage="production")

# Set aliases
registry.set_model_alias("my-classifier", alias="champion", version=1)

# Load model from registry
model = registry.load_model("my-classifier", stage="production")

Compare Runs

from podstack import registry

# Compare multiple runs
comparison = registry.compare_runs(
    run_ids=["run-id-1", "run-id-2", "run-id-3"],
    metric_keys=["loss", "accuracy"]
)

# Get metric history for a run
history = registry.get_metric_history("run-id-1", "loss")
for point in history:
    print(f"Step {point.step}: {point.value}")

# Search runs
runs = registry.search_runs(
    experiment_id="exp-id",
    status="completed",
    max_results=50
)

List and Browse

from podstack import registry

# List experiments
experiments = registry.list_experiments()

# List models
models = registry.list_models()

# Download artifacts
registry.download_artifact("run-id", "model/model.pt", "./downloads/")

GPU Runner - Direct Code Execution

For running code strings directly on GPUs without decorators:

import podstack

podstack.init(api_key="your-api-key", project_id="your-project-id")

# Run code on a remote GPU
result = podstack.run_on_gpu('''
import torch
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"Memory: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB")
''', gpu="L40S")

print(result.output)
print(f"Success: {result.success}")
print(f"Duration: {result.duration_seconds}s")

Client API

For direct API access to notebooks and executions:

from podstack import Client

client = Client(api_key="your-api-key")

# Create a notebook
notebook = client.sync_create_notebook(name="experiment", gpu_type="L40S")
print(f"JupyterLab: {notebook.jupyter_url}")

# Run code
result = client.sync_run("print('Hello GPU!')", gpu_type="L40S")
print(result.output)

Error Handling

from podstack import (
    PodstackError,
    AuthenticationError,
    GPUNotAvailableError,
    RateLimitError,
    ExecutionTimeoutError
)

try:
    result = train()
except AuthenticationError:
    print("Invalid API key")
except GPUNotAvailableError as e:
    print(f"GPU not available")
except RateLimitError as e:
    print(f"Rate limited, retry after {e.retry_after}s")
except ExecutionTimeoutError as e:
    print(f"Execution timed out: {e.execution_id}")
except PodstackError as e:
    print(f"Error: {e.message}")

Configuration

import podstack

# Option 1: Initialize explicitly
podstack.init(
    api_key="your-api-key",
    project_id="your-project-id",
    api_url="https://api.podstack.ai/v1",       # optional
    registry_url="https://registry.podstack.ai"  # optional
)

# Option 2: Environment variables
# PODSTACK_API_KEY=your-api-key
# PODSTACK_PROJECT_ID=your-project-id
# PODSTACK_API_URL=https://api.podstack.ai/v1
# PODSTACK_REGISTRY_URL=https://registry.podstack.ai

# Option 3: Auto-init (set PODSTACK_AUTO_INIT=1)
# SDK auto-initializes from env vars at import time

License

MIT License - see LICENSE for details.