stephen schuresko
8bba9ef2a6
Update the cleanup process in vast_start.py to require the user to press ENTER before destroying the instance. This change prevents immediate automatic termination after the subprocess ends, giving the user control over when to tear down the environment.
260 lines
9.1 KiB
Python
260 lines
9.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Vast.ai Instance Manager
|
|
|
|
This script automates the process of finding the cheapest available GPU on Vast.ai
|
|
that meets specific VRAM requirements, creating an instance, and setting up Ollama.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import time
|
|
from typing import Any, Dict, List
|
|
|
|
import requests
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Configuration
|
|
# -----------------------------------------------------------------------------
|
|
OLLAMA_MODEL = "qwen3-coder-next:q8_0"
|
|
MIN_VRAM_MB = 80000
|
|
MAX_BID_HOURLY = 2.5
|
|
DISK_SPACE_GB = int(MIN_VRAM_MB / 700 + 20)
|
|
|
|
VAST_API_BASE = "https://console.vast.ai/api/v0"
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Utility Functions
|
|
# -----------------------------------------------------------------------------
|
|
def load_api_key() -> str:
|
|
"""Loads the Vast.ai API key from the user's config directory."""
|
|
path = os.path.expanduser("~/.config/vastai/vast_api_key")
|
|
try:
|
|
with open(path, "r") as f:
|
|
return f.read().strip()
|
|
except FileNotFoundError:
|
|
raise Exception(f"API key not found at {path}. Please configure vastai CLI.")
|
|
|
|
def load_json_file(path: str) -> Dict[str, Any]:
|
|
"""Loads and parses a JSON file."""
|
|
try:
|
|
with open(path, "r") as f:
|
|
return json.load(f)
|
|
except FileNotFoundError:
|
|
raise Exception(f"Required configuration file '{path}' not found.")
|
|
|
|
# Load globals
|
|
API_KEY = load_api_key()
|
|
SEARCH_PARAMS = load_json_file("search_params.json")
|
|
CREATE_PAYLOAD = load_json_file("create_payload.json")
|
|
|
|
|
|
def vast_request(method: str, endpoint: str, params: Dict = None, json_data: Dict = None) -> Dict[str, Any]:
|
|
"""Sends an authenticated request to the Vast.ai API."""
|
|
url = f"{VAST_API_BASE}/{endpoint}"
|
|
headers = {"Authorization": f"Bearer {API_KEY}"}
|
|
|
|
print(f"[{method}] {url}")
|
|
response = requests.request(method, url, headers=headers, params=params, json=json_data)
|
|
|
|
if not response.ok:
|
|
print(f"API Error: {response.status_code} - {response.text}")
|
|
response.raise_for_status()
|
|
|
|
# Some endpoints like DELETE might return empty bodies
|
|
return response.json() if response.text else {}
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Core Operations
|
|
# -----------------------------------------------------------------------------
|
|
def get_cheapest_gpus() -> List[Dict[str, Any]]:
|
|
"""Searches for verified GPU offers matching our VRAM and price constraints."""
|
|
params = SEARCH_PARAMS.copy()
|
|
if "gpu_ram" not in params:
|
|
params["gpu_ram"] = {}
|
|
params["gpu_ram"]["gt"] = MIN_VRAM_MB
|
|
|
|
response = vast_request("GET", "bundles/", params=params)
|
|
offers = response.get("offers", [])
|
|
|
|
if not offers:
|
|
raise Exception("No GPU offers returned from Vast.ai.")
|
|
|
|
valid_offers = [
|
|
offer for offer in offers
|
|
if offer.get("gpu_ram", 0) > MIN_VRAM_MB
|
|
and offer.get("dph_total", float('inf')) < MAX_BID_HOURLY
|
|
and offer.get("verification") == "verified"
|
|
]
|
|
|
|
if not valid_offers:
|
|
raise Exception("No verified GPU offers found matching the VRAM and price constraints.")
|
|
|
|
# Ensure they are sorted by price (cheapest first)
|
|
valid_offers.sort(key=lambda x: x.get("dph_total", float('inf')))
|
|
return valid_offers
|
|
|
|
def create_instance(offer_id: int) -> int:
|
|
"""Creates a new instance from a specific offer ID."""
|
|
payload = CREATE_PAYLOAD.copy()
|
|
|
|
if "env" not in payload:
|
|
payload["env"] = {}
|
|
|
|
payload["env"]["OLLAMA_MODEL"] = OLLAMA_MODEL
|
|
payload["disk"] = DISK_SPACE_GB
|
|
|
|
response = vast_request("PUT", f"asks/{offer_id}/", json_data=payload)
|
|
if "new_contract" not in response:
|
|
raise Exception(f"Failed to create instance. API Response: {response}")
|
|
|
|
return response["new_contract"]
|
|
|
|
def destroy_instance(instance_id: int) -> None:
|
|
"""Destroys a running instance."""
|
|
print(f"Destroying instance {instance_id}...")
|
|
vast_request("DELETE", f"instances/{instance_id}/")
|
|
print(f"Instance {instance_id} destroyed.")
|
|
|
|
def get_instance_details(instance_id: int) -> Dict[str, Any]:
|
|
"""Fetches current details and status of an instance."""
|
|
response = vast_request("GET", f"instances/{instance_id}/")
|
|
instances = response.get("instances", [])
|
|
|
|
if not instances:
|
|
raise Exception(f"Instance {instance_id} not found.")
|
|
|
|
return instances[0]
|
|
|
|
def wait_until_ready(instance_id: int, poll_interval: int = 15) -> Dict[str, Any]:
|
|
"""Polls the instance status until it is running and ready."""
|
|
print(f"Waiting for instance {instance_id} to become ready...")
|
|
|
|
# Tags we care to monitor during startup
|
|
status_tags = [
|
|
"actual_status", "intended_status", "cur_state", "next_state",
|
|
"gpu_name", "gpu_util", "public_ipaddr", "ssh_port", "uptime_mins"
|
|
]
|
|
|
|
while True:
|
|
data = get_instance_details(instance_id)
|
|
status = data.get("actual_status", "unknown")
|
|
|
|
# Print a summarized status line
|
|
info = []
|
|
for tag in status_tags:
|
|
if tag in data and data[tag] is not None:
|
|
info.append(f"{tag}={data[tag]}")
|
|
|
|
print(f"Status: {status} | {' | '.join(info)}")
|
|
|
|
if status == "running":
|
|
# Extra check to ensure SSH port is populated
|
|
if data.get("ssh_port") and data.get("ssh_host"):
|
|
print("Instance is running and network is configured.")
|
|
return data
|
|
else:
|
|
print("Instance running, but waiting for SSH details to populate...")
|
|
|
|
time.sleep(poll_interval)
|
|
|
|
def ssh_and_setup(ip: str, port: int, user: str = "root") -> None:
|
|
"""Connects to the instance via SSH and sets up the environment."""
|
|
print(f"Setting up instance via SSH at {user}@{ip}:{port}")
|
|
|
|
ssh_base_cmd = [
|
|
"ssh", "-t",
|
|
"-o", "StrictHostKeyChecking=accept-new",
|
|
"-p", str(port),
|
|
f"{user}@{ip}"
|
|
]
|
|
|
|
# Disable auto-tmux
|
|
subprocess.run(ssh_base_cmd + ["touch", "~/.no_auto_tmux"], check=True)
|
|
|
|
# Pull the required Ollama model
|
|
pull_cmd = ssh_base_cmd + ["ollama", "pull", OLLAMA_MODEL]
|
|
print(f"Executing: {' '.join(pull_cmd)}")
|
|
subprocess.run(pull_cmd, check=True)
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Main Execution
|
|
# -----------------------------------------------------------------------------
|
|
def main():
|
|
print("Searching for the cheapest suitable GPU...")
|
|
try:
|
|
offers = get_cheapest_gpus()
|
|
except Exception as e:
|
|
print(f"Error finding GPUs: {e}")
|
|
return
|
|
|
|
instance_id = 0
|
|
for offer in offers:
|
|
offer_id = offer["id"]
|
|
gpu_name = offer.get("gpu_name", "Unknown GPU")
|
|
price = offer.get("dph_total", 0.0)
|
|
|
|
print(f"\nAttempting to claim Offer {offer_id} (GPU: {gpu_name}, Price: ${price:.3f}/hr)")
|
|
|
|
try:
|
|
instance_id = create_instance(offer_id)
|
|
print(f"Successfully claimed offer. Instance ID: {instance_id}")
|
|
break
|
|
except Exception as e:
|
|
print(f"Failed to claim offer {offer_id}: {e}")
|
|
|
|
if not instance_id:
|
|
print("Failed to claim any of the available offers.")
|
|
return
|
|
|
|
try:
|
|
# Wait for initialization
|
|
data = wait_until_ready(instance_id)
|
|
|
|
# Extract connection details
|
|
ssh_host = data.get("ssh_host")
|
|
ssh_port = data.get("ssh_port")
|
|
public_ip = data.get("public_ipaddr")
|
|
|
|
ports_config = data.get("ports", {})
|
|
pub_ssh_port = ports_config.get("22/tcp", [{}])[0].get("HostPort", ssh_port)
|
|
pub_ollama_port = ports_config.get("11434/tcp", [{}])[0].get("HostPort", "Unknown")
|
|
|
|
print(f"\nInstance is fully ready at {ssh_host}:{ssh_port}")
|
|
|
|
# Setup SSH and Ollama
|
|
ssh_and_setup(ssh_host, ssh_port, "root")
|
|
|
|
print("\n--- Setup Complete ---")
|
|
print(f"SSH Command: ssh root@{public_ip} -p {pub_ssh_port}")
|
|
print(f"Ollama Endpoint: http://{public_ip}:{pub_ollama_port} (Model: {OLLAMA_MODEL})")
|
|
|
|
# Run interactive shell or prompt
|
|
print(f"\nStarting interactive Ollama session for model: {OLLAMA_MODEL}")
|
|
run_cmd = [
|
|
"ssh", "-t",
|
|
"-o", "StrictHostKeyChecking=accept-new",
|
|
"-p", str(pub_ssh_port),
|
|
f"root@{public_ip}",
|
|
"ollama", "run", OLLAMA_MODEL
|
|
]
|
|
print(f"Executing: {' '.join(run_cmd)}")
|
|
subprocess.run(run_cmd)
|
|
|
|
except KeyboardInterrupt:
|
|
print("\nOperation cancelled by user.")
|
|
except Exception as e:
|
|
print(f"\nAn error occurred during instance setup: {e}")
|
|
finally:
|
|
if instance_id:
|
|
try:
|
|
input("\nPress ENTER to destroy the instance and clean up...")
|
|
except (EOFError, KeyboardInterrupt):
|
|
pass
|
|
print("\nCleaning up...")
|
|
destroy_instance(instance_id)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|