vastai-gpu-runner

Cloud GPU orchestration framework with direct Vast.ai deployment, managed declarative jobs, R2/GCS storage, worker lifecycle, and crash recovery. The package keeps its historical vastai-gpu-runner name; the managed-job surface is provider-neutral.

Features

• CloudRunner — direct VM lifecycle (Vast.ai, local) with retry and machine deduplication.
• VastaiRunner — hardened Vast.ai deployment with quality filters and ownership guards. Implements CloudRunner (direct VM SSH lifecycle), not ManagedJobRunner.
• LocalRunner — zero-cost local/CI backend: the same lifecycle as a subprocess, no cloud credentials.
• ManagedJobRunner (Protocol) — provider-neutral interface for declarative cloud batch workloads (GCP Batch, AWS Batch, Azure Batch). The contract behind GcpBatchRunner. Not the contract behind VastaiRunner (that's CloudRunner).
• Managed-job contract — typed lifecycle states, provider-neutral storage mounts, typed provider errors, duplicate-name conflicts, and idempotent cancellation/deletion.
• GcpBatchRunner — Google Cloud Batch backend implementing ManagedJobRunner. Submits jobs, polls status, and exposes task diagnostics.
• GcsSink — Google Cloud Storage artifact sink. Plain writes are create-only, CAS writes use generations, and staged JSON writes do not claim atomic rename or use a DONE marker.
• R2Sink — Cloudflare R2 / S3-compatible result storage with DONE markers and parallel downloads.
• BaseWorker — template method worker: GPU check, preflight gates, self-destruct.
• BatchState — atomic JSON persistence for crash-recoverable batch orchestration.
• Cost estimator — GPU speed factors, live pricing, scaling tables.
• CLI — credential checks, instance listing, cost estimation, orphan cleanup, batch, run, r2-lifecycle.

Installation

uv add "vastai-gpu-runner @ git+https://github.com/Lambda-Biolab/vastai-gpu-runner.git"

Requires Python >= 3.11, vastai CLI (pip install vastai), and R2 credentials in ~/.cloud-credentials.

For GCP Batch + GCS Sink functionality, install with the [gcp] extra:

uv add "vastai-gpu-runner[gcp] @ git+https://github.com/Lambda-Biolab/vastai-gpu-runner.git"
# or equivalently: pip install vastai-gpu-runner[gcp]

The [gcp] extra provides google-cloud-batch>=0.17.0, google-cloud-storage>=3.0.0, and google-cloud-logging>=3.10.0. Install with --all-extras to get all optional backends.

Quick start

from vastai_gpu_runner.providers.vastai import VastaiRunner
from vastai_gpu_runner.cleanup_policy import OwnershipPolicy
from vastai_gpu_runner.types import DeploymentConfig

runner = VastaiRunner(
    DeploymentConfig(gpu_model="RTX_4090", max_cost_per_hour=0.35),
    docker_image="my-org/my-image:latest",
    ownership=OwnershipPolicy(owned_images=frozenset({"my-org/my-image:latest"})),
)

result = runner.run_full_cycle(
    files={"worker.sh": script_path, "input.tar": data_path},
    local_output_dir=output_path,
    max_retries=3,
)

The ownership= path is canonical; allowed_images=frozenset(...) is a deprecated back-compat alias.

See docs/guide.md for workers, batch state, R2 storage, and cost estimation examples.

CLI

vastai-gpu-runner check                  # Verify Vast.ai + R2 credentials
vastai-gpu-runner instances              # List active instances
vastai-gpu-runner estimate -w 10         # Scaling table for 10h of GPU work
vastai-gpu-runner cleanup -l "myproject" # Destroy orphaned instances
vastai-gpu-runner run --provider local --file worker.sh --output outputs/local  # Run locally, no cloud credentials
vastai-gpu-runner batch --state-path batch.json --label myproject --image my-org/my-image:latest  # Compose a Vast.ai batch
vastai-gpu-runner r2-lifecycle show --bucket my-bucket --prefix outputs/ --credentials-file ~/.cloud-credentials.r2-admin
vastai-gpu-runner r2-lifecycle --bucket my-bucket --prefix outputs/  # R2 lifecycle operations

Architecture

vastai_gpu_runner/
├── types.py              # Enums and dataclasses
├── runner.py             # CloudRunner ABC
├── ssh.py                # SSH/SCP utilities
├── state.py              # Batch + job state persistence
├── orchestrator.py       # check_budget
├── cleanup_policy.py     # v4 candidate / verdict / ownership policy
├── cli.py                # CLI: check, instances, estimate, cleanup, run, batch, r2-lifecycle
├── providers/
│   ├── vastai.py         # Vast.ai implementation (CloudRunner)
│   ├── local.py          # Local subprocess backend (CloudRunner)
│   ├── destroy.py        # Provider-agnostic destroy orchestration
│   └── destroy_adapters/ # Per-provider destroy adapters
├── managed_jobs/         # ManagedJobRunner implementations (declarative cloud batch)
│   ├── base.py           # ManagedJobRunner Protocol + DTOs
│   ├── gcp_batch.py      # GcpBatchRunner + FakeGcpBatchClient + FakeGcsClient
│   └── state.py          # ManagedJobState JSON loader
├── storage/
│   ├── r2.py             # R2/S3 result storage
│   ├── r2_lifecycle.py   # R2 lifecycle operations
│   └── gcs.py            # GcsSink (GCS artifact sink)
├── hybrid.py             # Hybrid local+cloud orchestration
├── unit_lifecycle.py     # Unit-level lifecycle helpers
├── worker/base.py        # BaseWorker template method
├── worker/health.py      # GPU + R2 health checks
├── estimator/core.py     # Scaling tables, GPU speed factors
└── estimator/pricing.py  # Live Vast.ai pricing

See docs/architecture.md for design decisions.

Documentation

• User guide — deployment, workers, state, storage, estimation
• Extending — adding providers, custom storage, worker examples
• API reference — all classes, methods, and parameters
• Architecture — module layout and design decisions
• Changelog

The published package is currently 0.6.0. The managed-job contract stabilization is the recommended next 0.7.0 release; 0.7.0 has not been published.

Development

git clone https://github.com/Lambda-Biolab/vastai-gpu-runner.git
cd vastai-gpu-runner
uv sync --all-extras        # include GCP extras so pyright can resolve google.cloud.*
uv run pytest
uv run ruff check src/     # linting
uv run pyright src/ tests/ # type checking (includes tests)

License

Apache 2.0 — see LICENSE.