stephen schuresko
8a3d3449bf
- Update shebang to use `#!/usr/bin/env python3`. - Add module docstrings and type hinting across all functions. - Move configuration variables into clearly defined uppercase constants. - Improve error handling for file loading and API requests. - Restructure instance creation and SSH management logic into modular, readable functions. These changes make the script more robust, easier to maintain, and self-documenting.
258 lines
9.1 KiB
Python
258 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:
|
|
# Ask user if they want to clean up, or just do it automatically
|
|
# Based on original script, it immediately destroys the instance after the subprocess ends
|
|
print("\nCleaning up...")
|
|
if instance_id:
|
|
destroy_instance(instance_id)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|