Building a VastAI GPU Instance Management API

VastAI GPU Instance Management API: A Development Story

AI model training and serving require GPUs. Rather than maintaining our own GPU servers — which is cost-prohibitive — we decided to use VastAI, a GPU cloud marketplace. This post documents how we built an API to manage VastAI instances directly from the XGen platform.

What is VastAI?

VastAI is a marketplace for renting and leasing idle GPUs. It offers a wide variety of GPUs (A100, H100, RTX 4090, etc.) at prices significantly lower than AWS or GCP.

VastAI API Client

import httpx
from typing import Optional, List

class VastAIClient:
    BASE_URL = "https://console.vast.ai/api/v0"

    def __init__(self, api_key: str):
        self.api_key = api_key
        self.headers = {"Authorization": f"Bearer {api_key}"}

    async def search_offers(
        self,
        gpu_name: str = "RTX_4090",
        num_gpus: int = 1,
        min_ram: float = 16.0,
    ) -> List[dict]:
        async with httpx.AsyncClient() as client:
            params = {
                "gpu_name": gpu_name,
                "num_gpus": num_gpus,
                "gpu_ram": f">={min_ram}",
                "order": "dph_total",  # sort by hourly price
                "type": "on-demand",
            }
            resp = await client.get(
                f"{self.BASE_URL}/bundles",
                headers=self.headers,
                params=params,
            )
            resp.raise_for_status()
            return resp.json().get("offers", [])

    async def create_instance(
        self, offer_id: int, image: str, disk_gb: float = 50
    ) -> dict:
        async with httpx.AsyncClient() as client:
            payload = {
                "client_id": "me",
                "image": image,
                "disk": disk_gb,
                "onstart": "bash /opt/setup.sh",
            }
            resp = await client.put(
                f"{self.BASE_URL}/asks/{offer_id}/",
                headers=self.headers,
                json=payload,
            )
            resp.raise_for_status()
            return resp.json()

    async def destroy_instance(self, instance_id: int) -> bool:
        async with httpx.AsyncClient() as client:
            resp = await client.delete(
                f"{self.BASE_URL}/instances/{instance_id}/",
                headers=self.headers,
            )
            return resp.status_code == 200

Instance Health Monitoring

class InstanceMonitor:
    def __init__(self, vast_client: VastAIClient):
        self.client = vast_client

    async def health_check(self, instance_id: int) -> dict:
        instance = await self.client.get_instance(instance_id)
        status = instance.get("actual_status", "unknown")

        return {
            "instance_id": instance_id,
            "status": status,
            "gpu_util": instance.get("gpu_utilization", 0),
            "gpu_temp": instance.get("gpu_temp", 0),
            "uptime": instance.get("duration", 0),
            "cost_so_far": instance.get("total_cost", 0),
        }

    async def auto_destroy_idle(
        self, instance_id: int, idle_minutes: int = 30
    ):
        health = await self.health_check(instance_id)
        if health["gpu_util"] < 5 and health["uptime"] > idle_minutes * 60:
            await self.client.destroy_instance(instance_id)
            return True
        return False

FastAPI Endpoints

router = APIRouter(prefix="/api/gpu", tags=["gpu"])

@router.get("/offers")
async def search_gpu_offers(gpu_name: str = "RTX_4090", num_gpus: int = 1):
    offers = await vast_client.search_offers(gpu_name, num_gpus)
    return {"offers": offers[:10]}

@router.post("/instances")
async def create_gpu_instance(offer_id: int, image: str):
    result = await vast_client.create_instance(offer_id, image)
    return {"instance": result}

@router.delete("/instances/{instance_id}")
async def destroy_gpu_instance(instance_id: int):
    success = await vast_client.destroy_instance(instance_id)
    return {"success": success}

SSH Tunneling

We also implemented SSH tunneling to provide terminal access to instances, allowing team members to open a terminal directly from the XGen platform UI. With VastAI integrated into XGen, the team can now spin up and tear down GPU instances with just a few clicks in the web UI — and we've cut costs by roughly 60% compared to AWS.