__init__.py
import logging as _logging
import os as _os
_VAST_LOG_LEVELS = {
"critical": _logging.CRITICAL,
"error": _logging.ERROR,
"warning": _logging.WARNING,
"info": _logging.INFO,
"debug": _logging.DEBUG,
}
_env_level = _os.environ.get("VAST_LOG_LEVEL", "").lower()
_level = (
_VAST_LOG_LEVELS.get(_env_level, _logging.WARNING) if _env_level else _logging.INFO
)
logger = _logging.getLogger("vastai")
logger.setLevel(_level)
if not logger.handlers:
_handler = _logging.StreamHandler()
_handler.setFormatter(_logging.Formatter("%(levelname)s - %(message)s"))
logger.addHandler(_handler)
logger.propagate = False
from .sdk import VastAI
try:
from .serverless.client.client import (
Serverless,
CoroutineServerless,
ServerlessRequest,
)
from .serverless.client.request_status import RequestStatus
from .serverless.client.endpoint import Endpoint
from .serverless.server.worker import Worker
from .serverless.server.worker import WorkerConfig, HandlerConfig, LogActionConfig, BenchmarkConfig
from .sync.client import SyncClient
from .async_.client import AsyncClient
from .serverless.remote import Deployment
except ImportError:
# Serverless dependencies (aiohttp, etc.) not installed
Deployment = None
Serverless = None
CoroutineServerless = None
ServerlessRequest = None
RequestStatus = None
Endpoint = None
Worker = None
WorkerConfig = None
HandlerConfig = None
LogActionConfig = None
BenchmarkConfig = None
SyncClient = None
AsyncClient = None
__all__ = [
"VastAI",
"Deployment",
"Serverless",
"CoroutineServerless",
"ServerlessRequest",
"RequestStatus",
"Endpoint",
"Worker",
"WorkerConfig",
"HandlerConfig",
"LogActionConfig",
"BenchmarkConfig",
"SyncClient",
"AsyncClient",
]
_base.py
import os
from typing import Optional
from vastai.utils import VERSION
_APIKEY_SENTINEL = object()
def _resolve_api_key(api_key: object) -> str:
"""
Resolves the API key from (in order):
1. Explicit argument
2. VAST_API_KEY environment variable
3. XDG config file: $XDG_CONFIG_HOME/vastai/vast_api_key
(falls back to ~/.config/vastai/vast_api_key if xdg is unavailable)
4. Legacy file: ~/.vast_api_key
"""
if api_key is not _APIKEY_SENTINEL:
return api_key
env = os.getenv("VAST_API_KEY")
if env:
return env
try:
import xdg
config_dir = xdg.xdg_config_home()
except ImportError:
from pathlib import Path
config_dir = os.path.join(str(Path.home()), ".config")
xdg_path = os.path.join(config_dir, "vastai", "vast_api_key")
if os.path.exists(xdg_path):
return open(xdg_path).read().strip()
legacy_path = os.path.expanduser("~/.vast_api_key")
if os.path.exists(legacy_path):
return open(legacy_path).read().strip()
raise RuntimeError(
"No API key found. Pass api_key=, set VAST_API_KEY, "
f"or save your key to {xdg_path}"
)
class _BaseClient:
def __init__(self, api_key: object = _APIKEY_SENTINEL, vast_server: str = "https://console.vast.ai"):
self._api_key = _resolve_api_key(api_key)
self._vast_server = vast_server.rstrip("/")
def _headers(self) -> dict:
return {"Authorization": f"Bearer {self._api_key}", "User-Agent": f"vastai-sdk/{VERSION}"}
def _url(self, path: str) -> str:
return f"{self._vast_server}{path}"
api/__init__.py
"""Vast.ai SDK API modules."""
api/auth.py
"""Auth, secrets, templates, scheduled jobs, and TFA API functions for the Vast.ai SDK."""
def show_audit_logs(client):
"""Display account's history of important actions.
GET /audit_logs/
Args:
client: VastClient instance.
Returns:
list: Audit log entries.
"""
r = client.get("/audit_logs/")
r.raise_for_status()
return r.json()
def show_env_vars(client):
"""Show user environment variables.
GET /secrets/
Args:
client: VastClient instance.
Returns:
dict: Environment variables as key-value pairs.
"""
r = client.get("/secrets/")
r.raise_for_status()
return r.json().get("secrets", {})
def create_env_var(client, name, value):
"""Create a new user environment variable.
POST /secrets/
Args:
client: VastClient instance.
name (str): Environment variable name.
value (str): Environment variable value.
Returns:
dict: API response data.
"""
data = {"key": name, "value": value}
r = client.post("/secrets/", json_data=data)
r.raise_for_status()
return r.json()
def update_env_var(client, name, value):
"""Update an existing user environment variable.
PUT /secrets/
Args:
client: VastClient instance.
name (str): Environment variable name.
value (str): New environment variable value.
Returns:
dict: API response data.
"""
data = {"key": name, "value": value}
r = client.put("/secrets/", json_data=data)
r.raise_for_status()
return r.json()
def delete_env_var(client, name):
"""Delete a user environment variable.
DELETE /secrets/
Args:
client: VastClient instance.
name (str): Environment variable name to delete.
Returns:
dict: API response data.
"""
data = {"key": name}
r = client.delete("/secrets/", json_data=data)
r.raise_for_status()
return r.json()
def show_scheduled_jobs(client):
"""Display the list of scheduled jobs.
GET /commands/schedule_job/
Args:
client: VastClient instance.
Returns:
list: Scheduled job entries.
"""
r = client.get("/commands/schedule_job/")
r.raise_for_status()
return r.json()
def create_scheduled_job(client, start_time, end_time, api_endpoint, request_method,
request_body, frequency, instance_id,
day_of_the_week=None, hour_of_the_day=None):
"""Create a new scheduled job.
POST /commands/schedule_job/
Args:
client: VastClient instance.
start_time (float): Start time as unix timestamp.
end_time (float): End time as unix timestamp.
api_endpoint (str): API endpoint the job will call.
request_method (str): HTTP method (GET, POST, PUT, DELETE).
request_body (dict): JSON body for the scheduled request.
frequency (str): One of "HOURLY", "DAILY", "WEEKLY".
instance_id (int): Instance ID the job is associated with.
day_of_the_week (int, optional): Day of week (0=Sunday, 6=Saturday).
Required for WEEKLY, must be None for HOURLY/DAILY.
hour_of_the_day (int, optional): Hour of day (0-23).
Required for DAILY/WEEKLY, must be None for HOURLY.
Returns:
dict: API response data.
"""
json_blob = {
"start_time": start_time,
"end_time": end_time,
"api_endpoint": api_endpoint,
"request_method": request_method,
"request_body": request_body,
"day_of_the_week": day_of_the_week,
"hour_of_the_day": hour_of_the_day,
"frequency": frequency,
"instance_id": instance_id,
}
r = client.post("/commands/schedule_job/", json_data=json_blob)
r.raise_for_status()
return r.json()
def update_scheduled_job(client, id, request_body, start_time=None, end_time=None,
api_endpoint=None, request_method=None,
frequency=None, instance_id=None,
day_of_the_week=None, hour_of_the_day=None):
"""Update an existing scheduled job.
PUT /commands/schedule_job/{id}/
Args:
client: VastClient instance.
id (int): ID of the scheduled job to update.
request_body (dict): Updated JSON body for the scheduled request.
start_time (float, optional): Updated start time as unix timestamp.
end_time (float, optional): Updated end time as unix timestamp.
api_endpoint (str, optional): Updated API endpoint.
request_method (str, optional): Updated HTTP method.
frequency (str, optional): Updated frequency (HOURLY, DAILY, WEEKLY).
instance_id (int, optional): Updated instance ID.
day_of_the_week (int, optional): Updated day of week.
hour_of_the_day (int, optional): Updated hour of day.
Returns:
dict: API response data.
"""
json_blob = {"request_body": request_body}
if start_time is not None:
json_blob["start_time"] = start_time
if end_time is not None:
json_blob["end_time"] = end_time
if api_endpoint is not None:
json_blob["api_endpoint"] = api_endpoint
if request_method is not None:
json_blob["request_method"] = request_method
if frequency is not None:
json_blob["frequency"] = frequency
if instance_id is not None:
json_blob["instance_id"] = instance_id
if day_of_the_week is not None:
json_blob["day_of_the_week"] = day_of_the_week
if hour_of_the_day is not None:
json_blob["hour_of_the_day"] = hour_of_the_day
r = client.put(f"/commands/schedule_job/{id}/", json_data=json_blob)
r.raise_for_status()
return r.json()
def delete_scheduled_job(client, id):
"""Delete a scheduled job.
DELETE /commands/schedule_job/{id}/
Args:
client: VastClient instance.
id (int): ID of scheduled job to delete.
Returns:
dict: API response data.
"""
r = client.delete(f"/commands/schedule_job/{id}/")
r.raise_for_status()
return r.json()
# --- Two-Factor Authentication (TFA) Functions ---
def _build_tfa_verification_payload(**kwargs):
"""Build common payload for TFA verification requests.
Args:
**kwargs: TFA verification fields including:
tfa_method_id, tfa_method, code, backup_code, secret,
and any additional fields.
Returns:
dict: Payload with only non-None values.
"""
payload = {
"tfa_method_id": kwargs.get("tfa_method_id"),
"tfa_method": kwargs.get("tfa_method"),
"code": kwargs.get("code"),
"backup_code": kwargs.get("backup_code"),
"secret": kwargs.get("secret"),
}
for key, value in kwargs.items():
if key not in payload:
payload[key] = value
return {k: v for k, v in payload.items() if v is not None}
def tfa_activate(client, code, secret, method_type="totp", phone_number=None, label=None, method_id=None):
"""Activate a new 2FA method by verifying the code.
POST /api/v0/tfa/confirm-new/
Args:
client: VastClient instance.
code (str): 6-digit verification code from SMS or Authenticator app.
secret (str): Secret token from setup process.
method_type (str): 2FA method type ('sms' or 'totp'). Default 'totp'.
phone_number (str, optional): Phone number for SMS method (E.164 format).
label (str, optional): Label for the new 2FA method.
method_id (str, optional): 2FA Method ID if multiple of the same type.
Returns:
dict: API response data, may include backup_codes on first activation.
"""
payload = _build_tfa_verification_payload(
tfa_method_id=method_id,
tfa_method=method_type,
code=code,
secret=secret,
phone_number=phone_number,
label=label,
)
r = client.post("/api/v0/tfa/confirm-new/", json_data=payload)
r.raise_for_status()
return r.json()
def tfa_delete(client, code=None, backup_code=None, method_type=None, secret=None,
method_id=None, id_to_delete=None):
"""Remove a 2FA method from your account.
DELETE /api/v0/tfa/
Requires 2FA verification to prevent unauthorized removals.
Args:
client: VastClient instance.
code (str, optional): 2FA code from Authenticator app, SMS, or Email.
backup_code (str, optional): One-time backup code (alternative to code).
method_type (str, optional): 2FA method type ('email', 'sms', 'totp').
secret (str, optional): Secret token (required for SMS or Email).
method_id (str, optional): 2FA Method ID.
id_to_delete (int, optional): ID of the 2FA method to delete.
Returns:
dict: API response data including remaining_methods count.
"""
payload = _build_tfa_verification_payload(
tfa_method_id=method_id,
tfa_method=method_type,
code=code,
backup_code=backup_code,
secret=secret,
target_id=id_to_delete,
)
r = client.delete("/api/v0/tfa/", json_data=payload)
r.raise_for_status()
return r.json()
def tfa_login(client, code=None, backup_code=None, method_type=None, secret=None, method_id=None):
"""Complete 2FA login by verifying code.
POST /api/v0/tfa/
Args:
client: VastClient instance.
code (str, optional): 2FA code from Authenticator app, SMS, or Email.
backup_code (str, optional): One-time backup code (alternative to code).
method_type (str, optional): 2FA method type ('email', 'sms', 'totp').
secret (str, optional): Secret token (required for SMS or Email).
method_id (str, optional): 2FA Method ID.
Returns:
dict: API response data including session_key and
backup_codes_remaining.
"""
payload = _build_tfa_verification_payload(
tfa_method_id=method_id,
tfa_method=method_type,
code=code,
backup_code=backup_code,
secret=secret,
)
r = client.post("/api/v0/tfa/", json_data=payload)
r.raise_for_status()
return r.json()
def tfa_resend_sms(client, secret, phone_number=None):
"""Resend SMS 2FA code.
POST /api/v0/tfa/sms/resend/
Args:
client: VastClient instance.
secret (str): Secret token from the original 2FA request.
phone_number (str, optional): Phone number to receive SMS code
(E.164 format).
Returns:
dict: API response data.
"""
payload = _build_tfa_verification_payload(
secret=secret,
phone_number=phone_number,
)
r = client.post("/api/v0/tfa/sms/resend/", json_data=payload)
r.raise_for_status()
return r.json()
def tfa_regen_codes(client, code=None, backup_code=None, method_type=None, secret=None, method_id=None):
"""Regenerate backup codes for 2FA.
PUT /api/v0/tfa/regen-backup-codes/
Warning: This will invalidate all existing backup codes.
Args:
client: VastClient instance.
code (str, optional): 2FA code from Authenticator app, SMS, or Email.
backup_code (str, optional): One-time backup code (alternative to code).
method_type (str, optional): 2FA method type ('email', 'sms', 'totp').
secret (str, optional): Secret token (required for SMS or Email).
method_id (str, optional): 2FA Method ID.
Returns:
dict: API response data including new backup_codes list.
"""
payload = _build_tfa_verification_payload(
tfa_method_id=method_id,
tfa_method=method_type,
code=code,
backup_code=backup_code,
secret=secret,
)
r = client.put("/api/v0/tfa/regen-backup-codes/", json_data=payload)
r.raise_for_status()
return r.json()
def tfa_send_sms(client, phone_number=None):
"""Request a 2FA SMS verification code.
POST /api/v0/tfa/sms/
Args:
client: VastClient instance.
phone_number (str, optional): Phone number to receive SMS code
(E.164 format). If not provided, uses account phone number.
Returns:
dict: API response data including secret token.
"""
payload = {}
if phone_number:
payload["phone_number"] = phone_number
r = client.post("/api/v0/tfa/sms/", json_data=payload)
r.raise_for_status()
return r.json()
def tfa_status(client):
"""Show the current 2FA status and configured methods.
GET /tfa/status/
Args:
client: VastClient instance.
Returns:
dict: 2FA status including tfa_enabled, methods list,
and backup_codes_remaining.
"""
r = client.get("/tfa/status/")
r.raise_for_status()
return r.json()
def tfa_totp_setup(client):
"""Generate TOTP secret and QR code for Authenticator app setup.
POST /api/v0/tfa/totp-setup/
Args:
client: VastClient instance.
Returns:
dict: Setup data including secret and provisioning_uri.
"""
r = client.post("/api/v0/tfa/totp-setup/", json_data={})
r.raise_for_status()
return r.json()
def tfa_update(client, method_id, label=None, set_primary=None):
"""Update a 2FA method's settings.
PUT /api/v0/tfa/update/
Args:
client: VastClient instance.
method_id (int): ID of the 2FA method to update.
label (str, optional): New label/name for this 2FA method.
set_primary (bool, optional): Set this method as the primary
2FA method.
Returns:
dict: API response data including updated method info.
Raises:
ValueError: If neither label nor set_primary is provided.
"""
payload = {"tfa_method_id": method_id}
if label is not None:
payload["label"] = label
if set_primary is not None:
payload["is_primary"] = set_primary
if len(payload) == 1:
raise ValueError("Must specify at least one field to update (label or set_primary)")
r = client.put("/api/v0/tfa/update/", json_data=payload)
r.raise_for_status()
return r.json()
def tfa_auth_new(client, code=None, secret=None, backup_code=None,
method_type="email", method_id=None):
"""Authorize account to add a new 2FA method.
POST/PUT /api/v0/tfa/authorize-new-method/
Args:
client: VastClient instance.
code (str, optional): 2FA verification code.
secret (str, optional): Secret token from previous auth step.
backup_code (str, optional): One-time backup code.
method_type (str): 2FA method type ('email', 'sms', 'totp').
method_id (str, optional): Specific method ID to use.
Returns:
dict: API response data.
"""
url = "/api/v0/tfa/authorize-new-method/"
if secret and code:
payload = {"secret": secret, "code": code}
r = client.put(url, json_data=payload)
r.raise_for_status()
return r.json()
payload = {}
if backup_code:
payload["backup_code"] = backup_code
elif method_id:
payload["tfa_method_id"] = method_id
elif method_type:
payload["tfa_method"] = method_type
r = client.post(url, json_data=payload)
r.raise_for_status()
return r.json()
def tfa_send_email(client):
"""Request a 2FA email verification code.
POST /api/v0/tfa/email/
Args:
client: VastClient instance.
Returns:
dict: API response data including secret token.
"""
r = client.post("/api/v0/tfa/email/", json_data={})
r.raise_for_status()
return r.json()
api/billing.py
"""Billing, user, and account API functions for the Vast.ai SDK."""
import time
from dateutil import parser as dateutil_parser
def parse_date_arg(value, param):
"""Parse a user-supplied date, or raise naming the argument."""
try:
return dateutil_parser.parse(str(value))
except (ValueError, OverflowError) as exc:
raise ValueError(
f"{param}: could not parse {value!r} as a date. Use a format like "
"'2026-08-24' or '2026-08-24 13:00:00'."
) from exc
def show_invoices(client, start_date=None, end_date=None, only_charges=False, only_credits=False):
"""Get billing history reports (deprecated endpoint).
GET /users/me/invoices
Args:
client: VastClient instance.
start_date (str, optional): Start date/time for report.
end_date (str, optional): End date/time for report.
only_charges (bool): Show only charge items. Default False.
only_credits (bool): Show only credit items. Default False.
Returns:
dict: Invoice data including 'invoices' and 'current' charges.
"""
end_timestamp = time.time()
start_timestamp = time.time() - (24 * 60 * 60)
if end_date:
end_timestamp = parse_date_arg(end_date, "end_date").timestamp()
if start_date:
start_timestamp = parse_date_arg(start_date, "start_date").timestamp()
query_args = {
"owner": "me",
"sdate": start_timestamp,
"edate": end_timestamp,
"inc_charges": not only_credits,
}
r = client.get("/users/me/invoices", query_args=query_args)
r.raise_for_status()
rj = r.json()
rows = rj["invoices"]
current_charges = rj["current"]
if only_charges:
rows = [row for row in rows if row.get("type") == "charge"]
elif only_credits:
rows = [row for row in rows if row.get("type") == "payment"]
return {
"invoices": rows,
"current": current_charges,
}
def show_invoices_v1(client, params):
"""Get billing (invoices/charges) history reports with advanced filtering.
GET /api/v0/charges/ (for charges) or /api/v1/invoices/ (for invoices)
Args:
client: VastClient instance.
params (dict): Request parameters including:
charges (bool): If True, fetch charges; otherwise fetch invoices.
start_date (int/str, optional): Start date (YYYY-MM-DD or timestamp).
end_date (int/str, optional): End date (YYYY-MM-DD or timestamp).
limit (int): Number of results per page. Default 20, max 100.
next_token (str, optional): Pagination token for next page.
latest_first (bool): Sort by latest first. Default False.
charge_type (list, optional): Filter charge types (instance, volume, serverless).
invoice_type (list, optional): Filter invoice types.
format (str): Output format (table or tree). Default 'table'.
Returns:
dict: API response with 'results', 'count', 'total', and 'next_token'.
"""
from datetime import datetime, timezone
invoice_types_map = {
"transfers": "transfer",
"stripe": "stripe_payments",
"bitpay": "bitpay",
"coinbase": "coinbase",
"crypto.com": "crypto.com",
"reserved": "instance_prepay",
"payout_paypal": "paypal_manual",
"payout_wise": "wise_manual",
}
is_charges = params.get("charges", False)
start_date = params.get("start_date")
end_date = params.get("end_date")
# Handle default start and end date values
if not start_date and not end_date:
end_date = int(time.time())
if not start_date:
end_date_val = end_date if isinstance(end_date, int) else int(
datetime.strptime(str(end_date) + "+0000", '%Y-%m-%d%z').timestamp()
if isinstance(end_date, str) and not str(end_date).isdigit()
else int(end_date)
)
start_date = end_date_val - 7 * 24 * 60 * 60
elif not end_date:
start_date_val = start_date if isinstance(start_date, int) else int(
datetime.strptime(str(start_date) + "+0000", '%Y-%m-%d%z').timestamp()
if isinstance(start_date, str) and not str(start_date).isdigit()
else int(start_date)
)
end_date = start_date_val + 7 * 24 * 60 * 60
def to_timestamp(val):
if isinstance(val, int):
return val
if isinstance(val, str):
if val.isdigit():
return int(val)
return int(datetime.strptime(val + "+0000", '%Y-%m-%d%z').timestamp())
raise ValueError("Invalid date format")
start_timestamp = to_timestamp(start_date)
end_timestamp = to_timestamp(end_date)
# Build request parameters
date_col = 'day' if is_charges else 'when'
query_params = {
'select_filters': {date_col: {'gte': start_timestamp, 'lte': end_timestamp}},
'latest_first': params.get('latest_first', False),
'limit': min(params.get('limit', 20), 100),
}
if is_charges:
query_params['format'] = params.get('format', 'table')
charge_types = params.get('charge_type', [])
for ct in charge_types:
filters = query_params['select_filters'].setdefault('type', {}).setdefault('in', [])
if ct in {'i', 'instance'}:
filters.append('instance')
elif ct in {'v', 'volume'}:
filters.append('volume')
elif ct in {'s', 'serverless'}:
filters.append('serverless')
else:
invoice_type_list = params.get('invoice_type', [])
for it in invoice_type_list:
filters = query_params['select_filters'].setdefault('service', {}).setdefault('in', [])
if it in invoice_types_map:
filters.append(invoice_types_map[it])
next_token = params.get('next_token')
if next_token:
query_params['after_token'] = next_token
endpoint = '/api/v0/charges/' if is_charges else '/api/v1/invoices/'
r = client.get(endpoint, query_args=query_params)
r.raise_for_status()
return r.json()
def show_earnings(client, start_date=None, end_date=None, machine_id=None):
"""Get machine earning history reports.
GET /users/me/machine-earnings
Args:
client: VastClient instance.
start_date (str, optional): Start date/time for report.
end_date (str, optional): End date/time for report.
machine_id (int, optional): Machine ID to filter by.
Returns:
dict/list: Earnings data from the API.
"""
Minutes = 60.0
Hours = 60.0 * Minutes
Days = 24.0 * Hours
cday = time.time() / Days
sday = cday - 1.0
eday = cday - 1.0
if end_date:
eday = parse_date_arg(end_date, "end_date").timestamp() / Days
if start_date:
sday = parse_date_arg(start_date, "start_date").timestamp() / Days
query_args = {
"owner": "me",
"sday": sday,
"eday": eday,
"machid": machine_id,
}
r = client.get("/users/me/machine-earnings", query_args=query_args)
r.raise_for_status()
return r.json()
def show_deposit(client, id):
"""Display reserve deposit info for an instance.
GET /instances/balance/{id}/
Args:
client: VastClient instance.
id (int): ID of instance to get deposit info for.
Returns:
dict: Deposit/balance information.
"""
r = client.get(f"/instances/balance/{id}/", query_args={"owner": "me"})
r.raise_for_status()
return r.json()
def show_user(client):
"""Get current user data.
GET /users/current
Args:
client: VastClient instance.
Returns:
dict: User data (with api_key removed).
"""
r = client.get("/users/current")
r.raise_for_status()
user_blob = r.json()
user_blob.pop("api_key", None)
return user_blob
def set_user(client, params):
"""Update current user settings.
PUT /users/
Args:
client: VastClient instance.
params (dict): User settings to update. Possible keys include:
ssh_key, api_key, billaddress_name, billaddress_addr1,
billaddress_addr2, billaddress_city, billaddress_zip,
billaddress_country, billaddress_taxinfo,
balance_threshold_enabled, balance_threshold,
autobill_threshold, phone_number.
Returns:
dict: API response data.
"""
r = client.put("/users/", json_data=params)
r.raise_for_status()
return r.json()
def show_subaccounts(client):
"""Get current subaccounts.
GET /subaccounts
Args:
client: VastClient instance.
Returns:
list: Subaccount user data.
"""
r = client.get("/subaccounts", query_args={"owner": "me"})
r.raise_for_status()
return r.json()["users"]
def create_subaccount(client, email, username, password, host_only=False):
"""Create a subaccount.
POST /users/
Args:
client: VastClient instance.
email (str): Email address for the subaccount.
username (str): Username for the subaccount.
password (str): Password for the subaccount.
host_only (bool): If True, create as host-only account. Default False.
Returns:
dict: API response data.
"""
json_blob = {
"email": email,
"username": username,
"password": password,
"host_only": host_only,
"parent_id": "me",
}
r = client.post("/users/", json_data=json_blob)
r.raise_for_status()
return r.json()
def fetch_contracts(client, label=None, contract_ids=None):
"""Fetch contracts, optionally filtered by label.
POST /contracts/fetch/
Args:
client: VastClient instance.
label (str, optional): Instance label to filter by.
contract_ids (list, optional): List of contract IDs to filter.
Returns:
list: Matching contract dicts.
"""
json_blob = {}
if label is not None:
json_blob["label"] = label
if contract_ids is not None:
json_blob["contract_ids"] = list(contract_ids)
r = client.post("/contracts/fetch/", json_data=json_blob)
r.raise_for_status()
return r.json()["contracts"]
def show_ipaddrs(client):
"""Display user's history of IP addresses.
GET /users/me/ipaddrs
Args:
client: VastClient instance.
Returns:
list: IP address access history.
"""
r = client.get("/users/me/ipaddrs", query_args={"owner": "me"})
r.raise_for_status()
return r.json()["results"]
api/client.py
"""HTTP client for the Vast.ai API."""
import os
import re
import json
import sys
import shlex
import time
import requests
from urllib.parse import quote_plus, urlsplit
from typing import Dict, Optional
from vastai.utils import VERSION
# Emoji support
_HAS_EMOJI = sys.stdout.encoding and 'utf' in sys.stdout.encoding.lower()
INFO = "\u2139\ufe0f" if _HAS_EMOJI else "[i]"
server_url_default = os.getenv("VAST_URL") or "https://console.vast.ai"
# Our own environments redirect between hosts here; see VastSession.same_service.
VAST_DOMAIN = "vast.ai"
# Status codes that indicate a transient server-side issue worth retrying.
# 429 = rate limited; 502/503/504 = gateway/upstream errors that usually clear.
_RETRYABLE_STATUS = {429, 502, 503, 504}
# Longest server Retry-After we wait out inline; a longer one is returned to the caller.
_MAX_RETRY_AFTER_SECONDS = 60
# Transport-level exceptions worth retrying (network hiccups, slow peers).
# Other requests.RequestException subclasses (e.g. InvalidURL, TooManyRedirects)
# propagate immediately — retrying them would just burn time.
_RETRYABLE_EXC = (
requests.exceptions.ConnectionError,
requests.exceptions.Timeout,
)
# Default per-request timeout. Conservative enough for slow operations like
# log retrieval and instance creation; callers can override per-call.
_DEFAULT_TIMEOUT_SECONDS = 120
def as_curl_command(prep) -> str:
"""Render a prepared request as a runnable, one-flag-per-line curl command."""
parts = ["curl"]
if prep.method != "GET" or prep.body:
parts.append(f"-X {prep.method}")
auth = prep.headers.get("Authorization")
if auth:
parts.append("-H " + shlex.quote(f"Authorization: {auth}"))
if prep.body:
body = prep.body.decode() if isinstance(prep.body, bytes) else prep.body
parts.append("-H " + shlex.quote("Content-Type: application/json"))
parts.append("-d " + shlex.quote(body))
parts.append(shlex.quote(prep.url))
return " \\\n ".join(parts)
class VastSession(requests.Session):
"""A session that keeps our auth header across redirects within Vast's domain."""
def __init__(self, server_url: str):
super().__init__()
self._origin = urlsplit(server_url)
def rebuild_auth(self, prepared_request, response):
if not self.same_service(urlsplit(prepared_request.url)):
super().rebuild_auth(prepared_request, response)
def same_service(self, target) -> bool:
"""Whether a redirect target is still the service we were pointed at.
Same scheme and port, and either the identical host or another host
under vast.ai, which is where our own environments redirect
(candidate.vast.ai -> candidate-server.vast.ai). Everything else needs
an exact host, so a custom --url never widens the trust boundary.
"""
origin = self._origin
if target.scheme != origin.scheme or port_of(target) != port_of(origin):
return False
a, b = (origin.hostname or "").lower(), (target.hostname or "").lower()
return bool(a) and (a == b or (in_vast_domain(a) and in_vast_domain(b)))
def is_search_quota_exceeded(r) -> bool:
"""Whether ``r`` is a daily search quota 429.
The quota only resets at UTC midnight, so retrying it just burns requests.
"""
if r.status_code != 429:
return False
try:
return r.json().get("error") == "search_quota_exceeded"
except (ValueError, AttributeError):
return False
def retry_delay(r, backoff: float) -> Optional[float]:
"""Seconds to wait before retrying ``r``, or None when it should not be retried.
Honors the server's ``Retry-After`` (rate limit token refill) and falls back to
``backoff`` when it is absent or not a number. A search quota 429 or a wait over
``_MAX_RETRY_AFTER_SECONDS`` is not retried.
"""
if is_search_quota_exceeded(r):
return None
try:
wait = float(r.headers["Retry-After"])
except (KeyError, TypeError, ValueError):
return backoff
return wait if wait <= _MAX_RETRY_AFTER_SECONDS else None
def port_of(parts) -> Optional[int]:
return parts.port or {"https": 443, "http": 80}.get(parts.scheme)
def in_vast_domain(host: str) -> bool:
return host == VAST_DOMAIN or host.endswith("." + VAST_DOMAIN)
class VastClient:
"""HTTP client for Vast.ai API requests."""
def __init__(self, api_key=None, server_url=None, retry=3, explain=False, curl=False,
timeout=_DEFAULT_TIMEOUT_SECONDS, client_type="sdk"):
self.api_key = api_key
self.server_url = server_url or server_url_default
self.retry = retry
self.explain = explain
self.curl = curl
self.timeout = timeout
self.user_agent = f"vastai-{client_type}/{VERSION}"
def _build_url(self, subpath: str, query_args: Optional[Dict] = None) -> str:
"""Build full API URL from subpath and optional query args."""
if query_args is None:
query_args = {}
if not re.match(r"^/api/v(\d)+/", subpath):
subpath = "/api/v0" + subpath
if query_args:
query_json = "&".join(
"{x}={y}".format(x=x, y=quote_plus(y if isinstance(y, str) else json.dumps(y)))
for x, y in query_args.items()
)
result = self.server_url + subpath + "?" + query_json
else:
result = self.server_url + subpath
if self.explain:
print("query args:")
print(query_args)
print("")
print(f"base: {self.server_url + subpath + '?'} + query: ")
print(result)
print("")
return result
def _build_headers(self) -> Dict:
"""Build request headers with auth."""
result = {"User-Agent": self.user_agent}
if self.api_key is not None:
result["Authorization"] = "Bearer " + self.api_key
return result
def _request(self, method: str, url: str, headers: Dict, json_data=None,
timeout: Optional[float] = None) -> requests.Response:
"""Execute HTTP request with retry/timeout/exception handling.
Retries are attempted on:
- ``_RETRYABLE_STATUS`` codes (429, 502, 503, 504), waiting per ``retry_delay``
- ``_RETRYABLE_EXC`` transport exceptions (ConnectionError, Timeout)
After exhausting retries, the last response is returned (for status-code
retries) or the last exception is re-raised (for transport exceptions).
Non-retryable ``requests`` exceptions propagate immediately.
"""
effective_timeout = timeout if timeout is not None else self.timeout
t = 0.15
r = None
for i in range(0, self.retry):
req = requests.Request(method=method, url=url, headers=headers, json=json_data)
session = VastSession(self.server_url)
prep = session.prepare_request(req)
if self.explain:
print(f"\n{INFO} Prepared Request:")
print(f"{prep.method} {prep.url}")
print(f"Headers: {json.dumps(headers, indent=1)}")
print(f"Body: {json.dumps(json_data, indent=1)}" + "\n" + "_" * 100 + "\n")
if self.curl:
print("\n" + as_curl_command(prep) + "\n")
sys.exit(0)
try:
r = session.send(prep, timeout=effective_timeout)
except _RETRYABLE_EXC:
if i == self.retry - 1:
raise
time.sleep(t)
t *= 1.5
continue
delay = retry_delay(r, t) if r.status_code in _RETRYABLE_STATUS else None
if delay is not None and i < self.retry - 1:
time.sleep(delay)
t *= 1.5
continue
break
return r
def get(self, subpath: str, query_args: Optional[Dict] = None, json_data=None,
timeout: Optional[float] = None) -> requests.Response:
url = self._build_url(subpath, query_args)
headers = self._build_headers()
return self._request('GET', url, headers, json_data, timeout=timeout)
def post(self, subpath: str, query_args: Optional[Dict] = None, json_data=None,
timeout: Optional[float] = None) -> requests.Response:
url = self._build_url(subpath, query_args)
headers = self._build_headers()
return self._request('POST', url, headers, json_data if json_data is not None else {}, timeout=timeout)
def put(self, subpath: str, query_args: Optional[Dict] = None, json_data=None,
timeout: Optional[float] = None) -> requests.Response:
url = self._build_url(subpath, query_args)
headers = self._build_headers()
return self._request('PUT', url, headers, json_data if json_data is not None else {}, timeout=timeout)
def delete(self, subpath: str, query_args: Optional[Dict] = None, json_data=None,
timeout: Optional[float] = None) -> requests.Response:
url = self._build_url(subpath, query_args)
headers = self._build_headers()
return self._request('DELETE', url, headers, json_data if json_data is not None else {}, timeout=timeout)
api/clusters.py
"""Cluster and overlay network API functions for the Vast.ai SDK."""
def show_clusters(client):
"""Show clusters associated with your account.
GET /clusters/
Args:
client: VastClient instance.
Returns:
dict: Cluster data including nodes, subnets, and manager info.
"""
r = client.get("/clusters/")
r.raise_for_status()
return r.json()
def create_cluster(client, subnet, manager_id):
"""Create a Vast cluster.
POST /cluster/
Args:
client: VastClient instance.
subnet (str): Local subnet for cluster (e.g. '0.0.0.0/24').
manager_id (int): Machine ID of manager node. Must already exist.
Returns:
dict: API response data.
"""
json_blob = {
"subnet": subnet,
"manager_id": manager_id,
}
r = client.post("/cluster/", json_data=json_blob)
r.raise_for_status()
return r.json()
def delete_cluster(client, cluster_id):
"""Delete a Vast cluster.
DELETE /cluster/
Args:
client: VastClient instance.
cluster_id (int): ID of cluster to delete.
Returns:
dict: API response data.
"""
json_blob = {"cluster_id": cluster_id}
r = client.delete("/cluster/", json_data=json_blob)
r.raise_for_status()
return r.json()
def join_cluster(client, cluster_id, machine_ids):
"""Join machine(s) to a cluster.
PUT /cluster/
Args:
client: VastClient instance.
cluster_id (int): ID of cluster to join.
machine_ids (list[int]): Machine IDs to join to the cluster.
Returns:
dict: API response data.
"""
json_blob = {
"cluster_id": cluster_id,
"machine_ids": machine_ids,
}
r = client.put("/cluster/", json_data=json_blob)
r.raise_for_status()
return r.json()
def remove_machine_from_cluster(client, cluster_id, machine_id, new_manager_id=None):
"""Remove a machine from a cluster.
DELETE /cluster/remove_machine/
If removing the manager node, a new_manager_id must be specified.
Args:
client: VastClient instance.
cluster_id (int): ID of cluster.
machine_id (int): ID of machine to remove.
new_manager_id (int, optional): ID of machine to promote to manager.
Must already be in the cluster.
Returns:
dict: API response data.
"""
json_blob = {
"cluster_id": cluster_id,
"machine_id": machine_id,
}
if new_manager_id is not None:
json_blob["new_manager_id"] = new_manager_id
r = client.delete("/cluster/remove_machine/", json_data=json_blob)
r.raise_for_status()
return r.json()
def show_overlays(client):
"""Show overlays associated with your account.
GET /overlay/
Args:
client: VastClient instance.
Returns:
list: Overlay network data including instances and subnets.
"""
r = client.get("/overlay/")
r.raise_for_status()
return r.json()
def create_overlay(client, cluster_id, name):
"""Create an overlay network on top of a physical cluster.
POST /overlay/
Args:
client: VastClient instance.
cluster_id (int): ID of cluster to create overlay on.
name (str): Overlay network name.
Returns:
dict: API response data.
"""
json_blob = {
"cluster_id": cluster_id,
"name": name,
}
r = client.post("/overlay/", json_data=json_blob)
r.raise_for_status()
return r.json()
def delete_overlay(client, overlay_identifier):
"""Delete an overlay and remove all associated instances.
DELETE /overlay/
Args:
client: VastClient instance.
overlay_identifier: ID (int) or name (str) of overlay to delete.
Returns:
dict: API response data.
"""
try:
overlay_id = int(overlay_identifier)
json_blob = {"overlay_id": overlay_id}
except (ValueError, TypeError):
json_blob = {"overlay_name": overlay_identifier}
r = client.delete("/overlay/", json_data=json_blob)
r.raise_for_status()
return r.json()
def join_overlay(client, name, instance_id):
"""Add an instance to an overlay network.
PUT /overlay/
Args:
client: VastClient instance.
name (str): Overlay network name to join.
instance_id (int): Instance ID to add to the overlay.
Returns:
dict: API response data.
"""
json_blob = {
"name": name,
"instance_id": instance_id,
}
r = client.put("/overlay/", json_data=json_blob)
r.raise_for_status()
return r.json()
api/deployments.py
"""Deployment CRUD operations."""
from vastai.api.client import VastClient
def show_deployments(client: VastClient) -> list:
r = client.get("/deployments")
r.raise_for_status()
return r.json()["deployments"]
def show_deployment(client: VastClient, id: int) -> dict:
r = client.get(f"/deployment/{id}/")
r.raise_for_status()
return r.json()["deployment"]
def show_deployment_versions(client: VastClient, id: int) -> list:
r = client.get(f"/deployment/{id}/versions/")
r.raise_for_status()
rj = r.json()
if rj.get("success"):
return rj["versions"]
else:
raise RuntimeError(rj.get("msg", "Unknown error"))
def delete_deployment(client: VastClient, id: int) -> dict:
r = client.delete(f"/deployment/{id}/")
r.raise_for_status()
return r.json()
def stop_deployment(client: VastClient, id: int) -> dict:
r = client.post(f"/deployment/{id}/stop/")
r.raise_for_status()
return r.json()
def start_deployment(client: VastClient, id: int) -> dict:
r = client.post(f"/deployment/{id}/start/")
r.raise_for_status()
return r.json()
def delete_deployment_by_name(client: VastClient, name: str, tag: str = None) -> dict:
json_blob = {"name": name}
if tag is not None:
json_blob["tag"] = tag
r = client.delete("/deployments/", json_data=json_blob)
r.raise_for_status()
return r.json()
api/endpoints.py
"""Endpoint and workergroup API functions for the Vast.ai SDK."""
def show_endpoints(client):
"""Display user's current endpoint groups.
GET /endptjobs/
Args:
client: VastClient instance.
Returns:
list: Endpoint group results with sensitive fields removed.
"""
json_blob = {"client_id": "me", "api_key": client.api_key}
r = client.get("/endptjobs/", json_data=json_blob)
r.raise_for_status()
if r.status_code == 200:
rj = r.json()
if rj["success"]:
rows = rj["results"]
for row in rows:
row.pop("api_key", None)
row.pop("auto_delete_in_seconds", None)
row.pop("auto_delete_due_24h", None)
return rows
else:
return {"error": rj["msg"]}
return r.json()
def create_endpoint(client, **kwargs):
"""Create a new endpoint group.
POST /endptjobs/
Args:
client: VastClient instance.
**kwargs: Endpoint configuration options:
min_load (float): Minimum floor load in perf units/s. Default 0.0.
min_cold_load (float): Minimum floor load allowing cold workers. Default 0.0.
target_util (float): Target capacity utilization (max 1.0). Default 0.9.
cold_mult (float): Cold capacity target as multiple of hot. Default 2.5.
cold_workers (int): Min cold workers when no load. Default 5.
max_workers (int): Max workers for the endpoint group. Default 20.
endpoint_name (str): Deployment endpoint name.
auto_instance (str): Autoscaler instance type. Default "prod".
Returns:
dict: API response data.
"""
json_blob = {
"client_id": "me",
"min_load": kwargs.get("min_load", 0.0),
"min_cold_load": kwargs.get("min_cold_load", 0.0),
"target_util": kwargs.get("target_util", 0.9),
"cold_mult": kwargs.get("cold_mult", 2.5),
"cold_workers": kwargs.get("cold_workers", 5),
"max_workers": kwargs.get("max_workers", 20),
"endpoint_name": kwargs.get("endpoint_name"),
"max_queue_time": kwargs.get("max_queue_time"),
"target_queue_time": kwargs.get("target_queue_time"),
"inactivity_timeout": kwargs.get("inactivity_timeout"),
"autoscaler_instance": kwargs.get("auto_instance", "prod"),
}
r = client.post("/endptjobs/", json_data=json_blob)
r.raise_for_status()
return r.json()
def update_endpoint(client, id, **kwargs):
"""Update an existing endpoint group.
PUT /endptjobs/{id}/
Args:
client: VastClient instance.
id (int): ID of endpoint group to update.
**kwargs: Endpoint configuration options.
Returns:
dict: API response data.
"""
json_blob = {
"client_id": "me",
"endptjob_id": id,
"min_load": kwargs.get("min_load"),
"min_cold_load": kwargs.get("min_cold_load"),
"target_util": kwargs.get("target_util"),
"cold_mult": kwargs.get("cold_mult"),
"cold_workers": kwargs.get("cold_workers"),
"max_workers": kwargs.get("max_workers"),
"endpoint_name": kwargs.get("endpoint_name"),
"endpoint_state": kwargs.get("endpoint_state"),
"max_queue_time": kwargs.get("max_queue_time"),
"target_queue_time": kwargs.get("target_queue_time"),
"inactivity_timeout": kwargs.get("inactivity_timeout"),
"autoscaler_instance": kwargs.get("auto_instance", "prod"),
}
r = client.put(f"/endptjobs/{id}/", json_data=json_blob)
r.raise_for_status()
return r.json()
def delete_endpoint(client, id):
"""Delete an endpoint group.
DELETE /endptjobs/{id}/
Args:
client: VastClient instance.
id (int): ID of endpoint group to delete.
Returns:
dict: API response data.
"""
json_blob = {"client_id": "me", "endptjob_id": id}
r = client.delete(f"/endptjobs/{id}/", json_data=json_blob)
r.raise_for_status()
return r.json()
def _get_autoscaler_base_url(client):
"""Derive the autoscaler base URL (used for logs, worker status, and
update_workers; the autoscaler service is separate from the console API).
If the client is using the default console.vast.ai URL, the autoscaler
is at run.vast.ai. Otherwise the user's custom URL is used.
"""
from vastai.api.client import server_url_default
if client.server_url == server_url_default:
return "https://run.vast.ai"
return client.server_url
def get_endpt_logs(client, id, level=1, tail=None):
"""Fetch logs for a specific serverless endpoint group.
POST to <base>/get_endpoint_logs/
Args:
client: VastClient instance.
id (int): ID of endpoint group to fetch logs from.
level (int): Log detail level (0 to 3). Default 1.
0: info0, 1: info1, 2: trace, 3: debug
tail (int, optional): Number of tail lines.
Returns:
dict: Log data from the API response.
"""
base = _get_autoscaler_base_url(client)
url = base + "/get_endpoint_logs/"
json_blob = {"id": id, "api_key": client.api_key}
if tail is not None:
json_blob["tail"] = tail
headers = {}
if client.api_key is not None:
headers["Authorization"] = "Bearer " + client.api_key
r = client._request('POST', url, headers, json_blob)
r.raise_for_status()
if r.status_code == 200:
rj = r.json()
return rj
return {"error": r.text}
def show_workergroups(client):
"""Display user's current workergroups.
GET /autojobs/
Args:
client: VastClient instance.
Returns:
list: Workergroup results.
"""
json_blob = {"client_id": "me", "api_key": client.api_key}
r = client.get("/autojobs/", json_data=json_blob)
r.raise_for_status()
if r.status_code == 200:
rj = r.json()
if rj["success"]:
return rj["results"]
else:
return {"error": rj["msg"]}
return r.json()
def create_workergroup(client, **kwargs):
"""Create a new autoscale/workergroup.
POST /autojobs/
Args:
client: VastClient instance.
**kwargs: Workergroup configuration options:
template_hash (str): Template hash (required).
template_id (int): Template ID (optional).
search_params (str): Search param string for search offers.
launch_args (str): Launch args string for create instance.
endpoint_name (str): Deployment endpoint name.
endpoint_id (int): Deployment endpoint ID.
gpu_ram (float): Estimated GPU RAM requirement.
no_default (bool): Disable default search param query args.
auto_instance (str): Autoscaler instance type. Default "prod".
Returns:
dict: API response data.
"""
no_default = kwargs.get("no_default", False)
if no_default:
query = ""
else:
query = " verified=True rentable=True rented=False"
search_params_arg = kwargs.get("search_params")
search_params = ((search_params_arg if search_params_arg is not None else "") + query).strip()
json_blob = {
"client_id": "me",
"template_hash": kwargs.get("template_hash"),
"template_id": kwargs.get("template_id"),
"search_params": search_params,
"launch_args": kwargs.get("launch_args"),
"gpu_ram": kwargs.get("gpu_ram"),
"endpoint_name": kwargs.get("endpoint_name"),
"endpoint_id": kwargs.get("endpoint_id"),
"autoscaler_instance": kwargs.get("auto_instance", "prod"),
}
r = client.post("/autojobs/", json_data=json_blob)
r.raise_for_status()
return r.json()
def update_workergroup(client, id, **kwargs):
"""Update an existing workergroup.
PUT /autojobs/{id}/
Args:
client: VastClient instance.
id (int): ID of workergroup to update.
**kwargs: Workergroup configuration options.
Returns:
dict: API response data.
"""
no_default = kwargs.get("no_default", False)
if no_default:
query = ""
else:
query = " verified=True rentable=True rented=False"
search_params_arg = kwargs.get("search_params")
search_params = ((search_params_arg if search_params_arg is not None else "") + query).strip()
json_blob = {
"client_id": "me",
"autojob_id": id,
"template_hash": kwargs.get("template_hash"),
"template_id": kwargs.get("template_id"),
"search_params": search_params,
"launch_args": kwargs.get("launch_args"),
"gpu_ram": kwargs.get("gpu_ram"),
"endpoint_name": kwargs.get("endpoint_name"),
"endpoint_id": kwargs.get("endpoint_id"),
}
r = client.put(f"/autojobs/{id}/", json_data=json_blob)
r.raise_for_status()
return r.json()
def delete_workergroup(client, id):
"""Delete a workergroup.
DELETE /autojobs/{id}/
Note: Deleting a workergroup does not automatically destroy all
associated instances.
Args:
client: VastClient instance.
id (int): ID of workergroup to delete.
Returns:
dict: API response data.
"""
json_blob = {"client_id": "me", "autojob_id": id}
r = client.delete(f"/autojobs/{id}/", json_data=json_blob)
r.raise_for_status()
return r.json()
def get_wrkgrp_logs(client, id, level=1, tail=None):
"""Fetch logs for a specific serverless worker group.
POST to <base>/get_autogroup_logs/
Args:
client: VastClient instance.
id (int): ID of worker group to fetch logs from.
level (int): Log detail level (0 to 3). Default 1.
0: info0, 1: info1, 2: trace, 3: debug
tail (int, optional): Number of tail lines.
Returns:
dict: Log data from the API response.
"""
base = _get_autoscaler_base_url(client)
url = base + "/get_autogroup_logs/"
json_blob = {"id": id, "api_key": client.api_key}
if tail is not None:
json_blob["tail"] = tail
headers = {}
if client.api_key is not None:
headers["Authorization"] = "Bearer " + client.api_key
r = client._request('POST', url, headers, json_blob)
r.raise_for_status()
if r.status_code == 200:
rj = r.json()
return rj
return {"error": r.text}
def get_endpoint_workers(client, id):
"""List the live worker instances under a given endpoint, with runtime
status (creating/loading/idle/error/...) and ``measured_perf``.
Different from ``show_workergroups``: that one returns the user's
workergroup config records (template, search params, scaling policy, etc).
This returns the actual rented instances under one endpoint with their
current status.
POST to <base>/get_endpoint_workers/
Args:
client: VastClient instance.
id (int): ID of endpoint whose workers to list.
Returns:
list | dict: Worker data from the API response.
"""
base = _get_autoscaler_base_url(client)
url = base + "/get_endpoint_workers/"
json_blob = {"id": id}
headers = {}
if client.api_key is not None:
headers["Authorization"] = "Bearer " + client.api_key
r = client._request('POST', url, headers, json_blob)
r.raise_for_status()
return r.json()
def update_workers(client, id: int, cancel: bool = False):
"""Trigger a rolling update of all workers in a workergroup, or cancel an in-progress update.
POST /update_workers/
Args:
client: VastClient instance.
id: Workergroup ID.
cancel: If True, cancel an in-progress update.
Returns:
dict: Result from the autoscaler.
"""
from vastai.api.client import server_url_default
base_url = client.server_url
if base_url == server_url_default:
base_url = "https://run.vast.ai"
url = base_url + "/update_workers/"
json_blob = {"workergroup_id": id, "api_key": client.api_key}
if cancel:
json_blob["cancel_update"] = True
headers = {}
if client.api_key is not None:
headers["Authorization"] = "Bearer " + client.api_key
r = client._request('POST', url, headers, json_blob)
r.raise_for_status()
return r.json()
api/instances.py
"""Instance CRUD operations."""
import re
import time
import requests
from typing import Optional
from vastai.api.client import VastClient
from vastai.utils import parse_env
def _poll_result_url(result_url, retries=30, delay=0.3):
"""Poll a result URL until the content is ready. Total timeout ~9s."""
for _ in range(retries):
time.sleep(delay)
r = requests.get(result_url, timeout=10)
if r.status_code == 200:
return r.text
raise TimeoutError(f"Result not ready after {retries * delay}s: {result_url}")
def _strip_strings(value):
"""Recursively strip whitespace from string values."""
if isinstance(value, str):
return value.strip()
elif isinstance(value, dict):
return {k: _strip_strings(v) for k, v in value.items()}
elif isinstance(value, list):
return [_strip_strings(item) for item in value]
return value
def show_instances(client: VastClient, select_filters: Optional[dict] = None, order_by: Optional[list] = None) -> list:
"""Return all of the user's instances (optionally filtered/sorted) as a flat list.
Pages through the v1 ``/api/v1/instances/`` endpoint, following
``next_token`` until it is exhausted, and concatenates every page into one
flat list. ``select_cols`` is omitted on purpose: the backend then returns
full instance rows, matching the shape scripts and the SDK have always
depended on for this function's output.
"""
rows = []
params = {
"select_filters": select_filters or {},
"order_by": order_by or [{"col": "id", "dir": "asc"}],
"limit": 25,
}
while True:
data = show_instances_v1(client, params)
page = data.get("instances") or []
for row in page:
row = {k: _strip_strings(v) for k, v in row.items()}
row['duration'] = time.time() - row['start_date']
row['extra_env'] = {env_var[0]: env_var[1] for env_var in row['extra_env']}
rows.append(row)
next_token = data.get("next_token")
if not next_token:
break
params["after_token"] = next_token
return rows
def show_instances_v1(client: VastClient, params: dict) -> dict:
"""Fetch instances using the v1 paginated API.
Args:
client: VastClient instance.
params: Dict with select_filters, order_by, limit, after_token, select_cols.
Returns:
Full response dict (instances, next_token, total_instances, label_counts).
"""
r = client.get("/api/v1/instances/", query_args=params)
r.raise_for_status()
return r.json()
def show_instance_filters(client: VastClient) -> list:
"""Fetch distinct filterable values for instances."""
r = client.get("/instances/filters/")
r.raise_for_status()
return r.json().get("filters", [])
def show_instance(client: VastClient, id: int) -> Optional[dict]:
r = client.get(f"/instances/{id}/", query_args={"owner": "me"})
r.raise_for_status()
row = r.json()["instances"]
if row is None:
return None
row['duration'] = time.time() - row['start_date']
row['extra_env'] = {env_var[0]: env_var[1] for env_var in row['extra_env']}
return row
VALID_MOUNT_PATH = re.compile(r'^(/)?([^/\0]+(/)?)+$')
def resolve_runtype(runtype: Optional[str] = None, *, ssh: bool = False,
jupyter: bool = False, direct: bool = False,
args=None, jupyter_lab: bool = False,
jupyter_dir: Optional[str] = None,
default: Optional[str] = None) -> tuple:
"""Translate the friendly connection flags into a runtype string.
Returns ``(runtype, args)``. ``args`` comes back because an empty ``--args``
selects the args runtype and then has to be dropped from the payload.
``default`` is what an unflagged call resolves to: the CLI passes ``'ssh'``,
the SDK passes ``None`` so the key stays out of the payload entirely.
"""
if args in ('', [''], []):
args = None
empty_args = True
else:
empty_args = False
if runtype:
if ssh or jupyter or direct:
raise ValueError(
"Pass either runtype= or the ssh/jupyter/direct flags, not both."
)
return runtype, args
resolved = default
if args or empty_args:
resolved = 'args'
if not jupyter and (jupyter_dir or jupyter_lab):
jupyter = True
if jupyter and resolved == 'args':
raise ValueError(
"Can't use jupyter and args together. "
"Try onstart_cmd instead of args."
)
if jupyter:
resolved = 'jupyter_direc ssh_direc ssh_proxy' if direct else 'jupyter_proxy ssh_proxy'
elif ssh:
resolved = 'ssh_direc ssh_proxy' if direct else 'ssh_proxy'
elif direct and default is None:
# The CLI reaches here with default='ssh' and keeps it, which is the
# long-standing behaviour of a bare --direct. There is no such fallback
# for the SDK, so say so instead of dropping the flag.
raise ValueError("direct=True requires ssh=True or jupyter=True.")
return resolved, args
def build_volume_info(*, create_volume: Optional[int] = None,
link_volume: Optional[int] = None,
volume_size: Optional[float] = None,
mount_path: Optional[str] = None,
volume_label: Optional[str] = None) -> Optional[dict]:
"""Build the volume_info blob, or None when no volume was requested."""
if volume_size is not None and not create_volume:
raise ValueError("volume_size can only be used with create_volume.")
if not (create_volume or link_volume):
if mount_path or volume_label:
raise ValueError("mount_path and volume_label need create_volume or link_volume.")
return None
if not mount_path:
raise ValueError("mount_path is required when creating or linking a volume.")
if not VALID_MOUNT_PATH.match(mount_path):
raise ValueError(f"mount_path '{mount_path}' is not a valid Linux file path.")
volume_info = {
"mount_path": mount_path,
"create_new": True if create_volume else False,
"volume_id": create_volume if create_volume else link_volume,
}
if volume_label:
volume_info["name"] = volume_label
if volume_size:
volume_info["size"] = volume_size
elif create_volume:
volume_info["size"] = 15
return volume_info
def apply_portal_config(env: dict, runtype: Optional[str]) -> dict:
"""Return env with jupyter portal entries stripped on non-jupyter runtypes."""
if runtype and 'jupyter' in runtype:
return env
if "PORTAL_CONFIG" not in env:
return env
filtered = [c for c in env["PORTAL_CONFIG"].split("|") if 'jupyter' not in c.lower()]
if not filtered:
raise ValueError(
"env variable PORTAL_CONFIG must contain at least one non-jupyter "
"related config string if runtype is not jupyter"
)
env = dict(env)
env["PORTAL_CONFIG"] = "|".join(filtered)
return env
def build_create_instance_payload(
*, image: Optional[str] = None, disk: float = 10, env=None,
price: Optional[float] = None, bid_price: Optional[float] = None,
label: Optional[str] = None, extra: Optional[str] = None,
onstart_cmd: Optional[str] = None, login: Optional[str] = None,
python_utf8: bool = False, lang_utf8: bool = False,
jupyter_lab: bool = False, jupyter_dir: Optional[str] = None,
force: bool = False, cancel_unavail: bool = False,
template_hash: Optional[str] = None, user: Optional[str] = None,
runtype: Optional[str] = None, ssh: bool = False, jupyter: bool = False,
direct: bool = False, args=None, volume_info: Optional[dict] = None,
create_volume: Optional[int] = None, link_volume: Optional[int] = None,
volume_size: Optional[float] = None, mount_path: Optional[str] = None,
volume_label: Optional[str] = None,
) -> dict:
"""Build the JSON body for a create-instance request.
The one place the CLI, the SDK and the docs generator agree on what a
create-instance call takes. Pure: no client, no argparse, no printing.
"""
if price is not None and bid_price is not None:
raise ValueError("Pass either price= or bid_price=, not both.")
if price is None:
price = bid_price
if isinstance(env, str):
env = parse_env(env)
env = dict(env or {})
if template_hash is None:
runtype, args = resolve_runtype(
runtype, ssh=ssh, jupyter=jupyter, direct=direct, args=args,
jupyter_lab=jupyter_lab, jupyter_dir=jupyter_dir,
)
env = apply_portal_config(env, runtype)
if volume_info is None:
volume_info = build_volume_info(
create_volume=create_volume, link_volume=link_volume,
volume_size=volume_size, mount_path=mount_path,
volume_label=volume_label,
)
elif create_volume or link_volume:
raise ValueError("Pass either volume_info= or the create/link volume flags, not both.")
json_blob = {
"client_id": "me",
"image": image,
"env": env,
"price": price,
"disk": disk,
"label": label,
"extra": extra,
"onstart": onstart_cmd,
"image_login": login,
"python_utf8": python_utf8,
"lang_utf8": lang_utf8,
"use_jupyter_lab": jupyter_lab,
"jupyter_dir": jupyter_dir,
"force": force,
"cancel_unavail": cancel_unavail,
"template_hash_id": template_hash,
"user": user,
}
if runtype:
json_blob["runtype"] = runtype
if args is not None:
json_blob["args"] = args
if volume_info:
json_blob["volume_info"] = volume_info
return json_blob
def create_instance(client: VastClient, id, image=None, disk=10, env=None,
price=None, label=None, extra=None, onstart_cmd=None,
login=None, python_utf8=False, lang_utf8=False,
jupyter_lab=False, jupyter_dir=None, force=False,
cancel_unavail=False, template_hash=None, user=None,
runtype=None, args=None, volume_info=None,
# Appended, never inserted: the parameters above keep their
# original positions so existing positional calls still bind
# the same values.
bid_price=None, ssh=False, jupyter=False, direct=False,
create_volume=None, link_volume=None, volume_size=None,
mount_path=None, volume_label=None) -> dict:
"""Create one instance, or several when ``id`` is a list of offer ids."""
json_blob = build_create_instance_payload(
image=image, disk=disk, env=env, price=price, bid_price=bid_price,
label=label, extra=extra, onstart_cmd=onstart_cmd, login=login,
python_utf8=python_utf8, lang_utf8=lang_utf8, jupyter_lab=jupyter_lab,
jupyter_dir=jupyter_dir, force=force, cancel_unavail=cancel_unavail,
template_hash=template_hash, user=user, runtype=runtype, ssh=ssh,
jupyter=jupyter, direct=direct, args=args, volume_info=volume_info,
create_volume=create_volume, link_volume=link_volume,
volume_size=volume_size, mount_path=mount_path, volume_label=volume_label,
)
return create_instance_from_payload(client, id, json_blob)
def create_instance_from_payload(client: VastClient, id, json_blob: dict) -> dict:
"""Send an already-built create-instance payload for one or many offers."""
if isinstance(id, list):
json_blob["ids"] = id
r = client.post("/asks/bulk/", json_data=json_blob)
else:
r = client.put(f"/asks/{id}/", json_data=json_blob)
r.raise_for_status()
return r.json()
def destroy_instance(client: VastClient, id) -> dict:
json_blob = {}
if isinstance(id, list):
json_blob["instance_ids"] = id
r = client.delete("/instances/", json_data=json_blob)
else:
r = client.delete(f"/instances/{id}/", json_data=json_blob)
r.raise_for_status()
return r.json()
def start_instance(client: VastClient, id) -> dict:
json_blob = {"state": "running"}
if isinstance(id, list):
json_blob["ids"] = id
r = client.put("/instances/", json_data=json_blob)
else:
r = client.put(f"/instances/{id}/", json_data=json_blob)
r.raise_for_status()
return r.json()
def stop_instance(client: VastClient, id) -> dict:
json_blob = {"state": "stopped"}
if isinstance(id, list):
json_blob["ids"] = id
r = client.put("/instances/", json_data=json_blob)
else:
r = client.put(f"/instances/{id}/", json_data=json_blob)
r.raise_for_status()
return r.json()
def reboot_instance(client: VastClient, id: int) -> dict:
r = client.put(f"/instances/reboot/{id}/", json_data={})
r.raise_for_status()
return r.json()
def recycle_instance(client: VastClient, id: int) -> dict:
r = client.put(f"/instances/recycle/{id}/", json_data={})
r.raise_for_status()
return r.json()
def label_instance(client: VastClient, id: int, label: str) -> dict:
r = client.put(f"/instances/{id}/", json_data={"label": label})
r.raise_for_status()
return r.json()
def prepay_instance(client: VastClient, id: int, amount: float) -> dict:
r = client.put(f"/instances/prepay/{id}/", json_data={"amount": amount})
r.raise_for_status()
return r.json()
def change_bid(client: VastClient, id: int, price: float = None) -> dict:
r = client.put(f"/instances/bid_price/{id}/", json_data={"client_id": "me", "price": price})
r.raise_for_status()
return r.json()
def accept_price_increase(client: VastClient, id: int = None,
instance_ids=None, host_id: int = None) -> dict:
"""Deprecated shim that resolves to the per-row endpoint.
Prefer :func:`vastai.api.price_increase.accept`. The old
``id`` / ``instance_ids`` / ``host_id`` selectors are no longer
supported by the backend; ``id`` here is reinterpreted as an
instance id, looked up against the pending list, and forwarded
to the per-row accept. ``instance_ids`` and ``host_id`` raise.
Removed in the release after this one.
"""
from vastai.api import price_increase as _pi
if instance_ids is not None or host_id is not None:
raise TypeError(
"accept_price_increase: instance_ids and host_id are no longer "
"supported. Use vastai.api.price_increase.accept(client, pending_id) "
"per row, or vastai.sdk.VastAI.accept_price_increase(instance_id=…).")
if id is None:
raise TypeError("accept_price_increase: instance id is required")
try:
match = _pi.resolve_instance_to_pending(client, id)
except LookupError as err:
raise LookupError(
f"accept_price_increase: no pending price increase for instance {id}"
) from err
return _pi.accept(client, match["pending_price_increase_id"])
def execute(client: VastClient, id: int, command: str):
"""Execute a command on an instance and return the output."""
r = client.put(f"/instances/command/{id}/", json_data={"command": command})
r.raise_for_status()
rj = r.json()
result_url = rj.get("result_url")
if not result_url:
return rj
return _poll_result_url(result_url)
def logs(client: VastClient, instance_id: int, tail=None, filter=None, daemon_logs=False):
"""Request logs for an instance and return the log text."""
json_blob = {}
if filter:
json_blob['filter'] = filter
if tail:
json_blob['tail'] = tail
if daemon_logs:
json_blob['daemon_logs'] = 'true'
r = client.put(f"/instances/request_logs/{instance_id}/", json_data=json_blob)
r.raise_for_status()
rj = r.json()
result_url = rj.get("result_url")
if not result_url:
return rj
return _poll_result_url(result_url)
def update_instance(client: VastClient, id: int, template_id=None, template_hash_id=None,
image=None, args=None, env=None, onstart=None) -> dict:
json_blob = {"id": id}
if template_id is not None:
json_blob["template_id"] = template_id
if template_hash_id is not None:
json_blob["template_hash_id"] = template_hash_id
if image is not None:
json_blob["image"] = image
if args is not None:
json_blob["args"] = args
if env is not None:
json_blob["env"] = env
if onstart is not None:
json_blob["onstart"] = onstart
r = client.put(f"/instances/update_template/{id}/", json_data=json_blob)
r.raise_for_status()
return r.json()
def take_snapshot(client: VastClient, instance_id, repo=None, container_registry="docker.io",
docker_login_user=None, docker_login_pass=None, pause="true") -> dict:
req_json = {
"id": instance_id,
"container_registry": container_registry,
"personal_repo": repo,
"docker_login_user": docker_login_user,
"docker_login_pass": docker_login_pass,
"pause": pause
}
r = client.post(f"/instances/take_snapshot/{instance_id}/", json_data=req_json)
r.raise_for_status()
return r.json()
api/keys.py
"""SSH key and API key operations."""
from vastai.api.client import VastClient
def create_ssh_key(client: VastClient, ssh_key: str) -> dict:
"""Add an SSH public key to the account.
Args:
client: VastClient instance.
ssh_key: SSH public key content string.
Returns:
Response dict with created key info.
"""
r = client.post("/ssh/", json_data={"ssh_key": ssh_key})
r.raise_for_status()
return r.json()
def show_ssh_keys(client: VastClient) -> dict:
"""List SSH keys associated with the account.
Args:
client: VastClient instance.
Returns:
Response dict with SSH key info.
"""
r = client.get("/ssh/")
r.raise_for_status()
return r.json()
def update_ssh_key(client: VastClient, id: int, ssh_key: str) -> dict:
"""Update an existing SSH key.
Args:
client: VastClient instance.
id: SSH key ID.
ssh_key: New SSH public key content string.
Returns:
Response dict.
"""
payload = {
"id": id,
"ssh_key": ssh_key,
}
r = client.put(f"/ssh/{id}/", json_data=payload)
r.raise_for_status()
return r.json()
def delete_ssh_key(client: VastClient, id: int) -> dict:
"""Delete an SSH key from the account.
Args:
client: VastClient instance.
id: SSH key ID to delete.
Returns:
Response dict.
"""
r = client.delete(f"/ssh/{id}/")
r.raise_for_status()
return r.json()
def attach_ssh(client: VastClient, instance_id: int, ssh_key: str) -> dict:
"""Attach an SSH key to an instance.
Args:
client: VastClient instance.
instance_id: Instance ID to attach the key to.
ssh_key: SSH public key content string.
Returns:
Response dict.
"""
r = client.post(f"/instances/{instance_id}/ssh/", json_data={"ssh_key": ssh_key})
r.raise_for_status()
return r.json()
def detach_ssh(client: VastClient, instance_id: int, ssh_key_id: str) -> dict:
"""Detach an SSH key from an instance.
Args:
client: VastClient instance.
instance_id: Instance ID.
ssh_key_id: SSH key ID to detach.
Returns:
Response dict.
"""
r = client.delete(f"/instances/{instance_id}/ssh/{ssh_key_id}/")
r.raise_for_status()
return r.json()
def create_api_key(client: VastClient, name: str, permissions: dict = None,
key_params: str = None) -> dict:
"""Create a new API key.
Args:
client: VastClient instance.
name: Name for the new API key.
permissions: Dict of permissions for the key.
key_params: Optional key parameters string.
Returns:
Response dict with created API key info.
"""
json_blob = {"name": name}
if permissions is not None:
json_blob["permissions"] = permissions
if key_params is not None:
json_blob["key_params"] = key_params
r = client.post("/auth/apikeys/", json_data=json_blob)
r.raise_for_status()
return r.json()
def show_api_key(client: VastClient, id: int) -> dict:
"""Show details of a specific API key.
Args:
client: VastClient instance.
id: API key ID.
Returns:
API key details dict.
"""
r = client.get(f"/auth/apikeys/{id}/")
r.raise_for_status()
return r.json()
def show_api_keys(client: VastClient) -> dict:
"""List all API keys associated with the account.
Args:
client: VastClient instance.
Returns:
Envelope dict of the form ``{"apikeys": [...]}`` (what the backend
sends). The high-level ``VastAI.show_api_keys`` wrapper unwraps this
to a plain list.
"""
r = client.get("/auth/apikeys/")
r.raise_for_status()
return r.json()
def delete_api_key(client: VastClient, id: int) -> dict:
"""Delete an API key.
Args:
client: VastClient instance.
id: API key ID to delete.
Returns:
Response dict.
"""
r = client.delete(f"/auth/apikeys/{id}/")
r.raise_for_status()
return r.json()
def reset_api_key(client: VastClient) -> dict:
"""Reset the current API key (generates a new one).
Args:
client: VastClient instance.
Returns:
Response dict.
"""
json_blob = {"client_id": "me"}
r = client.put("/commands/reset_apikey/", json_data=json_blob)
r.raise_for_status()
return r.json()
api/machines.py
"""Machine management operations for hosts."""
from vastai.api.client import VastClient
def show_machine(client: VastClient, id: int) -> list:
"""Show a single hosted machine.
Args:
client: VastClient instance.
id: Machine ID.
Returns:
List of machine data dicts (API returns list even for single machine).
"""
r = client.get(f"/machines/{id}", query_args={"owner": "me"})
r.raise_for_status()
return r.json()
def show_machines(client: VastClient) -> list:
"""Show all hosted machines for the current user.
Args:
client: VastClient instance.
Returns:
List of machine dicts.
"""
r = client.get("/machines", query_args={"owner": "me"})
r.raise_for_status()
return r.json()["machines"]
def show_maints(client: VastClient, machine_ids: list) -> list:
"""Show maintenance information for host machines.
Args:
client: VastClient instance.
machine_ids: List of machine ID integers.
Returns:
List of maintenance info dicts.
"""
r = client.get("/machines/maintenances", query_args={"owner": "me", "machine_ids": machine_ids})
r.raise_for_status()
return r.json()
def list_machine(client: VastClient, id: int, price_gpu: float = None,
price_disk: float = None, price_inetu: float = None,
price_inetd: float = None, price_min_bid: float = None,
min_chunk: int = None, end_date: float = None,
discount_rate: float = None, duration: str = None,
vol_size: int = None, vol_price: float = None) -> dict:
"""List a machine for rent (create offers).
Args:
client: VastClient instance.
id: Machine ID.
price_gpu: Per GPU rental price in $/hour.
price_disk: Storage price in $/GB/month.
price_inetu: Price for upload bandwidth in $/GB.
price_inetd: Price for download bandwidth in $/GB.
price_min_bid: Per GPU minimum bid price floor in $/hour.
min_chunk: Minimum number of GPUs.
end_date: Contract offer expiration as unix epoch timestamp.
discount_rate: Max long-term prepay discount rate fraction.
duration: Duration string or seconds.
vol_size: Volume contract offer size in GB.
vol_price: Volume disk price.
Returns:
Response dict.
"""
json_blob = {
"machine": id,
"price_gpu": price_gpu,
"price_disk": price_disk,
"price_inetu": price_inetu,
"price_inetd": price_inetd,
"price_min_bid": price_min_bid,
"min_chunk": min_chunk,
"end_date": end_date,
"credit_discount_max": discount_rate,
"duration": duration,
"vol_size": vol_size,
"vol_price": vol_price,
}
r = client.put("/machines/create_asks/", json_data=json_blob)
r.raise_for_status()
return r.json()
def unlist_machine(client: VastClient, id: int) -> dict:
"""Unlist a listed machine (remove all offers).
Args:
client: VastClient instance.
id: Machine ID.
Returns:
Response dict.
"""
r = client.delete(f"/machines/{id}/asks/")
r.raise_for_status()
return r.json()
def cancel_maint(client: VastClient, id: int) -> dict:
"""Cancel scheduled maintenance window(s) for a machine.
Args:
client: VastClient instance.
id: Machine ID.
Returns:
Response dict.
"""
json_blob = {"client_id": "me", "machine_id": id}
r = client.put(f"/machines/{id}/cancel_maint/", json_data=json_blob)
r.raise_for_status()
return r.json()
def schedule_maint(client: VastClient, id: int, sdate: float = None,
duration: float = None, maintenance_category: str = None) -> dict:
"""Schedule a maintenance window for a machine.
Args:
client: VastClient instance.
id: Machine ID.
sdate: Start date as unix epoch timestamp.
duration: Duration in hours.
maintenance_category: Category of maintenance.
Returns:
Response dict.
"""
json_blob = {
"client_id": "me",
"sdate": sdate,
"duration": duration,
"maintenance_category": maintenance_category,
}
r = client.put(f"/machines/{id}/dnotify/", json_data=json_blob)
r.raise_for_status()
return r.json()
def cleanup_machine(client: VastClient, id: int) -> dict:
"""Remove all expired storage instances from a machine.
Args:
client: VastClient instance.
id: Machine ID.
Returns:
Response dict.
"""
r = client.put(f"/machines/{id}/cleanup/", json_data={})
r.raise_for_status()
return r.json()
def defrag_machines(client: VastClient, machine_ids: list) -> dict:
"""Defragment machines to rearrange GPU assignments.
Args:
client: VastClient instance.
machine_ids: List of machine ID integers.
Returns:
Response dict with defragmentation results.
"""
json_blob = {"machine_ids": machine_ids}
r = client.put("/machines/defrag_offers/", json_data=json_blob)
r.raise_for_status()
return r.json()
def delete_machine(client: VastClient, id: int) -> dict:
"""Force delete a machine if not in use by clients.
Args:
client: VastClient instance.
id: Machine ID.
Returns:
Response dict.
"""
r = client.post(f"/machines/{id}/force_delete/")
r.raise_for_status()
return r.json()
def reports(client: VastClient, id: int) -> dict:
"""Get user reports for a given machine.
Args:
client: VastClient instance.
id: Machine ID.
Returns:
Reports dict.
"""
json_blob = {"machine_id": id}
r = client.get(f"/machines/{id}/reports/", json_data=json_blob)
r.raise_for_status()
return r.json()
def set_defjob(client: VastClient, id: int, price_gpu: float = None,
price_inetu: float = None, price_inetd: float = None,
image: str = None, args: list = None) -> dict:
"""Create default jobs for a machine.
Args:
client: VastClient instance.
id: Machine ID.
price_gpu: Per GPU rental price in $/hour.
price_inetu: Price for upload bandwidth in $/GB.
price_inetd: Price for download bandwidth in $/GB.
image: Docker container image to launch.
args: List of arguments passed to container launch.
Returns:
Response dict.
"""
json_blob = {
"machine": id,
"price_gpu": price_gpu,
"price_inetu": price_inetu,
"price_inetd": price_inetd,
"image": image,
"args": args,
}
r = client.put("/machines/create_bids/", json_data=json_blob)
r.raise_for_status()
return r.json()
def remove_defjob(client: VastClient, id: int) -> dict:
"""Remove default jobs from a machine.
Args:
client: VastClient instance.
id: Machine ID.
Returns:
Response dict.
"""
r = client.delete(f"/machines/{id}/defjob/")
r.raise_for_status()
return r.json()
def set_min_bid(client: VastClient, id: int, price: float = None) -> dict:
"""Set the minimum bid/rental price for a machine.
Args:
client: VastClient instance.
id: Machine ID.
price: Per GPU min bid price in $/hour.
Returns:
Response dict.
"""
json_blob = {"client_id": "me", "price": price}
r = client.put(f"/machines/{id}/minbid/", json_data=json_blob)
r.raise_for_status()
return r.json()
api/metrics.py
"""Platform-wide GPU market metrics (available to admins and hosts)."""
from typing import Optional
from vastai.api.client import VastClient
def gpu_current(client: VastClient, verified: str = "all", hosting_type: str = "all",
num_gpus: str = "all") -> dict:
"""Current snapshot of GPU supply/demand, pricing, and perf per GPU type.
Args:
client: VastClient instance.
verified: "yes", "no", or "all".
hosting_type: "all", "secure_cloud", or "community".
num_gpus: machine GPU-count bucket — "all", "1", "2", "4", or "8".
Returns:
Parsed JSON response. Shape: {"success": True, "gpus": [...]} or
{"success": True, "gpus": [], "needs_machine": True} for hosts with no machines.
"""
r = client.get("/metrics/gpu/current/", query_args={"verified": verified, "hosting_type": hosting_type, "num_gpus": num_gpus})
r.raise_for_status()
return r.json()
def gpu_history(client: VastClient, gpu_name: str, verified: str = "all", hosting_type: str = "all",
start: Optional[int] = None, end: Optional[int] = None, step: Optional[int] = None,
num_gpus: str = "all") -> dict:
"""Time-series supply/demand, pricing, and stats per GPU type.
Args:
client: VastClient instance.
gpu_name: GPU name, comma-separated list, or "all".
verified: "yes", "no", or "all".
hosting_type: "all", "secure_cloud", or "community".
start: Start unix timestamp (defaults to end - 1 day server-side).
end: End unix timestamp (defaults to now server-side).
step: Step in seconds between data points.
num_gpus: machine GPU-count bucket — "all", "1", "2", "4", or "8".
Returns:
Parsed JSON response. Shape: {"success": True, "gpus": {gpu_name: {supply_demand, pricing, stats}}}
or {"success": True, "gpus": {}, "needs_machine": True}.
"""
params = {"gpu_name": gpu_name, "verified": verified, "hosting_type": hosting_type, "num_gpus": num_gpus}
if start is not None:
params["start"] = str(start)
if end is not None:
params["end"] = str(end)
if step is not None:
params["step"] = str(step)
r = client.get("/metrics/gpu/history/", query_args=params)
r.raise_for_status()
return r.json()
def gpu_locations(client: VastClient) -> dict:
"""Geographic locations of all GPUs on the platform.
Returns:
Parsed JSON response. Shape: {"success": True, "locations": [...]}
or {"success": True, "locations": [], "needs_machine": True}.
"""
r = client.get("/metrics/gpu/locations/")
r.raise_for_status()
return r.json()
api/offers.py
"""Search offers, templates, benchmarks, volumes, network volumes, and invoices."""
from typing import Optional
from vastai.api.billing import show_user
from vastai.api.client import VastClient
from vastai.api.instances import resolve_runtype
from vastai.api.query import (parse_order, parse_query, offers_alias,
offers_fields, offers_mult)
from vastai.utils import parse_env
def search_offers(client: VastClient, query: dict = None, offer_type: str = "on-demand",
order: list = None, limit: int = None, storage: float = 5.0,
no_default: bool = False, disable_bundling: bool = False) -> list:
"""Search for instance offers using a query dict.
Args:
client: VastClient instance.
query: Pre-parsed query dict of filters (e.g. {"gpu_name": {"eq": "RTX 3090"}}).
offer_type: One of "on-demand", "reserved", or "bid".
order: List of [field, direction] pairs, e.g. [["score", "desc"]].
limit: Max number of results.
storage: Allocated storage in GiB for pricing (default 5.0).
no_default: If True, skip default filters.
disable_bundling: Deprecated bundling flag.
Returns:
List of offer dicts.
"""
if no_default:
q = query or {}
else:
q = {"verified": {"eq": True}, "external": {"eq": False},
"rentable": {"eq": True}, "rented": {"eq": False}}
if query:
q.update(query)
if order is not None:
q["order"] = order
else:
q["order"] = [["score", "desc"]]
q["type"] = offer_type
if offer_type == "interruptible":
q["type"] = "bid"
if limit:
q["limit"] = int(limit)
q["allocated_storage"] = storage
if disable_bundling:
q["disable_bundling"] = True
r = client.post("/bundles/", json_data=q)
r.raise_for_status()
return r.json()["offers"]
def search_offers_new(client: VastClient, query: dict = None, offer_type: str = "on-demand",
order: list = None, limit: int = None, storage: float = 5.0,
no_default: bool = False, disable_bundling: bool = False) -> list:
"""Search for instance offers using the new /search/asks/ endpoint.
Args:
client: VastClient instance.
query: Pre-parsed query dict of filters (e.g. {"gpu_name": {"eq": "RTX 3090"}}).
offer_type: One of "on-demand", "reserved", or "bid".
order: List of [field, direction] pairs, e.g. [["score", "desc"]].
limit: Max number of results.
storage: Allocated storage in GiB for pricing (default 5.0).
no_default: If True, skip default filters.
disable_bundling: Deprecated bundling flag.
Returns:
List of offer dicts.
"""
if no_default:
q = query or {}
else:
q = {"verified": {"eq": True}, "external": {"eq": False},
"rentable": {"eq": True}, "rented": {"eq": False}}
if query:
q.update(query)
if order is not None:
q["order"] = order
else:
q["order"] = [["score", "desc"]]
q["type"] = offer_type
if offer_type == "interruptible":
q["type"] = "bid"
if limit:
q["limit"] = int(limit)
q["allocated_storage"] = storage
if disable_bundling:
q["disable_bundling"] = True
json_blob = {"select_cols": ["*"], "q": q}
r = client.put("/search/asks/", json_data=json_blob)
r.raise_for_status()
return r.json()["offers"]
def search_templates(client: VastClient, query: dict = None) -> list:
"""Search for templates using a query dict.
Args:
client: VastClient instance.
query: Pre-parsed query dict of select_filters.
Returns:
List of template dicts.
"""
query_args = {"select_cols": ["*"], "select_filters": query or {}}
r = client.get("/template/", query_args=query_args)
r.raise_for_status()
return r.json().get("templates", [])
def benchmarks_query_args(query: dict = None, order: list = None,
limit: int = None, after_token: str = None) -> dict:
"""Build the query args for a ``/benchmarks`` request."""
query_args = {"select_cols": ["*"], "select_filters": query or {}}
if order is not None:
query_args["order_by"] = order
if limit is not None:
query_args["limit"] = int(limit)
if after_token:
query_args["after_token"] = after_token
return query_args
def search_benchmarks(client: VastClient, query: dict = None, order: list = None,
limit: int = None, after_token: str = None) -> list:
"""Search for benchmarks using a query dict, as a flat list.
Pages through ``/benchmarks``, following ``next_token`` until it is
exhausted, and concatenates every page. The backend caps a page at 200 rows
by default, so callers that want one page at a time should use
:func:`search_benchmarks_v1` instead.
Args:
client: VastClient instance.
query: Pre-parsed query dict of select_filters.
order: List of {"col": ..., "dir": ...} dicts, e.g. [{"col": "last_update", "dir": "desc"}].
limit: Max number of results per page.
after_token: Pagination token to resume from.
Returns:
List of benchmark dicts.
"""
rows = []
params = benchmarks_query_args(query=query, order=order, limit=limit,
after_token=after_token)
while True:
data = search_benchmarks_v1(client, params)
rows.extend(data.get("benchmarks") or [])
next_token = data.get("next_token")
if not next_token:
return rows
params["after_token"] = next_token
def search_benchmarks_v1(client: VastClient, params: dict) -> dict:
"""Fetch one page of benchmarks using the paginated API.
Args:
client: VastClient instance.
params: Dict with select_cols, select_filters, order_by, limit, after_token.
Returns:
Full response dict (benchmarks, benchmarks_found, next_token).
"""
r = client.get("/benchmarks", query_args=params)
r.raise_for_status()
return r.json()
def search_volumes(client: VastClient, query: dict = None, order: list = None,
limit: int = None, storage: float = 1.0,
no_default: bool = False) -> list:
"""Search for volume offers.
Args:
client: VastClient instance.
query: Pre-parsed query dict of filters.
order: List of [field, direction] pairs.
limit: Max number of results.
storage: Allocated storage in GiB for pricing (default 1.0).
no_default: If True, skip default filters.
Returns:
List of volume offer dicts.
"""
if no_default:
q = query or {}
else:
q = {"verified": {"eq": True}, "external": {"eq": False}, "disk_space": {"gte": 1}}
if query:
q.update(query)
if order is not None:
q["order"] = order
else:
q["order"] = [["score", "desc"]]
if limit:
q["limit"] = int(limit)
q["allocated_storage"] = storage
r = client.post("/volumes/search/", json_data=q)
r.raise_for_status()
return r.json()["offers"]
def search_network_volumes(client: VastClient, query: dict = None, order: list = None,
limit: int = None, storage: float = 1.0,
no_default: bool = False) -> list:
"""Search for network volume offers.
Args:
client: VastClient instance.
query: Pre-parsed query dict of filters.
order: List of [field, direction] pairs.
limit: Max number of results.
storage: Allocated storage in GiB for pricing (default 1.0).
no_default: If True, skip default filters.
Returns:
List of network volume offer dicts.
"""
if no_default:
q = query or {}
else:
q = {"verified": {"eq": True}, "external": {"eq": False}, "disk_space": {"gte": 1}}
if query:
q.update(query)
if order is not None:
q["order"] = order
else:
q["order"] = [["score", "desc"]]
if limit:
q["limit"] = int(limit)
q["allocated_storage"] = storage
r = client.post("/network_volumes/search/", json_data=q)
r.raise_for_status()
return r.json()["offers"]
def search_invoices(client: VastClient, query: dict = None) -> list:
"""Search for invoices using a query dict.
Args:
client: VastClient instance.
query: Pre-parsed query dict of select_filters.
Returns:
List of invoice dicts.
"""
query_args = {"select_cols": ["*"], "select_filters": query or {}}
r = client.get("/invoices", query_args=query_args)
r.raise_for_status()
return r.json()
def template_fields_from_flags(*, ssh: bool = False, jupyter: bool = False,
direct: bool = False, jupyter_lab: bool = False,
login: Optional[str] = None,
hide_readme: bool = False, public: bool = False,
search_params: Optional[str] = None,
no_default: bool = False,
always_default: bool = False) -> dict:
"""Translate the friendly template flags into api-layer template fields.
Shared by the CLI commands and the SDK so both publish the same surface.
``always_default`` seeds the default offer filters even when no
``search_params`` were given, which is what the CLI has always done. The SDK
passes False: it used to send no filters at all on a template created
without search params, and quietly starting to attach three would change
templates that existing scripts create.
"""
default_search_query = {}
if not no_default and (always_default or search_params is not None):
default_search_query = {"verified": {"eq": True}, "external": {"eq": False},
"rentable": {"eq": True}}
return {
"jup_direct": jupyter and direct,
"ssh_direct": ssh and direct,
"use_ssh": ssh or jupyter,
"use_jupyter_lab": jupyter_lab,
"runtype": "jupyter" if jupyter else ("ssh" if ssh else "args"),
"docker_login_repo": login.split(" ")[0] if login else None,
"extra_filters": parse_query(search_params, default_search_query,
offers_fields, offers_alias, offers_mult),
"readme_visible": not hide_readme,
"private": not public,
}
def create_template(client: VastClient, name: str = None, image: str = None,
image_tag: str = None, href: str = None, repo: str = None,
env: str = None, onstart_cmd: str = None,
jup_direct: bool = False, ssh_direct: bool = False,
use_jupyter_lab: bool = False, runtype: str = "args",
use_ssh: bool = False, jupyter_dir: str = None,
docker_login_repo: str = None, extra_filters: dict = None,
disk_space: float = None, readme: str = None,
readme_visible: bool = True, desc: str = None,
private: bool = True) -> dict:
"""Create a new template.
Args:
client: VastClient instance.
name: Template name.
image: Docker image.
image_tag: Docker image tag.
href: Link to provide.
repo: Link to repository.
env: Docker options env string.
onstart_cmd: Onstart script contents.
jup_direct: Supports jupyter direct.
ssh_direct: Supports ssh direct.
use_jupyter_lab: Launch with jupyter lab.
runtype: Run type (jupyter, ssh, args).
use_ssh: Supports ssh.
jupyter_dir: Jupyter directory.
docker_login_repo: Docker login repository.
extra_filters: Search offer filters dict.
disk_space: Recommended disk space.
readme: Readme string.
readme_visible: Whether readme is visible.
desc: Description string.
private: Whether template is private.
Returns:
Response dict with template info.
"""
template = {
"name": name,
"image": image,
"tag": image_tag,
"href": href,
"repo": repo,
"env": env,
"onstart": onstart_cmd,
"jup_direct": jup_direct,
"ssh_direct": ssh_direct,
"use_jupyter_lab": use_jupyter_lab,
"runtype": runtype,
"use_ssh": use_ssh,
"jupyter_dir": jupyter_dir,
"docker_login_repo": docker_login_repo,
"extra_filters": extra_filters or {},
"recommended_disk_space": disk_space,
"readme": readme,
"readme_visible": readme_visible,
"desc": desc,
"private": private,
}
r = client.post("/template/", json_data=template)
r.raise_for_status()
return r.json()
def update_template(client: VastClient, hash_id: str, name: str = None,
image: str = None, image_tag: str = None, href: str = None,
repo: str = None, env: str = None, onstart_cmd: str = None,
jup_direct: bool = False, ssh_direct: bool = False,
use_jupyter_lab: bool = False, runtype: str = "args",
use_ssh: bool = False, jupyter_dir: str = None,
docker_login_repo: str = None, extra_filters: dict = None,
disk_space: float = None, readme: str = None,
readme_visible: bool = True, desc: str = None,
private: bool = True) -> dict:
"""Update an existing template.
Resolves the template's numeric id from its hash_id first; the v1 update
endpoint requires both. Edits always produce a new template version with a
new hash_id, so send the full template rather than only the changed fields.
Args:
client: VastClient instance.
hash_id: Hash ID of the template to update.
(remaining args same as create_template)
Returns:
Response dict with updated template info.
Raises:
ValueError: No template with this hash_id is visible to the caller.
"""
template = {
"id": resolve_template_id(client, hash_id),
"hash_id": hash_id,
"name": name,
"image": image,
"tag": image_tag,
"href": href,
"repo": repo,
"env": env,
"onstart": onstart_cmd,
"jup_direct": jup_direct,
"ssh_direct": ssh_direct,
"use_jupyter_lab": use_jupyter_lab,
"runtype": runtype,
"use_ssh": use_ssh,
"jupyter_dir": jupyter_dir,
"docker_login_repo": docker_login_repo,
"extra_filters": extra_filters or {},
"recommended_disk_space": disk_space,
"readme": readme,
"readme_visible": readme_visible,
"desc": desc,
"private": private,
}
r = client.put("/api/v1/template/", json_data=template)
r.raise_for_status()
return r.json()
def resolve_template_id(client: VastClient, hash_id: str) -> int:
"""Return the numeric id of the template whose hash_id matches.
The hash is content-based, so templates created by different users can share
one. When several match, the caller's own template wins over a public one.
Args:
client: VastClient instance.
hash_id: Hash ID of the template to look up.
Returns:
The template's numeric id.
Raises:
ValueError: No template with this hash_id is visible to the caller.
"""
matches = search_templates(client, {"hash_id": {"eq": hash_id}})
if not matches:
raise ValueError(f"No template found with hash_id {hash_id}")
if len(matches) > 1:
caller_id = show_user(client).get("id")
own = [match for match in matches if match.get("creator_id") == caller_id]
if own:
return own[0]["id"]
return matches[0]["id"]
def delete_template(client: VastClient, hash_id: str = None,
template_id: int = None) -> dict:
"""Delete a template by hash_id or template_id.
Args:
client: VastClient instance.
hash_id: Hash ID of the template to delete.
template_id: Numeric ID of the template to delete.
Returns:
Response dict.
"""
json_blob = {}
if hash_id:
json_blob["hash_id"] = hash_id
elif template_id:
json_blob["template_id"] = template_id
r = client.delete("/template/", json_data=json_blob)
r.raise_for_status()
return r.json()
def launch_instance(client: VastClient, gpu_name: str, num_gpus: str, image: str,
region: str = None, disk: float = 10, order: str = "score-",
limit: int = None, env: dict = None, label: str = None,
extra: str = None, onstart_cmd: str = None, login: str = None,
python_utf8: bool = False, lang_utf8: bool = False,
jupyter_lab: bool = False, jupyter_dir: str = None,
cancel_unavail: bool = False,
template_hash: str = None, runtype: str = None,
args: str = None, query: dict = None,
ssh: bool = False, jupyter: bool = False,
direct: bool = False) -> dict:
"""Launch the top instance from search offers matching the given criteria.
Searches for offers and launches the best match in a single API call.
Args:
client: VastClient instance.
gpu_name: GPU model name (e.g. "RTX_4090").
num_gpus: Number of GPUs required.
image: Docker image to launch.
region: Region name or country code list (e.g. "North_America" or "[US,CA]").
disk: Disk space in GB (default 10).
order: Sort order for offers (default "score-").
limit: Max number of offers to consider.
env: Environment variables dict.
label: Instance label.
extra: Extra docker options.
onstart_cmd: Onstart script contents.
login: Docker login credentials.
python_utf8: Enable Python UTF-8 mode.
lang_utf8: Enable lang UTF-8 mode.
jupyter_lab: Launch with Jupyter Lab.
jupyter_dir: Jupyter directory.
cancel_unavail: Cancel if unavailable.
template_hash: Template hash ID.
runtype: Run type (jupyter, ssh, args).
ssh: Launch as an ssh instance type.
jupyter: Launch as a jupyter instance instead of an ssh instance.
direct: Use (faster) direct connections for jupyter & ssh.
args: Container arguments.
query: Pre-built query dict (overrides auto-built query from gpu_name/num_gpus).
Returns:
Response dict with launch result.
"""
if isinstance(env, str):
env = parse_env(env)
if template_hash is None:
runtype, args = resolve_runtype(runtype, ssh=ssh, jupyter=jupyter,
direct=direct, args=args,
jupyter_lab=jupyter_lab,
jupyter_dir=jupyter_dir)
REGIONS = {
"North_America": "[AG, BS, BB, BZ, CA, CR, CU, DM, DO, SV, GD, GT, HT, HN, JM, MX, NI, PA, KN, LC, VC, TT, US]",
"South_America": "[AR, BO, BR, CL, CO, EC, FK, GF, GY, PY, PE, SR, UY, VE]",
"Europe": "[AL, AD, AT, BY, BE, BA, BG, HR, CY, CZ, DK, EE, FI, FR, DE, GR, HU, IS, IE, IT, LV, LI, LT, LU, MT, MD, MC, ME, NL, MK, NO, PL, PT, RO, RU, SM, RS, SK, SI, ES, SE, CH, UA, GB, VA, XK]",
"Asia": "[AF, AM, AZ, BH, BD, BT, BN, KH, CN, GE, IN, ID, IR, IQ, IL, JP, JO, KZ, KW, KG, LA, LB, MY, MV, MN, MM, NP, KP, OM, PK, PH, QA, SA, SG, KR, LK, SY, TW, TJ, TH, TL, TR, TM, AE, UZ, VN, YE, HK, MO]",
"Oceania": "[AS, AU, CK, FJ, PF, GU, KI, MH, FM, NR, NC, NZ, NU, MP, PW, PG, PN, WS, SB, TK, TO, TV, VU, WF]",
"Africa": "[DZ, AO, BJ, BW, BF, BI, CV, CM, CF, TD, KM, CG, CD, CI, DJ, EG, GQ, ER, SZ, ET, GA, GM, GH, GN, GW, KE, LS, LR, LY, MG, MW, ML, MR, MU, MA, MZ, NA, NE, NG, RW, ST, SN, SC, SL, SO, ZA, SS, SD, TZ, TG, TN, UG, ZM, ZW]",
}
if query is None:
args_query = f"num_gpus={num_gpus} gpu_name={gpu_name}"
if region:
region_query = REGIONS.get(region, region)
args_query += f" geolocation in {region_query}"
if disk:
args_query += f" disk_space>={disk}"
base_query = {"verified": {"eq": True}, "external": {"eq": False},
"rentable": {"eq": True}, "rented": {"eq": False}}
query = parse_query(args_query, base_query, offers_fields, offers_alias, offers_mult)
# Parse order string
order_list = parse_order(order) or []
query["order"] = order_list
query["type"] = "on-demand"
if limit:
query["limit"] = int(limit)
query["allocated_storage"] = disk
json_blob = {
"image": image,
"disk": disk,
"q": query,
"env": env or {},
"label": label,
"extra": extra,
"onstart": onstart_cmd,
"image_login": login,
"python_utf8": python_utf8,
"lang_utf8": lang_utf8,
"use_jupyter_lab": jupyter_lab,
"jupyter_dir": jupyter_dir,
"cancel_unavail": cancel_unavail,
"template_hash_id": template_hash,
}
if runtype:
json_blob["runtype"] = runtype
if args is not None:
json_blob["args"] = args
r = client.put("/launch_instance/", json_data=json_blob)
r.raise_for_status()
return r.json()
api/price_increase.py
"""Price-increase contract extension API.
Backend pairs (vast/web/views/instance.py:257-312,
vast/web/views/pydantic/models/instance.py):
GET /api/v0/instances/pending-price-increases/
PUT /api/v0/instances/accept-price-increase/
PUT /api/v0/instances/reject-price-increase/
Both PUT endpoints require body exactly
``{"pending_price_increase_id": <int>}`` (``_Base.Config.extra='forbid'``).
The backend has no batch endpoint, no ``instance_ids``, no ``host_id``,
and no ``snapshot`` field. The pending row id IS the identity.
Stale rows return ``HTTP 404`` with body
``{"success": false, "error": "no_pending_price_increase"}``; older
backends may still return ``HTTP 409`` (the frontend recognises both).
"""
from vastai.api.client import VastClient
NO_PENDING_PRICE_INCREASE = "no_pending_price_increase"
def list_pending(client: VastClient) -> dict:
"""Return the pending-price-increase envelope.
Shape (from web/price_increase_pending.py:_serialize_pending_row):
{
"success": True,
"count": <int>,
"truncated": <bool>,
"pending_price_increases": [
{
"pending_price_increase_id": <int>,
"contract_id": <int>,
"host_id": <int>,
"new_gpu_costpersec": <float|null>,
"new_disk_ram_costpersec": <float|null>,
"new_bwu_cost": <float|null>,
"new_bwd_cost": <float|null>,
"new_platform_fee": <float|null>,
"old_gpu_costpersec": <float|null>,
"old_disk_ram_costpersec": <float|null>,
"old_bwu_cost": <float|null>,
"old_bwd_cost": <float|null>,
"old_platform_fee": <float|null>,
"contract_end_date": <float|null>,
"ask_end_date": <float|null>,
"created_at": <float|null>
},
...
]
}
"""
r = client.get("/instances/pending-price-increases/")
r.raise_for_status()
return r.json()
def pending_rows(envelope: dict) -> list[dict]:
"""Return the pending rows from a pending-price-increase envelope."""
return envelope.get("pending_price_increases", []) or []
def find_pending_for_instance(rows: list[dict], instance_id: int) -> dict | None:
"""Find one pending row by contract/instance id."""
target = int(instance_id)
return next((row for row in rows if row.get("contract_id") == target), None)
def resolve_instance_to_pending(client: VastClient, instance_id: int) -> dict:
"""Resolve an instance id to its pending-price-increase row.
Raises ``LookupError`` when no pending row matches the instance id.
"""
envelope = list_pending(client)
match = find_pending_for_instance(pending_rows(envelope), instance_id)
if match is None:
raise LookupError(
f"no pending price increase for instance {instance_id}"
)
return match
def accept(client: VastClient, pending_id: int) -> dict:
"""Accept one pending price-increase row.
The body must be exactly ``{"pending_price_increase_id": int}``;
backend ``extra='forbid'`` returns 400 for any extra key. Returns
``{"success": True, "pending_price_increase_id": int, "contract_id": int}``.
"""
r = client.put(
"/instances/accept-price-increase/",
json_data={"pending_price_increase_id": int(pending_id)},
)
r.raise_for_status()
return r.json()
def reject(client: VastClient, pending_id: int) -> dict:
"""Reject one pending price-increase row.
Same body shape and response shape as :func:`accept`. The backend
tombstones the row (status=rejected); no cutover follows.
"""
r = client.put(
"/instances/reject-price-increase/",
json_data={"pending_price_increase_id": int(pending_id)},
)
r.raise_for_status()
return r.json()
api/query.py
"""Query parsing and field definitions for Vast.ai API.
This module is shared between the SDK and CLI layers.
"""
import re
import sys
import time
from datetime import datetime, timezone
from typing import Dict, List
def numeric_version(version_str):
try:
# Split the version string by the period
major, minor, patch = version_str.split('.')
# Pad each part with leading zeros to make it 3 digits
major = major.zfill(3)
minor = minor.zfill(3)
patch = patch.zfill(3)
# Concatenate the padded parts
numeric_version_str = f"{major}{minor}{patch}"
# Convert the concatenated string to an integer
result = int(numeric_version_str)
#print(result)
return result
except ValueError:
print("Invalid version string format. Expected format: X.X.X")
return None
def string_to_unix_epoch(date_string):
if date_string is None:
return None
try:
# Check if the input is a float or integer representing Unix time
return float(date_string)
except ValueError:
# Parse MM/DD/YYYY input; interpret as UTC midnight so filtering
# windows don't shift with the caller's local timezone.
date_object = datetime.strptime(date_string, "%m/%d/%Y").replace(tzinfo=timezone.utc)
return date_object.timestamp()
def fix_date_fields(query: Dict[str, Dict], date_fields: List[str]):
"""Takes in a query and date fields to correct and returns query with appropriate epoch dates"""
new_query: Dict[str, Dict] = {}
for field, sub_query in query.items():
# fix date values for given date fields
if field in date_fields:
new_sub_query = {k: string_to_unix_epoch(v) for k, v in sub_query.items()}
new_query[field] = new_sub_query
# else, use the original
else: new_query[field] = sub_query
return new_query
def version_string_sort(a, b) -> int:
"""
Accepts two version strings and decides whether a > b, a == b, or a < b.
This is meant as a sort function to be used for the driver versions in which only
the == operator currently works correctly. Not quite finished...
:param str a:
:param str b:
:return int:
"""
a_parts = a.split(".")
b_parts = b.split(".")
return 0
vol_offers_fields = {
"cpu_arch",
"cuda_vers",
"cluster_id",
"nw_disk_min_bw",
"nw_disk_avg_bw",
"nw_disk_max_bw",
"datacenter",
"disk_bw",
"disk_space",
"driver_version",
"duration",
"geolocation",
"gpu_arch",
"has_avx",
"host_id",
"id",
"inet_down",
"inet_up",
"machine_id",
"pci_gen",
"pcie_bw",
"reliability",
"storage_cost",
"static_ip",
"total_flops",
"ubuntu_version",
"verified",
}
offers_fields = {
"bw_nvlink",
"compute_cap",
"cpu_arch",
"cpu_cores",
"cpu_cores_effective",
"cpu_ghz",
"cpu_ram",
"cuda_max_good",
"datacenter",
"direct_port_count",
"driver_version",
"disk_bw",
"disk_space",
"dlperf",
"dlperf_per_dphtotal",
"dph_total",
"duration",
"external",
"flops_per_dphtotal",
"gpu_arch",
"gpu_display_active",
"gpu_frac",
# "gpu_ram_free_min",
"gpu_mem_bw",
"gpu_name",
"gpu_ram",
"gpu_total_ram",
"gpu_display_active",
"gpu_max_power",
"gpu_max_temp",
"has_avx",
"host_id",
"id",
"inet_down",
"inet_down_cost",
"inet_up",
"inet_up_cost",
"machine_id",
"min_bid",
"mobo_name",
"num_gpus",
"pci_gen",
"pcie_bw",
"reliability",
#"reliability2",
"rentable",
"rented",
"storage_cost",
"static_ip",
"total_flops",
"ubuntu_version",
"verification",
"verified",
"vms_enabled",
"geolocation",
"cluster_id"
}
offers_alias = {
"cuda_vers": "cuda_max_good",
"display_active": "gpu_display_active",
#"reliability": "reliability2",
"dlperf_usd": "dlperf_per_dphtotal",
"dph": "dph_total",
"flops_usd": "flops_per_dphtotal",
}
offers_mult = {
"cpu_ram": 1000,
"gpu_ram": 1000,
"gpu_total_ram" : 1000,
"duration": 24.0 * 60.0 * 60.0,
}
benchmarks_fields = {
"contract_id",# int ID of instance/contract reporting benchmark
"dph_base",# float gpu $/hr at benchmark time
"gpu_name",# string GPU model benchmarked (e.g. RTX_4090)
"id",# int benchmark unique ID
"image",# string image used for benchmark
"last_update",# float date of benchmark
"launch_args",# string launch args of the benchmarked worker
"machine_id",# int id of machine benchmarked
"num_gpus",# int number of gpus used in benchmark
"template_hash",# string hash of the template benchmarked (formerly 'name')
"template_id",# int id of the template benchmarked
"type",# string benchmark type
"value",# float benchmark result (formerly 'score')
}
invoices_fields = {
'id',# int,
'user_id',# int,
'when',# float,
'paid_on',# float,
'payment_expected',# float,
'amount_cents',# int,
'is_credit',# bool,
'is_delayed',# bool,
'balance_before',# float,
'balance_after',# float,
'original_amount',# int,
'event_id',# string,
'cut_amount',# int,
'cut_percent',# float,
'extra',# json,
'service',# string,
'stripe_charge',# json,
'stripe_refund',# json,
'stripe_payout',# json,
'error',# json,
'paypal_email',# string,
'transfer_group',# string,
'failed',# bool,
'refunded',# bool,
'is_check',# bool,
}
templates_fields = {
"creator_id",# int ID of creator
"created_at",# float time of initial template creation (UTC epoch timestamp)
"count_created",# int #instances created (popularity)
"default_tag",# string image default tag
"docker_login_repo",# string image docker repository
"id",# int template unique ID
"image",# string image used for benchmark
"jup_direct",# bool supports jupyter direct
"hash_id",# string unique hash ID of template
"private",# bool true: only your templates, None: public templates
"name",# string displayable name
"recent_create_date",# float last time of instance creation (UTC epoch timestamp)
"recommended_disk_space",# float min disk space required
"recommended",# bool is templated on our recommended list
"ssh_direct",# bool supports ssh direct
"tag",# string image tag
"use_ssh",# string supports ssh (direct or proxy)
}
def parse_order(order, field_alias: Dict = None) -> List:
"""Turn a comma-separated order string into [[field, direction], ...].
A leading or trailing ``-`` means descending, ``+`` ascending. A list is
returned unchanged so callers can pass either spelling.
"""
if order is None:
return None
if not isinstance(order, str):
return order
field_alias = offers_alias if field_alias is None else field_alias
parsed = []
for name in order.split(","):
name = name.strip()
if not name:
continue
direction = "asc"
field = name
if name.strip("-") != name:
direction = "desc"
field = name.strip("-")
if name.strip("+") != name:
direction = "asc"
field = name.strip("+")
parsed.append([field_alias.get(field, field), direction])
return parsed
def parse_query(query_str: str, res: Dict = None, fields = {}, field_alias = {}, field_multiplier = {}) -> Dict:
"""
Basically takes a query string (like the ones in the examples of commands for the search__offers function) and
processes it into a dict of URL parameters to be sent to the server.
:param str query_str:
:param Dict res:
:return Dict:
"""
if query_str is None:
return res
if res is None: res = {}
if type(query_str) == list:
query_str = " ".join(query_str)
query_str = query_str.strip()
# Revised regex pattern to accurately capture quoted strings, bracketed lists, and single words/numbers
#pattern = r"([a-zA-Z0-9_]+)\s*(=|!=|<=|>=|<|>| in | nin | eq | neq | not eq | not in )?\s*(\"[^\"]*\"|\[[^\]]+\]|[^ ]+)"
#pattern = "([a-zA-Z0-9_]+)( *[=><!]+| +(?:[lg]te?|nin|neq|eq|not ?eq|not ?in|in) )?( *)(\[[^\]]+\]|[^ ]+)?( *)"
pattern = r"([a-zA-Z0-9_]+)( *[=><!]+| +(?:[lg]te?|nin|neq|eq|not ?eq|not ?in|in) )?( *)(\[[^\]]+\]|\"[^\"]+\"|[^ ]+)?( *)"
opts = re.findall(pattern, query_str)
#print("parse_query regex:")
#print(opts)
#print(opts)
# res = {}
op_names = {
">=": "gte",
">": "gt",
"gt": "gt",
"gte": "gte",
"<=": "lte",
"<": "lt",
"lt": "lt",
"lte": "lte",
"!=": "neq",
"==": "eq",
"=": "eq",
"eq": "eq",
"neq": "neq",
"noteq": "neq",
"not eq": "neq",
"notin": "notin",
"not in": "notin",
"nin": "notin",
"in": "in",
}
joined = "".join("".join(x) for x in opts)
if joined != query_str:
raise ValueError(
"Unconsumed text. Did you forget to quote your query? " + repr(joined) + " != " + repr(query_str))
for field, op, _, value, _ in opts:
value = value.strip(",[]")
op = op.strip()
op_name = op_names.get(op)
if field in field_alias:
old_field = field
field = field_alias[field]
if old_field in res:
res[field] = res.pop(old_field)
v = res.setdefault(field, {})
if (field == "driver_version") and ('.' in value):
value = numeric_version(value)
if not field in fields:
print("Warning: Unrecognized field: {}, see list of recognized fields.".format(field), file=sys.stderr);
if not op_name:
raise ValueError("Unknown operator. Did you forget to quote your query? " + repr(op).strip("u"))
if op_name in ["in", "notin"]:
value = [x.strip() for x in value.split(",") if x.strip()]
if not value:
raise ValueError("Value cannot be blank. Did you forget to quote your query? " + repr((field, op, value)))
if not field:
raise ValueError("Field cannot be blank. Did you forget to quote your query? " + repr((field, op, value)))
if value in ["?", "*", "any"]:
if op_name != "eq":
raise ValueError("Wildcard only makes sense with equals.")
if field in v:
del v[field]
if field in res:
del res[field]
continue
if isinstance(value, str):
value = value.replace('_', ' ')
value = value.strip('\"')
elif isinstance(value, list):
value = [x.replace('_', ' ') for x in value]
value = [x.strip('\"') for x in value]
if field in field_multiplier:
value = float(value) * field_multiplier[field]
v[op_name] = value
else:
#print(value)
if (value == 'true') or (value == 'True'):
v[op_name] = True
elif (value == 'false') or (value == 'False'):
v[op_name] = False
elif (value == 'None') or (value == 'null'):
v[op_name] = None
else:
v[op_name] = value
if field not in res:
res[field] = v
else:
res[field].update(v)
#print(res)
return res
api/storage.py
"""Storage, volume, and data transfer API functions for the Vast.ai SDK."""
import time
def copy(client, src_id, dst_id, src_path, dst_path):
"""Copy directories between instances, volumes, cloud services, and local.
PUT /commands/copy_direct/ (remote-to-remote) or
PUT /commands/rsync/ (when one side is local/None)
Each of the source and destination locations can be either local or remote,
subject to the read and write permissions required to carry out the action.
Location IDs are the parsed form of a vast URL (see
``vastai.utils.parse_vast_url``):
- instance_id (int) legacy format, still supported
- \"C.instance_id\" container copy format
- \"cloud_service\" cloud service, e.g. \"drive\"
- \"cloud_service.id\" cloud service with ID, e.g. \"s3.101\" or \"hf.101\"
- \"local\" or None local machine
- \"V.volume_id\" volume copy, e.g. \"V.1234\"
Volume copy is currently only supported for copying to other volumes,
instances, or cloud services, not local.
You should not copy to /root or / as a destination directory, as this can
mess up the permissions on your instance ssh folder, breaking future copy
operations (as they use ssh authentication). See
https://vast.ai/docs/gpu-instances/data-movement#constraints for more
information about constraints.
Args:
client: VastClient instance.
src_id: Source location ID (or None/\"local\" for local).
dst_id: Destination location ID (or None/\"local\" for local).
src_path (str): Source path.
dst_path (str): Destination path.
Returns:
dict: API response data. For local transfers this includes the
``src_addr``/``src_port`` or ``dst_addr``/``dst_port`` the caller
uses to run rsync.
Examples:
copy(client, 6003036, 6003038, \"/workspace/\", \"/workspace/\")
copy(client, \"C.11824\", \"local\", \"/data/test\", \"data/test\")
copy(client, \"drive\", \"C.6003036\", \"/folder/file.txt\", \"/workspace/\")
copy(client, \"s3.101\", \"C.6003036\", \"/data/\", \"/workspace/\")
copy(client, \"V.1234\", \"C.5678\", \"/file\", \"/workspace/\")
copy(client, \"V.1234\", \"s3.101\", \"/file\", \"/workspace/\")
"""
req_json = {
"client_id": "me",
"src_id": src_id,
"dst_id": dst_id,
"src_path": src_path,
"dst_path": dst_path,
}
if src_id is None or dst_id is None:
r = client.put("/commands/rsync/", json_data=req_json)
else:
r = client.put("/commands/copy_direct/", json_data=req_json)
r.raise_for_status()
return r.json()
def cancel_copy(client, dst_id):
"""Cancel a remote copy in progress.
DELETE /commands/copy_direct/
Args:
client: VastClient instance.
dst_id: ID of copy destination to cancel.
Returns:
dict: API response data.
"""
req_json = {"client_id": "me", "dst_id": dst_id}
r = client.delete("/commands/copy_direct/", json_data=req_json)
r.raise_for_status()
return r.json()
def cancel_sync(client, dst_id):
"""Cancel a remote cloud sync in progress.
DELETE /commands/rclone/
Args:
client: VastClient instance.
dst_id: ID of cloud sync destination to cancel.
Returns:
dict: API response data.
"""
req_json = {"client_id": "me", "dst_id": dst_id}
r = client.delete("/commands/rclone/", json_data=req_json)
r.raise_for_status()
return r.json()
def rclone_flags(*, dry_run: bool = False, size_only: bool = False,
ignore_existing: bool = False, update: bool = False,
delete_excluded: bool = False) -> list:
"""Translate the friendly copy options into the rclone flag list."""
flags = []
if dry_run:
flags.append("--dry-run")
if size_only:
flags.append("--size-only")
if ignore_existing:
flags.append("--ignore-existing")
if update:
flags.append("--update")
if delete_excluded:
flags.append("--delete-excluded")
return flags
def cloud_copy(client, src, dst, instance, connection, transfer, flags=None):
"""Copy files/folders to and from cloud providers.
POST /commands/rclone/
Supported cloud providers include \"drive\", \"s3\", \"b2\", \"dropbox\", and
Hugging Face (\"hf\").
Args:
client: VastClient instance.
src (str): Path to source of object to copy.
dst (str): Path to target of copy operation. Default \"/workspace\".
instance (str): ID of the instance.
connection (str): ID of cloud connection on your account.
transfer (str): Type of transfer (e.g. \"Instance to Cloud\",
\"Cloud to Instance\").
flags (list, optional): Additional rclone flags such as
[\"--dry-run\", \"--size-only\", \"--ignore-existing\",
\"--update\", \"--delete-excluded\"].
Returns:
dict: API response data.
Hugging Face example (\"hf\" connection):
cloud_copy(client, \"my-bucket/data\", \"/workspace\", \"6003036\",
\"1234\", \"Cloud To Instance\")
"""
req_json = {
"src": src,
"dst": dst,
"instance_id": instance,
"selected": connection,
"transfer": transfer,
"flags": flags or [],
}
r = client.post("/commands/rclone/", json_data=req_json)
r.raise_for_status()
return r.json()
def clone_volume(client, source, dest, size=None, disable_compression=False):
"""Clone a volume to another volume.
POST /volumes/copy/
Args:
client: VastClient instance.
source: Source volume ID.
dest: Destination volume ID.
size (float, optional): Size in GB for destination volume.
disable_compression (bool): Disable compression during clone.
Default False.
Returns:
dict: API response data.
"""
json_blob = {
"src_id": source,
"dst_id": dest,
}
if size is not None:
json_blob["size"] = size
if disable_compression:
json_blob["disable_compression"] = True
r = client.post("/volumes/copy/", json_data=json_blob)
r.raise_for_status()
return r.json()
def show_volumes(client, type="all"):
"""Show stats on owned volumes.
GET /volumes
Args:
client: VastClient instance.
type (str): Volume type to display. Options: \"local\", \"network\",
\"all\". Default \"all\".
Returns:
list: Volume data with computed duration field.
"""
types = {
"local": "local_volume",
"network": "network_volume",
"all": "all_volume",
}
vol_type = types.get(type, "all_volume")
r = client.get("/volumes", query_args={"owner": "me", "type": vol_type})
r.raise_for_status()
rows = r.json()["volumes"]
processed = []
for row in rows:
row['duration'] = time.time() - row['start_date']
processed.append(row)
return processed
def show_instance_backups(client, contract_id=None):
"""List the backups Vast has taken of the caller's instances and volumes.
GET /instance_backups/
Args:
client: VastClient instance.
contract_id (int): Restrict to one instance or volume contract id.
Returns:
list: Backup rows, newest first.
"""
query_args = {}
if contract_id is not None:
query_args["contract_id"] = contract_id
r = client.get("/instance_backups/", query_args=query_args)
r.raise_for_status()
return r.json()["backups"]
def instance_backup_files(client, contract_id):
"""Objects stored for one backup, plus a short-lived token to fetch them.
GET /instance_backups/<contract_id>/download/
The returned URLs are not usable alone: each request must carry the
authorization_token, either as an `Authorization` header or appended as an
`?Authorization=<token>` query parameter (what the CLI does, so the printed
URLs are clickable). `truncated` is True when the backup holds more objects
than were listed.
Args:
client: VastClient instance.
contract_id (int): Instance or volume contract id.
Returns:
dict: authorization_token, expires_in_sec, truncated, files.
"""
r = client.get(f"/instance_backups/{contract_id}/download/")
r.raise_for_status()
return r.json()
def create_volume(client, id, size=15, name=None):
"""Create a new volume.
PUT /volumes/
Args:
client: VastClient instance.
id (int): ID of volume offer.
size (float): Size in GB. Default 15.
name (str, optional): Optional name of volume.
Returns:
dict: API response data.
"""
json_blob = {
"size": int(size),
"id": int(id),
}
if name:
json_blob["name"] = name
r = client.put("/volumes/", json_data=json_blob)
r.raise_for_status()
return r.json()
def delete_volume(client, id):
"""Delete a volume.
DELETE /volumes/
All instances using the volume must be destroyed before deletion.
Args:
client: VastClient instance.
id (int): ID of volume contract.
Returns:
dict: API response data.
"""
r = client.delete("/volumes/", query_args={"id": id})
r.raise_for_status()
return r.json()
def list_volume(client, id, size=15, price_disk=0.10, end_date=None):
"""[Host] List disk space for rent as a volume on a machine.
POST /volumes/
Args:
client: VastClient instance.
id (int): ID of machine to list.
size (int): Size of disk space in GB. Default 15.
price_disk (float): Storage price in $/GB/month. Default 0.10.
end_date (str, optional): Contract offer expiration date
(unix timestamp or MM/DD/YYYY format).
Returns:
dict: API response data.
"""
json_blob = {
"size": int(size),
"machine": int(id),
"price_disk": float(price_disk),
}
if end_date is not None:
json_blob["end_date"] = end_date
r = client.post("/volumes/", json_data=json_blob)
r.raise_for_status()
return r.json()
def list_volumes(client, ids, size=15, price_disk=0.10, end_date=None):
"""[Host] List disk space for rent as volumes on multiple machines.
POST /volumes/
Args:
client: VastClient instance.
ids (list[int]): IDs of machines to list.
size (int): Size of disk space in GB. Default 15.
price_disk (float): Storage price in $/GB/month. Default 0.10.
end_date (str, optional): Contract offer expiration date
(unix timestamp or MM/DD/YYYY format).
Returns:
dict: API response data.
"""
json_blob = {
"size": int(size),
"machine": [int(mid) for mid in ids],
"price_disk": float(price_disk),
}
if end_date is not None:
json_blob["end_date"] = end_date
r = client.post("/volumes/", json_data=json_blob)
r.raise_for_status()
return r.json()
def unlist_volume(client, id):
"""[Host] Unlist a volume offer.
POST /volumes/unlist
Args:
client: VastClient instance.
id (int): Volume ID to unlist.
Returns:
dict: API response data.
"""
json_blob = {"id": id}
r = client.post("/volumes/unlist", json_data=json_blob)
r.raise_for_status()
return r.json()
def create_network_volume(client, id, size=15, name=None):
"""Create a new network volume.
PUT /network_volumes/
Args:
client: VastClient instance.
id (int): ID of network volume offer.
size (float): Size in GB. Default 15.
name (str, optional): Optional name of network volume.
Returns:
dict: API response data.
"""
json_blob = {
"size": int(size),
"id": int(id),
}
if name:
json_blob["name"] = name
r = client.put("/network_volumes/", json_data=json_blob)
r.raise_for_status()
return r.json()
def list_network_volume(client, disk_id, price_disk=0.15, size=15, end_date=None):
"""[Host] List disk space for rent as a network volume.
POST /network_volumes/
Args:
client: VastClient instance.
disk_id (int): ID of disk to list.
price_disk (float): Storage price in $/GB/month. Default 0.15.
size (int): Size of disk space in GB. Default 15.
end_date (str, optional): Contract offer expiration date.
Returns:
dict: API response data.
"""
json_blob = {
"disk_id": disk_id,
"price_disk": price_disk,
"size": size,
}
if end_date is not None:
json_blob["end_date"] = end_date
r = client.post("/network_volumes/", json_data=json_blob)
r.raise_for_status()
return r.json()
def unlist_network_volume(client, id):
"""[Host] Unlist a network volume offer.
POST /network_volumes/unlist/
Args:
client: VastClient instance.
id (int): ID of network volume offer to unlist.
Returns:
dict: API response data.
"""
json_blob = {"id": id}
r = client.post("/network_volumes/unlist/", json_data=json_blob)
r.raise_for_status()
return r.json()
def show_network_disks(client):
"""[Host] Show network disks associated with your account.
GET /network_disk/
Args:
client: VastClient instance.
Returns:
dict: Network disk data including cluster and machine info.
"""
r = client.get("/network_disk/")
r.raise_for_status()
return r.json()
def add_network_disk(client, machines, mount_point, disk_id=None):
"""[Host] Add network disk to physical cluster.
POST /network_disk/
Args:
client: VastClient instance.
machines (list[int]): IDs of machines to add disk to.
mount_point (str): Mount path of disk to add.
disk_id (int, optional): ID of network disk to attach to machines.
Returns:
dict: API response data including disk_id.
"""
json_blob = {
"machines": [int(mid) for mid in machines],
"mount_point": mount_point,
}
if disk_id is not None:
json_blob["disk_id"] = disk_id
r = client.post("/network_disk/", json_data=json_blob)
r.raise_for_status()
return r.json()
def show_connections(client):
"""Display user's cloud connections/integrations.
GET /users/cloud_integrations/
Args:
client: VastClient instance.
Returns:
dict/list: Cloud integration data.
"""
r = client.get("/users/cloud_integrations/")
r.raise_for_status()
return r.json()
api/teams.py
"""Team CRUD operations, members, and roles."""
from vastai.api.client import VastClient
def create_team(client: VastClient, team_name: str, transfer_credit: float = 0) -> dict:
"""Create a new team.
Args:
client: VastClient instance.
team_name: Name of the team to create.
transfer_credit: Amount of personal credit to transfer to the new team.
Returns:
Response dict with team info.
"""
json_data = {"team_name": team_name, "transfer_credit": transfer_credit}
r = client.post("/team/", json_data=json_data)
r.raise_for_status()
return r.json()
def destroy_team(client: VastClient) -> dict:
"""Destroy the current team.
Args:
client: VastClient instance.
Returns:
Response dict.
"""
r = client.delete("/team/")
r.raise_for_status()
return r.json()
def show_members(client: VastClient) -> dict:
"""Show members of the current team.
Args:
client: VastClient instance.
Returns:
Response dict with member info.
"""
r = client.get("/team/members/")
r.raise_for_status()
return r.json()
def invite_member(client: VastClient, email: str, role: str) -> dict:
"""Invite a member to the current team.
Args:
client: VastClient instance.
email: Email address of the member to invite.
role: Role to assign to the invited member.
Returns:
Response dict.
"""
r = client.post("/team/invite/", query_args={"email": email, "role": role})
r.raise_for_status()
return r.json()
def remove_member(client: VastClient, id: int) -> dict:
"""Remove a member from the current team.
Args:
client: VastClient instance.
id: Member ID to remove.
Returns:
Response dict.
"""
r = client.delete(f"/team/members/{id}/")
r.raise_for_status()
return r.json()
def create_team_role(client: VastClient, name: str, permissions: dict) -> dict:
"""Add a new role to the current team.
Args:
client: VastClient instance.
name: Name of the role.
permissions: Dict of permissions for the role.
Returns:
Response dict.
"""
r = client.post("/team/roles/", json_data={"name": name, "permissions": permissions})
r.raise_for_status()
return r.json()
def show_team_role(client: VastClient, name: str) -> dict:
"""Show details of a specific team role.
Args:
client: VastClient instance.
name: Name of the role.
Returns:
Role details dict.
"""
r = client.get(f"/team/roles/{name}/")
r.raise_for_status()
return r.json()
def show_team_roles(client: VastClient) -> dict:
"""Show all roles for the current team.
Args:
client: VastClient instance.
Returns:
Response dict with roles info.
"""
r = client.get("/team/roles-full/")
r.raise_for_status()
return r.json()
def update_team_role(client: VastClient, id: int, name: str = None,
permissions: dict = None) -> dict:
"""Update an existing team role.
Args:
client: VastClient instance.
id: Role ID.
name: New name for the role.
permissions: Updated permissions dict.
Returns:
Response dict.
"""
json_blob = {}
if name is not None:
json_blob["name"] = name
if permissions is not None:
json_blob["permissions"] = permissions
r = client.put(f"/team/roles/{id}/", json_data=json_blob)
r.raise_for_status()
return r.json()
def remove_team_role(client: VastClient, name: str) -> dict:
"""Remove a role from the current team.
Args:
client: VastClient instance.
name: Name of the role to remove.
Returns:
Response dict.
"""
r = client.delete(f"/team/roles/{name}/")
r.raise_for_status()
return r.json()
def transfer_credit(client: VastClient, recipient: str, amount: float) -> dict:
"""Transfer credit to another account.
Args:
client: VastClient instance.
recipient: Recipient identifier (email or user ID).
amount: Amount of credit to transfer.
Returns:
Response dict.
"""
json_blob = {
"sender": "me",
"recipient": recipient,
"amount": amount,
}
r = client.put("/commands/transfer_credit/", json_data=json_blob)
r.raise_for_status()
return r.json()
async_/__init__.py
from .client import AsyncClient
from .results import AsyncOffer, AsyncInstance
async_/client.py
import aiohttp
import ssl
import os
import json
import tempfile
import asyncio
from typing import Any, Optional, Union
from vastai._base import _BaseClient, _APIKEY_SENTINEL
from vastai.data.query import Query
from vastai.data.offer import Offer
from vastai.data.instance import Instance, InstanceConfig, CreateInstanceResponse
from vastai.async_.results import AsyncOffer, AsyncInstance
from vastai.serverless.client.client import CoroutineServerless
class AsyncClient(_BaseClient):
"""
Asynchronous Vast.ai API client.
api_key resolution order:
1. Explicit api_key argument
2. VAST_API_KEY environment variable
3. $XDG_CONFIG_HOME/vastai/vast_api_key
4. ~/.vast_api_key (legacy)
"""
SSL_CERT_URL = "https://console.vast.ai/static/jvastai_root.cer"
def __init__(
self,
api_key: object = _APIKEY_SENTINEL,
vast_server: str = "https://console.vast.ai",
connection_limit: int = 100,
):
super().__init__(api_key, vast_server)
self._connection_limit = connection_limit
self._session: aiohttp.ClientSession | None = None
self._ssl_context: ssl.SSLContext | None = None
self._session_depth = 0
async def get_ssl_context(self) -> ssl.SSLContext:
"""Download Vast.ai root cert and build SSL context (cached).
The custom cert is added on top of the system default CAs, so this
context works for both console.vast.ai and serverless worker connections.
"""
if self._ssl_context is None:
async with aiohttp.ClientSession() as s:
async with s.get(self.SSL_CERT_URL) as resp:
if resp.status != 200:
raise Exception(f"Failed to fetch SSL cert: {resp.status}")
cert_bytes = await resp.read()
tmpfile = tempfile.NamedTemporaryFile(delete=False, suffix=".cer")
tmpfile.write(cert_bytes)
tmpfile.close()
ctx = ssl.create_default_context()
ctx.load_verify_locations(cafile=tmpfile.name)
self._ssl_context = ctx
os.unlink(tmpfile.name)
return self._ssl_context
async def _get_session(self, use_context=False) -> aiohttp.ClientSession:
if self._session is None or self._session.closed:
self._session = aiohttp.ClientSession(
connector=aiohttp.TCPConnector(
limit=self._connection_limit,
ssl=await self.get_ssl_context(),
)
)
self._session_depth += 1
if not use_context:
loop = asyncio.get_running_loop()
original_close = loop.close
def close():
loop.run_until_complete(self.close(force=True))
original_close()
loop.close = close
elif use_context:
self._session_depth += 1 # If we already are expected to keep a connection, and we expect to call close before exit, still flag so that close doesn't actually close.
return self._session
def is_open(self):
return self._session is not None and not self._session.closed
async def close(self, force=False):
if self._session and not self._session.closed:
if self._session_depth > 0:
self._session_depth -= 1
if force or not self._session_depth:
self._session_depth = 0
await self._session.close()
async def __aenter__(self):
await self._get_session(use_context=True)
return self
async def __aexit__(self, exc_type, exc, tb):
await self.close()
# ── Search ──────────────────────────────────────────────────────────────
async def search(
self,
query: Query,
type_: Optional[str] = None,
order: Optional[list] = [["score", "desc"]],
limit: Optional[int] = None,
allocated_storage: Optional[float] = None,
target_reliability: Optional[float] = None,
external: Optional[bool] = None,
extra_ids: Optional[list] = None,
has_avx: Optional[int] = None,
gpu_option: Optional[Any] = None,
show_incompatible: Optional[Any] = None,
sort_option: Optional[Any] = None,
template_id: Optional[Any] = None,
) -> list[AsyncOffer]:
"""Search available offers and return a list of AsyncOffer result wrappers."""
payload = dict(query.query)
for key, value in [
("type", type_),
("order", order),
("limit", limit),
("allocated_storage", allocated_storage),
("target_reliability", target_reliability),
("external", external),
("extra_ids", extra_ids),
("has_avx", has_avx),
("gpu_option", gpu_option),
("show_incompatible", show_incompatible),
("sort_option", sort_option),
("template_id", template_id),
]:
if value is not None:
payload[key] = value
session = await self._get_session()
async with session.post(
self._url("/api/v0/bundles/"),
json=payload,
headers=self._headers(),
) as resp:
resp.raise_for_status()
data = await resp.json(content_type=None)
return [AsyncOffer(Offer.from_dict(o), self) for o in data.get("offers", [])]
# ── Create Instance ─────────────────────────────────────────────────────
async def create_instance(
self,
offer_id: int,
config: InstanceConfig,
) -> AsyncInstance:
"""Accept an offer and create an instance. Returns an AsyncInstance wrapper."""
payload = config.to_dict()
session = await self._get_session()
async with session.put(
self._url(f"/api/v0/asks/{offer_id}/"),
json=payload,
headers=self._headers(),
) as resp:
resp.raise_for_status()
data = await resp.json(content_type=None)
if not data.get("success"):
raise Exception(
f"Failed to create instance: {data.get('error', '')} - {data.get('msg', '')}"
)
resp_data = CreateInstanceResponse.from_dict(data)
# Fetch the newly created instance
instances = await self.show_instances()
for inst in instances:
if inst.id == resp_data.new_contract:
return inst
# Fallback: return a minimal instance
return AsyncInstance(Instance.from_dict({"id": resp_data.new_contract}), self)
# ── Instances ───────────────────────────────────────────────────────────
async def show_instances(self) -> list[AsyncInstance]:
"""Return all instances owned by the authenticated user.
The legacy v0 ``/instances`` list endpoint is deprecated; this pages
through the v1 ``/api/v1/instances/`` endpoint (omitting ``select_cols``
so the backend returns full instance rows) and accumulates every page.
"""
session = await self._get_session()
instances = []
params = {
"select_filters": json.dumps({}),
"order_by": json.dumps([{"col": "id", "dir": "asc"}]),
"limit": 25,
}
while True:
async with session.get(
self._url("/api/v1/instances/"),
params=params,
headers=self._headers(),
) as resp:
resp.raise_for_status()
data = await resp.json(content_type=None)
instances.extend(data.get("instances") or [])
next_token = data.get("next_token")
if not next_token:
break
params["after_token"] = next_token
return [
AsyncInstance(Instance.from_dict(i), self)
for i in instances
]
async def destroy_instance(self, instance_or_id: Union[AsyncInstance, int]) -> None:
"""Destroy a running or stopped instance."""
instance_id = (
instance_or_id if isinstance(instance_or_id, int) else instance_or_id.id
)
session = await self._get_session()
async with session.delete(
self._url(f"/api/v0/instances/{instance_id}/"),
headers=self._headers(),
) as resp:
resp.raise_for_status()
# ── Serverless factory ──────────────────────────────────────────────────
def serverless(
self,
*,
instance: str = "prod",
autoscaler_url: Optional[str] = None,
debug: bool = False,
default_request_timeout: float = 600.0,
max_poll_interval: float = 5.0,
) -> CoroutineServerless:
"""Create a CoroutineServerless that shares this client's aiohttp session and SSL context."""
sl = CoroutineServerless(
api_key=self._api_key,
instance=instance,
autoscaler_url=autoscaler_url,
webserver_url=self._vast_server,
debug=debug,
connection_limit=self._connection_limit,
default_request_timeout=default_request_timeout,
max_poll_interval=max_poll_interval,
)
sl._transport_owner = self
return sl
async_/results.py
from __future__ import annotations
from typing import TYPE_CHECKING, Optional, Union
from vastai.data.offer import Offer
from vastai.data.instance import InstanceConfig
if TYPE_CHECKING:
from vastai.async_.client import AsyncClient
from vastai.data.instance import Instance
class AsyncOffer:
def __init__(self, data: Offer, client: AsyncClient):
self._data = data
self._client = client
async def create_instance(self, config: InstanceConfig) -> AsyncInstance:
"""Create an instance from this offer. Synonym for AsyncClient.create_instance(self.id, config)."""
return await self._client.create_instance(self._data.id, config)
def __getattr__(self, name: str):
return getattr(self._data, name)
def __repr__(self) -> str:
return f"AsyncOffer({self._data!r})"
class AsyncInstance:
def __init__(self, data: Instance, client: AsyncClient):
self._data = data
self._client = client
async def destroy(self) -> None:
"""Synonym for AsyncClient.destroy_instance(self)."""
return await self._client.destroy_instance(self._data.id)
def __getattr__(self, name: str):
return getattr(self._data, name)
def __repr__(self) -> str:
return f"AsyncInstance({self._data!r})"
cli/__init__.py
"""Vast.ai CLI layer."""
cli/commands/__init__.py
"""CLI command modules.
Each module registers its commands with the global parser at import time
via the ``@parser.command(...)`` decorator pattern. Importing a module
is sufficient to register all of its commands.
"""
def register_all_commands(parser):
"""Import all command modules to register their commands with the parser.
The imports themselves trigger the decorator registrations -- no
explicit ``register()`` call is needed per module.
"""
from vastai.cli.commands import ( # noqa: F401
instances, offers, machines, teams, keys, endpoints,
billing, storage, auth, misc, deployments, metrics,
benchmarks,
price_increase,
# clusters, # cluster/overlay commands disabled for now
)
cli/commands/auth.py
"""CLI commands for authentication and 2FA."""
import getpass
import json
import os
import re
import sys
import tempfile
from datetime import datetime, timezone
from vastai.cli.parser import argument
from vastai.cli.display import deindent, display_table
from vastai.api import auth as auth_api
from vastai.cli.util import SUCCESS, WARN, FAIL, format_key_suffix
# ---------------------------------------------------------------------------
# TFA helper functions (ported from old/vast.py)
# ---------------------------------------------------------------------------
TFA_METHOD_FIELDS = (
("id", "ID", "{}", None, True),
("user_id", "User ID", "{}", None, True),
("is_primary", "Primary", "{}", None, True),
("method", "Method", "{}", None, True),
("label", "Label", "{}", None, True),
("phone_number", "Phone Number", "{}", None, False),
("created_at", "Created (UTC)", "{}", lambda x: datetime.fromtimestamp(x, tz=timezone.utc).strftime('%Y-%m-%d %H:%M:%S') if x else "N/A", True),
("last_used", "Last Used (UTC)", "{}", lambda x: datetime.fromtimestamp(x, tz=timezone.utc).strftime('%Y-%m-%d %H:%M:%S') if x else "Never", True),
("fail_count", "Failures", "{}", None, True),
("locked_until", "Locked Until (UTC)", "{}", lambda x: datetime.fromtimestamp(x, tz=timezone.utc).strftime('%Y-%m-%d %H:%M:%S') if x else "N/A", True),
)
def display_tfa_methods(methods):
"""Helper function to display 2FA methods in a table."""
method_fields = TFA_METHOD_FIELDS
has_sms = any(m['method'] == 'sms' for m in methods)
if not has_sms: # Don't show Phone Number column if the user has no SMS methods
method_fields = tuple(field for field in TFA_METHOD_FIELDS if field[0] != 'phone_number')
display_table(methods, method_fields, replace_spaces=False)
def confirm_destructive_action(prompt="Are you sure? (y/n): "):
"""Prompt user for confirmation of destructive actions."""
try:
response = input(f" {prompt}").strip().lower()
return 'y' in response
except (EOFError, KeyboardInterrupt):
print("\nOperation cancelled.")
raise
def handle_failed_tfa_verification(args, e):
"""Parse TFA error responses and provide contextual error messages."""
try:
error_data = e.response.json()
except Exception:
print(f"\n{FAIL} Error: {e}")
return
error_msg = error_data.get("msg", str(e))
error_code = error_data.get("error", "")
if args.raw:
print(json.dumps(error_data, indent=2))
print(f"\n{FAIL} Error: {error_msg}")
# Provide helpful context for common errors
if error_code in {"tfa_locked", "2fa_verification_failed"}:
fail_count = error_data.get("fail_count", 0)
locked_until = error_data.get("locked_until")
if fail_count > 0:
print(f" Failed attempts: {fail_count}")
if locked_until:
lock_time_sec = (datetime.fromtimestamp(locked_until) - datetime.now()).seconds
minutes, seconds = divmod(lock_time_sec, 60)
print(f" Time Remaining for 2FA Lock: {minutes} minutes and {seconds} seconds...")
elif error_code == "2fa_expired":
# Note: Only SMS uses tfa challenges that expire when verifying
print("\n The SMS code and secret have expired. Please start over:")
print(" vastai tfa send-sms")
def format_backup_codes(backup_codes):
"""Format backup codes for display or file output."""
output_lines = [
"=" * 60, " VAST.AI 2FA BACKUP CODES", "=" * 60,
f"\nGenerated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
f"\n{WARN} WARNING: All previous backup codes are now invalid!",
"\nYour New Backup Codes (one-time use only):",
"-" * 40,
]
for i, code in enumerate(backup_codes, 1):
output_lines.append(f" {i:2d}. {code}")
output_lines.extend([
"-" * 40,
"\nIMPORTANT:",
" \u2022 Each code can only be used once",
" \u2022 Store them in a secure location",
" \u2022 Use these codes to log in if you lose access to your 2FA device",
"\n" + "=" * 60,
])
return "\n".join(output_lines)
def _save_to_file(content, filepath):
"""Save content to file, creating parent directories if needed."""
try:
filepath = os.path.abspath(os.path.expanduser(filepath))
parent_dir = os.path.dirname(filepath)
if parent_dir:
os.makedirs(parent_dir, exist_ok=True)
with open(filepath, "w") as f:
f.write(content)
return True
except (IOError, OSError) as e:
print(f"\n{FAIL} Error saving file: {e}")
return False
def _get_backup_codes_filename():
"""Generate a timestamped filename for backup codes."""
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
return f"vastai_backup_codes_{timestamp}.txt"
def save_backup_codes(backup_codes):
"""Save or display 2FA backup codes based on user choice."""
print(f"\nBackup codes regenerated successfully! {SUCCESS}")
print(f"\n{WARN} WARNING: All previous backup codes are now invalid!")
formatted_content = format_backup_codes(backup_codes)
filename = _get_backup_codes_filename()
while True:
print("\nHow would you like to save your new backup codes?")
print(f" 1. Save to default location (~/Downloads/{filename})")
print(f" 2. Save to a custom path")
print(f" 3. Print to screen ({WARN} potentially unsafe - visible to onlookers)")
try:
choice = input("\nEnter choice (1-3): ").strip()
if choice in {'1', '2'}:
if choice == '1':
downloads_dir = os.path.expanduser("~/Downloads")
filepath = os.path.join(downloads_dir, filename)
else: # choice == '2'
custom_path = input("\nEnter full path for backup codes file: ").strip()
if not custom_path:
print("Error: Path cannot be empty")
continue
filepath = os.path.abspath(os.path.expanduser(custom_path))
if os.path.isdir(filepath):
filepath = os.path.join(filepath, filename)
if _save_to_file(formatted_content, filepath):
print(f"\n{SUCCESS} Backup codes saved to: {filepath}")
print(f"\nIMPORTANT:")
print(f" \u2022 The file contains {len(backup_codes)} one-time use backup codes")
if choice == '1':
print(f" \u2022 Move this file to a secure location")
return
else:
print("Please try again with a different path.")
continue
elif choice == '3':
print(f"\n{WARN} WARNING: Printing sensitive codes to screen!")
confirm = input("\nAre you sure? Anyone nearby can see these codes. (yes/no): ").strip().lower()
if confirm in {'yes', 'y'}:
print("\n" + formatted_content + "\n")
return
else:
print("Cancelled. Please choose another option.")
continue
else:
print("Invalid choice. Please enter 1, 2, or 3.")
except (EOFError, KeyboardInterrupt):
print("\n\nOperation cancelled. Your backup codes were generated but not saved.")
print("You will need to regenerate them to get new codes.")
raise
from vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401
parser = _get_parser()
# ---------------------------------------------------------------------------
# set api-key
# ---------------------------------------------------------------------------
# Vast API keys are 64 hex characters.
API_KEY_PATTERN = re.compile(r"[0-9a-fA-F]{64}")
def read_api_key_input():
"""The key from a hidden prompt at a terminal; empty without one, so a script
whose stdin is an open pipe fails at once instead of waiting on it."""
if not sys.stdin.isatty():
return ""
try:
return getpass.getpass("Paste your API key (input is hidden): ").strip()
except (EOFError, KeyboardInterrupt):
print()
return ""
def write_private_file(path, text):
"""Write ``text`` to ``path`` readable by the owner only from the first
byte, via a temp file renamed into place so it is never half-written."""
fd, tmp = tempfile.mkstemp(dir=os.path.dirname(path), prefix=".vast_api_key.")
try:
with os.fdopen(fd, "wb") as f: # binary: no \n -> \r\n translation on Windows
f.write(text.encode())
os.replace(tmp, path)
except BaseException:
os.unlink(tmp)
raise
@parser.command(
argument("new_api_key", nargs="?", help="API key to save. Leave it out at a terminal to paste it at a hidden prompt, which keeps it out of your shell history"),
usage="vastai set api-key [APIKEY]",
help="Set api-key (get your api-key from the console/CLI)",
)
def set__api_key(args):
"""Set the api-key."""
from vastai.cli.util import APIKEY_FILE, APIKEY_FILE_HOME
# Exit rather than return 1: under --raw a returned value is printed as JSON and exits 0.
api_key = args.new_api_key.strip() if args.new_api_key is not None else read_api_key_input()
if not api_key:
print(f"\n{FAIL} Error: No API key given. Pass it (vastai set api-key KEY) or run this at a terminal to be prompted. Create one at https://cloud.vast.ai/manage-keys/?tab=api-keys")
sys.exit(1)
if not API_KEY_PATTERN.fullmatch(api_key):
print(f"\n{FAIL} Error: That isn't a Vast API key (expected 64 hex characters), so it wasn't saved.")
sys.exit(1)
write_private_file(APIKEY_FILE, api_key)
print("Your api key has been saved in {}".format(APIKEY_FILE))
if os.path.exists(APIKEY_FILE_HOME):
os.remove(APIKEY_FILE_HOME)
if getattr(args, "explain", False):
print("Removed legacy api key file at {} (no longer used).".format(APIKEY_FILE_HOME))
env_key = os.environ.get("VAST_API_KEY")
if env_key:
print(f"\n{WARN} VAST_API_KEY is set in your environment (ends in {format_key_suffix(env_key)}) and overrides the key you just saved.")
print("Unset VAST_API_KEY to use the saved key.")
# ---------------------------------------------------------------------------
# set role
# ---------------------------------------------------------------------------
@parser.command(
argument("role", help="'host' or 'client'", choices=["host", "client"]),
usage="vastai set role host|client",
help="Set host role to view host-only commands",
epilog=deindent("""
This only changes what --help and tab completion display — every command
still runs regardless of role, since the server enforces real permissions.
"""),
)
def set__role(args):
"""Set the client/host CLI display role."""
from vastai.cli.util import set_role_file
set_role_file(args.role)
print(f"CLI role set to '{args.role}'.")
if args.role == "client":
print("Host-only commands are now hidden from --help and tab completion.")
else:
print("All commands (client + host) are now shown in --help and tab completion.")
# ---------------------------------------------------------------------------
# Note: set__user and show__user are in billing.py.
# Note: show__ipaddrs is in billing.py.
# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------
# show audit-logs
# ---------------------------------------------------------------------------
@parser.command(
usage="vastai show audit-logs [--api-key API_KEY] [--raw]",
help="Display account's history of important actions",
)
def show__audit_logs(args):
"""Show the history of important actions."""
from vastai.cli.display import display_table, audit_log_fields
client = get_client(args)
rows = auth_api.show_audit_logs(client)
if args.raw:
return rows
else:
display_table(rows, audit_log_fields)
# ---------------------------------------------------------------------------
# environment variables
# ---------------------------------------------------------------------------
@parser.command(
argument("-s", "--show-values", action="store_true", help="Show the values of environment variables"),
usage="vastai show env-vars [-s]",
help="Show user environment variables",
)
def show__env_vars(args):
"""Show the environment variables for the current user."""
client = get_client(args)
env_vars = auth_api.show_env_vars(client)
if args.raw:
if not args.show_values:
masked_env_vars = {k: "*****" for k, v in env_vars.items()}
return masked_env_vars
else:
return env_vars
else:
if not env_vars:
print("No environment variables found.")
else:
for key, value in env_vars.items():
print(f"Name: {key}")
if args.show_values:
print(f"Value: {value}")
else:
print("Value: *****")
print("---")
if not args.show_values:
print("\nNote: Values are hidden. Use --show-values or -s option to display them.")
@parser.command(
argument("name", help="Environment variable name", type=str),
argument("value", help="Environment variable value", type=str),
usage="vastai create env-var <name> <value>",
help="Create a new user environment variable",
)
def create__env_var(args):
"""Create a new environment variable for the current user."""
client = get_client(args)
result = auth_api.create_env_var(client, name=args.name, value=args.value)
if result.get("success"):
print(result.get("msg", "Environment variable created successfully."))
else:
print(f"Failed to create environment variable: {result.get('msg', 'Unknown error')}")
@parser.command(
argument("name", help="Environment variable name to update", type=str),
argument("value", help="New environment variable value", type=str),
usage="vastai update env-var <name> <value>",
help="Update an existing user environment variable",
)
def update__env_var(args):
"""Update an existing environment variable for the current user."""
client = get_client(args)
result = auth_api.update_env_var(client, name=args.name, value=args.value)
if result.get("success"):
print(result.get("msg", "Environment variable updated successfully."))
else:
print(f"Failed to update environment variable: {result.get('msg', 'Unknown error')}")
@parser.command(
argument("name", help="Environment variable name to delete", type=str),
usage="vastai delete env-var <name>",
help="Delete a user environment variable",
)
def delete__env_var(args):
"""Delete an environment variable for the current user."""
client = get_client(args)
result = auth_api.delete_env_var(client, name=args.name)
if result.get("success"):
print(result.get("msg", "Environment variable deleted successfully."))
else:
print(f"Failed to delete environment variable: {result.get('msg', 'Unknown error')}")
# ---------------------------------------------------------------------------
# Note: show__scheduled_jobs and delete__scheduled_job are in billing.py.
# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------
# 2FA commands
# ---------------------------------------------------------------------------
@parser.command(
argument("code", help="6-digit verification code from SMS or Authenticator app", type=str),
argument("-t", "--method-type", choices=["sms", "totp"], help="New 2FA Method type to activate", type=str, default="totp"),
argument("--secret", help="Secret token from setup process (required)", type=str, required=True),
argument("--phone-number", help="Phone number for SMS method (E.164 format)", type=str, default=None),
argument("-l", "--label", help="Label for the new 2FA method", type=str, default=None),
usage="vastai tfa activate CODE --secret SECRET [--method-type METHOD_TYPE] [--phone-number PHONE_NUMBER] [--label LABEL]",
help="Activate a new 2FA method by verifying the code",
epilog=deindent(f"""
Complete the 2FA setup process by verifying your code.
{'*'*120}
NOTE: Prior to running this command, you must authorize your attempt to create a new 2FA method by following the instructions in the `vastai tfa auth-new` command.
This is required to ensure that only you can add new 2FA methods to your account.
{'*'*120}
For TOTP (Authenticator app):
1. Run 'vastai tfa totp-setup' to get the manual key/QR code and secret
2. Enter the manual key or scan the QR code with your Authenticator app
3. Run this command with the 6-digit code from your app and the secret token from step 1
For SMS:
1. Run 'vastai tfa send-sms --phone-number <PHONE_NUMBER>' to receive SMS and get secret token
2. Run this command with the code you received via SMS and the phone number it was sent to
If this is your first 2FA method, backup codes will be generated and displayed.
Save these backup codes in a secure location!
Examples:
vastai tfa activate --method-type totp --secret abc123def456 123456
vastai tfa activate --method-type sms --secret abc123def456 --phone-number +12345678901 123456
vastai tfa activate --method-type sms --secret abc123def456 --phone-number +12345678901 --label "Work Phone" 123456
"""),
)
def tfa__activate(args):
"""Activate a new 2FA method by confirming the verification code."""
client = get_client(args)
response_data = auth_api.tfa_activate(
client, code=args.code, secret=args.secret, method_type=args.method_type,
phone_number=args.phone_number, label=args.label,
)
method_name = "SMS" if args.phone_number or args.method_type == "sms" else "TOTP (Authenticator App)"
print(f"\n{SUCCESS} {method_name} 2FA method activated successfully!")
if "backup_codes" in response_data:
save_backup_codes(response_data["backup_codes"])
@parser.command(
argument("-id", "--id-to-delete", help="ID of the 2FA method to delete (see `vastai tfa status`)", type=int, default=None),
argument("-c", "--code", mutex_group='code_grp', required=True, help="2FA code from your Authenticator app, SMS, or Email to authorize deletion", type=str),
argument("-t", "--method-type", mutex_group="type_grp", choices=["email", "sms", "totp"],
help="2FA Method type. Only use when you only have one method of this type", type=str, default=None),
argument("-s", "--secret", help="Secret token (required for SMS or Email 2FA)", type=str, default=None),
argument("-bc", "--backup-code", mutex_group='code_grp', required=True, help="One-time backup code (alternative to regular 2FA code)", type=str, default=None),
argument("--method-id", mutex_group="type_grp", help="2FA Method ID to use if you have more than one of the same type ('id' from `tfa status`)", type=str, default=None),
usage="vastai tfa delete [--id-to-delete ID] [--code CODE] [--method-type METHOD_TYPE] [--secret SECRET] [--backup-code BACKUP_CODE] [--method-id ID]",
help="Remove a 2FA method from your account",
epilog=deindent(f"""
Remove a 2FA method from your account.
This action requires 2FA verification to prevent unauthorized removals.
{'*'*120}
NOTE: If you do not specify --id-to-delete, the system will attempt to delete the method you are using to authenticate.
However please be advised, it is much safer to specify the ID to avoid confusion if you have multiple methods.
{'*'*120}
Use `vastai tfa status` to see your active methods and their IDs.
Examples:
# Delete method #123, authorize with email code and secret from `tfa send-email`
vastai tfa delete --id-to-delete 123 --method-type email --secret abc123def456 -c 456789
# Delete method #123, authorize with TOTP/Authenticator code
vastai tfa delete --id-to-delete 123 --method-type totp --code 456789
# Delete method #123, authorize with SMS and secret from `tfa send-sms`
vastai tfa delete -id 123 --method-type sms --secret abc123def456 -c 456789
# Delete method #123, authorize with backup code
vastai tfa delete --id-to-delete 123 --backup-code ABCD-EFGH-IJKL
# Delete method #123, specify which TOTP method to use if you have multiple
vastai tfa delete -id 123 --method-id 456 -c 456789
# Delete the TOTP method you are using to authenticate (use with caution)
vastai tfa delete -c 456789
"""),
)
def tfa__delete(args):
"""Remove a 2FA method from the user's account."""
if args.method_type in ("sms", "email") and not args.secret:
print(f"\n{FAIL} Error: --secret is required for deletion authorization when using --method-type {args.method_type}.")
print("\nPlease use: `vastai tfa send-sms` or `vastai tfa send-email` to get the missing secret and try again.")
return 1
prompt = "\nAre you sure you want to delete this 2FA method? (y|n): "
if confirm_destructive_action(prompt) == False:
print("Operation cancelled.")
return
client = get_client(args)
try:
response_data = auth_api.tfa_delete(
client, code=args.code, backup_code=args.backup_code,
method_type=args.method_type, secret=args.secret, method_id=args.method_id,
id_to_delete=args.id_to_delete,
)
print(f"\n{SUCCESS} 2FA method deleted successfully.")
if "remaining_methods" in response_data:
remaining = response_data["remaining_methods"]
print(f"\nYou have {remaining} 2FA method{'s' if remaining != 1 else ''} remaining.")
else:
print(f"\n{WARN} WARNING: You have removed all 2FA methods from your account.")
print("Your backup codes have been invalidated and 2FA is now fully disabled.")
except Exception as e:
if hasattr(e, 'response'):
handle_failed_tfa_verification(args, e)
else:
print(f"Failed to delete 2FA method: {e}")
return 1
@parser.command(
argument("-c", "--code", mutex_group='code_grp', required=True, help="2FA code from Authenticator app, SMS, or Email", type=str),
argument("-t", "--method-type", mutex_group="type_grp", choices=["email", "sms", "totp"],
help="2FA Method type. Only use when you only have one method of this type", type=str, default=None),
argument("-s", "--secret", help="Secret token from previous login step (required for SMS or Email 2FA)", type=str, default=None),
argument("-bc", "--backup-code", mutex_group='code_grp', required=True, help="One-time backup code (alternative to regular 2FA code)", type=str, default=None),
argument("-id", "--method-id", mutex_group="type_grp", help="2FA Method ID if you have more than one of the same type ('id' from `tfa status`)", type=str, default=None),
usage="vastai tfa login [--code CODE] [--method-type METHOD_TYPE] [--secret SECRET] [--backup-code BACKUP_CODE]",
help="Complete 2FA login by verifying code",
epilog=deindent("""
Complete Two-Factor Authentication login by providing the 2FA code.
For Email: Include the --method-type email flag and provide -s/--secret from the `tfa send-email` command response
For TOTP: Include the --method-type totp and provide the 6-digit code from your Authenticator app
For SMS: Include the --method-type sms flag and provide -s/--secret from the `tfa send-sms` command response
For backup code: Use --backup-code instead of code (codes may only be used once)
Examples:
vastai tfa login --method-type totp -c 123456
vastai tfa login --method-type sms --code 123456 --secret abc123def456
vastai tfa login -t email -c 123456 -s abc123def456
vastai tfa login --backup-code ABCD-EFGH-IJKL
"""),
)
def tfa__login(args):
"""Complete 2FA login and store the session key."""
from vastai.cli.util import TFAKEY_FILE
client = get_client(args)
try:
response_data = auth_api.tfa_login(
client, code=args.code, backup_code=args.backup_code,
method_type=args.method_type, secret=args.secret, method_id=args.method_id,
)
if "session_key" in response_data:
session_key = response_data["session_key"]
if session_key != args.api_key:
with open(TFAKEY_FILE, "w") as f:
f.write(session_key)
print(f"{SUCCESS} 2FA login successful! Session key saved to {TFAKEY_FILE}")
else:
print(f"{SUCCESS} 2FA login successful! Your session key has been refreshed.")
# Display remaining backup codes if present
if "backup_codes_remaining" in response_data:
remaining = response_data["backup_codes_remaining"]
if remaining == 0:
print(f"{WARN} Warning: You have no backup codes remaining! Please generate new backup codes immediately to avoid being locked out of your account if you lose access to your 2FA device.")
elif remaining <= 3:
print(f"{WARN} Warning: You only have {remaining} backup codes remaining. Consider regenerating them.")
else:
print(f"Backup codes remaining: {remaining}")
else:
print("2FA login successful but a session key was not returned. Please check that you have an API Key that's properly set up")
except Exception as e:
if hasattr(e, 'response'):
handle_failed_tfa_verification(args, e)
else:
print(f"2FA login failed: {e}")
return 1
@parser.command(
argument("-p", "--phone-number", help="Phone number to receive SMS code (E.164 format, e.g., +1234567890)", type=str, default=None),
argument("-s", "--secret", help="Secret token from the original 2FA login attempt", type=str, required=True),
usage="vastai tfa resend-sms --secret SECRET [--phone-number PHONE_NUMBER]",
help="Resend SMS 2FA code",
epilog=deindent("""
Resend the SMS verification code to your phone.
This is useful if:
\u2022 You didn't receive the original SMS
\u2022 The code expired before you could use it
\u2022 You accidentally deleted the message
You must provide the same secret token from the original request.
Example:
vastai tfa resend-sms --secret abc123def456
"""),
)
def tfa__resend_sms(args):
"""Resend SMS 2FA code to the user's phone."""
client = get_client(args)
response_data = auth_api.tfa_resend_sms(client, secret=args.secret, phone_number=args.phone_number)
print("SMS code resent successfully!")
print(f"\n{response_data['msg']}")
@parser.command(
argument("-c", "--code", mutex_group='code_grp', required=True, help="2FA code from Authenticator app, SMS, or Email", type=str),
argument("-t", "--method-type", mutex_group="type_grp", choices=["email", "sms", "totp"],
help="2FA Method type. Only use when you only have one method of this type", type=str, default=None),
argument("-s", "--secret", help="Secret token from previous login step (required for SMS or Email 2FA)", type=str, default=None),
argument("-bc", "--backup-code", mutex_group='code_grp', required=True, help="One-time backup code (alternative to regular 2FA code)", type=str, default=None),
argument("-id", "--method-id", mutex_group="type_grp", help="2FA Method ID if you have more than one of the same type ('id' from `tfa status`)", type=str, default=None),
usage="vastai tfa regen-codes [--code CODE] [--method-type METHOD_TYPE] [--secret SECRET] [--backup-code BACKUP_CODE] [--method-id ID]",
help="Regenerate backup codes for 2FA",
epilog=deindent("""
Generate a new set of backup codes for your account.
This action requires 2FA verification to prevent unauthorized regeneration.
WARNING: This will invalidate all existing backup codes!
Any previously generated codes will no longer work.
Backup codes are one-time use codes that allow you to log in
if you lose access to your primary 2FA method (lost phone, etc).
You should regenerate your backup codes if:
\u2022 You've used several codes and are running low
\u2022 You think your codes may have been compromised
\u2022 You lost your saved codes and need new ones
Important: Save the new codes in a secure location immediately!
They will not be shown again.
Examples:
vastai tfa regen-codes --code --method-type totp 123456
vastai tfa regen-codes -c 123456 -t sms --secret abc123def456
vastai tfa regen-codes --backup-code ABCD-EFGH-IJKL
"""),
)
def tfa__regen_codes(args):
"""Regenerate backup codes for 2FA recovery."""
prompt = "\nThis will invalidate all existing backup codes. Continue? (y|n): "
if confirm_destructive_action(prompt) == False:
print("Operation cancelled.")
return
client = get_client(args)
try:
response_data = auth_api.tfa_regen_codes(
client, code=args.code, backup_code=args.backup_code,
method_type=args.method_type, secret=args.secret, method_id=args.method_id,
)
if "backup_codes" in response_data:
save_backup_codes(response_data["backup_codes"])
else:
print(f"\n{SUCCESS} Backup codes regenerated successfully!")
print("(No codes returned in response - this may be an error)")
except Exception as e:
if hasattr(e, 'response'):
handle_failed_tfa_verification(args, e)
else:
print(f"Failed to regenerate backup codes: {e}")
return 1
@parser.command(
argument("-p", "--phone-number", help="Phone number to receive SMS code (E.164 format, e.g., +1234567890)", type=str, default=None),
usage="vastai tfa send-sms [--phone-number PHONE_NUMBER]",
help="Request a 2FA SMS verification code",
epilog=deindent("""
Request a two-factor authentication code to be sent via SMS.
If --phone-number is not provided, uses the phone number on your account.
The secret token will be returned and must be used with 'vastai tfa activate'.
Examples:
vastai tfa send-sms
vastai tfa send-sms --phone-number +12345678901
"""),
)
def tfa__send_sms(args):
"""Request a 2FA SMS code to be sent to the user's phone."""
client = get_client(args)
response_data = auth_api.tfa_send_sms(client, phone_number=args.phone_number)
secret = response_data.get("secret", "")
print("SMS code sent successfully!")
print(f" Secret token: {secret}")
print(f"\nOnce you receive the SMS code:")
phone_num = f"--phone-number {args.phone_number}" if args.phone_number else "[--phone-number <PHONE_NUMBER>]"
print(f"\n If you are setting up SMS 2FA for the first time, run:")
print(f" vastai tfa activate --method-type sms --secret {secret} {phone_num} [--label <LABEL>] <CODE>")
print(f"\n Otherwise you can complete your 2FA log in with:")
print(f" vastai tfa login --method-type sms --secret {secret} -c <CODE>\n")
@parser.command(
usage="vastai tfa send-email",
help="Request a 2FA Email verification code",
epilog=deindent("""
Request a two-factor authentication code to be sent via Email.
The secret token will be returned and must be used with 'vastai tfa activate'.
Examples:
vastai tfa send-email
"""),
)
def tfa__send_email(args):
"""Request a 2FA email code to be sent to the user's email."""
client = get_client(args)
response_data = auth_api.tfa_send_email(client)
secret = response_data.get("secret", "")
print(f"{SUCCESS} Email code sent successfully!")
print(f" Secret token: {secret}")
print(f"\nOnce you receive the Email code:")
print(f"\n You can complete your 2FA log in with:")
print(f" vastai tfa login --method-type email --secret {secret} -c <CODE>\n")
def print_next_steps_after_new_method_auth():
print(f"\nNext Steps:"
"\n To add a new SMS 2FA method:"
"\n 1. Run `vastai tfa send-sms --phone-number <PHONE_NUMBER>` to receive SMS and get secret token"
"\n 2. Run `vastai tfa activate --method-type sms --secret <SECRET> --phone-number <PHONE_NUMBER> CODE` to activate the new method with the code you received via SMS\n"
"\n To add a new TOTP (Authenticator app) 2FA method:"
"\n 1. Run `vastai tfa totp-setup` to get the manual key/QR code and secret"
"\n 2. Enter the manual key or scan the QR code with your Authenticator app"
"\n 3. Run `vastai tfa activate --method-type totp --secret <SECRET> CODE` to activate the new method with the 6-digit code from your app")
@parser.command(
argument("-c", "--code", help="2FA code from Authenticator app, SMS, or Email", type=str),
argument("-s", "--secret", help="Secret token from previous auth step", type=str, default=None),
argument("-t", "--method-type", mutex_group="type_grp", choices=["email", "sms", "totp"],
help="2FA Method type. Only use when you only have one method of this type", type=str, default="email"),
argument("-bc", "--backup-code", mutex_group='type_grp',
help="One-time backup code (alternative to regular 2FA code)", type=str, default=None),
argument("-id", "--method-id", mutex_group="type_grp",
help="2FA Method ID if you have more than one of the same type ('id' from `tfa status`)", type=str, default=None),
usage="vastai tfa auth-new {[--method-type METHOD_TYPE | --method-id ID | --backup-code BACKUP_CODE] | [--secret SECRET --code CODE]}",
help="Authorize your account to add a new 2FA method",
epilog=deindent("""
Authorize your account to add a new 2FA method by verifying via email or an existing method.
This is a required step to ensure that only you can add new 2FA methods to your account.
Step 1. Run command with your chosen verification method:
- If you have an existing 2FA method set up, you can use '--backup-code BACKUP_CODE' to immediately authorize (skip Step 2)
- Use --method-type {sms|totp} or --method-id ID to specify which existing method to use for verification (see `vastai tfa status` for available methods and their IDs)
- Use --method-type email if you have a verified email address and/or no other 2FA methods set up to receive a code via email
Step 2. When prompted, enter the 2FA code from your 2FA method of choice to confirm authorization.
Note:
If you exit the command before being able to enter the code, you can run this command again
with --secret SECRET and --code CODE to complete the authorization step as long as the code has not expired.
Examples:
# Initiating New Method Authorization
vastai tfa auth-new (method type is email by default)
vastai tfa auth-new --method-type totp
vastai tfa auth-new -t sms
vastai tfa auth-new --method-id 456
vastai tfa auth-new --backup-code ABCD-EFGH-IJKL
# Completing Authorization with code and secret if not completed in previous run
vastai tfa auth-new --secret abc123def456 --code 123456
"""),
)
def tfa__auth_new(args):
"""Authorize the user to add a new 2FA method by verifying with an existing method."""
client = get_client(args)
secret, code = args.secret, args.code
if not secret and not code:
try:
response_data = auth_api.tfa_auth_new(
client, backup_code=args.backup_code,
method_type=args.method_type, method_id=args.method_id,
)
except Exception as e:
if hasattr(e, 'response'):
handle_failed_tfa_verification(args, e)
else:
print(f"Authorization failed: {e}")
return 1
if args.backup_code and response_data.get("msg") == "Authorization successful.":
print(f"\n{SUCCESS} Successfully authorized account for adding new 2FA method using backup code")
print_next_steps_after_new_method_auth()
return 0
secret = response_data.get("secret")
if not secret:
print(f"\n{FAIL} Error: No secret token received for authorization. Please try again.")
return 1
print(f"\n{SUCCESS} Authorization initiated successfully.")
print(f"2FA Secret: {secret}")
try:
code = input("Enter 2FA code to complete authorization: ").strip()
except (EOFError, KeyboardInterrupt):
print("\nOperation cancelled.")
print("You can still complete this authorization later by running:"
f"\n vastai tfa auth-new --secret {secret} --code <CODE>")
return 1
try:
auth_api.tfa_auth_new(client, code=code, secret=secret)
print(f"\n{SUCCESS} Successfully authorized account for adding new 2FA method!")
print_next_steps_after_new_method_auth()
except Exception as e:
if hasattr(e, 'response'):
handle_failed_tfa_verification(args, e)
else:
print(f"Authorization failed: {e}")
print(f"\n{FAIL} Authorization failed. Please try again.")
return 1
@parser.command(
help="Shows the current 2FA status and configured methods",
epilog=deindent("""
Show the current 2FA status for your account, including:
\u2022 Whether or not 2FA is enabled
\u2022 A list of active 2FA methods
\u2022 The number of backup codes remaining (if 2FA is enabled)
""")
)
def tfa__status(args):
"""Show the current 2FA status for the user."""
client = get_client(args)
response_data = auth_api.tfa_status(client)
if args.raw:
return response_data
tfa_enabled = response_data.get("tfa_enabled", False)
methods = response_data.get("methods", [])
backup_codes_remaining = response_data.get("backup_codes_remaining", 0)
if not tfa_enabled or not methods:
print(f"{WARN} No active 2FA methods found")
else:
print(f"2FA Status: Enabled {SUCCESS}")
print(f"\nActive 2FA Methods:")
display_tfa_methods(methods)
print(f"\nBackup codes remaining: {backup_codes_remaining}")
@parser.command(
usage="vastai tfa totp-setup",
help="Generate TOTP secret and QR code for Authenticator app setup",
epilog=deindent("""
Set up TOTP (Time-based One-Time Password) 2FA using an Authenticator app.
This command generates a new TOTP secret and displays:
\u2022 A QR code (for scanning with your app)
\u2022 A manual entry key (for typing into your app)
\u2022 A secret token (needed for the next step)
Workflow:
1. Run this command to generate the TOTP secret
2. Add the account to your Authenticator app by either:
\u2022 Scanning the displayed QR code, OR
\u2022 Manually entering the key shown
3. Once added, your app will display a 6-digit code
4. Complete setup by running:
vastai tfa activate --method-type totp --secret <SECRET> <CODE>
Supported Authenticator Apps:
\u2022 Google Authenticator
\u2022 Microsoft Authenticator
\u2022 Authy
\u2022 1Password
\u2022 Any TOTP-compatible app
Example:
vastai tfa totp-setup
"""),
)
def tfa__totp_setup(args):
"""Generate a TOTP secret and QR code for setting up Authenticator app 2FA."""
client = get_client(args)
response_data = auth_api.tfa_totp_setup(client)
if args.raw:
return response_data
secret = response_data.get("secret", "")
provisioning_uri = response_data.get("provisioning_uri", "")
print("\n" + "=" * 60)
print("TOTP (Authenticator App) 2FA Setup")
print("=" * 60)
print("\nScan this QR code with your Authenticator app:\n")
try:
import qrcode
qr = qrcode.QRCode(border=2)
qr.add_data(provisioning_uri)
qr.make()
qr.print_ascii(tty=True)
except ImportError:
print(" [QR code display requires 'qrcode' package]")
print(f" Install with: pip install qrcode")
print(f"\n Or manually enter this URI in your app:")
print(f" {provisioning_uri}")
print("\nOR Manual Entry Key (type this into your Authenticator app):")
print(f" {secret}")
print("\nNext Steps:")
print(" 1. Your Authenticator app should now display a 6-digit code")
print(" 2. Complete setup by running:")
print(f" vastai tfa activate --method-type totp --secret {secret} <CODE>")
print("\n" + "=" * 60 + "\n")
@parser.command(
argument("method_id", metavar="METHOD_ID", help="ID of the 2FA method to update (see `vastai tfa status`)", type=int),
argument("-l", "--label", help="New label/name for this 2FA method", type=str, default=None),
argument("-p", "--set-primary", help="Set this method as the primary/default 2FA method", default=None),
usage="vastai tfa update METHOD_ID [--label LABEL] [--set-primary]",
help="Update a 2FA method's settings",
epilog=deindent("""
Update the label or primary status of a 2FA method.
The label is a friendly name to help you identify different methods
(e.g. "Work Phone", "Personal Authenticator").
The primary method is your preferred/default 2FA method.
Examples:
vastai tfa update 123 --label "Work Phone"
vastai tfa update 456 --set-primary
vastai tfa update 789 --label "Backup Authenticator" --set-primary
"""),
)
def tfa__update(args):
"""Update settings for an existing 2FA method."""
label = args.label
set_primary = args.set_primary
if set_primary is not None:
if set_primary.lower() in {'true', 't'}:
set_primary = True
elif set_primary.lower() in {'false', 'f'}:
set_primary = False
else:
print("Error: --set-primary must be <t|true> or <f|false>")
return
if label is None and set_primary is None:
print("Error: You must specify at least one field to update (--label or --set-primary)")
return 1
client = get_client(args)
try:
response_data = auth_api.tfa_update(client, method_id=args.method_id, label=label, set_primary=set_primary)
except ValueError as e:
print(f"Error: {e}")
return 1
if args.raw:
return response_data
print("\n2FA method updated successfully!")
if args.label:
print(f" New label: {args.label}")
if set_primary is not None:
print(f" Set as primary method = {set_primary}")
cli/commands/benchmarks.py
"""
A CLI command to benchmark a template against multiple GPUs in parallel.
The template's pyworker benchmark determines the measured perf (ex. tokens/$).
vastai run benchmarks
"""
import argparse
import atexit
import functools
import json
import math
import random
import re
import signal
import statistics
import sys
import threading
import time
import uuid
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from rich.console import Console
from rich.live import Live
from rich.table import Table
from vastai import VastAI
from vastai.cli.parser import argument
from vastai.cli.display import deindent
from vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401 (used by test conftest)
from vastai.cli.util import _get_gpu_types
from vastai.data.query import OP_TO_STR
from vastai.data import query as _query_consts
parser = _get_parser()
# Default GPU sweep set + per-card VRAM fallback (MB). Catalog gpu_ram_mb
# overrides these values when populated; the keys stay the editorial sweep set.
DEFAULT_GPUS = {
"RTX 5090": 32607,
"RTX 4090": 24564,
"RTX 3090": 24576,
"RTX A6000": 49140,
}
# One worker per GPU, kept alive while idle so we can read its measured perf.
ENDPOINT_CONFIG = {"cold_workers": 1, "max_workers": 1, "min_load": 1.0}
RENTAL_TIMEOUT = 120
DEFAULT_BENCHMARK_TIMEOUT = 30 * 60
WORKER_POLL_INTERVAL = 10.0
PARALLEL_SUBMIT_DELAY = 4.0 # spacing between parallel endpoint submits
MAX_ABANDONS = 3 # bail after this many workers vanish without one reaching idle
# states in which the autoscaler shouldnt retry the worker
TERMINAL_STATES = {"stopped", "destroying", "unavail"}
TERMINAL_GRACE_SECONDS = 30
TERMINAL_STATUSES = {"done", "skipped", "timeout", "failed", "no_worker"}
def format_time_elapsed(seconds):
m, s = divmod(int(seconds), 60)
return f"{m}:{s:02d}"
@functools.lru_cache(maxsize=1)
def _canonical_gpu_names():
"""Map of lower-cased GPU name -> canonical name, for resolving user input
like "rtx_4090"/"RTX 4090" to the canonical "RTX 4090".
Union of the hardcoded query.py constants (fallback floor) with the
/api/v0/gpu_types/ catalog (picks up newly-onboarded SKUs). On a catalog
fetch failure the constants alone are used, so resolution never gets worse.
"""
canonical = {
v.lower(): v for k, v in vars(_query_consts).items()
if isinstance(v, str) and not k.startswith("_")
}
for gpu_type in (_get_gpu_types() or []):
name = gpu_type.get("canonical_name")
if name:
canonical[name.lower()] = name
return canonical
@functools.lru_cache(maxsize=1)
def _catalog_vram_map():
"""Map canonical_name -> gpu_ram_mb for entries the catalog has populated.
Empty when the catalog is unreachable or the column is unbackfilled.
"""
return {
gpu_type["canonical_name"]: gpu_type["gpu_ram_mb"]
for gpu_type in (_get_gpu_types() or [])
if gpu_type.get("canonical_name") and gpu_type.get("gpu_ram_mb")
}
def _per_card_vram_mb(gpu_name):
"""Per-card VRAM in MB: catalog gpu_ram_mb if populated, else the hardcoded
DEFAULT_GPUS value. None when neither source knows the GPU.
"""
return _catalog_vram_map().get(gpu_name) or DEFAULT_GPUS.get(gpu_name)
def parse_gpu_spec(token, default_num_gpus):
"""Parse a GPU token into (canonical_gpu_name, num_gpus).
Example: 2x RTX_4090" -> ("RTX 4090", 2).
"""
# so "rtx_4090", "RTX 4090", "Rtx_4090" all map to "RTX 4090"
canonical = _canonical_gpu_names()
token = token.strip()
# string to tuple. "2x RTX_4090" to ("RTX 4090", 2)
m = re.match(r"^(\d+)[\s_]*x[\s_]*(.+)$", token, re.IGNORECASE)
if m:
raw_name, n = m.group(2).strip().replace("_", " "), int(m.group(1))
else:
raw_name, n = token.replace("_", " "), default_num_gpus
return (canonical.get(raw_name.lower(), raw_name), n)
def format_filter_query(filters):
"""Render extra_filters dict into comparison string.
e.g. {"gpu_ram": {"gte": 16000}} -> "gpu_ram>=16000"
"""
parts = []
for field, ops in (filters or {}).items():
if not isinstance(ops, dict):
continue
for op, value in ops.items():
sym = OP_TO_STR.get(op, op)
parts.append(f"{field}{sym}{value}")
return ", ".join(parts)
# ------------------------------------------------------------------------------------
# Pre-flight helpers: check whether each (template, gpu_name, num_gpus) is rentable
# before we benchmark it, so a run can't fail late due to 0 matching offers. On a
# miss they build a user-friendly explanation of which filter blocked it.
# ------------------------------------------------------------------------------------
def has_matching_offer(vast, *, gpu_name, num_gpus, extra_filters):
"""True if any verified+rentable offer matches the template's filters"""
query = dict(extra_filters or {})
query["gpu_name"] = {"eq": gpu_name}
query["num_gpus"] = {"eq": num_gpus}
return bool(vast.search_offers(query=query, limit=1))
def format_skip_message(vast, *, gpu_name, num_gpus, extra_filters):
"""Build a user-friendly explanation for why a GPU has 0 matching offers.
Calls search_offers to find the blocking filter.
"""
if not extra_filters:
return (f"no offers for {gpu_name} (num_gpus={num_gpus})")
diag = find_blockers(
vast, gpu_name=gpu_name, num_gpus=num_gpus,
extra_filters=extra_filters,
)
if diag["base_count"] == 0:
return (f"no offers for {gpu_name} (num_gpus={num_gpus}), "
f"independent of template filters")
blockers = diag["single_blockers"]
if blockers:
sample = diag.get("sample_offer") or {}
blocker_lines = []
for key in blockers:
ops = extra_filters.get(key)
if not isinstance(ops, dict):
blocker_lines.append(f" {key}")
continue
gpu_value = sample.get(key)
for op, threshold in ops.items():
if check_template_filter(gpu_value, op, threshold) is True:
continue # search-engine quirk, skip
sym = OP_TO_STR.get(op, op)
if gpu_value is None:
detail = "no sample available to check"
else:
detail = f"{gpu_name} has {gpu_value}"
blocker_lines.append(f" {key}{sym}{threshold}: {detail}")
if blocker_lines:
plural = "s" if len(blocker_lines) > 1 else ""
lines = [f"blocked by template filter{plural}:"] + blocker_lines
# Only hint num_gpus when gpu_total_ram is the sole blocker; otherwise bumping it won't fix the run.
if blockers == ["gpu_total_ram"]:
ops = extra_filters.get("gpu_total_ram") or {}
if ops:
op, value = next(iter(ops.items()))
raw = min_gpus_for_ram(
value, op, diag.get("per_card_gpu_ram"))
if raw and raw > num_gpus:
viable = get_gpu_chunk_size(raw)
lines.append(
f" hint: try {viable}x {gpu_name} "
f"(host total then satisfies {value})")
return "\n".join(lines)
# All "blockers" turned out to be quirks; fall through to generic failure message
return ("0 offers match. No single filter is the culprit; the combination "
f"of template filters excludes {gpu_name} "
f"({format_filter_query(extra_filters)})")
def find_blockers(vast, *, gpu_name, num_gpus, extra_filters):
"""Find which template filter blocks this GPU from having offers.
Bails on the first real blocker. Returns a dict
{base_count, single_blockers, per_card_gpu_ram, sample_offer}.
"""
base_query = {"gpu_name": {"eq": gpu_name},
"num_gpus": {"eq": num_gpus}}
base_offers = vast.search_offers(query=base_query, limit=1)
base = len(base_offers)
sample = base_offers[0] if base_offers else None
per_card = sample.get("gpu_ram") if sample else None
blockers = []
if base > 0:
for key, ops in (extra_filters or {}).items():
# if GPU satisfies the given filter, its not a blocker so skip
if isinstance(ops, dict) and sample is not None:
gpu_value = sample.get(key)
if gpu_value is not None and all(
check_template_filter(gpu_value, op, threshold) is True
for op, threshold in ops.items()
):
continue
if not has_matching_offer(
vast, gpu_name=gpu_name, num_gpus=num_gpus,
extra_filters={key: ops},
):
blockers.append(key)
break
return {"base_count": base, "single_blockers": blockers,
"per_card_gpu_ram": per_card, "sample_offer": sample}
def check_template_filter(gpu_value, op, threshold):
"""Checks if gpu_value matches the template filters (e.g. gpu_total_ram>=16000)"""
if gpu_value is None or threshold is None:
return None
try:
a, t = float(gpu_value), float(threshold)
except (TypeError, ValueError):
return None
if op == "gt": return a > t
if op == "gte": return a >= t
if op == "lt": return a < t
if op == "lte": return a <= t
if op == "eq": return a == t
if op == "neq": return a != t
return None
def min_gpus_for_ram(threshold, op, per_card_gpu_ram):
"""Smallest ``num_gpus`` such that ``num_gpus * per_card_gpu_ram``
satisfies the template's ``gpu_total_ram`` filter. Returns None if any
input is missing or the operator isn't a comparison we can solve.
"""
if not per_card_gpu_ram or not threshold:
return None
if op == "gt":
return int(threshold // per_card_gpu_ram) + 1
if op == "gte":
return int(math.ceil(threshold / per_card_gpu_ram))
return None
def get_gpu_chunk_size(n):
"""Round n up to the next common GPU chunk size (1, 2, 4, 8). Returns n
unchanged for >8 since that's rare enough we'd rather be honest than guess.
"""
for c in (1, 2, 4, 8):
if n <= c:
return c
return n
def pick_num_gpus(gpu_name, extra_filters):
"""Smallest num_gpus where ``num_gpus * per_card_ram`` satisfies the
template's ``gpu_total_ram`` filter, rounded up to a chassis size.
Falls back to 1 when the template has no such filter or one card
already satisfies it.
"""
per_card = _per_card_vram_mb(gpu_name)
if per_card is None:
return 1
ops = (extra_filters or {}).get("gpu_total_ram") or {}
if not ops:
return 1
op, threshold = next(iter(ops.items()))
raw_min = min_gpus_for_ram(threshold, op, per_card)
if raw_min is None or raw_min <= 1:
return 1
return get_gpu_chunk_size(raw_min)
DEFAULT_CACHE_MAX_AGE_DAYS = 30
MAX_CACHE_ROWS = 100 # newest rows per template column; caps popular template+GPU combos
def lookup_cached_benchmark(vast, *, gpu_name, num_gpus, template_hash,
template_id, max_age_days):
"""Recent reported benchmarks for this exact spec (median + spread), or None.
Queried once per template column rather than filtered client-side: a
workergroup is created with either template_id or template_hash, so a row
can carry either one. Both columns are indexed with last_update on the
benchmarks table, so each query is index-backed instead of pulling every
perf row for the GPU and discarding the non-matching ones.
"""
base = {
"type": {"eq": "perf"},
"gpu_name": {"eq": gpu_name},
"num_gpus": {"eq": num_gpus},
"last_update": {"gte": time.time() - max_age_days * 86400},
}
rows = []
seen_ids = set()
for col, value in (("template_hash", template_hash), ("template_id", template_id)):
if not value:
continue
found = vast.search_benchmarks(
query=dict(base, **{col: {"eq": value}}),
order=[{"col": "last_update", "dir": "desc"}],
limit=MAX_CACHE_ROWS,
)
for r in found if isinstance(found, list) else []:
if not isinstance(r, dict):
continue
# A row carrying both columns comes back from both queries; dedupe on
# id, but keep rows without one rather than collapsing them together.
row_id = r.get("id")
if row_id is not None:
if row_id in seen_ids:
continue
seen_ids.add(row_id)
rows.append(r)
matched = [
r for r in rows
if isinstance(r.get("value"), (int, float)) and r["value"] > 0
]
if not matched:
return None
values = [r["value"] for r in matched]
newest = max((r.get("last_update") or 0) for r in matched)
# skip null rows (from before we tracked dph_base in benchmarks)
dphs = [
r["dph_base"]
for r in matched
if isinstance(r.get("dph_base"), (int, float)) and r["dph_base"] > 0
]
# use p5/p95 for the range
if len(values) >= 2:
qs = statistics.quantiles(values, n=20, method="inclusive")
low, high = qs[0], qs[-1]
else:
low = high = values[0]
return {
"median": statistics.median(values),
"low": low,
"high": high,
"n": len(matched),
"age_days": max(0.0, (time.time() - newest) / 86400),
"median_dph_base": statistics.median(dphs) if dphs else None,
}
def update_worker_states(worker_states, current_workers, gpu_name):
"""Update worker_states with current poll, and print only on worker
rotation (the live rich table already shows current status / elapsed).
Tracks state_started per worker for the terminal-state grace check.
Returns the count of workers abandoned during this poll.
"""
now = time.monotonic()
abandoned = 0
waiting_key = "__waiting__" # non-int sentinel: "no worker rented yet"
if not current_workers:
worker_states.setdefault(waiting_key, {"started": now})
return abandoned
worker_states.pop(waiting_key, None)
seen = set()
for w in current_workers:
wid = w.get("id")
if wid is None:
continue
seen.add(wid)
status = str(w.get("status", "")).lower()
prev = worker_states.get(wid)
if prev is None:
worker_states[wid] = {"status": status, "state_started": now}
continue
if prev["status"] != status:
prev["status"] = status
prev["state_started"] = now
for wid in list(worker_states):
if not isinstance(wid, int):
continue # skip the non-int waiting sentinel
if wid not in seen:
# Worker abandoned; live table only shows current worker, so rotation is otherwise invisible.
prev = worker_states[wid]
elapsed = format_time_elapsed(now - prev["state_started"])
print(f"[{gpu_name}] worker {wid} abandoned "
f"(last state={prev['status']} for {elapsed})",
file=sys.stderr)
del worker_states[wid]
abandoned += 1
return abandoned
def update_row(class_states, row_id, **fields):
"""Update one GPU's row in the live-table state dict. Sets
``run_started`` on first non-``queued`` status and ``run_ended`` on
terminal status (so elapsed freezes); each thread only writes its own
key so no lock needed.
"""
if class_states is None:
return
cur = class_states.setdefault(row_id, {})
new_status = fields.get("status")
now_ts = time.monotonic()
if new_status and new_status != cur.get("status"):
if "run_started" not in cur and new_status not in ("queued", "cached"):
cur["run_started"] = now_ts
if new_status in TERMINAL_STATUSES:
cur["run_ended"] = now_ts
else:
cur.pop("run_ended", None) # recovered from a transient terminal status; unfreeze elapsed
cur.update(fields)
def render_table(class_states):
"""Render the per-GPU progress as a rich.Table for the live display."""
table = Table(show_header=True, header_style="bold cyan")
table.add_column("GPU", no_wrap=True)
table.add_column("Status", no_wrap=True)
table.add_column("Endpoint", justify="right", no_wrap=True)
table.add_column("Worker", justify="right", no_wrap=True)
table.add_column("Elapsed", justify="right", no_wrap=True)
table.add_column("Perf", justify="right", no_wrap=True)
table.add_column("$/hr", justify="right", no_wrap=True)
table.add_column("Perf/$/hr", justify="right", no_wrap=True)
now = time.monotonic()
for gpu in sorted(class_states):
s = class_states[gpu]
status = s.get("status") or "-"
# Skipped rows are noise during the run; reason already printed at pre-flight, final summary still includes them.
if status == "skipped":
continue
run_started = s.get("run_started")
run_ended = s.get("run_ended")
# Blank for not-yet-running, frozen for terminal, live otherwise.
if run_started is None:
elapsed_str = "-"
elif run_ended is not None:
elapsed_str = format_time_elapsed(run_ended - run_started)
else:
elapsed_str = format_time_elapsed(now - run_started)
perf = s.get("perf")
dph = s.get("dph")
pps = (perf / dph) if (perf and dph) else None
table.add_row(
gpu,
status,
str(s.get("endpoint_id") or "-"),
str(s.get("worker_id") or "-"),
elapsed_str,
f"{perf:.1f}" if perf else "-",
f"${dph:.3f}" if dph else "-",
f"{pps:.1f}" if pps else "-",
)
return table
def benchmark_gpu(vast, *, gpu_name, num_gpus, timeout,
workergroups, endpoints,
template_hash=None, template_id=None,
auto_instance=None, autoscaler_url=None,
class_states=None, stop_event=None):
"""Rent one instance for ``gpu_name``, poll until idle + measured_perf,
tear down. Each call owns its own endpoint + workergroup so it can run
in parallel safely. Returns ``(gpu_name, num_gpus, status, perf, err, dph_base)``.
"""
endpoint_id = None
workergroup_id = None
# uuid suffix keeps the name unique when benchmarking the same GPU in parallel;
# no parens/shell chars, the backend rejects them in endpoint_name.
endpoint_name = f"benchmark {num_gpus}x {gpu_name} {uuid.uuid4().hex[:8]}"
row_id = f"{num_gpus}x {gpu_name}" # id for the live table so 1x/2x/4x of the same GPU gets its own row
start = time.monotonic()
update_row(class_states, row_id, status="provisioning")
search = f"gpu_name={gpu_name.replace(' ', '_')} num_gpus={num_gpus}"
if autoscaler_url:
status_vast = VastAI(api_key=vast.client.api_key, server_url=autoscaler_url, retry=vast.client.retry)
else:
status_vast = vast
try:
endpoint_kwargs = dict(ENDPOINT_CONFIG, endpoint_name=endpoint_name)
if auto_instance is not None:
endpoint_kwargs["auto_instance"] = auto_instance
ep_resp = vast.create_endpoint(**endpoint_kwargs)
# autoscaler returns the new id under either "result" or "id" depending on resource
endpoint_id = ep_resp.get("result", ep_resp.get("id")) if isinstance(ep_resp, dict) else None
if not isinstance(endpoint_id, int):
return (gpu_name, num_gpus, "error", None,
f"create_endpoint returned no id: {ep_resp!r}", None)
endpoints.add(endpoint_id)
update_row(class_states, row_id, endpoint_id=endpoint_id)
# cold_workers and min_load are set on ENDPOINT_CONFIG, not here; the autoscaler reads
# both from the endpoint group, not the workergroup (verified via autoscaler.cpp).
workergroup_kwargs = dict(
endpoint_id=endpoint_id,
endpoint_name=endpoint_name,
search_params=search,
)
if template_id is not None:
workergroup_kwargs["template_id"] = template_id
elif template_hash is not None:
workergroup_kwargs["template_hash"] = template_hash
if auto_instance is not None:
workergroup_kwargs["auto_instance"] = auto_instance
resp = vast.create_workergroup(**workergroup_kwargs)
workergroup_id = resp.get("result", resp.get("id")) if isinstance(resp, dict) else None
if not isinstance(workergroup_id, int):
return (gpu_name, num_gpus, "error", None,
f"create_workergroup returned no id: {resp!r}", None)
workergroups.add(workergroup_id)
update_row(class_states, row_id, status="waiting_for_worker")
worker_states = {}
dph_by_worker = {} # cache so we don't re-fetch dph on every poll
abandoned_count = 0 # bail when this hits MAX_ABANDONS; autoscaler is renting + giving up repeatedly, no progress
while time.monotonic() - start < timeout:
if stop_event is not None and stop_event.is_set():
update_row(class_states, row_id, status="aborted")
return (gpu_name, num_gpus, "aborted", None, "user aborted", None)
# calls get endpoint workers which includes measured perf(get workergroup workers doesnt)
resp = status_vast.get_endpoint_workers(endpoint_id)
if isinstance(resp, list):
workers = resp
elif isinstance(resp, dict) and isinstance(resp.get("workers"), list):
workers = resp["workers"]
else:
workers = []
abandoned_count += update_worker_states(worker_states, workers, gpu_name)
if abandoned_count >= MAX_ABANDONS:
update_row(class_states, row_id, status="failed")
return (gpu_name, num_gpus, "failed", None,
f"autoscaler abandoned {abandoned_count} workers without "
f"any reaching idle (likely template + GPU config not loadable "
f"on these hosts)", None)
if workers:
primary = workers[0]
primary_id = primary.get("id")
fields = {
"status": str(primary.get("status") or "?").lower(),
"worker_id": primary_id,
}
# Fetch dph once per new worker_id; surfaces $/hr in the live table as soon as the autoscaler rents (before model loads).
if primary_id and primary_id not in dph_by_worker:
try:
inst = vast.show_instance(id=primary_id)
dph_by_worker[primary_id] = inst.get("dph_base")
except Exception as e:
# Cache None so we don't re-attempt every poll, but log so a real bug isn't silently swallowed.
dph_by_worker[primary_id] = None
print(f"[warn] failed to fetch dph for worker {primary_id}: {e}",
file=sys.stderr)
if primary_id and dph_by_worker.get(primary_id) is not None:
fields["dph"] = dph_by_worker[primary_id]
update_row(class_states, row_id, **fields)
# measured_perf is only real once status==idle; before that it's a dlperf placeholder.
ready = [w for w in workers
if str(w.get("status", "")).lower() == "idle"
and (w.get("measured_perf") or 0) > 0]
if ready:
# dph was already fetched and cached when this worker_id first appeared in polling.
worker_id = ready[0].get("id")
dph = dph_by_worker.get(worker_id) if worker_id else None
update_row(class_states, row_id, status="done",
perf=ready[0]["measured_perf"], dph=dph)
return (gpu_name, num_gpus, "ok", ready[0]["measured_perf"], None, dph)
# Fail fast: every worker stuck terminal past the grace period, none reached idle.
now_ts = time.monotonic()
terminal_workers_long = workers and all(
str(w.get("status", "")).lower() in TERMINAL_STATES
and isinstance(w.get("id"), int)
and w["id"] in worker_states
and now_ts - worker_states[w["id"]]["state_started"]
> TERMINAL_GRACE_SECONDS
for w in workers
)
if terminal_workers_long:
states = sorted({str(w.get("status", "")).lower()
for w in workers})
update_row(class_states, row_id, status="failed")
return (gpu_name, num_gpus, "failed", None,
f"all workers terminal ({', '.join(states)}) for "
f">{TERMINAL_GRACE_SECONDS}s without reaching idle; "
f"autoscaler not rotating", None)
# Fail fast if no real worker (int id) has ever appeared.
if (not any(isinstance(k, int) for k in worker_states)
and time.monotonic() - start > RENTAL_TIMEOUT):
update_row(class_states, row_id, status="no_worker")
return (gpu_name, num_gpus, "no_worker", None,
f"autoscaler did not rent in {RENTAL_TIMEOUT}s "
f"(possible causes: insufficient credit, scoring issue, "
f"all candidates failed silently, or template+GPU mismatch "
f"missed by pre-flight)", None)
# Jitter so N parallel threads don't all hit the autoscaler at the same instant.
delay = WORKER_POLL_INTERVAL + random.uniform(0, 2.0)
if stop_event is not None:
if stop_event.wait(timeout=delay):
update_row(class_states, row_id, status="aborted")
return (gpu_name, num_gpus, "aborted", None, "user aborted", None)
else:
time.sleep(delay)
update_row(class_states, row_id, status="timeout")
return (gpu_name, num_gpus, "timeout", None,
f"no measured_perf in {timeout}s", None)
finally:
# Workergroup first (stops new workers), then endpoint. atexit sweeps anything we miss.
# 404s are silently absorbed since "already gone" is the desired end state.
if workergroup_id is not None:
try:
vast.delete_workergroup(id=workergroup_id)
workergroups.discard(workergroup_id)
except Exception as e:
if getattr(getattr(e, "response", None), "status_code", None) != 404:
print(f"[cleanup] failed to delete workergroup {workergroup_id}: {e}",
file=sys.stderr)
else:
workergroups.discard(workergroup_id)
if endpoint_id is not None:
try:
vast.delete_endpoint(id=endpoint_id)
endpoints.discard(endpoint_id)
except Exception as e:
if getattr(getattr(e, "response", None), "status_code", None) != 404:
print(f"[cleanup] failed to delete endpoint {endpoint_id}: {e}",
file=sys.stderr)
else:
endpoints.discard(endpoint_id)
@parser.command(
argument("--template_hash", type=str, default=None,
help="the template to benchmark; provide either --template_hash or --template_id"),
argument("--template_id", type=int, default=None,
help="the template to benchmark; provide either --template_hash or --template_id"),
argument("--gpus", type=str,
help="comma-separated GPU names to benchmark (e.g. RTX_4090,RTX_3090). Prefix a name with a count to set GPUs per instance, e.g. \"2x RTX_4090\" (overrides --num_gpus)"),
argument("--num_gpus", type=int, default=None,
help="GPUs per instance for tokens without an Nx prefix (default 1); overridden by inline Nx in --gpus"),
argument("--timeout", type=int, default=DEFAULT_BENCHMARK_TIMEOUT,
help=f"max seconds to wait for a benchmark before giving up (default {DEFAULT_BENCHMARK_TIMEOUT})"),
argument("--no-cache", dest="no_cache", action="store_true",
help="re-measure instead of reusing recent cached benchmark results"),
argument("-y", "--yes", action="store_true",
help="Skip confirmation prompt"),
argument("--auto_instance", type=str, default=None, help=argparse.SUPPRESS),
argument("--autoscaler_url", type=str, default=None, help=argparse.SUPPRESS),
usage="vastai run benchmarks [OPTIONS]",
help="Benchmark a template against one or more GPUs",
epilog=deindent("""
Rents one instance per GPU in parallel, measures perf, tears down.
Each rental runs for up to --timeout seconds and costs real money.
Specs benchmarked recently (same template, GPU, and count, by any
user) are served from the benchmarks table as an estimated perf, and
you're asked whether to reuse that or run a fresh benchmark. Pass
--no-cache to always re-measure, or -y to always reuse the cache.
Perf/$ uses the GPU $/hr, so cached and freshly measured rows are
comparable. Cached rows recorded before that price was tracked show
a perf with no price.
Examples:
# auto-sweep the default GPUs against TGI
vastai run benchmarks --template_hash 79ebdd2ebfb9d42cedf7a221c42d37a5
# specific GPUs against vLLM
vastai run benchmarks --template_hash 393fa8572e6c73c927c8275fe4dffd53 --gpus RTX_4090,RTX_3090
# multi-GPU configurations via inline Nx prefix against ComfyUI
vastai run benchmarks --template_hash 40ef49becc953aa910ee05bd4653b9b3 --gpus "2x RTX_4090, 2x RTX_3090"
# default count for tokens without an Nx prefix
vastai run benchmarks --template_hash 79ebdd2ebfb9d42cedf7a221c42d37a5 --gpus RTX_4090,RTX_3090 --num_gpus 2
# shorter timeout (10 min), skipping the cost prompt
vastai run benchmarks --template_hash 393fa8572e6c73c927c8275fe4dffd53 --timeout 600 -y
# ignore cached results and re-measure everything
vastai run benchmarks --template_hash 79ebdd2ebfb9d42cedf7a221c42d37a5 --no-cache
# raw JSON output for piping into another tool
vastai run benchmarks --template_hash 40ef49becc953aa910ee05bd4653b9b3 --raw
"""),
)
def run__benchmarks(args):
if args.template_id is None and args.template_hash is None:
print("error: one of --template_id or --template_hash is required "
"(run `vastai run benchmarks --help` for usage)",
file=sys.stderr)
return 1
if args.template_id is not None:
query = {"id": {"eq": args.template_id}}
ident = f"id={args.template_id}"
else:
query = {"hash_id": {"eq": args.template_hash}}
ident = f"hash={args.template_hash}"
# bump retry budget so parallel polls don't exhaust VastClient's tiny default (3 attempts, ~0.7s) under autoscaler rate limits.
retry = max(args.retry, 8)
vast = VastAI(api_key=args.api_key, server_url=args.url, retry=retry,
explain=getattr(args, 'explain', False),
curl=getattr(args, 'curl', False))
templates = vast.search_templates(query=query)
if not templates:
print(f"error: template not found ({ident})", file=sys.stderr)
return 1
template = templates[0]
try:
extra_filters = json.loads(template.get("extra_filters") or "{}")
except (ValueError, TypeError):
extra_filters = {}
template_name = template.get("name") or template.get("title") or "?"
filter_summary = format_filter_query(extra_filters) or "none"
print(f"\nTemplate id={template.get('id')} {template_name}, filters: {filter_summary}")
# if user provides specific GPU count like "2x RTX 4090", that takes precedence over --num_gpus
if args.gpus:
gpu_specs = [parse_gpu_spec(t, args.num_gpus or 1)
for t in args.gpus.split(",") if t.strip()]
elif args.num_gpus is not None:
gpu_specs = [(g, args.num_gpus) for g in DEFAULT_GPUS]
else:
gpu_specs = [
(g, pick_num_gpus(g, extra_filters))
for g in DEFAULT_GPUS
]
console = Console(stderr=True)
# deduplicates (gpu, num_gpus) so we dont benchmark the same GPU config twice
seen = set()
deduped = []
for g, n in gpu_specs:
if (g, n) in seen:
console.print(f"[yellow]skipping duplicate spec {n}x {g}[/yellow]",
highlight=False)
continue
seen.add((g, n))
deduped.append((g, n))
gpu_specs = deduped
# Pre-flight: for each spec, serve a recent cached benchmark if one exists
# (unless --no-cache), else skip specs with 0 matching offers, so the user
# sees cache hits and skip reasons before approving any rentals.
compatible_specs = []
skipped_results = []
cached_results = []
interactive = not (args.yes or args.raw)
for g, n in gpu_specs:
if not args.no_cache:
hit = lookup_cached_benchmark(
vast, gpu_name=g, num_gpus=n,
template_hash=template.get("hash_id"),
template_id=template.get("id"),
max_age_days=DEFAULT_CACHE_MAX_AGE_DAYS,
)
if hit:
rentable = has_matching_offer(
vast, gpu_name=g, num_gpus=n, extra_filters=extra_filters)
age = (f"{hit['age_days']:.0f}d" if hit["age_days"] >= 1
else "<1d")
note = None if rentable else "no offers available to rent right now"
median_dph_base = hit.get("median_dph_base")
pps = f", {hit['median'] / median_dph_base:.1f} perf/$" if median_dph_base else ""
msg = (f"[green][{n}x {g}] cached:[/green] median perf "
f"{hit['median']:.1f}{pps} "
f"(n={hit['n']}, range {hit['low']:.1f}-{hit['high']:.1f}, "
f"newest {age} ago)")
if note:
msg += f"; [yellow]{note}[/yellow]"
console.print(msg, highlight=False)
# Default is to reuse the cached result; only an interactive user
# who opts in pays to re-measure (and only when offers exist).
run_fresh = (
interactive and rentable
and input(
f" [{n}x {g}] run a fresh benchmark anyway? this rents a "
f"real GPU and charges your account [y/N] "
).strip().lower() in ("y", "yes")
)
if run_fresh:
compatible_specs.append((g, n)) # offers already confirmed
else:
cached_results.append((g, n, "cached", hit["median"], None, hit.get("median_dph_base")))
continue
if not has_matching_offer(
vast, gpu_name=g, num_gpus=n,
extra_filters=extra_filters,
):
msg = format_skip_message(
vast, gpu_name=g, num_gpus=n,
extra_filters=extra_filters,
)
console.print(f"[yellow][{n}x {g}] skipping:[/yellow] {msg}",
highlight=False)
skipped_results.append((g, n, "skipped", None, msg, None))
else:
compatible_specs.append((g, n))
# Live-table state. Pre-populate every GPU so the table is complete.
class_states = {}
for g, n in compatible_specs:
update_row(class_states, f"{n}x {g}", status="queued")
for sr in skipped_results:
update_row(class_states, f"{sr[1]}x {sr[0]}", status="skipped")
for cr in cached_results:
update_row(class_states, f"{cr[1]}x {cr[0]}", status="cached",
perf=cr[3], dph=cr[5])
timeout_minutes = args.timeout / 60.0
n = len(compatible_specs)
if n == 0:
if cached_results:
console.print()
console.print(render_table(class_states))
else:
console.print(
"\nNo compatible GPUs to benchmark for this template.",
style="bold red")
return print_results(args, skipped_results + cached_results)
spec_strs = [f"{cnt}x {gpu}" for gpu, cnt in compatible_specs]
spec_summary = ", ".join(spec_strs)
if n == 1:
print(f"\nWill rent 1 GPU configuration ({spec_summary}). "
f"Each runs ~2 to {timeout_minutes:.0f} min.")
else:
print(f"\nWill rent {n} GPU configurations in parallel "
f"({spec_summary}). Each runs ~2 to {timeout_minutes:.0f} min.")
print("This rents real GPUs and charges your account for usage.")
if not args.yes:
if input("Continue? [y/N] ").strip().lower() not in ("y", "yes"):
print("Aborted.", file=sys.stderr)
sys.exit(130)
print()
# Backing sets for the atexit sweep; per-GPU teardown discards on success.
workergroups = set()
endpoints = set()
stop_event = threading.Event()
def _cleanup():
# Mask SIGINT + catch BaseException so a second Ctrl+C can't leak records mid-iteration.
prev_handler = None
try:
prev_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
except (ValueError, OSError):
pass
try:
for workergroup_id in list(workergroups):
try:
vast.delete_workergroup(id=workergroup_id)
workergroups.discard(workergroup_id)
except BaseException:
pass
for ep_id in list(endpoints):
try:
vast.delete_endpoint(id=ep_id)
endpoints.discard(ep_id)
except BaseException:
pass
finally:
if prev_handler is not None:
try:
signal.signal(signal.SIGINT, prev_handler)
except (ValueError, OSError):
pass
atexit.register(_cleanup)
def _run_one_gpu(g, n):
try:
return benchmark_gpu(
vast,
gpu_name=g,
num_gpus=n,
timeout=args.timeout,
workergroups=workergroups,
endpoints=endpoints,
template_hash=args.template_hash,
template_id=args.template_id,
auto_instance=args.auto_instance,
autoscaler_url=args.autoscaler_url,
class_states=class_states,
stop_event=stop_event,
)
except Exception as e:
update_row(class_states, f"{n}x {g}", status="error")
# Surface response body for HTTPError so backend rate-limit / validation
# messages aren't swallowed (default __str__ on HTTPError just shows the URL).
body = ""
resp = getattr(e, "response", None)
if resp is not None:
try:
body = f" | body: {resp.text[:500]}"
except Exception:
pass
return (g, n, "error", None, f"{type(e).__name__}: {e}{body}", None)
run_results = []
executor = ThreadPoolExecutor(max_workers=len(compatible_specs))
try:
# Stagger submits to keep cold-start POSTs from tripping rate limits.
futures = []
for i, (g, n_for_class) in enumerate(compatible_specs):
futures.append(executor.submit(_run_one_gpu, g, n_for_class))
if i < len(compatible_specs) - 1:
time.sleep(PARALLEL_SUBMIT_DELAY)
try:
with Live(render_table(class_states), console=console,
refresh_per_second=2, transient=False) as live:
pending = set(futures)
while pending:
done, pending = wait(pending, timeout=0.5,
return_when=FIRST_COMPLETED)
for fut in done:
result = fut.result()
run_results.append(result)
# Surface failures live above the table.
g, n, status, _perf, err, _price = result
if status not in ("ok", "skipped") and err:
console.print(
f"[red][{n}x {g}] {status}:[/red] {err}",
highlight=False)
live.update(render_table(class_states))
except KeyboardInterrupt:
# Set stop_event so threads exit their poll/sleep immediately, then cleanup synchronously.
console.print(
"\n[yellow]Deleting endpoints, please wait... "
"If you exit now, you'll need to delete them manually with `vastai delete endpoint`.[/yellow]",
highlight=False)
stop_event.set()
executor.shutdown(wait=False, cancel_futures=True)
_cleanup()
sys.exit(130)
finally:
executor.shutdown(wait=True)
_cleanup()
results = skipped_results + cached_results + run_results
return print_results(args, results)
def print_results(args, results):
rows = []
for gpu, num_gpus, status, perf, err, price in results:
pps = (perf / price) if (perf and price) else None
rows.append({
"gpu_name": gpu,
"num_gpus": num_gpus,
"rental_dph": price,
"measured_perf": perf,
"status": status,
"perf_per_dollar": pps,
})
rows.sort(key=lambda r: (r["perf_per_dollar"] or -1), reverse=True)
if args.raw:
return rows
n_ok = sum(1 for r in rows if r["status"] == "ok")
n_cached = sum(1 for r in rows if r["status"] == "cached")
line = f"\nBenchmark complete: {n_ok + n_cached}/{len(rows)} GPUs measured"
print(line + (f" ({n_cached} from cache)." if n_cached else "."))
cli/commands/billing.py
"""CLI commands for billing, invoices, earnings, and user account management."""
import json
import sys
import time
from copy import deepcopy
from datetime import datetime, timezone
from io import StringIO
from vastai.cli.parser import argument
from vastai.cli.display import (
display_table, invoice_fields, user_fields, ipaddr_fields,
scheduled_jobs_fields, deindent, print_or_page,
)
from vastai.api import billing as billing_api
from vastai.api import auth as auth_api
from vastai.api import teams as teams_api
from vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401
parser = _get_parser()
# ---------------------------------------------------------------------------
# Helper utilities
# ---------------------------------------------------------------------------
def convert_timestamp_to_date(unix_timestamp):
utc_datetime = datetime.fromtimestamp(unix_timestamp, tz=timezone.utc)
return utc_datetime.strftime("%Y-%m-%d")
def to_timestamp_(val):
"""Convert a date string or int to a UNIX timestamp."""
if isinstance(val, int):
return val
if isinstance(val, str):
if val.isdigit():
return int(val)
return int(datetime.strptime(val + "+0000", '%Y-%m-%d%z').timestamp())
raise ValueError("Invalid date format")
def _sum(X, k):
y = 0
for x in X:
a = float(x.get(k, 0))
y += a
return y
def _select(X, k):
Y = set()
for x in X:
v = x.get(k, None)
if v is not None:
Y.add(v)
return Y
def filter_invoice_items(args, rows):
"""Filter invoice items by date range and charge/credit type."""
from datetime import date
try:
from dateutil import parser as dateutil_parser
except ImportError:
print("\nWARNING: Missing dateutil, can't parse time format")
end_timestamp = 9999999999
start_timestamp = 0
start_date_txt = ""
end_date_txt = ""
if args.end_date:
try:
end_date = dateutil_parser.parse(str(args.end_date))
end_date_txt = end_date.isoformat()
if end_date.tzinfo is None:
end_date = end_date.replace(tzinfo=timezone.utc)
end_timestamp = end_date.timestamp()
except ValueError:
print("Warning: Invalid end date format! Ignoring end date!")
if args.start_date:
try:
start_date = dateutil_parser.parse(str(args.start_date))
start_date_txt = start_date.isoformat()
if start_date.tzinfo is None:
start_date = start_date.replace(tzinfo=timezone.utc)
start_timestamp = start_date.timestamp()
except ValueError:
print("Warning: Invalid start date format! Ignoring start date!")
selector_flag = ""
if args.only_charges:
type_txt = "Only showing charges."
selector_flag = "only_charges"
def type_filter_fn(row):
return row["type"] == "charge"
elif args.only_credits:
type_txt = "Only showing credits."
selector_flag = "only_credits"
def type_filter_fn(row):
return row["type"] == "payment"
else:
type_txt = ""
def type_filter_fn(row):
return True
if args.end_date:
if args.start_date:
header_text = f'Invoice items after {start_date_txt} and before {end_date_txt}.'
else:
header_text = f'Invoice items before {end_date_txt}.'
elif args.start_date:
header_text = f'Invoice items after {start_date_txt}.'
else:
header_text = " "
header_text = header_text + " " + type_txt
rows = list(filter(
lambda row: end_timestamp >= (row["timestamp"] or 0.0) >= start_timestamp
and type_filter_fn(row) and float(row["amount"]) != 0,
rows
))
if start_date_txt:
start_date_txt = "S:" + start_date_txt
if end_date_txt:
end_date_txt = "E:" + end_date_txt
now = date.today()
invoice_number = now.year * 12 + now.month - 1
pdf_filename_fields = list(filter(
lambda fld: fld != "",
[str(invoice_number), start_date_txt, end_date_txt, selector_flag]
))
filename = "invoice_" + "-".join(pdf_filename_fields) + ".pdf"
return {"rows": rows, "header_text": header_text, "pdf_filename": filename}
# ---------------------------------------------------------------------------
# Invoices-v1 helpers (rich formatting)
# ---------------------------------------------------------------------------
charge_types = ['instance', 'volume', 'serverless', 'i', 'v', 's']
invoice_types = {
"transfers": "transfer",
"stripe": "stripe_payments",
"bitpay": "bitpay",
"coinbase": "coinbase",
"crypto.com": "crypto.com",
"reserved": "instance_prepay",
"payout_paypal": "paypal_manual",
"payout_wise": "wise_manual",
}
def format_invoices_charges_results(args, results):
indices_to_remove = []
for i, item in enumerate(results):
item['start'] = convert_timestamp_to_date(item['start']) if item['start'] else None
item['end'] = convert_timestamp_to_date(item['end']) if item['end'] else None
if item['amount'] == 0:
indices_to_remove.append(i)
elif args.invoices:
if item['type'] not in {'transfer', 'payout'}:
item['amount'] *= -1
item['amount_str'] = f"${item['amount']:.2f}" if item['amount'] > 0 else f"-${abs(item['amount']):.2f}"
else:
item['amount'] = f"${item['amount']:.3f}"
if args.charges:
if item['type'] in {'instance', 'volume'} and not args.verbose:
item['items'] = []
if item['source'] and '-' in item['source']:
item['type'], item['source'] = item['source'].capitalize().split('-')
item['items'] = format_invoices_charges_results(args, item['items'])
for i in reversed(indices_to_remove):
del results[i]
return results
def rich_object_to_string(rich_obj, no_color=True):
"""Render a Rich object (Table or Tree) to a string."""
from rich.console import Console
buffer = StringIO()
console = Console(record=True, file=buffer)
console.print(rich_obj)
return console.export_text(clear=True, styles=not no_color)
def create_charges_tree(results, parent=None, title="Charges Breakdown"):
"""Build and return a Rich Tree from nested charge results."""
from rich.text import Text
from rich.tree import Tree
from rich.panel import Panel
if parent is None:
root = Tree(Text(title, style="bold red"))
create_charges_tree(results, root)
return Panel(root, style="white on #000000", expand=False)
top_level = (parent.label.plain == title)
for item in results:
end_date = f" -> {item['end']}" if item['start'] != item['end'] else ""
label = Text.assemble(
(item["type"], "bold cyan"),
(f" {item['source']}" if item.get('source') else "", "gold1"), " -> ",
(f"{item['amount']}", 'bold green1' if top_level else 'green1'),
(f" -- {item['description']}", "bright_white" if top_level else "dim white"),
(f" ({item['start']}{end_date})", "bold bright_white" if top_level else "white")
)
node = parent.add(label, guide_style="blue3")
if item.get("items"):
create_charges_tree(item["items"], node)
return parent
def create_rich_table_for_charges(args, results):
"""Build and return a Rich Table from charge results."""
from rich.table import Table
from rich.text import Text
from rich import box
from rich.padding import Padding
table = Table(style="white", header_style="bold bright_yellow", box=box.DOUBLE_EDGE, row_styles=["on grey11", "none"])
table.add_column(Text("Type", justify="center"), style="bold steel_blue1", justify="center")
table.add_column(Text("ID", justify="center"), style="gold1", justify="center")
table.add_column(Text("Amount", justify="center"), style="sea_green2", justify="right")
table.add_column(Text("Start", justify="center"), style="bright_white", justify="center")
table.add_column(Text("End", justify="center"), style="bright_white", justify="center")
if not args.charge_type or 'serverless' in args.charge_type:
table.add_column(Text("Endpoint", justify="center"), style="bright_red", justify="center")
table.add_column(Text("Workergroup", justify="center"), style="orchid", justify="center")
for item in results:
row = [item['type'].capitalize(), item['source'], item['amount'], item['start'], item['end']]
if not args.charge_type or 'serverless' in args.charge_type:
row.append(str(item['metadata'].get('endpoint_id', '')))
row.append(str(item['metadata'].get('workergroup_id', '')))
table.add_row(*row)
return Padding(table, (1, 2), style="on #000000", expand=False)
def create_rich_table_for_invoices(results):
"""Build and return a Rich Table from invoice results."""
from rich.table import Table
from rich.text import Text
from rich import box
from rich.padding import Padding
invoice_type_to_color = {
"credit": "green1",
"transfer": "gold1",
"payout": "orchid",
"reserved": "sky_blue1",
"refund": "bright_red",
}
table = Table(style="white", header_style="bold bright_yellow", box=box.DOUBLE_EDGE, row_styles=["on grey11", "none"])
table.add_column(Text("ID", justify="center"), style="bright_white", justify="center")
table.add_column(Text("Created", justify="center"), style="yellow3", justify="center")
table.add_column(Text("Paid", justify="center"), style="yellow3", justify="center")
table.add_column(Text("Type", justify="center"), justify="center")
table.add_column(Text("Result", justify="center"), justify="right")
table.add_column(Text("Source", justify="center"), style="bright_cyan", justify="center")
table.add_column(Text("Description", justify="center"), style="bright_white", justify="left")
for item in results:
table.add_row(
str(item['metadata']['invoice_id']),
item['start'],
item['end'] if item['end'] else 'N/A',
Text(item['type'].capitalize(), style=invoice_type_to_color.get(item['type'], "white")),
Text(item['amount_str'], style="sea_green2" if item['amount'] > 0 else "bright_red"),
item['source'].capitalize() if item['type'] != 'transfer' else item['source'],
item['description'],
)
return Padding(table, (1, 2), style="on #000000", expand=False)
# ---------------------------------------------------------------------------
# show invoices (deprecated)
# ---------------------------------------------------------------------------
@parser.command(
argument("-q", "--quiet", action="store_true", help="only display numeric ids"),
argument("-s", "--start_date", help="start date and time for report. Many formats accepted (optional)", type=str),
argument("-e", "--end_date", help="end date and time for report. Many formats accepted (optional)", type=str),
argument("-c", "--only_charges", action="store_true", help="Show only charge items"),
argument("-p", "--only_credits", action="store_true", help="Show only credit items"),
argument("--instance_label", help="Filter charges on a particular instance label (useful for autoscaler groups)"),
usage="(DEPRECATED) vastai show invoices [OPTIONS]",
help="(DEPRECATED) Get billing history reports. Use `vastai show invoices-v1` instead.",
)
def show__invoices(args):
"""Show current payments and charges (deprecated)."""
if not args.quiet:
print("DEPRECATED: `vastai show invoices` will be removed in a future release. "
"Use `vastai show invoices-v1` for the new paginated command.", file=sys.stderr)
client = get_client(args)
result = billing_api.show_invoices(client, start_date=args.start_date, end_date=args.end_date,
only_charges=args.only_charges, only_credits=args.only_credits)
rows = result["invoices"]
current_charges = result["current"]
invoice_filter_data = filter_invoice_items(args, rows)
rows = invoice_filter_data["rows"]
filter_header = invoice_filter_data["header_text"]
if args.instance_label:
contract_ids = _select(rows, 'instance_id')
req_json = {
"label": args.instance_label,
"contract_ids": list(contract_ids)
}
if args.explain:
print("request json: ")
print(req_json)
result2 = client.post("/contracts/fetch/", json_data=req_json)
result2.raise_for_status()
filtered_rows = result2.json()["contracts"]
contract_ids = _select(filtered_rows, 'id')
rows = [row for row in rows if row.get("instance_id") in contract_ids]
if args.quiet:
for row in rows:
id = row.get("id", None)
if id is not None:
print(id)
elif args.raw:
return rows
else:
print(filter_header)
display_table(rows, invoice_fields)
print(f"Total: ${_sum(rows, 'amount')}")
print("Current: ", current_charges)
# ---------------------------------------------------------------------------
# show invoices-v1 (advanced with pagination and rich formatting)
# ---------------------------------------------------------------------------
@parser.command(
argument('-i', '--invoices', mutex_group='grp', action='store_true', required=True, help='Show invoices instead of charges'),
argument('-it', '--invoice-type', choices=invoice_types.keys(), nargs='+', metavar='type',
help=f'Filter which types of invoices to show: {{{", ".join(invoice_types.keys())}}}'),
argument('-c', '--charges', mutex_group='grp', action='store_true', required=True, help='Show charges instead of invoices'),
argument('-ct', '--charge-type', choices=charge_types, nargs='+', metavar='type',
help='Filter which types of charges to show: {i|instance, v|volume, s|serverless}'),
argument('-s', '--start-date', help='Start date (YYYY-MM-DD or timestamp)'),
argument('-e', '--end-date', help='End date (YYYY-MM-DD or timestamp)'),
argument('-l', '--limit', type=int, default=20, help='Number of results per page (default: 20, max: 100)'),
argument('-t', '--next-token', help='Pagination token for next page'),
argument('-f', '--format', choices=['table', 'tree'], default='table', help='Output format for charges (default: table)'),
argument('-v', '--verbose', action='store_true', help='Include full Instance Charge details and Invoice Metadata (tree view only)'),
argument('--latest-first', action='store_true', help='Sort by latest first'),
usage="vastai show invoices-v1 [OPTIONS]",
help="Get billing (invoices/charges) history reports with advanced filtering and pagination",
epilog=deindent("""
This command supports colored output and rich formatting if the 'rich' python module is installed!
Examples:
# Show the first 20 invoices in the last week (note: default window is a 7 day period ending today)
vastai show invoices-v1 --invoices
# Show the first 50 charges over a 7 day period starting from 2025-11-30 in tree format
vastai show invoices-v1 --charges -s 2025-11-30 -f tree -l 50
# Show the first 20 invoices of specific types for the month of November 2025
vastai show invoices-v1 -i -it stripe bitpay transfers --start-date 2025-11-01 --end-date 2025-11-30
# Show the first 20 charges for only volumes and serverless instances between two dates, including all details and metadata
vastai show invoices-v1 -c --charge-type v s -s 2025-11-01 -e 2025-11-05 --format tree --verbose
# Get the next page of paginated invoices, limit to 50 per page (note: type/date filters MUST match previous request for pagination to work)
vastai show invoices-v1 --invoices --limit 50 --next-token eyJ2YWx1ZXMiOiB7ImlkIjogMjUwNzgyMzR9LCAib3NfcGFnZSI6IDB9
# Show the last 10 instance (only) charges over a 7 day period ending in 2025-12-25, sorted by latest charges first
vastai show invoices-v1 --charges -ct instance --end-date 2025-12-25 -l 10 --latest-first
"""),
)
def show__invoices_v1(args):
"""Get billing (invoices/charges) history with advanced filtering and pagination."""
output_lines = []
try:
from rich.prompt import Confirm
has_rich = True
except ImportError:
output_lines.append("NOTE: To view results in color and table/tree format please install the 'rich' python module with 'pip install rich'\n")
has_rich = False
# Handle default start and end date values
if not args.start_date and not args.end_date:
args.end_date = int(time.time())
if not args.start_date:
args.start_date = args.end_date - 7 * 24 * 60 * 60
elif not args.end_date:
args.end_date = args.start_date + 7 * 24 * 60 * 60
try:
start_timestamp = to_timestamp_(args.start_date)
end_timestamp = to_timestamp_(args.end_date)
except Exception as e:
print(f"Error parsing dates: {e}")
print("Use format YYYY-MM-DD or UNIX timestamp")
return
if has_rich and not args.no_color and not args.raw:
print("(use --no-color to disable colored output)\n")
start_date = convert_timestamp_to_date(start_timestamp)
end_date = convert_timestamp_to_date(end_timestamp)
data_type = "Instance Charges" if args.charges else "Invoices"
output_lines.append(f"Fetching {data_type} from {start_date} to {end_date}...")
# Build SDK params
params = {
"charges": args.charges,
"start_date": args.start_date,
"end_date": args.end_date,
"limit": min(args.limit, 100) if args.limit > 0 else 20,
"latest_first": args.latest_first,
"format": args.format if args.charges else "table",
"charge_type": args.charge_type or [],
"invoice_type": args.invoice_type or [],
"next_token": args.next_token,
}
client = get_client(args)
if args.raw:
# Single-page raw response; dispatcher JSON-prints. Callers paginate
# by re-running with --next-token.
return billing_api.show_invoices_v1(client, params)
found_results, found_count = [], 0
looping = True
while looping:
response = billing_api.show_invoices_v1(client, params)
found_results += response.get('results', [])
found_count += response.get('count', 0)
total = response.get('total', 0)
next_token = response.get('next_token')
if has_rich is False:
output_lines.append("Raw response:\n" + json.dumps(response, indent=2))
if next_token:
print(f"Next page token: {next_token}\n")
elif not found_results:
output_lines.append("No results found")
else:
formatted_results = format_invoices_charges_results(args, deepcopy(found_results))
if args.invoices:
rich_obj = create_rich_table_for_invoices(formatted_results)
elif args.format == 'tree':
rich_obj = create_charges_tree(formatted_results)
else:
rich_obj = create_rich_table_for_charges(args, formatted_results)
output_lines.append(rich_object_to_string(rich_obj, no_color=args.no_color))
output_lines.append(f"Showing {found_count} of {total} results")
if next_token:
output_lines.append(f"Next page token: {next_token}\n")
paging = print_or_page(args, '\n'.join(output_lines))
if next_token and not paging:
if has_rich:
ans = Confirm.ask("Fetch next page?", show_default=False, default=False)
else:
ans = input("Fetch next page? (y/N): ").strip().lower() == 'y'
if ans:
params['next_token'] = next_token
output_lines.clear()
args.full = True
else:
looping = False
else:
looping = False
# ---------------------------------------------------------------------------
# show earnings
# ---------------------------------------------------------------------------
@parser.command(
argument("-q", "--quiet", action="store_true", help="only display numeric ids"),
argument("-s", "--start_date", help="start date and time for report. Many formats accepted", type=str),
argument("-e", "--end_date", help="end date and time for report. Many formats accepted ", type=str),
argument("-m", "--machine_id", help="Machine id (optional)", type=int),
usage="vastai show earnings [OPTIONS]",
help="Get machine earning history reports",
)
def show__earnings(args):
"""Show earnings history for a time range, optionally per machine."""
client = get_client(args)
rows = billing_api.show_earnings(client, start_date=args.start_date, end_date=args.end_date,
machine_id=args.machine_id)
if args.raw:
return rows
print(json.dumps(rows, indent=1, sort_keys=True))
# ---------------------------------------------------------------------------
# show deposit
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of instance to get info for", type=int),
usage="vastai show deposit ID [options]",
help="Display reserve deposit info for an instance",
)
def show__deposit(args):
"""Show reserve deposit info for an instance."""
client = get_client(args)
result = billing_api.show_deposit(client, id=args.id)
print(json.dumps(result, indent=1, sort_keys=True))
# ---------------------------------------------------------------------------
# show / set user
# ---------------------------------------------------------------------------
@parser.command(
argument("-q", "--quiet", action="store_true", help="display information about user"),
usage="vastai show user [OPTIONS]",
help="Get current user data",
epilog=deindent("""
Shows stats for logged-in user. These include user balance, email, and ssh key. Does not show API key.
"""),
)
def show__user(args):
"""Show stats for logged-in user."""
client = get_client(args)
user_blob = billing_api.show_user(client)
if args.raw:
return user_blob
else:
display_table([user_blob], user_fields)
@parser.command(
argument("--file", help="file path for params in json format", type=str),
usage="vastai set user --file FILE",
help="Update user data from json file",
epilog=deindent("""
Available fields:
Name Type Description
ssh_key string
paypal_email string
wise_email string
email string
normalized_email string
username string
fullname string
billaddress_line1 string
billaddress_line2 string
billaddress_city string
billaddress_zip string
billaddress_country string
billaddress_taxinfo string
balance_threshold_enabled string
balance_threshold string
autobill_threshold string
phone_number string
"""),
)
def set__user(args):
"""Update user data from a json file."""
params = None
with open(args.file, 'r') as file:
params = json.load(file)
client = get_client(args)
result = billing_api.set_user(client, params=params)
print(f"{result}")
# ---------------------------------------------------------------------------
# show subaccounts / create subaccount
# ---------------------------------------------------------------------------
@parser.command(
argument("-q", "--quiet", action="store_true", help="display subaccounts from current user"),
usage="vastai show subaccounts [OPTIONS]",
help="Get current subaccounts",
)
def show__subaccounts(args):
"""Show subaccounts."""
client = get_client(args)
rows = billing_api.show_subaccounts(client)
if args.raw:
return rows
else:
display_table(rows, user_fields)
@parser.command(
argument("--email", help="email address to use for login", type=str),
argument("--username", help="username to use for login", type=str),
argument("--password", help="password to use for login", type=str),
argument("--type", help="host/client", type=str),
usage="vastai create subaccount --email EMAIL --username USERNAME --password PASSWORD --type TYPE",
help="Create a subaccount",
epilog=deindent("""
Creates a new account that is considered a child of your current account as defined via the API key.
vastai create subaccount --email bob@gmail.com --username bob --password password --type host
vastai create subaccount --email vast@gmail.com --username vast --password password --type host
"""),
)
def create__subaccount(args):
"""Create a subaccount under your current account."""
host_only = False
if args.type:
host_only = args.type.lower() == "host"
if getattr(args, 'explain', False):
print("Request JSON would be: ")
print({
"email": args.email, "username": args.username,
"password": args.password, "host_only": host_only, "parent_id": "me",
})
return
client = get_client(args)
try:
rj = billing_api.create_subaccount(client, email=args.email, username=args.username,
password=args.password, host_only=host_only)
print(rj)
except Exception as e:
print(f"Failed with error: {e}")
# ---------------------------------------------------------------------------
# show ipaddrs
# ---------------------------------------------------------------------------
@parser.command(
usage="vastai show ipaddrs [--api-key API_KEY] [--raw]",
help="Display user's history of ip addresses",
)
def show__ipaddrs(args):
"""Show the history of ip address accesses."""
client = get_client(args)
rows = billing_api.show_ipaddrs(client)
if args.raw:
return rows
else:
display_table(rows, ipaddr_fields)
# ---------------------------------------------------------------------------
# transfer credit
# ---------------------------------------------------------------------------
@parser.command(
argument("recipient_pos", nargs="?", default=None, metavar="RECIPIENT",
help="email (or id) of recipient account (positional; use --recipient to name)"),
argument("amount_pos", nargs="?", default=None, type=float, metavar="AMOUNT",
help="dollars of credit to transfer (positional; use --amount to name)"),
argument("-r", "--recipient", help="email (or id) of recipient account", type=str, default=None),
argument("-a", "--amount", help="dollars of credit to transfer", type=float, default=None),
argument("--skip", help="skip confirmation", action="store_true", default=False),
usage="vastai transfer credit [--recipient EMAIL] [--amount DOLLARS] [RECIPIENT AMOUNT]",
help="Transfer credits to another account",
epilog=deindent("""
Transfer credits to another account. This action is irreversible.
Supports both named flags (recommended) and legacy positional arguments:
vastai transfer credit --recipient user@example.com --amount 25.50
vastai transfer credit -r user@example.com -a 25.50
vastai transfer credit user@example.com 25.50
"""),
)
def transfer__credit(args):
"""Transfer credits to another account."""
recipient = args.recipient if args.recipient is not None else args.recipient_pos
amount = args.amount if args.amount is not None else args.amount_pos
if recipient is None or amount is None:
print("Error: both recipient and amount are required (as flags or positional args).")
return
if not args.skip:
print(f"Transfer ${amount} credit to account {recipient}? This is irreversible.")
ok = input("Continue? [y/n] ")
if ok.strip().lower() != "y":
return
if args.explain:
print("request json: ")
print({"sender": "me", "recipient": recipient, "amount": amount})
client = get_client(args)
try:
rj = teams_api.transfer_credit(client, recipient=recipient, amount=amount)
if rj["success"]:
print(f"Sent {amount} to {recipient}")
else:
print(rj["msg"])
except Exception as e:
print(f"Failed with error: {e}")
# ---------------------------------------------------------------------------
# scheduled jobs
# ---------------------------------------------------------------------------
def normalize_schedule_fields(job):
"""Mutates the job dict to replace None values with readable scheduling labels."""
if job.get("day_of_the_week") is None:
job["day_of_the_week"] = "Everyday"
else:
days = ["Sunday", "Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday"]
job["day_of_the_week"] = days[int(job["day_of_the_week"])]
if job.get("hour_of_the_day") is None:
job["hour_of_the_day"] = "Every hour"
else:
hour = int(job["hour_of_the_day"])
suffix = "AM" if hour < 12 else "PM"
hour_12 = hour % 12
hour_12 = 12 if hour_12 == 0 else hour_12
job["hour_of_the_day"] = f"{hour_12}_{suffix}"
if job.get("min_of_the_hour") is None:
job["min_of_the_hour"] = "Every minute"
else:
job["min_of_the_hour"] = f"{int(job['min_of_the_hour']):02d}"
return job
@parser.command(
usage="vastai show scheduled-jobs [--api-key API_KEY] [--raw]",
help="Display the list of scheduled jobs",
)
def show__scheduled_jobs(args):
"""Show the list of scheduled jobs for the account."""
client = get_client(args)
rows = auth_api.show_scheduled_jobs(client)
if args.raw:
return rows
else:
rows = [normalize_schedule_fields(job) for job in rows]
display_table(rows, scheduled_jobs_fields)
@parser.command(
argument("id", help="id of scheduled job to remove", type=int),
usage="vastai delete scheduled-job ID",
help="Delete a scheduled job",
)
def delete__scheduled_job(args):
"""Delete a scheduled job."""
client = get_client(args)
result = auth_api.delete_scheduled_job(client, id=args.id)
print(result)
cli/commands/clusters.py
"""CLI commands for managing clusters and overlays."""
import json
from vastai.cli.parser import argument
from vastai.cli.display import (
display_table, cluster_fields, overlay_fields, deindent,
)
from vastai.api import clusters as clusters_api
from vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401
parser = _get_parser()
# ---------------------------------------------------------------------------
# clusters
# ---------------------------------------------------------------------------
@parser.command(
argument("subnet", help="local subnet for cluster, ex: '0.0.0.0/24'", type=str),
argument("manager_id", help="Machine ID of manager node in cluster. Must exist already.", type=int),
usage="vastai create cluster SUBNET MANAGER_ID",
help="Create Vast cluster",
epilog=deindent("""
Create Vast Cluster by defining a local subnet and manager id.""")
)
def create__cluster(args):
"""Create a Vast Cluster."""
if args.explain:
print("request json:", {"subnet": args.subnet, "manager_id": args.manager_id})
client = get_client(args)
result = clusters_api.create_cluster(client, subnet=args.subnet, manager_id=args.manager_id)
if args.raw:
return result
print(result.get("msg", result))
@parser.command(
usage="vastai show clusters",
help="Show clusters associated with your account.",
epilog=deindent("""
Show clusters associated with your account.
""")
)
def show__clusters(args):
"""Show clusters associated with your account."""
client = get_client(args)
response_data = clusters_api.show_clusters(client)
if args.raw:
return response_data
rows = []
for cluster_id, cluster_data in response_data['clusters'].items():
machine_ids = [node["machine_id"] for node in cluster_data["nodes"]]
manager_node = next(node for node in cluster_data['nodes'] if node['is_cluster_manager'])
row_data = {
'id': cluster_id,
'subnet': cluster_data['subnet'],
'node_count': len(cluster_data['nodes']),
'machine_ids': str(machine_ids),
'manager_id': str(manager_node['machine_id']),
'manager_ip': manager_node['local_ip'],
}
rows.append(row_data)
display_table(rows, cluster_fields, replace_spaces=False)
@parser.command(
argument("cluster_id", help="ID of cluster to delete", type=int),
usage="vastai delete cluster CLUSTER_ID",
help="Delete Cluster",
epilog=deindent("""
Delete Vast Cluster""")
)
def delete__cluster(args):
"""Delete a Vast Cluster."""
if args.explain:
print("request json:", {"cluster_id": args.cluster_id})
client = get_client(args)
result = clusters_api.delete_cluster(client, cluster_id=args.cluster_id)
if args.raw:
return result
print(result["msg"])
@parser.command(
argument("cluster_id", help="ID of cluster to add machine to", type=int),
argument("machine_ids", help="machine id(s) to join cluster", type=int, nargs="+"),
usage="vastai join cluster CLUSTER_ID MACHINE_IDS",
help="Join Machine to Cluster",
epilog=deindent("""
Join's Machine to Vast Cluster
""")
)
def join__cluster(args):
"""Join machine(s) to a Vast Cluster."""
if args.explain:
print("request json:", {"cluster_id": args.cluster_id, "machine_ids": args.machine_ids})
client = get_client(args)
result = clusters_api.join_cluster(client, cluster_id=args.cluster_id, machine_ids=args.machine_ids)
if args.raw:
return result
print(result["msg"])
@parser.command(
argument("cluster_id", help="ID of cluster you want to remove machine from.", type=int),
argument("machine_id", help="ID of machine to remove from cluster.", type=int),
argument("new_manager_id", help="ID of machine to promote to manager. Must already be in cluster", type=int, nargs="?"),
usage="vastai remove-machine-from-cluster CLUSTER_ID MACHINE_ID NEW_MANAGER_ID",
help="Removes machine from cluster",
epilog=deindent("""Removes machine from cluster and also reassigns manager ID,
if we're removing the manager node""")
)
def remove_machine_from_cluster(args):
"""Remove a machine from a cluster."""
if args.explain:
print("request json:", {"cluster_id": args.cluster_id, "machine_id": args.machine_id,
"new_manager_id": args.new_manager_id})
client = get_client(args)
result = clusters_api.remove_machine_from_cluster(
client, cluster_id=args.cluster_id, machine_id=args.machine_id,
new_manager_id=args.new_manager_id,
)
if args.raw:
return result
print(result.get("msg", result))
# ---------------------------------------------------------------------------
# overlays
# ---------------------------------------------------------------------------
@parser.command(
argument("cluster_id", help="ID of cluster to create overlay on top of", type=int),
argument("name", help="overlay network name"),
usage="vastai create overlay CLUSTER_ID OVERLAY_NAME",
help="Creates overlay network on top of a physical cluster",
epilog=deindent("""
Creates an overlay network to allow local networking between instances on a physical cluster"""),
)
def create__overlay(args):
"""Create an overlay network on a physical cluster."""
if args.explain:
print("request json:", {"cluster_id": args.cluster_id, "name": args.name})
client = get_client(args)
result = clusters_api.create_overlay(client, cluster_id=args.cluster_id, name=args.name)
if args.raw:
return result
print(result["msg"])
@parser.command(
usage="vastai show overlays",
help="Show overlays associated with your account.",
epilog=deindent("""
Show overlays associated with your account.
"""),
)
def show__overlays(args):
"""Show overlays associated with your account."""
client = get_client(args)
response_data = clusters_api.show_overlays(client)
if args.raw:
return response_data
rows = []
for overlay in response_data:
row_data = {
'overlay_id': overlay['overlay_id'],
'name': overlay['name'],
'subnet': overlay['internal_subnet'] if overlay['internal_subnet'] else 'N/A',
'cluster_id': overlay['cluster_id'],
'instance_count': len(overlay['instances']),
'instances': str(overlay['instances']),
}
rows.append(row_data)
display_table(rows, overlay_fields, replace_spaces=False)
@parser.command(
argument("overlay_identifier", help="ID (int) or name (str) of overlay to delete", nargs="?"),
usage="vastai delete overlay OVERLAY_IDENTIFIER",
help="Deletes overlay and removes all of its associated instances",
)
def delete__overlay(args):
"""Delete an overlay network."""
if args.explain:
print("request json:", {"overlay_identifier": args.overlay_identifier})
client = get_client(args)
result = clusters_api.delete_overlay(client, overlay_identifier=args.overlay_identifier)
if args.raw:
return result
print(result["msg"])
@parser.command(
argument("name", help="Overlay network name to join instance to.", type=str),
argument("instance_id", help="Instance ID to add to overlay.", type=int),
usage="vastai join overlay OVERLAY_NAME INSTANCE_ID",
help="Adds instance to an overlay network",
epilog=deindent("""
Adds an instance to a compatible overlay network."""),
)
def join__overlay(args):
"""Add an instance to an overlay network."""
if args.explain:
print("request json:", {"name": args.name, "instance_id": args.instance_id})
client = get_client(args)
result = clusters_api.join_overlay(client, name=args.name, instance_id=args.instance_id)
if args.raw:
return result
print(result["msg"])
cli/commands/deployments.py
"""CLI commands for managing deployments."""
from vastai.cli.parser import argument
from vastai.cli.display import display_table, deindent
from vastai.api import deployments as deployments_api
from vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401
parser = _get_parser()
deployment_fields = (
("id", "ID", "{}", None, True),
("name", "Name", "{}", None, True),
("tag", "Tag", "{}", None, True),
("image", "Image", "{}", None, True),
("endpoint_id", "Endpoint", "{}", None, True),
("file_hash", "Hash", "{}", None, True),
("current_version_id", "Version", "{}", None, True),
("ttl", "TTL", "{}", None, True),
("created_at", "Created", "{}", None, True),
("updated_at", "Updated", "{}", None, True),
)
deployment_detail_fields = (
("id", "ID", "{}", None, True),
("name", "Name", "{}", None, True),
("tag", "Tag", "{}", None, True),
("image", "Image", "{}", None, True),
("endpoint_id", "Endpoint", "{}", None, True),
("endpoint_state", "State", "{}", None, True),
("worker_count", "Workers", "{}", None, True),
("file_hash", "Hash", "{}", None, True),
("s3_key", "S3 Key", "{}", None, True),
("current_version_id", "Version", "{}", None, True),
("last_healthy_version_id", "Last Healthy", "{}", None, True),
("storage", "Storage", "{}", None, True),
("ttl", "TTL", "{}", None, True),
("created_at", "Created", "{}", None, True),
("updated_at", "Updated", "{}", None, True),
)
# ---------------------------------------------------------------------------
# show deployments
# ---------------------------------------------------------------------------
@parser.command(
argument("-q", "--quiet", action="store_true", help="only display numeric ids"),
usage="vastai show deployments [OPTIONS]",
help="Display user's current deployments",
epilog=deindent("""
Shows the user's current deployments.
Examples:
vastai show deployments
vastai show deployments --raw
vastai show deployments -q
"""),
)
def show__deployments(args):
"""Show the user's current deployments."""
client = get_client(args)
rows = deployments_api.show_deployments(client)
if args.quiet:
for row in rows:
id = row.get("id", None)
if id is not None:
print(id)
elif args.raw:
return rows
else:
display_table(rows, deployment_fields)
@parser.command(
argument("id", help="id of deployment to show info for", type=int),
argument("-q", "--quiet", action="store_true", help="only display numeric id"),
usage="vastai show deployment ID [OPTIONS]",
help="Display a single deployment",
epilog=deindent("""
Shows detailed info for a single deployment, including endpoint state and worker count.
Examples:
vastai show deployment 123
vastai show deployment 123 --raw
"""),
)
def show__deployment(args):
"""Show details of a single deployment."""
client = get_client(args)
result = deployments_api.show_deployment(client, id=args.id)
if args.raw:
return result
if args.quiet:
print(result.get("id", ""))
else:
display_table([result], deployment_detail_fields)
# ---------------------------------------------------------------------------
# show deployment-versions
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of deployment to show versions for", type=int),
usage="vastai show deployment-versions ID",
help="Display versions for a deployment",
epilog=deindent("""
Example: vastai show deployment-versions 1234
"""),
)
def show__deployment_versions(args):
"""Show version history for a deployment."""
client = get_client(args)
rows = deployments_api.show_deployment_versions(client, id=args.id)
if args.raw:
return rows
else:
import json
print(json.dumps(rows, indent=1, sort_keys=True))
# ---------------------------------------------------------------------------
# stop deployment
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of deployment to stop", type=int),
usage="vastai stop deployment ID [OPTIONS]",
help="Stop a running deployment",
epilog=deindent("""
Stops the endpoint associated with a deployment. The deployment is not
deleted and can be restarted with 'vastai start deployment'.
Examples:
vastai stop deployment 1234
vastai stop deployment 1234 --raw
"""),
)
def stop__deployment(args):
"""Stop a running deployment."""
client = get_client(args)
rj = deployments_api.stop_deployment(client, id=args.id)
if args.raw:
return rj
elif rj.get("success"):
print(f"Deployment {args.id} stopped.")
else:
print(rj.get("msg", rj))
# ---------------------------------------------------------------------------
# start deployment
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of deployment to start", type=int),
usage="vastai start deployment ID [OPTIONS]",
help="Start a stopped deployment",
epilog=deindent("""
Starts (or restarts) the endpoint associated with a deployment.
Examples:
vastai start deployment 1234
vastai start deployment 1234 --raw
"""),
)
def start__deployment(args):
"""Start a stopped deployment."""
client = get_client(args)
rj = deployments_api.start_deployment(client, id=args.id)
if args.raw:
return rj
elif rj.get("success"):
print(f"Deployment {args.id} started.")
else:
print(rj.get("msg", rj))
# ---------------------------------------------------------------------------
# delete deployment
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of deployment to delete", type=int, nargs="?", default=None),
argument("--name", help="name of deployment to delete (deletes all tags unless --tag is specified)", type=str, default=None),
argument("--tag", help="tag to filter by when deleting by name", type=str, default=None),
usage="vastai delete deployment [ID | --name NAME [--tag TAG]]",
help="Delete a deployment by id, or by name and optional tag",
epilog=deindent("""
Examples:
vastai delete deployment 1234
vastai delete deployment --name my-deployment
vastai delete deployment --name my-deployment --tag prod
"""),
)
def delete__deployment(args):
"""Delete a deployment."""
if args.id is not None and args.name is not None:
print("Error: specify either an id or --name, not both")
return
if args.tag is not None and args.name is None:
print("Error: --tag can only be used with --name")
return
client = get_client(args)
if args.id is not None:
rj = deployments_api.delete_deployment(client, id=args.id)
elif args.name is not None:
if args.explain:
json_blob = {"name": args.name}
if args.tag is not None:
json_blob["tag"] = args.tag
print("request json: ")
print(json_blob)
rj = deployments_api.delete_deployment_by_name(client, name=args.name, tag=args.tag)
else:
print("Error: must specify either an id or --name")
return
if args.raw:
return rj
elif rj.get("success"):
if "count" in rj:
print(f"Deleted {rj['count']} deployment(s)")
elif args.id is not None:
print(f"Deleted deployment {args.id}.")
else:
print("Deployment deleted successfully")
else:
print(rj.get("msg", rj))
cli/commands/endpoints.py
"""CLI commands for managing endpoints and workergroups."""
import json
import argparse
from vastai.cli.parser import argument
from vastai.cli.display import deindent
from vastai.api import endpoints as endpoints_api
from vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401
parser = _get_parser()
# ---------------------------------------------------------------------------
# endpoints
# ---------------------------------------------------------------------------
@parser.command(
argument("--min_load", help="minimum floor load in perf units/s (token/s for LLms)", type=float, default=0.0),
argument("--min_cold_load", help="minimum floor load in perf units/s (token/s for LLms), but allow handling with cold workers", type=float, default=0.0),
argument("--target_util", help="target capacity utilization (fraction, max 1.0, default 0.9)", type=float, default=0.9),
argument("--cold_mult", help="cold/stopped instance capacity target as multiple of hot capacity target (default 2.5)", type=float, default=2.5),
argument("--cold_workers", help="min number of workers to keep 'cold' when you have no load (default 5)", type=int, default=5),
argument("--max_workers", help="max number of workers your endpoint group can have (default 20)", type=int, default=20),
argument("--endpoint_name", help="deployment endpoint name (allows multiple autoscale groups to share same deployment endpoint)", type=str),
argument("--max_queue_time", help="maximum seconds requests may be queued on each worker (default 30.0)", type=float),
argument("--target_queue_time", help="target seconds for the queue to be cleared (default 10.0)", type=float),
argument("--inactivity_timeout", help="seconds of no traffic before the endpoint can scale to zero active workers", type=int),
argument("--auto_instance", help=argparse.SUPPRESS, type=str, default="prod"),
usage="vastai create endpoint [OPTIONS]",
help="Create a new endpoint group",
epilog=deindent("""
Create a new endpoint group to manage many autoscaling groups
Example: vastai create endpoint --target_util 0.9 --cold_mult 2.0 --endpoint_name "LLama"
"""),
)
def create__endpoint(args):
"""Create a new endpoint group."""
if args.explain:
print("request json: ")
print({
"client_id": "me", "min_load": args.min_load, "min_cold_load": args.min_cold_load,
"target_util": args.target_util, "cold_mult": args.cold_mult,
"cold_workers": args.cold_workers, "max_workers": args.max_workers,
"endpoint_name": args.endpoint_name, "autoscaler_instance": args.auto_instance,
})
client = get_client(args)
result = endpoints_api.create_endpoint(
client, min_load=args.min_load, min_cold_load=args.min_cold_load,
target_util=args.target_util, cold_mult=args.cold_mult,
cold_workers=args.cold_workers, max_workers=args.max_workers,
endpoint_name=args.endpoint_name, auto_instance=args.auto_instance,
max_queue_time=args.max_queue_time, target_queue_time=args.target_queue_time,
inactivity_timeout=args.inactivity_timeout,
)
if args.raw:
return result
print("create endpoint {}".format(result))
@parser.command(
usage="vastai show endpoints [--api-key API_KEY]",
help="Display user's current endpoint groups",
epilog=deindent("""
Example: vastai show endpoints
"""),
)
def show__endpoints(args):
"""Display user's current endpoint groups."""
if args.explain:
print("request json: ")
print({"client_id": "me", "api_key": args.api_key})
client = get_client(args)
result = endpoints_api.show_endpoints(client)
if isinstance(result, dict) and "error" in result:
print(result["error"])
return
if args.raw:
return result
else:
print(json.dumps(result, indent=1, sort_keys=True))
@parser.command(
argument("id", help="id of endpoint group to update", type=int),
argument("--min_load", help="minimum floor load in perf units/s (token/s for LLms)", type=float),
argument("--min_cold_load", help="minimum floor load in perf units/s (token/s for LLms), but allow handling with cold workers", type=float),
argument("--endpoint_state", help="active, suspended, or stopped", type=str),
argument("--auto_instance", help=argparse.SUPPRESS, type=str, default="prod"),
argument("--target_util", help="target capacity utilization (fraction, max 1.0, default 0.9)", type=float),
argument("--cold_mult", help="cold/stopped instance capacity target as multiple of hot capacity target (default 2.5)", type=float),
argument("--cold_workers", help="min number of workers to keep 'cold' when you have no load (default 5)", type=int),
argument("--max_workers", help="max number of workers your endpoint group can have (default 20)", type=int),
argument("--endpoint_name", help="deployment endpoint name (allows multiple workergroups to share same deployment endpoint)", type=str),
argument("--max_queue_time", help="maximum seconds requests may be queued on each worker (default 30.0)", type=float),
argument("--target_queue_time", help="target seconds for the queue to be cleared (default 10.0)", type=float),
argument("--inactivity_timeout", help="seconds of no traffic before the endpoint can scale to zero active workers", type=int),
usage="vastai update endpoint ID [OPTIONS]",
help="Update an existing endpoint group",
epilog=deindent("""
Example: vastai update endpoint 4242 --min_load 100 --target_util 0.9 --cold_mult 2.0 --endpoint_name "LLama"
"""),
)
def update__endpoint(args):
"""Update an existing endpoint group."""
client = get_client(args)
result = endpoints_api.update_endpoint(
client, id=args.id,
min_load=args.min_load, min_cold_load=args.min_cold_load,
target_util=args.target_util, cold_mult=args.cold_mult,
cold_workers=args.cold_workers, max_workers=args.max_workers,
endpoint_name=args.endpoint_name, endpoint_state=args.endpoint_state,
auto_instance=args.auto_instance,
max_queue_time=args.max_queue_time, target_queue_time=args.target_queue_time,
inactivity_timeout=args.inactivity_timeout,
)
if args.raw:
return result
print("update endpoint {}".format(result))
@parser.command(
argument("id", help="id of endpoint group to delete", type=int),
usage="vastai delete endpoint ID ",
help="Delete an endpoint group",
epilog=deindent("""
Example: vastai delete endpoint 4242
"""),
)
def delete__endpoint(args):
"""Delete an endpoint group."""
id = args.id
if args.explain:
print("request json: ")
print({"client_id": "me", "endptjob_id": args.id})
client = get_client(args)
result = endpoints_api.delete_endpoint(client, id=id)
print("delete endpoint {}".format(result))
@parser.command(
argument("id", help="id of endpoint group to fetch logs from", type=int),
argument("--level", help="log detail level (0 to 3)", type=int, default=1),
argument("--tail", help="", type=int, default=None),
usage="vastai get endpt-logs ID [--api-key API_KEY]",
help="Fetch logs for a specific serverless endpoint group",
epilog=deindent("""
Example: vastai get endpt-logs 382
"""),
)
def get__endpt_logs(args):
"""Fetch logs for a specific serverless endpoint group."""
if args.explain:
print(f"Fetching endpoint logs for id={args.id}")
client = get_client(args)
rj = endpoints_api.get_endpt_logs(client, id=args.id, level=args.level, tail=args.tail)
levels = {0: "info0", 1: "info1", 2: "trace", 3: "debug"}
if isinstance(rj, dict) and "error" in rj:
print(rj["error"])
return
if args.raw:
return rj
else:
dbg_lvl = levels[args.level]
if rj and dbg_lvl:
print(rj[dbg_lvl])
@parser.command(
argument("id", help="id of endpoint group whose workers to list", type=int),
usage="vastai get endpt-workers ID [--api-key API_KEY]",
help="List workers on an endpoint group with status and measured_perf",
epilog=deindent("""
Example: vastai get endpt-workers 21979
"""),
)
def get__endpt_workers(args):
"""List workers on a specific serverless endpoint group."""
client = get_client(args)
rj = endpoints_api.get_endpoint_workers(client, id=args.id)
if isinstance(rj, dict) and "error" in rj:
print(rj["error"])
return
workers = rj if isinstance(rj, list) else rj.get("workers", [])
if args.raw:
return workers
if not workers:
print("No workers found.")
return
for w in workers:
print(f"id={w.get('id')} status={w.get('status')} "
f"gpu_name={w.get('gpu_name')} measured_perf={w.get('measured_perf')}")
# ---------------------------------------------------------------------------
# workergroups
# ---------------------------------------------------------------------------
@parser.command(
argument("--template_hash", help="template hash (required, but **Note**: if you use this field, you can skip search_params, as they are automatically inferred from the template)", type=str),
argument("--template_id", help="template id (optional)", type=int),
argument("-n", "--no-default", action="store_true", help="Disable default search param query args"),
argument("--launch_args", help="launch args string for create instance ex: \"--onstart onstart_wget.sh --env '-e ONSTART_PATH=https://s3.amazonaws.com/vast.ai/onstart_OOBA.sh' --image atinoda/text-generation-webui:default-nightly --disk 64\"", type=str),
argument("--endpoint_name", help="deployment endpoint name (allows multiple workergroups to share same deployment endpoint)", type=str),
argument("--endpoint_id", help="deployment endpoint id (allows multiple workergroups to share same deployment endpoint)", type=int),
argument("--gpu_ram", help="estimated GPU RAM req (independent of search string)", type=float),
argument("--search_params", help="search param string for search offers ex: \"gpu_ram>=23 num_gpus=2 gpu_name=RTX_4090 inet_down>200 direct_port_count>2 disk_space>=64\"", type=str),
argument("--auto_instance", help=argparse.SUPPRESS, type=str, default="prod"),
usage="vastai create workergroup [OPTIONS]",
help="Create a new autoscale group",
epilog=deindent("""
Create a new autoscaling group to manage a pool of worker instances.
Example: vastai create workergroup --template_hash HASH --endpoint_name "LLama"
"""),
)
def create__workergroup(args):
"""Create a new workergroup."""
if args.explain:
print("request json: ")
print({"template_hash": args.template_hash, "search_params": args.search_params})
client = get_client(args)
try:
result = endpoints_api.create_workergroup(
client, template_hash=args.template_hash, template_id=args.template_id,
no_default=args.no_default, launch_args=args.launch_args,
endpoint_name=args.endpoint_name, endpoint_id=args.endpoint_id,
gpu_ram=args.gpu_ram, search_params=args.search_params,
auto_instance=args.auto_instance,
)
print("workergroup create {}".format(result))
except Exception as e:
print(f"Error creating workergroup: {e}")
@parser.command(
usage="vastai show workergroups [--api-key API_KEY]",
help="Display user's current workergroups",
epilog=deindent("""
Example: vastai show workergroups
"""),
)
def show__workergroups(args):
"""Display user's current workergroups."""
if args.explain:
print("request json: ")
print({"client_id": "me", "api_key": args.api_key})
client = get_client(args)
result = endpoints_api.show_workergroups(client)
if isinstance(result, dict) and "error" in result:
print(result["error"])
return
if args.raw:
return result
else:
print(json.dumps(result, indent=1, sort_keys=True))
@parser.command(
argument("id", help="id of autoscale group to update", type=int),
argument("--gpu_ram", help="estimated GPU RAM req (independent of search string)", type=float),
argument("--template_hash", help="template hash (**Note**: if you use this field, you can skip search_params, as they are automatically inferred from the template)", type=str),
argument("--template_id", help="template id", type=int),
argument("--search_params", help="search param string for search offers ex: \"gpu_ram>=23 num_gpus=2 gpu_name=RTX_4090 inet_down>200 direct_port_count>2 disk_space>=64\"", type=str),
argument("-n", "--no-default", action="store_true", help="Disable default search param query args"),
argument("--launch_args", help="launch args string for create instance ex: \"--onstart onstart_wget.sh --env '-e ONSTART_PATH=https://s3.amazonaws.com/public.vast.ai/onstart_OOBA.sh' --image atinoda/text-generation-webui:default-nightly --disk 64\"", type=str),
argument("--endpoint_name", help="deployment endpoint name (allows multiple workergroups to share same deployment endpoint)", type=str),
argument("--endpoint_id", help="deployment endpoint id (allows multiple workergroups to share same deployment endpoint)", type=int),
usage="vastai update workergroup WORKERGROUP_ID --endpoint_id ENDPOINT_ID [options]",
help="Update an existing autoscale group",
epilog=deindent("""
Example: vastai update workergroup 4242 --search_params \"gpu_ram>=23 num_gpus=2 gpu_name=RTX_4090 inet_down>200 direct_port_count>2 disk_space>=64\" --launch_args \"--onstart onstart_wget.sh --env '-e ONSTART_PATH=https://s3.amazonaws.com/public.vast.ai/onstart_OOBA.sh' --image atinoda/text-generation-webui:default-nightly --disk 64\" --gpu_ram 32.0 --endpoint_name "LLama" --endpoint_id 2
"""),
)
def update__workergroup(args):
"""Update an existing workergroup."""
client = get_client(args)
result = endpoints_api.update_workergroup(
client, id=args.id,
gpu_ram=args.gpu_ram,
template_hash=args.template_hash, template_id=args.template_id,
search_params=args.search_params, no_default=args.no_default,
launch_args=args.launch_args, endpoint_name=args.endpoint_name,
endpoint_id=args.endpoint_id,
)
if args.raw:
return result
print("workergroup update {}".format(result))
@parser.command(
argument("id", help="id of workergroup to update workers for", type=int),
argument("--cancel", action="store_true", help="cancel an in-progress update for the workergroup"),
usage="vastai update workers WORKERGROUP_ID [--cancel]",
help="Trigger a rolling update of all workers in a workergroup, or cancel an in-progress update",
epilog=deindent("""
Starts a rolling update of all workers in the specified workergroup. The autoscaler
will cycle through workers, updating them while maintaining capacity.
Use --cancel to cancel an update that is currently in progress.
Examples:
vastai update workers 4242
vastai update workers 4242 --cancel
"""),
)
def update__workers(args):
"""Trigger a rolling update of workers in a workergroup."""
client = get_client(args)
result = endpoints_api.update_workers(client, id=args.id, cancel=args.cancel)
if args.raw:
return result
if result.get("success"):
if result.get("cancelled"):
print(f"Update cancelled for workergroup {args.id}")
else:
print(f"Update started for workergroup {args.id} ({result.get('workers_to_update', 0)} workers)")
else:
print(f"Error: {result.get('error_msg', 'unknown error')}")
@parser.command(
argument("id", help="id of group to delete", type=int),
usage="vastai delete workergroup ID ",
help="Delete a workergroup group",
epilog=deindent("""
Note that deleting a workergroup doesn't automatically destroy all the instances that are associated with your workergroup.
Example: vastai delete workergroup 4242
"""),
)
def delete__workergroup(args):
"""Delete a workergroup."""
id = args.id
if args.explain:
print("request json: ")
print({"client_id": "me", "autojob_id": args.id})
client = get_client(args)
result = endpoints_api.delete_workergroup(client, id=id)
print("workergroup delete {}".format(result))
@parser.command(
argument("id", help="id of workergroup to fetch logs from", type=int),
argument("--level", help="log detail level (0 to 3)", type=int, default=1),
argument("--tail", help="", type=int, default=None),
usage="vastai get wrkgrp-logs ID [--api-key API_KEY]",
help="Fetch logs for a specific serverless worker group",
epilog=deindent("""
Example: vastai get endpt-logs 382
"""),
)
def get__wrkgrp_logs(args):
"""Fetch logs for a specific serverless worker group."""
if args.explain:
print(f"Fetching workergroup logs for id={args.id}")
client = get_client(args)
rj = endpoints_api.get_wrkgrp_logs(client, id=args.id, level=args.level, tail=args.tail)
levels = {0: "info0", 1: "info1", 2: "trace", 3: "debug"}
if isinstance(rj, dict) and "error" in rj:
print(rj["error"])
return
if args.raw:
return rj
else:
dbg_lvl = levels[args.level]
if rj and dbg_lvl:
print(rj[dbg_lvl])
cli/commands/instances.py
"""CLI commands for managing instances."""
import json
import sys
import time
import argparse
import subprocess
from vastai.cli.parser import argument, hidden_aliases, MyWideHelpFormatter
from vastai.cli.util import (
default_start_date, default_end_date,
parse_day_cron_style, parse_hour_cron_style,
validate_frequency_values, add_scheduled_job,
parse_env, split_list, exec_with_threads,
_get_gpu_names,
)
from vastai.api import instances as instances_api
from vastai.api import offers as offers_api
from vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401
# ---------------------------------------------------------------------------
# Instance display
# ---------------------------------------------------------------------------
from vastai.cli.display import display_table, print_or_page, instance_fields, deindent
parser = _get_parser()
# ---------------------------------------------------------------------------
# show instances (defined further down, after the display helpers it uses)
# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------
# show instance (single)
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of instance to get", type=int),
argument("-q", "--quiet", action="store_true", help="only display numeric id"),
usage="vastai show instance ID [options]",
help="Display user's current instances",
)
def show__instance(args):
"""Shows stats for a single instance."""
client = get_client(args)
result = instances_api.show_instance(client, id=args.id)
if result is None:
if args.raw:
return {"instances": None}
if not args.quiet:
print(f"Instance {args.id} not found or no longer exists.", file=sys.stderr)
return 1
if args.raw:
return result if result is not None else {"instances": None}
if result is None:
print(f"Instance {args.id} not found or no longer exists.")
return 1
if args.quiet:
print(result.get("id", ""))
else:
display_table([result], instance_fields)
# ---------------------------------------------------------------------------
# create instance
# ---------------------------------------------------------------------------
def cli_runtype(args):
"""Resolve the runtype from the CLI flags, keeping the CLI's 'ssh' default."""
if args.template_hash is not None:
return None
runtype, args.args = instances_api.resolve_runtype(
ssh=args.ssh, jupyter=args.jupyter, direct=args.direct, args=args.args,
jupyter_lab=args.jupyter_lab, jupyter_dir=args.jupyter_dir, default='ssh',
)
return runtype
def create_instance_impl(id, args):
if args.onstart:
with open(args.onstart, "r") as reader:
args.onstart_cmd = reader.read()
if args.onstart_cmd is None:
args.onstart_cmd = args.entrypoint
try:
runtype = cli_runtype(args)
except ValueError as e:
print(f"Error: {e}", file=sys.stderr)
return 1
json_blob = instances_api.build_create_instance_payload(
image=args.image, disk=args.disk, env=args.env, bid_price=args.bid_price,
label=args.label, extra=args.extra, onstart_cmd=args.onstart_cmd,
login=args.login, python_utf8=args.python_utf8, lang_utf8=args.lang_utf8,
jupyter_lab=args.jupyter_lab, jupyter_dir=args.jupyter_dir,
force=args.force, cancel_unavail=args.cancel_unavail,
template_hash=args.template_hash, user=args.user, runtype=runtype,
args=args.args, create_volume=args.create_volume,
link_volume=args.link_volume, volume_size=args.volume_size,
mount_path=args.mount_path, volume_label=args.volume_label,
)
if args.explain:
print("request json: ")
print(json_blob)
client = get_client(args)
rj = instances_api.create_instance_from_payload(client, id, json_blob)
if args.raw:
return rj
else:
print("Started. {}".format(rj))
return True
_create_instance_args = [
argument("--template_hash", help="Create instance from template info", type=str),
argument("--user", help="User to use with docker create. This breaks some images, so only use this if you are certain you need it.", type=str),
argument("--disk", help="size of local disk partition in GB", type=float, default=10),
argument("--image", help="docker container image to launch", type=str),
argument("--login", help="docker login arguments for private repo authentication, surround with '' ", type=str),
argument("--label", help="label to set on the instance", type=str),
argument("--onstart", help="filename to use as onstart script", type=str),
argument("--onstart-cmd", help="contents of onstart script as single argument", type=str),
argument("--entrypoint", help="override entrypoint for args launch instance", type=str),
argument("--ssh", help="Launch as an ssh instance type", action="store_true"),
argument("--jupyter", help="Launch as a jupyter instance instead of an ssh instance", action="store_true"),
argument("--direct", help="Use (faster) direct connections for jupyter & ssh", action="store_true"),
argument("--jupyter-dir", help="For runtype 'jupyter', directory in instance to use to launch jupyter. Defaults to image's working directory", type=str),
argument("--jupyter-lab", help="For runtype 'jupyter', Launch instance with jupyter lab", action="store_true"),
argument("--lang-utf8", help="Workaround for images with locale problems: install and generate locales before instance launch, and set locale to C.UTF-8", action="store_true"),
argument("--python-utf8", help="Workaround for images with locale problems: set python's locale to C.UTF-8", action="store_true"),
argument("--extra", help=argparse.SUPPRESS),
argument("--env", help="env variables and port mapping options, surround with '' ", type=str),
argument("--args", nargs=argparse.REMAINDER, help="list of arguments passed to container ENTRYPOINT. Onstart is recommended for this purpose. (must be last argument)"),
argument("--force", help="Skip sanity checks when creating from an existing instance", action="store_true"),
argument("--cancel-unavail", help="Return error if scheduling fails (rather than creating a stopped instance)", action="store_true"),
argument("--bid_price", help="(OPTIONAL) create an INTERRUPTIBLE instance with per machine bid price in $/hour", type=float),
argument("--create-volume", metavar="VOLUME_ASK_ID", help="Create a new local volume using an ID returned from the \"search volumes\" command and link it to the new instance", type=int),
argument("--link-volume", metavar="EXISTING_VOLUME_ID", help="ID of an existing rented volume to link to the instance during creation. (returned from \"show volumes\" cmd)", type=int),
argument("--volume-size", help="Size of the volume to create in GB. Only usable with --create-volume (default 15GB)", type=int),
argument("--mount-path", help="The path to the volume from within the new instance container. e.g. /root/volume", type=str),
argument("--volume-label", help="(optional) A name to give the new volume. Only usable with --create-volume", type=str),
]
@parser.command(
argument("id", help="id of instance type to launch (returned from search offers)", type=int),
*_create_instance_args,
usage="vastai create instance ID [OPTIONS] [--args ...]",
help="Create a new instance",
epilog=deindent("""
Performs the same action as pressing the "RENT" button on the website at https://console.vast.ai/create/
Creates an instance from an offer ID (which is returned from "search offers"). Each offer ID can only be used to create one instance.
Besides the offer ID, you must pass in an '--image' argument as a minimum.
If you use args/entrypoint launch mode, we create a container from your image as is, without attempting to inject ssh and or jupyter.
If you use the args launch mode, you can override the entrypoint with --entrypoint, and pass arguments to the entrypoint with --args.
If you use --args, that must be the last argument, as any following tokens are consumed into the args string.
For ssh/jupyter launch types, use --onstart-cmd to pass in startup script, instead of --entrypoint and --args.
Examples:
# create an on-demand instance with the PyTorch (cuDNN Devel) template and 64GB of disk
vastai create instance 384826 --template_hash 661d064bbda1f2a133816b6d55da07c3 --disk 64
# create an on-demand instance with the pytorch/pytorch image, 40GB of disk, open 8081 udp, direct ssh, set hostname to billybob, and a small onstart script
vastai create instance 6995713 --image pytorch/pytorch --disk 40 --env '-p 8081:8081/udp -h billybob' --ssh --direct --onstart-cmd "env | grep _ >> /etc/environment; echo 'starting up'";
# create an on-demand instance with the bobsrepo/pytorch:latest image, 20GB of disk, open 22, 8080, jupyter ssh, and set some env variables
vastai create instance 384827 --image bobsrepo/pytorch:latest --login '-u bob -p 9d8df!fd89ufZ docker.io' --jupyter --direct --env '-e TZ=PDT -e XNAME=XX4 -p 22:22 -p 8080:8080' --disk 20
# create an on-demand instance with the pytorch/pytorch image, 40GB of disk, override the entrypoint to bash and pass bash a simple command to keep the instance running. (args launch without ssh/jupyter)
vastai create instance 5801802 --image pytorch/pytorch --disk 40 --onstart-cmd 'bash' --args -c 'echo hello; sleep infinity;'
# create an interruptible (spot) instance with the PyTorch (cuDNN Devel) template, 64GB of disk, and a bid price of $0.10/hr
vastai create instance 384826 --template_hash 661d064bbda1f2a133816b6d55da07c3 --disk 64 --bid_price 0.1
Return value:
Returns a json reporting the instance ID of the newly created instance:
{'success': True, 'new_contract': 7835610}
"""),
)
def create__instance(args):
"""Create an instance from an offer ID."""
rj = create_instance_impl(args.id, args)
if args.raw:
return rj
@parser.command(
argument("ids", help="ids of instance types to launch (returned from search offers)", type=int, nargs='+'),
*_create_instance_args,
usage="vastai create instances [OPTIONS] ID0 ID1 ID2... [--args ...]",
help="Create instances from a list of offers",
)
def create__instances(args):
"""Bulk create instances."""
idlist = split_list(args.ids, 64)
exec_with_threads(lambda ids: create_instance_impl(ids, args), idlist, nt=8)
# ---------------------------------------------------------------------------
# destroy instance
# ---------------------------------------------------------------------------
def destroy_instance_impl(id, args):
client = get_client(args)
rj = instances_api.destroy_instance(client, id=id)
if args.raw:
return rj
elif rj.get("success"):
print("destroying instance {id}.".format(**(locals())))
return True
else:
print(rj.get("msg", rj))
return False
@parser.command(
argument("id", help="id of instance to delete", type=int),
argument("-y", "--yes", help="Skip confirmation prompt", action="store_true"),
usage="vastai destroy instance id [-h] [-y] [--api-key API_KEY] [--raw]",
help="Destroy an instance (irreversible, deletes data)",
epilog=deindent("""
Perfoms the same action as pressing the "DESTROY" button on the website at https://console.vast.ai/instances/
Example: vastai destroy instance 4242
Use -y or --yes to skip the confirmation prompt.
"""),
)
def destroy__instance(args):
"""Destroy a single instance."""
if not args.yes:
try:
confirm = input(f"Are you sure you want to destroy instance {args.id}? This is irreversible and will delete all data. [y/N] ").strip().lower()
except (EOFError, KeyboardInterrupt):
print("\nAborted.")
return
if confirm != "y":
print("Aborted.")
return
destroy_instance_impl(args.id, args)
@parser.command(
argument("ids", help="ids of instance to destroy", type=int, nargs='+'),
argument("-y", "--yes", help="Skip confirmation prompt", action="store_true"),
usage="vastai destroy instances [--raw] [-y] <id> [<id> ...]",
help="Destroy a list of instances (irreversible, deletes data)",
)
def destroy__instances(args):
"""Bulk destroy instances."""
if not args.yes:
id_list_str = ", ".join(str(i) for i in args.ids)
try:
confirm = input(f"Are you sure you want to destroy instances {id_list_str}? This is irreversible and will delete all data. [y/N] ").strip().lower()
except (EOFError, KeyboardInterrupt):
print("\nAborted.")
return
if confirm != "y":
print("Aborted.")
return
idlist = split_list(args.ids, 64)
exec_with_threads(lambda ids: destroy_instance_impl(ids, args), idlist, nt=8)
# ---------------------------------------------------------------------------
# start/stop/reboot/recycle instance
# ---------------------------------------------------------------------------
def start_instance_impl(id, args):
client = get_client(args)
rj = instances_api.start_instance(client, id=id)
if args.explain:
print("request json: ")
print({"state": "running"})
if rj.get("success"):
print("starting instance {id}.".format(**(locals())))
return True
else:
print(rj.get("msg", rj))
return False
@parser.command(
argument("id", help="ID of instance to start/restart", type=int),
usage="vastai start instance ID [OPTIONS]",
help="Start a stopped instance",
epilog=deindent("""
This command attempts to bring an instance from the "stopped" state into the "running" state. This is subject to resource availability on the machine that the instance is located on.
If your instance is stuck in the "scheduling" state for more than 30 seconds after running this, it likely means that the required resources on the machine to run your instance are currently unavailable.
Examples:
vastai start instances $(vastai show instances -q)
vastai start instance 329838
"""),
)
def start__instance(args):
"""Start a stopped instance."""
start_instance_impl(args.id, args)
@parser.command(
argument("ids", help="ids of instance to start", type=int, nargs='+'),
usage="vastai start instances [OPTIONS] ID0 ID1 ID2...",
help="Start a list of instances",
)
def start__instances(args):
"""Bulk start instances."""
idlist = split_list(args.ids, 64)
exec_with_threads(lambda ids: start_instance_impl(ids, args), idlist, nt=8)
def stop_instance_impl(id, args):
client = get_client(args)
rj = instances_api.stop_instance(client, id=id)
if args.explain:
print("request json: ")
print({"state": "stopped"})
if rj.get("success"):
print("stopping instance {id}.".format(**(locals())))
return True
else:
print(rj.get("msg", rj))
return False
@parser.command(
argument("id", help="id of instance to stop", type=int),
usage="vastai stop instance ID [OPTIONS]",
help="Stop a running instance",
epilog=deindent("""
This command brings an instance from the "running" state into the "stopped" state. When an instance is "stopped" all of your data on the instance is preserved,
and you can resume use of your instance by starting it again. Once stopped, starting an instance is subject to resource availability on the machine that the instance is located on.
There are ways to move data off of a stopped instance, which are described here: https://vast.ai/docs/gpu-instances/data-movement
"""),
)
def stop__instance(args):
"""Stop a running instance."""
stop_instance_impl(args.id, args)
@parser.command(
argument("ids", help="ids of instance to stop", type=int, nargs='+'),
usage="vastai stop instances [OPTIONS] ID0 ID1 ID2...",
help="Stop a list of instances",
epilog=deindent("""
Examples:
vastai stop instances $(vastai show instances -q)
vastai stop instances 329838 984849
"""),
)
def stop__instances(args):
"""Bulk stop instances."""
idlist = split_list(args.ids, 64)
exec_with_threads(lambda ids: stop_instance_impl(ids, args), idlist, nt=8)
@parser.command(
argument("id", help="id of instance to reboot", type=int),
argument("--schedule", choices=["HOURLY", "DAILY", "WEEKLY"], help="try to schedule a command to run hourly, daily, or weekly. Valid values are HOURLY, DAILY, WEEKLY For ex. --schedule DAILY"),
argument("--start_date", type=str, default=default_start_date(), help="Start date/time in format 'YYYY-MM-DD HH:MM:SS PM' (UTC). Default is now. (optional)"),
argument("--end_date", type=str, default=default_end_date(), help="End date/time in format 'YYYY-MM-DD HH:MM:SS PM' (UTC). Default is 7 days from now. (optional)"),
argument("--day", type=parse_day_cron_style, help="Day of week you want scheduled job to run on (0-6, where 0=Sunday) or \"*\". Default will be 0. For ex. --day 0", default=0),
argument("--hour", type=parse_hour_cron_style, help="Hour of day you want scheduled job to run on (0-23) or \"*\" (UTC). Default will be 0. For ex. --hour 16", default=0),
usage="vastai reboot instance ID [OPTIONS]",
help="Reboot (stop/start) an instance",
epilog=deindent("""
Stops and starts container without any risk of losing GPU priority.
"""),
)
def reboot__instance(args):
"""Reboot an instance."""
client = get_client(args)
rj = instances_api.reboot_instance(client, id=args.id)
if args.schedule:
validate_frequency_values(args.day, args.hour, args.schedule)
cli_command = "reboot instance"
api_endpoint = "/api/v0/instances/reboot/{id}/".format(id=args.id)
json_blob = {"instance_id": args.id}
add_scheduled_job(client, args, json_blob, cli_command, api_endpoint, "PUT", instance_id=args.id)
return
if rj.get("success"):
print("Rebooting instance {args.id}.".format(**(locals())))
else:
print(rj.get("msg", rj))
@parser.command(
argument("id", help="id of instance to recycle", type=int),
usage="vastai recycle instance ID [OPTIONS]",
help="Recycle (destroy/create) an instance",
epilog=deindent("""
Destroys and recreates container in place (from newly pulled image) without any risk of losing GPU priority.
"""),
)
def recycle__instance(args):
"""Recycle an instance."""
client = get_client(args)
rj = instances_api.recycle_instance(client, id=args.id)
if rj.get("success"):
print("Recycling instance {args.id}.".format(**(locals())))
else:
print(rj.get("msg", rj))
# ---------------------------------------------------------------------------
# update instance
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of instance to update", type=int),
argument("--template_id", help="new template ID to associate with the instance", type=int),
argument("--template_hash_id", help="new template hash ID to associate with the instance", type=str),
argument("--image", help="new image UUID for the instance", type=str),
argument("--args", help="new arguments for the instance", type=str),
argument("--env", help="new environment variables for the instance", type=json.loads),
argument("--onstart", help="new onstart script for the instance", type=str),
usage="vastai update instance ID [OPTIONS]",
help="Update recreate an instance from a new/updated template",
epilog=deindent("""
Example: vastai update instance 1234 --template_hash_id 661d064bbda1f2a133816b6d55da07c3
"""),
)
def update__instance(args):
"""Update/recreate an instance from a new or updated template."""
client = get_client(args)
rj = instances_api.update_instance(
client, id=args.id,
template_id=args.template_id,
template_hash_id=args.template_hash_id,
image=args.image,
args=args.args,
env=args.env,
onstart=args.onstart,
)
if args.raw:
return rj
if rj.get("success"):
print(f"Instance {args.id} updated successfully.")
if rj.get("updated_instance"):
print("Updated instance details:")
print(rj.get("updated_instance"))
else:
print(f"Failed to update instance {args.id}: {rj.get('msg')}")
# ---------------------------------------------------------------------------
# label / prepay / change bid
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of instance to label", type=int),
argument("label", help="label to set", type=str),
usage="vastai label instance <id> <label>",
help="Assign a string label to an instance",
)
def label__instance(args):
"""Set a label on an instance."""
if args.explain:
print("request json: ")
print({"label": args.label})
client = get_client(args)
rj = instances_api.label_instance(client, id=args.id, label=args.label)
if rj.get("success"):
print("label for {args.id} set to {args.label}.".format(**(locals())))
else:
print(rj.get("msg", rj))
@parser.command(
argument("id", help="id of instance to prepay for", type=int),
argument("amount", help="amount of instance credit prepayment (default discount func of 0.2 for 1 month, 0.3 for 3 months)", type=float),
usage="vastai prepay instance ID AMOUNT",
help="Deposit credits into reserved instance",
)
def prepay__instance(args):
"""Prepay for an instance."""
if args.explain:
print("request json: ")
print({"amount": args.amount})
client = get_client(args)
rj = instances_api.prepay_instance(client, id=args.id, amount=args.amount)
if rj.get("success"):
timescale = round(rj["timescale"], 3)
discount_rate = 100.0 * round(rj["discount_rate"], 3)
print("prepaid for {timescale} months of instance {args.id} applying ${args.amount} credits for a discount of {discount_rate}%".format(**(locals())))
else:
print(rj.get("msg", rj))
@parser.command(
argument("id", help="id of instance type to change bid", type=int),
argument("--price", help="per machine bid price in $/hour", type=float),
argument("--schedule", choices=["HOURLY", "DAILY", "WEEKLY"], help="try to schedule a command to run hourly, daily, or weekly. Valid values are HOURLY, DAILY, WEEKLY For ex. --schedule DAILY"),
argument("--start_date", type=str, default=default_start_date(), help="Start date/time in format 'YYYY-MM-DD HH:MM:SS PM' (UTC). Default is now. (optional)"),
argument("--end_date", type=str, default=default_end_date(), help="End date/time in format 'YYYY-MM-DD HH:MM:SS PM' (UTC). Default is 7 days from now. (optional)"),
argument("--day", type=parse_day_cron_style, help="Day of week you want scheduled job to run on (0-6, where 0=Sunday) or \"*\". Default will be 0. For ex. --day 0", default=0),
argument("--hour", type=parse_hour_cron_style, help="Hour of day you want scheduled job to run on (0-23) or \"*\" (UTC). Default will be 0. For ex. --hour 16", default=0),
usage="vastai change bid id [--price PRICE]",
help="Change the bid price for a spot/interruptible instance",
epilog=deindent("""
Change the current bid price of instance id to PRICE.
If PRICE is not specified, then a winning bid price is used as the default.
"""),
)
def change__bid(args):
"""Change the bid price for a spot instance."""
if args.explain:
print("request json: ")
print({"client_id": "me", "price": args.price})
client = get_client(args)
if args.schedule:
validate_frequency_values(args.day, args.hour, args.schedule)
cli_command = "change bid"
api_endpoint = "/api/v0/instances/bid_price/{id}/".format(id=args.id)
json_blob = {"client_id": "me", "price": args.price, "instance_id": args.id}
add_scheduled_job(client, args, json_blob, cli_command, api_endpoint, "PUT", instance_id=args.id)
return
instances_api.change_bid(client, id=args.id, price=args.price)
print("Per gpu bid price changed")
# The accept/reject/show subcommands for price-increase live in
# vastai/cli/commands/price_increase.py to keep this module focused on
# instance CRUD/bid changes.
# ---------------------------------------------------------------------------
# launch instance
# ---------------------------------------------------------------------------
@parser.command(
argument("-g", "--gpu-name", type=str, required=True, choices=_get_gpu_names(), help="Name of the GPU model, replace spaces with underscores"),
argument("-n", "--num-gpus", type=str, required=True, choices=["1", "2", "4", "8", "12", "14"], help="Number of GPUs required"),
argument("-r", "--region", type=str, help="Geographical location of the instance"),
argument("-i", "--image", required=True, help="Name of the image to use for instance"),
argument("-d", "--disk", type=float, default=16.0, help="Disk space required in GB"),
argument("--limit", default=3, type=int, help=""),
argument("-o", "--order", type=str, help="Comma-separated list of fields to sort on. postfix field with - to sort desc. ex: -o 'num_gpus,total_flops-'. default='score-'", default='score-'),
argument("--login", help="docker login arguments for private repo authentication, surround with '' ", type=str),
argument("--label", help="label to set on the instance", type=str),
argument("--onstart", help="filename to use as onstart script", type=str),
argument("--onstart-cmd", help="contents of onstart script as single argument", type=str),
argument("--entrypoint", help="override entrypoint for args launch instance", type=str),
argument("--ssh", help="Launch as an ssh instance type", action="store_true"),
argument("--jupyter", help="Launch as a jupyter instance instead of an ssh instance", action="store_true"),
argument("--direct", help="Use (faster) direct connections for jupyter & ssh", action="store_true"),
argument("--jupyter-dir", help="For runtype 'jupyter', directory in instance to use to launch jupyter. Defaults to image's working directory", type=str),
argument("--jupyter-lab", help="For runtype 'jupyter', Launch instance with jupyter lab", action="store_true"),
argument("--lang-utf8", help="Workaround for images with locale problems: install and generate locales before instance launch, and set locale to C.UTF-8", action="store_true"),
argument("--python-utf8", help="Workaround for images with locale problems: set python's locale to C.UTF-8", action="store_true"),
argument("--extra", help=argparse.SUPPRESS),
argument("--env", help="env variables and port mapping options, surround with '' ", type=str),
argument("--args", nargs=argparse.REMAINDER, help="list of arguments passed to container ENTRYPOINT. Onstart is recommended for this purpose. (must be last argument)"),
argument("--cancel-unavail", help="Return error if scheduling fails (rather than creating a stopped instance)", action="store_true"),
argument("--template_hash", help="template hash which contains all relevant information about an instance. This can be used as a replacement for other parameters describing the instance configuration", type=str),
usage="vastai launch instance [--help] [--api-key API_KEY] <gpu_name> <num_gpus> <image> [geolocation] [disk_space]",
help="Launch the top instance from the search offers based on the given parameters",
epilog=deindent("""
Launches an instance based on the given parameters. The instance will be created with the top offer from the search results.
Besides the gpu_name and num_gpus, you must pass in an '--image' argument as a minimum.
If you use args/entrypoint launch mode, we create a container from your image as is, without attempting to inject ssh and or jupyter.
If you use the args launch mode, you can override the entrypoint with --entrypoint, and pass arguments to the entrypoint with --args.
If you use --args, that must be the last argument, as any following tokens are consumed into the args string.
For ssh/jupyter launch types, use --onstart-cmd to pass in startup script, instead of --entrypoint and --args.
Examples:
# launch a single RTX 3090 instance with the pytorch image and 16 GB of disk space located anywhere
python vast.py launch instance -g RTX_3090 -n 1 -i pytorch/pytorch
# launch a 4x RTX 3090 instance with the pytorch image and 32 GB of disk space located in North America
python vast.py launch instance -g RTX_3090 -n 4 -i pytorch/pytorch -d 32.0 -r North_America
Available fields:
Name Type Description
num_gpus: int # of GPUs
gpu_name: string GPU model name
region: string Region of the instance
image: string Docker image name
disk_space: float Disk space in GB
ssh, jupyter, direct: bool Flags to specify the instance type and connection method.
env: str Environment variables and port mappings, encapsulated in single quotes.
args: list Arguments passed to the container's ENTRYPOINT, used only if '--args' is specified.
"""),
)
def launch__instance(args):
"""Launch the top instance from search offers."""
if args.onstart:
with open(args.onstart, "r") as reader:
args.onstart_cmd = reader.read()
if args.onstart_cmd is None:
args.onstart_cmd = args.entrypoint
try:
runtype = cli_runtype(args)
except ValueError as e:
print(f"Error: {e}", file=sys.stderr)
return 1
client = get_client(args)
try:
response_data = offers_api.launch_instance(
client,
gpu_name=args.gpu_name,
num_gpus=args.num_gpus,
image=args.image,
region=args.region,
disk=args.disk,
order=args.order,
limit=args.limit,
env=parse_env(args.env),
label=args.label,
extra=args.extra,
onstart_cmd=args.onstart_cmd,
login=args.login,
python_utf8=args.python_utf8,
lang_utf8=args.lang_utf8,
jupyter_lab=args.jupyter_lab,
jupyter_dir=args.jupyter_dir,
cancel_unavail=args.cancel_unavail,
template_hash=args.template_hash,
runtype=runtype,
args=args.args,
)
if args.raw:
return response_data
else:
print("Started. {}".format(response_data))
if response_data.get('success'):
print(f"Instance launched successfully: {response_data.get('new_contract')}")
else:
print(f"Failed to launch instance: {response_data.get('error')}, {response_data.get('message')}")
except Exception as err:
print(f"An error occurred: {err}")
# ---------------------------------------------------------------------------
# show instances -- display helpers (Rich tables, filtering, pagination)
# ---------------------------------------------------------------------------
_DEFAULT_INSTANCE_SELECT_COLS = [
"id", "actual_status", "label",
"num_gpus", "gpu_name", "gpu_util",
"disk_space", "disk_usage", "disk_util",
"volume_info",
"dph_total", "image_uuid",
"start_date", "verification",
]
_VERBOSE_INSTANCE_SELECT_COLS = _DEFAULT_INSTANCE_SELECT_COLS + [
"machine_id", "template_id", "template_name",
"geolocation", "inet_up", "inet_down",
"ssh_host", "ssh_port", "status_msg",
]
def _fmt_age(start_date):
"""Format seconds elapsed since start_date as e.g. '2d 3h' or '4h 15m'."""
if not start_date:
return "—"
secs = max(0, time.time() - start_date)
d, rem = divmod(int(secs), 86400)
h, rem = divmod(rem, 3600)
m, _ = divmod(rem, 60)
if d: return f"{d}d {h}h"
if h: return f"{h}h {m}m"
return f"{m}m"
def _fmt_disk(disk_usage, disk_space, disk_util):
"""Format disk as 'used/total GB (X%)' or '?/total GB'."""
total = f"{disk_space:.0f}" if disk_space is not None else "?"
if disk_usage is None or disk_usage < 0:
return f"?/{total} GB"
used = f"{disk_usage:.1f}"
if disk_util is not None and disk_util >= 0:
pct = disk_util * 100
return f"{used}/{total} GB ({pct:.0f}%)"
return f"{used}/{total} GB"
def _fmt_volumes(volume_info):
"""Format volume_info list as a compact string showing IDs and usage."""
if not volume_info:
return "—"
if len(volume_info) == 1:
v = volume_info[0]
vid = v.get("id", "?")
avail = v.get("avail_space")
total = v.get("total_space")
if avail is not None and total is not None:
used = total - avail
return f"#{vid} {used:.0f}/{total:.0f} GB"
return f"#{vid}"
return ", ".join(f"#{v.get('id', '?')}" for v in volume_info)
def _fmt_gpu(num_gpus, gpu_name, gpu_util):
"""Format as '4x RTX 3090' or '4x RTX 3090 (72%)'."""
base = f"{int(num_gpus)}x {gpu_name}" if num_gpus and gpu_name else (gpu_name or "—")
if gpu_util is not None and gpu_util >= 0:
return f"{base} ({gpu_util:.0f}%)"
return base
_STATUS_COLORS = {"running": "bold green", "loading": "bold yellow", "exited": "bright_red", "created": "bright_white"}
_VERIF_COLORS = {"verified": "sea_green2", "unverified": "gold1", "deverified": "bright_red"}
def _status_style(status):
return _STATUS_COLORS.get(status, "white")
def _verif_style(v):
return _VERIF_COLORS.get(v, "white")
_INSTANCE_COL_MAX_WIDTHS = {
"gpu": 20,
"image": 30,
"age": 8,
"volumes": 17,
"location": 22,
"net": 11,
"ssh": 22,
"template": 32,
"msg": 30,
}
# Column spec: (name, header, style, justify, min_width, drop_order, verbose_only)
_INSTANCE_COL_SPECS = [
("id", "ID", "bright_white", "right", 4, 0, False),
("status", "Status", None, "center", 7, 0, False),
("label", "Label", "bright_white", "left", 7, 0, False),
("gpu", "GPU", "steel_blue1", "left", 13, 0, False),
("disk", "Disk", "bright_white", "right", 8, 5, False),
("volumes", "Volumes", "bright_white", "left", 10, 6, False),
("dph", "$/hr", "sea_green2", "right", 7, 1, False),
("image", "Image", "orchid", "left", 30, 2, False),
("age", "Age", "bright_white", "left", 8, 3, False),
("verified", "Verified", None, "center", 10, 0, False),
("machine", "Machine", "gold1", "center", 5, 8, True),
("net", "Net Mbps", "bright_white", "left", 9, 9, True),
("location", "Location", "bright_white", "center", 10, 10, True),
("template", "Template", "bright_white", "center", 20, 11, True),
("ssh", "SSH", "cyan", "left", 21, 7, True),
("msg", "Msg", "dim white", "left", 15, 12, True),
]
_INSTANCE_COL_SPEC_BY_NAME = {s[0]: s for s in _INSTANCE_COL_SPECS}
try:
from rich.text import Text as _RichText
except ImportError:
_RichText = None
def _render_instance_col(name, inst):
"""Render a single cell value for the given column name."""
if name == "id":
return str(inst.get("id", "—"))
if name == "status":
s = inst.get("actual_status") or "—"
return _RichText(s, style=_status_style(s))
if name == "label":
return inst.get("label") or "—"
if name == "gpu":
return _fmt_gpu(inst.get("num_gpus"), inst.get("gpu_name"), inst.get("gpu_util"))
if name == "disk":
return _fmt_disk(inst.get("disk_usage"), inst.get("disk_space"), inst.get("disk_util"))
if name == "volumes":
return _fmt_volumes(inst.get("volume_info") or [])
if name == "dph":
dph = inst.get("dph_total")
return f"${dph:.4f}" if dph is not None else "—"
if name == "image":
return (inst.get("image_uuid") or "—")[:50]
if name == "age":
return _fmt_age(inst.get("start_date"))
if name == "verified":
v = inst.get("verification") or "—"
return _RichText(v, style=_verif_style(v))
if name == "ssh":
return f"{inst.get('ssh_host')}:{inst.get('ssh_port', '')}" if inst.get("ssh_host") else "—"
if name == "machine":
return str(inst.get("machine_id", "—"))
if name == "net":
up, down = inst.get("inet_up"), inst.get("inet_down")
return f"↑{up:.0f} ↓{down:.0f}" if (up is not None and down is not None) else "—"
if name == "location":
return inst.get("geolocation") or "—"
if name == "template":
tid, tname = inst.get("template_id"), inst.get("template_name") or ""
return (f"{tname[:28]} ({tid})" if tid else tname[:30] or "—")
if name == "msg":
return (inst.get("status_msg") or "—")[:40]
return "—"
def _estimate_table_width(specs):
"""Estimate rendered table width for a list of col specs."""
n = len(specs)
content = sum(
_INSTANCE_COL_MAX_WIDTHS.get(s[0]) or max(len(s[1]), s[4])
for s in specs
)
return 4 + 2 * n + (n - 1) + content
def _build_instances_table(instances, verbose=False, cols=None):
"""Build the Rich table for instances."""
import shutil
from rich.table import Table
from rich import box
show_volumes = any(inst.get("volume_info") for inst in instances)
term_width = shutil.get_terminal_size((120, 24)).columns
if cols is not None:
active = [_INSTANCE_COL_SPEC_BY_NAME[c] for c in cols if c in _INSTANCE_COL_SPEC_BY_NAME]
hidden = []
else:
candidate = [
s for s in _INSTANCE_COL_SPECS
if (not s[6] or verbose) and not (s[0] == "volumes" and not show_volumes)
]
droppable = sorted((s for s in candidate if s[5] > 0), key=lambda s: s[5], reverse=True)
active = list(candidate)
for drop_spec in droppable:
if _estimate_table_width(active) <= term_width:
break
active.remove(drop_spec)
hidden = [s[1] for s in candidate if s not in active]
tbl = Table(
style="white",
header_style="bold bright_yellow",
box=box.DOUBLE_EDGE,
row_styles=["on grey11", "none"],
)
for name, header, style, justify, min_width, *_ in active:
kwargs = dict(justify=justify, no_wrap=True, min_width=min_width)
if name in _INSTANCE_COL_MAX_WIDTHS:
kwargs["max_width"] = _INSTANCE_COL_MAX_WIDTHS[name]
if style:
kwargs["style"] = style
tbl.add_column(_RichText(header, justify="center"), **kwargs)
for inst in instances:
tbl.add_row(*[_render_instance_col(name, inst) for name, *_ in active])
return tbl, hidden
_FILTER_VALUE_COLORS = {
"status": _STATUS_COLORS,
"verification": _VERIF_COLORS,
}
def _render_filter_values(values, colors=None, bold=False, line_sep=False):
"""Render a sequence of filter values as a Rich Text, dot-separated."""
t = _RichText()
for i, v in enumerate(values):
if i: t.append("|" if line_sep else " · ", style="bright_white")
style = (colors or {}).get(v, "bright_white")
t.append(v, style=("bold " + style) if bold else style)
return t
def _build_summary_panel(total, label_counts, active_filters=None, order_by=None):
"""Build a Rich Panel summarising the instance query."""
from rich.panel import Panel
from rich.table import Table
from rich.text import Text
lines = []
lines.append(Text.assemble(("Total: ", "bold bright_yellow"), (f"{total} instances", "bold bright_white")))
if label_counts:
parts = []
for lbl, cnt in sorted(label_counts.items(), key=lambda x: -x[1]):
display = lbl if lbl else "(unlabeled)"
parts.append(f"{display}: {cnt}")
lines.append(Text.assemble(("Labels: ", "bold bright_yellow"), (" · ".join(parts), "bright_white")))
if active_filters:
filter_line = Text.assemble(("Filters: ", "bold bright_yellow"))
for i, (k, vals) in enumerate(active_filters.items()):
if i: filter_line.append(" ", style="dim")
filter_line.append(f"{k}=", style="bold bright_white")
filter_line.append_text(_render_filter_values(vals, _FILTER_VALUE_COLORS.get(k), bold=True, line_sep=True))
lines.append(filter_line)
if order_by:
order_line = Text.assemble(("Order by: ", "bold bright_yellow"))
for i, key in enumerate(order_by):
if i: order_line.append(" > ", style="bright_white")
order_line.append(key["col"], style="bold bright_white")
order_line.append(f" ({key['dir']})", style="bright_white")
lines.append(order_line)
grid = Table.grid(padding=(0, 0))
grid.add_column()
for line in lines:
grid.add_row(line)
return Panel(grid, title="[bold bright_yellow]Results Summary[/bold bright_yellow]", style="on #000000", border_style="bright_yellow", expand=False)
def _build_filters_panel(filters):
"""Build a Rich Panel showing the distinct filterable values."""
from rich.panel import Panel
from rich.table import Table
from rich.text import Text
statuses = sorted({f["actual_status"] for f in filters if f.get("actual_status")})
verifs = sorted({f["verification"] for f in filters if f.get("verification")})
gpus = sorted({f["gpu_name"] for f in filters if f.get("gpu_name")})
lines = [
Text.assemble(("--status: ", "bold bright_yellow"), _render_filter_values(statuses, _STATUS_COLORS)),
Text.assemble(("--verification: ", "bold bright_yellow"), _render_filter_values(verifs, _VERIF_COLORS)),
Text.assemble(("--gpu-name: ", "bold bright_yellow"), _render_filter_values(gpus)),
]
grid = Table.grid(padding=(0, 0))
grid.add_column()
for line in lines:
grid.add_row(line)
return Panel(grid, title="[bright_white]Filterable Values[/bright_white]", style="on #000000", border_style="bright_white", expand=False)
def _rich_object_to_string(rich_obj, no_color=True):
"""Render a Rich object to a string."""
from io import StringIO
from rich.console import Console
buffer = StringIO()
console = Console(record=True, file=buffer)
console.print(rich_obj)
return console.export_text(clear=True, styles=not no_color)
@parser.command(
argument("-q", "--quiet", action="store_true", help="only display numeric ids"),
argument("-v", "--verbose", action="store_true", help="show additional columns (SSH, location, template, etc.)"),
argument("-a", "--all", action="store_true", help="force fetching every page before printing, even with --limit/--next-token (fetching everything is already the default otherwise)"),
argument("-s", "--status", metavar="STATUS", nargs="+", help="filter by container status: running loading exited"),
argument("--label", metavar="LABEL", nargs="+", help="filter by instance label; pass empty string '' for unlabeled"),
argument("--gpu-name", metavar="GPU", nargs="+", dest="gpu_name", help="filter by GPU model name"),
argument("--verification", metavar="VERIF", nargs="+", choices=["verified", "unverified", "deverified"], help="filter by verification status"),
argument("-l", "--limit", type=int, default=None, help="max instances per page (1-25); passing this switches to single-page mode instead of fetching everything"),
argument("-t", "--next-token", dest="next_token", help="resume from a pagination token (implies single-page mode)"),
argument("--order-by", dest="order_by", metavar="COL [asc|desc]", action="append", help="sort by column; repeat for multiple keys"),
argument("--cols", metavar="COLS", help=f"override displayed columns (available: {','.join(s[0] for s in _INSTANCE_COL_SPECS)})"),
usage="vastai show instances [OPTIONS]",
help="Display user's current instances",
aliases=hidden_aliases(["show instances-v1"]),
epilog=deindent("""
Shows the instances the user is currently renting. By default this fetches
every page and prints one combined result, which is safe for scripts and
cron jobs. Passing --limit or --next-token switches to single-page mode for
manual pagination (add --all to still fetch every page in that mode).
Examples:
vastai show instances
vastai show instances --raw
vastai show instances -q
vastai show instances -v
vastai show instances --status running loading
vastai show instances --gpu-name 'RTX A5000' 'GTX 1070'
vastai show instances --label training --order-by start_date desc
vastai show instances --cols id,status,gpu,dph
vastai show instances --next-token eyJ2YWx1ZXMi...
"""),
)
def show__instances(args, extra_filters=None):
"""Show the user's current instances."""
if extra_filters and extra_filters.get('internal'):
client = get_client(args)
rows = instances_api.show_instances(client)
field = extra_filters.get('field', 'id')
return [str(r.get(field, '')) for r in rows]
try:
from rich.prompt import Confirm
from rich.text import Text
from rich.padding import Padding
has_rich = True
except ImportError:
has_rich = False
client = get_client(args)
# build select_filters
select_filters = {}
active_display_filters = {}
if args.status:
invalid_statuses = [s for s in args.status if s not in _STATUS_COLORS]
if invalid_statuses:
valid = ", ".join(sorted(_STATUS_COLORS))
print(f"Warning: unknown status value(s): {', '.join(invalid_statuses)}. Valid: {valid}", file=sys.stderr)
select_filters["actual_status"] = {"in": args.status}
active_display_filters["status"] = args.status
if args.label is not None:
vals = [None if l == "" else l for l in args.label]
select_filters["label"] = {"in": vals}
active_display_filters["label"] = [l or "(unlabeled)" for l in vals]
if args.gpu_name:
select_filters["gpu_name"] = {"in": args.gpu_name}
active_display_filters["gpu_name"] = args.gpu_name
if args.verification:
select_filters["verification"] = {"in": args.verification}
active_display_filters["verification"] = args.verification
# order_by
order_by = [{"col": "id", "dir": "asc"}]
if args.order_by:
order_by = []
seen_cols = set()
for entry in args.order_by:
parts = entry.split()
col = parts[0]
dirn = parts[1].lower() if len(parts) > 1 and parts[1].lower() in ("asc", "desc") else "asc"
order_by.append({"col": col, "dir": dirn})
seen_cols.add(col)
if "id" not in seen_cols:
order_by.append({"col": "id", "dir": "asc"})
# Legacy contract: fetch every page and print/return one combined result by
# default (safe for scripts). Only an explicit --limit or --next-token opts
# into single-page mode, matching the old paginated `instances-v1` behavior.
explicit_pagination = args.next_token is not None or args.limit is not None
fetch_all = args.all or not explicit_pagination
limit = max(1, min(args.limit, 25)) if args.limit is not None else 25
if args.raw and fetch_all and not args.quiet:
# Flat list of every matching instance, full unrestricted row shape --
# this is the contract scripts have always depended on for --raw.
return instances_api.show_instances(client, select_filters=select_filters, order_by=order_by)
params = {
"select_filters": json.dumps(select_filters),
"order_by": json.dumps(order_by),
"limit": limit,
}
if not args.raw:
select_cols = _VERBOSE_INSTANCE_SELECT_COLS if args.verbose else _DEFAULT_INSTANCE_SELECT_COLS
if not has_rich:
select_cols = [s[0] for s in instance_fields]
params["select_cols"] = json.dumps(select_cols)
if args.next_token:
params["after_token"] = args.next_token
if args.quiet and fetch_all:
while True:
data = instances_api.show_instances_v1(client, params)
for inst in data.get("instances") or []:
instance_id = inst.get("id")
if instance_id is not None:
print(instance_id)
next_token = data.get("next_token")
if not next_token:
return
params["after_token"] = next_token
# fetch filter breakdown
filter_combos = None
if not args.quiet and not args.raw:
try:
filter_combos = instances_api.show_instance_filters(client)
except Exception:
pass
user_cols = [c.strip() for c in args.cols.split(",")] if args.cols and has_rich else None
page = 0
offset = 0
looping = True
all_instances = []
while looping:
data = instances_api.show_instances_v1(client, params)
instances = data.get("instances", [])
next_token = data.get("next_token")
total = data.get("total_instances", 0)
label_cnts = data.get("label_counts", {})
page += 1
if args.quiet:
for inst in instances:
instance_id = inst.get("id")
if instance_id is not None:
print(instance_id)
break
if args.raw:
return data
if fetch_all:
all_instances.extend(instances)
if next_token:
sys.stderr.write(f"\rLoading more... (page {page + 1})")
sys.stderr.flush()
params["after_token"] = next_token
continue
sys.stderr.write("\r\033[K")
sys.stderr.flush()
instances = all_instances
output_parts = []
if has_rich:
if page == 1 or fetch_all:
if filter_combos:
output_parts.append(_rich_object_to_string(_build_filters_panel(filter_combos), no_color=args.no_color))
output_parts.append(_rich_object_to_string(_build_summary_panel(
total, label_cnts,
active_filters=active_display_filters,
order_by=order_by if args.order_by else None,
), no_color=args.no_color).rstrip("\n"))
else:
output_parts.append('')
if not instances:
empty_msg = "No instances matched your filters." if active_display_filters else "No instances found."
output_parts.append(_rich_object_to_string(Text(empty_msg, style="bright_white"), no_color=args.no_color).rstrip())
else:
tbl, hidden = _build_instances_table(instances, verbose=args.verbose, cols=user_cols)
caption = Text()
if not fetch_all:
if not args.next_token:
caption.append(f"[Page {page}]", style="bright_white")
caption.append(" · ", style="bright_white")
caption.append("Fetched Results: ", style="bright_white")
caption.append(f"{offset + 1} – {offset + len(instances)}", style="bright_white")
caption.append(f" of {total}", style="bright_white")
if hidden and (page == 1 or fetch_all):
caption.append(
f"\nColumns hidden to fit terminal width: {', '.join(hidden)}"
f" · use --cols to customize (see --help)",
style="dim",
)
if caption:
tbl.caption = caption
tbl.caption_justify = "left"
padded = Padding(tbl, (0, 1, 0, 1), style="on #000000", expand=False)
output_parts.append(_rich_object_to_string(padded, no_color=args.no_color))
if next_token:
output_parts.append(f"Next page token: {next_token}\n")
else:
if page == 1 or fetch_all:
if filter_combos:
statuses = sorted({f["actual_status"] for f in filter_combos if f.get("actual_status")})
verifs = sorted({f["verification"] for f in filter_combos if f.get("verification")})
gpus = sorted({f["gpu_name"] for f in filter_combos if f.get("gpu_name")})
print("Filterable Values:")
print(f" --status: {' | '.join(statuses) if statuses else '(none)'}")
print(f" --verification: {' | '.join(verifs) if verifs else '(none)'}")
print(f" --gpu-name: {' | '.join(gpus) if gpus else '(none)'}")
print()
summary_lines = [f"Total: {total} instances"]
if label_cnts:
lbl_parts = [f"{(lbl or '(unlabeled)')}: {cnt}" for lbl, cnt in sorted(label_cnts.items(), key=lambda x: -x[1])]
summary_lines.append(f"Labels: {' · '.join(lbl_parts)}")
if active_display_filters:
filter_parts = [f"{k}={' | '.join(str(v) for v in vals)}" for k, vals in active_display_filters.items()]
summary_lines.append(f"Filters: {' '.join(filter_parts)}")
if args.order_by:
order_parts = [f"{o['col']} ({o['dir']})" for o in order_by]
summary_lines.append(f"Order by: {' > '.join(order_parts)}")
print("Results Summary:")
for line in summary_lines:
print(f" {line}")
print()
if not instances:
print("No instances matched your filters." if active_display_filters else "No instances found.")
else:
display_table(instances, instance_fields)
if not fetch_all:
print(f"[Page {page}] Fetched Results: {offset + 1} – {offset + len(instances)} of {total}")
if next_token:
print(f"Next page token: {next_token}")
if page == 1:
print("\nNOTE: install the 'rich' module for colored output (pip install rich)")
if not fetch_all:
offset += len(instances)
if fetch_all or not next_token:
print_or_page(args, "\n".join(output_parts))
looping = False
elif sys.stdin.isatty() and sys.stdout.isatty():
print("\n".join(output_parts))
try:
if has_rich:
ans = Confirm.ask(f"Fetch next page? (page {page + 1})", default=False)
else:
ans = input(f"Fetch next page? (page {page + 1}) (y/N): ").strip().lower() == "y"
except (EOFError, KeyboardInterrupt):
ans = False
if ans:
params["after_token"] = next_token
else:
looping = False
else:
# Non-interactive invocation (no TTY): never block on a prompt.
print_or_page(args, "\n".join(output_parts))
looping = False
cli/commands/keys.py
"""CLI commands for managing API keys and SSH keys."""
import json
import os
from vastai.cli.parser import argument
from vastai.cli.display import deindent
from vastai.api import keys as keys_api
from vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401
parser = _get_parser()
# ---------------------------------------------------------------------------
# API keys
# ---------------------------------------------------------------------------
@parser.command(
argument("--name", help="name of the api-key", type=str),
argument("--permission_file", help="file path for json encoded permissions, see https://vast.ai/docs/cli/roles-and-permissions for more information", type=str),
argument("--key_params", help="optional wildcard key params for advanced keys", type=str),
usage="vastai create api-key --name NAME --permission_file PERMISSIONS",
help="Create a new api-key with restricted permissions. Can be sent to other users and teammates",
epilog=deindent("""
In order to create api keys you must understand how permissions must be sent via json format.
You can find more information about permissions here: https://vast.ai/docs/cli/roles-and-permissions
""")
)
def create__api_key(args):
"""Create a new api-key with restricted permissions."""
from vastai.cli.util import load_permissions_from_file
try:
client = get_client(args)
permissions = load_permissions_from_file(args.permission_file)
result = keys_api.create_api_key(client, name=args.name, permissions=permissions, key_params=args.key_params)
print("api-key created {}".format(result))
except FileNotFoundError:
print("Error: Permission file '{}' not found.".format(args.permission_file))
@parser.command(
argument("id", help="id of apikey to get", type=int),
usage="vastai show api-key",
help="Show an api-key",
)
def show__api_key(args):
"""Show an api-key."""
client = get_client(args)
result = keys_api.show_api_key(client, id=args.id)
print(result)
@parser.command(
usage="vastai show api-keys",
help="List your api-keys associated with your account",
)
def show__api_keys(args):
"""List api-keys associated with your account."""
client = get_client(args)
result = keys_api.show_api_keys(client)
if args.raw:
return result
else:
print(result)
@parser.command(
argument("id", help="id of apikey to remove", type=int),
usage="vastai delete api-key ID",
help="Remove an api-key",
)
def delete__api_key(args):
"""Remove an api-key."""
client = get_client(args)
result = keys_api.delete_api_key(client, id=args.id)
print(result)
@parser.command(
usage="vastai reset api-key",
help="Reset your api-key (get new key from website)",
)
def reset__api_key(args):
"""Reset your api-key."""
client = get_client(args)
rj = keys_api.reset_api_key(client)
if rj.get("success"):
print("New api-key: {}".format(rj.get("new_api_key", "")))
else:
print(rj.get("msg", "Failed to reset api-key"))
# ---------------------------------------------------------------------------
# SSH keys
# ---------------------------------------------------------------------------
@parser.command(
argument("ssh_key", help="add your existing ssh public key to your account (from the .pub file). If no public key is provided, a new key pair will be generated.", type=str, nargs='?'),
argument("-y", "--yes", help="automatically answer yes to prompts", action="store_true"),
usage="vastai create ssh-key [ssh_public_key] [-y]",
help="Create a new ssh-key",
epilog=deindent("""
You may use this command to add an existing public key, or create a new ssh key pair and add that public key, to your Vast account.
If you provide an ssh_public_key.pub argument, that public key will be added to your Vast account. All ssh public keys should be in OpenSSH format.
Example: $vastai create ssh-key 'ssh_public_key.pub'
If you don't provide an ssh_public_key.pub argument, a new Ed25519 key pair will be generated.
Example: $vastai create ssh-key
The generated keys are saved as ~/.ssh/id_ed25519 (private) and ~/.ssh/id_ed25519.pub (public). Any existing id_ed25519 keys are backed up as .backup_<timestamp>.
The public key will be added to your Vast account.
All ssh public keys are stored in your Vast account and can be used to connect to instances they've been added to.
""")
)
def create__ssh_key(args):
"""Create or add an SSH key to your account."""
from vastai.cli.util import get_ssh_key, generate_ssh_key
ssh_key_content = args.ssh_key
if not ssh_key_content:
ssh_key_content = generate_ssh_key(args.yes)
else:
print("Adding provided SSH public key to account...")
client = get_client(args)
result = keys_api.create_ssh_key(client, ssh_key=ssh_key_content)
print("ssh-key created {}\nNote: You may need to add the new public key to any pre-existing instances".format(result))
@parser.command(
usage="vastai show ssh-keys",
help="List your ssh keys associated with your account",
)
def show__ssh_keys(args):
"""List ssh keys associated with your account."""
client = get_client(args)
result = keys_api.show_ssh_keys(client)
if args.raw:
return result
else:
print(result)
@parser.command(
argument("id", help="id ssh key to delete", type=int),
usage="vastai delete ssh-key ID",
help="Remove an ssh-key",
)
def delete__ssh_key(args):
"""Remove an ssh-key."""
client = get_client(args)
result = keys_api.delete_ssh_key(client, id=args.id)
print(result)
@parser.command(
argument("id", help="id of the ssh key to update", type=int),
argument("ssh_key", help="new public key value", type=str),
usage="vastai update ssh-key ID SSH_KEY",
help="Update an existing SSH key",
)
def update__ssh_key(args):
"""Update an existing SSH key."""
from vastai.cli.util import get_ssh_key
ssh_key = get_ssh_key(args.ssh_key)
client = get_client(args)
result = keys_api.update_ssh_key(client, id=args.id, ssh_key=ssh_key)
print(result)
# ---------------------------------------------------------------------------
# attach / detach ssh
# ---------------------------------------------------------------------------
@parser.command(
argument("instance_id", help="id of instance to attach to", type=int),
argument("ssh_key", help="ssh key to attach to instance", type=str),
usage="vastai attach ssh instance_id ssh_key",
help="Attach an ssh key to an instance. This will allow you to connect to the instance with the ssh key",
epilog=deindent("""
Attach an ssh key to an instance. This will allow you to connect to the instance with the ssh key.
Examples:
vastai attach ssh 12371 AAAAB3NzaC1yc2EAAA...
vastai attach ssh 12371 $(cat ~/.ssh/id_rsa.pub)
vastai attach ssh 12371 ~/.ssh/id_rsa.pub
All examples attaches the ssh key to instance 12371
"""),
)
def attach__ssh(args):
"""Attach an ssh key to an instance."""
from vastai.cli.util import get_ssh_key
ssh_key = get_ssh_key(args.ssh_key)
client = get_client(args)
result = keys_api.attach_ssh(client, instance_id=args.instance_id, ssh_key=ssh_key)
print(result)
@parser.command(
argument("instance_id", help="id of the instance", type=int),
argument("ssh_key_id", help="id of the key to detach from the instance", type=str),
usage="vastai detach instance_id ssh_key_id",
help="Detach an ssh key from an instance",
epilog=deindent("""
Example: vastai detach 99999 12345
""")
)
def detach__ssh(args):
"""Detach an ssh key from an instance."""
client = get_client(args)
result = keys_api.detach_ssh(client, instance_id=args.instance_id, ssh_key_id=args.ssh_key_id)
print(result)
cli/commands/machines.py
"""CLI commands for managing host machines."""
import json
import os
import sys
import time
import warnings
import argparse
from contextlib import redirect_stdout, redirect_stderr
import requests
import urllib3
from vastai.cli.parser import argument
from vastai.cli.display import (
display_table, machine_fields, maintenance_fields,
network_disk_fields, network_disk_machine_fields, deindent,
)
from vastai.api import machines as machines_api
from vastai.api import instances as instances_api
from vastai.api import offers as offers_api
from vastai.api import storage as storage_api
from vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401
from vastai.cli.self_test.machine_diagnostics import (
base_result,
compact_offer_metadata,
failed_checks,
no_offer_failure,
preflight_requirement_checks,
render_preflight_advisories,
render_preflight_failure,
requirement_failure,
)
from vastai.cli.self_test.runtime_diagnostics import (
DAEMON_STARTUP_FAILED,
INSTANCE_CREATE_FAILED,
INSTANCE_CREATE_MISSING_CONTRACT,
INSTANCE_OFFLINE_BEFORE_TEST,
INSTANCE_START_TIMEOUT,
INTERRUPTED,
LegacyProgressParser,
MISSING_PUBLIC_IP,
PROGRESS_CONTAINER_PORT,
PROGRESS_EMPTY_TIMEOUT,
PROGRESS_ENDPOINT_LOST,
PROGRESS_ENDPOINT_UNREACHABLE,
PROGRESS_PORT_NOT_MAPPED,
RUNTIME_TEST_TIMEOUT,
classify_status_msg,
make_progress_endpoint_diagnostic,
make_failure,
redact_secret_text,
)
from vastai.cli.self_test.support_bundle import (
create_support_bundle,
format_bundle_summary,
support_bundles_enabled,
)
parser = _get_parser()
SELF_TEST_INSTANCE_LABEL_PREFIX = "vast-self-test-machine"
INSTANCE_LOG_TAIL_LINES = 1000
# ---------------------------------------------------------------------------
# show machine / machines
# ---------------------------------------------------------------------------
@parser.command(
argument("Machine", help="id of machine to display", type=int),
argument("-q", "--quiet", action="store_true", help="only display numeric ids"),
usage="vastai show machine ID [OPTIONS]",
help="[Host] Show hosted machines",
)
def show__machine(args):
"""Show a machine the host is offering for rent."""
client = get_client(args)
rows = machines_api.show_machine(client, id=args.Machine)
if args.raw:
return rows
else:
if args.quiet:
ids = [f"{row['id']}" for row in rows]
print(" ".join(id for id in ids))
else:
display_table(rows, machine_fields)
@parser.command(
argument("-q", "--quiet", action="store_true", help="only display numeric ids"),
usage="vastai show machines [OPTIONS]",
help="[Host] Show hosted machines",
)
def show__machines(args):
"""Show the machines user is offering for rent."""
client = get_client(args)
rows = machines_api.show_machines(client)
if args.raw:
return {"machines": rows}
else:
if args.quiet:
ids = [f"{row['id']}" for row in rows]
print(" ".join(id for id in ids))
else:
display_table(rows, machine_fields)
# ---------------------------------------------------------------------------
# show maints
# ---------------------------------------------------------------------------
@parser.command(
argument("-i", "--ids", help="comma separated string of machine_ids for which to get maintenance information", type=str),
argument("-q", "--quiet", action="store_true", help="only display numeric ids of the machines in maintenance"),
usage="vastai show maints --ids 'machine_id_1[,machine_id_2,...]' [OPTIONS]",
help="[Host] Show maintenance information for host machines",
)
def show__maints(args):
"""Show the maintenance information for the machines."""
machine_ids = args.ids.split(',')
machine_ids = list(map(int, machine_ids))
client = get_client(args)
rows = machines_api.show_maints(client, machine_ids=machine_ids)
if args.raw:
return rows
else:
if args.quiet:
ids = [f"{row['machine_id']}" for row in rows]
print(" ".join(id for id in ids))
else:
display_table(rows, maintenance_fields)
# ---------------------------------------------------------------------------
# show network-disks
# ---------------------------------------------------------------------------
@parser.command(
usage="vastai show network-disks",
help="[Host] Show network disks associated with your account.",
epilog=deindent("""
Show network disks associated with your account.
""")
)
def show__network_disks(args):
"""Show network disks associated with your account."""
client = get_client(args)
response_data = storage_api.show_network_disks(client)
if args.raw:
return response_data
for cluster_data in response_data['data']:
print(f"Cluster ID: {cluster_data['cluster_id']}")
display_table(cluster_data['network_disks'], network_disk_fields, replace_spaces=False)
machine_rows = []
for machine_id in cluster_data['machine_ids']:
machine_rows.append({
"machine_id": machine_id,
"mount_point": cluster_data['mounts'].get(str(machine_id), "N/A"),
})
print()
display_table(machine_rows, network_disk_machine_fields, replace_spaces=False)
print("\n")
# ---------------------------------------------------------------------------
# list machine / machines
# ---------------------------------------------------------------------------
def list_machine_impl(args, id):
"""Shared logic for listing a single machine."""
from vastai.cli.util import string_to_unix_epoch
client = get_client(args)
end_date = string_to_unix_epoch(args.end_date) if args.end_date else None
json_blob = {
'machine': id,
'price_gpu': args.price_gpu,
'price_disk': args.price_disk,
'price_inetu': args.price_inetu,
'price_inetd': args.price_inetd,
'price_min_bid': args.price_min_bid,
'min_chunk': args.min_chunk,
'end_date': end_date,
'credit_discount_max': args.discount_rate,
'duration': args.duration,
'vol_size': args.vol_size,
'vol_price': args.vol_price,
}
if args.explain:
print("request json: ")
print(json_blob)
rj = machines_api.list_machine(
client, id=id, price_gpu=args.price_gpu, price_disk=args.price_disk,
price_inetu=args.price_inetu, price_inetd=args.price_inetd,
price_min_bid=args.price_min_bid, min_chunk=args.min_chunk,
end_date=end_date, discount_rate=args.discount_rate,
duration=args.duration, vol_size=args.vol_size, vol_price=args.vol_price,
)
if rj.get("success"):
price_gpu_ = str(args.price_gpu) if args.price_gpu is not None else "def"
price_inetu_ = str(args.price_inetu)
price_inetd_ = str(args.price_inetd)
min_chunk_ = str(args.min_chunk)
discount_rate_ = str(args.discount_rate)
duration_ = str(args.duration)
if args.raw:
return rj
else:
print(f"offers created/updated for machine {id}, @ ${price_gpu_}/gpu/hr, ${price_inetu_}/GB up, ${price_inetd_}/GB down, {min_chunk_}/min gpus, max discount_rate {discount_rate_}, duration {duration_}")
num_extended = rj.get("extended", 0)
if num_extended > 0:
print(f"extended {num_extended} client contracts to {args.end_date}")
else:
if args.raw:
return rj
else:
print(rj.get("msg", rj))
@parser.command(
argument("id", help="id of machine to list", type=int),
argument("-g", "--price_gpu", help="per gpu rental price in $/hour (price for active instances)", type=float),
argument("-s", "--price_disk",
help="storage price in $/GB/month (price for inactive instances), default: $0.10/GB/month", type=float),
argument("-u", "--price_inetu", help="price for internet upload bandwidth in $/GB", type=float),
argument("-d", "--price_inetd", help="price for internet download bandwidth in $/GB", type=float),
argument("-b", "--price_min_bid", help="per gpu minimum bid price floor in $/hour", type=float),
argument("-r", "--discount_rate", help="Max long term prepay discount rate fraction, default: 0.4 ", type=float),
argument("-m", "--min_chunk", help="minimum amount of gpus (default: 1)", type=int, default=1),
argument("-e", "--end_date", help="contract offer expiration - the available until date (optional, in unix float timestamp or MM/DD/YYYY format)", type=str),
argument("-l", "--duration", help="Updates end_date daily to be duration from current date. Cannot be combined with end_date. Format is: `n days`, `n weeks`, `n months`, `n years`, or total intended duration in seconds."),
argument("-v", "--vol_size", help="Size for volume contract offer. Defaults to half of available disk. Set 0 to not create a volume contract offer.", type=int),
argument("-z", "--vol_price", help="Price for disk on volume contract offer. Defaults to price_disk. Invalid if vol_size is 0.", type=float),
usage="vastai list machine ID [options]",
help="[Host] list a machine for rent",
epilog=deindent("""
Performs the same action as pressing the "LIST" button on the site https://cloud.vast.ai/host/machines.
On the end date the listing will expire and your machine will unlist. However any existing client jobs will still remain until ended by their owners.
Once you list your machine and it is rented, it is extremely important that you don't interfere with the machine in any way.
If your machine has an active client job and then goes offline, crashes, or has performance problems, this could permanently lower your reliability rating.
We strongly recommend you test the machine first and only list when ready.
Raising any resource price above the current contract price writes a
pending row to `pending_price_increases` for each affected client.
Clients review and accept those rows via the console or
`vastai show pending-price-increases`; auto-extend stops at the old
price until they accept.
""")
)
def list__machine(args):
"""List a machine for rent."""
return list_machine_impl(args, args.id)
@parser.command(
argument("ids", help="ids of machines to list", type=int, nargs='+'),
argument("-g", "--price_gpu", help="per gpu on-demand rental price in $/hour (base price for active instances)", type=float),
argument("-s", "--price_disk",
help="storage price in $/GB/month (price for inactive instances), default: $0.10/GB/month", type=float),
argument("-u", "--price_inetu", help="price for internet upload bandwidth in $/GB", type=float),
argument("-d", "--price_inetd", help="price for internet download bandwidth in $/GB", type=float),
argument("-b", "--price_min_bid", help="per gpu minimum bid price floor in $/hour", type=float),
argument("-r", "--discount_rate", help="Max long term prepay discount rate fraction, default: 0.4 ", type=float),
argument("-m", "--min_chunk", help="minimum amount of gpus (default: 1)", type=int, default=1),
argument("-e", "--end_date", help="contract offer expiration - the available until date (optional, in unix float timestamp or MM/DD/YYYY format)", type=str),
argument("-l", "--duration", help="Updates end_date daily to be duration from current date. Cannot be combined with end_date. Format is: `n days`, `n weeks`, `n months`, `n years`, or total intended duration in seconds."),
argument("-v", "--vol_size", help="Size for volume contract offer. Defaults to half of available disk. Set 0 to not create a volume contract offer.", type=int),
argument("-z", "--vol_price", help="Price for disk on volume contract offer. Defaults to price_disk. Invalid if vol_size is 0.", type=float),
usage="vastai list machines IDs [options]",
help="[Host] list machines for rent",
epilog=deindent("""
This variant can be used to list or update the listings for multiple machines at once with the same args.
You could extend the end dates of all your machines using a command combo like this:
./vast.py list machines $(./vast.py show machines -q) -e 12/31/2024 --retry 6
""")
)
def list__machines(args):
"""List multiple machines for rent."""
return [list_machine_impl(args, id) for id in args.ids]
# ---------------------------------------------------------------------------
# unlist machine
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of machine to unlist", type=int),
usage="vastai unlist machine <id>",
help="[Host] Unlist a listed machine",
)
def unlist__machine(args):
"""Remove machine from list of machines for rent."""
client = get_client(args)
rj = machines_api.unlist_machine(client, id=args.id)
if rj.get("success"):
print("all offers for machine {machine_id} removed, machine delisted.".format(machine_id=args.id))
else:
print(rj.get("msg", rj))
# ---------------------------------------------------------------------------
# delete machine
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of machine to delete", type=int),
usage="vastai delete machine <id>",
help="[Host] Delete machine if the machine is not being used by clients. host jobs on their own machines are disregarded and machine is force deleted.",
)
def delete__machine(args):
"""Delete machine if the machine is not being used by clients."""
client = get_client(args)
rj = machines_api.delete_machine(client, id=args.id)
if rj.get("success"):
print("deleted machine_id ({machine_id}) and all related contracts.".format(machine_id=args.id))
else:
print(rj.get("msg", rj))
# ---------------------------------------------------------------------------
# cleanup machine
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of machine to cleanup", type=int),
usage="vastai cleanup machine ID [options]",
help="[Host] Remove all expired storage instances from the machine, freeing up space",
epilog=deindent("""
Instances expire on their end date. Expired instances still pay storage fees, but can not start.
Since hosts are still paid storage fees for expired instances, we do not auto delete them.
Instead you can use this CLI/API function to delete all expired storage instances for a machine.
This is useful if you are running low on storage, want to do maintenance, or are subsidizing storage, etc.
""")
)
def cleanup__machine(args):
"""Remove expired storage instances from a machine."""
client = get_client(args)
rj = machines_api.cleanup_machine(client, id=args.id)
if args.raw:
return rj
if rj.get("success"):
print(json.dumps(rj, indent=1))
else:
print(rj.get("msg", rj))
# ---------------------------------------------------------------------------
# defrag machines
# ---------------------------------------------------------------------------
@parser.command(
argument("IDs", help="ids of machines", type=int, nargs='+'),
usage="vastai defragment machines IDs ",
help="[Host] Defragment machines",
epilog=deindent("""
Defragment some of your machines. This will rearrange GPU assignments to try and make more multi-gpu offers available.
"""),
)
def defrag__machines(args):
"""Defragment machines to make more multi-gpu offers available."""
if args.explain:
print("request json: ")
print({"machine_ids": args.IDs})
client = get_client(args)
try:
result = machines_api.defrag_machines(client, machine_ids=args.IDs)
print(f"defragment result: {result}")
except Exception as e:
print(f"Error: {e}")
# ---------------------------------------------------------------------------
# set min_bid
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of machine to set min bid price for", type=int),
argument("--price", help="per gpu min bid price in $/hour", type=float),
usage="vastai set min_bid id [--price PRICE]",
help="[Host] Set the minimum bid/rental price for a machine",
epilog=deindent("""
Change the current min bid price of machine id to PRICE.
"""),
)
def set__min_bid(args):
"""Set the minimum bid/rental price for a machine."""
if args.explain:
print("request json: ")
print({"client_id": "me", "price": args.price})
client = get_client(args)
machines_api.set_min_bid(client, id=args.id, price=args.price)
print("Per gpu min bid price changed")
# ---------------------------------------------------------------------------
# set defjob
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of machine to launch default instance on", type=int),
argument("--price_gpu", help="per gpu rental price in $/hour", type=float),
argument("--price_inetu", help="price for internet upload bandwidth in $/GB", type=float),
argument("--price_inetd", help="price for internet download bandwidth in $/GB", type=float),
argument("--image", help="docker container image to launch", type=str),
argument("--args", nargs=argparse.REMAINDER, help="list of arguments passed to container launch"),
usage="vastai set defjob id [--api-key API_KEY] [--price_gpu PRICE_GPU] [--price_inetu PRICE_INETU] [--price_inetd PRICE_INETD] [--image IMAGE] [--args ...]",
help="[Host] Create default jobs for a machine",
epilog=deindent("""
Performs the same action as creating a background job at https://cloud.vast.ai/host/create.
""")
)
def set__defjob(args):
"""Create default jobs for a machine."""
if args.explain:
print("request json: ")
print({'machine': args.id, 'price_gpu': args.price_gpu, 'price_inetu': args.price_inetu,
'price_inetd': args.price_inetd, 'image': args.image, 'args': args.args})
client = get_client(args)
rj = machines_api.set_defjob(client, id=args.id, price_gpu=args.price_gpu,
price_inetu=args.price_inetu, price_inetd=args.price_inetd,
image=args.image, args=args.args)
if rj.get("success"):
print("bids created for machine {args.id}, @ ${args.price_gpu}/gpu/day, ${args.price_inetu}/GB up, ${args.price_inetd}/GB down".format(**locals()))
else:
print(rj.get("msg", rj))
# ---------------------------------------------------------------------------
# remove defjob
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of machine to remove default instance from", type=int),
usage="vastai remove defjob id",
help="[Host] Delete default jobs",
)
def remove__defjob(args):
"""Delete default jobs for a machine."""
client = get_client(args)
rj = machines_api.remove_defjob(client, id=args.id)
if rj.get("success"):
print("default instance for machine {machine_id} removed.".format(machine_id=args.id))
else:
print(rj.get("msg", rj))
# ---------------------------------------------------------------------------
# schedule / cancel maint
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of machine to schedule maintenance for", type=int),
argument("--sdate", help="maintenance start date in unix epoch time (UTC seconds)", type=float),
argument("--duration", help="maintenance duration in hours", type=float),
argument("--maintenance_category", help="(optional) can be one of [power, internet, disk, gpu, software, other]", type=str, default="not provided"),
usage="vastai schedule maintenance id [--sdate START_DATE --duration DURATION --maintenance_category MAINTENANCE_CATEGORY]",
help="[Host] Schedule upcoming maint window",
epilog=deindent("""
The proper way to perform maintenance on your machine is to wait until all active contracts have expired or the machine is vacant.
For unplanned or unscheduled maintenance, use this schedule maint command. That will notify the client that you have to take the machine down and that they should save their work.
You can specify a date, duration, reason and category for the maintenance.
Example: vastai schedule maint 8207 --sdate 1677562671 --duration 0.5 --maintenance_category "power"
"""),
)
def schedule__maint(args):
"""Schedule upcoming maintenance window."""
from datetime import datetime, timezone
from vastai.cli.util import string_to_unix_epoch
dt = datetime.fromtimestamp(args.sdate, tz=timezone.utc)
print(f"Scheduling maintenance window starting {dt} lasting {args.duration} hours")
print(f"This will notify all clients of this machine.")
ok = input("Continue? [y/n] ")
if ok.strip().lower() != "y":
return
if args.explain:
print("request json: ")
print({"client_id": "me", "sdate": args.sdate, "duration": args.duration,
"maintenance_category": args.maintenance_category})
client = get_client(args)
machines_api.schedule_maint(client, id=args.id, sdate=string_to_unix_epoch(args.sdate),
duration=args.duration, maintenance_category=args.maintenance_category)
print(f"Maintenance window scheduled for {dt} success")
@parser.command(
argument("id", help="id of machine to cancel maintenance(s) for", type=int),
usage="vastai cancel maint id",
help="[Host] Cancel maint window",
epilog=deindent("""
For deleting a machine's scheduled maintenance window(s), use this cancel maint command.
Example: vastai cancel maint 8207
"""),
)
def cancel__maint(args):
"""Cancel scheduled maintenance window(s)."""
print(f"Cancelling scheduled maintenance window(s) for machine {args.id}.")
ok = input("Continue? [y/n] ")
if ok.strip().lower() != "y":
return
if args.explain:
print("request json: ")
print({"client_id": "me", "machine_id": args.id})
client = get_client(args)
machines_api.cancel_maint(client, id=args.id)
print(f"Cancel maintenance window(s) scheduled for machine {args.id} success")
# ---------------------------------------------------------------------------
# add network-disk
# ---------------------------------------------------------------------------
@parser.command(
argument("machines", help="ids of machines to add disk to, that is networked to be on the same LAN as machine", type=int, nargs='+'),
argument("mount_point", help="mount path of disk to add", type=str),
argument("-d", "--disk_id", help="id of network disk to attach to machines in the cluster", type=int, nargs='?'),
usage="vastai add network-disk MACHINES MOUNT_PATH [options]",
help="[Host] Add Network Disk to Physical Cluster.",
epilog=deindent("""
This variant can be used to add a network disk to a physical cluster.
When you add a network disk for the first time, you just need to specify the machine(s) and mount_path.
When you add a network disk for the second time, you need to specify the disk_id.
Example:
vastai add network-disk 1 /mnt/disk1
vastai add network-disk 1 /mnt/disk1 -d 12345
""")
)
def add__network_disk(args):
"""Add network disk to a physical cluster."""
if args.explain:
print("request json: ")
print({"machines": [int(id) for id in args.machines], "mount_point": args.mount_point, "disk_id": args.disk_id})
client = get_client(args)
result = storage_api.add_network_disk(client, machines=args.machines, mount_point=args.mount_point,
disk_id=args.disk_id)
if args.raw:
return result
print("Attached network disk to machines. Disk id: " + str(result["disk_id"]))
# ---------------------------------------------------------------------------
# dump-logs
# ---------------------------------------------------------------------------
def _diagnostic_error(error):
return redact_secret_text(error) or ""
def _artifact_text(value):
if isinstance(value, str):
return value
return json.dumps(value, indent=2, sort_keys=True, default=str)
def collect_instance_log_artifacts(client, instance_id):
files = {}
errors = []
try:
instance_info = instances_api.show_instance(client, id=instance_id)
files["instance/show-instance.json"] = _artifact_text(instance_info)
except Exception as e:
errors.append({
"artifact": "instance/show-instance.json",
"error": _diagnostic_error(e),
})
for archive_name, daemon_logs in (
("instance/container.log", False),
("instance/daemon.log", True),
):
try:
logs = instances_api.logs(
client,
instance_id=instance_id,
tail=INSTANCE_LOG_TAIL_LINES,
daemon_logs=daemon_logs,
)
files[archive_name] = _artifact_text(logs)
except Exception as e:
errors.append({
"artifact": archive_name,
"error": _diagnostic_error(e),
})
return files, errors
@parser.command(
argument("machine_id", help="Machine ID", type=str),
argument("--instance-id", help="Instance ID to pull Vast instance logs from", type=int),
argument("--output-dir", help="Directory for the diagnostic bundle (default: /tmp)", type=str),
argument(
"--include-local-host-artifacts",
action="store_true",
help="Include local OS/kaalia artifacts; only use when running on the actual Vast host",
),
usage=(
"vastai dump-logs <machine_id> [--instance-id INSTANCE_ID] "
"[--output-dir DIR] [--include-local-host-artifacts]"
),
help="[Host] Bundle self-test diagnostics for support",
epilog=deindent("""
Creates a redacted diagnostic tarball containing CLI-visible self-test
evidence. If --instance-id is provided, the command also pulls container
and daemon logs from the Vast instance logs API.
Local OS/kaalia artifacts are only collected with
--include-local-host-artifacts. Use that option only when running this
command on the actual host machine; from a laptop, it would collect the
laptop's logs instead.
Example:
vastai dump-logs 12345 --instance-id 67890
"""),
)
def dump_logs(args):
run_started_at = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())
cli_output = [f"Manual diagnostic bundle requested for machine {args.machine_id}."]
extra_files = {}
extra_errors = []
if args.instance_id:
client = get_client(args)
cli_output.append(f"Requested Vast instance logs for instance {args.instance_id}.")
instance_files, instance_errors = collect_instance_log_artifacts(client, args.instance_id)
extra_files.update(instance_files)
extra_errors.extend(instance_errors)
else:
cli_output.append("No --instance-id provided; Vast instance logs were not requested.")
if not args.include_local_host_artifacts:
cli_output.append(
"Local host OS/kaalia artifacts were not collected. Use "
"--include-local-host-artifacts only when running on the actual host."
)
try:
bundle = create_support_bundle(
machine_id=args.machine_id,
output_dir=args.output_dir,
result={
"machine_id": args.machine_id,
"stage": "manual_dump_logs",
"reason": "Manual diagnostic bundle requested with vastai dump-logs.",
"instance_id": args.instance_id,
"includes_local_host_artifacts": args.include_local_host_artifacts,
},
cli_output=cli_output,
extra_files=extra_files,
extra_errors=extra_errors,
run_started_at=run_started_at,
command=sys.argv,
secrets=[getattr(args, "api_key", None), os.environ.get("VAST_API_KEY")],
include_local_host_artifacts=args.include_local_host_artifacts,
)
except Exception as e:
error = _diagnostic_error(e)
if getattr(args, "raw", False):
return {
"success": False,
"machine_id": args.machine_id,
"error": f"Failed to create diagnostic bundle: {error}",
}
print(f"WARNING: failed to create diagnostic bundle: {error}")
sys.exit(1)
if getattr(args, "raw", False):
return bundle
for line in format_bundle_summary(bundle):
print(line)
# ---------------------------------------------------------------------------
# self-test machine
# ---------------------------------------------------------------------------
@parser.command(
argument("machine_id", help="Machine ID", type=str),
argument("--debugging", action="store_true", help="Enable debugging output"),
argument("--ignore-requirements", action="store_true", help="Ignore the minimum system requirements and run the self test regardless"),
argument("--test-image", help="Use a custom self-test image for testing custom self-test images. Overrides VAST_SELF_TEST_IMAGE and CUDA mapping.", type=str),
argument("--support-bundle-dir", help="Directory for failure diagnostic bundles (default: /tmp)", type=str),
argument("--no-support-bundle", action="store_true", help="Do not create a diagnostic tarball when the self-test fails"),
usage="vastai self-test machine <machine_id> [--debugging] [--ignore-requirements] [--test-image IMAGE]",
help="[Host] Perform a self-test on the specified machine",
epilog=deindent("""
This command tests if a machine meets specific requirements and
runs a series of tests to ensure it's functioning correctly.
Examples:
vastai self-test machine 12345
vastai self-test machine 12345 --debugging
"""),
)
def self_test__machine(args):
"""
Performs a self-test on the specified machine to verify its compliance with
required specifications and functionality.
"""
instance_id = None
result = base_result(args.machine_id)
run_started_at = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())
cli_output = []
instance_log_files = {}
instance_log_errors = []
instance_log_collected = set()
ignore_requirements_warning = (
"WARNING: --ignore-requirements is set. Requirement checks are skipped as a "
"pass/fail gate, and passing this self-test does not qualify this machine for verification."
)
if not hasattr(args, 'debugging'):
args.debugging = False
if not hasattr(args, 'test_image'):
args.test_image = None
if not hasattr(args, 'support_bundle_dir'):
args.support_bundle_dir = None
if not hasattr(args, 'no_support_bundle'):
args.no_support_bundle = False
if getattr(args, "ignore_requirements", False):
result["warning"] = ignore_requirements_warning
result["diagnostics"]["requirements_ignored"] = True
def output_line(*args_to_print):
return " ".join(str(item) for item in args_to_print)
def progress_print(*args_to_print):
cli_output.append(output_line(*args_to_print))
if not args.raw:
print(*args_to_print)
def debug_print(*args_to_print):
if args.debugging:
cli_output.append(f"DEBUG: {output_line(*args_to_print)}")
if args.debugging and not args.raw:
print(*args_to_print)
def collect_instance_logs(inst_id):
if args.no_support_bundle or not support_bundles_enabled():
return
if not inst_id or inst_id in instance_log_collected:
return
instance_log_collected.add(inst_id)
files, errors = collect_instance_log_artifacts(client, inst_id)
instance_log_files.update(files)
instance_log_errors.extend(errors)
result["diagnostics"]["instance_log_collection"] = {
"instance_id": inst_id,
"files": sorted(instance_log_files.keys()),
"collection_errors": instance_log_errors,
}
def ensure_support_bundle():
if result.get("success"):
return None
if result.get("diagnostics", {}).get("support_bundle"):
return result["diagnostics"]["support_bundle"]
if args.no_support_bundle or not support_bundles_enabled():
return None
if instance_id:
collect_instance_logs(instance_id)
try:
bundle = create_support_bundle(
machine_id=args.machine_id,
output_dir=args.support_bundle_dir,
result=result,
cli_output=cli_output,
extra_files=instance_log_files,
extra_errors=instance_log_errors,
run_started_at=run_started_at,
command=sys.argv,
secrets=[getattr(args, "api_key", None), os.environ.get("VAST_API_KEY")],
include_local_host_artifacts=False,
)
except Exception as e:
error = redact_secret_text(e) or ""
result["diagnostics"]["support_bundle_error"] = error
progress_print(f"WARNING: failed to create self-test diagnostic bundle: {error}")
return None
result["diagnostics"]["support_bundle"] = bundle
return bundle
def finish_failure():
if args.raw:
ensure_support_bundle()
return result
if result.get("warning"):
print(result["warning"])
render_runtime_failure()
bundle = ensure_support_bundle()
if bundle:
for line in format_bundle_summary(bundle):
print(line)
print(f"Test failed: {result['reason']}")
sys.exit(1)
def set_runtime_failure(diagnostic, fallback_reason=None):
result["failure"] = diagnostic
result["failure_code"] = diagnostic["code"]
result["stage"] = diagnostic.get("stage") or result.get("stage")
result["reason"] = fallback_reason or diagnostic.get("summary") or ""
result["diagnostics"]["runtime_failure"] = diagnostic
def safe_error(error):
return redact_secret_text(error) or ""
def render_runtime_failure():
diagnostic = result.get("diagnostics", {}).get("runtime_failure")
if not diagnostic:
return
progress_print("Runtime failure diagnostics:")
progress_print(f"- code: {diagnostic.get('code')}")
if diagnostic.get("summary"):
progress_print(f"- summary: {diagnostic['summary']}")
if diagnostic.get("underlying_error"):
progress_print(f"- underlying error: {diagnostic['underlying_error']}")
endpoint = diagnostic.get("progress_endpoint") or result.get("diagnostics", {}).get("progress_endpoint")
if endpoint:
if endpoint.get("url"):
progress_print(f"- tried: {endpoint['url']}")
external_port = endpoint.get("external_port") or endpoint.get("host_port")
if external_port:
progress_print(f"- external port tested: {external_port}")
last_bits = []
if endpoint.get("last_status_code") is not None:
last_bits.append(f"HTTP {endpoint['last_status_code']}")
if endpoint.get("last_error_type"):
last_bits.append(str(endpoint["last_error_type"]))
if endpoint.get("last_error"):
last_bits.append(str(endpoint["last_error"]))
if last_bits:
progress_print(f"- last result: {' - '.join(last_bits)}")
if endpoint.get("mapped_ports"):
progress_print(f"- mapped container ports: {', '.join(endpoint['mapped_ports'])}")
if diagnostic.get("remediation"):
progress_print(f"- remediation: {diagnostic['remediation']}")
steps = diagnostic.get("suggested_steps") or []
if steps:
progress_print("- suggested steps:")
for step in steps:
progress_print(f" - {step}")
client = get_client(args)
try:
def http_status_code(error):
response = getattr(error, "response", None)
return getattr(response, "status_code", None)
def offer_search_error_summary(error):
return {
"status_code": http_status_code(error),
"error": safe_error(error),
}
def lookup_machine_for_offer_failure(machine_id):
try:
rows = machines_api.show_machine(client, id=machine_id)
if isinstance(rows, dict):
rows_for_status = [rows] if rows else []
else:
rows_for_status = rows or []
return {
"status": "visible" if rows_for_status else "empty",
"rows": rows_for_status,
"row_count": len(rows_for_status),
}
except requests.exceptions.HTTPError as e:
status_code = http_status_code(e)
status = "permission_denied" if status_code in (401, 403) else "not_found" if status_code == 404 else "error"
return {
"status": status,
"status_code": status_code,
"error": safe_error(e),
"rows": [],
"row_count": 0,
}
except requests.exceptions.RequestException as e:
return {
"status": "lookup_error",
"status_code": http_status_code(e),
"error": safe_error(e),
"rows": [],
"row_count": 0,
}
def compact_machine_lookup(machine_lookup):
if not machine_lookup:
return None
return {
"status": machine_lookup.get("status"),
"status_code": machine_lookup.get("status_code"),
"row_count": machine_lookup.get("row_count", len(machine_lookup.get("rows") or [])),
"error": machine_lookup.get("error"),
}
def selected_offer_for_self_test(machine_id):
strict_query = {
"machine_id": {"eq": machine_id},
"verified": {"eq": "any"},
"rentable": {"eq": True},
"rented": {"eq": "any"},
}
diagnostics = {
"strict_offer_count": None,
"broader_offer_count": None,
"broader_offers": [],
"search_error": None,
"machine_lookup": None,
}
try:
strict_offers = offers_api.search_offers(
client, query=strict_query, offer_type="on-demand",
order=[["score", "desc"]], storage=5.0, no_default=True,
)
except requests.exceptions.HTTPError as e:
if http_status_code(e) in (401, 403):
diagnostics["search_error"] = offer_search_error_summary(e)
check, failure = no_offer_failure(machine_id, [], search_error=e)
return None, (check, failure), diagnostics
raise
debug_print("Captured strict offers from search_offers:", strict_offers)
diagnostics["strict_offer_count"] = len(strict_offers or [])
if strict_offers:
sorted_offers = sorted(strict_offers, key=lambda x: x.get("dlperf", 0), reverse=True)
selected = dict(sorted_offers[0])
selected["machine_id"] = selected.get("machine_id") or machine_id
debug_print("Selected offer found:", selected)
return selected, None, diagnostics
broader_query = {
"machine_id": {"eq": machine_id},
"verified": {"eq": "any"},
"rentable": {"eq": "any"},
"rented": {"eq": "any"},
}
try:
broader_offers = offers_api.search_offers(
client, query=broader_query, offer_type="on-demand",
order=[["score", "desc"]], storage=5.0, no_default=True,
)
except requests.exceptions.HTTPError as e:
if http_status_code(e) in (401, 403):
diagnostics["search_error"] = offer_search_error_summary(e)
check, failure = no_offer_failure(machine_id, [], search_error=e)
return None, (check, failure), diagnostics
raise
diagnostics["broader_offer_count"] = len(broader_offers or [])
diagnostics["broader_offers"] = [
compact_offer_metadata(dict(offer, machine_id=offer.get("machine_id", machine_id)))
for offer in (broader_offers or [])[:5]
]
machine_lookup = None
if not broader_offers:
machine_lookup = lookup_machine_for_offer_failure(machine_id)
diagnostics["machine_lookup"] = compact_machine_lookup(machine_lookup)
check, failure = no_offer_failure(machine_id, broader_offers, machine_lookup=machine_lookup)
return None, (check, failure), diagnostics
selected_offer, offer_failure, offer_diagnostics = selected_offer_for_self_test(args.machine_id)
result["diagnostics"]["offer_search"] = offer_diagnostics
if offer_failure:
check, failure = offer_failure
result["checks"] = [check]
result["failure"] = failure
result["failure_code"] = failure["code"]
result["stage"] = "select_offer"
result["reason"] = failure["summary"]
render_preflight_failure(args.machine_id, result["checks"], failure, progress_print)
return finish_failure()
result["offer"] = compact_offer_metadata(selected_offer)
checks = preflight_requirement_checks(selected_offer)
result["checks"] = checks
unmet_checks = failed_checks(checks)
if unmet_checks:
failure = requirement_failure(checks)
result["diagnostics"]["preflight_failure"] = failure
render_preflight_failure(args.machine_id, checks, failure, progress_print)
render_preflight_advisories(args.machine_id, checks, progress_print)
if not args.ignore_requirements:
result["failure"] = failure
result["failure_code"] = failure["code"]
result["stage"] = "preflight_requirements"
result["reason"] = failure["summary"]
return finish_failure()
progress_print("Continuing despite unmet requirements because --ignore-requirements is set.")
else:
progress_print(f"Machine ID {args.machine_id} meets all the requirements.")
render_preflight_advisories(args.machine_id, checks, progress_print)
if args.ignore_requirements:
progress_print(ignore_requirements_warning)
# ----- CUDA version to docker image mapping -----
def cuda_map_to_image(cuda_version, compute_cap=None):
"""Return (image, reason). Reason explains why this image was picked."""
docker_repo = "vastai/test"
image_tag_prefix = "self-test-v2-cuda"
def image_for(version):
return f"{docker_repo}:{image_tag_prefix}-{version}"
if isinstance(cuda_version, float):
cuda_version = str(cuda_version)
original_cuda = cuda_version
# cuda-12.8 (torch 2.10) still ships sm_70 (Volta); cuda-13.0
# (torch 2.11) never did. Neither builds sm_50/sm_60 kernels.
# Anything pre-Volta (compute_cap < 700) must use the cuda-11.8
# legacy image.
if compute_cap is not None and compute_cap < 700:
return (
image_for("11.8"),
f"compute_cap={compute_cap} below sm_70 -> forced {image_tag_prefix}-11.8",
)
# Volta sm_70/sm_72 hosts: cuda-12.8 wheels include sm_70 but
# cuda-13.0 wheels never did. Cap the driver-reported CUDA version
# at 12.8 so the map below resolves to cuda-12.8 even if the
# operator has installed a CUDA 13 driver on a V100.
clamped_for_volta = False
if compute_cap is not None and compute_cap < 750:
if float(cuda_version) > 12.8:
cuda_version = "12.8"
clamped_for_volta = True
docker_tag_map = {
"11.8": image_for("11.8"),
"12.8": image_for("12.8"),
"13.0": image_for("13.0"),
"13.3": image_for("13.3"),
}
cap_hint = f"compute_cap={compute_cap}" if compute_cap is not None else "compute_cap=unknown"
cuda_float = float(cuda_version)
compatible_versions = sorted(float(version) for version in docker_tag_map)
selected_version = max(
(version for version in compatible_versions if version <= cuda_float),
default=None,
)
if selected_version is not None:
selected_version_str = f"{selected_version:.1f}"
image = docker_tag_map[selected_version_str]
if clamped_for_volta:
reason = f"{cap_hint} (Volta) + cuda_max_good={original_cuda} -> clamped to {cuda_version} -> {image}"
elif selected_version_str == cuda_version:
reason = f"{cap_hint}, cuda_max_good={cuda_version} -> exact match -> {image}"
else:
reason = (
f"{cap_hint}, cuda_max_good={original_cuda} -> "
f"selected newest image <= host CUDA ({selected_version_str}) -> {image}"
)
return image, reason
raise KeyError(f"No CUDA version found for {cuda_version} or any lower version")
top_offer = selected_offer
if not top_offer:
progress_print(f"No valid offers found for Machine ID {args.machine_id}")
result["reason"] = "No valid offers found."
else:
ask_contract_id = top_offer["id"]
cuda_version = top_offer["cuda_max_good"]
compute_cap = top_offer.get("compute_cap")
image_override = args.test_image or os.environ.get("VAST_SELF_TEST_IMAGE")
if image_override:
docker_image = image_override
image_reason = "custom self-test image override"
else:
docker_image, image_reason = cuda_map_to_image(cuda_version, compute_cap)
result["diagnostics"]["image"] = {
"image": docker_image,
"reason": image_reason,
"override": bool(image_override),
}
# ----- create the test instance -----
try:
result["phase"] = "rental"
result["stage"] = "create_instance"
from vastai.cli.util import parse_env
env = parse_env("-e TZ=PDT -e XNAME=XX4 -p 5000:5000 -p 1234:1234")
runtype = "ssh_direc ssh_proxy"
self_test_label = f"{SELF_TEST_INSTANCE_LABEL_PREFIX}-{args.machine_id}"
result["diagnostics"]["launch"] = {
"runtype": runtype,
"jupyter_lab": False,
"ports": ["5000/tcp", "1234/tcp"],
"label": self_test_label,
}
progress_print(f"Starting test with {docker_image} ({image_reason})")
rj = instances_api.create_instance(
client,
id=ask_contract_id,
image=docker_image,
disk=40,
env=env,
price=None,
label=self_test_label,
extra=None,
onstart_cmd="/verification/remote.sh",
login=None,
python_utf8=False,
lang_utf8=False,
jupyter_lab=False,
jupyter_dir=None,
force=False,
cancel_unavail=False,
template_hash=None,
user=None,
runtype=runtype,
args=None,
)
debug_print("Captured instance_info from create_instance:", rj)
except Exception as e:
error = safe_error(e)
progress_print(f"Error creating instance: {error}")
set_runtime_failure(
make_failure(
INSTANCE_CREATE_FAILED,
stage="create_instance",
summary="Failed to create instance. Check the docker configuration.",
error=error,
underlying_error=error,
)
)
result["error"] = error
return finish_failure()
instance_id = rj.get("new_contract")
if not instance_id:
progress_print("Instance creation response did not contain 'new_contract'.")
set_runtime_failure(
make_failure(
INSTANCE_CREATE_MISSING_CONTRACT,
stage="create_instance",
details=f"Create-instance response: {rj}",
),
"Instance creation failed.",
)
else:
# ----- helper: check if instance exists -----
def instance_exist(inst_id):
try:
info = instances_api.show_instance(client, id=inst_id)
if not info:
return False
status = info.get('intended_status') or info.get('actual_status')
if status in ['destroyed', 'terminated', 'offline']:
return False
return True
except requests.exceptions.HTTPError as e:
if e.response.status_code == 404:
return False
debug_print(f"HTTPError when checking instance existence: {safe_error(e)}")
return False
except Exception as e:
debug_print(f"No instance found or Unexpected error checking instance existence: {safe_error(e)}")
return False
# ----- helper: destroy instance silently with retries -----
def destroy_instance_silent(inst_id, collect_logs=False):
if collect_logs:
collect_instance_logs(inst_id)
max_retries = 10
for attempt in range(1, max_retries + 1):
try:
if args.raw:
with open(os.devnull, 'w') as devnull, redirect_stdout(devnull), redirect_stderr(devnull):
instances_api.destroy_instance(client, id=inst_id)
else:
instances_api.destroy_instance(client, id=inst_id)
if not args.raw:
print(f"Instance {inst_id} destroyed successfully on attempt {attempt}.")
return {"success": True}
except Exception as e:
if not args.raw:
print(f"Error destroying instance {inst_id}: {safe_error(e)}")
if attempt < max_retries:
if not args.raw:
print(f"Retrying in 10 seconds... (Attempt {attempt}/{max_retries})")
time.sleep(10)
else:
if not args.raw:
print(f"Failed to destroy instance {inst_id} after {max_retries} attempts.")
return {"success": False, "error": "Max retries exceeded"}
# ----- wait for instance to start -----
def wait_for_instance(inst_id, timeout=900, interval=10):
start_time = time.time()
debug_print("Starting wait_for_instance with ID:", inst_id)
while time.time() - start_time < timeout:
try:
instance_info = instances_api.show_instance(client, id=inst_id)
if not instance_info:
progress_print(f"No information returned for instance {inst_id}. Retrying...")
time.sleep(interval)
continue
status_msg = instance_info.get('status_msg', '')
status_msg_clean = status_msg.strip() if isinstance(status_msg, str) else ""
status_msg_lower = status_msg_clean.lower()
status_msg_is_error = any(
token in status_msg_lower
for token in (
"error",
"failed",
"failure",
"exception",
"traceback",
"oci runtime",
"permission denied",
)
)
if status_msg_clean and status_msg_is_error:
diagnostic = classify_status_msg(status_msg_clean) or make_failure(
DAEMON_STARTUP_FAILED,
stage="startup",
error=status_msg_clean,
underlying_error=status_msg_clean,
)
reason = f"Instance {inst_id} encountered an error: {status_msg_clean}"
progress_print(reason)
if instance_exist(inst_id):
destroy_instance_silent(inst_id, collect_logs=True)
progress_print(f"Instance {inst_id} has been destroyed due to error.")
else:
progress_print(f"Instance {inst_id} could not be destroyed or does not exist.")
return False, reason, diagnostic
actual_status = instance_info.get('actual_status', 'unknown')
intended_status = instance_info.get('intended_status', 'unknown')
if actual_status == 'offline':
reason = "Instance offline during testing"
diagnostic = make_failure(INSTANCE_OFFLINE_BEFORE_TEST, stage="startup")
progress_print(reason)
if instance_exist(inst_id):
destroy_instance_silent(inst_id, collect_logs=True)
progress_print(f"Instance {inst_id} has been destroyed due to being offline.")
else:
progress_print(f"Instance {inst_id} could not be destroyed or does not exist.")
return False, reason, diagnostic
if intended_status in ('stopped', 'exited') or actual_status in ('stopped', 'exited'):
reason = f"Instance {inst_id} stopped before reaching running status"
if status_msg_clean:
reason = f"{reason}: {status_msg_clean}"
diagnostic = classify_status_msg(status_msg_clean) if status_msg_clean else None
if diagnostic is None:
diagnostic = make_failure(
DAEMON_STARTUP_FAILED,
stage="startup",
summary="Instance stopped before the self-test runtime started.",
details=reason,
error=status_msg_clean or None,
underlying_error=status_msg_clean or None,
)
progress_print(reason)
if instance_exist(inst_id):
destroy_instance_silent(inst_id, collect_logs=True)
progress_print(f"Instance {inst_id} has been destroyed due to startup failure.")
else:
progress_print(f"Instance {inst_id} could not be destroyed or does not exist.")
return False, reason, diagnostic
if intended_status == 'running' and actual_status == 'running':
debug_print(f"Instance {inst_id} is now running.")
return instance_info, None, None
progress_print(f"Instance {inst_id} status: {actual_status}... waiting for 'running' status.")
time.sleep(interval)
except Exception as e:
error = safe_error(e)
progress_print(f"Error retrieving instance info for {inst_id}: {error}. Retrying...")
debug_print(f"Exception details: {error}")
time.sleep(interval)
reason = f"Instance did not become running within {timeout} seconds. Verify network configuration. Use the self-test machine function in the vastai CLI"
progress_print(reason)
return False, reason, make_failure(
INSTANCE_START_TIMEOUT,
stage="startup",
details=reason,
)
# ----- run machine tester -----
def run_machinetester(ip_address, port, inst_id, machine_id, delay, mapped_ports=None):
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
delay = int(delay)
message = ''
legacy_parser = LegacyProgressParser()
progress_url = f"https://{ip_address}:{port}/progress"
request_timeout = 10
attempt_count = 0
last_error_type = None
last_error = None
last_status_code = None
first_connection_established = False
def update_progress_endpoint():
endpoint = make_progress_endpoint_diagnostic(
url=progress_url,
public_ip=ip_address,
container_port=PROGRESS_CONTAINER_PORT,
host_port=port,
timeout_seconds=request_timeout,
attempt_count=attempt_count,
first_connection_established=first_connection_established,
last_error_type=last_error_type,
last_error=last_error,
last_status_code=last_status_code,
mapped_ports=mapped_ports,
)
result["diagnostics"]["progress_endpoint"] = endpoint
return endpoint
update_progress_endpoint()
def is_instance(iid):
try:
info = instances_api.show_instance(client, id=iid)
debug_print(f"is_instance(): Output from show instance: {info}")
if not info or not isinstance(info, dict):
debug_print("is_instance(): No valid instance information received.")
return 'unknown'
actual_status = info.get('actual_status', 'unknown')
return actual_status if actual_status in ['running', 'offline', 'exited', 'created'] else 'unknown'
except Exception as e:
debug_print(f"is_instance(): Error: {safe_error(e)}")
return 'unknown'
if delay > 0:
debug_print(f"Sleeping for {delay} seconds before starting tests.")
time.sleep(delay)
start_time = time.time()
no_response_seconds = 0
printed_lines = set()
instance_destroyed = False
try:
while time.time() - start_time < 600:
status = is_instance(inst_id)
debug_print(f"Instance {inst_id} status: {status}")
if status == 'offline':
reason = "Instance offline during testing"
progress_print(f"Instance {inst_id} went offline. {reason}")
destroy_instance_silent(inst_id, collect_logs=True)
instance_destroyed = True
return False, reason, make_failure(INSTANCE_OFFLINE_BEFORE_TEST, stage="runtime")
try:
debug_print(f"Sending GET request to {progress_url}")
attempt_count += 1
response = requests.get(progress_url, verify=False, timeout=request_timeout)
last_status_code = response.status_code
last_error_type = None
last_error = None
if response.status_code == 200 and not first_connection_established:
progress_print("Successfully established HTTPS connection to the server.")
first_connection_established = True
if response.status_code == 200:
message = response.text.strip()
else:
last_error_type = "HTTPStatus"
last_error = f"HTTP {response.status_code} from progress endpoint"
message = ''
debug_print(f"Received message: '{message}'")
except requests.exceptions.RequestException as e:
last_status_code = None
last_error_type = e.__class__.__name__
last_error = safe_error(e)
debug_print(f"Error making HTTPS request: {last_error}")
message = ''
update_progress_endpoint()
if message:
lines = message.split('\n')
new_lines = [line for line in lines if line not in printed_lines]
for line in new_lines:
diagnostic = legacy_parser.process_line(line)
if line == 'DONE':
progress_print("Test completed successfully.")
progress_print("Test passed.")
destroy_instance_silent(inst_id)
instance_destroyed = True
return True, "", None
elif line.startswith('ERROR'):
progress_print(line)
progress_print(f"Test failed with error: {line}.")
destroy_instance_silent(inst_id, collect_logs=True)
instance_destroyed = True
return False, line, diagnostic
else:
progress_print(line)
printed_lines.add(line)
no_response_seconds = 0
else:
no_response_seconds += 20
debug_print(f"No message received. Incremented no_response_seconds to {no_response_seconds}.")
if status == 'running' and no_response_seconds >= 120:
endpoint = update_progress_endpoint()
if not first_connection_established:
progress_print("No response for 120s — port was never reachable.")
progress_print(f"Tried: {endpoint.get('url')}")
if endpoint.get("external_port"):
progress_print(f"External port tested: {endpoint.get('external_port')}")
if endpoint.get("last_error_type") or endpoint.get("last_status_code") is not None:
last_result = endpoint.get("last_error_type") or f"HTTP {endpoint.get('last_status_code')}"
if endpoint.get("last_error"):
last_result = f"{last_result}: {endpoint.get('last_error')}"
progress_print(f"Last result: {last_result}")
progress_print("Possible causes:")
progress_print(" 1. TCP firewall/NAT forwarding is blocking the mapped public port")
progress_print(" 2. Container did not start or did not bind the progress server")
progress_print(" 3. NAT loopback/hairpinning may fail when testing from the same LAN as the host")
progress_print(f" 4. direct_port_count below the 3 ports/GPU minimum - check with: vastai search offers 'machine_id={machine_id} rentable=any rented=any'")
return_reason = "Port never reachable within 120 seconds"
diagnostic = make_failure(
PROGRESS_ENDPOINT_UNREACHABLE,
stage="runtime_connect",
details=return_reason,
progress_endpoint=endpoint,
)
elif endpoint.get("last_status_code") == 200 and not endpoint.get("last_error_type"):
progress_print("Progress endpoint was reachable but returned no output for 120s.")
progress_print(f"Tried: {endpoint.get('url')}")
if endpoint.get("external_port"):
progress_print(f"External port tested: {endpoint.get('external_port')}")
progress_print("Possible causes:")
progress_print(" 1. Runtime script stalled before writing progress")
progress_print(" 2. Container process is alive but the test worker is hung")
progress_print(" 3. Host stall, GPU hang, or slow I/O prevented progress updates")
return_reason = "Progress endpoint returned no output for 120 seconds"
diagnostic = make_failure(
PROGRESS_EMPTY_TIMEOUT,
stage="runtime_progress",
details=return_reason,
progress_endpoint=endpoint,
)
else:
progress_print("Connection lost after initial success — instance may have crashed.")
progress_print(f"Tried: {endpoint.get('url')}")
if endpoint.get("external_port"):
progress_print(f"External port tested: {endpoint.get('external_port')}")
if endpoint.get("last_error_type") or endpoint.get("last_status_code") is not None:
last_result = endpoint.get("last_error_type") or f"HTTP {endpoint.get('last_status_code')}"
if endpoint.get("last_error"):
last_result = f"{last_result}: {endpoint.get('last_error')}"
progress_print(f"Last result: {last_result}")
progress_print("Possible causes:")
progress_print(" 1. Container crash, OOM, or runtime process exit — check docker logs")
progress_print(" 2. GPU errors or reset — check dmesg for Xid errors")
progress_print(" 3. Host stall or network loss after the progress server started")
return_reason = "Connection lost after 120 seconds — possible crash during stress test"
diagnostic = make_failure(
PROGRESS_ENDPOINT_LOST,
stage="runtime_connect",
details=return_reason,
progress_endpoint=endpoint,
)
destroy_instance_silent(inst_id, collect_logs=True)
instance_destroyed = True
return False, return_reason, diagnostic
debug_print("Waiting for 20 seconds before the next check.")
time.sleep(20)
debug_print(f"Time limit reached. Destroying instance {inst_id}.")
return False, "Test did not complete within the time limit", make_failure(
RUNTIME_TEST_TIMEOUT,
stage="runtime",
)
finally:
if not instance_destroyed and inst_id and instance_exist(inst_id):
destroy_instance_silent(inst_id, collect_logs=True)
progress_print(f"Machine: {machine_id} Done with testing remote.py results {message}")
warnings.simplefilter('default')
# ----- main orchestration: wait then test -----
result["phase"] = "wait"
result["stage"] = "wait_for_instance"
instance_info, wait_reason, wait_diagnostic = wait_for_instance(instance_id)
if not instance_info:
set_runtime_failure(wait_diagnostic, wait_reason)
else:
ip_address = instance_info.get("public_ipaddr")
if not ip_address:
set_runtime_failure(
make_failure(MISSING_PUBLIC_IP, stage="instance_network"),
"Failed to retrieve public IP address.",
)
else:
all_ports = instance_info.get("ports", {})
port_mappings = all_ports.get("5000/tcp", [])
port = port_mappings[0].get("HostPort") if port_mappings else None
if not port:
endpoint = make_progress_endpoint_diagnostic(
public_ip=ip_address,
container_port=PROGRESS_CONTAINER_PORT,
host_port=None,
timeout_seconds=10,
mapped_ports=all_ports,
)
result["diagnostics"]["progress_endpoint"] = endpoint
progress_print(f"Port 5000/tcp not found in mapped ports. Available ports: {list(all_ports.keys())}")
progress_print("Possible causes:")
progress_print(" 1. The instance launch did not map the self-test progress port")
progress_print(f" 2. direct_port_count below the 3 ports/GPU minimum - check with: vastai search offers 'machine_id={args.machine_id} rentable=any rented=any'")
progress_print(" 3. Container is not exposing port 5000/tcp")
set_runtime_failure(
make_failure(
PROGRESS_PORT_NOT_MAPPED,
stage="instance_network",
details=f"Available ports: {list(all_ports.keys())}",
progress_endpoint=endpoint,
),
"Failed to retrieve mapped port.",
)
else:
delay = "15"
result["phase"] = "test"
result["stage"] = "run_machinetester"
success, reason, runtime_diagnostic = run_machinetester(
ip_address, port, instance_id, args.machine_id, delay, mapped_ports=all_ports,
)
result["success"] = success
result["reason"] = reason
if success:
result["phase"] = "complete"
result["stage"] = "complete"
else:
set_runtime_failure(runtime_diagnostic, reason)
except KeyboardInterrupt:
result["success"] = False
set_runtime_failure(
make_failure(INTERRUPTED, stage=result.get("stage")),
"Interrupted by user (Ctrl+C)",
)
result["error"] = result["reason"]
progress_print("\nInterrupted — cleaning up test instance...")
except Exception as e:
error = safe_error(e)
result["success"] = False
result["reason"] = error
result["failure_code"] = "unexpected_error"
result["failure"] = {
"code": "unexpected_error",
"summary": error,
"remediation": "Retry with --debugging and inspect the error details.",
}
result["error"] = error
finally:
# Always attempt to destroy the test instance, including on Ctrl+C.
# KeyboardInterrupt is BaseException (not Exception) so it skips the
# typed except above and lands here. Surface failures loudly: a
# silently-leaked instance keeps billing the host.
if instance_id:
try:
info = instances_api.show_instance(client, id=instance_id)
if not info:
debug_print(f"Test instance {instance_id} already gone during cleanup.")
info = {}
status = (info or {}).get('intended_status') or (info or {}).get('actual_status')
if info and status not in ('destroyed', 'terminated', 'offline'):
if not result.get("success"):
collect_instance_logs(instance_id)
progress_print(f"Destroying test instance {instance_id} (status: {status})...")
instances_api.destroy_instance(client, id=instance_id)
progress_print(f"Test instance {instance_id} destroyed.")
except KeyboardInterrupt:
progress_print(
f"\nSecond interrupt during cleanup — instance {instance_id} may still be running.\n"
f" Destroy it manually: vastai destroy instance {instance_id}"
)
raise
except requests.exceptions.HTTPError as e:
if e.response is not None and e.response.status_code == 404:
debug_print(f"Test instance {instance_id} already gone during cleanup.")
else:
progress_print(
f"WARNING: failed to destroy test instance {instance_id}: {safe_error(e)}\n"
f" Destroy it manually: vastai destroy instance {instance_id}"
)
except Exception as e:
progress_print(
f"WARNING: failed to destroy test instance {instance_id}: {safe_error(e)}\n"
f" Destroy it manually: vastai destroy instance {instance_id}"
)
if args.raw:
if not result.get("success"):
ensure_support_bundle()
return result
else:
if result.get("warning"):
print(result["warning"])
if result["success"]:
if result.get("diagnostics", {}).get("preflight_failure"):
print("Runtime checks passed, but minimum requirement checks were ignored.")
print("This run does not qualify the machine for verification.")
else:
print("Test completed successfully.")
sys.exit(0)
else:
render_runtime_failure()
bundle = ensure_support_bundle()
if bundle:
for line in format_bundle_summary(bundle):
print(line)
print(f"Test failed: {result['reason']}")
sys.exit(1)
cli/commands/metrics.py
"""CLI commands for platform-wide GPU market metrics (host/admin access)."""
import json
import sys
from datetime import datetime, timezone
from vastai.cli.parser import argument, hidden_aliases
from vastai.cli.display import display_table, deindent
from vastai.cli.utils import get_parser as _get_parser, get_client
from vastai.api import metrics as metrics_api
parser = _get_parser()
_VERIFIED_MAP = {"true": "yes", "false": "no", "all": "all"}
_HOSTING_MAP = {"true": "secure_cloud", "false": "community", "all": "all"}
_NEEDS_MACHINE_MSG = "No metrics available. This endpoint is for hosts with active machines."
def notice_if_old_name(args, old, new):
if getattr(args, "command", None) == old:
print(f"NOTICE: `vastai {old}` is now `vastai {new}`. The old name will be removed in the next release.",
file=sys.stderr)
_GPU_CURRENT_FIELDS = (
("gpu_name", "GPU", "{}", None, True),
("total", "Total", "{}", None, False),
("available", "Avail", "{}", None, False),
("unavailable", "Unrent", "{}", None, False),
("usage", "Usage%", "{:.1f}", None, False),
("rented_verified", "Rnt/Ver", "{}", None, False),
("avail_verified", "Avl/Ver", "{}", None, False),
("rented_unverified", "Rnt/Unv", "{}", None, False),
("avail_unverified", "Avl/Unv", "{}", None, False),
("dlperf", "DLPerf", "{:.1f}", None, False),
("tflops", "TFLOPS", "{:.1f}", None, False),
("tflops_per_dollar", "TFLOPS/$", "{:.1f}", None, False),
("price_p10", "P10 $/hr", "{:.4f}", None, False),
("price_median", "Med $/hr", "{:.4f}", None, False),
("price_p90", "P90 $/hr", "{:.4f}", None, False),
)
_GPU_TRENDS_FIELDS = (
("date", "Date", "{}", None, True),
("rented_verified", "Rnt/Ver", "{}", None, False),
("avail_verified", "Avl/Ver", "{}", None, False),
("rented_unverified", "Rnt/Unv", "{}", None, False),
("avail_unverified", "Avl/Unv", "{}", None, False),
("unavail_verified", "Unr/Ver", "{}", None, False),
("unavail_unverified","Unr/Unv", "{}", None, False),
("total", "Total", "{}", None, False),
("rented_p10", "R.p10", "{:.4f}", None, False),
("rented_median", "R.med", "{:.4f}", None, False),
("rented_p90", "R.p90", "{:.4f}", None, False),
("avail_p10", "A.p10", "{:.4f}", None, False),
("avail_median", "A.med", "{:.4f}", None, False),
("avail_p90", "A.p90", "{:.4f}", None, False),
)
_GPU_LOCATION_FIELDS = (
("gpu_name", "GPU", "{}", None, True),
("city", "City", "{}", None, True),
("country_code", "CC", "{}", None, True),
("num_gpus", "GPUs", "{}", None, False),
("state", "State", "{}", None, False),
("verified", "Verified", "{}", None, False),
("latitude", "Lat", "{:.4f}", None, False),
("longitude", "Lon", "{:.4f}", None, False),
)
@parser.command(
argument("--verified", type=str, choices=["true", "false", "all"], default="all",
help="Filter GPUs by verification status"),
argument("--datacenter", type=str, choices=["true", "false", "all"], default="all",
help="Filter GPUs by datacenter hosting type"),
argument("--num-gpus", type=str, choices=["all", "1", "2", "4", "8", "other"], default="all",
help="Filter by machine GPU-count bucket (1x, 2x, 4x, 8x, or other)"),
usage="vastai show gpu-metrics [OPTIONS]",
aliases=hidden_aliases(["metrics gpu"]),
help="[Host] Get current GPU market metrics",
epilog=deindent("""
Get current GPU metrics with counts, usage, performance, and pricing.
For historical metrics, see the `show gpu-trends` command.
Requires host or admin access.
Examples:
vastai show gpu-metrics
vastai show gpu-metrics --verified true --datacenter true
vastai show gpu-metrics --num-gpus 8
vastai show gpu-metrics --raw
"""),
)
def show__gpu_metrics(args):
"""Get current GPU metrics."""
notice_if_old_name(args, "metrics gpu", "show gpu-metrics")
client = get_client(args)
resp = metrics_api.gpu_current(
client,
verified=_VERIFIED_MAP[args.verified],
hosting_type=_HOSTING_MAP[args.datacenter],
num_gpus=args.num_gpus,
)
if resp.get("needs_machine"):
print(_NEEDS_MACHINE_MSG)
return
if args.raw:
print(json.dumps(resp, indent=1))
return
display_table(resp.get("gpus", []), _GPU_CURRENT_FIELDS)
@parser.command(
argument("name", type=str, nargs="?", default="RTX 5090,RTX 4090,RTX 3090",
help="GPU name, comma-separated list, or 'all'. Underscores are accepted in place of spaces."),
argument("--verified", type=str, choices=["true", "false", "all"], default="all",
help="Filter by verified status"),
argument("--datacenter", type=str, choices=["true", "false", "all"], default="all",
help="Filter by datacenter hosting type"),
argument("--num-gpus", type=str, choices=["all", "1", "2", "4", "8", "other"], default="all",
help="Filter by machine GPU-count bucket (1x, 2x, 4x, 8x, or other)"),
argument("--start", type=int, default=None, help="Start unix timestamp"),
argument("--end", type=int, default=None, help="End unix timestamp"),
argument("--step", type=int, default=None,
help="Time between data points in seconds (e.g. 3600 for hourly). Minimum 60s; step may be raised server-side to cap points returned."),
argument("--full", action="store_true", default=False,
help="Show all data points instead of sampling ~20"),
usage="vastai show gpu-trends [NAME] [OPTIONS]",
aliases=hidden_aliases(["metrics gpu-trends"]),
help="[Host] Get GPU market history",
epilog=deindent("""
Show GPU supply/demand and pricing trends over time. Defaults to RTX 5090, 4090, 3090
for the last 24 hours. Requires host or admin access.
Examples:
vastai show gpu-trends
vastai show gpu-trends "RTX 4090"
vastai show gpu-trends "RTX 4090" --full
vastai show gpu-trends "RTX 4090" --raw
vastai show gpu-trends all --verified true --datacenter true
vastai show gpu-trends "RTX 5090" --num-gpus 8
vastai show gpu-trends "RTX 4090,H100_SXM" --start 1773298800 --end 1773817200 --step 3600
"""),
)
def show__gpu_trends(args):
"""Get GPU metrics history."""
notice_if_old_name(args, "metrics gpu-trends", "show gpu-trends")
# Accept underscores as space aliases: "H100_SXM" -> "H100 SXM"
args.name = args.name.replace("_", " ")
client = get_client(args)
resp = metrics_api.gpu_history(
client,
gpu_name=args.name,
verified=_VERIFIED_MAP[args.verified],
hosting_type=_HOSTING_MAP[args.datacenter],
start=args.start,
end=args.end,
step=args.step,
num_gpus=args.num_gpus,
)
if resp.get("needs_machine"):
print(_NEEDS_MACHINE_MSG)
return
if "gpus" in resp:
requested = [g.strip() for g in args.name.split(",")]
gpu_items = [(name, resp["gpus"][name]) for name in requested if name in resp["gpus"]]
seen = set(requested)
for name, data in resp["gpus"].items():
if name not in seen:
gpu_items.append((name, data))
else:
gpu_items = [(args.name, resp)]
sd_keys = ["rented_verified", "avail_verified", "rented_unverified", "avail_unverified",
"unavail_verified", "unavail_unverified", "total"]
pr_keys = ["rented_p10", "rented_median", "rented_p90", "avail_p10", "avail_median", "avail_p90"]
if args.raw:
raw_out = {}
for gpu_name, gpu_data in gpu_items:
sd = gpu_data.get("supply_demand", {})
pr = gpu_data.get("pricing", {})
timestamps = sd.get("timestamps", [])
rows = []
for i, ts in enumerate(timestamps):
row = {"date": datetime.fromtimestamp(ts, tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S"), "timestamp": ts}
for key in sd_keys:
arr = sd.get(key, [])
row[key] = arr[i] if i < len(arr) else None
for key in pr_keys:
arr = pr.get(key, [])
row[key] = arr[i] if i < len(arr) else None
rows.append(row)
raw_out[gpu_name] = {"stats": gpu_data.get("stats", {}), "rows": rows}
print(json.dumps(raw_out, indent=1))
return
for gpu_name, gpu_data in gpu_items:
stats = gpu_data.get("stats", {})
print(f"\n=== {gpu_name} ===")
print(f" tflops: {stats.get('tflops', '-')} dlperf: {stats.get('dlperf', '-')} tflops/$: {stats.get('tflops_per_dollar', '-')}")
sd = gpu_data.get("supply_demand", {})
pr = gpu_data.get("pricing", {})
timestamps = sd.get("timestamps", [])
if not timestamps:
print(" (no data)")
continue
n = len(timestamps)
sample_step = 1 if args.full else max(1, n // 20)
rows = []
for i in range(0, n, sample_step):
row = {"date": datetime.fromtimestamp(timestamps[i], tz=timezone.utc).strftime("%Y-%m-%d %H:%M")}
for key in sd_keys:
arr = sd.get(key, [])
row[key] = arr[i] if i < len(arr) else None
for key in pr_keys:
arr = pr.get(key, [])
row[key] = arr[i] if i < len(arr) else None
rows.append(row)
display_table(rows, _GPU_TRENDS_FIELDS)
if sample_step > 1:
print(f"\n ({n} data points, showing every {sample_step}th. Use --full for all.)")
@parser.command(
argument("--verified", type=str, choices=["true", "false", "all"], default="all",
help="Filter by verification status"),
argument("--datacenter", type=str, choices=["true", "false", "all"], default="all",
help="Filter by datacenter hosting type"),
argument("--rented", type=str, choices=["true", "false", "all"], default="all",
help="Filter by rented status. 'false' covers both available and unrentable GPUs; use --state to tell them apart"),
argument("--state", type=str, choices=["all", "rented", "available", "unavailable"], default="all",
help="Filter by listing state. 'unavailable' means listed but neither rentable nor rented"),
argument("--gpu", type=str, default=None,
help="Filter by GPU name (comma-separated list). Underscores are accepted in place of spaces."),
usage="vastai show gpu-locations [OPTIONS]",
aliases=hidden_aliases(["metrics gpu-locations"]),
help="[Host] Get GPU location metrics",
epilog=deindent("""
Show geographic locations of GPUs on the platform. Filtering is applied
client-side — the endpoint returns one shared dataset and the CLI narrows
the rows locally. Requires host or admin access.
Examples:
vastai show gpu-locations
vastai show gpu-locations --verified true --datacenter true
vastai show gpu-locations --gpu "RTX 4090,H100_SXM"
vastai show gpu-locations --rented false --raw
"""),
)
def show__gpu_locations(args):
"""Get GPU location metrics."""
notice_if_old_name(args, "metrics gpu-locations", "show gpu-locations")
client = get_client(args)
resp = metrics_api.gpu_locations(client)
if resp.get("needs_machine"):
print(_NEEDS_MACHINE_MSG)
return
locations = resp.get("locations", [])
for field in ("verified", "datacenter"):
choice = getattr(args, field)
if choice != "all":
want = choice == "true"
locations = [loc for loc in locations if bool(loc.get(field)) == want]
# current servers send a three-valued state; older ones only the rented boolean
def state_of(loc):
return loc.get("state") or ("rented" if loc.get("rented") else "available")
if args.rented != "all":
locations = [loc for loc in locations if (state_of(loc) == "rented") == (args.rented == "true")]
if args.state != "all":
locations = [loc for loc in locations if state_of(loc) == args.state]
if args.gpu:
wanted_gpus = {g.strip().replace("_", " ") for g in args.gpu.split(",") if g.strip()}
locations = [loc for loc in locations if loc.get("gpu_name") in wanted_gpus]
if args.raw:
print(json.dumps({"success": True, "locations": locations}, indent=1))
return
display_table(locations, _GPU_LOCATION_FIELDS)
cli/commands/misc.py
"""CLI commands for miscellaneous operations (execute, logs, reports, ssh/scp urls, snapshots, templates, self-test).
Note: create/update/delete template commands are registered in offers.py.
self_test__machine is registered in machines.py.
"""
import json
import re
import time
import subprocess
import argparse
import os
import requests as req_lib
from vastai.cli.parser import argument
from vastai.cli.display import deindent
from vastai.cli.util import (
default_start_date, default_end_date,
parse_day_cron_style, parse_hour_cron_style,
validate_frequency_values, add_scheduled_job,
)
from vastai.api import instances as instances_api
from vastai.api import machines as machines_api
from vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401
parser = _get_parser()
# ---------------------------------------------------------------------------
# execute
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of instance to execute on", type=int),
argument("COMMAND", help="bash command surrounded by single quotes", type=str),
argument("--schedule", choices=["HOURLY", "DAILY", "WEEKLY"], help="try to schedule a command to run hourly, daily, or weekly. Valid values are HOURLY, DAILY, WEEKLY For ex. --schedule DAILY"),
argument("--start_date", type=str, default=default_start_date(), help="Start date/time in format 'YYYY-MM-DD HH:MM:SS PM' (UTC). Default is now. (optional)"),
argument("--end_date", type=str, default=default_end_date(), help="End date/time in format 'YYYY-MM-DD HH:MM:SS PM' (UTC). Default is 7 days from now. (optional)"),
argument("--day", type=parse_day_cron_style, help="Day of week you want scheduled job to run on (0-6, where 0=Sunday) or \"*\". Default will be 0. For ex. --day 0", default=0),
argument("--hour", type=parse_hour_cron_style, help="Hour of day you want scheduled job to run on (0-23) or \"*\" (UTC). Default will be 0. For ex. --hour 16", default=0),
usage="vastai execute id COMMAND",
help="Execute a (constrained) remote command on a machine",
epilog=deindent("""
Examples:
vastai execute 99999 'ls -l -o -r'
vastai execute 99999 'rm -r home/delete_this.txt'
vastai execute 99999 'du -d2 -h'
available commands:
ls List directory contents
rm Remote files or directories
du Summarize device usage for a set of files
Return value:
Returns the output of the command which was executed on the instance, if successful. May take a few seconds to retrieve the results.
"""),
)
def execute(args):
"""Execute a (constrained) remote command on a machine."""
json_blob = {"command": args.COMMAND}
if args.explain:
print("request json: ")
print(json_blob)
client = get_client(args)
rj = instances_api.execute(client, id=args.id, command=args.COMMAND)
if args.schedule:
validate_frequency_values(args.day, args.hour, args.schedule)
cli_command = "execute"
api_endpoint = "/api/v0/instances/command/{id}/".format(id=args.id)
json_blob["instance_id"] = args.id
add_scheduled_job(client, args, json_blob, cli_command, api_endpoint, "PUT", instance_id=args.id)
return
if rj.get("success"):
for i in range(0, 30):
time.sleep(0.3)
url = rj["result_url"]
r2 = req_lib.get(url)
if r2.status_code == 200:
filtered_text = r2.text.replace(rj["writeable_path"], '')
print(filtered_text)
break
else:
print(rj)
# ---------------------------------------------------------------------------
# logs
# ---------------------------------------------------------------------------
@parser.command(
argument("INSTANCE_ID", help="id of instance", type=int),
argument("--tail", help="Number of lines to show from the end of the logs (default '1000')", type=str),
argument("--filter", help="Grep filter for log entries", type=str),
argument("--daemon-logs", help="Fetch daemon system logs instead of container logs", action="store_true"),
usage="vastai logs INSTANCE_ID [OPTIONS] ",
help="Get the logs for an instance",
)
def logs(args):
"""Get the logs for an instance."""
if args.explain:
print("request json: ")
json_blob = {}
if args.filter:
json_blob['filter'] = args.filter
if args.tail:
json_blob['tail'] = args.tail
if args.daemon_logs:
json_blob['daemon_logs'] = 'true'
print(json_blob)
client = get_client(args)
rj = instances_api.logs(client, instance_id=args.INSTANCE_ID, tail=args.tail,
filter=args.filter, daemon_logs=args.daemon_logs)
if isinstance(rj, str):
cleaned_text = re.sub(r'\n\s*\n', '\n', rj)
print(cleaned_text)
else:
print(rj.get("msg", rj))
# ---------------------------------------------------------------------------
# reports
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="machine id", type=int),
usage="vastai reports ID",
help="Get the user reports for a given machine",
)
def reports(args):
"""Get the user reports for a given machine."""
if args.explain:
print("request json: ")
print({"machine_id": args.id})
client = get_client(args)
result = machines_api.reports(client, id=args.id)
print(f"reports: {json.dumps(result, indent=2)}")
# ---------------------------------------------------------------------------
# ssh-url / scp-url
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of instance", type=int),
usage="vastai ssh-url ID",
help="ssh url helper",
)
def ssh_url(args):
"""Get the SSH URL for an instance."""
return _ssh_url(args, "ssh://")
@parser.command(
argument("id", help="id", type=int),
usage="vastai scp-url ID",
help="scp url helper",
)
def scp_url(args):
"""Get the SCP URL for an instance."""
return _ssh_url(args, "scp://")
def _ssh_url(args, protocol):
"""Internal helper for ssh/scp URL generation."""
from vastai.cli.util import DIRS
json_object = None
# Try reading cached SSH info
try:
with open(f"{DIRS['temp']}/ssh_{args.id}.json", 'r') as openfile:
json_object = json.load(openfile)
except Exception:
pass
port = None
ipaddr = None
if json_object is not None:
ipaddr = json_object["ipaddr"]
port = json_object["port"]
if ipaddr is None or ipaddr.endswith('.vast.ai'):
client = get_client(args)
rows = instances_api.show_instances(client)
if args.id:
matches = [row for row in rows if row['id'] == args.id]
if not matches:
print(f"error: no instance found with id {args.id}")
return 1
instance = matches[0]
elif len(rows) > 1:
print("Found multiple running instances")
return 1
else:
instance = rows[0]
ports = instance.get("ports", {})
port_22d = ports.get("22/tcp", None)
port = -1
try:
if port_22d is not None:
ipaddr = instance["public_ipaddr"]
port = int(port_22d[0]["HostPort"])
else:
ipaddr = instance["ssh_host"]
port = int(instance["ssh_port"]) + 1 if "jupyter" in instance["image_runtype"] else int(instance["ssh_port"])
except Exception:
port = -1
if port > 0:
print(f'{protocol}root@{ipaddr}:{port}')
else:
print('error: ssh port not found')
# Cache SSH info
try:
with open(f"{DIRS['temp']}/ssh_{args.id}.json", "w") as outfile:
json.dump({"ipaddr": ipaddr, "port": port}, outfile)
except Exception:
pass
# ---------------------------------------------------------------------------
# take snapshot
# ---------------------------------------------------------------------------
@parser.command(
argument("instance_id", help="instance_id of the container instance to snapshot", type=str),
argument("--container_registry", help="Container registry to push the snapshot to. Default will be docker.io", type=str, default="docker.io"),
argument("--repo", help="repo to push the snapshot to", type=str),
argument("--docker_login_user", help="Username for container registry with repo", type=str),
argument("--docker_login_pass", help="Password or token for container registry with repo", type=str),
argument("--pause", help="Pause container's processes being executed by the CPU to take snapshot (true/false). Default will be true", type=str, default="true"),
usage="vastai take snapshot INSTANCE_ID "
"--repo REPO --docker_login_user USER --docker_login_pass PASS"
"[--container_registry REGISTRY] [--pause true|false]",
help="Schedule a snapshot of a running container and push it to your repo in a container registry",
epilog=deindent("""
Takes a snapshot of a running container instance and pushes snapshot to the specified repository in container registry.
Use pause=true to pause the container during commit (safer but slower),
or pause=false to leave it running (faster but may produce a less safe snapshot).
"""),
)
def take__snapshot(args):
"""Take a container snapshot and push to registry."""
instance_id = args.instance_id
repo = args.repo
container_registry = args.container_registry
user = args.docker_login_user
password = args.docker_login_pass
pause_flag = args.pause
print(f"Taking snapshot for instance {instance_id} and pushing to repo {repo} in container registry {container_registry}")
if args.explain:
print("Request JSON:")
print(json.dumps({
"id": instance_id, "container_registry": container_registry,
"personal_repo": repo, "docker_login_user": user,
"docker_login_pass": password, "pause": pause_flag,
}, indent=2))
client = get_client(args)
data = instances_api.take_snapshot(
client, instance_id=instance_id, repo=repo,
container_registry=container_registry,
docker_login_user=user, docker_login_pass=password,
pause=pause_flag,
)
if data.get("success"):
print(f"Snapshot request sent successfully. Please check your repo {repo} in container registry {container_registry} in 5-10 mins. It can take longer than 5-10 mins to push your snapshot image to your repo depending on the size of your image.")
else:
print(data.get("msg", "Unknown error with snapshot request"))
# ---------------------------------------------------------------------------
# Note: create__template, update__template, delete__template are in offers.py.
# Note: self_test__machine is in machines.py.
# ---------------------------------------------------------------------------
cli/commands/offers.py
"""CLI commands for searching offers, templates, and benchmarks."""
import json
import sys
import time
from vastai.cli.parser import argument, hidden_aliases
from vastai.cli.display import display_table, displayable_fields, displayable_fields_reserved, benchmarks_displayable_fields, deindent
from vastai.api import offers as offers_api
from vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401
parser = _get_parser()
# ---------------------------------------------------------------------------
# search offers
# ---------------------------------------------------------------------------
@parser.command(
argument("-t", "--type", default="on-demand", help="Show 'on-demand', 'reserved', or 'bid'(interruptible) pricing. default: on-demand"),
argument("-i", "--interruptible", dest="type", const="bid", action="store_const", help="Alias for --type=bid"),
argument("-b", "--bid", dest="type", const="bid", action="store_const", help="Alias for --type=bid"),
argument("-r", "--reserved", dest="type", const="reserved", action="store_const", help="Alias for --type=reserved"),
argument("-d", "--on-demand", dest="type", const="on-demand", action="store_const", help="Alias for --type=on-demand"),
argument("-n", "--no-default", action="store_true", help="Disable default query"),
argument("--new", action="store_true", help="New search exp"),
argument("--limit", type=int, help=""),
argument("--storage", type=float, default=5.0, help="Amount of storage to use for pricing, in GiB. default=5.0GiB"),
argument("-o", "--order", type=str, help="Comma-separated list of fields to sort on. postfix field with - to sort desc. ex: -o 'num_gpus,total_flops-'. default='score-'", default='score-'),
argument("query", help="Query to search for. default: 'external=false rentable=true verified=true cpu_arch in [amd64,arm64]', pass -n to ignore default", nargs="*", default=None),
usage="vastai search offers [--help] [--api-key API_KEY] [--raw] <query>",
help="Search for instance types using custom query",
epilog=deindent("""
Query syntax:
query = comparison comparison...
comparison = field op value
field = <name of a field>
op = one of: <, <=, ==, !=, >=, >, in, notin
value = <bool, int, float, string> | 'any' | [value0, value1, ...]
bool: True, False
note: to pass '>' and '<' on the command line, make sure to use quotes
note: to encode a string query value (ie for gpu_name), replace any spaces ' ' with underscore '_'
Examples:
# search for somewhat reliable single RTX 3090 instances
vastai search offers 'reliability > 0.98 num_gpus=1 gpu_name=RTX_3090'
# search for datacenter gpus with minimal compute_cap and total_flops
vastai search offers 'compute_cap > 610 total_flops > 5 datacenter=True'
# search for reliable 4 gpu offers in Taiwan or Sweden
vastai search offers 'reliability>0.99 num_gpus=4 geolocation in [TW,SE]'
# search for reliable RTX 3090 or 4090 gpus NOT in China or Vietnam
vastai search offers 'reliability>0.99 gpu_name in ["RTX 4090", "RTX 3090"] geolocation notin [CN,VN]'
# search for machines with nvidia drivers 535.86.05 or greater (and various other options)
vastai search offers 'disk_space>146 duration>24 gpu_ram>10 cuda_vers>=12.1 direct_port_count>=2 driver_version >= 535.86.05'
# search for reliable machines with at least 4 gpus, unverified, order by num_gpus
vastai search offers 'reliability > 0.99 num_gpus>=4 verified=False' -o 'num_gpus-'
# search for arm64 cpu architecture
vastai search offers 'cpu_arch=arm64'
Available fields:
Name Type Description
bw_nvlink float bandwidth NVLink
compute_cap: int cuda compute capability*100 (ie: 650 for 6.5, 700 for 7.0)
cpu_arch string host machine cpu architecture (e.g. amd64, arm64)
cpu_cores: int # virtual cpus
cpu_ghz: Float # cpu clock speed GHZ
cpu_cores_effective: float # virtual cpus you get
cpu_ram: float system RAM in gigabytes
cuda_vers: float machine max supported cuda version (based on driver version)
datacenter: bool show only datacenter offers
direct_port_count int open ports on host's router
disk_bw: float disk read bandwidth, in MB/s
disk_space: float disk storage space, in GB
dlperf: float DL-perf score (see FAQ for explanation)
dlperf_usd: float DL-perf/$
dph: float $/hour rental cost
driver_version: string machine's nvidia/amd driver version as 3 digit string ex. "535.86.05,"
duration: float max rental duration in days
external: bool show external offers in addition to datacenter offers
flops_usd: float TFLOPs/$
geolocation: string Two letter country code. Works with operators =, !=, in, notin (e.g. geolocation not in ['XV','XZ'])
gpu_arch string host machine gpu architecture (e.g. nvidia, amd)
gpu_max_power float GPU power limit (watts)
gpu_max_temp float GPU temp limit (C)
gpu_mem_bw: float GPU memory bandwidth in GB/s
gpu_name: string GPU model name (no quotes, replace spaces with underscores, ie: RTX_3090 rather than 'RTX 3090')
gpu_ram: float per GPU RAM in GB
gpu_total_ram: float total GPU RAM in GB
gpu_frac: float Ratio of GPUs in the offer to gpus in the system
gpu_display_active: bool True if the GPU has a display attached
has_avx: bool CPU supports AVX instruction set.
id: int instance unique ID
inet_down: float internet download speed in Mb/s
inet_down_cost: float internet download bandwidth cost in $/GB
inet_up: float internet upload speed in Mb/s
inet_up_cost: float internet upload bandwidth cost in $/GB
machine_id int machine id of instance
min_bid: float current minimum bid price in $/hr for interruptible
num_gpus: int # of GPUs
pci_gen: float PCIE generation
pcie_bw: float PCIE bandwidth (CPU to GPU)
reliability: float machine reliability score (see FAQ for explanation)
rentable: bool is the instance currently rentable
storage_cost: float storage cost in $/GB/month
static_ip: bool is the IP addr static/stable
total_flops: float total TFLOPs from all GPUs
ubuntu_version string host machine ubuntu OS version
verified: bool is the machine verified
vms_enabled: bool is the machine a VM instance
"""),
aliases=hidden_aliases(["search instances"]),
)
def search__offers(args):
"""Creates a query based on search parameters as in the examples above.
:param argparse.Namespace args: should supply all the command-line options
"""
from vastai.api.query import parse_query, offers_fields, offers_alias, offers_mult
try:
if args.no_default:
query = {}
else:
query = {"verified": {"eq": True}, "external": {"eq": False}, "rentable": {"eq": True}}
if args.query is not None:
query = parse_query(args.query, query, offers_fields, offers_alias, offers_mult)
if not args.no_default:
# Request every cpu arch like the console (CLN-3461); explicit cpu_arch wins.
query.setdefault("cpu_arch", {"in": ["amd64", "arm64"]})
order = []
for name in args.order.split(","):
name = name.strip()
if not name:
continue
direction = "asc"
field = name
if name.strip("-") != name:
direction = "desc"
field = name.strip("-")
if name.strip("+") != name:
direction = "asc"
field = name.strip("+")
if field in offers_alias:
field = offers_alias[field]
order.append([field, direction])
query["order"] = order
query["type"] = args.type
if args.limit:
query["limit"] = int(args.limit)
query["allocated_storage"] = args.storage
if query["type"] == 'interruptible':
query["type"] = 'bid'
except ValueError as e:
print("Error: ", e)
return 1
json_blob = query
client = get_client(args)
if args.new:
json_blob = {"select_cols": ['*'], "q": query}
if args.explain:
print("request json: ")
print(json_blob)
r = client.put("/search/asks/", json_data=json_blob)
else:
if args.explain:
print("request json: ")
print(json_blob)
r = client.post("/bundles/", json_data=json_blob)
r.raise_for_status()
if r.headers.get('Content-Type') != 'application/json':
print(f"invalid return Content-Type: {r.headers.get('Content-Type')}")
return
rows = r.json()["offers"]
if args.raw:
return rows
else:
if args.type == "reserved":
display_table(rows, displayable_fields_reserved)
else:
display_table(rows, displayable_fields)
# ---------------------------------------------------------------------------
# search benchmarks
# ---------------------------------------------------------------------------
@parser.command(
argument("query", help="Search query in simple query syntax (see below)", nargs="*", default=None),
argument("-a", "--all", action="store_true", help="force fetching every page, even with --limit/--next-token (fetching everything is already the default otherwise)"),
argument("-l", "--limit", type=int, default=None, help="max results per page; passing this switches to single-page mode instead of fetching everything"),
argument("-t", "--next-token", dest="next_token", help="resume from a pagination token (implies single-page mode)"),
usage="vastai search benchmarks [--help] [--api-key API_KEY] [--raw] [--limit N] [--next-token TOKEN] <query>",
help="Search for benchmark results using custom query",
epilog=deindent("""
By default this fetches every page and prints one combined result. Passing
--limit or --next-token switches to single-page mode, which prints a page
and asks whether to fetch the next one (add --all to still fetch every page
in that mode). The prompt is skipped when output is piped or --raw is used;
those print the next page token so you can resume with --next-token.
--raw prints a flat list when fetching every page, and the full response
(benchmarks plus next_token) in single-page mode.
Query syntax:
query = comparison comparison...
comparison = field op value
field = <name of a field>
op = one of: <, <=, ==, !=, >=, >, in, notin
value = <bool, int, float, string> | 'any' | [value0, value1, ...]
bool: True, False
note: to pass '>' and '<' on the command line, make sure to use quotes
note: to encode a string query value (ie for gpu_name), replace any spaces ' ' with underscore '_'
Examples:
# search for benchmarks with value > 100 on RTX 4090s for 2 specific machines
vastai search benchmarks 'value > 100.0 gpu_name=RTX_4090 machine_id in [302,402]'
# a page at a time, prompting between pages
vastai search benchmarks 'gpu_name=RTX_5090' --limit 50
Available fields:
Name Type Description
contract_id int ID of instance/contract reporting benchmark
dph_base float gpu $/hr at the time of the benchmark
gpu_name string GPU model benchmarked (e.g. RTX_4090)
id int benchmark unique ID
image string image used for benchmark
last_update float date of benchmark
launch_args string launch args of the benchmarked worker
machine_id int id of machine benchmarked
num_gpus int number of gpus used in benchmark
template_hash string hash of the template benchmarked (formerly 'name')
template_id int id of the template benchmarked
type string benchmark type
value float benchmark result (formerly 'score')
"""),
aliases=hidden_aliases(["search benchmarks"]),
)
def search__benchmarks(args):
"""Creates a query based on search parameters as in the examples above.
:param argparse.Namespace args: should supply all the command-line options
"""
from vastai.api.query import parse_query, benchmarks_fields, fix_date_fields
try:
query = {}
if args.query is not None:
query = parse_query(args.query, query, benchmarks_fields)
query = fix_date_fields(query, ['last_update'])
except ValueError as e:
print("Error: ", e)
return 1
client = get_client(args)
# Fetch every page by default; --limit/--next-token opt into manual pagination.
fetch_all = args.all or (args.limit is None and args.next_token is None)
if fetch_all:
rows = offers_api.search_benchmarks(
client, query=query, limit=args.limit, after_token=args.next_token)
if args.raw:
# Flat list of every matching row -- the contract scripts have
# always depended on for --raw. Single-page mode returns the
# envelope instead, since only it has a token worth handing back.
return rows
display_table(rows, benchmarks_displayable_fields)
return
params = offers_api.benchmarks_query_args(
query=query, limit=args.limit, after_token=args.next_token)
page = 0
while True:
page += 1
data = offers_api.search_benchmarks_v1(client, params)
rows = data.get("benchmarks") or []
next_token = data.get("next_token")
if args.raw:
return {"success": True, "benchmarks_found": len(rows),
"benchmarks": rows, "next_token": next_token}
display_table(rows, benchmarks_displayable_fields)
if not next_token:
return
print(f"Next page token: {next_token}")
# Never block on a prompt when piped or scripted.
if not (sys.stdin.isatty() and sys.stdout.isatty()):
return
try:
if input(f"Fetch next page? (page {page + 1}) (y/N): ").strip().lower() != "y":
return
except (EOFError, KeyboardInterrupt):
return
params["after_token"] = next_token
# ---------------------------------------------------------------------------
# search templates
# ---------------------------------------------------------------------------
@parser.command(
argument("query", help="Search query in simple query syntax (see below)", nargs="*", default=None),
usage="vastai search templates [--help] [--api-key API_KEY] [--raw] <query>",
help="Search for template results using custom query",
epilog=deindent("""
Query syntax:
query = comparison comparison...
comparison = field op value
field = <name of a field>
op = one of: <, <=, ==, !=, >=, >, in, notin
value = <bool, int, float, string> | 'any' | [value0, value1, ...]
bool: True, False
note: to pass '>' and '<' on the command line, make sure to use quotes
note: to encode a string query value (ie for gpu_name), replace any spaces ' ' with underscore '_'
Examples:
# search for somewhat reliable single RTX 3090 instances, filter out any duplicates or offers that conflict with our existing stopped instances
vastai search templates 'count_created > 100 creator_id in [38382,48982]'
Available fields:
Name Type Description
creator_id int ID of creator
created_at float time of initial template creation (UTC epoch timestamp)
count_created int #instances created (popularity)
default_tag string image default tag
docker_login_repo string image docker repository
id int template unique ID
image string image used for template
jup_direct bool supports jupyter direct
hash_id string unique hash ID of template
name string displayable name
recent_create_date float last time of instance creation (UTC epoch timestamp)
recommended_disk_space float min disk space required
recommended bool is templated on our recommended list
ssh_direct bool supports ssh direct
tag string image tag
use_ssh bool supports ssh (direct or proxy) """),
aliases=hidden_aliases(["search templates"]),
)
def search__templates(args):
"""Creates a query based on search parameters as in the examples above.
:param argparse.Namespace args: should supply all the command-line options
"""
from vastai.api.query import parse_query, templates_fields, fix_date_fields
try:
query = {}
if args.query is not None:
query = parse_query(args.query, query, templates_fields)
query = fix_date_fields(query, ['created_at', 'recent_create_date'])
except ValueError as e:
print("Error: ", e)
return 1
client = get_client(args)
try:
rows = offers_api.search_templates(client, query=query)
print(json.dumps(rows, indent=1, sort_keys=True))
except Exception as e:
print(f"Error: {e}")
# ---------------------------------------------------------------------------
# search invoices
# ---------------------------------------------------------------------------
@parser.command(
argument("query", help="Search query in simple query syntax (see below)", nargs="*", default=None),
usage="vastai search invoices [--help] [--api-key API_KEY] [--raw] <query>",
help="Search for invoices using custom query",
epilog=deindent("""
Query syntax:
query = comparison comparison...
comparison = field op value
field = <name of a field>
op = one of: <, <=, ==, !=, >=, >, in, notin
value = <bool, int, float, string> | 'any' | [value0, value1, ...]
bool: True, False
note: to pass '>' and '<' on the command line, make sure to use quotes
note: to encode a string query value (ie for gpu_name), replace any spaces ' ' with underscore '_'
Examples:
# search for somewhat reliable single RTX 3090 instances, filter out any duplicates or offers that conflict with our existing stopped instances
vastai search invoices 'amount_cents>3000 '
Available fields:
Name Type Description
id int,
user_id int,
when float, utc epoch timestamp of initial invoice creation
paid_on float, actual payment date (utc epoch timestamp )
payment_expected float, expected payment date (utc epoch timestamp )
amount_cents int, amount of payment in cents
is_credit bool, is a credit purchase
is_delayed bool, is not yet paid
balance_before float, balance before
balance_after float, balance after
original_amount int, original amount of payment
event_id string,
cut_amount int,
cut_percent float,
extra json,
service string, type of payment
stripe_charge json,
stripe_refund json,
stripe_payout json,
error json,
paypal_email string, email for paypal/wise payments
transfer_group string,
failed bool,
refunded bool,
is_check bool,
"""),
aliases=hidden_aliases(["search invoices"]),
)
def search__invoices(args):
"""Search for invoices using custom query."""
from vastai.api.query import parse_query, invoices_fields, fix_date_fields
try:
query = {}
if args.query is not None:
query = parse_query(args.query, query, invoices_fields)
query = fix_date_fields(query, ['when', 'paid_on', 'payment_expected'])
except ValueError as e:
print("Error: ", e)
return 1
client = get_client(args)
rows = offers_api.search_invoices(client, query=query)
if args.raw:
return rows
print(json.dumps(rows, indent=1, sort_keys=True))
# ---------------------------------------------------------------------------
# create / update / delete template
# ---------------------------------------------------------------------------
@parser.command(
argument("--name", help="name of the template", type=str),
argument("--image", help="docker container image to launch", type=str),
argument("--image_tag", help="docker image tag", type=str),
argument("--href", help="link you want to provide", type=str),
argument("--repo", help="link to repository", type=str),
argument("--login", help="docker login arguments for private repo authentication, surround with ''", type=str),
argument("--env", help="Contents of the 'Docker options' field", type=str),
argument("--ssh", help="Launch as an ssh instance type", action="store_true"),
argument("--jupyter", help="Launch as a jupyter instance instead of an ssh instance", action="store_true"),
argument("--direct", help="Use (faster) direct connections for jupyter & ssh", action="store_true"),
argument("--jupyter-dir", help="For runtype 'jupyter', directory in instance to use to launch jupyter", type=str),
argument("--jupyter-lab", help="For runtype 'jupyter', Launch instance with jupyter lab", action="store_true"),
argument("--onstart-cmd", help="contents of onstart script as single argument", type=str),
argument("--search_params", help="search offers filters", type=str),
argument("-n", "--no-default", action="store_true", help="Disable default search param query args"),
argument("--disk_space", help="disk storage space, in GB", type=str),
argument("--readme", help="readme string", type=str),
argument("--hide-readme", help="hide the readme from users", action="store_true"),
argument("--desc", help="description string", type=str),
argument("--public", help="make template available to public", action="store_true"),
usage="vastai create template",
help="Create a new template",
)
def create__template(args):
"""Create a new template."""
fields = offers_api.template_fields_from_flags(
ssh=args.ssh, jupyter=args.jupyter, direct=args.direct,
jupyter_lab=args.jupyter_lab, login=args.login,
hide_readme=args.hide_readme, public=args.public,
search_params=args.search_params, no_default=args.no_default,
always_default=True,
)
extra_filters = fields["extra_filters"]
if args.explain:
print("request json: ")
print({"name": args.name, "image": args.image, "extra_filters": extra_filters})
client = get_client(args)
try:
rj = offers_api.create_template(
client, name=args.name, image=args.image, image_tag=args.image_tag,
href=args.href, repo=args.repo, env=args.env, onstart_cmd=args.onstart_cmd,
jupyter_dir=args.jupyter_dir, disk_space=args.disk_space,
readme=args.readme, desc=args.desc, **fields,
)
if rj.get("success"):
print(f"New Template: {rj['template']}")
else:
print(rj.get('msg', rj))
except Exception:
print("The response is not valid JSON.")
@parser.command(
argument("HASH_ID", help="hash id of the template", type=str),
argument("--name", help="name of the template", type=str),
argument("--image", help="docker container image to launch", type=str),
argument("--image_tag", help="docker image tag", type=str),
argument("--href", help="link you want to provide", type=str),
argument("--repo", help="link to repository", type=str),
argument("--login", help="docker login arguments for private repo authentication, surround with ''", type=str),
argument("--env", help="Contents of the 'Docker options' field", type=str),
argument("--ssh", help="Launch as an ssh instance type", action="store_true"),
argument("--jupyter", help="Launch as a jupyter instance instead of an ssh instance", action="store_true"),
argument("--direct", help="Use (faster) direct connections for jupyter & ssh", action="store_true"),
argument("--jupyter-dir", help="For runtype 'jupyter', directory in instance to use to launch jupyter", type=str),
argument("--jupyter-lab", help="For runtype 'jupyter', Launch instance with jupyter lab", action="store_true"),
argument("--onstart-cmd", help="contents of onstart script as single argument", type=str),
argument("--search_params", help="search offers filters", type=str),
argument("-n", "--no-default", action="store_true", help="Disable default search param query args"),
argument("--disk_space", help="disk storage space, in GB", type=str),
argument("--readme", help="readme string", type=str),
argument("--hide-readme", help="hide the readme from users", action="store_true"),
argument("--desc", help="description string", type=str),
argument("--public", help="make template available to public", action="store_true"),
usage="vastai update template HASH_ID",
help="Update an existing template",
epilog=deindent("""
Edits create a new version of the template with a new hash id; the original stays
valid for anyone already using it. Send the full template, not only the changed fields.
Examples:
vastai update template c81e7ab0e928a508510d1979346de10d --name "tgi-llama2-7B-quantized" --image "ghcr.io/huggingface/text-generation-inference:1.0.3" --env "-p 3000:3000 -e MODEL_ARGS='--model-id TheBloke/Llama-2-7B-chat-GPTQ --quantize gptq'" --onstart-cmd 'wget -O - https://raw.githubusercontent.com/vast-ai/vast-pyworker/main/scripts/launch_tgi.sh | bash' --search_params "gpu_ram>=23 num_gpus=1 gpu_name=RTX_3090 inet_down>128 direct_port_count>3 disk_space>=192 driver_version>=535086005 rented=False" --disk_space 8.0 --ssh --direct
"""),
)
def update__template(args):
"""Update an existing template."""
fields = offers_api.template_fields_from_flags(
ssh=args.ssh, jupyter=args.jupyter, direct=args.direct,
jupyter_lab=args.jupyter_lab, login=args.login,
hide_readme=args.hide_readme, public=args.public,
search_params=args.search_params, no_default=args.no_default,
always_default=True,
)
extra_filters = fields["extra_filters"]
if args.explain:
print("request json: ")
print({"hash_id": args.HASH_ID, "name": args.name, "image": args.image})
client = get_client(args)
try:
rj = offers_api.update_template(
client, hash_id=args.HASH_ID, name=args.name, image=args.image,
image_tag=args.image_tag, href=args.href, repo=args.repo, env=args.env,
onstart_cmd=args.onstart_cmd, jupyter_dir=args.jupyter_dir,
disk_space=args.disk_space, readme=args.readme, desc=args.desc,
**fields,
)
if rj.get("success"):
print(f"updated template: {json.dumps(rj['template'], indent=1)}")
else:
print(rj.get("msg", "template update failed"))
return 1
except Exception as e:
print(_response_msg(getattr(e, "response", None)) or str(e))
return 1
def _response_msg(response):
"""Return the API's msg from an error response body, or None when there is none."""
if response is None:
return None
try:
body = response.json()
except ValueError:
return None
return body.get("msg") if isinstance(body, dict) else None
@parser.command(
argument("--template-id", help="Template ID of Template to Delete", type=int),
argument("--hash-id", help="Hash ID of Template to Delete", type=str),
usage="vastai delete template [--template-id <id> | --hash-id <hash_id>]",
help="Delete a Template",
)
def delete__template(args):
"""Delete a template."""
if not args.hash_id and not args.template_id:
print('ERROR: Must Specify either Template ID or Hash ID to delete a template')
return
if args.explain:
print("request json: ")
print({"hash_id": args.hash_id, "template_id": args.template_id})
client = get_client(args)
try:
rj = offers_api.delete_template(client, hash_id=args.hash_id, template_id=args.template_id)
print(rj.get('msg', rj))
except Exception as e:
print(f"Error: {e}")
cli/commands/price_increase.py
"""CLI commands for the per-row price-increase contract extension flow.
Backend source of truth: ``vast/web/views/instance.py:257-312`` and
``vast/web/price_increase_pending.py``. All three endpoints are per-row
keyed on ``pending_price_increase_id``; the body is exactly
``{"pending_price_increase_id": int}`` with ``extra='forbid'``.
The CLI argument stays a positional **instance ID** to preserve UX
parity with ``show instances``; this module resolves it to the row
id with a single ``GET /pending-price-increases/`` per command
invocation, mirroring what the frontend modal does implicitly.
Stale signal is ``HTTP 404`` with body
``{"success": false, "error": "no_pending_price_increase"}`` (legacy
``HTTP 409`` is treated identically — the FE still recognises it).
"""
import json
import sys
from datetime import datetime, timezone
from requests.exceptions import HTTPError
from vastai.api import price_increase as price_increase_api
from vastai.cli.display import (
display_table,
deindent,
pending_price_increase_fields,
)
from vastai.cli.parser import argument
from vastai.cli.utils import get_parser as _get_parser, get_client
parser = _get_parser()
# Display rescale: backend writes per-second / fractional values; the human
# table renders the conventions the frontend uses on the in-app modal.
_SEC_PER_HOUR = 3600
_SEC_PER_MONTH = 3600 * 24 * 30
# ---------------------------------------------------------------------------
# Formatters
# ---------------------------------------------------------------------------
def _fmt_epoch(ts):
"""Render an epoch float as ``YYYY-MM-DD HH:MM UTC``; ``-`` if absent."""
if ts is None:
return "-"
return datetime.fromtimestamp(float(ts), tz=timezone.utc).strftime("%Y-%m-%d %H:%M UTC")
def _fmt_delta(old, new, scale, precision):
"""Render ``old → new`` for a resource; ``-`` when the resource did not change.
A ``null`` ``new`` means the backend did not change the resource on this
row; the column collapses to ``-`` so the table does not invent a delta.
"""
if new is None:
return "-"
old_part = f"{old * scale:.{precision}f}" if old is not None else "?"
new_part = f"{new * scale:.{precision}f}"
return f"{old_part} -> {new_part}"
def _prepare_row(row):
"""Project a backend row into the table dict ``pending_price_increase_fields`` expects."""
return {
"pending_id": row.get("pending_price_increase_id"),
"instance_id": row.get("contract_id"),
"host_id": row.get("host_id"),
"current_end": _fmt_epoch(row.get("contract_end_date")),
"new_end": _fmt_epoch(row.get("ask_end_date")),
"gpu": _fmt_delta(
row.get("old_gpu_costpersec"), row.get("new_gpu_costpersec"),
_SEC_PER_HOUR, 4,
),
"storage": _fmt_delta(
row.get("old_disk_ram_costpersec"), row.get("new_disk_ram_costpersec"),
_SEC_PER_MONTH, 4,
),
"bw_up": _fmt_delta(
row.get("old_bwu_cost"), row.get("new_bwu_cost"), 1, 4,
),
"bw_down": _fmt_delta(
row.get("old_bwd_cost"), row.get("new_bwd_cost"), 1, 4,
),
"fee": _fmt_delta(
row.get("old_platform_fee"), row.get("new_platform_fee"), 100, 2,
),
}
# ---------------------------------------------------------------------------
# Stale-row detection
# ---------------------------------------------------------------------------
def _is_stale_error(err):
"""Detect the backend's stale-row signal across the documented variants.
Modern backend: ``HTTP 404`` with body
``{"error": "no_pending_price_increase"}``.
Legacy: ``HTTP 409`` (the frontend still recognises it, so do we).
"""
response = getattr(err, "response", None)
if response is None:
return False
status = getattr(response, "status_code", None)
if status == 409:
return True
if status != 404:
return False
try:
body = response.json() or {}
except Exception:
return False
return body.get("error") == price_increase_api.NO_PENDING_PRICE_INCREASE
# ---------------------------------------------------------------------------
# show pending-price-increases
# ---------------------------------------------------------------------------
@parser.command(
argument("-q", "--quiet", action="store_true",
help="print only pending_price_increase_id values, one per line"),
usage="vastai show pending-price-increases [OPTIONS]",
help="List pending price-increase challenges for the authenticated user",
epilog=deindent("""
Lists every pending price-increase challenge the backend has open for
you. Each row shows the per-resource old -> new prices and the cutover
boundary: the new rate applies only after each contract's current
end_date — your remaining time on the current term is billed at the
original price.
Pipe the pending IDs to `vastai accept price-increase` or
`vastai reject price-increase` (use --quiet for one ID per line):
vastai show pending-price-increases --quiet
"""),
)
def show__pending_price_increases(args):
"""List pending price-increase challenges in a human or raw form."""
if args.explain:
print("GET /instances/pending-price-increases/")
client = get_client(args)
envelope = price_increase_api.list_pending(client)
if args.raw:
return envelope
rows = price_increase_api.pending_rows(envelope)
if args.quiet:
for row in rows:
print(row.get("pending_price_increase_id"))
return
if not rows:
print("No pending price increases.")
return
display_table([_prepare_row(r) for r in rows], pending_price_increase_fields)
if envelope.get("truncated"):
print("\nResult truncated by the server; refine with `vastai show "
"pending-price-increases --raw` to retrieve the full payload.")
# ---------------------------------------------------------------------------
# Shared accept/reject confirmation + fan-out
# ---------------------------------------------------------------------------
def _unique_instance_ids(ids):
"""Return integer ids in first-seen order."""
unique = []
seen = set()
for raw_id in ids:
instance_id = int(raw_id)
if instance_id in seen:
continue
seen.add(instance_id)
unique.append(instance_id)
return unique
def _resolve_instance_ids_to_pending(rows, instance_ids):
"""Return ``[(instance_id, pending_id_or_None)]`` preserving CLI argument order."""
return [
(
iid,
(
price_increase_api.find_pending_for_instance(rows, iid) or {}
).get("pending_price_increase_id"),
)
for iid in instance_ids
]
def _prompt_or_exit(present_imperative, instance_ids, rows, *, yes, raw=False):
"""Confirm before mutating; honours the spec's TTY + --yes rules.
Returns ``True`` to proceed, ``False`` to abort cleanly with exit 0.
Non-TTY without ``--yes`` exits 1 with a clear message (matches the
Two-Step Confirmation requirement).
"""
if yes:
return True
if raw:
print(
"--yes is required with --raw so stdout remains machine-readable.",
file=sys.stderr,
)
sys.exit(1)
if not sys.stdin.isatty():
print(
'--yes is required when stdin is not a TTY. Run "vastai show '
'pending-price-increases" to review first.',
file=sys.stderr,
)
sys.exit(1)
instance_id_set = set(instance_ids)
matching = [r for r in rows if r.get("contract_id") in instance_id_set]
if matching:
display_table(
[_prepare_row(r) for r in matching], pending_price_increase_fields,
)
answer = input(
f"{present_imperative} these price increases? [y/N]: "
).strip().lower()
return answer == "y"
def _run_per_row(args, *, route_verb, api_call, present_imperative, past_tense, cutover_note):
"""Shared accept/reject driver: resolve, confirm, fan-out, summarise."""
instance_ids = _unique_instance_ids(args.ids)
if args.explain:
for iid in instance_ids:
print(f"PUT /instances/{route_verb}-price-increase/ body for instance {iid}")
client = get_client(args)
envelope = price_increase_api.list_pending(client)
rows = price_increase_api.pending_rows(envelope)
pairs = _resolve_instance_ids_to_pending(rows, instance_ids)
has_matching_rows = any(pid is not None for _, pid in pairs)
if has_matching_rows and not _prompt_or_exit(
present_imperative, instance_ids, rows, yes=args.yes, raw=args.raw,
):
print("Aborted.")
return
acted_ids, stale, failed = [], 0, 0
results = []
for iid, pid in pairs:
if pid is None:
if not args.raw:
print(f"pending price increase no longer available for instance {iid}")
results.append({
"instance_id": iid,
"pending_price_increase_id": None,
"status": "stale",
"error": price_increase_api.NO_PENDING_PRICE_INCREASE,
})
stale += 1
continue
try:
result = api_call(client, pid)
except HTTPError as err:
if _is_stale_error(err):
if not args.raw:
print(
"pending price increase no longer available — re-run "
"vastai show pending-price-increases"
)
results.append({
"instance_id": iid,
"pending_price_increase_id": pid,
"status": "stale",
"error": price_increase_api.NO_PENDING_PRICE_INCREASE,
})
stale += 1
continue
if not args.raw:
print(f"failed for instance {iid}: {err}", file=sys.stderr)
results.append({
"instance_id": iid,
"pending_price_increase_id": pid,
"status": "failed",
"error": str(err),
})
failed += 1
continue
if not args.raw:
print(
f"{past_tense} pending_id={result.get('pending_price_increase_id')} "
f"contract_id={result.get('contract_id')}"
)
results.append({
"instance_id": iid,
"pending_price_increase_id": result.get("pending_price_increase_id"),
"contract_id": result.get("contract_id"),
"status": past_tense.lower(),
"response": result,
})
acted_ids.append(result.get("contract_id"))
total = len(instance_ids)
n_ok = len(acted_ids)
if args.raw:
exit_code = 1 if failed else 2 if stale else 0
payload = {
"success": exit_code == 0,
"exit_code": exit_code,
"requested": total,
"acted": n_ok,
"stale": stale,
"failed": failed,
"results": results,
}
print(json.dumps(payload, indent=1, sort_keys=True))
sys.exit(exit_code)
print(
f"\n{past_tense} {n_ok} / Stale {stale} / Failed {failed} "
f"of {total} requested."
)
if n_ok:
ids_str = ", ".join(str(cid) for cid in acted_ids)
print(f"{past_tense} price increase for {n_ok} instance(s): {ids_str}")
if cutover_note:
print(cutover_note)
# Spec: non-stale failures take precedence over stale (so a script can
# distinguish "all stale" from "mixed failures").
if failed:
sys.exit(1)
if stale:
sys.exit(2)
# ---------------------------------------------------------------------------
# accept price-increase
# ---------------------------------------------------------------------------
@parser.command(
argument("ids", help="instance IDs to accept (one or more).", type=int, nargs="+"),
argument("--yes", "-y", action="store_true",
help="skip the interactive prompt; required when stdin is not a TTY"),
usage="vastai accept price-increase ID [ID ...] [--yes]",
help="Accept one or more pending host price increases",
epilog=deindent("""
Review pending price increases with `vastai show pending-price-increases`
before accepting. The CLI fans out one PUT per instance ID (no batch
endpoint exists on the backend), and the new rate applies only after
each contract's current end_date — your remaining time on the current
term is billed at the original price.
Examples:
vastai accept price-increase 123 --yes
vastai accept price-increase 1 2 3 --yes
Exit codes:
0 — every requested row was accepted
1 — at least one row failed (non-stale)
2 — at least one row was stale (re-run `show pending-price-increases`)
"""),
)
def accept__price_increase(args):
"""Per-row accept for pending price-increase challenges."""
return _run_per_row(
args,
route_verb="accept",
api_call=price_increase_api.accept,
present_imperative="Accept",
past_tense="Accepted",
cutover_note="New rate applies after each contract's current end_date.",
)
# ---------------------------------------------------------------------------
# reject price-increase
# ---------------------------------------------------------------------------
@parser.command(
argument("ids", help="instance IDs to reject (one or more).", type=int, nargs="+"),
argument("--yes", "-y", action="store_true",
help="skip the interactive prompt; required when stdin is not a TTY"),
usage="vastai reject price-increase ID [ID ...] [--yes]",
help="Reject one or more pending host price increases",
epilog=deindent("""
Tombstones the matching pending rows on the backend (no cutover follows).
Use `vastai show pending-price-increases` to review the rows first; the
CLI fans out one PUT per instance ID.
Exit codes mirror `accept price-increase`:
0 — every requested row was rejected
1 — at least one row failed (non-stale)
2 — at least one row was stale
"""),
)
def reject__price_increase(args):
"""Per-row reject for pending price-increase challenges."""
return _run_per_row(
args,
route_verb="reject",
api_call=price_increase_api.reject,
present_imperative="Reject",
past_tense="Rejected",
cutover_note=None,
)
cli/commands/repl.py
"""`vastai repl` — an interactive shell over the CLI's own commands."""
from vastai.cli.display import deindent
from vastai.cli.utils import get_parser as _get_parser
parser = _get_parser()
@parser.command(
usage="vastai repl",
help="Start an interactive vastai shell",
epilog=deindent("""
Runs vastai commands from one long-lived prompt, with Tab completion
over commands, flags and your live instance ids:
vast> show instances
vast> destroy instance <tab>
Commands behave exactly as they do from the shell — same auth, same
flags, same output.
"""),
)
def repl(args):
from vastai.cli.repl.session import run_repl
run_repl(args)
cli/commands/storage.py
"""CLI commands for managing volumes, storage, copies, and cloud sync."""
import json
import os
import time
import subprocess
from urllib.parse import quote, urlparse
from vastai.cli.parser import argument
from vastai.cli.display import (
display_table, vol_displayable_fields, nw_vol_displayable_fields,
volume_fields, connection_fields, deindent,
instance_backup_fields, instance_backup_file_fields,
)
from vastai.cli.display import strip_strings
from vastai.cli.util import (
default_start_date,
parse_day_cron_style, parse_hour_cron_style,
validate_frequency_values, add_scheduled_job,
)
from vastai.api import storage as storage_api
from vastai.api import offers as offers_api
from vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401
parser = _get_parser()
# ---------------------------------------------------------------------------
# copy / cancel copy / cancel sync
# ---------------------------------------------------------------------------
@parser.command(
argument("src", help="Source location for copy operation (supports multiple formats)", type=str),
argument("dst", help="Target location for copy operation (supports multiple formats)", type=str),
argument("-i", "--identity", help="Location of ssh private key", type=str),
usage="vastai copy SRC DST",
help="Copy directories between instances and/or local",
epilog=deindent("""
Copies a directory from a source location to a target location. Each of source and destination
directories can be either local or remote, subject to appropriate read and write
permissions required to carry out the action.
Supported location formats:
- [instance_id:]path (legacy format, still supported)
- C.instance_id:path (container copy format)
- cloud_service:path (cloud service format)
- cloud_service.cloud_service_id:path (cloud service with ID)
- local:path (explicit local path)
- V.volume_id:path (volume copy)
You should not copy to /root or / as a destination directory, as this can mess up the permissions on your instance ssh folder, breaking future copy operations (as they use ssh authentication)
You can see more information about constraints here: https://vast.ai/docs/gpu-instances/data-movement#constraints
Volume copy is currently only supported for copying to other volumes, instances, or cloud services, not local.
Examples:
vastai copy 6003036:/workspace/ 6003038:/workspace/
vastai copy C.11824:/data/test local:data/test
vastai copy local:data/test C.11824:/data/test
vastai copy drive:/folder/file.txt C.6003036:/workspace/
vastai copy s3.101:/data/ C.6003036:/workspace/
vastai copy hf.101:/my-bucket/data/ C.6003036:/workspace/
vastai copy V.1234:/file C.5678:/workspace/
vastai copy V.1234:/file s3.101:/workspace/
The first example copy syncs all files from the absolute directory '/workspace' on instance 6003036 to the directory '/workspace' on instance 6003038.
The second example copy syncs files from container 11824 to the local machine using structured syntax.
The third example copy syncs files from local to container 11824 using structured syntax.
The fourth example copy syncs files from Google Drive to an instance.
The fifth example copy syncs files from S3 bucket with id 101 to an instance.
The sixth example copy syncs files from a Hugging Face bucket using connection id 101 to an instance.
The seventh example copy syncs files from volume 1234 to instance 5678.
The eighth example copy syncs files from volume 1234 to S3 bucket with id 101.
"""),
)
def copy(args):
"""Transfer data from one instance to another."""
from vastai.utils import parse_vast_url
client = get_client(args)
(src_id, src_path) = parse_vast_url(args.src)
(dst_id, dst_path) = parse_vast_url(args.dst)
print(f"copying {str(src_id)+':' if src_id else ''}{src_path} {str(dst_id)+':' if dst_id else ''}{dst_path}")
if args.explain:
print("request json: ")
print({"client_id": "me", "src_id": src_id, "dst_id": dst_id, "src_path": src_path, "dst_path": dst_path})
rj = storage_api.copy(client, src_id=src_id, dst_id=dst_id, src_path=src_path, dst_path=dst_path)
if (rj.get("success")) and ((src_id is None or src_id == "local") or (dst_id is None or dst_id == "local")):
ssh_args = ["ssh"]
if args.identity is not None:
ssh_args.extend(["-i", args.identity])
if (src_id is None or src_id == "local"):
remote_port = rj["dst_port"]
remote_addr = rj["dst_addr"]
ssh_cmd = " ".join(ssh_args + ["-p", str(remote_port), "-o", "StrictHostKeyChecking=no"])
cmd = [
"rsync", "-arz", "-v", "--progress", "--rsh=ssh", "-e", ssh_cmd,
src_path, f"vastai_kaalia@{remote_addr}::{dst_id}/{dst_path}",
]
print(" ".join(cmd))
subprocess.run(cmd)
elif (dst_id is None or dst_id == "local"):
os.makedirs(dst_path, exist_ok=True)
remote_port = rj["src_port"]
remote_addr = rj["src_addr"]
ssh_cmd = " ".join(ssh_args + ["-p", str(remote_port), "-o", "StrictHostKeyChecking=no"])
cmd = [
"rsync", "-arz", "-v", "--progress", "--rsh=ssh", "-e", ssh_cmd,
f"vastai_kaalia@{remote_addr}::{src_id}/{src_path}", dst_path,
]
print(" ".join(cmd))
subprocess.run(cmd)
else:
if rj.get("success"):
print("Remote to Remote copy initiated - check instance status bar for progress updates (~30 seconds delayed).")
else:
if rj.get("msg") == "src_path not supported VMs.":
print("copy between VM instances does not currently support subpaths (only full disk copy)")
elif rj.get("msg") == "dst_path not supported for VMs.":
print("copy between VM instances does not currently support subpaths (only full disk copy)")
else:
print(rj.get("msg", rj))
@parser.command(
argument("dst", help="instance_id:/path to target of copy operation", type=str),
usage="vastai cancel copy DST",
help="Cancel a remote copy in progress, specified by DST id",
epilog=deindent("""
Use this command to cancel any/all current remote copy operations copying to a specific named instance, given by DST.
Examples:
vastai cancel copy 12371
The first example cancels all copy operations currently copying data into instance 12371
"""),
)
def cancel__copy(args):
"""Cancel a remote copy in progress, specified by DST id."""
client = get_client(args)
dst_id = args.dst
if dst_id is None:
print("invalid arguments")
return
print(f"canceling remote copies to {dst_id} ")
rj = storage_api.cancel_copy(client, dst_id=dst_id)
if rj.get("success"):
print("Remote copy canceled - check instance status bar for progress updates (~30 seconds delayed).")
else:
print(rj.get("msg", rj))
@parser.command(
argument("dst", help="instance_id:/path to target of sync operation", type=str),
usage="vastai cancel sync DST",
help="Cancel a remote copy in progress, specified by DST id",
epilog=deindent("""
Use this command to cancel any/all current remote cloud sync operations copying to a specific named instance, given by DST.
Examples:
vastai cancel sync 12371
The first example cancels all copy operations currently copying data into instance 12371
"""),
)
def cancel__sync(args):
"""Cancel a remote cloud sync in progress, specified by DST id."""
client = get_client(args)
dst_id = args.dst
if dst_id is None:
print("invalid arguments")
return
print(f"canceling remote copies to {dst_id} ")
rj = storage_api.cancel_sync(client, dst_id=dst_id)
if rj.get("success"):
print("Remote copy canceled - check instance status bar for progress updates (~30 seconds delayed).")
else:
print(rj.get("msg", rj))
# ---------------------------------------------------------------------------
# cloud copy
# ---------------------------------------------------------------------------
@parser.command(
argument("--src", help="path to source of object to copy", type=str),
argument("--dst", help="path to target of copy operation", type=str, default="/workspace"),
argument("--instance", help="id of the instance", type=str),
argument("--connection", help="id of cloud connection on your account (get from calling 'vastai show connections')", type=str),
argument("--transfer", help="type of transfer, possible options include Instance To Cloud and Cloud To Instance", type=str, default="Instance to Cloud"),
argument("--dry-run", help="show what would have been transferred", action="store_true"),
argument("--size-only", help="skip based on size only, not mod-time or checksum", action="store_true"),
argument("--ignore-existing", help="skip all files that exist on destination", action="store_true"),
argument("--update", help="skip files that are newer on the destination", action="store_true"),
argument("--delete-excluded", help="delete files on dest excluded from transfer", action="store_true"),
argument("--schedule", choices=["HOURLY", "DAILY", "WEEKLY"], help="try to schedule a command to run hourly, daily, or weekly. Valid values are HOURLY, DAILY, WEEKLY For ex. --schedule DAILY"),
argument("--start_date", type=str, default=default_start_date(), help="Start date/time in format 'YYYY-MM-DD HH:MM:SS PM' (UTC). Default is now. (optional)"),
argument("--end_date", type=str, help="End date/time in format 'YYYY-MM-DD HH:MM:SS PM' (UTC). Default is contract's end. (optional)"),
argument("--day", type=parse_day_cron_style, help="Day of week you want scheduled job to run on (0-6, where 0=Sunday) or \"*\". Default will be 0. For ex. --day 0", default=0),
argument("--hour", type=parse_hour_cron_style, help="Hour of day you want scheduled job to run on (0-23) or \"*\" (UTC). Default will be 0. For ex. --hour 16", default=0),
usage="vastai cloud copy --src SRC --dst DST --instance INSTANCE_ID --connection CONNECTION_ID --transfer TRANSFER_TYPE",
help="Copy files/folders to and from cloud providers",
epilog=deindent("""
Copies a directory from a source location to a target location. Each of source and destination
directories can be either local or remote, subject to appropriate read and write
permissions required to carry out the action. The format for both src and dst is [instance_id:]path.
You can find more information about the cloud copy operation here: https://vast.ai/docs/gpu-instances/cloud-sync
Examples:
vastai show connections
ID NAME Cloud Type
1001 test_dir drive
1003 data_dir drive
1004 hf_data hf
vastai cloud copy --src /folder --dst /workspace --instance 6003036 --connection 1001 --transfer "Instance To Cloud"
vastai cloud copy --src my-bucket/data --dst /workspace --instance 6003036 --connection 1004 --transfer "Cloud To Instance"
The first cloud copy example copies all contents of /folder on instance 6003036 into /workspace in gdrive connection 'test_dir'.
The second cloud copy example copies my-bucket/data from Hugging Face connection 'hf_data' into /workspace on the instance.
Note: Hugging Face cloud destinations must already exist as bucket paths; directories are not created automatically.
"""),
)
def cloud__copy(args):
"""Transfer data to/from cloud providers."""
client = get_client(args)
if (args.src is None) and (args.dst is None):
print("invalid arguments")
return
flags = storage_api.rclone_flags(
dry_run=args.dry_run, size_only=args.size_only,
ignore_existing=args.ignore_existing, update=args.update,
delete_excluded=args.delete_excluded,
)
print(f"copying {args.src} {args.dst} {args.instance} {args.connection} {args.transfer}")
req_json = {
"src": args.src,
"dst": args.dst,
"instance_id": args.instance,
"selected": args.connection,
"transfer": args.transfer,
"flags": flags
}
if args.explain:
print("request json: ")
print(req_json)
if args.schedule:
validate_frequency_values(args.day, args.hour, args.schedule)
r = client.get(f"/instances/{args.instance}/", query_args={"owner": "me"})
r.raise_for_status()
row = r.json()["instances"]
if args.transfer.lower() == "instance to cloud":
if row:
up_cost = row.get("internet_up_cost_per_tb", None)
if up_cost is not None:
confirm = input(
f"Internet upload cost is ${up_cost} per TB. "
"Are you sure you want to schedule a cloud backup? (y/n): "
).strip().lower()
if confirm != "y":
print("Cloud backup scheduling aborted.")
return
else:
print("Warning: Could not retrieve internet upload cost. Proceeding without confirmation. You can use show scheduled-jobs and delete scheduled-job commands to delete scheduled cloud backup job.")
cli_command = "cloud copy"
api_endpoint = "/api/v0/commands/rclone/"
contract_end_date = row.get("end_date", None)
add_scheduled_job(client, args, req_json, cli_command, api_endpoint, "POST", instance_id=args.instance, contract_end_date=contract_end_date)
return
else:
print("Instance not found. Please check the instance ID.")
return
rj = storage_api.cloud_copy(client, src=args.src, dst=args.dst, instance=args.instance,
connection=args.connection, transfer=args.transfer, flags=flags)
if rj:
print("Cloud Copy Started - check instance status bar for progress updates (~30 seconds delayed).")
print("When the operation is finished you should see 'Cloud Copy Operation Finished' in the instance status bar.")
# ---------------------------------------------------------------------------
# show connections
# ---------------------------------------------------------------------------
@parser.command(
usage="vastai show connections [--api-key API_KEY] [--raw]",
help="Display user's cloud connections",
)
def show__connections(args):
"""Show the user's cloud connections."""
client = get_client(args)
rows = storage_api.show_connections(client)
if args.raw:
return rows
else:
display_table(rows, connection_fields)
# ---------------------------------------------------------------------------
# Note: show__network_disks and add__network_disk are in machines.py.
# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------
# search volumes
# ---------------------------------------------------------------------------
@parser.command(
argument("-n", "--no-default", action="store_true", help="Disable default query"),
argument("--limit", type=int, help=""),
argument("--storage", type=float, default=1.0, help="Amount of storage to use for pricing, in GiB. default=1.0GiB"),
argument("-o", "--order", type=str, help="Comma-separated list of fields to sort on. postfix field with - to sort desc. ex: -o 'disk_space,inet_up-'. default='score-'", default='score-'),
argument("query", help="Query to search for. default: 'external=false verified=true disk_space>=1', pass -n to ignore default", nargs="*", default=None),
usage="vastai search volumes [--help] [--api-key API_KEY] [--raw] <query>",
help="Search for volume offers using custom query",
epilog=deindent("""
Query syntax:
query = comparison comparison...
comparison = field op value
field = <name of a field>
op = one of: <, <=, ==, !=, >=, >, in, notin
value = <bool, int, float, string> | 'any' | [value0, value1, ...]
bool: True, False
note: to pass '>' and '<' on the command line, make sure to use quotes
note: to encode a string query value (ie for gpu_name), replace any spaces ' ' with underscore '_'
Examples:
# search for volumes with greater than 50GB of available storage and greater than 500 Mb/s upload and download speed
vastai search volumes "disk_space>50 inet_up>500 inet_down>500"
Available fields:
Name Type Description
cpu_arch: string host machine cpu architecture (e.g. amd64, arm64)
cuda_vers: float machine max supported cuda version (based on driver version)
datacenter: bool show only datacenter offers
disk_bw: float disk read bandwidth, in MB/s
disk_space: float disk storage space, in GB
driver_version: string machine's nvidia/amd driver version as 3 digit string ex. "535.86.05"
duration: float max rental duration in days
geolocation: string Two letter country code. Works with operators =, !=, in, notin (e.g. geolocation not in ['XV','XZ'])
gpu_arch: string host machine gpu architecture (e.g. nvidia, amd)
gpu_name: string GPU model name (no quotes, replace spaces with underscores, ie: RTX_3090 rather than 'RTX 3090')
has_avx: bool CPU supports AVX instruction set.
id: int volume offer unique ID
inet_down: float internet download speed in Mb/s
inet_up: float internet upload speed in Mb/s
machine_id: int machine id of volume offer
pci_gen: float PCIE generation
pcie_bw: float PCIE bandwidth (CPU to GPU)
reliability: float machine reliability score (see FAQ for explanation)
storage_cost: float storage cost in $/GB/month
static_ip: bool is the IP addr static/stable
total_flops: float total TFLOPs from all GPUs
ubuntu_version: string host machine ubuntu OS version
verified: bool is the machine verified
"""),
)
def search__volumes(args):
"""Search for volume offers using custom query."""
from vastai.api.query import parse_query, vol_offers_fields, offers_alias, offers_mult
try:
if args.no_default:
query = {}
else:
query = {"verified": {"eq": True}, "external": {"eq": False}, "disk_space": {"gte": 1}}
if args.query is not None:
query = parse_query(args.query, query, vol_offers_fields, {}, offers_mult)
order = []
for name in args.order.split(","):
name = name.strip()
if not name:
continue
direction = "asc"
field = name
if name.strip("-") != name:
direction = "desc"
field = name.strip("-")
if name.strip("+") != name:
direction = "asc"
field = name.strip("+")
if field in offers_alias:
field = offers_alias[field]
order.append([field, direction])
except ValueError as e:
print("Error: ", e)
return 1
client = get_client(args)
if args.explain:
print("request json: ")
print(query)
rows = offers_api.search_volumes(
client, query=query, order=order,
limit=args.limit, storage=args.storage,
no_default=True, # defaults already applied above
)
if args.raw:
return rows
else:
display_table(rows, vol_displayable_fields)
# ---------------------------------------------------------------------------
# search network volumes
# ---------------------------------------------------------------------------
@parser.command(
argument("-n", "--no-default", action="store_true", help="Disable default query"),
argument("--limit", type=int, help=""),
argument("--storage", type=float, default=1.0, help="Amount of storage to use for pricing, in GiB. default=1.0GiB"),
argument("-o", "--order", type=str, help="Comma-separated list of fields to sort on. postfix field with - to sort desc. ex: -o 'disk_space,inet_up-'. default='score-'", default='score-'),
argument("query", help="Query to search for. default: 'external=false verified=true disk_space>=1', pass -n to ignore default", nargs="*", default=None),
usage="vastai search network volumes [--help] [--api-key API_KEY] [--raw] <query>",
help="Search for network volume offers using custom query",
epilog=deindent("""
Query syntax:
query = comparison comparison...
comparison = field op value
field = <name of a field>
op = one of: <, <=, ==, !=, >=, >, in, notin
value = <bool, int, float, string> | 'any' | [value0, value1, ...]
bool: True, False
note: to pass '>' and '<' on the command line, make sure to use quotes
note: to encode a string query value (ie for gpu_name), replace any spaces ' ' with underscore '_'
Examples:
# search for volumes with greater than 50GB of available storage and greater than 500 Mb/s upload and download speed
vastai search volumes "disk_space>50 inet_up>500 inet_down>500"
Available fields:
Name Type Description
duration: float max rental duration in days
geolocation: string Two letter country code. Works with operators =, !=, in, notin (e.g. geolocation not in ['XV','XZ'])
id: int volume offer unique ID
inet_down: float internet download speed in Mb/s
inet_up: float internet upload speed in Mb/s
reliability: float machine reliability score (see FAQ for explanation)
storage_cost: float storage cost in $/GB/month
verified: bool is the machine verified
"""),
)
def search__network_volumes(args):
"""Search for network volume offers using custom query."""
from vastai.api.query import parse_query, vol_offers_fields, offers_alias, offers_mult
try:
if args.no_default:
query = {}
else:
query = {"verified": {"eq": True}, "external": {"eq": False}, "disk_space": {"gte": 1}}
if args.query is not None:
query = parse_query(args.query, query, vol_offers_fields, {}, offers_mult)
order = []
for name in args.order.split(","):
name = name.strip()
if not name:
continue
direction = "asc"
field = name
if name.strip("-") != name:
direction = "desc"
field = name.strip("-")
if name.strip("+") != name:
direction = "asc"
field = name.strip("+")
if field in offers_alias:
field = offers_alias[field]
order.append([field, direction])
except ValueError as e:
print("Error: ", e)
return 1
client = get_client(args)
if args.explain:
print("request json: ")
print(query)
rows = offers_api.search_network_volumes(
client, query=query, order=order,
limit=args.limit, storage=args.storage,
no_default=True, # defaults already applied above
)
if args.raw:
return rows
else:
display_table(rows, nw_vol_displayable_fields)
# ---------------------------------------------------------------------------
# create / delete / clone volume
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of volume offer", type=int),
argument("-s", "--size",
help="size in GB of volume. Default %(default)s GB.", default=15, type=float),
argument("-n", "--name", help="Optional name of volume.", type=str),
usage="vastai create volume ID [options]",
help="Create a new volume",
epilog=deindent("""
Creates a volume from an offer ID (which is returned from "search volumes"). Each offer ID can be used to create multiple volumes,
provided the size of all volumes does not exceed the size of the offer.
"""),
)
def create__volume(args):
"""Create a new volume from an offer ID."""
if args.explain:
print("request json: ")
print({"size": int(args.size), "id": int(args.id)})
client = get_client(args)
result = storage_api.create_volume(client, id=args.id, size=args.size, name=args.name)
if args.raw:
return result
else:
print("Created. {}".format(result))
@parser.command(
argument("id", help="id of network volume offer", type=int),
argument("-s", "--size",
help="size in GB of network volume. Default %(default)s GB.", default=15, type=float),
argument("-n", "--name", help="Optional name of network volume.", type=str),
usage="vastai create network volume ID [options]",
help="Create a new network volume",
epilog=deindent("""
Creates a network volume from an offer ID (which is returned from "search network volumes"). Each offer ID can be used to create multiple volumes,
provided the size of all volumes does not exceed the size of the offer.
"""),
)
def create__network_volume(args):
"""Create a new network volume from an offer ID."""
if args.explain:
print("request json: ")
print({"size": int(args.size), "id": int(args.id)})
client = get_client(args)
result = storage_api.create_network_volume(client, id=args.id, size=args.size, name=args.name)
if args.raw:
return result
else:
print("Created. {}".format(result))
@parser.command(
argument("id", help="id of volume contract", type=int),
usage="vastai delete volume ID",
help="Delete a volume",
epilog=deindent("""
Deletes volume with the given ID. All instances using the volume must be destroyed before the volume can be deleted.
"""),
)
def delete__volume(args):
"""Delete a volume."""
client = get_client(args)
result = storage_api.delete_volume(client, id=args.id)
if args.raw:
return result
else:
print("Deleted. {}".format(result))
@parser.command(
argument("source", help="id of volume contract being cloned", type=int),
argument("dest", help="id of volume offer volume is being copied to", type=int),
argument("-s", "--size", help="Size of new volume contract, in GB. Must be greater than or equal to the source volume, and less than or equal to the destination offer.", type=float),
argument("-d", "--disable_compression", action="store_true", help="Do not compress volume data before copying."),
usage="vastai clone volume <source_id> <dest_id> [options]",
help="Clone an existing volume",
epilog=deindent("""
Create a new volume with the given offer, by copying the existing volume.
Size defaults to the size of the existing volume, but can be increased if there is available space.
"""),
)
def clone__volume(args):
"""Clone an existing volume."""
if args.explain:
print("request json: ")
print({"src_id": args.source, "dst_id": args.dest})
client = get_client(args)
result = storage_api.clone_volume(client, source=args.source, dest=args.dest,
size=args.size, disable_compression=args.disable_compression)
if args.raw:
return result
else:
print("Created. {}".format(result))
# ---------------------------------------------------------------------------
# show volumes
# ---------------------------------------------------------------------------
@parser.command(
argument("-t", "--type", help="volume type to display. Default to all. Possible values are \"local\", \"all\", \"network\"", type=str, default="all"),
usage="vastai show volumes [OPTIONS]",
help="Show stats on owned volumes.",
epilog=deindent("""
Show stats on owned volumes
"""),
)
def show__volumes(args):
"""Show stats on owned volumes."""
client = get_client(args)
processed = storage_api.show_volumes(client, type=args.type)
for row in processed:
for k, v in row.items():
row[k] = strip_strings(v)
if args.raw:
return processed
else:
display_table(processed, volume_fields, replace_spaces=False)
# ---------------------------------------------------------------------------
# show instance-backups / show instance-backup-files
# ---------------------------------------------------------------------------
def _authorized_url(url, token):
"""Fold the download token into `url` as a query parameter.
B2 takes the token either in the Authorization header or as an `Authorization`
query parameter, and the second form is what makes a URL usable on its own -- paste
it into a browser, or curl it without -H. Returns the url untouched when either
piece is missing, so a malformed payload still lists its files.
"""
if not url or not token:
return url
sep = "&" if urlparse(url).query else "?"
return "{}{}Authorization={}".format(url, sep, quote(token, safe=""))
@parser.command(
argument("--contract-id", help="only show backups of this instance or volume contract id", type=int),
usage="vastai show instance-backups [OPTIONS]",
help="Show backups Vast has taken of your instances.",
epilog=deindent("""
Lists the backups Vast has taken of your instances and volumes, newest first. These
are created by Vast (for example before a machine is decommissioned), not by you.
Status is pending while the copy runs, then success or error. A backup outlives the
instance it came from, so entries stay listed after the instance is gone.
Restore one into a running instance with `vastai copy`, naming vast.CONTRACT_ID as
the source. An instance's own filesystem sits under upper/:
vastai show instance-backups
vastai copy vast.12345678:/upper/workspace 99887766:/workspace
A backed-up volume has no upper/ -- its contents sit at the root of its own
contract id:
vastai copy vast.12345679:/ 99887766:/volumes/data
To pull the objects down directly instead of restoring into an instance, see
`vastai show instance-backup-files CONTRACT_ID`.
Example: vastai show instance-backups --contract-id 12345678
"""),
)
def show__instance_backups(args):
"""Show backups Vast has taken of your instances."""
client = get_client(args)
rows = storage_api.show_instance_backups(client, contract_id=args.contract_id)
if args.raw:
return rows
now = time.time()
for row in rows:
row["age_hours"] = (now - (row.get("created_at") or now)) / 3600.0
# A backup whose contract has been archived out of the live tables reports no type.
row["type"] = row.get("type") or "unknown"
row["info"] = strip_strings(row.get("info") or "")
display_table(rows, instance_backup_fields, replace_spaces=False)
@parser.command(
argument("contract_id", help="instance or volume contract id whose backup to list", type=int),
usage="vastai show instance-backup-files CONTRACT_ID [OPTIONS]",
help="List the stored objects of one backup, with direct download URLs.",
epilog=deindent("""
Lists the objects stored for one backup. Each URL already carries a download token
scoped to this backup, so it works on its own -- open it in a browser, or:
curl -o model.safetensors "<url>"
The token is also printed on its own for the header form, `curl -H "Authorization:
<token>"`, which keeps it out of shell history and proxy logs. Either way it expires
(the lifetime is printed), so re-run this command rather than saving the URLs.
If the listing reports itself truncated, the backup holds more objects than are
shown; restore with `vastai copy vast.CONTRACT_ID:/ INSTANCE_ID:/path` instead of
fetching file by file.
Example: vastai show instance-backup-files 12345678
"""),
)
def show__instance_backup_files(args):
"""List the stored objects of one backup, with direct download URLs."""
client = get_client(args)
blob = storage_api.instance_backup_files(client, args.contract_id)
if args.raw:
return blob
files = blob.get("files") or []
# --raw stays a passthrough of the API payload; only the table gets ready-to-click URLs.
token = blob.get("authorization_token")
for row in files:
row["size_mb"] = (row.get("size") or 0) / 1e6
row["url"] = _authorized_url(row.get("url"), token)
print("authorization token: {} (already embedded in the URLs below)".format(token))
print("expires in: {} seconds".format(blob.get("expires_in_sec")))
if blob.get("truncated"):
print("WARNING: listing truncated -- this backup holds more objects than shown below.")
if not files:
print("No objects stored for this backup yet.")
return
display_table(files, instance_backup_file_fields, replace_spaces=False)
# ---------------------------------------------------------------------------
# list / unlist volumes (host)
# ---------------------------------------------------------------------------
@parser.command(
argument("id", help="id of machine to list", type=int),
argument("-p", "--price_disk",
help="storage price in $/GB/month, default: $%(default).2f/GB/month", default=.10, type=float),
argument("-e", "--end_date", help="contract offer expiration - the available until date (optional, in unix float timestamp or MM/DD/YYYY format), default 3 months", type=str),
argument("-s", "--size", help="size of disk space allocated to offer in GB, default %(default)s GB", default=15),
usage="vastai list volume ID [options]",
help="[Host] list disk space for rent as a volume on a machine",
epilog=deindent("""
Allocates a section of disk on a machine to be used for volumes.
"""),
)
def list__volume(args):
"""List disk space for rent as a volume on a machine."""
from vastai.cli.util import string_to_unix_epoch
end_date = string_to_unix_epoch(args.end_date) if args.end_date else None
if args.explain:
print("request json: ")
print({"size": int(args.size), "machine": int(args.id), "price_disk": float(args.price_disk)})
client = get_client(args)
result = storage_api.list_volume(client, id=args.id, size=int(args.size),
price_disk=float(args.price_disk), end_date=end_date)
if args.raw:
return result
else:
print("Created. {}".format(result))
@parser.command(
argument("ids", help="id of machines list", type=int, nargs='+'),
argument("-p", "--price_disk",
help="storage price in $/GB/month, default: $%(default).2f/GB/month", default=.10, type=float),
argument("-e", "--end_date", help="contract offer expiration - the available until date (optional, in unix float timestamp or MM/DD/YYYY format), default 3 months", type=str),
argument("-s", "--size", help="size of disk space allocated to offer in GB, default %(default)s GB", default=15),
usage="vastai list volume IDs [options]",
help="[Host] list disk space for rent as a volume on machines",
epilog=deindent("""
Allocates a section of disk on machines to be used for volumes.
"""),
)
def list__volumes(args):
"""List disk space for rent as a volume on multiple machines."""
from vastai.cli.util import string_to_unix_epoch
end_date = string_to_unix_epoch(args.end_date) if args.end_date else None
if args.explain:
print("request json: ")
print({"size": int(args.size), "machine": [int(id) for id in args.ids], "price_disk": float(args.price_disk)})
client = get_client(args)
result = storage_api.list_volumes(client, ids=args.ids, size=int(args.size),
price_disk=float(args.price_disk), end_date=end_date)
if args.raw:
return result
else:
print("Created. {}".format(result))
@parser.command(
argument("disk_id", help="id of network disk to list", type=int),
argument("-p", "--price_disk", help="storage price in $/GB/month, default: $%(default).2f/GB/month", default=.15, type=float),
argument("-e", "--end_date", help="contract offer expiration - the available until date (optional, in unix float timestamp or MM/DD/YYYY format), default 1 month", type=str, default=None),
argument("-s", "--size", help="size of disk space allocated to offer in GB, default %(default)s GB", default=15, type=int),
usage="vastai list network volume DISK_ID [options]",
help="[Host] list disk space for rent as a network volume",
)
def list__network_volume(args):
"""List disk space for rent as a network volume."""
from vastai.cli.util import string_to_unix_epoch
end_date = string_to_unix_epoch(args.end_date) if args.end_date else None
if args.explain:
print("request json: ")
print({"disk_id": args.disk_id, "price_disk": args.price_disk, "size": args.size})
client = get_client(args)
result = storage_api.list_network_volume(client, disk_id=args.disk_id, price_disk=args.price_disk,
size=args.size, end_date=end_date)
if args.raw:
return result
print(result["msg"])
@parser.command(
argument("id", help="volume ID you want to unlist", type=int),
usage="vastai unlist volume ID",
help="[Host] unlist volume offer",
)
def unlist__volume(args):
"""Unlist a volume offer."""
if args.explain:
print("request json:", {"id": args.id})
client = get_client(args)
result = storage_api.unlist_volume(client, id=args.id)
if args.raw:
return result
else:
print(result["msg"])
@parser.command(
argument("id", help="id of network volume offer to unlist", type=int),
usage="vastai unlist network volume OFFER_ID",
help="[Host] Unlists network volume offer",
)
def unlist__network_volume(args):
"""Unlist a network volume offer."""
if args.explain:
print("request json: ")
print({"id": args.id})
client = get_client(args)
result = storage_api.unlist_network_volume(client, id=args.id)
if args.raw:
return result
print(result["msg"])
cli/commands/teams.py
"""CLI commands for managing teams."""
import json
from vastai.cli.parser import argument
from vastai.cli.display import deindent
from vastai.api import teams as teams_api
from vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401
parser = _get_parser()
# ---------------------------------------------------------------------------
# create / destroy team
# ---------------------------------------------------------------------------
@parser.command(
argument("--team-name", help="name of the team", type=str),
argument("--transfer-credit", help="amount of personal credit to transfer to the new team", type=float, default=0),
usage="vastai create-team --team-name TEAM_NAME [--transfer-credit AMOUNT]",
help="Create a new team",
epilog=deindent("""
Creates a new team under your account.
Unlike legacy teams, this command does NOT convert your personal account into a team.
Each team is created as a separate account, and you can be a member of multiple teams.
When you create a team:
- You become the team owner.
- The team starts as an independent account with its own billing, credits, and resources.
- Default roles (owner, manager, member) are automatically created.
- You can invite others, assign roles, and manage resources within the team.
Optional:
You can transfer a portion of your existing personal credits to the team by using
the `--transfer-credit` flag. Example:
vastai create-team --team-name myteam --transfer-credit 25
Notes:
- You cannot create a team from within another team account.
- Your API key keeps acting on your personal account. To work in the team,
switch to it in the console and create an API key there.
For more details, see:
https://vast.ai/docs/teams-quickstart
"""),
)
def create__team(args):
"""Create a new team."""
client = get_client(args)
result = teams_api.create_team(client, team_name=args.team_name, transfer_credit=args.transfer_credit)
print(result)
@parser.command(
usage="vastai destroy team",
help="Destroy your team",
)
def destroy__team(args):
"""Destroy your team."""
client = get_client(args)
result = teams_api.destroy_team(client)
print(result)
# ---------------------------------------------------------------------------
# team roles
# ---------------------------------------------------------------------------
@parser.command(
argument("--name", help="name of the role", type=str),
argument("--permissions", help="file path for json encoded permissions, look in the docs for more information", type=str),
usage="vastai create team-role --name NAME --permissions PERMISSIONS",
help="Add a new role to your team",
epilog=deindent("""
Creating a new team role involves understanding how permissions must be sent via json format.
You can find more information about permissions here: https://vast.ai/docs/cli/roles-and-permissions
""")
)
def create__team_role(args):
"""Create a new team role."""
from vastai.cli.util import load_permissions_from_file
client = get_client(args)
permissions = load_permissions_from_file(args.permissions)
result = teams_api.create_team_role(client, name=args.name, permissions=permissions)
print(result)
@parser.command(
argument("NAME", help="name of the role", type=str),
usage="vastai show team-role NAME",
help="Show your team role",
)
def show__team_role(args):
"""Show a specific team role."""
client = get_client(args)
result = teams_api.show_team_role(client, name=args.NAME)
print(json.dumps(result, indent=1, sort_keys=True))
@parser.command(
usage="vastai show team-roles",
help="Show roles for a team",
)
def show__team_roles(args):
"""Show all team roles."""
client = get_client(args)
result = teams_api.show_team_roles(client)
if args.raw:
return result
else:
print(result)
@parser.command(
argument("id", help="id of the role", type=int),
argument("--name", help="name of the role", type=str),
argument("--permissions", help="file path for json encoded permissions, look in the docs for more information", type=str),
usage="vastai update team-role ID --name NAME --permissions PERMISSIONS",
help="Update an existing team role",
)
def update__team_role(args):
"""Update an existing team role."""
from vastai.cli.util import load_permissions_from_file
client = get_client(args)
permissions = load_permissions_from_file(args.permissions)
result = teams_api.update_team_role(client, id=args.id, name=args.name, permissions=permissions)
if args.raw:
return result
else:
print(json.dumps(result, indent=1))
@parser.command(
argument("NAME", help="name of the role", type=str),
usage="vastai remove team-role NAME",
help="Remove a role from your team",
)
def remove__team_role(args):
"""Remove a team role."""
client = get_client(args)
result = teams_api.remove_team_role(client, name=args.NAME)
print(result)
# ---------------------------------------------------------------------------
# team members
# ---------------------------------------------------------------------------
@parser.command(
argument("--email", help="email of user to be invited", type=str),
argument("--role", help="role of user to be invited", type=str),
usage="vastai invite member --email EMAIL --role ROLE",
help="Invite a team member",
)
def invite__member(args):
"""Invite a team member."""
client = get_client(args)
result = teams_api.invite_member(client, email=args.email, role=args.role)
if result:
print(f"successfully invited {args.email} to your current team")
@parser.command(
usage="vastai show members",
help="Show your team members",
)
def show__members(args):
"""Show team members."""
client = get_client(args)
result = teams_api.show_members(client)
if args.raw:
return result
else:
print(result)
@parser.command(
argument("id", help="id of user to remove", type=int),
usage="vastai remove member ID",
help="Remove a team member",
)
def remove__member(args):
"""Remove a team member."""
client = get_client(args)
result = teams_api.remove_member(client, id=args.id)
print(result)
cli/commands/uninstall.py
"""`vastai uninstall` — remove a managed (curl | bash installer) install.
pip installs aren't managed by the installer; for those this command only
prints the correct pip uninstall instructions and exits non-zero — it never
shells out to pip, matching `update`'s same-method discipline.
"""
import sys
from vastai.cli.parser import argument
from vastai.cli.display import deindent
from vastai.cli.util import DIRS, VERSION
from vastai.cli.selfupdate import (
INSTALL_SH_HINT, UninstallError, is_managed_install, install_root, perform_uninstall,
)
from vastai.cli.utils import get_parser as _get_parser
parser = _get_parser()
PIP_UNINSTALL_HINT = "pip uninstall vastai"
@parser.command(
argument("-y", "--yes", action="store_true", help="skip the confirmation prompt"),
usage="vastai uninstall [-y]",
help="Remove the managed CLI install",
hidden=True, # still under testing — remove once ready to announce
epilog=deindent(f"""
Removes a CLI installed with the managed installer
({INSTALL_SH_HINT}): the install root and its symlinks in
~/.local/bin. Config in {DIRS['config']} (your API key) is left
untouched, so re-running the installer keeps you logged in.
For pip installs, uninstall with: {PIP_UNINSTALL_HINT}
"""),
)
def uninstall(args):
if not is_managed_install():
print(
"This CLI was not installed with the managed installer, so this "
f"command can't remove it.\n Installed via pip? Run: {PIP_UNINSTALL_HINT}",
file=sys.stderr,
)
return 1
root = install_root()
if not args.yes:
if not sys.stdin.isatty():
print(
"Refusing to uninstall without confirmation in a non-interactive "
"shell. Re-run with --yes to skip the prompt.",
file=sys.stderr,
)
return 1
try:
answer = input(f"Remove vastai {VERSION} from {root}? [y/N] ").strip().lower()
except EOFError:
answer = ""
if answer not in ("y", "yes"):
print("Aborted.")
return 1
try:
removed = perform_uninstall()
except UninstallError as e:
print(str(e), file=sys.stderr)
return 1
print(f"vastai {VERSION} uninstalled from {removed}.")
print(f"Config in {DIRS['config']} was left untouched.")
return 0
cli/commands/update.py
"""`vastai update` — self-update for managed (curl | bash installer) installs.
pip installs aren't managed by the installer; for those this command only
prints the correct pip upgrade instructions and exits non-zero — it never
shells out to pip, since the CLI cannot know which pip/venv owns its
environment.
"""
import sys
from vastai.cli.parser import argument
from vastai.cli.display import deindent
from vastai.cli.util import VERSION
from vastai.cli.selfupdate import (
INSTALL_SH_HINT, PIP_UPGRADE_HINT, UpdateError,
fetch_manifest, is_managed_install, is_newer, perform_update,
)
from vastai.cli.utils import get_parser as _get_parser
parser = _get_parser()
EXIT_STALE = 10 # `update --check` exit code when a newer version exists
@parser.command(
argument("--check", action="store_true", help="only check for a newer version; exit 0 if current, 10 if an update is available"),
argument("--version", dest="target_version", metavar="VERSION", help="install a specific version instead of the latest (also how you pin or roll back)"),
usage="vastai update [--check | --version VERSION]",
help="Update the CLI to the latest version",
epilog=deindent("""
Updates a CLI installed with the managed installer
(curl -fsSL https://vast.ai/install.sh | bash). The new version is
installed into a fresh environment and verified before it is swapped
in, so an interrupted update never leaves a broken CLI. To roll back a
bad release, pin the previous one: vastai update --version 1.2.3
For pip installs, update with: pip install --upgrade vastai
"""),
)
def update(args):
try:
if args.check:
latest = fetch_manifest()["latest"]
if is_newer(latest, VERSION):
print(f"Update available: {latest} (you have {VERSION})")
return EXIT_STALE
print(f"vastai {VERSION} is up to date (latest: {latest})")
return 0
if not is_managed_install():
print(
"This CLI was not installed with the managed installer, so it "
f"cannot self-update.\n Installed via pip? Run: {PIP_UPGRADE_HINT}\n"
f" Or switch to the managed install: {INSTALL_SH_HINT}",
file=sys.stderr,
)
return 1
manifest = fetch_manifest()
target = args.target_version or manifest["latest"]
current = VERSION
if target == current:
latest = manifest["latest"]
if args.target_version:
print(f"vastai {current} is already installed (latest available: {latest}).")
else:
print(f"vastai {current} is up to date (latest: {latest}).")
return 0
print(f"Updating vastai {current} -> {target} ...")
perform_update(target, manifest)
print(f"Done. vastai is now {target}.")
return 0
except UpdateError as e:
print(str(e), file=sys.stderr)
if not args.check:
print("The current install was left untouched.", file=sys.stderr)
return 1
cli/display.py
"""Display formatting and field definitions for the Vast.ai CLI."""
import re
import shutil
import subprocess
from datetime import datetime, timezone
from typing import Dict, Tuple
def strip_strings(value):
if isinstance(value, str):
return value.strip()
elif isinstance(value, dict):
return {k: strip_strings(v) for k, v in value.items()}
elif isinstance(value, list):
return [strip_strings(item) for item in value]
return value # Return as is if not a string, list, or dict
def translate_null_strings_to_blanks(d: Dict) -> Dict:
"""Map over a dict and translate any null string values into ' '.
Leave everything else as is. This is needed because you cannot add TableCell
objects with only a null string or the client crashes.
:param Dict d: dict of item values.
:rtype Dict:
"""
# Beware: locally defined function.
def translate_nulls(s):
if s == "":
return " "
return s
new_d = {k: translate_nulls(v) for k, v in d.items()}
return new_d
def deindent(message: str, add_separator: bool = True) -> str:
"""
Deindent a quoted string. Scans message and finds the smallest number of whitespace characters in any line and
removes that many from the start of every line.
:param str message: Message to deindent.
:rtype str:
"""
message = re.sub(r" *$", "", message, flags=re.MULTILINE)
indents = [len(x) for x in re.findall("^ *(?=[^ ])", message, re.MULTILINE) if len(x)]
a = min(indents)
message = re.sub(r"^ {," + str(a) + "}", "", message, flags=re.MULTILINE)
if add_separator:
# For help epilogs - cleanly separating extra help from options
line_width = min(150, shutil.get_terminal_size((80, 20)).columns)
message = "_"*line_width + "\n"*2 + message.strip() + "\n" + "_"*line_width
return message.strip()
# These are the fields that are displayed when a search is run
displayable_fields = (
# ("bw_nvlink", "Bandwidth NVLink", "{}", None, True),
("id", "ID", "{}", None, True),
("cuda_max_good", "CUDA", "{:0.1f}", None, True),
("num_gpus", "N", "{}x", None, False),
("gpu_name", "Model", "{}", None, True),
("cpu_arch", "arch", "{}", None, True),
("pcie_bw", "PCIE", "{:0.1f}", None, True),
("cpu_ghz", "cpu_ghz", "{:0.1f}", None, True),
("cpu_cores_effective", "vCPUs", "{:0.1f}", None, True),
("cpu_ram", "RAM", "{:0.1f}", lambda x: x / 1000, False),
("gpu_ram", "VRAM", "{:0.1f}", lambda x: x / 1000, False),
("disk_space", "Disk", "{:.0f}", None, True),
("dph_total", "$/hr", "{:0.4f}", None, True),
("dlperf", "DLP", "{:0.1f}", None, True),
("dlperf_per_dphtotal", "DLP/$", "{:0.2f}", None, True),
("score", "score", "{:0.1f}", None, True),
("driver_version", "NV Driver", "{}", None, True),
("inet_up", "Net_up", "{:0.1f}", None, True),
("inet_down", "Net_down", "{:0.1f}", None, True),
("reliability", "R", "{:0.1f}", lambda x: x * 100, True),
("duration", "Max_Days", "{:0.1f}", lambda x: x / (24.0 * 60.0 * 60.0), True),
("machine_id", "mach_id", "{}", None, True),
("verification", "status", "{}", None, True),
("host_id", "host_id", "{}", None, True),
("direct_port_count", "ports", "{}", None, True),
("geolocation", "country", "{}", None, True),
# ("direct_port_count", "Direct Port Count", "{}", None, True),
)
displayable_fields_reserved = (
# ("bw_nvlink", "Bandwidth NVLink", "{}", None, True),
("id", "ID", "{}", None, True),
("cuda_max_good", "CUDA", "{:0.1f}", None, True),
("num_gpus", "N", "{}x", None, False),
("gpu_name", "Model", "{}", None, True),
("cpu_arch", "arch", "{}", None, True),
("pcie_bw", "PCIE", "{:0.1f}", None, True),
("cpu_ghz", "cpu_ghz", "{:0.1f}", None, True),
("cpu_cores_effective", "vCPUs", "{:0.1f}", None, True),
("cpu_ram", "RAM", "{:0.1f}", lambda x: x / 1000, False),
("disk_space", "Disk", "{:.0f}", None, True),
("discounted_dph_total", "$/hr", "{:0.4f}", None, True),
("dlperf", "DLP", "{:0.1f}", None, True),
("dlperf_per_dphtotal", "DLP/$", "{:0.2f}", None, True),
("driver_version", "NV Driver", "{}", None, True),
("inet_up", "Net_up", "{:0.1f}", None, True),
("inet_down", "Net_down", "{:0.1f}", None, True),
("reliability", "R", "{:0.1f}", lambda x: x * 100, True),
("duration", "Max_Days", "{:0.1f}", lambda x: x / (24.0 * 60.0 * 60.0), True),
("machine_id", "mach_id", "{}", None, True),
("verification", "status", "{}", None, True),
("host_id", "host_id", "{}", None, True),
("direct_port_count", "ports", "{}", None, True),
("geolocation", "country", "{}", None, True),
# ("direct_port_count", "Direct Port Count", "{}", None, True),
)
vol_displayable_fields = (
("id", "ID", "{}", None, True),
("cuda_max_good", "CUDA", "{:0.1f}", None, True),
("cpu_ghz", "cpu_ghz", "{:0.1f}", None, True),
("disk_bw", "Disk B/W", "{:0.1f}", None, True),
("disk_space", "Disk", "{:.0f}", None, True),
("disk_name", "Disk Name", "{}", None, True),
("storage_cost", "$/Gb/Month", "{:.2f}", None, True),
("driver_version", "NV Driver", "{}", None, True),
("inet_up", "Net_up", "{:0.1f}", None, True),
("inet_down", "Net_down", "{:0.1f}", None, True),
("reliability", "R", "{:0.1f}", lambda x: x * 100, True),
("duration", "Max_Days", "{:0.1f}", lambda x: x / (24.0 * 60.0 * 60.0), True),
("machine_id", "mach_id", "{}", None, True),
("verification", "status", "{}", None, True),
("host_id", "host_id", "{}", None, True),
("geolocation", "country", "{}", None, True),
)
nw_vol_displayable_fields = (
("id", "ID", "{}", None, True),
("disk_space", "Disk", "{:.0f}", None, True),
("storage_cost", "$/Gb/Month", "{:.2f}", None, True),
("inet_up", "Net_up", "{:0.1f}", None, True),
("inet_down", "Net_down", "{:0.1f}", None, True),
("reliability", "R", "{:0.1f}", lambda x: x * 100, True),
("duration", "Max_Days", "{:0.1f}", lambda x: x / (24.0 * 60.0 * 60.0), True),
("verification", "status", "{}", None, True),
("host_id", "host_id", "{}", None, True),
("cluster_id", "cluster_id", "{}", None, True),
("geolocation", "country", "{}", None, True),
("nw_disk_min_bw", "Min BW MiB/s", "{}", None, True),
("nw_disk_max_bw", "Max BW MiB/s", "{}", None, True),
("nw_disk_avg_bw", "Avg BW MiB/s", "{}", None, True),
)
# These fields are displayed for 'search benchmarks' (benchmarks table, AUTO-1418 schema)
benchmarks_displayable_fields = (
("id", "ID", "{}", None, True),
("machine_id", "mach_id", "{}", None, True),
("gpu_name", "Model", "{}", None, True),
("num_gpus", "N", "{}x", None, False),
("value", "Value", "{:0.2f}", None, True),
("dph_base", "$/hr", "{:0.4f}", None, True),
("type", "Type", "{}", None, True),
("template_hash", "Template", "{}", None, True),
("template_id", "Templ_ID", "{}", None, True),
("contract_id", "contract", "{}", None, True),
("image", "Image", "{}", None, True),
("last_update", "Date", "{}", lambda x: datetime.fromtimestamp(x, timezone.utc).strftime("%Y-%m-%d %H:%M"), True),
)
# These fields are displayed when you do 'show instances'
instance_fields = (
("id", "ID", "{}", None, True),
("machine_id", "Machine", "{}", None, True),
("actual_status", "Status", "{}", None, True),
("num_gpus", "Num", "{}x", None, False),
("gpu_name", "Model", "{}", None, True),
("gpu_util", "Util. %", "{:0.1f}", None, True),
("cpu_cores_effective", "vCPUs", "{:0.1f}", None, True),
("cpu_ram", "RAM", "{:0.1f}", lambda x: x / 1000, False),
("disk_space", "Storage", "{:.0f}", None, True),
("ssh_host", "SSH Addr", "{}", None, True),
("ssh_port", "SSH Port", "{}", None, True),
("dph_total", "$/hr", "{:0.4f}", None, True),
("image_uuid", "Image", "{}", None, True),
# ("dlperf", "DLPerf", "{:0.1f}", None, True),
# ("dlperf_per_dphtotal", "DLP/$", "{:0.1f}", None, True),
("inet_up", "Net up", "{:0.1f}", None, True),
("inet_down", "Net down", "{:0.1f}", None, True),
("reliability2", "R", "{:0.1f}", lambda x: x * 100, True),
("label", "Label", "{}", None, True),
("duration", "age(hours)", "{:0.2f}", lambda x: x/(3600.0), True),
("uptime_mins", "uptime(mins)", "{:0.2f}", None, True),
)
cluster_fields = (
("id", "ID", "{}", None, True),
("subnet", "Subnet", "{}", None, True),
("node_count", "Nodes", "{}", None, True),
("manager_id", "Manager ID", "{}", None, True),
("manager_ip", "Manager IP", "{}", None, True),
("machine_ids", "Machine ID's", "{}", None, True)
)
network_disk_fields = (
("network_disk_id", "Network Disk ID", "{}", None, True),
("free_space", "Free Space (GB)", "{}", None, True),
("total_space", "Total Space (GB)", "{}", None, True),
)
network_disk_machine_fields = (
("machine_id", "Machine ID", "{}", None, True),
("mount_point", "Mount Point", "{}", None, True),
)
overlay_fields = (
("overlay_id", "Overlay ID", "{}", None, True),
("name", "Name", "{}", None, True),
("subnet", "Subnet", "{}", None, True),
("cluster_id", "Cluster ID", "{}", None, True),
("instance_count", "Instances", "{}", None, True),
("instances", "Instance IDs", "{}", None, True),
)
instance_backup_fields = (
("id", "ID", "{}", None, True),
("contract_id", "Contract", "{}", None, True),
("type", "Type", "{}", None, True),
("machine_id", "Machine", "{}", None, True),
("status", "Status", "{}", None, True),
("age_hours", "Age(hours)", "{:0.2f}", None, True),
("info", "Info", "{}", None, True),
)
instance_backup_file_fields = (
("name", "Name", "{}", None, True),
("size_mb", "Size(MB)", "{:0.2f}", None, False),
("url", "URL", "{}", None, True),
)
volume_fields = (
("id", "ID", "{}", None, True),
("cluster_id", "Cluster ID", "{}", None, True),
("label", "Name", "{}", None, True),
("disk_space", "Disk", "{:.0f}", None, True),
("status", "status", "{}", None, True),
("disk_name", "Disk Name", "{}", None, True),
("driver_version", "NV Driver", "{}", None, True),
("inet_up", "Net_up", "{:0.1f}", None, True),
("inet_down", "Net_down", "{:0.1f}", None, True),
("reliability2", "R", "{:0.1f}", lambda x: x * 100, True),
("duration", "age(hours)", "{:0.2f}", lambda x: x/(3600.0), True),
("machine_id", "mach_id", "{}", None, True),
("verification", "Verification", "{}", None, True),
("host_id", "host_id", "{}", None, True),
("geolocation", "country", "{}", None, True),
("instances", "instances","{}", None, True)
)
# These fields are displayed when you do 'show machines'
machine_fields = (
("id", "ID", "{}", None, True),
("num_gpus", "#gpus", "{}", None, True),
("gpu_name", "gpu_name", "{}", None, True),
("disk_space", "disk", "{}", None, True),
("hostname", "hostname", "{}", lambda x: x[:16], True),
("driver_version", "driver", "{}", None, True),
("reliability2", "reliab", "{:0.4f}", None, True),
("verification", "veri", "{}", None, True),
("public_ipaddr", "ip", "{}", None, True),
("geolocation", "geoloc", "{}", None, True),
("num_reports", "reports", "{}", None, True),
("listed_gpu_cost", "gpuD_$/h", "{:0.2f}", None, True),
("min_bid_price", "gpuI$/h", "{:0.2f}", None, True),
("credit_discount_max", "rdisc", "{:0.2f}", None, True),
("listed_inet_up_cost", "netu_$/TB", "{:0.2f}", lambda x: x * 1024, True),
("listed_inet_down_cost", "netd_$/TB", "{:0.2f}", lambda x: x * 1024, True),
("gpu_occupancy", "occup", "{}", None, True),
)
# These fields are displayed when you do 'show maints'
maintenance_fields = (
("machine_id", "Machine ID", "{}", None, True),
("start_time", "Start (Date/Time in UTC)", "{}", lambda x: datetime.fromtimestamp(x, tz=timezone.utc).strftime('%Y-%m-%d/%H:%M'), True),
("end_time", "End (Date/Time in UTC)", "{}", lambda x: datetime.fromtimestamp(x, tz=timezone.utc).strftime('%Y-%m-%d/%H:%M'), True),
("duration_hours", "Duration (Hrs)", "{}", None, True),
("maintenance_category", "Category", "{}", None, True),
)
ipaddr_fields = (
("ip", "ip", "{}", None, True),
("first_seen", "first_seen", "{}", None, True),
("first_location", "first_location", "{}", None, True),
)
audit_log_fields = (
("ip_address", "ip_address", "{}", None, True),
("api_key_id", "api_key_id", "{}", None, True),
("created_at", "created_at", "{}", None, True),
("api_route", "api_route", "{}", None, True),
("args", "args", "{}", None, True),
)
# Backend writes per-second costs (web/price_increase_pending.py:_serialize_pending_row).
# Display matches the frontend convention: GPU per hour, storage per month, bandwidth per
# GB, platform fee as a percentage. Pre-format `old → new` strings in the command before
# rendering so a null `new_*` collapses to "-" via the empty-row fallback in display_table.
pending_price_increase_fields = (
("pending_id", "Pending ID", "{}", None, True),
("instance_id", "Instance", "{}", None, True),
("host_id", "Host", "{}", None, True),
("current_end", "Current End (UTC)", "{}", None, True),
("new_end", "New End (UTC)", "{}", None, True),
("gpu", "GPU ($/GPU/hr)", "{}", None, True),
("storage", "Storage ($/GB/mo)", "{}", None, True),
("bw_up", "BW Up ($/GB)", "{}", None, True),
("bw_down", "BW Down ($/GB)", "{}", None, True),
("fee", "Platform Fee (%)", "{}", None, True),
)
scheduled_jobs_fields = (
("id", "Scheduled Job ID", "{}", None, True),
("instance_id", "Instance ID", "{}", None, True),
("api_endpoint", "API Endpoint", "{}", None, True),
("start_time", "Start (Date/Time in UTC)", "{}", lambda x: datetime.fromtimestamp(x, tz=timezone.utc).strftime('%Y-%m-%d/%H:%M'), True),
("end_time", "End (Date/Time in UTC)", "{}", lambda x: datetime.fromtimestamp(x, tz=timezone.utc).strftime('%Y-%m-%d/%H:%M'), True),
("day_of_the_week", "Day of the Week", "{}", None, True),
("hour_of_the_day", "Hour of the Day in UTC", "{}", None, True),
("min_of_the_hour", "Minute of the Hour", "{}", None, True),
("frequency", "Frequency", "{}", None, True),
)
invoice_fields = (
("description", "Description", "{}", None, True),
("quantity", "Quantity", "{}", None, True),
("rate", "Rate", "{}", None, True),
("amount", "Amount", "{}", None, True),
("timestamp", "Timestamp", "{:0.1f}", None, True),
("type", "Type", "{}", None, True)
)
user_fields = (
# ("api_key", "api_key", "{}", None, True),
("balance", "Balance", "{}", None, True),
("balance_threshold", "Bal. Thld", "{}", None, True),
("balance_threshold_enabled", "Bal. Thld Enabled", "{}", None, True),
("billaddress_city", "City", "{}", None, True),
("billaddress_country", "Country", "{}", None, True),
("billaddress_line1", "Addr Line 1", "{}", None, True),
("billaddress_line2", "Addr line 2", "{}", None, True),
("billaddress_zip", "Zip", "{}", None, True),
("billed_expected", "Billed Expected", "{}", None, True),
("billed_verified", "Billed Vfy", "{}", None, True),
("billing_creditonly", "Billing Creditonly", "{}", None, True),
("can_pay", "Can Pay", "{}", None, True),
("credit", "Credit", "{:0.2f}", None, True),
("email", "Email", "{}", None, True),
("email_verified", "Email Vfy", "{}", None, True),
("fullname", "Full Name", "{}", None, True),
("got_signup_credit", "Got Signup Credit", "{}", None, True),
("has_billing", "Has Billing", "{}", None, True),
("has_payout", "Has Payout", "{}", None, True),
("id", "Id", "{}", None, True),
("last4", "Last4", "{}", None, True),
("paid_expected", "Paid Expected", "{}", None, True),
("paid_verified", "Paid Vfy", "{}", None, True),
("password_resettable", "Pwd Resettable", "{}", None, True),
("paypal_email", "Paypal Email", "{}", None, True),
("ssh_key", "Ssh Key", "{}", None, True),
("user", "User", "{}", None, True),
("username", "Username", "{}", None, True)
)
connection_fields = (
("id", "ID", "{}", None, True),
("name", "NAME", "{}", None, True),
("cloud_type", "Cloud Type", "{}", None, True),
)
# ANSI escape codes for background/foreground colors
BG_DARK_GRAY = '\033[40m' # Dark gray background
BG_LIGHT_GRAY = '\033[48;5;240m' # Light gray background
FG_WHITE = '\033[97m' # Bright white text
BG_RESET = '\033[0m' # Reset all formatting
def display_table(rows: list, fields: Tuple, replace_spaces: bool = True, auto_width: bool = True) -> None:
"""Basically takes a set of field names and rows containing the corresponding data and prints a nice tidy table
of it.
:param list rows: Each row is a dict with keys corresponding to the field names (first element) in the fields tuple.
:param Tuple fields: 5-tuple describing a field. First element is field name, second is human readable version, third is format string, fourth is a lambda function run on the data in that field, fifth is a bool determining text justification. True = left justify, False = right justify. Here is an example showing the tuples in action.
:rtype None:
Example of 5-tuple: ("cpu_ram", "RAM", "{:0.1f}", lambda x: x / 1000, False)
"""
header = [name for _, name, _, _, _ in fields]
out_rows = [header]
lengths = [len(x) for x in header]
for instance in rows:
row = []
out_rows.append(row)
for key, name, fmt, conv, _ in fields:
conv = conv or (lambda x: x)
val = instance.get(key, None)
if val is None:
s = "-"
else:
val = conv(val)
s = fmt.format(val)
if replace_spaces:
s = s.replace(' ', '_')
idx = len(row)
lengths[idx] = max(len(s), lengths[idx])
row.append(s)
if auto_width:
width = shutil.get_terminal_size((80, 20)).columns
start_col_idxs = [0]
total_len = 4 # +6ch for row label and -2ch for missing last sep in " ".join()
for i, l in enumerate(lengths):
total_len += l + 2
if total_len > width:
start_col_idxs.append(i) # index for the start of the next group
total_len = l + 6 # l + 2 + the 4 from the initial length
groups = {}
for row in out_rows:
grp_num = 0
for i in range(len(start_col_idxs)):
start = start_col_idxs[i]
end = start_col_idxs[i+1] if i+1 < len(start_col_idxs) else len(lengths)
groups.setdefault(grp_num, []).append(row[start:end])
grp_num += 1
for i, group in groups.items():
idx = start_col_idxs[i]
group_lengths = lengths[idx:idx+len(group[0])]
for row_num, row in enumerate(group):
bg_color = BG_DARK_GRAY if (row_num - 1) % 2 else BG_LIGHT_GRAY
row_label = " #" if row_num == 0 else f"{row_num:3d}"
out = [row_label]
for l, s, f in zip(group_lengths, row, fields[idx:idx+len(row)]):
_, _, _, _, ljust = f
if ljust: s = s.ljust(l)
else: s = s.rjust(l)
out.append(s)
print(bg_color + FG_WHITE + " ".join(out) + BG_RESET)
print()
else:
for row in out_rows:
out = []
for l, s, f in zip(lengths, row, fields):
_, _, _, _, ljust = f
if ljust:
s = s.ljust(l)
else:
s = s.rjust(l)
out.append(s)
print(" ".join(out))
def print_or_page(args, text):
""" Print text to terminal, or pipe to pager_cmd if too long. """
line_threshold = shutil.get_terminal_size(fallback=(80, 24)).lines
lines = text.splitlines()
if not args.full and len(lines) > line_threshold:
pager_cmd = ['less', '-R'] if shutil.which('less') else None
if pager_cmd:
proc = subprocess.Popen(pager_cmd, stdin=subprocess.PIPE)
proc.communicate(input=text.encode())
return True
else:
print(text)
return False
else:
print(text)
return False
cli/main.py
"""CLI entry point for vastai."""
import sys
import os
import json
import requests
from datetime import datetime, timezone
from vastai.api.client import is_search_quota_exceeded
from vastai.cli.parser import apwrap, argument, MyWideHelpFormatter, set_completers
from vastai.cli.util import (
APIKEY_FILE, TFAKEY_FILE, VERSION, server_url_default, api_key_guard,
format_key_suffix,
)
try:
JSONDecodeError = json.JSONDecodeError
except AttributeError:
JSONDecodeError = ValueError
def _is_tfa_session_expired(status_code, errmsg):
"""True if this error signals an expired 2FA session key.
401 "Invalid user key" is sent when a deleted API key is used; an expired
TFA session instead returns 404 "Session expired. Please log in again."
(matches legacy vast.py's handling of both signals).
"""
return (
(status_code == 401 and errmsg == "Invalid user key")
or (status_code == 404 and errmsg == "Session expired. Please log in again.")
)
def with_quota_reset(response, errmsg):
"""Append the reset time from ``X-RateLimit-Reset`` to a search quota denial's message."""
if not is_search_quota_exceeded(response):
return errmsg
try:
reset = datetime.fromtimestamp(int(response.headers["X-RateLimit-Reset"]), timezone.utc)
except (KeyError, ValueError):
return errmsg
return f"{errmsg}. Resets at {reset:%Y-%m-%d %H:%M} UTC."
def read_key(path):
"""The key stored at a path, or None if it is absent or unreadable."""
try:
with open(path) as f:
return f.read().strip() or None
except OSError:
return None
def _emit_error(args, status_code, message):
"""Emit a command error in the appropriate format.
In ``--raw`` mode, prints a JSON error object to stderr so scripts can
parse it; otherwise prints a human-readable line. Always goes to stderr
so stdout stays clean for scripting consumers.
"""
if getattr(args, "raw", False):
payload = {"error": True, "status_code": status_code, "msg": message}
print(json.dumps(payload), file=sys.stderr)
return
if status_code:
print(f"Failed with error {status_code}: {message}", file=sys.stderr)
else:
print(message, file=sys.stderr)
if status_code == 401:
if "Two Factor Authentication" in message:
print(
"\nThis endpoint requires a 2FA session. Authenticate with the method you have set up:\n"
" vastai tfa login --method-type totp -c <CODE>\n"
" vastai tfa login --method-type sms --secret <SECRET> -c <CODE>\n"
" vastai tfa login --method-type email --secret <SECRET> -c <CODE>\n"
" vastai tfa login --backup-code ABCD-EFGH-IJKL\n"
"For SMS or email, first run 'vastai tfa send-sms' or 'vastai tfa send-email' to get <SECRET>.",
file=sys.stderr,
)
return
env = os.environ.get("VAST_API_KEY")
file_key = read_key(APIKEY_FILE)
tfa_key = read_key(TFAKEY_FILE)
# A --api-key outranks the rest, so a value matching none came from there.
sent = getattr(args, "api_key", None)
key_missing = not sent and not env and not file_key
if sent and sent == tfa_key:
print(f" Sent your 2FA session key from {TFAKEY_FILE} (ends in {format_key_suffix(sent)}).", file=sys.stderr)
print(" If it has expired, start a new one: vastai tfa login", file=sys.stderr)
elif sent and sent != env and sent != file_key:
server = getattr(args, "url", None) or server_url_default
print(f" Sent key from --api-key (ends in {format_key_suffix(sent)}).", file=sys.stderr)
print(f" Check that key is valid for {server}.", file=sys.stderr)
elif key_missing:
print(" No API key is configured.", file=sys.stderr)
print(" Run: vastai set api-key <KEY>", file=sys.stderr)
print(" Create a key at https://console.vast.ai/manage-keys/", file=sys.stderr)
elif env and file_key:
print(f" Sent key from $VAST_API_KEY (ends in {format_key_suffix(env)}). Env var overrides the file.", file=sys.stderr)
print(f" Unset the VAST_API_KEY env var to use the saved key in {APIKEY_FILE} (ends in {format_key_suffix(file_key)}) instead.", file=sys.stderr)
elif env:
print(f" Sent key from $VAST_API_KEY (ends in {format_key_suffix(env)}). Update the env var with a new key.", file=sys.stderr)
elif file_key:
print(f" Sent key from {APIKEY_FILE} (ends in {format_key_suffix(file_key)}). Update with: vastai set api-key <KEY>", file=sys.stderr)
# Create the global parser instance
parser = apwrap(
epilog="Use 'vastai COMMAND --help' for more info about a command. AI agent? See https://raw.githubusercontent.com/vast-ai/vast-cli/master/vastai/SKILL.md",
formatter_class=MyWideHelpFormatter
)
def main():
# Import all command modules - the import itself triggers decorator
# registrations on the global parser via _get_parser().
from vastai.cli.commands import ( # noqa: F401
instances, offers, machines, teams, keys, endpoints,
billing, storage, auth, misc, deployments, metrics,
benchmarks,
price_increase,
repl,
update,
uninstall,
# clusters, # cluster/overlay commands disabled for now
)
# Wire up tab completers now that command modules are loaded
set_completers(
instance_machine_fn=lambda **kw: instances.show__instances(
type('Args', (), {'api_key': os.getenv('VAST_API_KEY'), 'url': server_url_default,
'retry': 3, 'explain': False, 'curl': False,
'quiet': False, 'raw': False})(),
{'internal': True, 'field': 'machine_id'}),
instance_fn=lambda **kw: instances.show__instances(
type('Args', (), {'api_key': os.getenv('VAST_API_KEY'), 'url': server_url_default,
'retry': 3, 'explain': False, 'curl': False,
'quiet': False, 'raw': False})(),
{'internal': True, 'field': 'id'}),
)
# Add global arguments
parser.add_argument("--url", help="Server REST API URL", default=server_url_default)
parser.add_argument("--retry", help="Retry limit", type=int, default=3)
parser.add_argument("--explain", action="store_true", help="Output verbose explanation of mapping of CLI calls to HTTPS API endpoints")
parser.add_argument("--raw", action="store_true", help="Output machine-readable json")
parser.add_argument("--full", action="store_true", help="Print full results instead of paging with `less` for commands that support it")
parser.add_argument("--curl", action="store_true", help="Show a curl equivalency to the call")
parser.add_argument("--api-key", help="API Key to use. defaults to using the one stored in {}".format(APIKEY_FILE), type=str, required=False, default=os.getenv("VAST_API_KEY", api_key_guard))
parser.add_argument("--version", help="Show CLI version", action="version", version=VERSION)
parser.add_argument("--no-color", action="store_true", help="Disable colored output for commands that support it")
# Tab completion
try:
import argcomplete
import argparse
from typing import List, Optional
from vastai.cli.parser import build_command_maps, two_stage_command_completions
class MyAutocomplete(argcomplete.CompletionFinder):
# Two-stage (verb -> object) completion over the flat "verb object"
# subparser names. Logic in parser.py (pure + unit-tested).
def _command_maps(self):
cached = getattr(self, "_cmd_maps", None)
if cached is None:
from vastai.cli.util import get_role
cached = self._cmd_maps = build_command_maps(self._parser, role=get_role())
return cached
def _get_completions(self, comp_words, cword_prefix, cword_prequote, last_wordbreak_pos):
verbs, verb_objs, singles = self._command_maps()
cands, merged = two_stage_command_completions(
comp_words, cword_prefix, verbs, verb_objs, singles)
if cands is not None:
return cands
return super()._get_completions(merged, cword_prefix, cword_prequote, last_wordbreak_pos)
def quote_completions(self, completions: List[str], cword_prequote: str, last_wordbreak_pos: Optional[int]) -> List[str]:
pre = super().quote_completions(completions, cword_prequote, last_wordbreak_pos)
return sorted(pre, key=lambda x: x.startswith('-'))
myautocc = MyAutocomplete()
myautocc(parser.parser)
except ImportError:
pass
args = parser.parse_args()
# Passive upgrade nudge: best-effort, ≤1 manifest GET + ≤1 stderr line per
# 24h, silent when offline/piped/CI. Never raises. Opt out with
# VASTAI_NO_UPDATE_CHECK=1. Runs first so it's visible before any command
# output, not buried after it. See selfupdate.py and §7.
from vastai.cli.selfupdate import notify_update
notify_update(args)
# API key resolution
if args.api_key is api_key_guard:
key_file = TFAKEY_FILE if os.path.exists(TFAKEY_FILE) else APIKEY_FILE
if os.path.exists(key_file):
with open(key_file, "r") as reader:
args.api_key = reader.read().strip()
else:
args.api_key = None
run_command(args)
def run_command(args):
"""Execute ``args.func(args)``, retrying once on an expired 2FA session key."""
while True:
try:
res = args.func(args)
if args.raw and res is not None:
try:
print(json.dumps(res, indent=1, sort_keys=True))
except (TypeError, ValueError):
print(json.dumps(res.json(), indent=1, sort_keys=True))
sys.exit(0)
sys.exit(res)
except requests.exceptions.HTTPError as e:
try:
errmsg = e.response.json().get("msg")
except JSONDecodeError:
if e.response.status_code == 401:
errmsg = "Please log in or sign up"
else:
errmsg = "(no detail message supplied)"
# 2FA session key expired -> retry with long-lived API key.
if _is_tfa_session_expired(e.response.status_code, errmsg) and os.path.exists(TFAKEY_FILE):
print(f"Failed with error {e.response.status_code}: Your 2FA session has expired.")
os.remove(TFAKEY_FILE)
if os.path.exists(APIKEY_FILE):
with open(APIKEY_FILE, "r") as reader:
args.api_key = reader.read().strip()
print(f"Trying again with your normal API Key from {APIKEY_FILE}...")
print("To start a new 2FA session, run: vastai tfa login")
continue
else:
print("Run `vastai tfa login` to start a new 2FA session and try again.")
break
_emit_error(args, e.response.status_code, with_quota_reset(e.response, errmsg))
break
except ValueError as e:
_emit_error(args, 0, str(e))
break
if __name__ == "__main__":
try:
main()
except (KeyboardInterrupt, BrokenPipeError):
pass
cli/parser.py
"""
Argument parser infrastructure for the Vast.ai CLI.
Extracted from vast.py - contains the core parser wrapper classes
that handle command registration, argument parsing, and tab completion.
"""
from __future__ import unicode_literals, print_function
import sys
import argparse
from pathlib import Path
from vastai.cli.util import get_role, is_client_view
# ---------------------------------------------------------------------------
# Tab-completion helpers
#
# These reference show__instances which lives in the CLI commands layer.
# We keep them as lazy references (initially None) so this module can be
# imported without pulling in the whole command tree. The CLI entry-point
# populates them once everything is wired up.
# ---------------------------------------------------------------------------
# Populated later by the CLI entry point
_complete_instance_machine = None
_complete_instance = None
def complete_instance_machine(prefix=None, action=None, parser=None, parsed_args=None):
if _complete_instance_machine is not None:
return _complete_instance_machine(prefix=prefix, action=action, parser=parser, parsed_args=parsed_args)
return []
def complete_instance(prefix=None, action=None, parser=None, parsed_args=None):
if _complete_instance is not None:
return _complete_instance(prefix=prefix, action=action, parser=parser, parsed_args=parsed_args)
return []
def complete_sshkeys(prefix=None, action=None, parser=None, parsed_args=None):
return [str(m) for m in Path.home().joinpath('.ssh').glob('*.pub')]
# ---------------------------------------------------------------------------
# Two-stage (verb -> object) command-name completion
#
# Pure, unit-testable helpers that turn the flat "verb object" subparser names
# into staged completion. The argcomplete adapter in cli/main.py calls them.
# ---------------------------------------------------------------------------
def build_command_maps(parser, role=None):
"""Derive (verbs, verb_objs, singles) from a parser's subparser names, excluding hidden and (unless role='host') host-only commands."""
names = []
for action in parser._actions:
if isinstance(action, argparse._SubParsersAction):
for name, sp in action.choices.items():
if getattr(sp, 'hidden', False) or name != getattr(sp, 'command_name', name):
continue
if is_client_view(role) and getattr(sp, 'host_only', False):
continue
names.append(name)
break
verbs, verb_objs, singles = set(), {}, set()
for name in names:
verb, sep, obj = name.partition(" ")
if sep:
verbs.add(verb)
verb_objs.setdefault(verb, set()).add(obj)
else:
singles.add(verb)
return verbs, verb_objs, singles
def two_stage_command_completions(comp_words, cword_prefix, verbs, verb_objs, singles):
"""Stage command-name completion. Returns (candidates, merged): candidates is
the verb/object list, or None to delegate flag completion using merged."""
n = len(comp_words)
if n == 1: # completing the verb / bare command
return sorted(c for c in (verbs | singles) if c.startswith(cword_prefix)), None
if n == 2 and comp_words[1] in verbs: # completing the object for a verb
objs = verb_objs.get(comp_words[1], set())
return sorted(o for o in objs if o.startswith(cword_prefix)), None
# Past the command name: fuse "verb object" so the right subparser is found.
# Never fuse across a flag, mirroring apwrap.parse_args.
merged = comp_words
if n >= 3 and comp_words[1] in verbs and not comp_words[2].startswith("-"):
merged = [comp_words[0], comp_words[1] + " " + comp_words[2]] + list(comp_words[3:])
return None, merged
def set_completers(instance_machine_fn=None, instance_fn=None):
"""
Wire up the tab-completion functions. Called once from the CLI
entry-point after all command modules have been loaded.
"""
global _complete_instance_machine, _complete_instance
if instance_machine_fn is not None:
_complete_instance_machine = instance_machine_fn
if instance_fn is not None:
_complete_instance = instance_fn
# ---------------------------------------------------------------------------
# Argument wrapper
# ---------------------------------------------------------------------------
class argument(object):
"""Thin wrapper that stores positional and keyword args for later
application to an argparse subparser."""
def __init__(self, *args, mutex_group=None, **kwargs):
self.args = args
self.kwargs = kwargs
self.mutex_group = mutex_group # Name of the mutually exclusive group this arg belongs to
class hidden_aliases(object):
"""A list-like that is falsy so argparse won't display aliases in help."""
def __init__(self, l):
self.l = l
def __iter__(self):
return iter(self.l)
def __bool__(self):
return False
def __nonzero__(self):
return False
def append(self, x):
self.l.append(x)
class MyWideHelpFormatter(argparse.RawTextHelpFormatter):
def __init__(self, prog):
super().__init__(prog, width=128, max_help_position=50, indent_increment=1)
COMMAND_GROUPS = {
'vastai.cli.commands.offers': 'Search',
'vastai.cli.commands.instances': 'Instances',
'vastai.cli.commands.benchmarks': 'Instances',
'vastai.cli.commands.machines': 'Host machines',
'vastai.cli.commands.teams': 'Teams',
'vastai.cli.commands.keys': 'Auth & keys',
'vastai.cli.commands.auth': 'Auth & keys',
'vastai.cli.commands.billing': 'Billing & account',
'vastai.cli.commands.endpoints': 'Serverless',
'vastai.cli.commands.deployments': 'Serverless',
'vastai.cli.commands.metrics': 'Metrics',
'vastai.cli.commands.storage': 'Storage volumes',
'vastai.cli.commands.misc': 'Other',
'vastai.cli.commands.repl': 'Other',
'vastai.cli.commands.update': 'Other',
}
# Some misc.py commands belong with their target resource section, not Other.
COMMAND_OVERRIDES = {
'execute': 'Instances',
'logs': 'Instances',
'ssh-url': 'Instances',
'scp-url': 'Instances',
'take snapshot': 'Instances',
'reports': 'Host machines',
}
# Host-only command classification (CLN-3582): display-only, hides host-only commands from --help/completion/hints for the client role; unclassified defaults to visible.
# Modules where every command is host-only.
HOST_ONLY_MODULES = {
'vastai.cli.commands.machines',
'vastai.cli.commands.metrics',
}
# Individual host-only commands in modules that are mostly client-facing (same shape as COMMAND_OVERRIDES above).
HOST_ONLY_COMMANDS = {
'reports', # misc.py: machine reports, not instance logs
'show earnings', # billing.py: machine earning history
'list network-volume', # storage.py: list disk space for rent
'list volume',
'list volumes',
'unlist network-volume',
'unlist volume',
}
def is_host_only_command(name, module, explicit=None):
"""Whether a command is hidden from the client CLI view: explicit host_only= kwarg > per-command override > per-module default > False (visible)."""
if explicit is not None:
return bool(explicit)
if name in HOST_ONLY_COMMANDS:
return True
return module in HOST_ONLY_MODULES
GROUP_ORDER = [
'Search', 'Instances', 'Host machines', 'Teams', 'Auth & keys',
'Billing & account', 'Serverless', 'Metrics', 'Storage volumes', 'Other',
]
# ---------------------------------------------------------------------------
# Unreleased/feature-flagged commands
#
# Hidden from --help and tab completion regardless of anything else — a
# discoverability gate, not an access gate: the command still runs if typed
# directly, so internal testing keeps working. Remove an entry once the
# feature is ready to announce.
# ---------------------------------------------------------------------------
HIDDEN_COMMANDS = {
'repl', # interactive shell: internal testing before announcing
'search network-volumes', # network volumes are not yet released
'create network-volume',
'list network-volume',
'unlist network-volume',
'show network-disks', # network disks are not yet released
'add network-disk',
}
def is_hidden_command(name, explicit=None):
"""Whether a registered command should be hidden from --help/completion.
Priority: explicit ``hidden=`` kwarg on the decorator > ``HIDDEN_COMMANDS``.
Defaults to ``False`` (visible).
"""
if explicit is not None:
return bool(explicit)
return name in HIDDEN_COMMANDS
class GroupedArgumentParser(argparse.ArgumentParser):
"""ArgumentParser that renders subcommands grouped by registering module."""
def format_help(self):
sub_action = self._subparsers_action()
if sub_action is None or not sub_action.choices:
return super().format_help()
formatter = self._get_formatter()
formatter.add_usage(self.usage, self._actions, self._mutually_exclusive_groups)
formatter.add_text(self.description)
# Skip the subparsers action; we render commands grouped below.
for action_group in self._action_groups:
non_sub_actions = [
a for a in action_group._group_actions
if not isinstance(a, argparse._SubParsersAction)
]
if not non_sub_actions:
continue
formatter.start_section(action_group.title)
formatter.add_text(action_group.description)
formatter.add_arguments(non_sub_actions)
formatter.end_section()
self._add_grouped_commands(formatter, sub_action)
formatter.add_text(self.epilog)
return formatter.format_help()
def _subparsers_action(self):
for a in self._actions:
if isinstance(a, argparse._SubParsersAction):
return a
return None
def _add_grouped_commands(self, formatter, sub_action):
"""Render grouped commands."""
role = get_role()
grouped = {}
for pseudo in sub_action._choices_actions:
sp = sub_action.choices.get(pseudo.dest)
if sp is None:
continue
if getattr(sp, 'hidden', False):
continue
if is_client_view(role) and getattr(sp, 'host_only', False):
continue
label = COMMAND_OVERRIDES.get(pseudo.dest)
if label is None:
func = sp.get_default('func')
module = getattr(func, '__module__', None)
label = COMMAND_GROUPS.get(module, 'Other')
grouped.setdefault(label, []).append(pseudo)
seen = set()
for label in GROUP_ORDER:
if label in grouped:
formatter.start_section(label)
formatter.add_arguments(grouped[label])
formatter.end_section()
seen.add(label)
for label, actions in grouped.items():
if label in seen:
continue
formatter.start_section(label)
formatter.add_arguments(actions)
formatter.end_section()
# ---------------------------------------------------------------------------
# Main parser wrapper
# ---------------------------------------------------------------------------
class apwrap(object):
"""Wraps :class:`argparse.ArgumentParser` with convenience methods for
two-word command registration (``verb object``) and tab-completion."""
def __init__(self, *args, **kwargs):
if "formatter_class" not in kwargs:
kwargs["formatter_class"] = MyWideHelpFormatter
self.parser = GroupedArgumentParser(*args, **kwargs)
self.parser.set_defaults(func=self.fail_with_help)
self.subparsers_ = None
self.subparser_objs = []
self.added_help_cmd = False
self.post_setup = []
self.verbs = set()
self.objs = set()
def fail_with_help(self, *a, **kw):
self.parser.print_help(sys.stderr)
raise SystemExit
def add_argument(self, *a, **kw):
if not kw.get("parent_only"):
for x in self.subparser_objs:
try:
# Create a global options group for better visual separation
if not hasattr(x, '_global_options_group'):
x._global_options_group = x.add_argument_group('Global options (available for all commands)')
# Use SUPPRESS as default for subparsers so they don't overwrite
# values already set by the main parser when the argument is placed
# before the subcommand (e.g., `vastai --url <url> get wrkgrp-logs`)
subparser_kw = kw.copy()
subparser_kw['default'] = argparse.SUPPRESS
x._global_options_group.add_argument(*a, **subparser_kw)
except argparse.ArgumentError:
# duplicate - or maybe other things, hopefully not
pass
return self.parser.add_argument(*a, **kw)
def subparsers(self, *a, **kw):
if self.subparsers_ is None:
kw["metavar"] = "command"
kw["help"] = "command to run. one of:"
kw["dest"] = "command" # the name as typed, so a command can tell when its alias was used
self.subparsers_ = self.parser.add_subparsers(*a, **kw)
return self.subparsers_
def get_name(self, verb, obj):
if obj:
self.verbs.add(verb)
self.objs.add(obj)
name = verb + ' ' + obj
else:
self.objs.add(verb)
name = verb
return name
def command(self, *arguments, aliases=(), help=None, hidden=None, host_only=None, **kwargs):
help_ = help
if not self.added_help_cmd:
self.added_help_cmd = True
@self.command(argument("subcommand", default=None, nargs="?"), help="print this help message")
def help(*a, **kw):
self.fail_with_help()
def inner(func):
dashed_name = func.__name__.replace("_", "-")
verb, _, obj = dashed_name.partition("--")
name = self.get_name(verb, obj)
aliases_transformed = [] if aliases else hidden_aliases([])
for x in aliases:
verb, _, obj = x.partition(" ")
aliases_transformed.append(self.get_name(verb, obj))
if "formatter_class" not in kwargs:
kwargs["formatter_class"] = MyWideHelpFormatter
sp = self.subparsers().add_parser(name, aliases=aliases_transformed, help=help_, **kwargs)
sp.hidden = is_hidden_command(name, explicit=hidden)
sp.host_only = is_host_only_command(name, func.__module__, explicit=host_only)
sp.command_name = name
# TODO: Sometimes the parser.command has a help parameter. Ideally
# I'd extract this during the sdk phase but for the life of me
# I can't find it.
setattr(func, "mysignature", sp)
setattr(func, "mysignature_help", help_)
self.subparser_objs.append(sp)
self._process_arguments_with_groups(sp, arguments)
sp.set_defaults(func=func)
return func
if len(arguments) == 1 and type(arguments[0]) != argument:
func = arguments[0]
arguments = []
return inner(func)
return inner
def parse_args(self, argv=None, *a, **kw):
if argv is None:
argv = sys.argv[1:]
argv_ = []
for x in argv:
# Merge "verb object" into one command token, but never swallow a
# flag: a verb that is also a bare command (e.g. `update`) must
# accept options (`vastai update --check`).
if argv_ and argv_[-1] in self.verbs and not x.startswith("-"):
argv_[-1] += " " + x
else:
argv_.append(x)
args = self.parser.parse_args(argv_, *a, **kw)
for func in self.post_setup:
func(args)
return args
def _process_arguments_with_groups(self, parser_obj, arguments):
"""Process arguments and handle mutually exclusive groups"""
mutex_groups_to_required = {}
arg_to_group = {}
# Determine if any mutex groups are required
for arg in arguments:
key = arg.args[0]
if arg.mutex_group:
is_required = arg.kwargs.pop('required', False)
group_name = arg.mutex_group
arg_to_group[key] = group_name
if mutex_groups_to_required.get(group_name):
continue # if marked as required then it stays required
else:
mutex_groups_to_required[group_name] = is_required
name_to_group_parser = {} # Create mutually exclusive group parsers
for group_name, is_required in mutex_groups_to_required.items():
mutex_group = parser_obj.add_mutually_exclusive_group(required=is_required)
name_to_group_parser[group_name] = mutex_group
for arg in arguments: # Add args via the appropriate parser
key = arg.args[0]
if arg_to_group.get(key):
group_parser = name_to_group_parser[arg_to_group[key]]
tsp = group_parser.add_argument(*arg.args, **arg.kwargs)
else:
tsp = parser_obj.add_argument(*arg.args, **arg.kwargs)
self._add_completer(tsp, arg)
def _add_completer(self, tsp, arg):
"""Helper function to add completers based on argument names"""
myCompleter = None
comparator = arg.args[0].lower()
if comparator.startswith('machine'):
myCompleter = complete_instance_machine
elif comparator.startswith('id') or comparator.endswith('id'):
myCompleter = complete_instance
elif comparator.startswith('ssh'):
myCompleter = complete_sshkeys
if myCompleter:
setattr(tsp, 'completer', myCompleter)
cli/repl/__init__.py
"""Interactive shell for the vast CLI (`vastai repl`).
One process, the same commands: `session.py` owns the loop, `bridge.py` runs a
line through the real parser, `completion.py` stages Tab completion and
`catalog.py` indexes the command tree for both.
"""
from vastai.cli.repl.session import run_repl # noqa: F401
cli/repl/bridge.py
"""Run one CLI line inside the REPL process.
The REPL reuses the exact parser and command functions that ``vastai ...``
uses, so ``show instances`` at the prompt behaves identically to
``vastai show instances`` — same auth, same output, same flags. The only
difference is control flow: a failing command ends the line, never the session.
"""
import argparse
import json
import os
import shlex
import sys
import requests
from vastai.cli import main as cli_main
from vastai.cli.repl.catalog import option_action
# Globals the session carries onto every line, so a bare `show instances`
# doesn't have to repeat --api-key/--url/--raw. A flag typed on the line wins.
SESSION_GLOBALS = (
"api_key", "url", "retry", "explain", "curl", "raw", "full", "no_color",
)
# Outcomes of an expired-2FA-session recovery attempt.
RETRY = "retry" # a fresh key was loaded; run the command again
HANDLED = "handled" # the expiry was reported; don't report it twice
def run_line(parser, line, session_args):
"""Parse and execute one command line. Returns the command's return value,
or None if the line failed to parse or the command errored."""
try:
argv = tokenize(line)
except ValueError as exc:
print(f"parse error: {exc}", file=sys.stderr)
return None
if not argv:
return None
try:
args = parser.parse_args(argv)
except SystemExit:
return None # argparse already printed usage or an error
apply_session_globals(parser, args, session_args, argv)
func = getattr(args, "func", None)
if func is None:
return None
return _invoke(args, func, session_args)
def tokenize(line):
"""Split a line into argv the way a shell would.
POSIX rules everywhere except Windows, where a backslash is a path
separator rather than an escape: POSIX splitting would quietly turn
``--onstart C:\\tmp\\go.sh`` into ``C:tmpgo.sh``.
"""
if os.name != "nt":
return shlex.split(line)
lexer = shlex.shlex(line, posix=False)
lexer.whitespace_split = True
return [_unquote(token) for token in lexer]
def _unquote(token):
if len(token) > 1 and token[0] == token[-1] and token[0] in "\"'":
return token[1:-1]
return token
def apply_session_globals(parser, args, session_args, argv=()):
"""Carry the session's global flags onto one line's args.
Presence is read from the line itself rather than by comparing parsed
values against defaults: an option typed with exactly its default value
(``--retry 3`` while the session runs with ``--retry 10``) must still win.
"""
inner = getattr(parser, "parser", parser)
sub = getattr(getattr(args, "func", None), "mysignature", None)
typed = explicit_dests(inner, sub, argv)
for name in SESSION_GLOBALS:
if name in typed or not hasattr(args, name) or not hasattr(session_args, name):
continue
setattr(args, name, getattr(session_args, name))
def explicit_dests(root, sub, argv):
"""The dests of the options a line actually names, in any form argparse
accepts (``--flag``, ``--flag=value``, a short alias, an abbreviation).
Each option is read against the parser argparse itself would use: the root
parser before the command name, the command's own parser after it. An
abbreviation can be unambiguous in one and ambiguous in the other — ``--c``
is ``--curl`` up front, but could be ``--cols`` on ``show instances``.
Scanning stops where argparse stops treating words as options: at ``--`` or
at a REMAINDER option, so ``create instance --args --raw`` passes ``--raw``
to the container rather than counting it as a global.
"""
dests = set()
parser, skip = root, 0
for token in argv:
if skip:
skip -= 1
continue
if token == "--":
break
if not token.startswith("-") or token == "-":
parser = sub or root # past the command name from here on
continue
action = option_action(parser, token)
if action is None:
continue
if action.nargs == argparse.REMAINDER:
break
dests.add(action.dest)
if parser is root and action.nargs != 0 and "=" not in token:
skip = 1 # its value is not the command name
return dests
def _invoke(args, func, session_args=None):
try:
res = func(args)
except SystemExit:
# Commands (and --help) exit the process in one-shot mode; here that
# just ends the line.
return None
except requests.exceptions.HTTPError as exc:
outcome = _recover_expired_tfa_session(args, exc, session_args)
if outcome == RETRY:
return _invoke(args, func) # retry once, as the one-shot CLI does
if outcome != HANDLED:
_emit_http_error(args, exc)
return None
except ValueError as exc:
cli_main._emit_error(args, 0, str(exc))
return None
except KeyboardInterrupt:
print("^C", file=sys.stderr)
return None
except Exception as exc: # a broken command must not kill the session
print(f"{type(exc).__name__}: {exc}", file=sys.stderr)
return None
if getattr(args, "raw", False) and res is not None:
_print_raw(res)
return res
def _recover_expired_tfa_session(args, exc, session_args):
"""Fall back to the saved API key when a 2FA session expires, as
``main.run_command`` does — and keep the new key on the session, so the
rest of the REPL's lines work too rather than failing one by one.
Returns RETRY, HANDLED (reported, nothing left to say) or None (not a 2FA
expiry, so the caller reports it).
"""
status, msg = _error_detail(exc)
if not cli_main._is_tfa_session_expired(status, msg):
return None
if not os.path.exists(cli_main.TFAKEY_FILE):
return None
print(f"Failed with error {status}: Your 2FA session has expired.")
os.remove(cli_main.TFAKEY_FILE)
if not os.path.exists(cli_main.APIKEY_FILE):
print("Run `vastai tfa login` to start a new 2FA session and try again.")
return HANDLED
with open(cli_main.APIKEY_FILE, "r") as reader:
key = reader.read().strip()
args.api_key = key
if session_args is not None:
session_args.api_key = key
print(f"Trying again with your normal API Key from {cli_main.APIKEY_FILE}...")
print("To start a new 2FA session, run: vastai tfa login")
return RETRY
def _error_detail(exc):
"""(status, message) from an API error, however malformed the response."""
resp = exc.response
status = getattr(resp, "status_code", 0)
try:
msg = resp.json().get("msg")
except (ValueError, AttributeError):
msg = None
if not isinstance(msg, str):
# A body that is valid JSON but carries no "msg" would otherwise hand
# None to _emit_error, which tests it with `in` and raises TypeError.
msg = "Please log in or sign up" if status == 401 else "(no detail message supplied)"
return status, msg
def _emit_http_error(args, exc):
"""Format an API error the same way the one-shot CLI's main loop does."""
status, msg = _error_detail(exc)
cli_main._emit_error(args, status, msg)
def _print_raw(res):
try:
print(json.dumps(res, indent=1, sort_keys=True))
except (TypeError, ValueError):
print(json.dumps(res.json(), indent=1, sort_keys=True))
cli/repl/catalog.py
"""What the REPL knows about the CLI's command tree.
Commands are registered on one argparse parser as flat ``"verb object"``
subparser names ("show instances"). The catalog derives the two-stage view of
that — verbs, the objects each verb takes, bare commands, per-command options —
once at startup, so completion and command resolution stay dict lookups instead
of a walk over ~150 subparsers on every keystroke.
"""
import argparse
import difflib
from vastai.cli.parser import build_command_maps
def option_action(parser, token):
"""The action a flag token names on a parser: exact match, ``--flag=value``,
or the unambiguous abbreviation argparse itself would accept."""
options = parser._option_string_actions
name = token.split("=", 1)[0]
action = options.get(name)
if action is not None:
return action
if name.startswith("--"):
matches = {id(a): a for opt, a in options.items() if opt.startswith(name)}
if len(matches) == 1:
return next(iter(matches.values()))
return None
class CommandCatalog:
"""Two-stage (verb -> object) index over a live ``apwrap`` parser."""
def __init__(self, parser):
self._parser = getattr(parser, "parser", parser)
self.verbs, self.verb_objects, self.singles = build_command_maps(self._parser)
self._choices = self._subparser_choices()
# Hidden commands are left out of the maps above — they are gated from
# discovery, not from use — but a line naming one must still resolve,
# or the REPL would reject a command the one-shot CLI runs.
self._runnable_verbs = {name.split(" ")[0] for name in self._choices if " " in name}
self._flag_cache = {}
self.first_words = sorted(self.verbs | self.singles)
self.names = sorted(self.singles | {
f"{verb} {obj}" for verb, objs in self.verb_objects.items() for obj in objs})
def _subparser_choices(self):
for action in self._parser._actions:
if isinstance(action, argparse._SubParsersAction):
return action.choices
return {}
def objects(self, verb):
"""The objects registered for a verb, e.g. show -> instances, user, ..."""
return sorted(self.verb_objects.get(verb, ()))
def strip_options(self, tokens):
"""Drop the global options a line leads with (``--raw show user``), so
what's left starts at the command name."""
i = 0
while i < len(tokens) and tokens[i].startswith("-") and tokens[i] != "-":
token = tokens[i]
action = option_action(self._parser, token)
i += 1
if action is not None and action.nargs != 0 and "=" not in token:
i += 1 # the option's value
return tokens[i:]
def resolve(self, tokens):
"""The command name a tokenised line names, or None if it names none.
Mirrors ``apwrap.parse_args``: a verb absorbs the next token as its
object, but never a flag, so ``update --check`` stays the bare command.
"""
if not tokens:
return None
if tokens[0] in self._runnable_verbs and len(tokens) > 1 and not tokens[1].startswith("-"):
fused = f"{tokens[0]} {tokens[1]}"
if fused in self._choices:
return fused
# `update bogus` is not the bare `update` command. Only a bare
# command that takes an argument can legitimately own a second word.
if not self.takes_positional(tokens[0]):
return None
return tokens[0] if tokens[0] in self._choices else None
def suggest(self, tokens):
"""Close command names for a line that resolves to nothing — 'did you mean'."""
if not tokens:
return []
head = tokens[0]
if head in self.verbs: # known verb, missing or misspelled object
objs = self.objects(head)
near = difflib.get_close_matches(tokens[1], objs, n=3, cutoff=0.5) if len(tokens) > 1 else []
return [f"{head} {obj}" for obj in (near or objs)]
return (difflib.get_close_matches(" ".join(tokens[:2]), self.names, n=3, cutoff=0.5)
or difflib.get_close_matches(head, self.first_words, n=3, cutoff=0.5))
def flags(self, name):
"""``{"--flag": action}`` for one command, cached after the first lookup.
Long forms only: they are what completion offers. Use ``option`` to look
up a token the user actually typed, which may be a short alias.
"""
if name not in self._flag_cache:
sub = self._choices.get(name)
flags = {}
if sub is not None:
for action in sub._actions:
for opt in action.option_strings:
if opt.startswith("--"):
flags[opt] = action
self._flag_cache[name] = flags
return self._flag_cache[name]
def option(self, name, token):
"""The action a flag token names on one command, short form included."""
sub = self._choices.get(name)
return option_action(sub, token) if sub is not None else None
def global_option(self, token):
"""The action a flag token names among the global options."""
return option_action(self._parser, token)
def positionals(self, name):
sub = self._choices.get(name)
return [a for a in sub._actions if not a.option_strings] if sub is not None else []
def takes_positional(self, name):
return bool(self.positionals(name))
def positional_completer(self, name, index=0):
"""The completer for a command's index-th positional, if it has one.
``apwrap._add_completer`` tags id/machine/ssh positionals with a
completer function; that is what turns ``destroy instance <tab>`` into a
list of live instance ids, and the second argument of
``update ssh-key <id> <tab>`` into local .pub paths instead.
"""
positionals = self.positionals(name)
if not positionals:
return None
if index >= len(positionals):
# A trailing nargs='+'/'*' positional keeps accepting values.
last = positionals[-1]
return getattr(last, "completer", None) if last.nargs in ("+", "*") else None
return getattr(positionals[index], "completer", None)
cli/repl/completion.py
"""Tab completion for the REPL.
Inside our own readline loop the completer sees the whole line, so completion
can be staged the way the command tree actually is:
sh<tab> -> show, search, set, ssh-url, ... (verbs)
show <tab> -> instances, user, invoices, ... (objects)
show instances --<tab> -> --raw, --quiet, ... (flags)
destroy instance <tab> -> live instance ids (values)
It is also the fast path: `vastai` under argcomplete forks a fresh interpreter
per Tab press, while here the parser is already in memory and ids are cached,
so completion is a dict lookup.
"""
import argparse
import threading
import time
from vastai.cli.repl.catalog import CommandCatalog
VALUE_CACHE_TTL = 30.0 # seconds a fetched id list stays warm
COMPLETION_BUDGET = 2.0 # seconds a Tab press will wait for that fetch
class LiveValues:
"""Cache over the id completers the parser attaches to positionals.
Each lookup would otherwise be an API round trip on a keystroke, so the
full list is fetched once per completer and filtered locally. Failures are
cached too — a down API or a missing key must not hang every Tab press.
"""
def __init__(self, ttl=VALUE_CACHE_TTL, clock=time.monotonic,
budget=COMPLETION_BUDGET):
self._ttl = ttl
self._clock = clock
self._budget = budget
self._cache = {}
self._pending = set()
self._generation = 0
def clear(self):
"""Forget every cached list — called when the session's credentials
change, so Tab can't offer the previous account's ids."""
self._cache.clear()
self._pending.clear()
self._generation += 1 # discard whatever is still in flight
def matching(self, completer, prefix):
cached = self._cache.get(completer)
if cached is None or self._clock() - cached[0] > self._ttl:
self._fetch(completer)
cached = self._cache.get(completer)
if cached is None:
return [] # still in flight; a later Tab press will have it
return [v for v in cached[1] if v.startswith(prefix)]
def _fetch(self, completer):
"""Fetch in a worker, and wait only a moment for it.
The API client allows 120s per request and retries, so fetching on the
input thread would let one Tab press freeze the prompt for minutes
against an unreachable endpoint. Whatever the worker eventually returns
lands in the cache for the next press.
"""
if completer in self._pending:
return
self._pending.add(completer)
generation = self._generation
def work():
try:
values = [str(v) for v in (completer(prefix="") or [])]
except Exception:
values = []
if generation != self._generation:
return # the account changed while we were fetching
self._cache[completer] = (self._clock(), values)
self._pending.discard(completer)
worker = threading.Thread(target=work, daemon=True)
worker.start()
worker.join(self._budget)
class ReplCompleter:
"""A readline completer over the live command tree.
``suggestions`` is the whole of the logic and is pure with respect to
readline (it takes the line, not the terminal state); ``complete`` is just
the readline protocol adapter around it.
"""
def __init__(self, catalog, values=None):
self.catalog = catalog if isinstance(catalog, CommandCatalog) else CommandCatalog(catalog)
self.values = values if values is not None else LiveValues()
self.matches = []
# -- readline protocol -------------------------------------------------
def complete(self, text, state):
if state == 0:
try:
import readline
buf = readline.get_line_buffer()[: readline.get_endidx()]
except Exception:
buf = text
try:
self.matches = self.suggestions(buf, text)
except Exception:
self.matches = [] # a broken completion must never break typing
try:
return self.matches[state]
except IndexError:
return None
def install(self):
"""Bind this completer to readline. No-op where readline is missing."""
try:
import readline
except ImportError:
return False
readline.set_completer(self.complete)
# Split on whitespace only, so the completer sees '--flag' and
# '--flag=value' as single tokens and decides what they mean.
readline.set_completer_delims(" \t\n")
if "libedit" in (getattr(readline, "__doc__", "") or ""):
readline.parse_and_bind("bind ^I rl_complete") # macOS libedit
else:
readline.parse_and_bind("tab: complete")
return True
# -- the actual logic --------------------------------------------------
def suggestions(self, buf, text=None):
"""Candidates for the token being typed at the end of ``buf``."""
if text is None:
text = "" if (not buf or buf[-1].isspace()) else buf.split()[-1]
stripped = buf.lstrip()
tokens = stripped.split()
typed = tokens if (not stripped or buf[-1:].isspace()) else tokens[:-1]
typed = self.catalog.strip_options(typed) # `--raw show <tab>` completes too
if not typed:
return self._starting_with(self.catalog.first_words, text)
if len(typed) == 1 and typed[0] in self.catalog.verbs and not text.startswith("-"):
return self._starting_with(self.catalog.objects(typed[0]), text)
name = self.catalog.resolve(typed)
if name is None:
return []
if text.startswith("-"):
return self._starting_with(self.catalog.flags(name), text)
return self._values(name, typed[len(name.split()):], text)
def _values(self, name, args, text):
"""Complete an argument value: the choices or ids of whichever option or
positional the cursor sits on."""
option, index = self._position(name, args)
if option is not None:
if option.choices:
return self._starting_with([str(c) for c in option.choices], text)
return self._live(getattr(option, "completer", None), text)
return self._live(self.catalog.positional_completer(name, index), text)
def _position(self, name, args):
"""Walk a command's arguments as argparse would.
Returns the option still awaiting values (the cursor is typing one of
them) and the index of the positional being typed — so a second
positional gets its own completer, and a multi-value option keeps
offering its choices.
"""
option, index, i = None, 0, 0
while i < len(args):
token = args[i]
i += 1
if token.startswith("-") and token != "-":
option = None
action = self.catalog.option(name, token)
if action is None or action.nargs == 0 or "=" in token:
continue
taken = 0
while (i < len(args) and _accepts_more(action, taken)
and not args[i].startswith("-")):
taken += 1
i += 1
if _accepts_more(action, taken):
option = action # still hungry: the cursor's word is its value
continue
option = None
index += 1
return option, index
def _live(self, completer, text):
return self.values.matching(completer, text) if completer else []
@staticmethod
def _starting_with(candidates, text):
return sorted(c for c in candidates if c.startswith(text))
def _accepts_more(action, taken):
"""Whether an option can still absorb another value after `taken` of them."""
nargs = action.nargs
if nargs in ("+", "*", argparse.REMAINDER):
return True
if nargs == "?":
return taken < 1
if isinstance(nargs, int):
return taken < nargs
return taken < 1 # nargs None: exactly one value
cli/repl/session.py
"""The `vastai repl` loop: the same CLI, one live process.
vast> show instances
vast> search offers 'gpu_name=RTX_4090 num_gpus=1' -o dph
vast> destroy instance <tab> # completes from your live instances
Every line is handed to the same parser `vastai` uses, so behaviour is
identical to the one-shot CLI — only faster, because the interpreter and the
command tree load once. Type `help` for the commands, `clear` to clear the
screen, `exit` or Ctrl-D to leave.
"""
import copy
import os
import shutil
import sys
import textwrap
from vastai.cli.repl.bridge import run_line
from vastai.cli.repl.catalog import CommandCatalog
from vastai.cli.repl.completion import ReplCompleter
from vastai.cli.util import APIKEY_FILE, TFAKEY_FILE
PROMPT = "vast> "
EXIT_WORDS = ("exit", "quit", "q")
CLEAR_WORDS = ("clear", "cls")
CLEAR_SCREEN = "\033[H\033[2J\033[3J" # cursor home, clear screen, clear scrollback
class Repl:
def __init__(self, parser, session_args, stdout=None):
self.parser = parser
# A copy: the session's own state is not the caller's.
self.args = copy.copy(session_args)
self.out = stdout or sys.stdout
self._stored_key = _stored_api_key()
# Whether this session's key came from the config file. A key from
# $VAST_API_KEY or --api-key outranks the file (the CLI says so, and
# `set api-key` warns about it), so it must not be replaced by one.
self._from_disk = getattr(session_args, "api_key", None) in (None, self._stored_key)
self.catalog = CommandCatalog(parser)
self.completer = ReplCompleter(self.catalog)
# -- one line ----------------------------------------------------------
def handle(self, line):
"""Run one input line. Returns False when the session should end."""
s = line.strip()
if not s:
return True
if s in EXIT_WORDS:
return False
if s in CLEAR_WORDS:
self._clear_screen()
return True
# Leading global flags (`--raw show user`) are the parser's business,
# not part of the command name.
tokens = self.catalog.strip_options(s.split())
if tokens and tokens[0] == "repl":
self._print("Already in the REPL. Use exit or Ctrl-D to leave.")
return True
# Resolve first so a typo gets a short 'did you mean' instead of
# argparse dumping all ~150 command names. A line that is only flags
# (`--version`) names no command and goes straight to the parser.
if tokens and self.catalog.resolve(tokens) is None:
self._unknown(tokens)
return True
run_line(self.parser, s, self.args)
self._refresh_credentials()
return True
def _unknown(self, tokens):
"""Report an unrecognised line. argparse would answer with all ~150
command names; a verb's own objects, or a few near misses, is the
answer worth reading."""
if len(tokens) == 1 and tokens[0] in self.catalog.verbs:
self._print_wrapped(f"{tokens[0]} takes an object: ",
self.catalog.objects(tokens[0]))
return
print(f"unknown command: {' '.join(tokens[:2])}", file=sys.stderr)
suggestions = self.catalog.suggest(tokens)[:6]
if suggestions:
self._print_wrapped("did you mean: ", suggestions)
print("(help lists every command)", file=sys.stderr)
def _print_wrapped(self, lead, items):
width = shutil.get_terminal_size(fallback=(80, 24)).columns
print(textwrap.fill(", ".join(items), width=max(40, width),
initial_indent=lead, subsequent_indent=" " * len(lead)),
file=sys.stderr)
def _clear_screen(self):
# Legacy Windows consoles don't honour ANSI escapes; cls always works there.
if os.name == "nt":
os.system("cls")
else:
self.out.write(CLEAR_SCREEN)
self.out.flush()
def _print(self, text):
print(text, file=self.out)
def _refresh_credentials(self):
"""Pick up a key written mid-session by `set api-key` or `tfa login`.
Those commands write the config file but not our namespace, so without
this the session would keep sending the key it started with — often
none at all, leaving every later line to fail on auth. Only a session
that took its key from the file follows the file; one running on
$VAST_API_KEY or --api-key keeps what it was given, as the one-shot CLI
would. A key that disappears is dropped rather than reused: an expired
2FA session with no fallback should report "no API key", not repeat the
same failure on every line.
"""
if not self._from_disk:
return
key = _stored_api_key()
if key == self._stored_key:
return
self.args.api_key = key
self._stored_key = key
self.completer.values.clear() # ids belong to the old account
# -- the loop ----------------------------------------------------------
def run(self):
self._setup_terminal()
while True:
try:
line = input(PROMPT)
except KeyboardInterrupt:
self._print("^C")
continue
except EOFError:
self._print("")
break
if not self.handle(line):
break
def _setup_terminal(self):
"""Wire completion. History lives in readline for the session only —
persisting it would put whatever was typed, credentials included, in
plain text on disk, which wants a redaction design of its own."""
self._wire_live_completions()
self.completer.install()
def _wire_live_completions(self):
"""Point the parser's id completers at this session.
The one-shot CLI wires these for argcomplete using only $VAST_API_KEY;
rebinding them here means `destroy instance <tab>` works for a key that
came from the config file or from `--api-key` on the repl command.
"""
from vastai.cli.commands import instances
from vastai.cli.parser import set_completers
def ids(field):
return lambda **kw: instances.show__instances(
self._completion_args(), {"internal": True, "field": field})
set_completers(instance_machine_fn=ids("machine_id"), instance_fn=ids("id"))
def _completion_args(self):
"""Args for completion's own API calls: this session's auth, endpoint
and retries, with the output modes off. `--curl` in particular makes
the client print a curl command and exit, which would otherwise land in
the middle of the prompt on every Tab press.
"""
args = copy.copy(self.args)
for flag in ("explain", "curl", "raw", "full"):
setattr(args, flag, False)
args.retry = 1 # a Tab press must not queue minutes of retries
return args
def _stored_api_key():
"""The key on disk right now, resolved as the one-shot CLI resolves it."""
path = TFAKEY_FILE if os.path.exists(TFAKEY_FILE) else APIKEY_FILE
try:
with open(path, "r") as reader:
return reader.read().strip() or None
except OSError:
return None
def run_repl(args):
from vastai.cli.main import parser
Repl(parser, args).run()
cli/self_test/__init__.py
"""Helpers for CLI self-test diagnostics."""
cli/self_test/machine_diagnostics.py
"""Structured diagnostics for ``vastai self-test machine``."""
from vastai.cli.util import required_inet_mbps
DIAGNOSTICS_VERSION = "1"
ROOT_CURRENTLY_RENTED = "currently_rented"
ROOT_DEVERIFIED_OR_BELOW_THRESHOLD = "deverified_or_below_threshold"
ROOT_API_PERMISSION_FAILED = "api_permission_failed"
ROOT_ZERO_ACTIVE_OFFERS = "zero_active_offers"
ROOT_OFFLINE_OR_NOT_LISTED = "offline_or_not_listed"
ROOT_UNKNOWN_NO_RENTABLE_OFFER = "unknown_no_rentable_offer"
NO_OFFER_ROOT_STATES = (
ROOT_CURRENTLY_RENTED,
ROOT_DEVERIFIED_OR_BELOW_THRESHOLD,
ROOT_API_PERMISSION_FAILED,
ROOT_ZERO_ACTIVE_OFFERS,
ROOT_OFFLINE_OR_NOT_LISTED,
ROOT_UNKNOWN_NO_RENTABLE_OFFER,
)
# Very large GPU hosts, such as 8x B300 systems, can exceed 2 TB of total VRAM.
# Once a host has about 2 TB of system RAM, do not disqualify it only because it
# is slightly below 95% of total VRAM.
SYSTEM_RAM_REQUIREMENT_CAP_MIB = 2_000_000
def safe_float(value):
if value is None:
return 0.0
try:
return float(value)
except (ValueError, TypeError):
return 0.0
def per_gpu_vram_gib(offer):
"""Return per-GPU VRAM in GiB from canonical offer fields when possible."""
gpu_total_ram = safe_float(offer.get("gpu_total_ram"))
num_gpus = safe_float(offer.get("num_gpus"))
if gpu_total_ram > 0 and num_gpus > 0:
return gpu_total_ram / num_gpus / 1024
gpu_ram = safe_float(offer.get("gpu_ram"))
if gpu_ram >= 1024:
return gpu_ram / 1024
return gpu_ram
def compact_offer_metadata(offer):
if not offer:
return None
fields = (
"id",
"machine_id",
"gpu_name",
"num_gpus",
"dph_total",
"dlperf",
"cuda_max_good",
"compute_cap",
"reliability",
"direct_port_count",
"pcie_bw",
"inet_down",
"inet_up",
"gpu_ram",
"gpu_total_ram",
"cpu_ram",
"cpu_cores",
"rentable",
"rented",
"vericode",
"verified",
"verification",
"error_description",
)
return {field: offer.get(field) for field in fields if field in offer}
def base_result(machine_id):
return {
"success": False,
"reason": "",
"warning": None,
"phase": "preflight",
"diagnostics_version": DIAGNOSTICS_VERSION,
"machine_id": machine_id,
"offer": None,
"checks": [],
"failure": None,
"failure_code": None,
"stage": None,
"error": None,
"diagnostics": {},
}
def _check(
check_id,
title,
actual,
required,
operator,
unit,
passed,
purpose,
remediation,
):
status = "pass" if passed else "fail"
return {
"id": check_id,
"title": title,
"status": status,
"actual": actual,
"required": required,
"operator": operator,
"unit": unit,
"summary": f"{title}: actual {actual} {unit}, required {operator} {required} {unit}",
"purpose": purpose,
"remediation": remediation,
}
def _info_check(
check_id,
title,
actual,
recommended,
operator,
unit,
summary,
purpose,
remediation,
):
return {
"id": check_id,
"title": title,
"status": "info",
"actual": actual,
"required": recommended,
"operator": operator,
"unit": unit,
"summary": summary,
"purpose": purpose,
"remediation": remediation,
}
def preflight_requirement_checks(offer):
gpu_total_ram = safe_float(offer.get("gpu_total_ram"))
per_gpu_ram_gib = per_gpu_vram_gib(offer)
required_mbps = required_inet_mbps(gpu_total_ram)
cpu_ram = safe_float(offer.get("cpu_ram"))
cpu_cores = int(safe_float(offer.get("cpu_cores")))
num_gpus = int(safe_float(offer.get("num_gpus")))
listed_gpus = num_gpus if num_gpus > 0 else 1
direct_port_count = safe_float(offer.get("direct_port_count"))
required_min_ports = 3 * listed_gpus
recommended_max_ports = 64 * listed_gpus
uncapped_required_cpu_ram = 0.95 * gpu_total_ram
required_cpu_ram = min(uncapped_required_cpu_ram, SYSTEM_RAM_REQUIREMENT_CAP_MIB)
required_cpu_cores = 2 * num_gpus
checks = [
_check(
"cuda.version",
"CUDA version",
safe_float(offer.get("cuda_max_good")),
11.8,
">=",
"CUDA",
safe_float(offer.get("cuda_max_good")) >= 11.8,
"The self-test image needs a host driver stack compatible with CUDA 11.8 or newer.",
"Update the NVIDIA driver/CUDA stack, then confirm with: vastai search offers "
f"'machine_id={offer.get('machine_id') or 'unknown'} rentable=any rented=any'",
),
_check(
"reliability",
"Reliability",
safe_float(offer.get("reliability")),
0.90,
">",
"ratio",
safe_float(offer.get("reliability")) > 0.90,
"Self-test uses reliability as a guardrail before launching a temporary instance.",
"Let the host stabilize and resolve recent failures before retrying the self-test.",
),
_check(
"network.direct_ports",
"Direct port count",
direct_port_count,
required_min_ports,
">=",
"ports",
direct_port_count >= required_min_ports,
(
"The tester needs at least 3 directly mapped ports per listed GPU for remote "
"progress, SSH checks, and runtime port allocation."
),
(
f"Open at least {required_min_ports} direct ports for {listed_gpus} listed GPU(s), "
"check firewall/NAT forwarding, and make TCP/UDP forwarding symmetric where required."
),
),
_check(
"pcie.bandwidth",
"PCIe bandwidth",
safe_float(offer.get("pcie_bw")),
2.85,
">",
"GB/s",
safe_float(offer.get("pcie_bw")) > 2.85,
"Low PCIe bandwidth can make GPU stress and transfer checks fail or time out.",
"Check BIOS PCIe generation/lane settings and confirm GPUs are seated in full-speed slots.",
),
_check(
"network.download",
"Download speed",
safe_float(offer.get("inet_down")),
round(required_mbps, 2),
">=",
"Mb/s",
safe_float(offer.get("inet_down")) >= required_mbps,
"The bandwidth floor scales with total VRAM so large GPU hosts can complete data movement tests.",
"Improve host download bandwidth or reduce contention, then rerun the Vast host verification.",
),
_check(
"network.upload",
"Upload speed",
safe_float(offer.get("inet_up")),
round(required_mbps, 2),
">=",
"Mb/s",
safe_float(offer.get("inet_up")) >= required_mbps,
"The tester needs enough upload bandwidth to report progress and complete network checks.",
"Improve host upload bandwidth or reduce contention, then rerun the Vast host verification.",
),
_check(
"gpu.ram",
"GPU RAM",
round(per_gpu_ram_gib, 2),
7,
">",
"GiB",
per_gpu_ram_gib > 7,
"The verification workload requires more than 7 GB of VRAM per GPU.",
"Use a GPU with more VRAM for this self-test.",
),
_check(
"system.ram",
"System RAM",
cpu_ram,
round(required_cpu_ram, 2),
">=",
"MiB",
cpu_ram >= required_cpu_ram,
(
"System RAM must be close to total VRAM so CPU-side staging does not starve the tests. "
"For very large GPU hosts, this requirement is capped at about 2 TB."
),
(
"Add system RAM or reduce the listed GPU set so system RAM is at least 95% of "
"total VRAM, up to the 2 TB cap."
),
),
_check(
"cpu.cores",
"CPU cores",
cpu_cores,
required_cpu_cores,
">=",
"cores",
cpu_cores >= required_cpu_cores,
"The tester expects at least two CPU cores per GPU for stable orchestration.",
"Expose more CPU cores to the host or reduce the GPU count for this offer.",
),
]
if direct_port_count > recommended_max_ports:
checks.append(
_info_check(
"network.direct_ports.recommended_max",
"Direct port count advisory",
direct_port_count,
recommended_max_ports,
"<=",
"ports",
(
f"Direct port count: actual {direct_port_count} ports, recommended <= "
f"{recommended_max_ports} ports for {listed_gpus} GPU(s)"
),
(
"Vast instances can use at most 64 open ports each. Mapping more than "
"64 ports per listed GPU is usually wasted effort."
),
(
"This is advisory only, not a self-test gate. Keep enough direct ports for "
"self-test and normal workloads, but avoid mapping very large unused ranges."
),
)
)
return checks
def failed_checks(checks):
return [check for check in checks if check.get("status") == "fail"]
def informational_checks(checks):
return [check for check in checks if check.get("status") == "info"]
def _http_status(error):
if not error:
return None
if isinstance(error, dict):
return error.get("status_code")
response = getattr(error, "response", None)
return getattr(response, "status_code", None)
def _machine_lookup_status(machine_lookup):
if not machine_lookup:
return None, []
if isinstance(machine_lookup, dict) and "status" in machine_lookup:
status = machine_lookup.get("status")
rows = machine_lookup.get("rows") or []
if isinstance(rows, dict):
rows = [rows] if rows else []
return status, rows
if isinstance(machine_lookup, list):
return ("visible", machine_lookup) if machine_lookup else ("empty", [])
if isinstance(machine_lookup, dict):
return ("visible", [machine_lookup]) if machine_lookup else ("empty", [])
return None, []
def diagnose_no_offer_state(machine_id, broader_offers, machine_lookup=None, search_error=None):
"""Infer the most likely no-rentable root state from CLI-visible data.
This is intentionally heuristic. Public offer search alone cannot prove
whether a machine is offline, hidden, or has zero active offers, so the
return value includes a confidence field for UI/host copy.
"""
offers = broader_offers or []
evidence = []
suggested_steps = [
f"Inspect all visible offers with: vastai search offers 'machine_id={machine_id} rentable=any rented=any'."
]
search_status = _http_status(search_error)
lookup_status, lookup_rows = _machine_lookup_status(machine_lookup)
lookup_status_code = machine_lookup.get("status_code") if isinstance(machine_lookup, dict) else None
if search_status in (401, 403) or lookup_status in ("permission_denied", "unauthorized") or lookup_status_code in (401, 403):
if search_status:
evidence.append(f"Offer search failed with HTTP {search_status}.")
if lookup_status_code:
evidence.append(f"Machine lookup failed with HTTP {lookup_status_code}.")
evidence.append("The API key or account could not read the required machine/offer state.")
suggested_steps = [
"Check that the API key is valid and has permission to read offers and hosted machines.",
"Retry after logging in with an account that owns or can inspect this host.",
]
return {
"root_state": ROOT_API_PERMISSION_FAILED,
"confidence": "high",
"evidence": evidence,
"suggested_steps": suggested_steps,
}
if offers:
evidence.append(f"Found {len(offers)} visible offer(s), but none were currently rentable.")
rented_offers = [offer for offer in offers if offer.get("rented") is True]
if rented_offers:
evidence.append(f"{len(rented_offers)} visible offer(s) are marked rented=true.")
suggested_steps = [
"Wait for the current rental to end, then retry the self-test.",
suggested_steps[0],
]
return {
"root_state": ROOT_CURRENTLY_RENTED,
"confidence": "medium",
"evidence": evidence,
"suggested_steps": suggested_steps,
}
state_evidence = []
for offer in offers:
reliability = safe_float(offer.get("reliability"))
vericode = offer.get("vericode")
verified = offer.get("verified")
verification = offer.get("verification")
error_description = offer.get("error_description")
if reliability and reliability <= 0.90:
state_evidence.append(f"Offer reliability is {reliability}, which is at or below the 0.90 self-test threshold.")
if vericode not in (None, "", 0, "0"):
state_evidence.append(f"Offer reports vericode={vericode}.")
if verified is False:
state_evidence.append("Offer is marked verified=false.")
if isinstance(verification, str) and verification.lower() not in ("", "verified", "true"):
state_evidence.append(f"Offer verification state is {verification}.")
if error_description:
state_evidence.append(f"Offer error_description is {error_description}.")
if state_evidence:
evidence.extend(state_evidence)
suggested_steps = [
"Resolve the host verification or reliability issue, then wait for the offer state to refresh.",
suggested_steps[0],
]
return {
"root_state": ROOT_DEVERIFIED_OR_BELOW_THRESHOLD,
"confidence": "medium",
"evidence": evidence,
"suggested_steps": suggested_steps,
}
evidence.append("Visible offers exist, but their payload did not expose a specific non-rentable reason.")
suggested_steps = [
"Refresh the marketplace/host state and retry shortly.",
suggested_steps[0],
]
return {
"root_state": ROOT_UNKNOWN_NO_RENTABLE_OFFER,
"confidence": "low",
"evidence": evidence,
"suggested_steps": suggested_steps,
}
evidence.append("No visible on-demand offers were returned for this machine.")
if lookup_status == "visible" and lookup_rows:
evidence.append("Machine lookup returned a visible machine record.")
suggested_steps = [
"List or relist an on-demand offer for this machine, then retry self-test.",
"Check host listing state with: vastai show machines.",
]
return {
"root_state": ROOT_ZERO_ACTIVE_OFFERS,
"confidence": "medium",
"evidence": evidence,
"suggested_steps": suggested_steps,
}
if lookup_status in ("empty", "not_found") or lookup_status_code == 404:
if lookup_status_code == 404:
evidence.append("Machine lookup returned HTTP 404.")
else:
evidence.append("Machine lookup returned no visible machine record.")
elif lookup_status:
evidence.append(f"Machine lookup status was {lookup_status}.")
else:
evidence.append("No authoritative machine lookup result was available.")
suggested_steps = [
"Confirm the machine ID is correct and the host is online.",
"If this is your host, check that it is listed and visible with: vastai show machines.",
suggested_steps[0],
]
return {
"root_state": ROOT_OFFLINE_OR_NOT_LISTED,
"confidence": "low",
"evidence": evidence,
"suggested_steps": suggested_steps,
}
def no_offer_failure(machine_id, broader_offers, machine_lookup=None, search_error=None):
diagnosis = diagnose_no_offer_state(
machine_id,
broader_offers,
machine_lookup=machine_lookup,
search_error=search_error,
)
broad_count = len(broader_offers or [])
if diagnosis["root_state"] == ROOT_API_PERMISSION_FAILED:
summary = f"Could not inspect rentable offers for machine {machine_id} because API authorization failed."
causes = [
"The API key may be invalid, expired, or missing permission to read offer/machine state.",
]
code = "api_permission_failed"
remediation = "Check the API key permissions, then retry with an account that can inspect this host."
elif broad_count:
summary = f"No currently rentable on-demand offer found for machine {machine_id}."
causes = [
"The machine may be listed but not rentable right now.",
"The offer may be unavailable to new rentals because host state, verification, or pricing changed.",
"The machine may already be occupied or temporarily unavailable.",
]
code = "no_rentable_offer"
remediation = (
"Check host state with: vastai show machines. Then inspect offers with: "
f"vastai search offers 'machine_id={machine_id} rentable=any rented=any'."
)
else:
summary = f"No on-demand offer found for machine {machine_id}."
causes = [
"The machine may be offline.",
"The machine may not have an active listed offer.",
"The machine ID may be incorrect or not visible to this account.",
]
code = "no_offer"
remediation = (
"Check host state with: vastai show machines. Then inspect offers with: "
f"vastai search offers 'machine_id={machine_id} rentable=any rented=any'."
)
check = {
"id": "offer.available",
"title": "Rentable offer available",
"status": "fail",
"actual": 0,
"required": 1,
"operator": ">=",
"unit": "offers",
"summary": summary,
"purpose": "The self-test needs a selected offer so it can rent one temporary diagnostic instance.",
"remediation": remediation,
}
failure = {
"code": code,
"summary": summary,
"likely_causes": causes,
"remediation": remediation,
"root_state": diagnosis["root_state"],
"confidence": diagnosis["confidence"],
"evidence": diagnosis["evidence"],
"suggested_steps": diagnosis["suggested_steps"],
}
return check, failure
def requirement_failure(checks):
failures = failed_checks(checks)
summary = f"{len(failures)} preflight requirement check(s) failed."
return {
"code": "preflight_requirements_failed",
"summary": summary,
"failed_check_ids": [check["id"] for check in failures],
"remediation": "Resolve the failed checks below, or rerun with --ignore-requirements to dogfood anyway.",
}
def render_preflight_failure(machine_id, checks, failure=None, print_fn=print):
print_fn(f"Preflight diagnostics for machine {machine_id} failed:")
for check in failed_checks(checks):
print_fn(f"- {check['title']}")
print_fn(f" actual: {check['actual']} {check['unit']}")
print_fn(f" required: {check['operator']} {check['required']} {check['unit']}")
print_fn(f" purpose: {check['purpose']}")
print_fn(f" remediation: {check['remediation']}")
if failure and failure.get("likely_causes"):
print_fn("Likely causes:")
for cause in failure["likely_causes"]:
print_fn(f"- {cause}")
if failure and failure.get("root_state"):
print_fn(f"Root state: {failure['root_state']} (confidence: {failure.get('confidence', 'unknown')})")
evidence = failure.get("evidence") if failure else None
if evidence:
print_fn("Evidence:")
for item in evidence:
print_fn(f"- {item}")
steps = failure.get("suggested_steps") if failure else None
if steps:
print_fn("Suggested steps:")
for step in steps:
print_fn(f"- {step}")
def render_preflight_advisories(machine_id, checks, print_fn=print):
info_checks = informational_checks(checks)
if not info_checks:
return
print_fn(f"Preflight advisory for machine {machine_id}:")
for check in info_checks:
print_fn(f"- {check['title']}")
print_fn(f" actual: {check['actual']} {check['unit']}")
print_fn(f" recommended: {check['operator']} {check['required']} {check['unit']}")
print_fn(f" purpose: {check['purpose']}")
print_fn(f" note: {check['remediation']}")
cli/self_test/runtime_diagnostics.py
"""Runtime diagnostic helpers for CLI self-test machine flows.
The functions in this module are intentionally pure so the CLI can use them for
raw output shaping without coupling parser/classifier behavior to live Vast API
calls or the existing self-test orchestration.
"""
from __future__ import annotations
from dataclasses import dataclass
import re
from typing import Iterable
INSTANCE_CREATE_FAILED = "instance_create_failed"
INSTANCE_CREATE_MISSING_CONTRACT = "instance_create_missing_contract"
INSTANCE_STATUS_ERROR = "instance_status_error"
INSTANCE_STATUS_POLL_FAILED = "instance_status_poll_failed"
INSTANCE_START_TIMEOUT = "instance_start_timeout"
INSTANCE_OFFLINE_BEFORE_TEST = "instance_offline_before_test"
MISSING_PUBLIC_IP = "missing_public_ip"
PROGRESS_PORT_NOT_MAPPED = "progress_port_not_mapped"
PROGRESS_ENDPOINT_UNREACHABLE = "progress_endpoint_unreachable"
PROGRESS_ENDPOINT_LOST = "progress_endpoint_lost"
PROGRESS_EMPTY_TIMEOUT = "progress_empty_timeout"
RUNTIME_TEST_TIMEOUT = "runtime_test_timeout"
LEGACY_PROGRESS_ERROR = "legacy_progress_error"
DOCKER_PULL_FAILED = "docker_pull_failed"
DAEMON_STARTUP_FAILED = "daemon_startup_failed"
NVML_FAILED = "nvml_failed"
RESNET_FAILED = "resnet_failed"
ECC_FAILED = "ecc_failed"
NCCL_FAILED = "nccl_failed"
STRESS_GPU_BURN_FAILED = "stress_gpu_burn_failed"
INTERRUPTED = "interrupted"
CLEANUP_FAILED = "cleanup_failed"
PROGRESS_CONTAINER_PORT = "5000/tcp"
RUNTIME_FAILURE_CODES = (
INSTANCE_CREATE_FAILED,
INSTANCE_CREATE_MISSING_CONTRACT,
INSTANCE_STATUS_ERROR,
INSTANCE_STATUS_POLL_FAILED,
INSTANCE_START_TIMEOUT,
INSTANCE_OFFLINE_BEFORE_TEST,
MISSING_PUBLIC_IP,
PROGRESS_PORT_NOT_MAPPED,
PROGRESS_ENDPOINT_UNREACHABLE,
PROGRESS_ENDPOINT_LOST,
PROGRESS_EMPTY_TIMEOUT,
RUNTIME_TEST_TIMEOUT,
LEGACY_PROGRESS_ERROR,
DOCKER_PULL_FAILED,
DAEMON_STARTUP_FAILED,
NVML_FAILED,
RESNET_FAILED,
ECC_FAILED,
NCCL_FAILED,
STRESS_GPU_BURN_FAILED,
INTERRUPTED,
CLEANUP_FAILED,
)
@dataclass(frozen=True)
class FailureCatalogEntry:
code: str
summary: str
remediation: str
suggested_steps: tuple[str, ...]
FAILURE_CATALOG: dict[str, FailureCatalogEntry] = {
INSTANCE_CREATE_FAILED: FailureCatalogEntry(
INSTANCE_CREATE_FAILED,
"Failed to create the runtime test instance.",
"Check the offer, docker image, and instance creation response.",
("Retry with --debugging enabled.", "Inspect the create-instance API error."),
),
INSTANCE_CREATE_MISSING_CONTRACT: FailureCatalogEntry(
INSTANCE_CREATE_MISSING_CONTRACT,
"Instance creation did not return a new contract id.",
"Treat the create response as malformed or incomplete.",
("Inspect the raw create-instance response.", "Retry after confirming the offer is still rentable."),
),
INSTANCE_STATUS_ERROR: FailureCatalogEntry(
INSTANCE_STATUS_ERROR,
"The instance reported an error while starting.",
"Inspect the instance status message and host/container logs.",
("Run show instance for the contract.", "Check docker logs on the host if available."),
),
INSTANCE_STATUS_POLL_FAILED: FailureCatalogEntry(
INSTANCE_STATUS_POLL_FAILED,
"Failed to poll instance status.",
"Confirm API connectivity and retry the status check.",
("Retry the CLI command.", "Check network/API errors from the status request."),
),
INSTANCE_START_TIMEOUT: FailureCatalogEntry(
INSTANCE_START_TIMEOUT,
"The instance did not reach running state before timeout.",
"Check host capacity, docker startup, and network configuration.",
("Inspect instance status_msg.", "Try the test again after confirming the host is healthy."),
),
INSTANCE_OFFLINE_BEFORE_TEST: FailureCatalogEntry(
INSTANCE_OFFLINE_BEFORE_TEST,
"The instance went offline before the runtime test could run.",
"Investigate host availability and instance lifecycle events.",
("Check machine status.", "Review host daemon and container logs."),
),
MISSING_PUBLIC_IP: FailureCatalogEntry(
MISSING_PUBLIC_IP,
"The running instance did not expose a public IP address.",
"Confirm the instance network configuration and public IP assignment.",
("Inspect show instance output.", "Retry on a machine with public networking available."),
),
PROGRESS_PORT_NOT_MAPPED: FailureCatalogEntry(
PROGRESS_PORT_NOT_MAPPED,
"The runtime progress port was not mapped.",
"Confirm port 5000/tcp is exposed and direct ports are available.",
("Check the available mapped ports in the diagnostic output.", "Verify the machine has enough direct ports."),
),
PROGRESS_ENDPOINT_UNREACHABLE: FailureCatalogEntry(
PROGRESS_ENDPOINT_UNREACHABLE,
"The runtime progress endpoint was never reachable.",
"Check TCP firewall/NAT forwarding, direct port mapping, container startup, and NAT hairpinning.",
(
"Confirm the mapped public TCP port forwards to the host LAN IP.",
"Inspect docker logs to confirm the progress server bound port 5000/tcp.",
"If testing from the same LAN as the host, retry from an outside network to rule out NAT loopback/hairpinning.",
),
),
PROGRESS_ENDPOINT_LOST: FailureCatalogEntry(
PROGRESS_ENDPOINT_LOST,
"The runtime progress endpoint became unreachable after connecting.",
"Look for container crashes, OOM, GPU errors, or host instability.",
(
"Inspect docker logs for a crash or missing progress server.",
"Check dmesg for Xid, GPU reset, OOM, or host stall messages.",
"Check for network loss between the CLI and host public endpoint.",
),
),
PROGRESS_EMPTY_TIMEOUT: FailureCatalogEntry(
PROGRESS_EMPTY_TIMEOUT,
"The progress endpoint returned no new output before timeout.",
"Check whether the runtime script stalled or stopped writing progress.",
("Inspect runtime logs.", "Retry with debugging enabled."),
),
RUNTIME_TEST_TIMEOUT: FailureCatalogEntry(
RUNTIME_TEST_TIMEOUT,
"The runtime test did not complete before timeout.",
"Investigate long-running or stalled test stages.",
("Check the last reported progress stage.", "Run individual tests to isolate the stall."),
),
LEGACY_PROGRESS_ERROR: FailureCatalogEntry(
LEGACY_PROGRESS_ERROR,
"The legacy runtime progress stream reported an unclassified error.",
"Use the raw error text and active stage to decide the next action.",
("Inspect the original ERROR line.", "Retry with debugging enabled if the cause is unclear."),
),
DOCKER_PULL_FAILED: FailureCatalogEntry(
DOCKER_PULL_FAILED,
"The test image could not be pulled.",
"Check image name, tag availability, registry access, and credentials.",
("Verify the docker image tag exists.", "Check for registry unauthorized or denied errors."),
),
DAEMON_STARTUP_FAILED: FailureCatalogEntry(
DAEMON_STARTUP_FAILED,
"The container or daemon failed during startup.",
"Inspect docker daemon, OCI runtime, and container startup logs.",
("Check docker logs.", "Verify NVIDIA container runtime and host daemon health."),
),
NVML_FAILED: FailureCatalogEntry(
NVML_FAILED,
"NVML or nvidia-smi failed during system checks.",
"Check NVIDIA driver, NVML, and GPU visibility on the host.",
("Run nvidia-smi on the host.", "Check driver/library mismatch or GPU reset errors."),
),
RESNET_FAILED: FailureCatalogEntry(
RESNET_FAILED,
"The ResNet runtime test failed.",
"Check PyTorch/CUDA health, available VRAM, and GPU compute stability.",
("Look for CUDA OOM, cuDNN, or torch runtime errors.", "Run a smaller isolated torch workload."),
),
ECC_FAILED: FailureCatalogEntry(
ECC_FAILED,
"The ECC runtime test failed.",
"Check GPU ECC counters and hardware health.",
("Inspect ECC error counters.", "Review dmesg and nvidia-smi health output."),
),
NCCL_FAILED: FailureCatalogEntry(
NCCL_FAILED,
"The NCCL distributed runtime test failed.",
"Check multi-GPU connectivity, NCCL transport, and network fabric.",
("Inspect NCCL error output.", "Verify peer-to-peer and multi-GPU communication."),
),
STRESS_GPU_BURN_FAILED: FailureCatalogEntry(
STRESS_GPU_BURN_FAILED,
"The stress-ng or gpu-burn runtime test failed.",
"Check thermals, power stability, GPU Xid errors, and host stress logs.",
("Inspect dmesg for Xid errors.", "Review gpu-burn and stress-ng output."),
),
INTERRUPTED: FailureCatalogEntry(
INTERRUPTED,
"The runtime test was interrupted.",
"Ensure cleanup completed or destroy the test instance manually.",
("Check whether the test instance still exists.", "Destroy leaked instances if needed."),
),
CLEANUP_FAILED: FailureCatalogEntry(
CLEANUP_FAILED,
"Runtime test cleanup failed.",
"Destroy the temporary test instance manually to avoid continued billing.",
("Run destroy instance for the temporary contract.", "Retry cleanup after checking API connectivity."),
),
}
STAGE_SYSTEM_REQUIREMENTS = "system_requirements"
STAGE_RESNET = "resnet"
STAGE_ECC = "ecc"
STAGE_NCCL = "nccl"
STAGE_STRESS_GPU_BURN = "stress_gpu_burn"
STAGE_STARTUP = "startup"
_STAGE_PATTERNS: tuple[tuple[re.Pattern[str], str], ...] = (
(re.compile(r"^\s*Running system requirements test\.\.\.\s*$", re.IGNORECASE), STAGE_SYSTEM_REQUIREMENTS),
(re.compile(r"^\s*Running ResNet50/ResNet18 test\.\.\.\s*$", re.IGNORECASE), STAGE_RESNET),
(re.compile(r"^\s*Running ECC test\.\.\.\s*$", re.IGNORECASE), STAGE_ECC),
(re.compile(r"^\s*Running NCCL distributed test\.\.\.\s*$", re.IGNORECASE), STAGE_NCCL),
(re.compile(r"^\s*Running stress-ng and gpu-burn\.\.\.\s*$", re.IGNORECASE), STAGE_STRESS_GPU_BURN),
)
_NVML_RE = re.compile(
r"nvml|nvidia-smi|driver/library version mismatch|failed to initialize.*nvidia",
re.IGNORECASE,
)
_RESNET_RE = re.compile(
r"resnet|torch|pytorch|cuda out of memory|outofmemory|cudnn|cublas|cuda error|runtimeerror",
re.IGNORECASE,
)
_ECC_RE = re.compile(r"\becc\b|volatile double bit|aggregate single bit", re.IGNORECASE)
_NCCL_RE = re.compile(r"\bnccl\b|unhandled system error|connection timed out|allreduce|peer access", re.IGNORECASE)
_STRESS_RE = re.compile(r"stress-ng|gpu-burn|xid|thermal|power limit|burn-in|hardware error", re.IGNORECASE)
_DOCKER_PULL_RE = re.compile(
r"pull|manifest|not found|unauthorized|denied|repository does not exist|no such image|"
r"pull access denied|requested access to the resource is denied",
re.IGNORECASE,
)
_STARTUP_RE = re.compile(
r"daemon|container|startup|start container|failed to start|oci runtime|runc|nvidia-container|"
r"docker|exited|exec format error|permission denied|mount|entrypoint",
re.IGNORECASE,
)
_API_KEY_RE = re.compile(r"([?&]api_key=)[^&\s]+")
def redact_secret_text(value: object) -> str | None:
"""Return a host-visible string with obvious API-key query values redacted."""
if value is None:
return None
return _API_KEY_RE.sub(r"\1REDACTED", str(value))
def _mapped_port_names(mapped_ports) -> list[str]:
if not mapped_ports:
return []
if isinstance(mapped_ports, dict):
return sorted(str(key) for key in mapped_ports.keys())
return sorted(str(port) for port in mapped_ports)
def make_progress_endpoint_diagnostic(
*,
url: str | None = None,
public_ip: str | None = None,
container_port: str = PROGRESS_CONTAINER_PORT,
host_port: str | int | None = None,
timeout_seconds: int | float | None = None,
attempt_count: int = 0,
first_connection_established: bool = False,
last_error_type: str | None = None,
last_error: object = None,
last_status_code: int | None = None,
mapped_ports=None,
) -> dict[str, object]:
"""Shape progress endpoint state for raw output and UI consumption."""
if url is None and public_ip and host_port:
url = f"https://{public_ip}:{host_port}/progress"
diagnostic: dict[str, object] = {
"url": url,
"public_ip": public_ip,
"container_port": container_port,
"external_port": str(host_port) if host_port is not None else None,
"host_port": str(host_port) if host_port is not None else None,
"timeout_seconds": timeout_seconds,
"attempt_count": int(attempt_count or 0),
"first_connection_established": bool(first_connection_established),
"last_error_type": last_error_type,
"last_error": redact_secret_text(last_error),
"last_status_code": last_status_code,
"mapped_ports": _mapped_port_names(mapped_ports),
}
return diagnostic
def failure_catalog() -> dict[str, dict[str, object]]:
"""Return a JSON-serializable copy of the runtime failure catalog."""
return {
code: {
"code": entry.code,
"summary": entry.summary,
"remediation": entry.remediation,
"suggested_steps": list(entry.suggested_steps),
}
for code, entry in FAILURE_CATALOG.items()
}
def get_failure_entry(code: str) -> FailureCatalogEntry:
"""Return the catalog entry for a stable runtime failure code."""
try:
return FAILURE_CATALOG[code]
except KeyError as exc:
raise ValueError(f"Unknown runtime failure code: {code}") from exc
def make_failure(
code: str,
*,
stage: str | None = None,
summary: str | None = None,
details: str | None = None,
error: str | None = None,
remediation: str | None = None,
suggested_steps: Iterable[str] | None = None,
underlying_error: str | None = None,
progress_endpoint: dict[str, object] | None = None,
) -> dict[str, object]:
"""Build a raw-output-friendly diagnostic dictionary."""
entry = get_failure_entry(code)
diagnostic: dict[str, object] = {
"code": code,
"stage": stage,
"summary": summary or entry.summary,
"remediation": remediation or entry.remediation,
"suggested_steps": list(suggested_steps) if suggested_steps is not None else list(entry.suggested_steps),
}
if details:
diagnostic["details"] = details
if error:
diagnostic["error"] = error
if underlying_error:
diagnostic["underlying_error"] = underlying_error
if progress_endpoint:
diagnostic["progress_endpoint"] = progress_endpoint
return diagnostic
def stage_from_progress_line(line: str) -> str | None:
"""Return the legacy runtime stage introduced by a progress line."""
for pattern, stage in _STAGE_PATTERNS:
if pattern.search(line):
return stage
return None
def classify_legacy_error_line(line: str, stage: str | None = None) -> dict[str, object]:
"""Classify a legacy progress ``ERROR`` line into a runtime diagnostic."""
stripped = line.strip()
lowered_stage = stage.lower() if stage else None
code = LEGACY_PROGRESS_ERROR
if _NVML_RE.search(stripped):
code = NVML_FAILED
elif _NCCL_RE.search(stripped) or lowered_stage == STAGE_NCCL:
code = NCCL_FAILED
elif _ECC_RE.search(stripped) or lowered_stage == STAGE_ECC:
code = ECC_FAILED
elif _STRESS_RE.search(stripped) or lowered_stage == STAGE_STRESS_GPU_BURN:
code = STRESS_GPU_BURN_FAILED
elif _RESNET_RE.search(stripped) or lowered_stage == STAGE_RESNET:
code = RESNET_FAILED
return make_failure(
code,
stage=stage,
error=stripped,
details=f"Legacy progress stream reported: {stripped}",
underlying_error=stripped,
)
class LegacyProgressParser:
"""Stateful parser for the legacy runtime progress text stream."""
def __init__(self) -> None:
self.stage: str | None = None
def process_line(self, line: str) -> dict[str, object] | None:
"""Update parser state and return a diagnostic for ``ERROR`` lines."""
stage = stage_from_progress_line(line)
if stage:
self.stage = stage
return None
if line.strip().upper().startswith("ERROR"):
return classify_legacy_error_line(line, self.stage)
return None
def parse(self, text: str) -> list[dict[str, object]]:
"""Parse progress text and return diagnostics for any ``ERROR`` lines."""
diagnostics: list[dict[str, object]] = []
for line in text.splitlines():
diagnostic = self.process_line(line)
if diagnostic is not None:
diagnostics.append(diagnostic)
return diagnostics
def parse_legacy_progress(text: str) -> list[dict[str, object]]:
"""Parse a legacy progress payload with a fresh parser instance."""
return LegacyProgressParser().parse(text)
def classify_status_msg(status_msg: str | None) -> dict[str, object] | None:
"""Classify startup/status messages reported while the instance starts."""
if not status_msg:
return None
msg = status_msg.strip()
if not msg:
return None
code = INSTANCE_STATUS_ERROR
if _DOCKER_PULL_RE.search(msg):
code = DOCKER_PULL_FAILED
elif _STARTUP_RE.search(msg):
code = DAEMON_STARTUP_FAILED
return make_failure(
code,
stage=STAGE_STARTUP,
error=msg,
details=f"Instance status message reported: {msg}",
underlying_error=msg,
)
__all__ = [
"CLEANUP_FAILED",
"DAEMON_STARTUP_FAILED",
"DOCKER_PULL_FAILED",
"ECC_FAILED",
"FAILURE_CATALOG",
"FailureCatalogEntry",
"INSTANCE_CREATE_FAILED",
"INSTANCE_CREATE_MISSING_CONTRACT",
"INSTANCE_OFFLINE_BEFORE_TEST",
"INSTANCE_START_TIMEOUT",
"INSTANCE_STATUS_ERROR",
"INSTANCE_STATUS_POLL_FAILED",
"INTERRUPTED",
"LEGACY_PROGRESS_ERROR",
"LegacyProgressParser",
"MISSING_PUBLIC_IP",
"NCCL_FAILED",
"NVML_FAILED",
"PROGRESS_EMPTY_TIMEOUT",
"PROGRESS_CONTAINER_PORT",
"PROGRESS_ENDPOINT_LOST",
"PROGRESS_ENDPOINT_UNREACHABLE",
"PROGRESS_PORT_NOT_MAPPED",
"RESNET_FAILED",
"RUNTIME_FAILURE_CODES",
"RUNTIME_TEST_TIMEOUT",
"STAGE_ECC",
"STAGE_NCCL",
"STAGE_RESNET",
"STAGE_STARTUP",
"STAGE_STRESS_GPU_BURN",
"STAGE_SYSTEM_REQUIREMENTS",
"STRESS_GPU_BURN_FAILED",
"classify_legacy_error_line",
"classify_status_msg",
"failure_catalog",
"get_failure_entry",
"make_failure",
"make_progress_endpoint_diagnostic",
"parse_legacy_progress",
"redact_secret_text",
"stage_from_progress_line",
]
cli/self_test/support_bundle.py
"""Support bundle helpers for host self-test failures."""
from __future__ import annotations
import json
import io
import os
import re
import subprocess
import tarfile
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from vastai.cli.self_test.runtime_diagnostics import redact_secret_text
DEFAULT_BUNDLE_DIR = "/tmp"
MAX_TEXT_BYTES = 256 * 1024
MAX_LOG_BYTES = 256 * 1024
MAX_COMMAND_SECONDS = 10
JSON_ARTIFACT_RE = re.compile(r"(^|/)[^/]+\.json$")
SENSITIVE_KEY_RE = re.compile(
r"(api[_-]?key|token|secret|password|passwd|authorization|jupyter_token|instance_api_key|host_port_range)",
re.IGNORECASE,
)
SENSITIVE_ASSIGNMENT_RE = re.compile(
r"(?i)\b(api[_-]?key|token|secret|password|passwd|authorization|jupyter_token|instance_api_key|host_port_range)"
r"\b\s*[:=]\s*([^\s,;]+)"
)
DMESG_KEYWORDS_RE = re.compile(r"\b(warn|warning|err|error|fail|fault|crit|panic|oops|bug)\b", re.IGNORECASE)
def utc_timestamp() -> str:
return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
def support_bundles_enabled() -> bool:
value = os.environ.get("VAST_SELF_TEST_SUPPORT_BUNDLE", "")
return value.strip().lower() not in ("0", "false", "no", "off")
def _safe_machine_id(machine_id: object) -> str:
return re.sub(r"[^A-Za-z0-9_.-]+", "_", str(machine_id or "unknown"))[:80]
def _redact_text(text: object, secrets: list[str] | None = None) -> str:
redacted = redact_secret_text("" if text is None else str(text)) or ""
for secret in secrets or []:
if secret:
redacted = redacted.replace(str(secret), "REDACTED")
return SENSITIVE_ASSIGNMENT_RE.sub(lambda match: f"{match.group(1)}=REDACTED", redacted)
def _redact_json(value: Any, secrets: list[str] | None = None) -> Any:
if isinstance(value, dict):
return {
str(key): "REDACTED" if SENSITIVE_KEY_RE.search(str(key)) else _redact_json(item, secrets)
for key, item in value.items()
}
if isinstance(value, list):
return [_redact_json(item, secrets) for item in value]
if isinstance(value, str):
return _redact_text(value, secrets)
return value
def _truncate_text(text: str, max_bytes: int) -> str:
encoded = text.encode("utf-8", errors="replace")
if len(encoded) <= max_bytes:
return text
tail = encoded[-max_bytes:].decode("utf-8", errors="replace")
return f"[truncated to last {max_bytes} bytes]\n{tail}"
def _decode_bytes(data: bytes) -> str:
return data.decode("utf-8", errors="replace")
def _read_file_tail(path: Path, max_bytes: int) -> str:
size = path.stat().st_size
with path.open("rb") as handle:
if size > max_bytes:
handle.seek(-max_bytes, os.SEEK_END)
data = handle.read()
text = _decode_bytes(data)
if size > max_bytes:
return f"[truncated from {size} bytes to last {max_bytes} bytes]\n{text}"
return text
def _run_command(args: list[str], *, timeout: int = MAX_COMMAND_SECONDS) -> tuple[str, str | None]:
try:
proc = subprocess.run(args, capture_output=True, text=True, timeout=timeout, check=False)
except FileNotFoundError:
return "", f"command not found: {args[0]}"
except subprocess.TimeoutExpired:
return "", f"command timed out after {timeout}s: {' '.join(args)}"
output = proc.stdout or ""
if proc.stderr:
output = f"{output}\n[stderr]\n{proc.stderr}".strip()
error = None
if proc.returncode != 0:
error = f"command exited {proc.returncode}: {' '.join(args)}"
return output, error
def _safe_archive_name(name: str) -> str:
return re.sub(r"(^/+|\.\.|\\)", "_", name)
def _summarize_ss_output(text: str) -> str:
lines = [line for line in text.splitlines() if line.strip()]
if not lines:
return ""
summarized = []
for idx, line in enumerate(lines):
parts = line.split()
if idx == 0 and parts:
summarized.append(" ".join(parts[:5]))
continue
if len(parts) >= 5:
summarized.append(" ".join(parts[:5]))
else:
summarized.append(line)
return "\n".join(summarized)
def _collect_host_artifacts(secrets: list[str] | None = None) -> tuple[dict[str, str], list[dict[str, str]]]:
files: dict[str, str] = {}
errors: list[dict[str, str]] = []
for path in sorted(Path("/var/lib/vastai_kaalia").glob("kaalia.log*")):
if not path.is_file():
continue
archive_name = f"host/kaalia/{path.name}"
try:
files[archive_name] = _redact_text(_read_file_tail(path, MAX_LOG_BYTES), secrets)
except Exception as exc:
errors.append({"artifact": archive_name, "error": _redact_text(exc, secrets)})
daemon_json = Path("/etc/docker/daemon.json")
if daemon_json.exists():
try:
files["host/etc-docker-daemon.json"] = _redact_text(_read_file_tail(daemon_json, MAX_TEXT_BYTES), secrets)
except Exception as exc:
errors.append({"artifact": "host/etc-docker-daemon.json", "error": _redact_text(exc, secrets)})
else:
errors.append({"artifact": "host/etc-docker-daemon.json", "error": "file not found"})
command_specs = [
("host/dmesg-filtered.log", ["dmesg", "-T"], "dmesg"),
(
"host/journalctl-docker-vastai_kaalia.log",
["journalctl", "-b", "-p", "err..alert", "-u", "docker", "-u", "vastai_kaalia", "--since", "-24h", "--no-pager"],
"journalctl",
),
("host/nvidia-smi.txt", ["nvidia-smi"], "nvidia-smi"),
("host/nvidia-smi-L.txt", ["nvidia-smi", "-L"], "nvidia-smi -L"),
("host/docker-info.txt", ["docker", "info"], "docker info"),
("host/ip-addr.txt", ["ip", "addr"], "ip addr"),
("host/ss-tlnp-summary.txt", ["ss", "-tlnp"], "ss -tlnp summary"),
]
for archive_name, command, label in command_specs:
output, error = _run_command(command)
if archive_name == "host/dmesg-filtered.log":
output = "\n".join(line for line in output.splitlines() if DMESG_KEYWORDS_RE.search(line))
elif archive_name == "host/ss-tlnp-summary.txt":
output = _summarize_ss_output(output)
files[archive_name] = _truncate_text(_redact_text(output, secrets), MAX_TEXT_BYTES)
if error:
errors.append({"artifact": label, "error": _redact_text(error, secrets)})
mounts = Path("/proc/mounts")
if mounts.exists():
try:
mount_lines = [
line for line in _read_file_tail(mounts, MAX_TEXT_BYTES).splitlines()
if "/var/lib/docker" in line
]
files["host/docker-mounts.txt"] = _redact_text("\n".join(mount_lines), secrets)
except Exception as exc:
errors.append({"artifact": "host/docker-mounts.txt", "error": _redact_text(exc, secrets)})
else:
errors.append({"artifact": "host/docker-mounts.txt", "error": "/proc/mounts not found"})
return files, errors
def create_support_bundle(
*,
machine_id: object,
output_dir: str | None = None,
result: dict[str, Any] | None = None,
cli_output: list[str] | None = None,
extra_files: dict[str, str] | None = None,
extra_errors: list[dict[str, str]] | None = None,
run_started_at: str | None = None,
command: list[str] | None = None,
secrets: list[str] | None = None,
include_local_host_artifacts: bool = False,
include_host_logs: bool | None = None,
) -> dict[str, Any]:
"""Create a redacted self-test support tarball and return metadata."""
if include_host_logs is not None:
include_local_host_artifacts = include_host_logs
created_at = utc_timestamp()
safe_id = _safe_machine_id(machine_id)
bundle_dir = Path(output_dir or DEFAULT_BUNDLE_DIR).expanduser()
bundle_dir.mkdir(parents=True, exist_ok=True)
bundle_path = bundle_dir / f"vast_selftest_{safe_id}_{created_at}.tar.gz"
files: dict[str, str] = {}
errors: list[dict[str, str]] = []
if cli_output:
files["self-test-output.log"] = _redact_text("\n".join(cli_output), secrets)
if result is not None:
files["self-test-result.json"] = json.dumps(_redact_json(result, secrets), indent=2, sort_keys=True)
if extra_files:
for name, content in extra_files.items():
files[name] = _redact_text(content, secrets)
if extra_errors:
errors.extend(_redact_json(extra_errors, secrets))
if include_local_host_artifacts:
host_files, host_errors = _collect_host_artifacts(secrets)
files.update(host_files)
errors.extend(host_errors)
manifest = {
"bundle_version": 1,
"machine_id": str(machine_id),
"created_at_utc": created_at,
"run_started_at_utc": run_started_at,
"command": _redact_json(command or [], secrets),
"max_text_bytes_per_artifact": MAX_TEXT_BYTES,
"max_log_bytes_per_artifact": MAX_LOG_BYTES,
"includes_local_host_artifacts": include_local_host_artifacts,
"files": sorted(files.keys()) + ["manifest.json", "collection-errors.json"],
"collection_errors": errors,
"note": (
"This bundle contains CLI-visible self-test evidence. Host OS artifacts "
"are included only when includes_local_host_artifacts=true; otherwise "
"host daemon/kaalia/system logs need host-side daemon/backend collection "
"or an explicit dump run on the actual host."
),
}
files["manifest.json"] = json.dumps(manifest, indent=2, sort_keys=True)
files["collection-errors.json"] = json.dumps(errors, indent=2, sort_keys=True)
with tarfile.open(bundle_path, "w:gz") as tar:
for name, content in sorted(files.items()):
safe_name = _safe_archive_name(name)
artifact_text = content if JSON_ARTIFACT_RE.search(name) else _truncate_text(content, MAX_LOG_BYTES)
data = artifact_text.encode("utf-8", errors="replace")
info = tarfile.TarInfo(safe_name)
info.size = len(data)
info.mtime = int(time.time())
tar.addfile(info, fileobj=io.BytesIO(data))
os.chmod(bundle_path, 0o600)
size_bytes = bundle_path.stat().st_size
return {
"path": str(bundle_path),
"created_at_utc": created_at,
"size_bytes": size_bytes,
"files": sorted(files.keys()),
"collection_errors": errors,
}
def format_bundle_summary(bundle: dict[str, Any]) -> list[str]:
lines = [
"Self-test diagnostic bundle saved to:",
f" {bundle.get('path')}",
"Bundle contents:",
]
for name in bundle.get("files", []):
lines.append(f" - {name}")
errors = bundle.get("collection_errors") or []
if errors:
lines.append("Some artifacts could not be collected:")
for item in errors[:10]:
lines.append(f" - {item.get('artifact')}: {item.get('error')}")
if len(errors) > 10:
lines.append(f" - ... {len(errors) - 10} more")
lines.append("Review this tarball before sharing it with support.")
return lines
cli/selfupdate.py
"""Self-update engine for the managed (curl | bash) CLI install.
Layout (docs/install-design.md, XDG Base Directory Specification): the
install root defaults to $XDG_DATA_HOME/vastai (~/.local/share/vastai),
override with $VASTAI_INSTALL_DIR — a single active venv at current/, with
bin/vastai a fixed symlink into it. Updating rebuilds current/ in place
(build temp → verify → swap); there is no version retention.
"Managed install?" is detected structurally (``is_managed_install``): a managed
CLI runs from <root>/current with a sibling <root>/bin/uv. That's ground
truth — no on-disk marker to drift or hand-copy. pip installs fail the check
and the updater never touches them.
The passive nudge is throttled to one manifest GET and one notice per week,
capped at 1s, and swallows every failure — the CLI must never get slower or
noisier because the manifest endpoint is unreachable.
"""
import json
import os
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path
import requests
from vastai.cli.util import DATA_HOME, DIRS, VERSION
DEFAULT_MANIFEST_URL = "https://vast.ai/cli/manifest.json"
INSTALL_SH_HINT = "curl -fsSL https://vast.ai/install.sh | bash"
PIP_UPGRADE_HINT = "pip install --upgrade vastai"
# Oldest version whose bundled selfupdate.py understands this layout
# (<root>/current + <root>/bin/uv, XDG-based root). Versions before this pin
# assumed ~/.vastai/env and would install fine but then see themselves as unmanaged,
# permanently losing self-update until the install is rebuilt from scratch.
MIN_DOWNGRADE_VERSION = "1.4.0"
UPDATE_CHECK_FILE = os.path.join(DIRS['state'], "update_check.json")
CHECK_INTERVAL_S = 7 * 24 * 60 * 60
NUDGE_TIMEOUT_S = 1.0
# Console scripts shipped in the wheel; each gets a swapped symlink in bin/.
MANAGED_BINARIES = ("vastai", "serve-vast-deployment", "register-python-argcomplete")
# Fixed, not XDG-overridable — every other XDG-aware CLI lands its symlink here too.
LOCAL_BIN = Path.home() / ".local" / "bin"
class UpdateError(Exception):
"""A self-update step failed; the current install is untouched."""
class UninstallError(Exception):
"""A step of vastai uninstall failed; nothing further was attempted."""
# ---------------------------------------------------------------------------
# Install detection / manifest / versions
# ---------------------------------------------------------------------------
def install_root() -> Path:
"""Where the managed install lives.
Prefers ground truth from where this interpreter is actually running: if
we're already inside a managed install (``<root>/current`` next to
``<root>/bin/uv``), that's authoritative and immune to
$VASTAI_INSTALL_DIR/$XDG_DATA_HOME differing between install time and
now — e.g. a custom install dir that isn't set in the current shell would
otherwise make a real managed install look unmanaged. Only when we're
*not* currently running from one (a pip install checking
``is_managed_install``) do we fall back to computing where a managed
install would live.
"""
prefix = Path(sys.prefix).resolve()
if prefix.name == "current" and (prefix.parent / "bin" / "uv").exists():
return prefix.parent
if os.environ.get("VASTAI_INSTALL_DIR"):
return Path(os.environ["VASTAI_INSTALL_DIR"]).expanduser()
return Path(DATA_HOME) / "vastai"
def is_managed_install() -> bool:
"""True iff this CLI was installed by the managed installer.
Detected from where the interpreter actually runs: a managed CLI lives in
``<root>/current`` next to ``<root>/bin/uv``. Ground truth, no marker file.
pip installs run from elsewhere and fail the check.
"""
try:
root = install_root()
return (
Path(sys.prefix).resolve() == (root / "current").resolve()
and (root / "bin" / "uv").exists()
)
except Exception:
return False
def fetch_manifest(timeout: float = 10.0) -> dict:
url = os.environ.get("VASTAI_MANIFEST_URL") or DEFAULT_MANIFEST_URL
try:
r = requests.get(url, timeout=timeout)
r.raise_for_status()
data = r.json()
except (requests.exceptions.RequestException, ValueError) as e:
raise UpdateError(f"Could not fetch release manifest from {url}: {e}")
if not isinstance(data, dict) or "latest" not in data:
raise UpdateError(f"Malformed release manifest at {url}")
if data.get("schema") != 1:
raise UpdateError(
f"Release manifest schema {data.get('schema')!r} is not supported "
f"by this CLI version.\nRe-run the installer: {INSTALL_SH_HINT}"
)
return data
def version_key(v: str):
"""Sort key tolerant of dev/rc (below release), post (above), +local (ignored)."""
v = (v or "").strip().lstrip("v").split("+", 1)[0]
nums, extra = [], 0
for part in v.split("."):
if part.isdigit():
nums.append(int(part))
continue
m = re.match(r"\d+", part)
if m:
nums.append(int(m.group()))
extra = 1 if part.startswith("post") else -1
break
while len(nums) < 3:
nums.append(0)
return tuple(nums), extra
def is_newer(candidate: str, current: str) -> bool:
return version_key(candidate) > version_key(current)
# ---------------------------------------------------------------------------
# Passive nudge
# ---------------------------------------------------------------------------
def _load_check_state() -> dict:
try:
with open(UPDATE_CHECK_FILE) as f:
data = json.load(f)
return data if isinstance(data, dict) else {}
except (OSError, ValueError):
return {}
def _save_check_state(state: dict) -> None:
try:
with open(UPDATE_CHECK_FILE, "w") as f:
json.dump(state, f)
except OSError:
pass
def _stderr_is_tty() -> bool:
try:
return sys.stderr.isatty()
except Exception:
return False
def notify_update(args=None) -> None:
"""Pre-command, best-effort upgrade nudge. Must never raise or block."""
try:
_notify_update(args)
except Exception:
pass
def _notify_update(args) -> None:
suppressed = (
os.environ.get("VASTAI_NO_UPDATE_CHECK")
or os.environ.get("CI")
or getattr(args, "raw", False)
or not _stderr_is_tty()
# the updater itself doesn't need a nudge
or getattr(getattr(args, "func", None), "__module__", "") == "vastai.cli.commands.update"
)
if suppressed:
return
now = time.time()
state = _load_check_state()
if now - state.get("checked_at", 0) > CHECK_INTERVAL_S:
state["checked_at"] = now # stamped even on failure, so errors back off too
try:
state["latest"] = fetch_manifest(timeout=NUDGE_TIMEOUT_S).get("latest")
except UpdateError:
pass
_save_check_state(state)
latest = state.get("latest")
if not latest or not is_newer(latest, VERSION):
return
if now - state.get("notified_at", 0) <= CHECK_INTERVAL_S:
return
hint = "vastai update" if is_managed_install() else PIP_UPGRADE_HINT
arrow = "↑" if "utf" in (sys.stderr.encoding or "").lower() else "*"
# Trailing blank line separates the nudge from the command output that follows it.
print(
f"{arrow} vastai {latest} is available (you have {VERSION}). Run `{hint}`.\n",
file=sys.stderr,
)
state["notified_at"] = now
_save_check_state(state)
# ---------------------------------------------------------------------------
# Update / rollback (managed installs only)
# ---------------------------------------------------------------------------
def _run(cmd, *, env=None):
cmd = [str(c) for c in cmd]
try:
return subprocess.run(cmd, env=env, check=True, capture_output=True, text=True)
except subprocess.CalledProcessError as e:
# Surface the tool's own stderr — that's the real error (e.g. uv's
# resolution failure), not a category we invent.
detail = (e.stderr or e.stdout or "").strip()
raise UpdateError(f"Command failed: {' '.join(cmd)}" + (f"\n{detail}" if detail else ""))
except OSError as e:
# Carry the actual OSError through (wrong path, missing binary, …)
# rather than relabeling everything "Required tool not found".
raise UpdateError(f"Could not run {' '.join(cmd)}: {e}")
def _prune_uv_cache(uv: Path, env: dict) -> None:
"""Best-effort prune of uv's cache after a successful swap.
Each update hardlinks the new env out of the cache, so the cache
accumulates every version's wheels forever unless pruned. Prune (not
clean): the cache is the user-global one, shared with any other uv on
the machine. The update already succeeded — never fail because of this.
"""
try:
_run([uv, "cache", "prune", "--quiet"], env=env)
except UpdateError:
pass
def _link_binaries(root: Path) -> None:
"""Point bin/<name> at the fixed current/bin/<name> for each shipped
console script. The target path is constant across updates (always
``current/``), so this is not symlink *retargeting* — just (re)creating
stable links."""
for name in MANAGED_BINARIES:
if not (root / "current" / "bin" / name).exists():
continue
link, tmp = root / "bin" / name, root / "bin" / f".{name}.tmp"
if tmp.is_symlink() or tmp.exists():
tmp.unlink()
os.symlink(Path("..") / "current" / "bin" / name, tmp)
os.replace(tmp, link)
def wheel_spec(target: str, manifest: dict) -> str:
"""Requirement spec that installs ``target``.
Latest installs the manifest's release wheel: hash-pinned bytes straight
from the GitHub Release, no PyPI propagation window. Other targets are
pins/rollbacks to long-published versions, where the PyPI pin is
race-free and covers releases that predate wheel_url.
"""
install = manifest.get("install") or {}
url, sha = install.get("wheel_url"), install.get("wheel_sha256")
if target == manifest.get("latest") and url and sha:
return f"vastai @ {url}#sha256={sha}"
return f"vastai=={target}"
def perform_update(target: str, manifest: dict) -> None:
"""Install ``target`` into a fresh env, verify it, then swap it in.
Single active install (the deno model): the new env is built in a temp dir
and only swapped over the live one after it verifies, so an interrupted
update can never leave a half-written or broken ``vastai``. There is no
version retention — a re-pin or rollback is just ``update --version X``,
which reinstalls. Wheel integrity comes from ``wheel_spec``; we
additionally confirm the installed version matches ``target``.
"""
if version_key(target) < version_key(MIN_DOWNGRADE_VERSION):
raise UpdateError(
f"Can't roll back to {target}: versions before {MIN_DOWNGRADE_VERSION} "
"don't recognize this install's layout and would lose the ability "
f"to self-update. The oldest version you can roll back to is {MIN_DOWNGRADE_VERSION}."
)
root = install_root()
uv = root / "bin" / "uv"
if not uv.exists():
raise UpdateError(f"Bootstrap tool missing at {uv}.\nRe-run the installer: {INSTALL_SH_HINT}")
python_pin = (manifest.get("install") or {}).get("python") or "3.12"
env_dir, new_dir, old_dir = root / "current", root / ".current.new", root / ".current.old"
shutil.rmtree(new_dir, ignore_errors=True)
uv_env = dict(os.environ,
UV_PYTHON_INSTALL_DIR=str(root / "python"),
UV_PYTHON_PREFERENCE="only-managed")
try:
# --relocatable: shebangs must not hardcode the build path (env is renamed into place).
_run([uv, "venv", new_dir, "--python", python_pin, "--relocatable", "--quiet"], env=uv_env)
_run([uv, "pip", "install", "--python", new_dir / "bin" / "python",
"--quiet", wheel_spec(target, manifest)], env=uv_env)
installed = (_run([new_dir / "bin" / "vastai", "--version"]).stdout or "").strip()
if target not in installed:
raise UpdateError(
f"Verification failed: new install reports version {installed!r}, expected {target}"
)
except UpdateError:
shutil.rmtree(new_dir, ignore_errors=True)
raise
# Swap via two renames (~instant); a crash mid-build only dirties .current.new.
shutil.rmtree(old_dir, ignore_errors=True)
if env_dir.exists():
os.replace(env_dir, old_dir)
os.replace(new_dir, env_dir)
_link_binaries(root)
shutil.rmtree(old_dir, ignore_errors=True)
_prune_uv_cache(uv, uv_env)
# ---------------------------------------------------------------------------
# Uninstall (managed installs only)
# ---------------------------------------------------------------------------
def perform_uninstall() -> Path:
"""Remove the managed install: the install root and its bin/ symlinks.
Config, cache, and state (``vastai.cli.util.DIRS`` — each independently
XDG-overridable) are left alone — the same guarantee documented for a
manual ``rm -rf`` (docs/install-design.md §3). Safe to call while running
from inside ``root``: like the update swap above, removing the directory
doesn't disturb an already-running interpreter on POSIX.
Raises ``UninstallError`` on any real failure (permissions, a symlink
loop, …) instead of silently reporting success — callers should treat
this as fallible, not a fire-and-forget best-effort cleanup.
"""
root = install_root().resolve()
for name in MANAGED_BINARIES:
link = LOCAL_BIN / name
try:
if link.is_symlink() and root in link.resolve().parents:
link.unlink()
except OSError as e:
raise UninstallError(f"Could not remove {link}: {e}")
try:
shutil.rmtree(root)
except FileNotFoundError:
pass # already gone — nothing left to remove
except OSError as e:
raise UninstallError(f"Could not remove {root}: {e}")
return root
cli/util.py
"""
CLI utility functions and constants for the Vast.ai CLI.
Extracted from vast.py - contains version checking, config directory setup,
constants, and various helper functions used by CLI commands.
"""
from __future__ import unicode_literals, print_function
import re
import json
import sys
import argparse
import os
import time
import math
import subprocess
import shutil
import requests
import getpass
from pathlib import Path
from datetime import date, datetime, timedelta, timezone
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, List, Optional
# Re-export string_to_unix_epoch so CLI command modules can import from here
from vastai.api.query import string_to_unix_epoch # noqa: F401
from vastai.utils import VERSION, parse_version, parse_env, smart_split # noqa: F401
# Server URL default — canonical definition lives in api/client.py
from vastai.api.client import server_url_default # noqa: E402, F401
# Sentinel for API key argument default (to distinguish "not provided" from None)
api_key_guard = object()
# ---------------------------------------------------------------------------
# App constants
# ---------------------------------------------------------------------------
APP_NAME = "vastai"
# Define emoji support and fallbacks
_HAS_EMOJI = sys.stdout.encoding and 'utf' in sys.stdout.encoding.lower()
SUCCESS = "\u2705" if _HAS_EMOJI else "[OK]"
WARN = "\u26a0\ufe0f" if _HAS_EMOJI else "[!]"
FAIL = "\u274c" if _HAS_EMOJI else "[X]"
INFO = "\u2139\ufe0f" if _HAS_EMOJI else "[i]"
# ---------------------------------------------------------------------------
# Config directory setup
# ---------------------------------------------------------------------------
try:
# Although xdg-base-dirs is the newer name, there's
# python compatibility issues with dependencies that
# can be unresolvable using things like python 3.9
# So we actually use the older name, thus older
# version for now. This is as of now (2024/11/15)
# the safer option. -cjm
import xdg
DIRS = {
'config': xdg.xdg_config_home(),
'temp': xdg.xdg_cache_home(),
'state': xdg.xdg_state_home(),
}
# Not part of DIRS: DIRS entries are this CLI's own runtime data and get
# auto-created below regardless of install method. DATA_HOME is where a
# *managed install* lives (selfupdate.install_root()) — a pip install
# must never side-effect that directory into existence just by running.
DATA_HOME = xdg.xdg_data_home()
except Exception:
# Reasonable defaults.
from pathlib import Path
_home = str(Path.home())
DIRS = {
'config': os.path.join(_home, '.config'),
'temp': os.path.join(_home, '.cache'),
'state': os.path.join(_home, '.local', 'state'),
}
DATA_HOME = os.path.join(_home, '.local', 'share')
for key in DIRS.keys():
DIRS[key] = path = os.path.join(DIRS[key], APP_NAME)
if not os.path.exists(path):
os.makedirs(path)
CACHE_FILE = os.path.join(DIRS['temp'], "gpu_names_cache.json")
GPU_TYPES_CACHE_FILE = os.path.join(DIRS['temp'], "gpu_types_cache.json")
CACHE_DURATION = timedelta(hours=24)
def _get_gpu_names() -> Optional[List[str]]:
"""Returns a set of GPU names available on Vast.ai, with results cached for 24 hours."""
def is_cache_valid() -> bool:
"""Checks if the cache file exists and is less than 24 hours old."""
if not os.path.exists(CACHE_FILE):
return False
cache_age = datetime.now() - datetime.fromtimestamp(os.path.getmtime(CACHE_FILE))
return cache_age < CACHE_DURATION
if is_cache_valid():
try:
with open(CACHE_FILE, "r") as file:
gpu_names = json.load(file)
except (OSError, json.JSONDecodeError):
gpu_names = None
else:
endpoint = "/api/v0/gpu_names/unique/"
url = f"{server_url_default}{endpoint}"
try:
r = requests.get(url, headers={})
r.raise_for_status()
gpu_names = r.json()
except (requests.exceptions.RequestException, ValueError):
return None
try:
with open(CACHE_FILE, "w") as file:
json.dump(gpu_names, file)
except OSError:
pass
try:
return [
name.replace(" ", "_").replace("-", "_") for name in gpu_names['gpu_names']
]
except (TypeError, KeyError):
return None
def _get_gpu_types() -> Optional[List[dict]]:
"""Returns the GPU type catalog from /api/v0/gpu_types/, cached for 24 hours.
Source of truth for canonical GPU names and per-card VRAM. Returns None on
any error so callers fall back to the hardcoded constants (never empty).
"""
def is_cache_valid() -> bool:
"""Checks if the cache file exists and is less than 24 hours old."""
if not os.path.exists(GPU_TYPES_CACHE_FILE):
return False
cache_age = datetime.now() - datetime.fromtimestamp(os.path.getmtime(GPU_TYPES_CACHE_FILE))
return cache_age < CACHE_DURATION
if is_cache_valid():
try:
with open(GPU_TYPES_CACHE_FILE, "r") as file:
payload = json.load(file)
except (OSError, json.JSONDecodeError):
payload = None
else:
endpoint = "/api/v0/gpu_types/"
url = f"{server_url_default}{endpoint}"
try:
r = requests.get(url, headers={})
r.raise_for_status()
payload = r.json()
except (requests.exceptions.RequestException, ValueError):
return None
try:
with open(GPU_TYPES_CACHE_FILE, "w") as file:
json.dump(payload, file)
except OSError:
pass
try:
return payload['gpu_types']
except (TypeError, KeyError):
return None
APIKEY_FILE = os.path.join(DIRS['config'], "vast_api_key")
APIKEY_FILE_HOME = os.path.expanduser("~/.vast_api_key") # Legacy
TFAKEY_FILE = os.path.join(DIRS['config'], "vast_tfa_key")
if not os.path.exists(APIKEY_FILE) and os.path.exists(APIKEY_FILE_HOME):
#print(f'copying key from {APIKEY_FILE_HOME} -> {APIKEY_FILE}')
shutil.copyfile(APIKEY_FILE_HOME, APIKEY_FILE)
def format_key_suffix(k):
"""Format the last 4 chars of an API key for display, e.g. '...a3f9'."""
if k and len(k) >= 4:
return f"...{k[-4:]}"
return "(empty)"
# Client/host CLI role (CLN-3582): display-only (--help/completion/hints), never gates execution; client is the default until the role explicitly resolves to 'host'.
ROLE_FILE = os.path.join(DIRS['config'], "vast_role")
ROLE_CLIENT = "client"
ROLE_HOST = "host"
VALID_ROLES = {ROLE_CLIENT, ROLE_HOST}
def get_role():
"""Return the stored CLI role ('client' or 'host'), or None if ROLE_FILE is missing/unrecognized — never raises."""
try:
with open(ROLE_FILE) as f:
role = f.read().strip().lower()
except OSError:
return None
return role if role in VALID_ROLES else None
def is_client_view(role):
"""Whether role (as returned by get_role) means "show the client view" — everything except explicit 'host', including None."""
return role != ROLE_HOST
def set_role_file(role):
"""Persist role ('client' or 'host') to ROLE_FILE; raises ValueError for anything else."""
if role not in VALID_ROLES:
raise ValueError(f"role must be one of {sorted(VALID_ROLES)}, got {role!r}")
with open(ROLE_FILE, "w") as f:
f.write(role)
def detect_role(client):
"""Best-effort: silently resolve and permanently cache the role via client if it isn't known yet; never raises (a failed check just leaves it undetected for retry)."""
if get_role() is not None:
return
try:
from vastai.api import billing as billing_api
is_host = bool(billing_api.show_user(client).get("host_agreement_accepted"))
set_role_file(ROLE_HOST if is_host else ROLE_CLIENT)
except Exception:
return
# ---------------------------------------------------------------------------
# Simple utility class
# ---------------------------------------------------------------------------
class Object(object):
pass
# ---------------------------------------------------------------------------
# Validation helpers
# ---------------------------------------------------------------------------
def validate_seconds(value):
"""Validate that the input value is a valid number for seconds between yesterday and Jan 1, 2100."""
try:
val = int(value)
# Calculate min_seconds as the start of yesterday in seconds
yesterday = datetime.now() - timedelta(days=1)
min_seconds = int(yesterday.timestamp())
# Calculate max_seconds for Jan 1st, 2100 in seconds
max_date = datetime(2100, 1, 1, 0, 0, 0)
max_seconds = int(max_date.timestamp())
if not (min_seconds <= val <= max_seconds):
raise argparse.ArgumentTypeError(f"{value} is not a valid second timestamp.")
return val
except ValueError:
raise argparse.ArgumentTypeError(f"{value} is not a valid integer.")
# ---------------------------------------------------------------------------
# VRL (Vast Resource Locator) parsing
# ---------------------------------------------------------------------------
class VRLException(Exception):
pass
def parse_vast_url(url_str):
"""
Breaks up a vast-style url in the form instance_id:path and does
some basic sanity type-checking.
:param url_str:
:return:
"""
instance_id = None
path = url_str
#print(f'url_str: {url_str}')
if (":" in url_str):
url_parts = url_str.split(":", 2)
if len(url_parts) == 2:
(instance_id, path) = url_parts
else:
raise VRLException("Invalid VRL (Vast resource locator).")
else:
try:
instance_id = int(path)
path = "/"
except:
pass
valid_unix_path_regex = re.compile('^(/)?([^/\0]+(/)?)+$')
# Got this regex from https://stackoverflow.com/questions/537772/what-is-the-most-correct-regular-expression-for-a-unix-file-path
if (path != "/") and (valid_unix_path_regex.match(path) is None):
raise VRLException(f"Path component: {path} of VRL is not a valid Unix style path.")
#print(f'instance_id: {instance_id}')
#print(f'path: {path}')
return (instance_id, path)
# ---------------------------------------------------------------------------
# SSH key helpers
# ---------------------------------------------------------------------------
def get_ssh_key(argstr):
# Import deindent from display module (avoids circular imports)
from vastai.cli.display import deindent
ssh_key = argstr
# Including a path to a public key is pretty reasonable.
if os.path.exists(argstr):
with open(argstr) as f:
ssh_key = f.read()
if "PRIVATE KEY" in ssh_key:
raise ValueError(deindent("""
\U0001f434 Woah, hold on there, partner!
That's a *private* SSH key. You need to give the *public*
one. It usually starts with 'ssh-rsa', is on a single line,
has around 200 or so "base64" characters and ends with
some-user@some-where. "Generate public ssh key" would be
a good search term if you don't know how to do this.
""", add_separator=False))
if not ssh_key.lower().startswith('ssh'):
raise ValueError(deindent("""
Are you sure that's an SSH public key?
Usually it starts with the stanza 'ssh-(keytype)'
where the keytype can be things such as rsa, ed25519-sk,
or dsa. What you passed me was:
{}
And welp, that just don't look right.
""".format(ssh_key), add_separator=False))
return ssh_key
def generate_ssh_key(auto_yes=False):
"""
Generate a new SSH key pair using ssh-keygen and return the public key content.
Args:
auto_yes (bool): If True, automatically answer yes to prompts
Returns:
str: The content of the generated public key
Raises:
SystemExit: If ssh-keygen is not available or key generation fails
"""
print("No SSH key provided. Generating a new SSH key pair and adding public key to account...")
# Define paths
ssh_dir = Path.home() / '.ssh'
private_key_path = ssh_dir / 'id_ed25519'
public_key_path = ssh_dir / 'id_ed25519.pub'
# Create .ssh directory if it doesn't exist
try:
ssh_dir.mkdir(mode=0o700, exist_ok=True)
except OSError as e:
print(f"Error creating .ssh directory: {e}", file=sys.stderr)
sys.exit(1)
# Check if any part of the key pair already exists and backup if needed
if private_key_path.exists() or public_key_path.exists():
print(f"An SSH key pair 'id_ed25519' already exists in {ssh_dir}")
if auto_yes:
print("Auto-answering yes to backup existing key pair.")
response = 'y'
else:
response = input("Would you like to generate a new key pair and backup your existing id_ed25519 key pair. [y/N]: ").lower()
if response not in ['y', 'yes']:
print("Aborted. No new key generated.")
sys.exit(0)
# Generate timestamp for backup
timestamp = int(time.time())
backup_private_path = ssh_dir / f'id_ed25519.backup_{timestamp}'
backup_public_path = ssh_dir / f'id_ed25519.pub.backup_{timestamp}'
try:
# Backup existing private key if it exists
if private_key_path.exists():
private_key_path.rename(backup_private_path)
print(f"Backed up existing private key to: {backup_private_path}")
# Backup existing public key if it exists
if public_key_path.exists():
public_key_path.rename(backup_public_path)
print(f"Backed up existing public key to: {backup_public_path}")
except OSError as e:
print(f"Error backing up existing SSH keys: {e}", file=sys.stderr)
sys.exit(1)
print("Generating new SSH key pair and adding public key to account...")
# Check if ssh-keygen is available
try:
subprocess.run(['ssh-keygen', '--help'], capture_output=True, check=False)
except FileNotFoundError:
print("Error: ssh-keygen not found. Please install OpenSSH client tools.", file=sys.stderr)
sys.exit(1)
# Generate the SSH key pair
try:
cmd = [
'ssh-keygen',
'-t', 'ed25519', # Ed25519 key type
'-f', str(private_key_path), # Output file path
'-N', '', # Empty passphrase
'-C', f'{os.getenv("USER") or os.getenv("USERNAME", "user")}-vast.ai' # User
]
result = subprocess.run(
cmd,
capture_output=True,
text=True,
input='y\n', # Automatically answer 'yes' to overwrite prompts
check=True
)
except subprocess.CalledProcessError as e:
print(f"Error generating SSH key: {e}", file=sys.stderr)
if e.stderr:
print(f"ssh-keygen error: {e.stderr}", file=sys.stderr)
sys.exit(1)
except Exception as e:
print(f"Unexpected error during key generation: {e}", file=sys.stderr)
sys.exit(1)
# Set proper permissions for the private key
try:
private_key_path.chmod(0o600) # Read/write for owner only
except OSError as e:
print(f"Warning: Could not set permissions for private key: {e}", file=sys.stderr)
# Read and return the public key content
try:
with open(public_key_path, 'r') as f:
public_key_content = f.read().strip()
return public_key_content
except IOError as e:
print(f"Error reading generated public key: {e}", file=sys.stderr)
sys.exit(1)
# ---------------------------------------------------------------------------
# List / threading helpers
# ---------------------------------------------------------------------------
def split_list(lst, k):
"""
Splits a list into sublists of maximum size k.
"""
return [lst[i:i + k] for i in range(0, len(lst), k)]
def exec_with_threads(f, args, nt=16, max_retries=5):
def worker(sub_args):
for arg in sub_args:
retries = 0
while retries <= max_retries:
try:
result = None
if isinstance(arg, tuple):
result = f(*arg)
else:
result = f(arg)
if result: # Assuming a truthy return value means success
break
except Exception as e:
print(str(e))
pass
retries += 1
stime = 0.25 * 1.3 ** retries
print(f"retrying in {stime}s")
time.sleep(stime) # Exponential backoff
# Split args into nt sublists
args_per_thread = math.ceil(len(args) / nt)
sublists = [args[i:i + args_per_thread] for i in range(0, len(args), args_per_thread)]
with ThreadPoolExecutor(max_workers=nt) as executor:
executor.map(worker, sublists)
# ---------------------------------------------------------------------------
# Date / scheduling helpers
# ---------------------------------------------------------------------------
def default_start_date():
return datetime.now(timezone.utc).strftime("%Y-%m-%d")
def default_end_date():
return (datetime.now(timezone.utc) + timedelta(days=7)).strftime("%Y-%m-%d")
def convert_timestamp_to_date(unix_timestamp):
utc_datetime = datetime.fromtimestamp(unix_timestamp, tz=timezone.utc)
return utc_datetime.strftime("%Y-%m-%d")
def parse_day_cron_style(value):
"""
Accepts an integer string 0-6 or '*' to indicate 'Every day'.
Returns 0-6 as int, or None if '*'.
"""
val = str(value).strip()
if val == "*":
return None
try:
day = int(val)
if 0 <= day <= 6:
return day
except ValueError:
pass
raise argparse.ArgumentTypeError("Day must be 0-6 (0=Sunday) or '*' for every day.")
def parse_hour_cron_style(value):
"""
Accepts an integer string 0-23 or '*' to indicate 'Every hour'.
Returns 0-23 as int, or None if '*'.
"""
val = str(value).strip()
if val == "*":
return None
try:
hour = int(val)
if 0 <= hour <= 23:
return hour
except ValueError:
pass
raise argparse.ArgumentTypeError("Hour must be 0-23 or '*' for every hour.")
def convert_dates_to_timestamps(args):
end_timestamp = time.time()
start_timestamp = time.time() - (24 * 60 * 60)
from dateutil import parser as dateutil_parser
if args.end_date:
try:
end_timestamp = _parse_date_to_utc_timestamp(args.end_date, dateutil_parser)
except ValueError as e:
print(f"Warning: Invalid end date format! Ignoring end date! \n {str(e)}")
if args.start_date:
try:
start_timestamp = _parse_date_to_utc_timestamp(args.start_date, dateutil_parser)
except ValueError as e:
print(f"Warning: Invalid start date format! Ignoring start date! \n {str(e)}")
return start_timestamp, end_timestamp
def _parse_date_to_utc_timestamp(value, dateutil_parser):
"""Parse a user-supplied date string into a UNIX timestamp.
Naive inputs (e.g. 'YYYY-MM-DD') are interpreted as UTC so that filter
windows don't shift with the caller's local timezone. Aware inputs keep
their declared offset.
"""
dt = dateutil_parser.parse(str(value))
if dt.tzinfo is None:
dt = dt.replace(tzinfo=timezone.utc)
return dt.timestamp()
def validate_frequency_values(day_of_the_week, hour_of_the_day, frequency):
# Helper to raise an error with a consistent message.
def raise_frequency_error():
msg = ""
if frequency == "HOURLY":
msg += "For HOURLY jobs, day and hour must both be \"*\"."
elif frequency == "DAILY":
msg += "For DAILY jobs, day must be \"*\" and hour must have a value between 0-23."
elif frequency == "WEEKLY":
msg += "For WEEKLY jobs, day must have a value between 0-6 and hour must have a value between 0-23."
sys.exit(msg)
if frequency == "HOURLY":
if not (day_of_the_week is None and hour_of_the_day is None):
raise_frequency_error()
if frequency == "DAILY":
if not (day_of_the_week is None and hour_of_the_day is not None):
raise_frequency_error()
if frequency == "WEEKLY":
if not (day_of_the_week is not None and hour_of_the_day is not None):
raise_frequency_error()
def add_scheduled_job(client, args, req_json, cli_command, api_endpoint, request_method, instance_id, contract_end_date=None):
start_timestamp, end_timestamp = convert_dates_to_timestamps(args)
if args.end_date is None:
end_timestamp = contract_end_date
args.end_date = convert_timestamp_to_date(contract_end_date)
if start_timestamp >= end_timestamp:
raise ValueError("--start_date must be less than --end_date.")
day, hour, frequency = args.day, args.hour, args.schedule
request_body = {
"start_time": start_timestamp,
"end_time": end_timestamp,
"api_endpoint": api_endpoint,
"request_method": request_method,
"request_body": req_json,
"day_of_the_week": day,
"hour_of_the_day": hour,
"frequency": frequency,
"instance_id": instance_id
}
response = client.post("/commands/schedule_job/", json_data=request_body)
if args.explain:
print("request json: ")
print(request_body)
if response.status_code == 200:
print(f"add_scheduled_job insert: success - Scheduling {frequency} job to {cli_command} from {args.start_date} UTC to {args.end_date} UTC")
elif response.status_code == 401:
print(f"add_scheduled_job insert: failed status_code: {response.status_code}. It could be because you aren't using a valid api_key.")
elif response.status_code == 422:
user_input = input("Existing scheduled job found. Do you want to update it (y|n)? ")
if user_input.strip().lower() == "y":
scheduled_job_id = response.json()["scheduled_job_id"]
response = update_scheduled_job(client, cli_command, f"/commands/schedule_job/{scheduled_job_id}/", frequency, args.start_date, args.end_date, request_body)
else:
print("Job update aborted by the user.")
else:
print(f"add_scheduled_job insert: failed error: {response.status_code}. Response body: {response.text}")
def update_scheduled_job(client, cli_command, schedule_job_path, frequency, start_date, end_date, request_body):
response = client.put(schedule_job_path, json_data=request_body)
response.raise_for_status()
if response.status_code == 200:
print(f"add_scheduled_job update: success - Scheduling {frequency} job to {cli_command} from {start_date} UTC to {end_date} UTC")
print(response.json())
elif response.status_code == 401:
print(f"add_scheduled_job update: failed status_code: {response.status_code}. It could be because you aren't using a valid api_key.")
else:
print(f"add_scheduled_job update: failed status_code: {response.status_code}.")
print(response.json())
return response
# ---------------------------------------------------------------------------
# Permissions
# ---------------------------------------------------------------------------
def load_permissions_from_file(file_path):
with open(file_path, 'r') as file:
return json.load(file)
def get_template_arguments():
from vastai.cli.parser import argument
return [
argument("--name", help="name of the template", type=str),
argument("--image", help="docker container image to launch", type=str),
argument("--image_tag", help="docker image tag (can also be appended to end of image_path)", type=str),
argument("--href", help="link you want to provide", type=str),
argument("--repo", help="link to repository", type=str),
argument("--login", help="docker login arguments for private repo authentication, surround with ''", type=str),
argument("--env", help="Contents of the 'Docker options' field", type=str),
argument("--ssh", help="Launch as an ssh instance type", action="store_true"),
argument("--jupyter", help="Launch as a jupyter instance instead of an ssh instance", action="store_true"),
argument("--direct", help="Use (faster) direct connections for jupyter & ssh", action="store_true"),
argument("--jupyter-dir", help="For runtype 'jupyter', directory in instance to use to launch jupyter. Defaults to image's working directory", type=str),
argument("--jupyter-lab", help="For runtype 'jupyter', Launch instance with jupyter lab", action="store_true"),
argument("--onstart-cmd", help="contents of onstart script as single argument", type=str),
argument("--search_params", help="search offers filters", type=str),
argument("-n", "--no-default", action="store_true", help="Disable default search param query args"),
argument("--disk_space", help="disk storage space, in GB", type=str),
argument("--readme", help="readme string", type=str),
argument("--hide-readme", help="hide the readme from users", action="store_true"),
argument("--desc", help="description string", type=str),
argument("--public", help="make template available to public", action="store_true"),
]
# ---------------------------------------------------------------------------
# Verification thresholds
# ---------------------------------------------------------------------------
def required_inet_mbps(gpu_total_ram_mib):
"""Machine-total VRAM-scaled bandwidth floor for the self-test pre-flight check.
Returns the minimum inet_down / inet_up (Mb/s) a machine needs to qualify
for the GPU verification pipeline. Floors at 100, caps at 500, scales
linearly with total VRAM against a 192 GiB reference point.
Falsy / missing gpu_total_ram falls to the 100 Mb/s floor. The column is
MiB (binary), so a B200 reporting 183359 MiB (~179 GiB) lands at 466
Mb/s; the cap is reached once total VRAM crosses 192 GiB, e.g. multi-GPU
machines.
"""
total_vram_gib = (gpu_total_ram_mib or 0) / 1024.0
return min(500.0, max(100.0, 500.0 * total_vram_gib / 192.0))
cli/utils.py
"""Shared helpers used across all CLI command modules."""
def get_parser():
from vastai.cli.main import parser
return parser
def get_client(args):
"""Create a VastClient from parsed CLI args."""
from vastai.api.client import VastClient
from vastai.cli.util import detect_role
client = VastClient(
api_key=args.api_key,
server_url=args.url,
retry=args.retry,
explain=getattr(args, 'explain', False),
curl=getattr(args, 'curl', False),
client_type="cli",
)
detect_role(client)
return client
data/__init__.py
from .query import Query, Column
from .offer import Offer
from .endpoint import EndpointConfig, EndpointData
from .deployment import DeploymentConfig, DeploymentData, DeploymentPutResponse
from .workergroup import WorkergroupConfig
from .instance import Instance, InstanceConfig, CreateInstanceResponse
data/deployment.py
import dataclasses
from dataclasses import dataclass
from typing import Any, Optional
@dataclass
class DeploymentConfig:
"""Configuration for creating or updating a deployment.
Required: name, image, file_hash, file_size.
Endpoint scaling params are optional and follow server defaults.
"""
name: str
image: str
file_hash: str
file_size: int
env: Optional[str] = None
tag: Optional[str] = None
search_params: Optional[str] = None
storage: Optional[float] = None
ttl: Optional[float] = None
version_label: Optional[str] = None
docker_login_user: Optional[str] = None
docker_login_pass: Optional[str] = None
docker_login_repo: Optional[str] = None
# Endpoint scaling params
cold_workers: Optional[int] = None
max_workers: Optional[int] = None
min_load: Optional[float] = None
min_cold_load: Optional[float] = None
target_util: Optional[float] = None
cold_mult: Optional[float] = None
max_queue_time: Optional[float] = None
target_queue_time: Optional[float] = None
inactivity_timeout: Optional[float] = None
autoscaler_instance: Optional[str] = None
def to_dict(self) -> dict:
return {k: v for k, v in dataclasses.asdict(self).items() if v is not None}
@dataclass
class DeploymentPutResponse:
"""Response from PUT /api/v0/deployments/."""
success: bool
action: str # "created", "soft_update", "autoscale_update", "exists"
deployment_id: int
endpoint_id: int
workergroup_id: Optional[int] = None
upload_url: Optional[str] = None
upload_fields: Optional[dict] = None
evicted_versions: Optional[list] = None
@classmethod
def from_dict(cls, d: dict) -> "DeploymentPutResponse":
return cls(
success=d["success"],
action=d["action"],
deployment_id=d["deployment_id"],
endpoint_id=d["endpoint_id"],
workergroup_id=d.get("workergroup_id"),
upload_url=d.get("upload_url"),
upload_fields=d.get("upload_fields"),
evicted_versions=d.get("evicted_versions"),
)
@dataclass
class DeploymentData:
"""Full deployment entity as returned by GET /api/v0/deployment/{id}/."""
id: int
name: str
tag: str
endpoint_id: Optional[int]
endpoint_state: Optional[str]
worker_count: Optional[int]
s3_key: Optional[str]
env: Optional[str]
image: str
storage: float
search_params: Optional[str]
file_hash: str
current_version_id: int
last_healthy_version_id: Optional[int]
ttl: Optional[float]
last_client_heartbeat: Optional[float]
created_at: float
updated_at: float
@classmethod
def from_dict(cls, d: dict) -> "DeploymentData":
return cls(
id=d["id"],
name=d["name"],
tag=d.get("tag", "default"),
endpoint_id=d.get("endpoint_id"),
endpoint_state=d.get("endpoint_state"),
worker_count=d.get("worker_count"),
s3_key=d.get("s3_key"),
env=d.get("env"),
image=d["image"],
storage=d.get("storage", 50.0),
search_params=d.get("search_params"),
file_hash=d["file_hash"],
current_version_id=d["current_version_id"],
last_healthy_version_id=d.get("last_healthy_version_id"),
ttl=d.get("ttl"),
last_client_heartbeat=d.get("last_client_heartbeat"),
created_at=d["created_at"],
updated_at=d["updated_at"],
)
data/endpoint.py
import dataclasses
from dataclasses import dataclass
from typing import Optional
@dataclass
class EndpointConfig:
"""Configuration for creating or updating an endpoint job.
Only endpoint_name is required. All other fields default to None,
meaning the server will use its own defaults.
"""
endpoint_name: str
cold_workers: Optional[int] = None
max_workers: Optional[int] = None
min_load: Optional[float] = None
min_cold_load: Optional[float] = None
target_util: Optional[float] = None
cold_mult: Optional[float] = None
max_queue_time: Optional[float] = None
target_queue_time: Optional[float] = None
autoscaler_instance: Optional[str] = None
endpoint_state: Optional[str] = None
inactivity_timeout: Optional[float] = None
def to_dict(self) -> dict:
return {k: v for k, v in dataclasses.asdict(self).items() if v is not None}
@dataclass
class EndpointData:
"""Full endpoint entity as returned by the API."""
id: int
config: EndpointConfig
api_key: str
user_id: int
created_at: float
auto_delete_in_seconds: Optional[float]
auto_delete_due_24h: bool
@classmethod
def from_dict(cls, d: dict) -> "EndpointData":
config = EndpointConfig(
endpoint_name=d["endpoint_name"],
cold_workers=d.get("cold_workers"),
max_workers=d.get("max_workers"),
min_load=d.get("min_load"),
min_cold_load=d.get("min_cold_load"),
target_util=d.get("target_util"),
cold_mult=d.get("cold_mult"),
max_queue_time=d.get("max_queue_time"),
target_queue_time=d.get("target_queue_time"),
endpoint_state=d.get("endpoint_state"),
inactivity_timeout=d.get("inactivity_timeout"),
)
return cls(
id=d["id"],
config=config,
api_key=d["api_key"],
user_id=d["user_id"],
created_at=d["created_at"],
auto_delete_in_seconds=d.get("auto_delete_in_seconds"),
auto_delete_due_24h=d.get("auto_delete_due_24h", False),
)
data/instance.py
import dataclasses
from dataclasses import dataclass
from typing import Any, Optional
@dataclass
class InstanceConfig:
"""Configuration for creating an instance from an offer.
image is required. All other fields are optional.
"""
image: str
disk: Optional[float] = None
label: Optional[str] = None
onstart: Optional[str] = None
env: Optional[dict[str, str]] = None
image_login: Optional[str] = None
runtype: Optional[str] = None
args: Optional[list] = None
args_str: Optional[str] = None
python_utf8: Optional[bool] = None
lang_utf8: Optional[bool] = None
use_jupyter_lab: Optional[bool] = None
jupyter_dir: Optional[str] = None
template_hash_id: Optional[str] = None
template_id: Optional[int] = None
extra: Optional[str] = None
target_state: Optional[str] = None
def to_dict(self) -> dict:
return {k: v for k, v in dataclasses.asdict(self).items() if v is not None}
@dataclass
class CreateInstanceResponse:
"""Response from PUT /api/v0/asks/{id}/ on success."""
success: bool
new_contract: int
instance_api_key: Optional[str] = None
ask_id: Optional[int] = None
@classmethod
def from_dict(cls, d: dict) -> "CreateInstanceResponse":
return cls(
success=d["success"],
new_contract=d["new_contract"],
instance_api_key=d.get("instance_api_key"),
ask_id=d.get("ask_id"),
)
@dataclass
class Instance:
"""Instance entity as returned by GET /api/v1/instances/."""
id: int
machine_id: Optional[int]
actual_status: Optional[str]
cur_state: Optional[str]
next_state: Optional[str]
intended_status: Optional[str]
image_uuid: Optional[str]
image_args: Optional[str]
image_runtype: Optional[str]
label: Optional[str]
num_gpus: Optional[int]
gpu_name: Optional[str]
gpu_util: Optional[float]
gpu_arch: Optional[str]
gpu_temp: Optional[float]
cuda_max_good: Optional[str]
driver_version: Optional[str]
cpu_cores_effective: Optional[float]
cpu_util: Optional[float]
cpu_ram: Optional[int]
disk_space: Optional[float]
disk_util: Optional[float]
disk_usage: Optional[float]
mem_usage: Optional[float]
mem_limit: Optional[float]
vmem_usage: Optional[float]
inet_up: Optional[float]
inet_down: Optional[float]
ssh_host: Optional[str]
ssh_port: Optional[int]
ssh_idx: Optional[str]
public_ipaddr: Optional[str]
local_ipaddrs: Optional[list]
direct_port_start: Optional[int]
direct_port_end: Optional[int]
status_msg: Optional[str]
jupyter_token: Optional[str]
extra_env: Optional[list]
onstart: Optional[str]
uptime_mins: Optional[float]
start_date: Optional[float]
end_date: Optional[float]
reliability2: Optional[float]
dph_total: Optional[float]
host_id: Optional[int]
template_id: Optional[int]
template_hash_id: Optional[str]
template_name: Optional[str]
ports: Optional[str]
@classmethod
def from_dict(cls, d: dict) -> "Instance":
"""Construct an Instance from an API response dict, ignoring unknown keys."""
known = {f.name: f for f in dataclasses.fields(cls)}
kwargs = {}
for name, field in known.items():
if name in d:
kwargs[name] = d[name]
elif field.default is not dataclasses.MISSING:
pass
elif field.default_factory is not dataclasses.MISSING:
pass
else:
kwargs[name] = None
return cls(**kwargs)
data/offer.py
from dataclasses import dataclass
from typing import Optional
@dataclass
class Offer:
"""
A search result offer from ask_contract_offers, as returned by /api/v0/bundles/.
Field names match the API response keys.
"""
id: int
ask_contract_id: Optional[int]
machine_id: int
host_id: int
num_gpus: int
gpu_name: str
gpu_ram: int # MB
gpu_total_ram: Optional[int] # MB
gpu_frac: Optional[float]
gpu_arch: Optional[str]
gpu_lanes: Optional[int]
gpu_mem_bw: Optional[float]
gpu_max_power: Optional[float]
gpu_max_temp: Optional[float]
compute_cap: Optional[int]
cuda_max_good: Optional[float]
bw_nvlink: Optional[float]
cpu_name: Optional[str]
cpu_cores: Optional[int]
cpu_cores_effective: Optional[float]
cpu_ghz: Optional[float]
cpu_ram: Optional[int] # MB
cpu_arch: Optional[str]
disk_space: Optional[float]
disk_bw: Optional[float]
disk_name: Optional[str]
inet_up: Optional[float]
inet_down: Optional[float]
inet_up_cost: Optional[float]
inet_down_cost: Optional[float]
direct_port_count: Optional[int]
pci_gen: Optional[float]
pcie_bw: Optional[float]
mobo_name: Optional[str]
os_version: Optional[str]
driver_vers: Optional[int]
has_avx: Optional[int]
dph_base: Optional[float]
dph_total: Optional[float]
min_bid: Optional[float]
dlperf: Optional[float]
dlperf_per_dphtotal: Optional[float]
flops_per_dphtotal: Optional[float]
total_flops: Optional[float]
score: Optional[float]
vram_costperhour: Optional[float]
storage_cost: Optional[float]
storage_total_cost: Optional[float]
credit_discount_max: Optional[float]
rentable: Optional[bool]
rented: Optional[bool]
is_bid: Optional[bool]
external: Optional[int]
hosting_type: Optional[int]
hostname: Optional[str]
public_ipaddr: Optional[str]
geolocation: Optional[str]
geolocode: Optional[int]
reliability: Optional[float]
reliability_mult: Optional[float]
expected_reliability: Optional[float]
verification: Optional[str]
vericode: Optional[int]
static_ip: Optional[bool]
vms_enabled: Optional[bool]
is_vm_deverified: Optional[bool]
gpu_display_active: Optional[bool]
logo: Optional[str]
webpage: Optional[str]
resource_type: Optional[str]
start_date: Optional[float]
end_date: Optional[float]
duration: Optional[float]
bundle_id: Optional[int]
bundled_results: Optional[int]
cluster_id: Optional[int]
avail_vol_ask_id: Optional[int]
avail_vol_dph: Optional[float]
avail_vol_size: Optional[float]
nw_disk_min_bw: Optional[int]
nw_disk_max_bw: Optional[int]
nw_disk_avg_bw: Optional[int]
platform_fee: Optional[float]
@classmethod
def from_dict(cls, d: dict) -> "Offer":
"""Construct an Offer from a raw API response dict, ignoring unknown keys.
Missing Optional fields default to None."""
import dataclasses
known = {f.name: f for f in dataclasses.fields(cls)}
kwargs = {}
for name, field in known.items():
if name in d:
kwargs[name] = d[name]
elif hasattr(field.default, '__call__') or field.default is not dataclasses.MISSING:
pass # has a default, let dataclass handle it
elif field.default_factory is not dataclasses.MISSING:
pass # has a default_factory
else:
# No default — use None for Optional fields, skip otherwise
kwargs[name] = None
return cls(**kwargs)
data/query.py
# Implements an object oriented query language for Vast standard query (vast/web/apiquery.py)
from typing import Any, Optional
# Op-to-string symbol mapping. Used by Query.unparse_query() and by
# vastai/cli/commands/benchmarks.py for human-readable filter rendering.
OP_TO_STR = {
"eq": "=",
"neq": "!=",
"gt": ">",
"gte": ">=",
"lt": "<",
"lte": "<=",
"in": " in ",
"notin": " notin ",
}
class Query:
def __init__(self, query_dict: dict[str, dict[str, Any]]):
# self.query has format {column : {op : value}}; may have multiple ops per column and multiple columns.
self.query = query_dict
@classmethod
def search_defaults(
cls,
verified: Optional[bool] = True,
rentable: Optional[bool] = True,
rented: Optional[bool] = False,
) -> "Query":
"""Creates a Query pre-populated with default search filters for end-user use."""
q = cls({})
for name, value in [
("verified", verified),
("rentable", rentable),
("rented", rented),
]:
if value is not None:
q.query[name] = {"eq": value}
return q
def unparse_query(self) -> str:
"""Converts this Query back into a query string that parse_query() would parse into the same dict.
Inverse of vast/web/parse.py:parse_query, i.e. parse_query(query.unparse_query()) == query.query.
"""
# Fields whose values get multiplied by these factors during parsing
mult_fields = {
"cpu_ram": 1000,
"gpu_ram": 1000,
"gpu_total_ram": 1000,
"duration": 24.0 * 60.0 * 60.0,
}
parts = []
for field, ops in self.query.items():
for op, value in ops.items():
op_str = OP_TO_STR[op]
# Reverse the multiplier applied during parsing
if field in mult_fields:
value = value / mult_fields[field]
# Use int if it's a whole number
if value == int(value):
value = int(value)
# Format the value
if isinstance(value, bool):
val_str = "true" if value else "false"
elif value is None:
val_str = "None"
elif isinstance(value, list):
# in/notin: reverse the space→underscore replacement per element
items = [str(v).replace(" ", "_") for v in value]
val_str = ",".join(items)
elif isinstance(value, str):
# Reverse the underscore→space replacement done by parse_query
val_str = value.replace(" ", "_")
else:
val_str = str(value)
parts.append(f"{field}{op_str}{val_str}")
return " ".join(parts)
def extend(self, other: "Query") -> "Query":
"""
Extends self.query with filters in other.
self.query should now have all columns that exist in either self.query or other.query;
and each column in self.query should have all corresponding ops in either query.
"""
for column, ops in other.query.items():
if column in self.query:
conflicts = self.query[column].keys() & ops.keys()
if conflicts:
raise ValueError(
f"Cannot redefine op(s) {conflicts} for column '{column}'"
)
self.query[column].update(ops)
else:
self.query[column] = dict(ops)
return self
class Column:
def __init__(self, name: str):
self.name = name
def __eq__(self, value: Any) -> Query:
return Query({self.name: {"eq": value}})
def __ne__(self, value: Any) -> Query:
return Query({self.name: {"neq": value}})
def __lt__(self, value: Any) -> Query:
return Query({self.name: {"lt": value}})
def __le__(self, value: Any) -> Query:
return Query({self.name: {"lte": value}})
def __gt__(self, value: Any) -> Query:
return Query({self.name: {"gt": value}})
def __ge__(self, value: Any) -> Query:
return Query({self.name: {"gte": value}})
def in_(self, values) -> Query:
return Query({self.name: {"in": values}})
def notin_(self, values) -> Query:
return Query({self.name: {"notin": values}})
# Pre-initialized Column objects for all columns in ask_contract_offers
id_ = Column("id")
ask_contract_id = Column("ask_contract_id")
bundle_id = Column("bundle_id")
bundled_results = Column("bundled_results")
bw_nvlink = Column("bw_nvlink")
compute_cap = Column("compute_cap")
cpu_arch = Column("cpu_arch")
cpu_cores = Column("cpu_cores")
cpu_cores_effective = Column("cpu_cores_effective")
cpu_ghz = Column("cpu_ghz")
cpu_name = Column("cpu_name")
cpu_ram = Column("cpu_ram")
credit_discount_max = Column("credit_discount_max")
cuda_max_good = Column("cuda_max_good")
direct_port_count = Column("direct_port_count")
disk_bw = Column("disk_bw")
disk_name = Column("disk_name")
disk_space = Column("disk_space")
dlperf = Column("dlperf")
dlperf_per_dphtotal = Column("dlperf_per_dphtotal")
dph_base = Column("dph_base")
dph_total = Column("dph_total")
driver_version = Column("driver_version")
driver_vers = Column("driver_vers")
duration = Column("duration")
end_date = Column("end_date")
external = Column("external")
flops_per_dphtotal = Column("flops_per_dphtotal")
geolocation = Column("geolocation")
geolocode = Column("geolocode")
gpu_arch = Column("gpu_arch")
gpu_display_active = Column("gpu_display_active")
gpu_frac = Column("gpu_frac")
gpu_ids = Column("gpu_ids")
gpu_lanes = Column("gpu_lanes")
gpu_mem_bw = Column("gpu_mem_bw")
gpu_name = Column("gpu_name")
gpu_ram = Column("gpu_ram")
gpu_total_ram = Column("gpu_total_ram")
gpu_max_power = Column("gpu_max_power")
gpu_max_temp = Column("gpu_max_temp")
has_avx = Column("has_avx")
host_id = Column("host_id")
hosting_type = Column("hosting_type")
hostname = Column("hostname")
inet_down = Column("inet_down")
inet_down_cost = Column("inet_down_cost")
inet_up = Column("inet_up")
inet_up_cost = Column("inet_up_cost")
is_bid = Column("is_bid")
logo = Column("logo")
machine_id = Column("machine_id")
min_bid = Column("min_bid")
mobo_name = Column("mobo_name")
num_gpus = Column("num_gpus")
os_version = Column("os_version")
pci_gen = Column("pci_gen")
pcie_bw = Column("pcie_bw")
public_ipaddr = Column("public_ipaddr")
reliability = Column("reliability")
reliability_mult = Column("reliability_mult")
rentable = Column("rentable")
rented = Column("rented")
score = Column("score")
start_date = Column("start_date")
static_ip = Column("static_ip")
storage_cost = Column("storage_cost")
storage_total_cost = Column("storage_total_cost")
total_flops = Column("total_flops")
verification = Column("verification")
vericode = Column("vericode")
vram_costperhour = Column("vram_costperhour")
webpage = Column("webpage")
vms_enabled = Column("vms_enabled")
expected_reliability = Column("expected_reliability")
is_vm_deverified = Column("is_vm_deverified")
resource_type = Column("resource_type")
cluster_id = Column("cluster_id")
avail_vol_ask_id = Column("avail_vol_ask_id")
avail_vol_dph = Column("avail_vol_dph")
avail_vol_size = Column("avail_vol_size")
nw_disk_min_bw = Column("nw_disk_min_bw")
nw_disk_max_bw = Column("nw_disk_max_bw")
nw_disk_avg_bw = Column("nw_disk_avg_bw")
platform_fee = Column("platform_fee")
# Virtual columns accepted by the API that are converted to real column(s) by search_asks_
# (see vast/web/views/misc.py)
geolocation = Column("geolocation") # hashed → geolocode
datacenter = Column("datacenter") # bool → hosting_type {eq: 0|1}
duration = Column("duration") # seconds from now → end_date
verified = Column("verified") # bool → vericode
allocated_storage = Column("allocated_storage") # GB → disk_space {gte: value}
target_reliability = Column(
"target_reliability"
) # float → expected_reliability {gte: value}
# Column aliases (remapped to real columns before querying)
reliability2 = Column("reliability2") # → reliability
dphtotal = Column("dphtotal") # → dph_total
gpu_totalram = Column("gpu_totalram") # → gpu_total_ram
gpu_totalram_GB = Column("gpu_totalram_GB") # → gpu_total_ram (value * 1000)
cpu_ram_GB = Column("cpu_ram_GB") # → cpu_ram (value * 1000)
gpu_ram_GB = Column("gpu_ram_GB") # → gpu_ram (value * 1000)
ubuntu_version = Column("ubuntu_version") # → os_version
driver_version = Column("driver_version") # → driver_vers
# ---------------------------------------------------------------------------
# GPU name constants (from vast/controller/gpulist.json, column [1])
#
# Usage: gpu_name == RTX_4090 or gpu_name.in_([RTX_4090, RTX_4090_Ti])
# ---------------------------------------------------------------------------
# NVIDIA Data-Center / AI
A10 = "A10"
A10g = "A10g"
A16 = "A16"
A30 = "A30"
A40 = "A40"
A100_PCIE = "A100 PCIE"
A100_SXM4 = "A100 SXM4"
A100_SXM = "A100 SXM"
A100X = "A100X"
A800_PCIE = "A800 PCIE"
B200 = "B200"
GH200_SXM = "GH200 SXM"
H100_PCIE = "H100 PCIE"
H100_SXM = "H100 SXM"
H100_NVL = "H100 NVL"
H200 = "H200"
H200_NVL = "H200 NVL"
L4 = "L4"
L40 = "L40"
L40S = "L40S"
Tesla_K20c = "Tesla K20c"
Tesla_K80 = "Tesla K80"
Tesla_M40 = "Tesla M40"
Tesla_P4 = "Tesla P4"
Tesla_P6 = "Tesla P6"
Tesla_P40 = "Tesla P40"
Tesla_P100 = "Tesla P100"
Tesla_T4 = "Tesla T4"
Tesla_V100 = "Tesla V100"
# NVIDIA Consumer – GeForce / GTX
GeForce_970 = "GeForce 970"
GTX_750_Ti = "GTX 750 Ti"
GTX_1050 = "GTX 1050"
GTX_1050_Ti = "GTX 1050 Ti"
GTX_1060 = "GTX 1060"
GTX_1070 = "GTX 1070"
GTX_1070_Ti = "GTX 1070 Ti"
GTX_1080 = "GTX 1080"
GTX_1080_Ti = "GTX 1080 Ti"
GTX_1650 = "GTX 1650"
GTX_1650_S = "GTX 1650 S"
GTX_1650_Ti = "GTX 1650 Ti"
GTX_1660 = "GTX 1660"
GTX_1660_S = "GTX 1660 S"
GTX_1660_Ti = "GTX 1660 Ti"
GTX_TITAN_X = "GTX TITAN X"
# NVIDIA Consumer – RTX 20-series
RTX_2060 = "RTX 2060"
RTX_2060S = "RTX 2060S"
RTX_2070 = "RTX 2070"
RTX_2070S = "RTX 2070S"
RTX_2080 = "RTX 2080"
RTX_2080S = "RTX 2080S"
RTX_2080_Ti = "RTX 2080 Ti"
# NVIDIA Consumer – RTX 30-series
RTX_3050 = "RTX 3050"
RTX_3060 = "RTX 3060"
RTX_3060_laptop = "RTX 3060 laptop"
RTX_3060_Ti = "RTX 3060 Ti"
RTX_3070 = "RTX 3070"
RTX_3070_laptop = "RTX 3070 laptop"
RTX_3070_Ti = "RTX 3070 Ti"
RTX_3080 = "RTX 3080"
RTX_3080_Ti = "RTX 3080 Ti"
RTX_3090 = "RTX 3090"
RTX_3090_Ti = "RTX 3090 Ti"
# NVIDIA Consumer – RTX 40-series
RTX_4060 = "RTX 4060"
RTX_4060_laptop = "RTX 4060 laptop"
RTX_4060_Ti = "RTX 4060 Ti"
RTX_4070 = "RTX 4070"
RTX_4070_laptop = "RTX 4070 laptop"
RTX_4070S = "RTX 4070S"
RTX_4070S_Ti = "RTX 4070S Ti"
RTX_4070_Ti = "RTX 4070 Ti"
RTX_4080 = "RTX 4080"
RTX_4080_laptop = "RTX 4080 laptop"
RTX_4080S = "RTX 4080S"
RTX_4090 = "RTX 4090"
RTX_4090D = "RTX 4090D"
RTX_4090_Ti = "RTX 4090 Ti"
# NVIDIA Consumer – RTX 50-series
RTX_5060 = "RTX 5060"
RTX_5060_Ti = "RTX 5060 Ti"
RTX_5070 = "RTX 5070"
RTX_5070_Ti = "RTX 5070 Ti"
RTX_5080 = "RTX 5080"
RTX_5090 = "RTX 5090"
RTX_5090_Ti = "RTX 5090 Ti"
# NVIDIA Consumer – Titan
Titan_RTX = "Titan RTX"
Titan_V = "Titan V"
TITAN_V_CEO = "TITAN V CEO"
Titan_X = "Titan X"
Titan_Xp = "Titan Xp"
# NVIDIA Professional – Quadro
Quadro_K620 = "Quadro K620"
Quadro_K2200 = "Quadro K2200"
Quadro_P2000 = "Quadro P2000"
Quadro_P4000 = "Quadro P4000"
Quadro_P5000 = "Quadro P5000"
Quadro_P6000 = "Quadro P6000"
Quadro_GP100 = "Quadro GP100"
Quadro_GV100 = "Quadro GV100"
Q_RTX_4000 = "Q RTX 4000"
Q_RTX_5000 = "Q RTX 5000"
Q_RTX_6000 = "Q RTX 6000"
Q_RTX_8000 = "Q RTX 8000"
# NVIDIA Professional – RTX A-series
RTX_A2000 = "RTX A2000"
RTX_A4000 = "RTX A4000"
RTX_A4500 = "RTX A4500"
RTX_A5000 = "RTX A5000"
RTX_A6000 = "RTX A6000"
# NVIDIA Professional – RTX Ada Generation
RTX_2000Ada = "RTX 2000Ada"
RTX_4000Ada = "RTX 4000Ada"
RTX_4500Ada = "RTX 4500Ada"
RTX_5000Ada = "RTX 5000Ada"
RTX_5880Ada = "RTX 5880Ada"
RTX_6000Ada = "RTX 6000Ada"
# NVIDIA Professional – RTX PRO (Blackwell)
RTX_PRO_4000 = "RTX PRO 4000"
RTX_PRO_4000_Mobile = "RTX PRO 4000 Mobile"
RTX_PRO_4500 = "RTX PRO 4500"
RTX_PRO_5000 = "RTX PRO 5000"
RTX_PRO_6000 = "RTX PRO 6000"
RTX_PRO_6000_S = "RTX PRO 6000 S"
RTX_PRO_6000_WS = "RTX PRO 6000 WS"
RTX_PRO_6000_Max_Q = "RTX PRO 6000 Max-Q"
# NVIDIA Mining / Other
CMP_50HX = "CMP 50HX"
CMP_170HX = "CMP 170HX"
P104_100 = "P104-100"
P106_100 = "P106-100"
# AMD Instinct
InstinctMI50 = "InstinctMI50"
InstinctMI100 = "InstinctMI100"
InstinctMI210 = "InstinctMI210"
InstinctMI250X = "InstinctMI250X"
# AMD Radeon Pro
PRO_W6800 = "PRO W6800"
PRO_W7800 = "PRO W7800"
PRO_W7900 = "PRO W7900"
Pro_V620 = "Pro V620"
Radeon_VII = "Radeon VII"
Radeon_Pro_VII = "Radeon Pro VII"
# AMD Radeon RX
RX_6800 = "RX 6800"
RX_6800_XT = "RX 6800 XT"
RX_6900_XT = "RX 6900 XT"
RX_6950_XT = "RX 6950 XT"
RX_7600 = "RX 7600"
RX_7700_XT = "RX 7700 XT"
RX_7800_XT = "RX 7800 XT"
RX_7900_GRE = "RX 7900 GRE"
RX_7900_XT = "RX 7900 XT"
RX_7900_XTX = "RX 7900 XTX"
data/workergroup.py
import dataclasses
from dataclasses import dataclass
from typing import Optional
@dataclass
class WorkergroupConfig:
"""Configuration for creating a workergroup (autoscale job).
Request-only. Workergroup responses come back as raw dicts from
show_workergroups, so fields the API returns but does not accept as input
do not belong here.
Either endpoint_id or endpoint_name is required (endpoint_id is set
automatically when using ManagedEndpoint.add_workergroup()).
Either template_hash or template_id is required (template resolves image, search_query, etc.).
search_params is always required by the server; defaults to empty dict.
"""
search_params: str | dict = ""
endpoint_id: Optional[int] = None
endpoint_name: Optional[str] = None
template_hash: Optional[str] = None
template_id: Optional[int] = None
launch_args: Optional[str] = None
min_cold_load: Optional[float] = None
max_workers: Optional[int] = None
gpu_ram: Optional[float] = None
autoscaler_instance: Optional[str] = None
docker_login_user: Optional[str] = None
docker_login_pass: Optional[str] = None
docker_login_repo: Optional[str] = None
def to_dict(self) -> dict:
d = {k: v for k, v in dataclasses.asdict(self).items() if v is not None}
# search_params is always required by the server
d.setdefault("search_params", self.search_params)
return d
pdf/__init__.py
pdf/vast_pdf.py
#!/usr/bin/env python3
# DEPRECATED: This file is kept for backwards compatibility.
# Please use the vastai package instead.
# vast_pdf: Library of functions to create PDF reports of various type.
# Currently only makes invoices.
import argparse
import os
import datetime
import math
import time
import random
from decimal import Decimal
import typing
import PIL.Image
try:
from borb.pdf.canvas.color.color import HexColor, X11Color
from borb.pdf.canvas.layout.image.image import Image
from borb.pdf.canvas.layout.layout_element import Alignment
from borb.pdf.canvas.layout.page_layout.multi_column_layout import SingleColumnLayout
from borb.pdf.canvas.layout.page_layout.page_layout import PageLayout
from borb.pdf.canvas.layout.table.fixed_column_width_table import FixedColumnWidthTable
from borb.pdf.canvas.layout.table.flexible_column_width_table import FlexibleColumnWidthTable
from borb.pdf.canvas.layout.table.table import Table, TableCell
from borb.pdf.canvas.layout.text.paragraph import Paragraph
from borb.pdf.document import Document
from borb.pdf.page.page import Page
from borb.pdf.pdf import PDF
except ImportError:
print("""\nERROR: This library depends on a Python package called Borb to make the PDF files. To install this
package do 'pip3 install borb'.\n""")
## Globals
num_rows_first_page: int = 18
num_rows_subsequents_pages: int = 30
invoice_total: float = 0
now = datetime.date.today()
invoice_number: int = now.year * 12 + now.month - 1
page_count: int = 0
no_table_borders = False
no_table_borders = True
# def Paragraph_wr(text: str, *args, **kwargs):
# """This is just a wrapper for the Paragraph function that prevents the text from being the null string "".
#
# :param str text: The text to be checked for null.
# """
# if text == "":
# text = " "
# return Paragraph(text, *args, **kwargs)
def build_2nd_block_table() -> FixedColumnWidthTable:
"""
This function creates a Table containing invoice information.
This information spans the page and is the second large block of
text on the page.
:rtype FixedColumnWidthTable: a Table containing information such as the company address, payment terms, date, and sum of charge_fields/payments.
"""
global invoice_total, now
table = FixedColumnWidthTable(number_of_rows=3, number_of_columns=3)
table.add(Paragraph("Vast.ai Inc."))
table.add(Paragraph("Date", font="Helvetica-Bold", horizontal_alignment=Alignment.RIGHT))
table.add(Paragraph("%d/%d/%d" % (now.day, now.month, now.year), horizontal_alignment=Alignment.RIGHT))
table.add(Paragraph("100 Van Ness Ave."))
table.add(Paragraph("Payment Terms:", font="Helvetica-Bold", horizontal_alignment=Alignment.RIGHT))
table.add(Paragraph("Charged - Do Not Pay", horizontal_alignment=Alignment.RIGHT))
table.add(Paragraph("San Francisco CA 94102"))
table.add(
Paragraph("Total", font="Helvetica-Bold", font_size=Decimal(20), horizontal_alignment=Alignment.RIGHT))
table.add(Paragraph(format_float_val_as_currency(invoice_total), font="Helvetica-Bold", font_size=Decimal(20),
horizontal_alignment=Alignment.RIGHT))
table.set_padding_on_all_cells(Decimal(2), Decimal(2), Decimal(2), Decimal(2))
if no_table_borders: table.no_borders()
return table
def blank_row(table: Table, col_num: int, row_num: int = 1) -> None:
"""Just a set of blank rows to act as filler.
:param table: The table we want to modify
:param int col_num: How many columns in the table?
:param int row_num: How many blank rows do we need?
:rtype None:
"""
for j in range(row_num):
for i in range(col_num):
table.add(Paragraph(" "))
return None
def field_and_filler(user_blob: typing.Dict, table: Table, fieldname: str) -> None:
"""Adds a field to the table along with a filler cell to span the table.
:param str fieldname:
:param Dict user_blob: A dict containing the user's info.
:param Table table: The table to be modified.
:rtype None:
"""
table.add(Paragraph(str(user_blob[fieldname])))
# table.add(Paragraph(" "))
return None
def build_billto_table(user_blob: dict) -> FixedColumnWidthTable:
"""
This function builds a Table containing billing and shipping information
It spans the page using a single wide column.
:param Dict user_blob: A dict containing the user's info.
:rtype Table: a Table containing shipping and billing information
"""
table = FixedColumnWidthTable(number_of_rows=5, number_of_columns=1)
table.add(Paragraph("BILL TO", font="Helvetica-Bold"))
# table.add(Paragraph(" ", font="Helvetica-Bold"))
field_ids = ["fullname", "billaddress_line1", "billaddress_line2"]
list(map(lambda fieldname: field_and_filler(user_blob, table, fieldname), field_ids))
table.add(Paragraph("%s, %s" % (user_blob["billaddress_city"], user_blob["billaddress_zip"]))) # BILLING
# table.add(Paragraph(" ")) # SHIPPING
# table.add(Paragraph(" ")) # BILLING
# table.add(Paragraph(" ")) # SHIPPING
table.set_padding_on_all_cells(Decimal(2), Decimal(2), Decimal(2), Decimal(2))
if no_table_borders: table.no_borders()
return table
class Charge:
"""
This class represents a charge line.
"""
def __init__(self, name: str, quantity: float, rate: float,
amount: float, type: str, last4: str, timestamp: float):
"""
Charge constructor. Basically just turns a row into a Charge object.
:param str name:
:param float quantity:
:param float rate:
:param float amount:
:param str type:
:param str last4:
:param float timestamp:
"""
self.name: str = name
# assert quantity >= 0
self.quantity: float = quantity
# assert rate >= 0
self.rate: float = rate
# assert amount >= 0
self.amount: float = amount
self.type = type
self.last4 = last4
self.timestamp = timestamp
def format_float_val_as_currency(v) -> str:
"""
Given a float value, format it as currency. Note that the sign is backwards due to the fact that the charge_fields in
the data that comes back are positive numbers but are written as negative on the invoice. There's also the fact
that in our format string, if the value is actually negative it will print as "$-x.xx" instead of the correct
"-$x.xx". I don't know if this can be fixed with format strings so we rely on this hack.
:param float v:
:rtype str:
"""
if v > 0:
return " -${:10.2f}".format(v)
else:
return " ${:10.2f}".format(-v)
def build_charge_table(charges: typing.List[Charge], page_number: int)\
-> FlexibleColumnWidthTable:
"""
This function builds a Table containing itemized billing information
:param typing.List[Charge] charges: the rows on the invoice
:param int page_number: Current page number.
:rtype FlexibleColumnWidthTable: Borb Table containing the information.
"""
global invoice_total
num_rows = len(charges)
table = FlexibleColumnWidthTable(number_of_rows=(num_rows + 3), number_of_columns=4)
table_header_padding_top = 4
table_header_padding_bottom = 5
table_header_padding_left = 3
table_header_padding_right = 3
item_font_size = Decimal(11)
for h in ["Item", "Quantity", "Rate", "Amount"]:
table.add(
TableCell(
Paragraph(h, font_color=X11Color("White"),
padding_top=Decimal(table_header_padding_top),
padding_bottom=Decimal(table_header_padding_bottom),
padding_left=Decimal(table_header_padding_left),
padding_right=Decimal(table_header_padding_right),
vertical_alignment=Alignment.TOP),
background_color=HexColor("0b3954")
)
)
odd_color = HexColor("BBBBBB")
even_color = HexColor("FFFFFF")
for row_number, item in enumerate(charges):
c = even_color if row_number % 2 == 0 else odd_color
table.add(TableCell(Paragraph(item.name, font="Helvetica-Bold",
font_size=item_font_size), background_color=c))
if item.type == "payment":
table.add(TableCell(Paragraph(" ", horizontal_alignment=Alignment.RIGHT,
font_size=item_font_size),
background_color=c))
table.add(TableCell(Paragraph(" ", horizontal_alignment=Alignment.RIGHT,
font_size=item_font_size), background_color=c))
else:
table.add(TableCell(Paragraph(" {:10.2f}".format(item.quantity),
horizontal_alignment=Alignment.RIGHT,
font_size=item_font_size),
background_color=c))
table.add(TableCell(Paragraph(format_float_val_as_currency(item.rate),
horizontal_alignment=Alignment.RIGHT,
font_size=item_font_size),
background_color=c))
table.add(TableCell(
Paragraph(format_float_val_as_currency(item.amount),
horizontal_alignment=Alignment.RIGHT,
font_size=item_font_size), background_color=c))
if page_number == page_count:
table.add(TableCell(
Paragraph(" ", font="Helvetica-Bold", horizontal_alignment=Alignment.RIGHT),
col_span=3))
table.add(TableCell(Paragraph(" ", horizontal_alignment=Alignment.RIGHT)))
table.add(TableCell(
Paragraph("Total", font="Helvetica-Bold", horizontal_alignment=Alignment.RIGHT),
col_span=3))
table.add(TableCell(
Paragraph(format_float_val_as_currency(invoice_total), horizontal_alignment=Alignment.RIGHT)))
else:
blank_row(table, 4, 2)
table.set_padding_on_all_cells(Decimal(2), Decimal(5), Decimal(2), Decimal(5))
if no_table_borders: table.no_borders()
return table
def product_row(charge_fields) -> Charge:
"""Makes a single row with charge information in it.
:param charge_fields:
:rtype Charge:
"""
type = charge_fields["type"]
is_credit = charge_fields["is_credit"] if "is_credit" in charge_fields else False
timestamp = charge_fields["timestamp"]
last4 = charge_fields["last4"] if "last4" in charge_fields and charge_fields["last4"] is not None else " "
def credit_or_auto_billing():
formatted = datetime.datetime.fromtimestamp(timestamp, tz=datetime.timezone.utc).strftime('%Y-%m-%d %H:%M')
if is_credit:
return "Add Credit: *" + last4 + " : " + formatted
return "Auto-Billing: *" + last4 + " : " + formatted
description = charge_fields[
"description"] if "description" in charge_fields else credit_or_auto_billing() # "Add Credit: *" + last4 + ":" + str(timestamp)
amount = float(charge_fields["amount"]) if "amount" in charge_fields else 0.0
quantity = float(charge_fields["quantity"]) if "quantity" in charge_fields else 1.0
rate = float(charge_fields["rate"]) if "rate" in charge_fields else amount
return Charge(description, quantity, rate, amount, type, last4, timestamp)
def product_rows(rows_invoice) -> typing.List[Charge]:
"""
Maps a list of invoice dicts to a list of Charge objects using product_row function.
:param rows_invoice:
:rtype typing.List[Charge]:
"""
if rows_invoice is None:
rows_invoice = []
return list(map(lambda charges: product_row(charges), rows_invoice))
def build_invoice_charges_table(rows_invoice: typing.List[typing.Dict],
charges_per_page: int, page_number: int) -> FlexibleColumnWidthTable:
"""
This function creates a page of invoice charge_fields and depletes the list of charge_fields by the number it
prints out. Essentially a pop function that removes multiple items at once.
:param int page_number:
:param typing.List[typing.Dict] rows_invoice: List of rows in the invoice
:param int charges_per_page: How many rows of charge_fields we put on this page.
:rtype FlexibleColumnWidthTable:
"""
rows_invoice_chunk = rows_invoice[0:charges_per_page]
# sums = {"amount": compute_column_sum(rows_invoice, "amount", True)}
del rows_invoice[0:charges_per_page]
return build_charge_table(product_rows(rows_invoice_chunk), page_number)
def compute_column_sum(rows_invoice: typing.List[typing.Dict],
column_name: str,
values_are_negative: bool = False) -> float:
"""Sum over one of the columns in the invoice.
:param typing.List[typing.Dict] rows_invoice: List of rows in the invoice
:param str column_name: Name of column to sum over.
:param bool values_are_negative: are we summing values that are actually negative?
:rtype: float
"""
s: float = 0
for row in rows_invoice:
v: float = math.floor(float(row[column_name]) * 100) / 100
# v = round(float(row[column_name]) * 100) / 100
if values_are_negative:
s = s - v
else:
s = s + v
return math.floor(s * 100) / 100
# return round(s * 100) / 100
def generate_invoice_page(user_blob: typing.Dict,
rows_invoice: typing.List[typing.Dict],
page_number: int, date_header_text: str = "") -> Page:
"""Makes a single page of the invoice.
:param str date_header_text:
:param Dict user_blob: Dict of info about the user
:param typing.List[typing.Dict] rows_invoice: List of rows in the invoice
:param int page_number: The page number for this page.
:rtype Page:
"""
global invoice_number
if page_number == 1:
rows_per_page = num_rows_first_page
else:
rows_per_page = num_rows_subsequents_pages
page: Page = Page()
# set PageLayout
page_layout: PageLayout = SingleColumnLayout(page,
vertical_margin=page.get_page_info().get_height() * Decimal(0.02))
table_logo_and_invoice_num = build_logo_and_invoice_num_table(page_number)
page_layout.add(table_logo_and_invoice_num)
page_layout.add(Paragraph(" "))
if page_number == 1:
# Invoice information table
page_layout.add(build_2nd_block_table())
# Empty paragraph for spacing
page_layout.add(Paragraph(" "))
# Billing and shipping information table
page_layout.add(build_billto_table(user_blob))
page_layout.add(Paragraph(date_header_text + " "))
if len(rows_invoice) == 0:
# If we don't handle this case the client crashes with no output.
page_layout.add(Paragraph("NO DATA"))
else:
table_invoice_rows = build_invoice_charges_table(rows_invoice, rows_per_page, page_number)
page_layout.add(table_invoice_rows)
return page
def build_logo_and_invoice_num_table(page_number) -> FixedColumnWidthTable:
"""
At the top of every page is a table with our logo, the invoice number, and little text reading "page X of Y".
This function creates that table and returns it.
:param page_number:
:rtype FixedColumnWidthTable:
"""
if page_number == 1:
invoice_number_font_size = Decimal(20)
invoice_word_font_size = Decimal(50)
logo_img_filename: str = os.path.join(os.path.dirname(__file__), 'vast.ai-logo.png')
logo_img_width: int = 72
logo_img_height: int = 105
else:
invoice_number_font_size = Decimal(14)
invoice_word_font_size = Decimal(20)
logo_img_filename: str = os.path.join(os.path.dirname(__file__), 'vast.ai-logo-50pct.png')
logo_img_width: int = 36
logo_img_height: int = 53
logo_img = PIL.Image.open(logo_img_filename)
# add corporate logo
# now = datetime.datetime.now()
table_logo_and_invoice_num = FixedColumnWidthTable(number_of_rows=2, number_of_columns=4)
table_logo_and_invoice_num.add(
TableCell(
Image(
logo_img,
width=Decimal(logo_img_width),
height=Decimal(logo_img_height),
), row_span=2)
)
table_logo_and_invoice_num.add(
Paragraph("Page %d of %d" % (page_number, page_count), font="Helvetica", horizontal_alignment=Alignment.RIGHT))
table_logo_and_invoice_num.add(Paragraph(" ", font="Helvetica-Bold", horizontal_alignment=Alignment.RIGHT))
table_logo_and_invoice_num.add(Paragraph("Invoice", font="Helvetica", font_size=invoice_word_font_size,
horizontal_alignment=Alignment.RIGHT))
table_logo_and_invoice_num.add(Paragraph(" ", font="Helvetica-Bold", horizontal_alignment=Alignment.RIGHT))
table_logo_and_invoice_num.add(Paragraph(" ", font="Helvetica-Bold", horizontal_alignment=Alignment.RIGHT))
table_logo_and_invoice_num.add(Paragraph("# %d" % invoice_number, font="Helvetica",
font_size=invoice_number_font_size, horizontal_alignment=Alignment.RIGHT))
if no_table_borders: table_logo_and_invoice_num.no_borders()
return table_logo_and_invoice_num
def compute_pages_needed(rows_invoice: typing.List[typing.Dict]) -> int:
"""Function to work out how many pages we need so that the page_count can be filled in at the top of every page.
:param typing.List[typing.Dict] rows_invoice: The list of dicts we use elsewhere.
:rtype int:
"""
num_rows_invoice = len(rows_invoice)
num_rows_invoice = num_rows_invoice - num_rows_first_page
page_count: int = math.ceil(num_rows_invoice / num_rows_subsequents_pages) + 1
return page_count
def generate_invoice(user_blob: typing.Dict,
rows_invoice: typing.List[typing.Dict], filter_data: typing.Dict) -> None:
"""
This is the main function in this library. It calls everything else and makes the invoice page by page. The
resulting invoice is written as a single PDF file.
:param filter_data: Parameters for client side filter.
:param user_blob: info about the user
:param typing.List[typing.Dict] rows_invoice: The list of dicts we use elsewhere.
:rtype None:
"""
pdf: Document = Document()
global page_count
page_count = compute_pages_needed(rows_invoice)
global invoice_total
invoice_total = compute_column_sum(rows_invoice, "amount")
page_number = 1
if len(rows_invoice) == 0:
page = generate_invoice_page(user_blob, rows_invoice, page_number, filter_data["header_text"])
print("Adding Empty page ", str(page_number))
pdf.append_page(page)
else:
while len(rows_invoice) > 0:
page = generate_invoice_page(user_blob, rows_invoice, page_number, filter_data["header_text"])
print("Adding page ", str(page_number))
pdf.append_page(page)
page_number += 1
# We write out the latest PDF so that we can watch the file change with `evince` or similar viewer even though
# parameters may differ from run to run.
with open("latest-invoice.pdf", "wb") as debug_file_handle:
PDF.dumps(debug_file_handle, pdf)
debug_file_handle.close()
with open(filter_data["pdf_filename"], "wb") as users_file_handle:
PDF.dumps(users_file_handle, pdf)
users_file_handle.close()
pdf/watcher.sh
#!/usr/bin/env bash
# This little utility notices when the timestamps change in this directory
# and then reruns the pdf generator script. If you have something like
# evince running that reloads the PDF when it changes this can be very
# convenient.
while true
do
watch -g ls -l *.py *.sh *.png *.pdf; echo "Something changed. Sleeping 1..."; sleep 1; vastai generate pdf-invoices
done
sdk.py
"""Main VastAI SDK class that wraps all api/ module functions."""
from __future__ import annotations
import warnings
from typing import Dict, List, Optional, Union
from vastai._base import _resolve_api_key, _APIKEY_SENTINEL
from vastai.api.client import VastClient
from vastai.api.query import (parse_order, parse_query, benchmarks_fields,
fix_date_fields, invoices_fields,
offers_mult, templates_fields,
vol_offers_fields)
from vastai.api import instances, offers, machines, teams, keys, endpoints, billing, storage, clusters, auth, deployments
def benchmarks_query(query):
"""Parse a benchmarks query string into the filter dict the API wants."""
if isinstance(query, str):
query = parse_query(query, {}, benchmarks_fields)
return fix_date_fields(query, ["last_update"])
return query
def volume_query(query):
"""Parse a volume query string; the api layer seeds the defaults."""
if isinstance(query, str):
return parse_query(query, {}, vol_offers_fields, {}, offers_mult)
return query
def template_fields(*, runtype=None, use_ssh=None, jup_direct=None,
ssh_direct=None, use_jupyter_lab=None,
docker_login_repo=None, extra_filters=None,
readme_visible=None, private=None, **flags) -> dict:
"""Friendly template flags, with the raw api fields taking precedence.
update_template forwarded **kwargs straight to the api layer, so the raw
field names are what every working script uses today -- the friendly ones
all raised TypeError. Both spellings have to keep working.
"""
fields = offers.template_fields_from_flags(**flags)
overrides = {
"runtype": runtype, "use_ssh": use_ssh, "jup_direct": jup_direct,
"ssh_direct": ssh_direct, "use_jupyter_lab": use_jupyter_lab,
"docker_login_repo": docker_login_repo, "extra_filters": extra_filters,
"readme_visible": readme_visible, "private": private,
}
fields.update({k: v for k, v in overrides.items() if v is not None})
return fields
class VastAI:
"""High-level SDK for the Vast.ai platform.
Wraps the low-level api/ modules behind a single object with one method
per API operation. Every method delegates to the corresponding module
function, passing ``self.client`` (a :class:`VastClient`) along with any
caller-supplied arguments.
Args:
api_key: Vast.ai API key. When omitted, the key is resolved from
``VAST_API_KEY``, then ``$XDG_CONFIG_HOME/vastai/vast_api_key``
(falling back to ``~/.config/vastai/vast_api_key``), then the
legacy ``~/.vast_api_key``.
server_url: Base URL of the Vast.ai API server.
retry: Number of retries on transient HTTP errors.
raw: If *True*, return raw JSON dicts instead of formatted output.
explain: If *True*, print request details for debugging.
quiet: If *True*, suppress informational output.
curl: If *True*, print equivalent curl commands instead of executing.
"""
def __init__(
self,
api_key: Optional[str] = None,
server_url: Optional[str] = None,
retry: int = 3,
raw: bool = False,
explain: bool = False,
quiet: bool = False,
curl: bool = False,
):
resolved_key = _resolve_api_key(_APIKEY_SENTINEL if api_key is None else api_key)
self.client = VastClient(resolved_key, server_url, retry, explain, curl)
self.raw = raw
self.quiet = quiet
# ------------------------------------------------------------------
# Instance methods
# ------------------------------------------------------------------
def show_instances(self) -> list[dict]:
"""Return all of the authenticated user's instances as a flat list."""
return instances.show_instances(self.client)
def show_instances_v1(self, params: dict) -> dict:
"""Return instances using the paginated v1 API; for filtering, sorting, and manual pagination."""
return instances.show_instances_v1(self.client, params)
def show_instance_filters(self) -> list:
"""Return distinct filterable values for instances."""
return instances.show_instance_filters(self.client)
def show_instance(self, id: int) -> Optional[dict]:
"""Return details of a single instance."""
return instances.show_instance(self.client, id)
def create_instance(
self, id: int, image: Optional[str] = None, disk: float = 10,
env: Optional[Union[str, dict]] = None, price: Optional[float] = None,
label: Optional[str] = None, extra: Optional[str] = None,
onstart_cmd: Optional[str] = None, login: Optional[str] = None,
python_utf8: bool = False, lang_utf8: bool = False,
jupyter_lab: bool = False, jupyter_dir: Optional[str] = None,
force: bool = False, cancel_unavail: bool = False,
template_hash: Optional[str] = None, user: Optional[str] = None,
runtype: Optional[str] = None, args: Optional[list] = None,
volume_info: Optional[dict] = None,
bid_price: Optional[float] = None, ssh: bool = False,
jupyter: bool = False, direct: bool = False,
create_volume: Optional[int] = None, link_volume: Optional[int] = None,
volume_size: Optional[float] = None, mount_path: Optional[str] = None,
volume_label: Optional[str] = None,
) -> dict:
"""Create a new instance from a contract offer ID.
``runtype`` is left out of the request when neither it nor the
ssh/jupyter/direct flags are given, so the backend picks the default.
"""
return instances.create_instance(
self.client, id,
image=image, disk=disk, env=env, price=price, bid_price=bid_price,
label=label, extra=extra, onstart_cmd=onstart_cmd, login=login,
python_utf8=python_utf8, lang_utf8=lang_utf8, jupyter_lab=jupyter_lab,
jupyter_dir=jupyter_dir, force=force, cancel_unavail=cancel_unavail,
template_hash=template_hash, user=user, runtype=runtype, ssh=ssh,
jupyter=jupyter, direct=direct, args=args, volume_info=volume_info,
create_volume=create_volume, link_volume=link_volume,
volume_size=volume_size, mount_path=mount_path, volume_label=volume_label,
)
def destroy_instance(self, id: int) -> dict:
"""Destroy an instance."""
return instances.destroy_instance(self.client, id)
def start_instance(self, id: int) -> dict:
"""Start a stopped instance."""
return instances.start_instance(self.client, id)
def stop_instance(self, id: int) -> dict:
"""Stop a running instance."""
return instances.stop_instance(self.client, id)
def reboot_instance(self, id: int) -> dict:
"""Reboot an instance."""
return instances.reboot_instance(self.client, id)
def recycle_instance(self, id: int) -> dict:
"""Recycle an instance."""
return instances.recycle_instance(self.client, id)
def label_instance(self, id: int, label: str) -> dict:
"""Set a label on an instance."""
return instances.label_instance(self.client, id, label)
def prepay_instance(self, id: int, amount: float) -> dict:
"""Prepay for an instance."""
return instances.prepay_instance(self.client, id, amount)
def change_bid(self, id: int, price: Optional[float] = None) -> dict:
"""Change the bid price for an instance."""
return instances.change_bid(self.client, id, price=price)
def show_pending_price_increases(self) -> dict:
"""Return the pending-price-increase envelope for the authenticated user.
See ``vastai.api.price_increase.list_pending`` for the response shape.
"""
from vastai.api import price_increase
return price_increase.list_pending(self.client)
def accept_price_increase(self, instance_id: int) -> dict:
"""Accept the pending price increase on ``instance_id``.
Resolves the instance's pending row via the pending-list endpoint
and PUTs the per-row accept. The new prices and extended end date
apply only after the contract's current ``end_date``.
Raises ``LookupError`` when no pending row matches the instance id.
"""
from vastai.api import price_increase
match = price_increase.resolve_instance_to_pending(self.client, instance_id)
return price_increase.accept(self.client, match["pending_price_increase_id"])
def reject_price_increase(self, instance_id: int) -> dict:
"""Reject the pending price increase on ``instance_id``.
Tombstones the matching pending row on the backend; no cutover
follows. Raises ``LookupError`` when no pending row matches.
"""
from vastai.api import price_increase
match = price_increase.resolve_instance_to_pending(self.client, instance_id)
return price_increase.reject(self.client, match["pending_price_increase_id"])
def execute(self, id: int, command: str) -> dict:
"""Execute a command on an instance."""
return instances.execute(self.client, id, command)
def logs(
self,
instance_id: int,
tail: Optional[str] = None,
filter: Optional[str] = None,
daemon_logs: bool = False,
) -> str:
"""Retrieve logs for an instance."""
return instances.logs(self.client, instance_id, tail=tail, filter=filter, daemon_logs=daemon_logs)
def ssh_url(self, id: int) -> str:
"""Get the SSH URL for an instance."""
inst = instances.show_instance(self.client, id)
if isinstance(inst, list):
inst = inst[0] if inst else {}
if not isinstance(inst, dict):
return ""
port = inst.get("ssh_port") or (inst.get("ports") or {}).get("22/tcp", [{}])[0].get("HostPort")
ip = inst.get("ssh_host") or inst.get("public_ipaddr")
return f"ssh://root@{ip}:{port}" if port and ip else ""
def scp_url(self, id: int) -> str:
"""Get the SCP URL for an instance."""
inst = instances.show_instance(self.client, id)
if isinstance(inst, list):
inst = inst[0] if inst else {}
if not isinstance(inst, dict):
return ""
port = inst.get("ssh_port") or (inst.get("ports") or {}).get("22/tcp", [{}])[0].get("HostPort")
ip = inst.get("ssh_host") or inst.get("public_ipaddr")
return f"scp://root@{ip}:{port}" if port and ip else ""
def take_snapshot(self, instance_id, **kwargs) -> dict:
"""Take a container snapshot and push to a registry."""
return instances.take_snapshot(self.client, instance_id, **kwargs)
# ------------------------------------------------------------------
# Deployment methods
# ------------------------------------------------------------------
def show_deployments(self) -> list[dict]:
"""Return all deployments for the authenticated user."""
return deployments.show_deployments(self.client)
def show_deployment(self, id: int) -> dict:
"""Return details of a single deployment."""
return deployments.show_deployment(self.client, id)
def show_deployment_versions(self, id: int) -> list:
"""Return version history for a deployment."""
return deployments.show_deployment_versions(self.client, id)
def delete_deployment(self, id: int) -> dict:
"""Delete a deployment."""
return deployments.delete_deployment(self.client, id)
# ------------------------------------------------------------------
# Offer / search methods
# ------------------------------------------------------------------
def search_offers(
self,
query=None,
type: str = "on-demand",
order: str = "score-",
limit: Optional[int] = None,
storage: float = 5.0,
no_default: bool = False,
) -> list:
"""Search for GPU offers.
Args:
query: Query string (e.g. "num_gpus=1 gpu_name=RTX_4090") or pre-parsed dict.
type: One of "on-demand", "reserved", or "bid".
order: Comma-separated sort fields, e.g. "score-" or "dph_total".
limit: Max results.
storage: Allocated storage in GiB for pricing.
no_default: Skip default filters (verified, rentable, etc.).
"""
from vastai.api.query import parse_query, offers_fields, offers_alias, offers_mult
from vastai.utils import preprocess_search_query, postprocess_search_results
# Expand georegion/chunked directives before parsing.
# Seed defaults before parsing so `field=any` correctly removes them,
# matching CLI behavior. Pass no_default=True to the helper afterward
# so it does not reapply the same defaults.
georegion_active, chunked = False, False
defaults_applied = False
if isinstance(query, str):
georegion_active, chunked, query = preprocess_search_query(query)
base = {} if no_default else {
"verified": {"eq": True}, "external": {"eq": False},
"rentable": {"eq": True},
}
query = parse_query(query, base, offers_fields, offers_alias, offers_mult)
defaults_applied = True
# Parse order string into list
order_list = parse_order(order)
results = offers.search_offers(
self.client, query=query, offer_type=type, order=order_list,
limit=limit, storage=storage,
no_default=(no_default or defaults_applied),
)
if isinstance(results, list):
results = postprocess_search_results(results, georegion_active=georegion_active, chunked=chunked)
return results
def search_templates(self, query: Optional[Union[str, dict]] = None) -> list[dict]:
"""Search for templates.
``query`` takes the same string syntax as ``vastai search templates``
(e.g. "name=pytorch"), or an already-parsed filter dict.
"""
if isinstance(query, str):
query = parse_query(query, {}, templates_fields)
query = fix_date_fields(query, ["created_at", "recent_create_date"])
return offers.search_templates(self.client, query=query)
def search_benchmarks(self, query: Optional[Union[str, dict]] = None,
order: Optional[list] = None,
limit: Optional[int] = None,
after_token: Optional[str] = None,
all_pages: bool = False) -> list[dict]:
"""Search for benchmarks.
Returns a single page (the backend caps it at 200 rows) so a broad query
doesn't scan the whole table. Pass all_pages=True to follow pagination
and get every matching row, or use search_benchmarks_v1() to page
manually via next_token.
"""
query = benchmarks_query(query)
if all_pages:
return offers.search_benchmarks(self.client, query=query, order=order,
limit=limit, after_token=after_token)
params = offers.benchmarks_query_args(query=query, order=order,
limit=limit, after_token=after_token)
data = offers.search_benchmarks_v1(self.client, params)
return data.get("benchmarks") or []
def search_benchmarks_v1(self, params: dict) -> dict:
"""Return benchmarks using the paginated API; for manual pagination."""
return offers.search_benchmarks_v1(self.client, params)
def search_volumes(
self, query: Optional[Union[str, dict]] = None,
order: Optional[Union[str, list]] = None, limit: Optional[int] = None,
storage: float = 1.0, no_default: bool = False,
) -> list[dict]:
"""Search for volume offers.
``query`` takes the same string syntax as ``vastai search volumes``,
or an already-parsed filter dict.
"""
return offers.search_volumes(
self.client, query=volume_query(query),
order=parse_order(order), limit=limit, storage=storage,
no_default=no_default,
)
def search_network_volumes(
self, query: Optional[Union[str, dict]] = None,
order: Optional[Union[str, list]] = None, limit: Optional[int] = None,
storage: float = 1.0, no_default: bool = False,
) -> list[dict]:
"""Search for network volume offers.
``query`` takes the same string syntax as
``vastai search network volumes``, or an already-parsed filter dict.
"""
return offers.search_network_volumes(
self.client, query=volume_query(query),
order=parse_order(order), limit=limit, storage=storage,
no_default=no_default,
)
def search_invoices(self, query: Optional[Union[str, dict]] = None) -> list[dict]:
"""Search for invoices.
``query`` takes the same string syntax as ``vastai search invoices``,
or an already-parsed filter dict.
"""
if isinstance(query, str):
query = parse_query(query, {}, invoices_fields)
query = fix_date_fields(query, ["when", "paid_on", "payment_expected"])
return offers.search_invoices(self.client, query=query)
def search_offers_new(
self,
query=None,
type: str = "on-demand",
order: str = "score-",
limit: Optional[int] = None,
storage: float = 5.0,
no_default: bool = False,
) -> list:
"""Search for GPU offers using the new /search/asks/ endpoint.
Args:
query: Query string (e.g. "num_gpus=1 gpu_name=RTX_4090") or pre-parsed dict.
type: One of "on-demand", "reserved", or "bid".
order: Comma-separated sort fields, e.g. "score-" or "dph_total".
limit: Max results.
storage: Allocated storage in GiB for pricing.
no_default: Skip default filters (verified, rentable, etc.).
"""
from vastai.api.query import parse_query, offers_fields, offers_alias, offers_mult
from vastai.utils import preprocess_search_query, postprocess_search_results
georegion_active, chunked = False, False
defaults_applied = False
if isinstance(query, str):
georegion_active, chunked, query = preprocess_search_query(query)
base = {} if no_default else {
"verified": {"eq": True}, "external": {"eq": False},
"rentable": {"eq": True},
}
query = parse_query(query, base, offers_fields, offers_alias, offers_mult)
defaults_applied = True
order_list = parse_order(order)
results = offers.search_offers_new(
self.client, query=query, offer_type=type, order=order_list,
limit=limit, storage=storage,
no_default=(no_default or defaults_applied),
)
if isinstance(results, list):
results = postprocess_search_results(results, georegion_active=georegion_active, chunked=chunked)
return results
def launch_instance(
self, gpu_name: str, num_gpus: str, image: str,
region: Optional[str] = None, disk: float = 10, order: str = "score-",
limit: Optional[int] = None, env: Optional[Union[str, dict]] = None,
label: Optional[str] = None, extra: Optional[str] = None,
onstart_cmd: Optional[str] = None, login: Optional[str] = None,
python_utf8: bool = False, lang_utf8: bool = False,
jupyter_lab: bool = False, jupyter_dir: Optional[str] = None,
cancel_unavail: bool = False, template_hash: Optional[str] = None,
runtype: Optional[str] = None, args: Optional[list] = None,
query: Optional[dict] = None,
ssh: bool = False, jupyter: bool = False, direct: bool = False,
) -> dict:
"""Launch the top instance from search offers matching the given criteria."""
return offers.launch_instance(
self.client, gpu_name, num_gpus, image, region=region, disk=disk,
order=order, limit=limit, env=env, label=label, extra=extra,
onstart_cmd=onstart_cmd, login=login, python_utf8=python_utf8,
lang_utf8=lang_utf8, jupyter_lab=jupyter_lab, jupyter_dir=jupyter_dir,
cancel_unavail=cancel_unavail, template_hash=template_hash,
runtype=runtype, ssh=ssh, jupyter=jupyter, direct=direct, args=args,
query=query,
)
# ------------------------------------------------------------------
# Machine methods
# ------------------------------------------------------------------
def show_machines(self) -> list[dict]:
"""Return all hosted machines."""
return machines.show_machines(self.client)
def show_machine(self, id: int) -> dict:
"""Return details of a single machine.
The underlying ``GET /machines/{id}`` endpoint returns a one-element
list; this wrapper unwraps it so callers get a single machine dict.
Raises ``ValueError`` if the backend returns zero or multiple rows.
"""
result = machines.show_machine(self.client, id)
if not isinstance(result, list):
return result
if not result:
raise ValueError(f"Machine {id} not found")
if len(result) > 1:
raise ValueError(f"Expected 1 machine for id={id}, got {len(result)}")
return result[0]
def show_maints(self, ids) -> list[dict]:
"""Show maintenance information for machines."""
return machines.show_maints(self.client, ids)
def list_machine(self, id: int, **kwargs) -> dict:
"""List a machine for rent with optional pricing parameters."""
return machines.list_machine(self.client, id, **kwargs)
def list_machines(self, ids, **kwargs) -> list[dict]:
"""List multiple machines for rent."""
results = []
for mid in ids:
results.append(machines.list_machine(self.client, mid, **kwargs))
return results
def unlist_machine(self, id: int) -> dict:
"""Unlist a machine from being available."""
return machines.unlist_machine(self.client, id)
def set_defjob(self, id: int, **kwargs) -> dict:
"""Set the default job on a machine."""
return machines.set_defjob(self.client, id, **kwargs)
def remove_defjob(self, id: int) -> dict:
"""Remove the default job from a machine."""
return machines.remove_defjob(self.client, id)
def set_min_bid(self, id: int, price: Optional[float] = None) -> dict:
"""Set the minimum bid price for a machine."""
return machines.set_min_bid(self.client, id, price=price)
def schedule_maint(self, id: int, sdate, duration, category: str = "not provided") -> dict:
"""Schedule maintenance for a machine."""
return machines.schedule_maint(self.client, id, sdate, duration, maintenance_category=category)
def cancel_maint(self, id: int) -> dict:
"""Cancel scheduled maintenance for a machine."""
return machines.cancel_maint(self.client, id)
def cleanup_machine(self, id: int) -> dict:
"""Clean up a machine's configuration and resources."""
return machines.cleanup_machine(self.client, id)
def defrag_machines(self, ids) -> dict:
"""Defragment machines."""
return machines.defrag_machines(self.client, ids)
def delete_machine(self, id: int) -> dict:
"""Delete a machine if not being used by clients."""
return machines.delete_machine(self.client, id)
def reports(self, id: int) -> list[dict]:
"""Generate reports for a machine."""
return machines.reports(self.client, id)
# ------------------------------------------------------------------
# Team methods
# ------------------------------------------------------------------
def create_team(self, team_name: str) -> dict:
"""Create a new team."""
return teams.create_team(self.client, team_name)
def destroy_team(self) -> dict:
"""Destroy the current team."""
return teams.destroy_team(self.client)
def show_members(self) -> list[dict]:
"""Show all team members."""
return teams.show_members(self.client)
def invite_member(self, email: str, role: str) -> dict:
"""Invite a new member to the team."""
return teams.invite_member(self.client, email, role)
def remove_member(self, id: int) -> dict:
"""Remove a member from the team."""
return teams.remove_member(self.client, id)
def show_team_roles(self) -> list[dict]:
"""Show all team roles."""
return teams.show_team_roles(self.client)
def show_team_role(self, name: str) -> dict:
"""Show details of a specific team role."""
return teams.show_team_role(self.client, name)
def create_team_role(self, name: str, permissions) -> dict:
"""Create a new team role."""
return teams.create_team_role(self.client, name, permissions)
def remove_team_role(self, name: str) -> dict:
"""Remove a team role."""
return teams.remove_team_role(self.client, name)
# ------------------------------------------------------------------
# SSH / API key methods
# ------------------------------------------------------------------
def show_ssh_keys(self) -> list[dict]:
"""Show all SSH keys."""
return keys.show_ssh_keys(self.client)
def create_ssh_key(self, ssh_key: Optional[str] = None) -> dict:
"""Create a new SSH key."""
return keys.create_ssh_key(self.client, ssh_key=ssh_key)
def delete_ssh_key(self, id: int) -> dict:
"""Delete an SSH key."""
return keys.delete_ssh_key(self.client, id)
def attach_ssh(self, instance_id: int, ssh_key: str) -> dict:
"""Attach an SSH key to an instance."""
return keys.attach_ssh(self.client, instance_id, ssh_key)
def detach_ssh(self, instance_id: int, ssh_key_id: str) -> dict:
"""Detach an SSH key from an instance."""
return keys.detach_ssh(self.client, instance_id, ssh_key_id)
def show_api_keys(self) -> list[dict]:
"""Return all API keys associated with the account.
The underlying ``GET /auth/apikeys/`` endpoint returns an envelope dict
``{"apikeys": [...]}``; this wrapper unwraps it so callers get a plain
list of API key dicts.
"""
result = keys.show_api_keys(self.client)
if isinstance(result, dict) and "apikeys" in result:
return result["apikeys"]
return result
def show_api_key(self, id: int) -> dict:
"""Show details of an API key."""
return keys.show_api_key(self.client, id)
def create_api_key(self, name, permissions, key_params=None) -> dict:
"""Create a new API key."""
return keys.create_api_key(self.client, name, permissions, key_params=key_params)
def delete_api_key(self, id: int) -> dict:
"""Delete an API key."""
return keys.delete_api_key(self.client, id)
def reset_api_key(self) -> dict:
"""Reset the API key."""
return keys.reset_api_key(self.client)
def update_ssh_key(self, id: int, ssh_key: str) -> dict:
"""Update an SSH key."""
return keys.update_ssh_key(self.client, id, ssh_key)
# ------------------------------------------------------------------
# Endpoint methods
# ------------------------------------------------------------------
def show_endpoints(self) -> list[dict]:
"""Show all serverless endpoints."""
return endpoints.show_endpoints(self.client)
def create_endpoint(self, **kwargs) -> dict:
"""Create a new serverless endpoint."""
return endpoints.create_endpoint(self.client, **kwargs)
def delete_endpoint(self, id: int) -> dict:
"""Delete a serverless endpoint."""
return endpoints.delete_endpoint(self.client, id)
def get_endpt_logs(self, id: int, level: int = 1, tail: Optional[int] = None) -> dict:
"""Fetch logs for a serverless endpoint."""
return endpoints.get_endpt_logs(self.client, id, level=level, tail=tail)
def show_workergroups(self) -> list[dict]:
"""Show all worker groups."""
return endpoints.show_workergroups(self.client)
def create_workergroup(self, **kwargs) -> dict:
"""Create a new autoscale worker group."""
return endpoints.create_workergroup(self.client, **kwargs)
def delete_workergroup(self, id: int) -> dict:
"""Delete a worker group."""
return endpoints.delete_workergroup(self.client, id)
def get_wrkgrp_logs(self, id: int, level: int = 1, tail: Optional[int] = None) -> dict:
"""Fetch logs for a worker group."""
return endpoints.get_wrkgrp_logs(self.client, id, level=level, tail=tail)
def get_endpoint_workers(self, id: int):
"""List workers under a given endpoint, with live status and measured_perf."""
return endpoints.get_endpoint_workers(self.client, id)
# ------------------------------------------------------------------
# Billing methods
# ------------------------------------------------------------------
def show_invoices(
self, start_date: Optional[str] = None, end_date: Optional[str] = None,
only_charges: bool = False, only_credits: bool = False,
) -> dict:
"""Show invoice details (deprecated; use show_invoices_v1).
Returns dict with 'invoices' list and 'current' charges.
"""
warnings.warn(
"VastAI.show_invoices() is deprecated; use VastAI.show_invoices_v1(**params) for the paginated v1 API.",
DeprecationWarning, stacklevel=2,
)
return billing.show_invoices(
self.client, start_date=start_date, end_date=end_date,
only_charges=only_charges, only_credits=only_credits,
)
def show_invoices_v1(self, **kwargs) -> dict:
"""Get billing history reports with advanced filtering and pagination."""
return billing.show_invoices_v1(self.client, kwargs)
def show_earnings(
self, start_date: Optional[str] = None, end_date: Optional[str] = None,
machine_id: Optional[int] = None,
) -> list[dict]:
"""Show earnings information."""
return billing.show_earnings(
self.client, start_date=start_date, end_date=end_date,
machine_id=machine_id,
)
def show_deposit(self, id: int) -> dict:
"""Show deposit details."""
return billing.show_deposit(self.client, id)
def show_user(self) -> dict:
"""Show current user details."""
return billing.show_user(self.client)
def set_user(self, params) -> dict:
"""Set user parameters."""
return billing.set_user(self.client, params)
def show_subaccounts(self) -> list[dict]:
"""Show all subaccounts."""
return billing.show_subaccounts(self.client)
def create_subaccount(
self,
email: str,
username: str,
password: str,
type: Optional[str] = None,
) -> dict:
"""Create a new subaccount."""
host_only = type is not None and type.lower() == "host"
return billing.create_subaccount(self.client, email, username, password, host_only=host_only)
def show_ipaddrs(self) -> list[dict]:
"""Show IP addresses."""
return billing.show_ipaddrs(self.client)
def fetch_contracts(self, label: Optional[str] = None, contract_ids: Optional[list] = None) -> list:
"""Fetch contracts, optionally filtered by label."""
return billing.fetch_contracts(self.client, label=label, contract_ids=contract_ids)
# ------------------------------------------------------------------
# Storage methods
# ------------------------------------------------------------------
def copy(self, src: str, dst: str) -> dict:
"""Copy files between instances, volumes, cloud services, and local.
Supported location formats:
- ``[instance_id:]path`` legacy format, still supported
- ``C.instance_id:path`` container copy format
- ``cloud_service:path`` cloud service format
- ``cloud_service.id:path`` cloud service with ID
- ``local:path`` explicit local path
- ``V.volume_id:path`` volume copy
Volume copy is currently only supported for copying to other volumes,
instances, or cloud services, not local.
Do not copy to /root or / as a destination directory, as this can mess
up the permissions on the instance ssh folder, breaking future copy
operations (they use ssh authentication). See
https://vast.ai/docs/gpu-instances/data-movement#constraints for more
information about constraints.
Args:
src: Source in vast URL format, e.g. "instance_id:/path" or just "/local/path".
dst: Destination in vast URL format.
Examples:
vast.copy("6003036:/workspace/", "6003038:/workspace/")
vast.copy("C.11824:/data/test", "local:data/test")
vast.copy("local:data/test", "C.11824:/data/test")
vast.copy("drive:/folder/file.txt", "C.6003036:/workspace/")
vast.copy("s3.101:/data/", "C.6003036:/workspace/")
vast.copy("V.1234:/file", "C.5678:/workspace/")
vast.copy("V.1234:/file", "s3.101:/workspace/")
"""
from vastai.utils import parse_vast_url
src_id, src_path = parse_vast_url(src)
dst_id, dst_path = parse_vast_url(dst)
return storage.copy(self.client, src_id, dst_id, src_path, dst_path)
def cancel_copy(self, dst_id) -> dict:
"""Cancel a file copy operation."""
return storage.cancel_copy(self.client, dst_id)
def cancel_sync(self, dst_id) -> dict:
"""Cancel a file sync operation."""
return storage.cancel_sync(self.client, dst_id)
def cloud_copy(
self, src: str, dst: str = "/workspace", instance: Optional[str] = None,
connection: Optional[str] = None, transfer: str = "Instance to Cloud",
dry_run: bool = False, size_only: bool = False,
ignore_existing: bool = False, update: bool = False,
delete_excluded: bool = False, flags: Optional[list] = None,
) -> dict:
"""Copy files between cloud and instance.
Scheduling (the CLI's --schedule/--day/--hour) is not available here;
it posts to a different endpoint than this one.
"""
if flags is None:
flags = storage.rclone_flags(
dry_run=dry_run, size_only=size_only,
ignore_existing=ignore_existing, update=update,
delete_excluded=delete_excluded,
)
return storage.cloud_copy(
self.client, src=src, dst=dst, instance=instance,
connection=connection, transfer=transfer, flags=flags,
)
def clone_volume(self, source: int, dest: int, **kwargs) -> dict:
"""Clone an existing volume."""
return storage.clone_volume(self.client, source, dest, **kwargs)
def show_volumes(self, type: str = "all") -> list[dict]:
"""Show stats on owned volumes."""
return storage.show_volumes(self.client, type=type)
def show_instance_backups(self, contract_id: Optional[int] = None) -> list[dict]:
"""List the backups Vast has taken of your instances and volumes."""
return storage.show_instance_backups(self.client, contract_id=contract_id)
def instance_backup_files(self, contract_id: int) -> dict:
"""Objects stored for one backup, plus a short-lived token to fetch them."""
return storage.instance_backup_files(self.client, contract_id)
def create_volume(self, id: int, size: float = 15, name: Optional[str] = None) -> dict:
"""Create a new volume from an offer ID."""
return storage.create_volume(self.client, id, size=size, name=name)
def delete_volume(self, id: int) -> dict:
"""Delete a volume."""
return storage.delete_volume(self.client, id)
def list_volume(self, id: int, **kwargs) -> dict:
"""List disk space for rent as a volume."""
return storage.list_volume(self.client, id, **kwargs)
def unlist_volume(self, id: int) -> dict:
"""Unlist a volume offer."""
return storage.unlist_volume(self.client, id)
def create_network_volume(self, id: int, size: float = 15, name: Optional[str] = None) -> dict:
"""Create a new network volume."""
return storage.create_network_volume(self.client, id, size=size, name=name)
def list_network_volume(self, disk_id: int, **kwargs) -> dict:
"""List disk space for rent as a network volume."""
return storage.list_network_volume(self.client, disk_id, **kwargs)
def unlist_network_volume(self, id: int) -> dict:
"""Unlist a network volume offer."""
return storage.unlist_network_volume(self.client, id)
def show_network_disks(self) -> dict:
"""Show network disks associated with your account."""
return storage.show_network_disks(self.client)
def add_network_disk(self, machines: List[int], mount_point: str, disk_id: Optional[int] = None) -> dict:
"""Add a network disk to a physical cluster."""
return storage.add_network_disk(self.client, machines, mount_point, disk_id=disk_id)
def show_connections(self) -> list[dict]:
"""Show all connections."""
return storage.show_connections(self.client)
# ------------------------------------------------------------------
# Cluster methods
# ------------------------------------------------------------------
def show_clusters(self) -> dict:
"""Show clusters associated with your account."""
return clusters.show_clusters(self.client)
def create_cluster(self, subnet: str, manager_id: int) -> dict:
"""Create a Vast cluster."""
return clusters.create_cluster(self.client, subnet, manager_id)
def delete_cluster(self, cluster_id: int) -> dict:
"""Delete a cluster."""
return clusters.delete_cluster(self.client, cluster_id)
def join_cluster(self, cluster_id: int, machine_ids: List[int]) -> dict:
"""Join machines to a cluster."""
return clusters.join_cluster(self.client, cluster_id, machine_ids)
def remove_machine_from_cluster(
self,
cluster_id: int,
machine_id: int,
new_manager_id: Optional[int] = None,
) -> dict:
"""Remove a machine from a cluster."""
return clusters.remove_machine_from_cluster(
self.client, cluster_id, machine_id, new_manager_id=new_manager_id,
)
def show_overlays(self) -> list[dict]:
"""Show overlays associated with your account."""
return clusters.show_overlays(self.client)
def create_overlay(self, cluster_id: int, name: str) -> dict:
"""Create an overlay network on a physical cluster."""
return clusters.create_overlay(self.client, cluster_id, name)
def delete_overlay(self, overlay_identifier: Optional[str] = None) -> dict:
"""Delete an overlay and remove all associated instances."""
return clusters.delete_overlay(self.client, overlay_identifier=overlay_identifier)
def join_overlay(self, name: str, instance_id: int) -> dict:
"""Add an instance to an overlay network."""
return clusters.join_overlay(self.client, name, instance_id)
# ------------------------------------------------------------------
# Auth / account methods
# ------------------------------------------------------------------
def set_api_key(self, api_key: str) -> None:
"""Update the API key used by this client."""
self.client.api_key = api_key
def show_audit_logs(self) -> list[dict]:
"""Display account audit logs."""
return auth.show_audit_logs(self.client)
def show_env_vars(self, show_values: bool = False) -> dict:
"""Show user environment variables.
Args:
show_values: If True, return actual values. If False, mask them.
"""
secrets = auth.show_env_vars(self.client)
if not show_values and isinstance(secrets, dict):
return {k: "****" for k in secrets}
return secrets
def create_env_var(self, name: str, value: str) -> dict:
"""Create a new user environment variable."""
return auth.create_env_var(self.client, name, value)
def delete_env_var(self, name: str) -> dict:
"""Delete a user environment variable."""
return auth.delete_env_var(self.client, name)
def show_scheduled_jobs(self) -> list[dict]:
"""Show scheduled jobs for the account."""
return auth.show_scheduled_jobs(self.client)
def create_scheduled_job(self, start_time, end_time, api_endpoint, request_method,
request_body, frequency, instance_id, **kwargs) -> dict:
"""Create a new scheduled job."""
return auth.create_scheduled_job(
self.client, start_time, end_time, api_endpoint, request_method,
request_body, frequency, instance_id, **kwargs,
)
def update_scheduled_job(self, id: int, request_body, **kwargs) -> dict:
"""Update an existing scheduled job."""
return auth.update_scheduled_job(self.client, id, request_body, **kwargs)
def delete_scheduled_job(self, id: int) -> dict:
"""Delete a scheduled job."""
return auth.delete_scheduled_job(self.client, id)
def create_template(
self, name: Optional[str] = None, image: Optional[str] = None,
image_tag: Optional[str] = None, href: Optional[str] = None,
repo: Optional[str] = None, login: Optional[str] = None,
env: Optional[str] = None, ssh: bool = False, jupyter: bool = False,
direct: bool = False, jupyter_dir: Optional[str] = None,
jupyter_lab: bool = False, onstart_cmd: Optional[str] = None,
search_params: Optional[str] = None, no_default: bool = False,
disk_space: Optional[float] = None, readme: Optional[str] = None,
hide_readme: bool = False, desc: Optional[str] = None,
public: bool = False,
runtype: Optional[str] = None, use_ssh: Optional[bool] = None,
jup_direct: Optional[bool] = None, ssh_direct: Optional[bool] = None,
use_jupyter_lab: Optional[bool] = None,
docker_login_repo: Optional[str] = None,
extra_filters: Optional[dict] = None,
readme_visible: Optional[bool] = None, private: Optional[bool] = None,
) -> dict:
"""Create a new template.
``search_params`` is parsed into the template's offer filters, the same
way ``vastai create template --search_params`` does it.
"""
return offers.create_template(
self.client,
name=name, image=image, image_tag=image_tag, href=href,
repo=repo, env=env, onstart_cmd=onstart_cmd, jupyter_dir=jupyter_dir,
disk_space=disk_space, readme=readme, desc=desc,
**template_fields(
ssh=ssh, jupyter=jupyter, direct=direct, jupyter_lab=jupyter_lab,
login=login, hide_readme=hide_readme, public=public,
search_params=search_params, no_default=no_default,
runtype=runtype, use_ssh=use_ssh, jup_direct=jup_direct,
ssh_direct=ssh_direct, use_jupyter_lab=use_jupyter_lab,
docker_login_repo=docker_login_repo, extra_filters=extra_filters,
readme_visible=readme_visible, private=private,
),
)
def delete_template(self, template_id: Optional[int] = None, hash_id: Optional[str] = None) -> dict:
"""Delete a template by ID or hash."""
return offers.delete_template(self.client, template_id=template_id, hash_id=hash_id)
def update_template(
self, hash_id: str, name: Optional[str] = None, image: Optional[str] = None,
image_tag: Optional[str] = None, href: Optional[str] = None,
repo: Optional[str] = None, login: Optional[str] = None,
env: Optional[str] = None, ssh: bool = False, jupyter: bool = False,
direct: bool = False, jupyter_dir: Optional[str] = None,
jupyter_lab: bool = False, onstart_cmd: Optional[str] = None,
search_params: Optional[str] = None, no_default: bool = False,
disk_space: Optional[float] = None, readme: Optional[str] = None,
hide_readme: bool = False, desc: Optional[str] = None,
public: bool = False,
runtype: Optional[str] = None, use_ssh: Optional[bool] = None,
jup_direct: Optional[bool] = None, ssh_direct: Optional[bool] = None,
use_jupyter_lab: Optional[bool] = None,
docker_login_repo: Optional[str] = None,
extra_filters: Optional[dict] = None,
readme_visible: Optional[bool] = None, private: Optional[bool] = None,
) -> dict:
"""Update an existing template by creating a new version of it with a new hash_id.
Raises:
ValueError: No template with this hash_id is visible to the caller.
"""
return offers.update_template(
self.client, hash_id=hash_id,
name=name, image=image, image_tag=image_tag, href=href,
repo=repo, env=env, onstart_cmd=onstart_cmd, jupyter_dir=jupyter_dir,
disk_space=disk_space, readme=readme, desc=desc,
**template_fields(
ssh=ssh, jupyter=jupyter, direct=direct, jupyter_lab=jupyter_lab,
login=login, hide_readme=hide_readme, public=public,
search_params=search_params, no_default=no_default,
runtype=runtype, use_ssh=use_ssh, jup_direct=jup_direct,
ssh_direct=ssh_direct, use_jupyter_lab=use_jupyter_lab,
docker_login_repo=docker_login_repo, extra_filters=extra_filters,
readme_visible=readme_visible, private=private,
),
)
# ------------------------------------------------------------------
# Batch instance methods
# ------------------------------------------------------------------
def create_instances(
self, ids: List[int], image: Optional[str] = None, disk: float = 10,
env: Optional[Union[str, dict]] = None, price: Optional[float] = None,
label: Optional[str] = None, extra: Optional[str] = None,
onstart_cmd: Optional[str] = None, login: Optional[str] = None,
python_utf8: bool = False, lang_utf8: bool = False,
jupyter_lab: bool = False, jupyter_dir: Optional[str] = None,
force: bool = False, cancel_unavail: bool = False,
template_hash: Optional[str] = None, user: Optional[str] = None,
runtype: Optional[str] = None, args: Optional[list] = None,
volume_info: Optional[dict] = None,
bid_price: Optional[float] = None, ssh: bool = False,
jupyter: bool = False, direct: bool = False,
create_volume: Optional[int] = None, link_volume: Optional[int] = None,
volume_size: Optional[float] = None, mount_path: Optional[str] = None,
volume_label: Optional[str] = None,
) -> dict:
"""Create multiple instances from a list of offer IDs."""
return instances.create_instance(
self.client, ids,
image=image, disk=disk, env=env, price=price, bid_price=bid_price,
label=label, extra=extra, onstart_cmd=onstart_cmd, login=login,
python_utf8=python_utf8, lang_utf8=lang_utf8, jupyter_lab=jupyter_lab,
jupyter_dir=jupyter_dir, force=force, cancel_unavail=cancel_unavail,
template_hash=template_hash, user=user, runtype=runtype, ssh=ssh,
jupyter=jupyter, direct=direct, args=args, volume_info=volume_info,
create_volume=create_volume, link_volume=link_volume,
volume_size=volume_size, mount_path=mount_path, volume_label=volume_label,
)
def destroy_instances(self, ids: List[int]) -> dict:
"""Destroy multiple instances."""
return instances.destroy_instance(self.client, id=ids)
def start_instances(self, ids: List[int]) -> dict:
"""Start multiple instances."""
return instances.start_instance(self.client, id=ids)
def stop_instances(self, ids: List[int]) -> dict:
"""Stop multiple instances."""
return instances.stop_instance(self.client, id=ids)
# ------------------------------------------------------------------
# Update methods
# ------------------------------------------------------------------
def update_endpoint(self, id: int, **kwargs) -> dict:
"""Update an existing endpoint."""
return endpoints.update_endpoint(self.client, id=id, **kwargs)
def update_env_var(self, name: str, value: str) -> dict:
"""Update an existing user environment variable."""
return auth.update_env_var(self.client, name, value)
def update_instance(self, id: int, **kwargs) -> dict:
"""Update/recreate an instance from a new/updated template."""
return instances.update_instance(self.client, id=id, **kwargs)
def update_workergroup(self, id: int, **kwargs) -> dict:
"""Update an existing autoscale worker group."""
return endpoints.update_workergroup(self.client, id=id, **kwargs)
def update_workers(self, id: int, cancel: bool = False) -> dict:
"""Trigger a rolling update of all workers in a workergroup, or cancel an in-progress update."""
return endpoints.update_workers(self.client, id=id, cancel=cancel)
def update_team_role(self, id: int, **kwargs) -> dict:
"""Update an existing team role."""
return teams.update_team_role(self.client, id=id, **kwargs)
# ------------------------------------------------------------------
# TFA (Two-Factor Authentication) methods
# ------------------------------------------------------------------
def tfa_activate(self, **kwargs) -> dict:
"""Activate a new 2FA method by verifying the code."""
return auth.tfa_activate(self.client, **kwargs)
def tfa_delete(self, **kwargs) -> dict:
"""Remove a 2FA method from your account."""
return auth.tfa_delete(self.client, **kwargs)
def tfa_login(self, **kwargs) -> dict:
"""Complete 2FA login by verifying code."""
return auth.tfa_login(self.client, **kwargs)
def tfa_regen_codes(self, **kwargs) -> dict:
"""Regenerate backup codes for 2FA."""
return auth.tfa_regen_codes(self.client, **kwargs)
def tfa_resend_sms(self, **kwargs) -> dict:
"""Resend SMS 2FA code."""
return auth.tfa_resend_sms(self.client, **kwargs)
def tfa_send_sms(self, **kwargs) -> dict:
"""Request a 2FA SMS verification code."""
return auth.tfa_send_sms(self.client, **kwargs)
def tfa_status(self) -> dict:
"""Show the current 2FA status and configured methods."""
return auth.tfa_status(self.client)
def tfa_totp_setup(self) -> dict:
"""Generate TOTP secret and QR code for Authenticator app setup."""
return auth.tfa_totp_setup(self.client)
def tfa_update(self, **kwargs) -> dict:
"""Update a 2FA method's settings."""
return auth.tfa_update(self.client, **kwargs)
# ------------------------------------------------------------------
# Additional billing methods
# ------------------------------------------------------------------
def generate_pdf_invoices(self, **kwargs):
"""Generate PDF invoices based on filters."""
raise NotImplementedError("generate_pdf_invoices is not yet implemented")
def transfer_credit(self, recipient: str, amount: float) -> dict:
"""Transfer credit to another account."""
return teams.transfer_credit(self.client, recipient=recipient, amount=amount)
# ------------------------------------------------------------------
# Additional methods
# ------------------------------------------------------------------
def list_volumes(self, ids, **kwargs) -> dict:
"""List disk space for rent as volumes on multiple machines."""
return storage.list_volumes(self.client, ids=ids, **kwargs)
def self_test_machine(self, machine_id, **kwargs):
"""Perform a self-test on the specified machine."""
raise NotImplementedError("self_test_machine requires CLI")
# ------------------------------------------------------------------
# Backward-compatible aliases
# ------------------------------------------------------------------
invite_team_member = invite_member
remove_team_member = remove_member
show_team_members = show_members
serverless/__init__.py
from .client import (
Serverless,
CoroutineServerless,
ServerlessRequest,
RequestStatus,
Endpoint,
Endpoint_,
ManagedEndpoint,
ManagedDeployment,
Session,
)
serverless/client/__init__.py
from .client import Serverless, CoroutineServerless, ServerlessRequest, _ServerlessBase
from .request_status import RequestStatus
from .endpoint import Endpoint, Endpoint_
from .managed import ManagedEndpoint, ManagedDeployment
from .session import Session
from .worker import Worker as ServerlessWorker
serverless/client/client.py
# client.py
from .connection import _make_request
from .endpoint import Endpoint, Endpoint_
from .managed import ManagedEndpoint, ManagedDeployment
from .worker import Worker
from .session import Session
from .request_status import RequestStatus
from vastai.data.endpoint import EndpointConfig, EndpointData
from vastai.data.deployment import (
DeploymentConfig,
DeploymentData,
DeploymentPutResponse,
)
from vastai.data.workergroup import WorkergroupConfig
import asyncio
import aiohttp
import ssl
import os
import tempfile
import logging
import random
import time
import collections
from abc import abstractmethod
from typing import (
Any,
Awaitable,
Callable,
Coroutine,
Deque,
Dict,
Generic,
Optional,
TypeVar,
Union,
List,
)
logger = logging.getLogger("vastai")
R = TypeVar("R", bound=Awaitable)
class ServerlessRequest(asyncio.Future):
"""A fire-and-observe request to a Serverless endpoint.
Composes a RequestStatus tracker for status/timestamp tracking.
Properties delegate to the tracker for backward compatibility.
"""
def __init__(self, tracker: Optional[RequestStatus] = None):
super().__init__()
self.tracker = tracker if tracker is not None else RequestStatus()
@property
def status(self):
return self.tracker.status
@status.setter
def status(self, value):
self.tracker.status = value
@property
def create_time(self):
return self.tracker.create_time
@create_time.setter
def create_time(self, value):
self.tracker.create_time = value
@property
def start_time(self):
return self.tracker.start_time
@start_time.setter
def start_time(self, value):
self.tracker.start_time = value
@property
def complete_time(self):
return self.tracker.complete_time
@complete_time.setter
def complete_time(self, value):
self.tracker.complete_time = value
@property
def req_idx(self):
return self.tracker.req_idx
@req_idx.setter
def req_idx(self, value):
self.tracker.req_idx = value
def then(self, callback) -> "ServerlessRequest":
def _done(fut):
if fut.exception() is not None:
print(fut.exception())
return
callback(fut.result())
self.add_done_callback(_done)
return self
class _ServerlessBase(Generic[R]):
SSL_CERT_URL = "https://console.vast.ai/static/jvastai_root.cer"
VAST_WEB_URL = "https://console.vast.ai"
VAST_SERVERLESS_URL = "https://run.vast.ai"
def __init__(
self,
api_key: Optional[str] = os.environ.get("VAST_API_KEY", None),
*,
debug: bool = False,
instance: str = "prod",
autoscaler_url: Optional[str] = None,
webserver_url: str = VAST_WEB_URL,
connection_limit: int = 500,
default_request_timeout: float = 600.0,
max_poll_interval: float = 5.0,
):
if api_key is None or api_key == "":
raise AttributeError(
"API key missing. Please set VAST_API_KEY in your environment variables."
)
self.api_key = api_key
self.vast_web_url = webserver_url
if autoscaler_url:
self.autoscaler_url = autoscaler_url
else:
match instance:
case "prod":
self.autoscaler_url = "https://run.vast.ai"
case "alpha":
self.autoscaler_url = "https://run-alpha.vast.ai"
case "candidate":
self.autoscaler_url = "https://run-candidate.vast.ai"
case "local":
self.autoscaler_url = "http://localhost:8080"
case _:
self.autoscaler_url = "https://run.vast.ai"
self.latencies = collections.deque(maxlen=50)
self.debug = debug
self.default_request_timeout = float(default_request_timeout)
self.max_poll_interval = float(max_poll_interval)
self.logger = logging.getLogger(self.__class__.__name__)
if not self.logger.handlers:
self.logger.addHandler(logging.NullHandler())
if debug:
handler = logging.StreamHandler()
formatter = logging.Formatter(
"[%(asctime)s] %(name)s - %(levelname)s - %(message)s"
)
handler.setFormatter(formatter)
handler.setLevel(logging.DEBUG)
self.logger.addHandler(handler)
self.logger.setLevel(logging.DEBUG)
# If we attach our own handler, avoid double-logging via root handlers.
self.logger.propagate = False
else:
# Let the application decide if/where logs go.
self.logger.propagate = True
self.connection_limit = connection_limit
self._session: aiohttp.ClientSession | None = None
self._ssl_context: ssl.SSLContext | None = None
self._transport_owner = (
None # if set, delegate _get_session/get_ssl_context and skip close
)
async def __aenter__(self):
await self._get_session()
return self
async def __aexit__(self, exc_type, exc, tb):
await self.close()
async def _get_session(self) -> aiohttp.ClientSession:
if self._transport_owner is not None:
return await self._transport_owner._get_session()
if self._session is None or self._session.closed:
self.logger.info("Started aiohttp ClientSession")
connector = aiohttp.TCPConnector(
limit=self.connection_limit, ssl=await self.get_ssl_context()
)
self._session = aiohttp.ClientSession(connector=connector)
return self._session
def is_open(self):
if self._transport_owner is not None:
return self._transport_owner.is_open()
return self._session is not None and not self._session.closed
async def close(self):
if self._transport_owner is not None:
return # transport owned by another client
if self._session and not self._session.closed:
await self._session.close()
self.logger.info("Closed aiohttp ClientSession")
async def get_ssl_context(self) -> ssl.SSLContext:
"""Download Vast.ai root cert and build SSL context (cached)."""
if self._transport_owner is not None:
return await self._transport_owner.get_ssl_context()
if self._ssl_context is None:
async with aiohttp.ClientSession() as s:
async with s.get(self.SSL_CERT_URL) as resp:
if resp.status != 200:
raise Exception(f"Failed to fetch SSL cert: {resp.status}")
cert_bytes = await resp.read()
tmpfile = tempfile.NamedTemporaryFile(delete=False, suffix=".cer")
tmpfile.write(cert_bytes)
tmpfile.close()
ctx = ssl.create_default_context()
ctx.load_verify_locations(cafile=tmpfile.name)
# The Vast.ai root CA cert has pathlen set without keyCertSign in
# Key Usage. OpenSSL 3.x (Python ≥3.10) rejects this by default.
# Clearing VERIFY_X509_STRICT relaxes that single check while
# keeping full chain-of-trust and signature verification intact.
ctx.verify_flags &= ~ssl.VERIFY_X509_STRICT
self.logger.info("Loaded Vast.ai SSL certificate")
self._ssl_context = ctx
os.unlink(tmpfile.name)
return self._ssl_context
async def get_endpoint(self, name="") -> Endpoint_[R]:
endpoints = await self.get_endpoints()
for e in endpoints:
if e.name == name:
return e
raise Exception(f"Endpoint {name} could not be found")
async def get_endpoints(self) -> list[Endpoint_[R]]:
try:
result = await _make_request(
client=self,
url=self.vast_web_url,
route="/api/v0/endptjobs/",
api_key=self.api_key,
params={"client_id": "me"},
)
except Exception as ex:
raise Exception(f"Failed to get endpoints:\nReason={ex}")
if not result.get("ok"):
raise Exception(
f"Failed to get endpoints: HTTP {result.get('status')} - {result.get('text', '')[:512]}"
)
response = result.get("json") or {}
endpoints = []
for e in response.get("results", []):
data = EndpointData.from_dict(e)
endpoints.append(Endpoint(client=self, data=data))
self.logger.info(f"Found {len(endpoints)} endpoints")
return endpoints
async def fetch_endpoint(self, endpoint_id: int) -> EndpointData:
"""Fetch full endpoint data by ID. Used internally by ManagedEndpoint.get()."""
try:
result = await _make_request(
client=self,
url=self.vast_web_url,
route=f"/api/v0/endptjobs/{endpoint_id}/",
api_key=self.api_key,
)
except Exception as ex:
raise Exception(f"Failed to fetch endpoint {endpoint_id}:\nReason={ex}")
if not result.get("ok"):
raise Exception(
f"Failed to fetch endpoint {endpoint_id}: HTTP {result.get('status')} - {result.get('text', '')[:512]}"
)
return EndpointData.from_dict(result["json"]["result"])
async def fetch_deployment(self, deployment_id: int) -> DeploymentData:
"""Fetch full deployment data by ID. Used internally by ManagedDeployment.get()."""
try:
result = await _make_request(
client=self,
url=self.vast_web_url,
route=f"/api/v0/deployment/{deployment_id}/",
api_key=self.api_key,
)
except Exception as ex:
raise Exception(f"Failed to fetch deployment {deployment_id}:\nReason={ex}")
if not result.get("ok"):
raise Exception(
f"Failed to fetch deployment {deployment_id}: HTTP {result.get('status')} - {result.get('text', '')[:512]}"
)
return DeploymentData.from_dict(result["json"]["deployment"])
async def create_endpoint(self, config: EndpointConfig) -> ManagedEndpoint[R]:
"""Create a new endpoint job. Returns a ManagedEndpoint (data fetched lazily)."""
try:
result = await _make_request(
client=self,
url=self.vast_web_url,
route="/api/v0/endptjobs/",
api_key=self.api_key,
body=config.to_dict(),
method="POST",
)
except Exception as ex:
raise Exception(f"Failed to create endpoint:\nReason={ex}")
if not result.get("ok"):
raise Exception(
f"Failed to create endpoint: HTTP {result.get('status')} - {result.get('text', '')[:512]}"
)
endpoint_id = result["json"]["result"]
self.logger.info(f"Created endpoint {endpoint_id}")
return ManagedEndpoint(id=endpoint_id, client=self)
async def put_deployment(self, config: DeploymentConfig) -> ManagedDeployment[R]:
"""Create or update a deployment. Returns a ManagedDeployment (data fetched lazily)."""
logger.debug(f"putting deployment: \n {config.to_dict()}")
try:
result = await _make_request(
client=self,
url=self.vast_web_url,
route="/api/v0/deployments/",
api_key=self.api_key,
body=config.to_dict(),
method="PUT",
)
except Exception as ex:
raise Exception(f"Failed to put deployment:\nReason={ex}")
if not result.get("ok"):
raise Exception(
f"Failed to put deployment: HTTP {result.get('status')} - {result.get('text', '')[:512]}"
)
resp = DeploymentPutResponse.from_dict(result["json"])
self.logger.info(
f"Deployment {resp.deployment_id} ({resp.action}), endpoint {resp.endpoint_id}"
)
return ManagedDeployment(
id=resp.deployment_id,
endpoint_id=resp.endpoint_id,
client=self,
put_response=resp,
)
async def create_workergroup(self, config: WorkergroupConfig) -> int:
"""Create a workergroup (autoscale job). Returns the workergroup ID."""
try:
result = await _make_request(
client=self,
url=self.vast_web_url,
route="/api/v0/workergroups/",
api_key=self.api_key,
body=config.to_dict(),
method="POST",
)
except Exception as ex:
raise Exception(f"Failed to create workergroup:\nReason={ex}")
if not result.get("ok"):
raise Exception(
f"Failed to create workergroup: HTTP {result.get('status')} - {result.get('text', '')[:512]}"
)
wg_id = result["json"]["id"]
self.logger.info(f"Created workergroup {wg_id}")
return wg_id
async def delete_workergroup(self, workergroup_id: int) -> None:
"""Delete a workergroup (autoscale job) by ID."""
try:
result = await _make_request(
client=self,
url=self.vast_web_url,
route=f"/api/v0/workergroups/{workergroup_id}/",
api_key=self.api_key,
method="DELETE",
)
except Exception as ex:
raise Exception(
f"Failed to delete workergroup {workergroup_id}:\nReason={ex}"
)
if not result.get("ok"):
raise Exception(
f"Failed to delete workergroup {workergroup_id}: HTTP {result.get('status')} - {result.get('text', '')[:512]}"
)
self.logger.info(f"Deleted workergroup {workergroup_id}")
async def update_workers(self, workergroup_id: int) -> dict:
"""Trigger a rolling update for all workers in a workergroup via the autoscaler."""
try:
result = await _make_request(
client=self,
url=self.autoscaler_url,
route="/update_workers/",
api_key=self.api_key,
body={"workergroup_id": workergroup_id, "api_key": self.api_key},
method="POST",
)
except Exception as ex:
raise Exception(
f"Failed to trigger worker update for workergroup {workergroup_id}:\nReason={ex}"
)
if not result.get("ok"):
raise Exception(
f"Failed to trigger worker update for workergroup {workergroup_id}: HTTP {result.get('status')} - {result.get('text', '')[:512]}"
)
self.logger.info(f"Triggered rolling update for workergroup {workergroup_id}")
return result.get("json", {})
async def find_workergroup_for_endpoint(self, endpoint_id: int) -> Optional[int]:
"""Fetch all workergroups and return the first one matching endpoint_id, or None."""
try:
result = await _make_request(
client=self,
url=self.vast_web_url,
route="/api/v0/workergroups/",
api_key=self.api_key,
method="GET",
)
except Exception:
return None
if not result.get("ok"):
return None
for wg in result.get("json", {}).get("results", []):
if wg.get("endpoint_id") == endpoint_id:
return wg["id"]
return None
async def delete_endpoint(self, endpoint_id: int) -> None:
"""Delete an endpoint job by ID."""
try:
result = await _make_request(
client=self,
url=self.vast_web_url,
route=f"/api/v0/endptjobs/{endpoint_id}/",
api_key=self.api_key,
method="DELETE",
)
except Exception as ex:
raise Exception(f"Failed to delete endpoint {endpoint_id}:\nReason={ex}")
if not result.get("ok"):
raise Exception(
f"Failed to delete endpoint {endpoint_id}: HTTP {result.get('status')} - {result.get('text', '')[:512]}"
)
self.logger.info(f"Deleted endpoint {endpoint_id}")
async def get_endpoint_workers(self, endpoint: Endpoint_[R]) -> List[Worker]:
if not isinstance(endpoint, Endpoint_):
raise ValueError("endpoint must be an Endpoint")
url = f"{self.autoscaler_url}/get_endpoint_workers/"
payload = {"id": endpoint.id, "api_key": self.api_key}
async with self._session.post(
url, json=payload, timeout=aiohttp.ClientTimeout(total=30)
) as resp:
if resp.status != 200:
text = await resp.text()
raise RuntimeError(
f"get_endpoint_workers failed: HTTP {resp.status} - {text}"
)
data = await resp.json(content_type=None)
# If error message from authenticate_endpoint_apikey_by_id occurs, there is a possibility that
# the endpoint's worker instances are not ready to be queried. If an error message occurs,
# return an empty list and print the error message to the user. The endpoint get_endpoint_workers
# should normally return a list of dictionaries containing worker instance information.
if isinstance(data, dict):
if "error_msg" in data.keys():
self.logger.warning(
f"Received the following error from get_endpoint_workers:{data['error_msg']}.\nEndpoint may not be ready for query. Check credentials or wait a few minutes and try again."
)
return []
if not isinstance(data, list):
raise RuntimeError(
f"Unexpected response type (wanted list): {type(data)}"
)
return [Worker.from_dict(item) for item in data]
async def get_endpoint_session(
self, endpoint, session_id: int, session_auth: str, timeout: float = 10.0
):
try:
result = await _make_request(
client=self,
url=session_auth.get("url"),
api_key="",
route="/session/get",
body={"session_id": session_id, "session_auth": session_auth},
method="POST",
retries=1,
timeout=timeout,
stream=False,
)
if not result.get("ok"):
error_msg = f"Error on /session/get: {result.get('json', {}).get('error', result.get('text', 'Unknown error'))}"
raise Exception(error_msg)
worker_response = result.get("json") or {}
auth_data = worker_response.get("auth_data")
if auth_data is None:
raise Exception("Missing auth_data in response")
session = Session(
endpoint=endpoint,
session_id=session_id,
lifetime=worker_response.get("lifetime"),
expiration=worker_response.get("expiration"),
auth_data=auth_data,
url=auth_data.get("url"),
)
return session
except asyncio.TimeoutError:
raise
except Exception as ex:
error_msg = f"Failed to get session {session_id}: {ex}"
self.logger.error(error_msg)
raise Exception(error_msg)
async def end_endpoint_session(self, session: Session, timeout: float = 10.0):
try:
self.logger.debug(
f"Attempting to end session {session.session_id} at {session.url}"
)
result = await _make_request(
client=self,
url=session.url,
api_key="",
route="/session/end",
body={
"session_id": session.session_id,
"session_auth": session.auth_data,
},
method="POST",
retries=1,
timeout=timeout,
stream=False,
)
if not result.get("ok"):
error_msg = f"Error on /session/end: {result.get('json', {}).get('error', result.get('text', 'Unknown error'))}"
raise Exception(error_msg)
self.logger.debug(f"Successfully ended session {session.session_id}")
return
except asyncio.TimeoutError:
raise
except Exception as ex:
error_msg = f"Failed to end session {session.session_id}: {ex}"
self.logger.error(error_msg)
raise Exception(error_msg)
async def start_endpoint_session(
self,
endpoint: Endpoint_[R],
cost: int = 100,
lifetime: float = 60,
on_close_route: str = None,
on_close_payload: dict = None,
timeout: float = None,
) -> Session:
try:
session_start_response = await self.queue_endpoint_request(
endpoint=endpoint,
worker_route="/session/create",
worker_payload={
"lifetime": lifetime,
"on_close_route": on_close_route,
"on_close_payload": on_close_payload,
},
cost=cost,
timeout=timeout,
worker_timeout=10.0,
)
if not session_start_response.get("ok"):
error_msg = f"Error on /session/create: {session_start_response.get('json', {}).get('error', session_start_response.get('text', 'Unknown error'))}"
raise Exception(error_msg)
response = session_start_response.get("response")
if response is None:
raise Exception("No response from /session/create")
if not isinstance(response, dict):
raise Exception(
"Invalid response from /session/create: expected mapping"
)
session_id = response.get("session_id")
if session_id is None:
raise Exception("Missing session id")
expiration = response.get("expiration")
url = session_start_response.get("url")
if url is None:
raise Exception("Missing URL")
auth_data = session_start_response.get("auth_data")
if auth_data is None:
raise Exception("Missing auth data")
return Session(
endpoint=endpoint,
session_id=session_id,
lifetime=lifetime,
expiration=expiration,
url=url,
auth_data=auth_data,
)
except asyncio.TimeoutError:
raise
except Exception as ex:
error_msg = f"Failed to create session: {ex}"
self.logger.error(error_msg)
raise Exception(error_msg)
async def _do_request(
self,
endpoint: Endpoint_[R],
worker_route: str,
worker_payload: dict,
session: Optional[Session] = None,
tracker: Optional[RequestStatus] = None,
cost: int = 100,
timeout: Optional[float] = None,
worker_timeout: Optional[float] = 600,
retry: bool = True,
max_retries: Optional[int] = None,
stream: bool = False,
) -> dict:
"""Core request logic: route to a worker, execute, retry on failure. Returns the response dict."""
if tracker is None:
tracker = RequestStatus()
request_idx: int = 0
total_attempts = 0
start_time = time.time()
try:
while True:
total_attempts += 1
tracker.status = "Queued"
worker_url = ""
auth_data = {}
session_id = None
# Check total elapsed time
if timeout is not None and (time.time() - start_time) >= timeout:
raise asyncio.TimeoutError(
f"Timed out after {time.time() - start_time:.1f}s waiting for worker"
)
if session is None:
if request_idx == 0:
self.logger.debug(
f"Sending initial route call for request_idx {request_idx}"
)
else:
self.logger.debug(
f"Sending retry route call for request_idx {request_idx}"
)
route = await endpoint._route(
cost=cost, req_idx=request_idx, timeout=60.0
)
request_idx = route.request_idx
if request_idx:
self.logger.debug(f"Got request index {request_idx}")
else:
self.logger.info( # usually non-fatal due to get_serverless_groups polling delay; we should fix that and make this an error again.
"Did not get request_idx from initial route"
)
poll_interval = 1
poll_elapsed = 0
attempt = 0
while route.status != "READY":
tracker.status = "Polling"
# Check total elapsed time
if (
timeout is not None
and (time.time() - start_time) >= timeout
):
raise asyncio.TimeoutError(
f"Timed out after {time.time() - start_time:.1f}s waiting for worker to become ready"
)
await asyncio.sleep(poll_interval)
poll_elapsed += poll_interval
route = await endpoint._route(
cost=cost, req_idx=request_idx, timeout=60.0
)
request_idx = route.request_idx or request_idx
attempt += 1
poll_interval = random.uniform(
0.1,
min(
2 ** min(attempt, 20) + random.uniform(0, 1),
self.max_poll_interval,
),
)
self.logger.debug(f"Polling route, attempt {attempt}")
worker_url = route.get_url()
auth_data = route.body
else:
if session.url is not None:
worker_url = session.url
auth_data = session.auth_data
session_id = session.session_id
payload = worker_payload
worker_request_body = {
"auth_data": auth_data,
"session_id": session_id,
"payload": payload,
}
self.logger.debug("Found worker machine, starting work")
if tracker.status != "Retrying":
tracker.status = "In Progress"
tracker.start_time = time.time()
# Transport/JSON failures may raise; HTTP errors return ok=False.
try:
result = await _make_request(
client=self,
url=worker_url,
route=worker_route,
api_key=endpoint.api_key,
body=worker_request_body,
method="POST",
retries=1, # avoid stacking retries with the outer loop
timeout=worker_timeout,
stream=stream,
allow_non_json=True,
)
except (
aiohttp.ClientConnectorError,
aiohttp.ServerDisconnectedError,
) as ex:
# Worker is gone
if session is not None:
# Session is bound to this worker - can't re-route
self.logger.error(
f"Session worker unavailable ({type(ex).__name__})"
)
session.open = False
raise ConnectionError(
f"Session worker unavailable: {ex}"
) from ex
# No session - reset request_idx to force a fresh route
self.logger.warning(
f"Worker unavailable ({type(ex).__name__}), re-routing to new worker"
)
tracker.status = "Retrying"
continue
except Exception as ex:
self.logger.error(f"Worker request failed: {ex}")
tracker.status = "Retrying"
continue
if not result.get("ok"):
if (
retry
and result.get("retryable")
and (max_retries is None or total_attempts < max_retries)
):
# Check if we have time left before retrying
if (
timeout is not None
and (time.time() - start_time) >= timeout
):
raise asyncio.TimeoutError(
f"Request timed out after {time.time() - start_time:.1f}s"
)
tracker.status = "Retrying"
await asyncio.sleep(
min(
(2 ** min(total_attempts, 20)) + random.uniform(0, 1),
self.max_poll_interval,
)
)
continue
# Return the raw HTTP result to the caller
tracker.status = "Complete"
tracker.complete_time = time.time()
if tracker.start_time is not None:
self.latencies.append(
tracker.complete_time - tracker.start_time
)
return {
"response": result.get("json")
if result.get("json") is not None
else {"error": result.get("text", "")},
"ok": result.get("ok"),
"status": result.get("status"),
"text": result.get("text"),
"latency": (tracker.complete_time - tracker.start_time)
if tracker.start_time
else None,
"url": worker_url,
"request_idx": request_idx,
"auth_data": auth_data,
}
# Success
if stream:
worker_response = result.get("stream")
elif result.get("content") is not None:
worker_response = result.get("content")
else:
worker_response = result.get("json")
tracker.status = "Complete"
tracker.complete_time = time.time()
self.latencies.append(tracker.complete_time - tracker.start_time)
self.logger.info("Endpoint request task completed")
return {
"response": worker_response,
"ok": result.get("ok"),
"status": result.get("status"),
"text": result.get("text"),
"content_type": result.get("content_type"),
"latency": tracker.complete_time - tracker.start_time,
"url": worker_url,
"request_idx": request_idx,
"auth_data": auth_data,
}
except asyncio.CancelledError:
tracker.status = "Cancelled"
raise
except Exception as ex:
tracker.status = "Errored"
self.logger.error(f"Request errored: {ex}")
raise
@abstractmethod
def queue_endpoint_request(
self,
endpoint: Endpoint_[R],
worker_route: str,
worker_payload: dict,
session: Optional[Session] = None,
serverless_request: Optional[Union[ServerlessRequest, RequestStatus]] = None,
cost: int = 100,
timeout: Optional[float] = None,
worker_timeout: Optional[float] = 600,
retry: bool = True,
max_retries: Optional[int] = None,
stream: bool = False,
) -> R:
"""Dispatch a request to a serverless endpoint. Return type depends on subclass."""
...
class Serverless(_ServerlessBase[ServerlessRequest]):
"""Fire-and-observe client: queue_endpoint_request returns a ServerlessRequest (Future)."""
def queue_endpoint_request(
self,
endpoint: Endpoint,
worker_route: str,
worker_payload: dict,
session: Optional[Session] = None,
serverless_request: Optional[Union[ServerlessRequest, RequestStatus]] = None,
cost: int = 100,
timeout: Optional[float] = None,
worker_timeout: Optional[float] = 600,
retry: bool = True,
max_retries: Optional[int] = None,
stream: bool = False,
) -> ServerlessRequest:
"""Return a Future that will resolve once the request completes."""
if serverless_request is None:
serverless_request = ServerlessRequest()
elif isinstance(serverless_request, RequestStatus):
serverless_request = ServerlessRequest(tracker=serverless_request)
async def task(request: ServerlessRequest):
try:
result = await self._do_request(
endpoint=endpoint,
worker_route=worker_route,
worker_payload=worker_payload,
session=session,
tracker=request.tracker,
cost=cost,
timeout=timeout,
worker_timeout=worker_timeout,
retry=retry,
max_retries=max_retries,
stream=stream,
)
request.set_result(result)
except asyncio.CancelledError:
pass # tracker already set by _do_request
except Exception as ex:
request.set_exception(ex)
bg_task = asyncio.create_task(task(serverless_request))
def _propagate_cancel(fut: ServerlessRequest):
if fut.cancelled():
bg_task.cancel()
serverless_request.add_done_callback(_propagate_cancel)
self.logger.info("Queued endpoint request")
return serverless_request
class CoroutineServerless(_ServerlessBase[Coroutine[Any, Any, dict]]):
"""Async/await client: queue_endpoint_request returns a coroutine."""
def queue_endpoint_request(
self,
endpoint: Endpoint,
worker_route: str,
worker_payload: dict,
session: Optional[Session] = None,
serverless_request: Optional[Union[ServerlessRequest, RequestStatus]] = None,
cost: int = 100,
timeout: Optional[float] = None,
worker_timeout: Optional[float] = 600,
retry: bool = True,
max_retries: Optional[int] = None,
stream: bool = False,
) -> Coroutine[Any, Any, dict]:
"""Return a coroutine that resolves to the response dict when awaited."""
if isinstance(serverless_request, ServerlessRequest):
tracker = serverless_request.tracker
elif isinstance(serverless_request, RequestStatus):
tracker = serverless_request
else:
tracker = None
return self._do_request(
endpoint=endpoint,
worker_route=worker_route,
worker_payload=worker_payload,
session=session,
tracker=tracker,
cost=cost,
timeout=timeout,
worker_timeout=worker_timeout,
retry=retry,
max_retries=max_retries,
stream=stream,
)
serverless/client/connection.py
# connection.py
import aiohttp
import asyncio
import random
import json
from typing import AsyncIterator, Dict, Optional, Any
from vastai.utils import VERSION
_JITTER_CAP_SECONDS = 5.0
def _retryable(status: int) -> bool:
return status in (408, 429) or (500 <= status < 600)
def _backoff_delay(attempt: int) -> float:
# capped exponential backoff with jitter
return min((2 ** attempt) + random.uniform(0, 1), _JITTER_CAP_SECONDS)
def _build_kwargs(
*,
headers: Dict[str, str],
params: Dict[str, str],
ssl_context,
timeout: Optional[float],
body: Optional[dict],
method: str,
stream: bool,
) -> Dict:
return {
"headers": headers,
"params": params,
"ssl": ssl_context,
"timeout": aiohttp.ClientTimeout(total=None if stream else timeout),
**({"json": body} if method != "GET" and body else {}),
}
async def _iter_sse_json(resp: aiohttp.ClientResponse) -> AsyncIterator[dict]:
"""
Yield JSON objects from an SSE/text stream. Accepts lines starting with 'data:' or raw JSONL.
"""
buffer = b""
async for chunk in resp.content.iter_any():
if not chunk:
continue
buffer += chunk
while b"\n" in buffer:
line, buffer = buffer.split(b"\n", 1)
line = line.strip()
if not line:
continue
if line.startswith(b"data:"):
line = line[5:].strip()
try:
yield json.loads(line.decode("utf-8"))
except Exception:
# Ignore keepalives/bad fragments
continue
# flush tail if present
tail = buffer.strip()
if tail:
try:
yield json.loads(tail.decode("utf-8"))
except Exception:
pass
async def _open_once(
*,
session: aiohttp.ClientSession,
method: str,
url: str,
route: str,
kwargs: Dict,
):
"""
Execute one HTTP attempt and return the aiohttp response object.
Caller is responsible for reading/closing via 'async with' or resp.release()/resp.close().
"""
request_fn = {"GET": session.get, "POST": session.post, "PUT": session.put, "DELETE": session.delete}.get(method, session.post)
return await request_fn(url + route, **kwargs)
async def _make_request(
client,
route: str,
api_key: str,
url: str = "",
body: Optional[dict] = None,
params: Optional[dict] = None,
method: str = "GET",
retries: int = 5,
timeout: float = 30,
stream: bool = False,
allow_non_json: bool = False,
) -> Dict[str, Any]:
"""
Make an HTTP request with capped exponential backoff + jitter, returning a structured result.
- Never raises for HTTP non-2xx responses. Instead returns result with ok=False and status/text/json.
- Raises only for "mechanical" failures (aiohttp/transport) and invalid JSON on successful (2xx) responses.
- allow_non_json: a 2xx text/* or media body that is not a JSON object or array is returned
(str for text/*, bytes otherwise) instead of raising.
Return shape (non-stream):
{
"ok": bool,
"status": int|None,
"url": str,
"headers": dict,
"text": str,
"json": Any|None,
"content": bytes|str|None, # allow_non_json only: str for text/*, else bytes
"content_type": str, # the declared type, without parameters ("" if none)
"retryable": bool,
"attempt": int
}
Return shape (stream=True, ok=True adds "stream"):
{ ... , "stream": AsyncIterator[dict] }
"""
method = method.upper()
body = body or {}
params = {**(params or {})}
params["api_key"] = api_key
headers = {"Authorization": f"Bearer {api_key}", "User-Agent": f"vastai-sdk/{VERSION}"}
session = await client._get_session()
ssl_context = await client.get_ssl_context() if client else None
full_url = url + route
last_result: Dict[str, Any] = {
"ok": False,
"status": None,
"url": full_url,
"headers": {},
"text": "",
"json": None,
"retryable": False,
"attempt": 0,
}
# STREAMING PATH (open upfront so we can return status + a stream iterator)
if stream:
for attempt in range(1, retries + 1):
last_result["attempt"] = attempt
try:
kwargs = _build_kwargs(
headers=headers,
params=params,
ssl_context=ssl_context,
timeout=timeout,
body=body,
method=method,
stream=True,
)
resp = await _open_once(
session=session,
method=method,
url=url,
route=route,
kwargs=kwargs,
)
status = resp.status
last_result["status"] = status
last_result["headers"] = dict(resp.headers)
last_result["retryable"] = _retryable(status)
if status < 200 or status >= 300:
text = await resp.text()
last_result["text"] = text
# best-effort JSON parse (do not raise)
try:
if text and text.lstrip().startswith(("{", "[")):
last_result["json"] = json.loads(text)
except Exception:
pass
resp.release()
if last_result["retryable"] and attempt < retries:
await asyncio.sleep(_backoff_delay(attempt))
continue
# final non-2xx: return what we saw
return last_result
# 2xx: return stream iterator that owns the response lifetime
last_result["ok"] = True
last_result["text"] = "" # streaming: we don't pre-read body
async def _stream_iter() -> AsyncIterator[dict]:
try:
async for obj in _iter_sse_json(resp):
yield obj
finally:
try:
resp.release()
except Exception:
pass
last_result["stream"] = _stream_iter()
return last_result
except asyncio.TimeoutError as ex:
if attempt == retries:
raise TimeoutError(f"Request to {full_url} timed out after {timeout}s") from ex
await asyncio.sleep(_backoff_delay(attempt))
except Exception as ex:
if attempt == retries:
raise ex
await asyncio.sleep(_backoff_delay(attempt))
return last_result
# NON-STREAMING PATH
for attempt in range(1, retries + 1):
last_result["attempt"] = attempt
try:
kwargs = _build_kwargs(
headers=headers,
params=params,
ssl_context=ssl_context,
timeout=timeout,
body=body,
method=method,
stream=False,
)
request_fn = {"GET": session.get, "POST": session.post, "PUT": session.put, "DELETE": session.delete}.get(method, session.post)
async with await request_fn(full_url, **kwargs) as resp:
status = resp.status
content_type = (resp.headers.get("Content-Type") or "").split(";", 1)[0].lower()
is_text = content_type.startswith("text/")
is_media = (allow_non_json and bool(content_type) and not is_text
and "json" not in content_type)
raw = await resp.read() if is_media else None
text = "" if is_media and 200 <= status < 300 else await resp.text(errors="replace")
result: Dict[str, Any] = {
"ok": 200 <= status < 300,
"status": status,
"url": full_url,
"headers": dict(resp.headers),
"text": text,
"json": None,
"content": None,
"content_type": content_type,
"retryable": _retryable(status),
"attempt": attempt,
}
if result["ok"]:
if is_media or (allow_non_json and is_text):
# The body, unless it is an object or array (mislabelled JSON).
body = raw if is_media else text
try:
if body.lstrip()[:1] in ("{", "[", b"{", b"["):
result["json"] = json.loads(body)
return result
except ValueError:
pass
result["content"] = body
return result
# Successful responses are expected to be JSON; invalid JSON is a hard failure
try:
result["json"] = await resp.json(content_type=None)
except Exception:
raise Exception(f"Invalid JSON from {full_url}:\n{text}")
# Debug: log the exact response
if hasattr(client, 'logger'):
client.logger.debug(f"_make_request response from {route}: {result}")
return result
# Non-2xx: best-effort JSON parse (do not raise)
try:
if text and text.lstrip().startswith(("{", "[")):
result["json"] = json.loads(text)
except Exception:
pass
# Retry only if retryable and attempts remain
if result["retryable"] and attempt < retries:
last_result = result
await asyncio.sleep(_backoff_delay(attempt))
continue
# Debug: log the exact response for non-2xx
if hasattr(client, 'logger'):
client.logger.debug(f"_make_request response from {route}: {result}")
return result
except asyncio.TimeoutError as ex:
if attempt == retries:
raise TimeoutError(f"Request to {full_url} timed out after {timeout}s") from ex
await asyncio.sleep(_backoff_delay(attempt))
except Exception as ex:
if attempt == retries:
raise ex
await asyncio.sleep(_backoff_delay(attempt))
return last_result
serverless/client/endpoint.py
# endpoint.py
import asyncio
from dataclasses import dataclass
import os
import time
from .connection import _make_request
from typing import Awaitable, Generic, Optional, TypeVar, Union, TYPE_CHECKING
import logging
from vastai.data.endpoint import EndpointData
logger = logging.getLogger("vastai")
if TYPE_CHECKING:
from .client import _ServerlessBase, ServerlessRequest
from .request_status import RequestStatus
from .session import Session
R = TypeVar("R", bound=Awaitable)
@dataclass
class EndpointDataPartial:
name: str
id: int
api_key: str
class Endpoint_(Generic[R]):
client: "_ServerlessBase[R]"
def __repr__(self):
return f"<Endpoint {self.name} (id={self.data.id})>"
def __init__(
self,
client: "_ServerlessBase[R]",
name: Optional[str] = None,
id: Optional[int] = None,
api_key: Optional[str] = None,
data: Optional[EndpointData] = None,
soft_refresh_threshold=3 * 24 * 3600,
hard_refresh_threshold=6 * 24 * 3600,
):
if client is None:
raise ValueError("client is required")
self.client = client
if isinstance(data, EndpointData):
self.data: "EndpointData | EndpointDataPartial" = data
elif isinstance(name, str) and isinstance(id, int) and isinstance(api_key, str):
self.data: "EndpointData | EndpointDataPartial" = EndpointDataPartial(
name, id, api_key
)
else:
raise ValueError(
"Either data or all of name, id, and api_key are required!"
)
self.refresh_task: Optional[asyncio.Task[EndpointData]] = None
self.last_refresh = time.time()
self.soft_refresh_threshold = soft_refresh_threshold
self.hard_refresh_threshold = hard_refresh_threshold
@property
def name(self):
return (
self.data.config.endpoint_name
if isinstance(self.data, EndpointData)
else self.data.name
)
@property
def id(self):
return self.data.id
@property
def api_key(self):
return self.data.api_key
async def refresh(self, block=False):
logger.debug("refresh")
if self.refresh_task is not None:
if block:
try:
self.data = await self.refresh_task
finally:
self.refresh_task = None
else:
try:
self.data = self.refresh_task.result()
self.refresh_task = None
except asyncio.InvalidStateError:
pass
except Exception as e: # refresh task errored, reset it.
self.refresh_task = None
raise e
else:
if block:
self.data = await self.client.fetch_endpoint(self.data.id)
else:
self.refresh_task = asyncio.create_task(
self.client.fetch_endpoint(self.data.id)
)
def request(
self,
route,
payload,
serverless_request: Optional[
Union["ServerlessRequest", "RequestStatus"]
] = None,
cost: int = 100,
retry: bool = True,
stream: bool = False,
timeout: Optional[float] = None,
session: Optional["Session"] = None,
) -> R:
return self.client.queue_endpoint_request(
endpoint=self,
worker_route=route,
worker_payload=payload,
serverless_request=serverless_request,
cost=cost,
retry=retry,
stream=stream,
timeout=timeout,
session=session,
)
def close_session(self, session: "Session"):
return self.client.end_endpoint_session(session=session)
async def session_healthcheck(self, session: "Session"):
result = await self.client.get_endpoint_session(
endpoint=self, session_id=session.session_id, session_auth=session.auth_data
)
return result is not None
def get_session(self, session_id: int, session_auth: dict, timeout: float = 10):
return self.client.get_endpoint_session(
endpoint=self,
session_id=session_id,
session_auth=session_auth,
timeout=timeout,
)
def session(
self,
cost: int = 100,
lifetime: float = 60,
on_close_route: str = None,
on_close_payload: dict = None,
timeout: float = None,
) -> "Session":
return self.client.start_endpoint_session(
endpoint=self,
cost=cost,
lifetime=lifetime,
on_close_route=on_close_route,
on_close_payload=on_close_payload,
timeout=timeout,
)
def get_workers(self):
return self.client.get_endpoint_workers(self)
async def _route(
self, cost: float = 0.0, req_idx: int = 0, timeout: float = 60.0
) -> "RouteResponse":
VAST_DEBUG_WORKER_URL = os.environ.get("VAST_DEBUG_WORKER_URL")
if VAST_DEBUG_WORKER_URL:
return RouteResponse({"request_idx": 1, "url": VAST_DEBUG_WORKER_URL})
await self.client._get_session()
if time.time() > self.last_refresh + self.soft_refresh_threshold:
await self.refresh(
block=time.time() > self.last_refresh + self.hard_refresh_threshold
)
try:
result = await _make_request(
client=self.client,
url=self.client.autoscaler_url,
route="/route/",
api_key=self.data.api_key,
body={
"endpoint": self.name,
"api_key": self.data.api_key,
"cost": cost,
"request_idx": req_idx,
"replay_timeout": timeout,
},
method="POST",
timeout=10.0,
retries=1,
stream=False,
)
except Exception as ex:
raise RuntimeError(f"Failed to route endpoint: {ex}") from ex
if not result.get("ok"):
raise RuntimeError(
f"Failed to route endpoint: HTTP {result.get('status')} - {result.get('text', '')[:512]}"
)
return RouteResponse(result.get("json") or {})
class RouteResponse:
status: str
body: dict
request_idx: int
def __repr__(self):
return f"<RouteResponse status={self.status}>"
def __init__(self, body: dict):
if "request_idx" in body.keys():
self.request_idx = body.get("request_idx")
else:
self.request_idx = 0
if "url" in body.keys():
self.status = "READY"
self.body = body
else:
self.status = "WAITING"
self.body = body
def get_url(self):
return self.body.get("url")
# Backward-compatible alias — works for instantiation at runtime.
# Static analysis infers the R parameter from the client argument to __init__.
Endpoint = Endpoint_
serverless/client/managed.py
from __future__ import annotations
from typing import TYPE_CHECKING, Awaitable, Generic, Optional, Union, TypeVar
from vastai.data.endpoint import EndpointData
from vastai.data.deployment import DeploymentData, DeploymentPutResponse
from vastai.data.workergroup import WorkergroupConfig
from .endpoint import Endpoint_
import requests
if TYPE_CHECKING:
from .client import _ServerlessBase
R = TypeVar("R", bound=Awaitable)
class ManagedEndpoint(Generic[R]):
"""Wraps an endpoint with lazy data fetching, caching, and request dispatch.
The .id property is always available. All other data is fetched on first
access via .get() and cached. Call .invalidate() to force a re-fetch.
"""
def __init__(
self, id: int, client: _ServerlessBase[R], data: Optional[EndpointData] = None
):
self._id = id
self._client = client
self._routing_endpoint: Optional[Endpoint_[R]] = (
Endpoint_[R](client, data=data) if isinstance(data, EndpointData) else None
)
@property
def id(self) -> int:
return self._id
def invalidate(self):
"""Clear cached data so the next .get() re-fetches from the API."""
self._data = None
self._routing_endpoint = None
async def _get_routing_endpoint(self) -> Endpoint_[R]:
if self._routing_endpoint is None:
data = await self._client.fetch_endpoint(self._id)
self._routing_endpoint = Endpoint_(client=self._client, data=data)
return self._routing_endpoint
async def request(self, route: str, payload: dict, **kwargs) -> dict:
"""Route a request to this endpoint's workers. Fetches endpoint data if needed."""
ep = await self._get_routing_endpoint()
return await self._client.queue_endpoint_request(
endpoint=ep,
worker_route=route,
worker_payload=payload,
**kwargs,
)
async def add_workergroup(
self,
config_or_template_hash: Union[WorkergroupConfig, str, None] = None,
**kwargs,
) -> int:
"""Add a workergroup to this endpoint. Returns the workergroup ID.
Accepts either:
- A WorkergroupConfig object
- A template_hash string (common case)
- Keyword arguments passed to WorkergroupConfig
"""
if isinstance(config_or_template_hash, WorkergroupConfig):
config = config_or_template_hash
if config.endpoint_id is None:
config.endpoint_id = self._id
elif isinstance(config_or_template_hash, str):
config = WorkergroupConfig(
endpoint_id=self._id, template_hash=config_or_template_hash, **kwargs
)
else:
config = WorkergroupConfig(endpoint_id=self._id, **kwargs)
return await self._client.create_workergroup(config)
async def delete(self) -> None:
"""Delete this endpoint."""
await self._client.delete_endpoint(self._id)
def __repr__(self) -> str:
status = "loaded" if self._routing_endpoint is not None else "lazy"
return f"<ManagedEndpoint id={self._id} ({status})>"
class ManagedDeployment(Generic[R]):
"""Wraps a deployment with lazy data fetching, caching, and endpoint access.
The .id and .endpoint_id properties are always available.
If the deployment was just created/updated, .needs_upload and .upload()
are available to handle S3 blob uploads.
"""
def __init__(
self,
id: int,
endpoint_id: int,
client: _ServerlessBase[R],
data: Optional[DeploymentData] = None,
put_response: Optional[DeploymentPutResponse] = None,
):
self._id = id
self._endpoint_id = endpoint_id
self._client = client
self._data = data
self._put_response = put_response
self._endpoint: Optional[ManagedEndpoint[R]] = None
@property
def id(self) -> int:
return self._id
@property
def endpoint_id(self) -> int:
return self._endpoint_id
@property
def action(self) -> Optional[str]:
"""The action taken by put_deployment (e.g. 'created', 'soft_update', 'exists')."""
return self._put_response.action if self._put_response else None
@property
def workergroup_id(self) -> Optional[int]:
"""The workergroup ID associated with this deployment, if returned by the server."""
return self._put_response.workergroup_id if self._put_response else None
@property
def needs_upload(self) -> bool:
"""True if the deployment requires a blob upload to S3 to complete setup."""
return (
self._put_response is not None
and self._put_response.upload_url is not None
and self._put_response.upload_fields is not None
)
async def upload(self, file_path: str) -> None:
"""Upload a tarball to S3 using the presigned POST from put_deployment.
Raises ValueError if no upload is needed, RuntimeError on upload failure.
"""
if not self.needs_upload:
raise ValueError("No upload needed for this deployment")
import aiohttp
import os
upload_url = self._put_response.upload_url
upload_fields = self._put_response.upload_fields
data = aiohttp.FormData()
for key, value in upload_fields.items():
data.add_field(key, value)
data.add_field(
"file",
open(file_path, "rb"),
filename=os.path.basename(file_path),
content_type="application/gzip",
)
session = await self._client._get_session()
async with session.post(upload_url, data=data) as resp:
if resp.status not in (200, 201, 204):
text = await resp.text()
raise RuntimeError(
f"S3 upload failed: HTTP {resp.status} - {text[:512]}"
)
self._put_response.upload_url = None
self._put_response.upload_fields = None
@property
def endpoint(self) -> ManagedEndpoint[R]:
"""Get the ManagedEndpoint for this deployment (lazy — data fetched on first use)."""
if self._endpoint is None:
self._endpoint = ManagedEndpoint(id=self._endpoint_id, client=self._client)
return self._endpoint
async def get(self) -> DeploymentData:
"""Fetch and cache the full deployment data. Returns cached data on subsequent calls."""
if self._data is None:
self._data = await self._client.fetch_deployment(self._id)
return self._data
def invalidate(self):
"""Clear cached data so the next .get() re-fetches from the API."""
self._data = None
def sync_heartbeat(
self,
): # for use in background thread outside of main aiohttp connection
headers = {"Authorization": f"Bearer {self._client.api_key}"}
requests.post(
url=self._client.vast_web_url.rstrip("/")
+ f"/api/v0/deployment/{self.id}/heartbeat/",
headers=headers,
)
def __repr__(self) -> str:
status = "loaded" if self._data is not None else "lazy"
upload = " needs_upload" if self.needs_upload else ""
return f"<ManagedDeployment id={self._id} endpoint_id={self._endpoint_id} ({status}{upload})>"
serverless/client/request_status.py
import time
from dataclasses import dataclass, field
from typing import Optional
@dataclass
class RequestStatus:
"""Observable status tracker for a serverless request lifecycle."""
status: str = "New"
create_time: float = field(default_factory=time.time)
start_time: Optional[float] = None
complete_time: Optional[float] = None
req_idx: int = 0
serverless/client/session.py
from .endpoint import Endpoint
from typing import Callable, Optional
import inspect
import logging
log = logging.getLogger(__name__)
class Session:
endpoint: Endpoint
session_id: str
open: bool
url: str
lifetime: float
auth_data: dict
on_close_route: str
on_close_payload: dict
def __init__(self, endpoint: Endpoint, session_id: str, lifetime: float, expiration: str, url: str, auth_data: dict, on_close_route: str = None, on_close_payload: dict = None):
if endpoint is None:
raise ValueError("Session cannot be created with empty endpoint")
if session_id is None:
raise ValueError("Session cannot be created with empty session_id")
if url is None:
raise ValueError("Session cannot be created with empty url")
self.endpoint = endpoint
self.session_id = session_id
self.lifetime = lifetime
self.expiration = expiration
self.url = url
self.auth_data = auth_data
self.open = True
self.on_close_route = on_close_route
self.on_close_payload = on_close_payload
self._closing = False # guard to ensure one-time execution
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
await self.close()
return False
async def is_open(self):
result = await self.endpoint.session_healthcheck(self)
self.open = result
return result
async def close(self):
"""
Explicit close for non-async contexts.
Returns an awaitable if async work is required.
"""
if not self.open or self._closing:
return None
self._closing = True
try:
await self.endpoint.close_session(self)
except Exception as e:
log.warning(f"Error closing session {self.session_id}: {e}")
finally:
self.open = False
return
def request(
self,
route,
payload,
serverless_request=None,
cost: int = 100,
retry: bool = True,
stream: bool = False,
):
"""Forward requests to the endpoint"""
if not self.open:
raise ValueError("Cannot make request on closed session.")
async def _wrapped_request():
result = await self.endpoint.request(
route=route,
payload=payload,
serverless_request=serverless_request,
cost=cost,
retry=retry,
stream=stream,
session=self,
)
if result.get("status") == 410:
self.open = False
raise ValueError("Cannot make request on closed session.")
return result
return _wrapped_request()serverless/client/worker.py
from dataclasses import dataclass
from typing import Any, Dict, List, Tuple, Optional
@dataclass
class Worker:
id: int
status: str
cur_load: float
new_load: float
cur_load_rolling_avg: float
cur_perf: float
perf: float
measured_perf: float
dlperf: float
reliability: float
reqs_working: int
disk_usage: float
loaded_at: float
started_at: float
@staticmethod
def from_dict(d: Dict[str, Any]) -> "Worker":
# Be resilient to missing / extra fields
status = d.get("status") or "UNKNOWN"
return Worker(
id=int(d.get("id")),
status=status,
cur_load=float(d.get("cur_load", 0.0)),
new_load=float(d.get("new_load", 0.0)),
cur_load_rolling_avg=float(d.get("cur_load_rolling_avg", 0.0)),
cur_perf=float(d.get("cur_perf", 0.0)),
perf=float(d.get("perf", 0.0)),
measured_perf=float(d.get("measured_perf", 0.0)),
dlperf=float(d.get("dlperf", 0.0)),
reliability=float(d.get("reliability", 0.0)),
reqs_working=int(d.get("reqs_working", 0)),
disk_usage=float(d.get("disk_usage", 0.0)),
loaded_at=float(d.get("loaded_at", 0.0)),
started_at=float(d.get("started_at", 0.0)),
)
serverless/remote/__init__.py
import os
from .base import DockerLogin, Image, Autoscaling
if os.environ.get("IS_DEPLOYMENT", False):
from .serve import Deployment
else:
from .deploy import Deployment
__all__ = ["Deployment", "DockerLogin", "Image", "Autoscaling"]
serverless/remote/base.py
from abc import ABC, abstractmethod
from typing import (
Callable,
Any,
Awaitable,
ParamSpec,
Optional,
TypeVar,
Type,
AsyncContextManager,
TypedDict,
)
from typing_extensions import Unpack
from vastai.data import Query
from dataclasses import dataclass
class DockerLogin(TypedDict, total=False):
docker_login_user: str
docker_login_pass: str
docker_login_repo: str
class Image:
def __init__(
self,
from_image: str,
storage: float = 50,
worker_port: tuple[int, str] = (3000, "tcp"),
**docker_login: Unpack[DockerLogin],
):
self._image = from_image
self._envs: dict[str, str] = {}
self._runs: list[str | tuple[str, ...]] = []
self._venv: str | None = (
None # None=SDK default, ""=system python, str=custom venv path
)
self._pip_installs: list[str] = []
self._apt_gets: list[str] = []
self._docker_login: DockerLogin = docker_login
self._requires = Query.search_defaults()
self._storage = storage
self._copies: list[tuple[str, str]] = []
self._ports: set[tuple[int, str]] = {worker_port}
def pip_install(self, *args: str) -> "Image":
for arg in args:
self._pip_installs.append(arg)
return self
def apt_get(self, *args: str) -> "Image":
for arg in args:
self._apt_gets.append(arg)
return self
def run_script(self, script) -> "Image":
self._runs.append(script)
return self
def run_cmd(self, *args) -> "Image":
self._runs.append(args)
return self
def env(self, **kwargs: str) -> "Image":
for k, v in kwargs.items():
self._envs[k] = v
return self
def require(self, *args: Query) -> "Image":
for arg in args:
self._requires.extend(arg)
return self
def copy(self, src: str, dst: str) -> "Image":
self._copies.append((src, dst))
return self
def venv(self, path: str) -> "Image":
"""Use an existing venv at the given path instead of the SDK managed venv."""
self._venv = path
return self
def use_system_python(self) -> "Image":
"""Use the image's system Python instead of a venv."""
self._venv = ""
return self
def publish_port(self, number: int, type_: str = "tcp") -> "Image":
self._ports.add((number, type_))
return self
class Autoscaling(TypedDict, total=False):
# Scaling params
cold_workers: int
max_workers: int
min_load: int
min_cold_load: int
target_util: float
cold_mult: int
max_queue_time: float
target_queue_time: float
inactivity_timeout: int
autoscaler_instance: str
@dataclass
class Config:
name: str
pip_installs: list[str]
apt_gets: list[str]
envs: list[list[str]] | list[tuple[str, ...]] # format [(KEY,VALUE)]
runs: (
list[str | list[str]] | list[str | tuple[str, ...]]
) # tuples get turned into list after JSON serialize->deserialize round trip
T = TypeVar("T")
DeploymentT = TypeVar("DeploymentT", bound="Deployment_")
class Deployment_(ABC):
by_name: dict[Type, dict[str, "Deployment_"]] = {}
def __init__(
self,
name: Optional[str] = None,
tag: str = "default",
version_label: Optional[str] = None,
**kwargs,
):
self.name = name
self.tag = tag
self.version_label: str | None = version_label
self.root_module: str | None = None
self.is_main = False
def context(
self, *args, **kwargs
) -> Callable[
[Type[AsyncContextManager[T]]], Type[AsyncContextManager[T]]
]: # registers context_class as deployment context
return lambda c: c # no-op when not in serve mode
def get_context(self, context_class: Type[AsyncContextManager[T]]) -> T:
raise NotImplementedError(
f"{context_class}"
) # should not be run when not in serve mode
@property
def name(self):
return self.name_
@name.setter
def name(self, name: str | None):
self.name_ = name
if isinstance(name, str):
if self.__class__ not in self.by_name:
self.by_name[self.__class__] = {}
self.by_name[self.__class__][name] = self
@classmethod
def lookup(cls: Type[DeploymentT], name: str) -> DeploymentT | None:
if cls not in cls.by_name:
return None
ret = cls.by_name[cls].get(name)
return ret # type: ignore
def relativize(self, f: Callable[..., Any]) -> tuple[str]:
mod = f.__module__
name = f.__name__
file = f.__globals__.get("__file__")
# assume single file deployment for now.
if self.root_module is None:
self.root_module = mod
if self.name_ is None:
self.name = ""
if file is None:
raise Exception(
f"Cowardly refusing to deploy function {name} defined in your interactive Python session! We need it to be defined in a .py file to deploy from."
)
self.file = file
if f.__module__ != self.root_module:
raise Exception(
f"Using a single-file deployment object to deploy remote functions from multiple .py files! Remote function {name} is defined in {file} but we're already trying to deploy from {self.file}."
)
return (name,)
P = ParamSpec("P")
# TODO: add type checked version of remote
@abstractmethod
def remote(
self,
f: Callable[P, Awaitable[Any]] | None = None,
*,
benchmark_dataset: list[dict] | None = None,
benchmark_generator: Callable[[], dict] | None = None,
benchmark_runs: int = 10,
workload_calculator: Callable[..., float] | None = None,
) -> (
Callable[P, Awaitable[Any]]
| Callable[[Callable[P, Awaitable[Any]]], Callable[P, Awaitable[Any]]]
): ...
# Stubs for methods that are expected to be present in serve mode, but shouldn't do anything.
def configure_autoscaling(self, **kwargs: Unpack[Autoscaling]):
pass
def image(self, from_image: str, storage: int) -> Image:
return Image("")
def ensure_ready(self): # NO-OP in serve mode
pass
serverless/remote/deploy.py
import inspect
import threading
import time
import json
import os
from typing import Optional, Any, Callable, Awaitable, ParamSpec, BinaryIO
from typing_extensions import Unpack
from dataclasses import asdict, dataclass
from vastai.data import query
from vastai import AsyncClient
from vastai._base import _APIKEY_SENTINEL
from vastai.data.deployment import DeploymentConfig
from vastai.serverless.client import ManagedDeployment
from . import serialization
from .base import Deployment_, Config, DockerLogin, Image, Autoscaling
from .utils import create_deployment_tarball, compute_deployment_hash
from os.path import getsize
import tempfile
import asyncio
import logging
logger = logging.getLogger("vastai")
logger.debug("mode: deploy")
# TODO: implement heartbeat, sync ready.
"""
Design:
Deployment has a root module.
Single file vs package deployments:
single file:
- from vastai.serverless import deployments
- deployment = deployments.Deployment(Image)
- @deployment.remote...
- deployment.ensure_ready()
multi file: # skip for now
- @deployments.post_configure(name = name, image = image)
def post_configure(deployment):
deployment.ensure_ready()
with post_configure:
import mod1
import mod2
Serve:
- pull tarball
- read env, onstart, installs, submodules, name from config. # non essential
- inserts envs, runs installs, runs onstart. # non essential
- from vastai.serverless import deployments # essential
- import deployment # for package deployments require __init__.py to explicitly import remote function modules
- my_deployment = deployments.get(name) # name is None if not manually set here.
- serve(my_deployment) -> # essential
async loop
async serve(fname, args) -> serialize(my_deployment.registry[fname](deserialize(args)))
- fname relative to root module
Deploy:
- make endpoint args and get a tarball upload link
- make tarball and post it
- when functions are called:
- convert fname to be relative to root module.
Client -> Serve depends:
- POST -> app.endpoint/api/f serialized args -> returns f(args).
- serialization: root module agrees (i.e., `deployment` in serve is the root_module of deployment)
Serve depends:
- handle_f -> app.get(f_name)(args,kwargs)
- app <- deployments.get_app(name)
- import deployment -> registers app and app functions
- deployment module in ./deployment/ or ./deployment.py, untarred in tarball
- name from config.json
- environment:
- On Image defined by from_image
- with machine requirements set by requires_
- envs set to envs (from config.json)
- pip installs installed ""
- apt gets installed ""
- runs ran ""
- vastai sdk installed (by bootstrap script)
- where config.json, deployments are in tarball
- bootstrap script knows deployment_id to grab from tarball with from S3
Serve <-> Webserver:
- deployment_id in env
- bootstrap script in onstart
- Image, search params from deployments
Serve <-> Deploy:
- Tarball as promised
Deploy <-> webserver
"""
DEBUG_DEPLOYMENT_TAR = os.environ.get("DEBUG_DEPLOYMENT_TAR")
@dataclass
class _FullDeployment:
root_module: str
deployment: ManagedDeployment
P = ParamSpec("P")
class Deployment(Deployment_): # TODO: Async Context Manager compatible with client
def __init__(
self,
api_key: str | object = _APIKEY_SENTINEL,
ttl: float | None = None,
autoscaler_instance="prod",
autoscaler_url: Optional[str] = None,
webserver_url="https://console.vast.ai",
**kwargs,
):
super().__init__(**kwargs)
self.async_client = AsyncClient(api_key, vast_server=webserver_url)
self.client = self.async_client.serverless(
instance=autoscaler_instance, autoscaler_url=autoscaler_url
)
self.autoscaler_instance = autoscaler_instance
self.autoscaler_url = self.client.autoscaler_url
self.webserver_url = webserver_url
self._image: Image | None = None
self._autoscaling: Autoscaling | None = None
self._ttl = ttl
self._inner: _FullDeployment | None = None
def _compile_env(self, checked_image: Image) -> str:
envs = [f"-p {port}:{port}/{type_}" for port, type_ in checked_image._ports]
if self.autoscaler_url is not None:
envs.append(f"-e REPORT_ADDR={self.autoscaler_url}")
if self.webserver_url != "https://console.vast.ai":
envs.append(f"-e VAST_API_URL={self.webserver_url}")
if checked_image._venv == "":
envs.append("-e USE_SYSTEM_PYTHON=true")
envs.append("-e UV_SYSTEM_PYTHON=1")
envs.append("-e UV_BREAK_SYSTEM_PACKAGES=1")
elif checked_image._venv is not None:
envs.append(f"-e ENV_PATH={checked_image._venv}")
return " ".join(envs)
def _into_deployment_config_and_tarball(self, tar_path: str) -> DeploymentConfig:
# should error if _image, _autoscaling, file, or any other field needed to calculate DeploymentConfig is None
if not isinstance(self._image, Image):
raise Exception(
"Trying to deploy a deployment without an image configured."
)
if not isinstance(self.name, str):
raise Exception(
"Trying to deploy an unbound deployment. Have any remote functions been registered?"
)
if not isinstance(self.root_module, str):
raise Exception(
"Trying to deploy an unbound deployment. Have any remote functions been registered?"
)
if not isinstance(self._autoscaling, dict):
raise Exception(
"Trying to deploy a deployment without autoscaling configured."
)
logger.debug(f"Building deployment tarball at {tar_path}")
hash, size = self._compute_hash_and_filesize_and_make_tar(
tar_path, self.name, self._image
)
logger.info(f"Deployment tarball built: hash={hash}, size={size} bytes")
return DeploymentConfig(
name=self.name
if self.name
else self.root_module, # if "", is main deployment for module
image=self._image._image,
env=self._compile_env(self._image),
file_hash=hash,
file_size=size,
tag=self.tag,
search_params=self._image._requires.unparse_query(),
storage=self._image._storage,
ttl=self._ttl,
version_label=self.version_label,
**self._image._docker_login,
**self._autoscaling,
)
def configure_autoscaling(self, **kwargs: Unpack[Autoscaling]):
if self._autoscaling is None:
self._autoscaling = kwargs
else:
self._autoscaling.update(kwargs)
def image(
self, from_image: str, storage: int, **docker_login: Unpack[DockerLogin]
) -> Image:
self._image = Image(from_image, storage, **docker_login)
return self._image
def _collate_config(self, checked_name: str, checked_image: Image) -> Config:
# should error if any required fields are still None
return Config(
checked_name,
checked_image._pip_installs,
checked_image._apt_gets,
list(checked_image._envs.items()),
checked_image._runs,
)
def _compute_hash_and_filesize_and_make_tar(
self, tar_path: str, checked_name: str, checked_image: Image
) -> tuple[str, int]:
if not isinstance(self.file, str):
raise Exception(
"Trying to deploy a deployment not yet bound to a Python module. Have any remote functions been registered?"
)
config = self._collate_config(checked_name, checked_image)
hash = compute_deployment_hash(config, self.file, checked_image._copies)
create_deployment_tarball(tar_path, config, self.file, checked_image._copies)
size = getsize(tar_path)
return (hash, size)
def _heartbeat_thread(self, deployment: ManagedDeployment, ttl: float):
interval = min(ttl / 2, 60)
while True:
deployment.sync_heartbeat()
time.sleep(interval)
async def async_ensure_ready(self):
if not isinstance(self.root_module, str):
raise Exception(
"Trying to deploy a deployment not yet bound to a Python module. Have any remote functions been registered?"
)
logger.info(f"Preparing deployment '{self.name}' (module={self.root_module})")
with tempfile.NamedTemporaryFile(
delete_on_close=False
) as f: # deletes at end of context manager instead of at f.close
tar_path = f.name if not DEBUG_DEPLOYMENT_TAR else DEBUG_DEPLOYMENT_TAR
f.close() # _into_deployment_config_and_tarball will reopen it when it makes the tarball
config = self._into_deployment_config_and_tarball(tar_path)
logger.debug(f"Registering deployment with server")
deployment = await self.client.put_deployment(config)
if deployment.needs_upload:
logger.info(f"Uploading deployment tarball")
await deployment.upload(tar_path)
logger.info(f"Upload complete")
else:
logger.info(f"Deployment tarball already up to date, skipping upload")
if deployment.action == "soft_update":
wg_id = deployment.workergroup_id
if not wg_id:
wg_id = await self.client.find_workergroup_for_endpoint(
deployment.endpoint_id
)
if wg_id:
logger.info(f"Triggering rolling update for workergroup {wg_id}")
await self.client.update_workers(wg_id)
else:
logger.warning(
f"soft_update but no workergroup found for endpoint {deployment.endpoint_id}, skipping update_workers"
)
self._inner = _FullDeployment(self.root_module, deployment)
if self._ttl is not None:
threading.Thread(
target=self._heartbeat_thread,
args=(deployment, self._ttl),
daemon=True,
).start()
logger.info(f"Deployment '{self.name}' is ready (id={deployment.id})")
def ensure_ready(self):
asyncio.run(self.async_ensure_ready())
def _unwrap_worker_response(self, response: dict[str, Any]) -> serialization.JSON:
try:
return response["response"]["result"]
except KeyError:
raise Exception(
f"Remote function call failed with status code {response['status']}: {response['text']}"
)
async def _dispatch(self, f_name, globals, sig, args, kwargs) -> Any:
if not isinstance(self._inner, _FullDeployment):
raise Exception("Deployment is not ready. Call .ensure_ready() first!")
bound_args = sig.bind(*args, **kwargs)
bound_args.apply_defaults()
route = "/remote/" + "/".join(f_name)
logger.debug(f"Dispatching remote call to {route}")
return serialization.deserialize_unwrap_error(
self._unwrap_worker_response(
await self._inner.deployment.endpoint.request(
route,
{
"kwargs": {
k: serialization.serialize(v, self._inner.root_module)
for k, v in bound_args.arguments.items()
}
},
)
),
self._inner.root_module,
globals,
)
def remote(
self,
f: Callable[P, Awaitable[Any]] | None = None,
*,
benchmark_dataset: list[dict] | None = None,
benchmark_generator: Callable[[], dict] | None = None,
benchmark_runs: int = 10,
workload_calculator: Callable[..., float] | None = None,
) -> (
Callable[P, Awaitable[Any]]
| Callable[[Callable[P, Awaitable[Any]]], Callable[P, Awaitable[Any]]]
):
def decorator(
f: Callable[P, Awaitable[Any]], **_
) -> Callable[P, Awaitable[Any]]:
f_rel_name = self.relativize(f)
logger.debug(f"Registered remote function: {f.__name__}")
f_globals = f.__globals__
sig = inspect.signature(f)
def inner(*args: P.args, **kwargs: P.kwargs) -> Awaitable[Any]:
return self._dispatch(f_rel_name, f_globals, sig, args, kwargs)
return inner
if f is not None:
return decorator(f)
else:
return decorator
serverless/remote/serialization.py
import base64
from typing import Any, TypeAlias
from importlib import import_module
import logging
logger = logging.getLogger("vastai")
def relativize_module(mod: str, root_module):
if mod == root_module:
return {"type": "relative", "path": []}
elif mod.startswith(root_module):
return {
"type": "relative",
"path": mod[len(root_module) :].lstrip(".").split("."),
} # will be non-empty
else:
return {"type": "absolute", "path": mod.split(".")}
def derelativize_module(mod, root_module: str):
if mod["type"] == "relative":
prefix = [] if root_module == "__main__" else root_module.split(".")
return ".".join(prefix + mod["path"])
else:
return ".".join(mod["path"])
JSON: TypeAlias = int | str | float | list["JSON"] | dict[str, "JSON"]
def serialize(obj, root_module: str) -> JSON:
if type(obj) in [int, str, float, type(None)]:
return obj
elif type(obj) is bytes:
return {"type": "bytes", "contents": base64.b64encode(obj).decode("utf-8")}
elif type(obj) in [list, tuple]:
return {
"type": type(obj).__name__,
"contents": [serialize(child, root_module) for child in obj],
}
elif (
type(obj) is dict
): # JSON dicts must have string keys, so instead we do this as list of tuples
return {
"type": "dict",
"contents": [
[serialize(k, root_module), serialize(v, root_module)]
for k, v in obj.items()
],
}
elif hasattr(obj, "__dict__") and hasattr(
obj, "__class__"
): # is a member of a normal Python class
return {
"type": "obj",
"module": relativize_module(obj.__class__.__module__, root_module),
"class": obj.__class__.__qualname__,
"contents": {k: serialize(v, root_module) for k, v in obj.__dict__.items()},
}
raise TypeError(f"Object of type {type(obj)} is not JSON serializable!")
def deserialize(json, root_module: str, globals):
if type(json) in [int, str, float, type(None)]:
return json
elif json["type"] == "bytes":
return base64.b64decode(json["contents"])
elif json["type"] in ["list", "tuple"]:
return __builtins__[json["type"]](
deserialize(child, root_module, globals) for child in json["contents"]
)
elif json["type"] == "dict":
return dict(
(deserialize(k, root_module, globals), deserialize(v, root_module, globals))
for k, v in json["contents"]
)
elif json["type"] == "obj":
modname = derelativize_module(json["module"], root_module)
namespace = globals if modname == "" else import_module(modname).__dict__
cls = namespace[json["class"]]
obj = cls.__new__(cls)
for k, v in json["contents"].items():
setattr(obj, k, deserialize(v, root_module, globals))
return obj
raise TypeError(
f"JSON does not correspond to known Vast deployment datatype: {json}"
)
def serialize_ok(obj, root_module: str):
return {"ok": serialize(obj, root_module)}
def serialize_err(err, root_module: str):
return {"err": serialize(err, root_module)}
def deserialize_unwrap_error(json, root_module: str, globals):
logger.debug(f"got json: \n{json}")
for k, v in json.items():
if (k) == "ok":
return deserialize(v, root_module, globals)
if (k) == "err":
err = deserialize(v, root_module, globals)
if isinstance(err, Exception):
raise err
raise TypeError(
'Expected json to have format {"ok": deserializable} or {"err": deserializes_to_error_object}: '
+ json
)
serverless/remote/serve_deployment.py
#!/usr/bin/python
"""
Deployment bootstrap script invoked by start_server.sh when IS_DEPLOYMENT=true.
Downloads and extracts the deployment tarball, applies config.json
(env vars, apt packages, pip packages, run scripts), then starts the
deployment worker.
"""
import json
import os
import subprocess
import sys
import tarfile
import tempfile
import urllib.request
VAST_API_URL = os.environ.get("VAST_API_URL", "https://console.vast.ai")
def get_api_key() -> str:
"""Resolve an API key from the container environment.
Inside a Vast instance, Kaalia sets CONTAINER_API_KEY with the
instance-level API key that has permissions to fetch deployment blobs.
"""
key = os.environ.get("CONTAINER_API_KEY") or os.environ.get("VAST_API_KEY")
if not key:
raise RuntimeError(
"No API key found in environment (CONTAINER_API_KEY or VAST_API_KEY)"
)
return key
def get_download_url(deployment_id: str, api_key: str) -> str:
"""Fetch the presigned S3 download URL for the deployment blob."""
url = f"{VAST_API_URL}/api/v0/deployment/{deployment_id}/download_url/"
req = urllib.request.Request(url, headers={"Authorization": f"Bearer {api_key}"})
with urllib.request.urlopen(req) as resp:
data = json.loads(resp.read())
if not data.get("success"):
raise RuntimeError(f"Failed to get download URL: {data}")
return data["download_url"]
def download_tarball(download_url: str) -> str:
"""Download the deployment tarball to a temp file, returning its path."""
fd, path = tempfile.mkstemp(suffix=".tar.gz")
try:
with urllib.request.urlopen(download_url) as resp, os.fdopen(fd, "wb") as f:
while True:
chunk = resp.read(1 << 20) # 1 MB
if not chunk:
break
f.write(chunk)
except Exception:
os.unlink(path)
raise
return path
def extract_tarball(tarball_path: str):
cwd = os.getcwd()
with tarfile.open(tarball_path, "r:gz") as tf:
for member in tf.getmembers():
tf.extract(
member, filter="fully_trusted"
) # preserve absolute paths; deployment .tar files are created by client via deployments SDK and are trusted inside the container.
def get_config(path: str):
"""Load and return config.json as a Config dataclass."""
from vastai.serverless.remote.base import Config
with open(path) as f:
raw = json.load(f)
return Config(
name=raw["name"],
pip_installs=raw.get("pip_installs", []),
apt_gets=raw.get("apt_gets", []),
envs=raw.get("envs", []),
runs=raw.get("runs", []),
)
def export_envs(envs: list):
"""Export environment variables and append to /etc/environment."""
etc_lines = []
for entry in envs:
key, value = entry[0], entry[1]
os.environ[key] = value
etc_lines.append(f'{key}="{value}"\n')
if etc_lines:
with open("/etc/environment", "a") as f:
f.writelines(etc_lines)
def run_apt_gets(packages: list[str]):
"""Install apt packages."""
if not packages:
return
subprocess.run(
["apt-get", "update", "-qq"],
check=True,
)
subprocess.run(
["apt-get", "install", "-y", "-qq"] + packages,
check=True,
)
def run_pip_installs(packages: list[str]):
"""Install pip packages using uv."""
if not packages:
return
subprocess.run(
["uv", "pip", "install"] + packages,
check=True,
)
def run_scripts(runs: list):
"""Run scripts from config.
Each entry is either a string (passed to sh -c) or a list of args (execvp style).
"""
for entry in runs:
if isinstance(entry, str):
subprocess.run(["sh", "-c", entry], check=True)
elif isinstance(entry, (list, tuple)):
subprocess.run(list(entry), check=True)
else:
raise ValueError(f"Invalid run entry: {entry!r}")
def main():
deployment_id = os.environ.get("DEPLOYMENT_ID")
if not deployment_id:
raise RuntimeError("DEPLOYMENT_ID environment variable not set")
debug_deployment_tar = os.environ.get("DEBUG_DEPLOYMENT_TAR")
if debug_deployment_tar:
tarball_path = debug_deployment_tar
else:
api_key = get_api_key()
# Download deployment tarball
print(f"Fetching download URL for deployment {deployment_id}")
download_url = get_download_url(deployment_id, api_key)
print("Downloading deployment tarball")
tarball_path = download_tarball(download_url)
# Extract preserving absolute paths — the deployment runs in a
# container owned by the client, and we trust their path choices.
print("Extracting deployment tarball")
try:
extract_tarball(tarball_path)
finally:
os.unlink(tarball_path)
# Load config
config = get_config("config.json")
print(f"Loaded config for deployment: {config.name}")
# (1) Export envs
if config.envs:
print(f"Exporting {len(config.envs)} environment variables")
export_envs(config.envs)
# (2) Install packages
if config.apt_gets:
print(f"Installing {len(config.apt_gets)} apt packages")
run_apt_gets(config.apt_gets)
if config.pip_installs:
print(f"Installing {len(config.pip_installs)} pip packages")
run_pip_installs(config.pip_installs)
# (3) Run scripts
if config.runs:
print(f"Running {len(config.runs)} setup scripts")
run_scripts(config.runs)
# Look up and start deployment
from vastai.serverless.remote.serve import Deployment
# ensure deployment import time code is run in serve mode.
os.environ["IS_DEPLOYMENT"] = "1"
# deployment module/package is guaranteed from tarball
sys.path.insert(0, os.getcwd())
import deployment # running this import has the side effect of registering deployments and remote functions with vastai.serverless.remote.serve.Deployment
our_deployment = Deployment.lookup(config.name)
if our_deployment is None:
raise RuntimeError(f"Failed to lookup registered deployment: {config.name}")
print(f"Starting deployment worker: {config.name}")
worker = our_deployment.into_worker()
worker.run()
if __name__ == "__main__":
main()
serverless/remote/serve.py
from .base import Config, Deployment_
from ..server.worker import Worker, WorkerConfig, HandlerConfig, BenchmarkConfig
from .serialization import serialize, deserialize, serialize_ok, serialize_err
from typing import (
ParamSpec,
Type,
Any,
Callable,
Awaitable,
TypeVar,
AsyncContextManager,
Optional,
)
from dataclasses import dataclass
import asyncio
import logging
logger = logging.getLogger("vastai")
logger.debug("mode: serve")
@dataclass
class RemoteFunc:
func: Callable[..., Awaitable[Any]]
globals: dict[str, Any]
allow_parallel_requests: bool
max_queue_time: float
benchmark_dataset: Optional[list[dict[str, Any]]]
benchmark_generator: Optional[Callable[[], dict[str, Any]]]
benchmark_runs: int
workload_calculator: Optional[Callable[..., float]]
T = TypeVar("T")
P = ParamSpec("P")
class Deployment(Deployment_, AsyncContextManager):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.contexts: dict[Type, Any] = {}
self.context_managers: dict[Type, AsyncContextManager] = {}
self.remote_funcs: dict[tuple[str], RemoteFunc] = {}
def context(
self, *args, **kwargs
) -> Callable[
[Type[AsyncContextManager[T]]], Type[AsyncContextManager[T]]
]: # registers context_class as deployment context
def decorator(
context_class: Type[AsyncContextManager[T]],
) -> Type[AsyncContextManager[T]]:
self.context_managers[context_class] = context_class(*args, **kwargs)
return context_class
return decorator
def get_context(self, context_class: Type[AsyncContextManager[T]]) -> T:
if context_class not in self.contexts:
raise KeyError(f"Context {context_class} not registered or not yet entered")
return self.contexts[context_class]
async def __aenter__(self):
self.contexts = dict(
zip(
self.context_managers.keys(),
await asyncio.gather(
*(
manager.__aenter__()
for manager in self.context_managers.values()
)
),
)
)
async def __aexit__(self, exc_type, exc, tb):
await asyncio.gather(
*(
self.context_managers[cls_].__aexit__(exc_type, exc, tb)
for cls_ in self.contexts.keys() # only run aexits for contexts we've aentered
)
)
self.contexts = {} # invalidate contexts
def _wrap_remote_func(
self, root_module: str, func: Callable[..., Awaitable[Any]], func_globals: dict
) -> Callable[..., Awaitable[Any]]:
"""Wrap a remote function with serialization/deserialization.
Expects payload: {"args": [serialized_args], "kwargs": {serialized_kwargs}}
Returns: {"ok": serialized_result} on success, {"err": serialized_exception} on failure.
"""
async def wrapper(*, args: list = [], kwargs: dict = {}) -> dict:
deserialized_args = [
deserialize(a, root_module, func_globals) for a in args
]
deserialized_kwargs = {
k: deserialize(v, root_module, func_globals) for k, v in kwargs.items()
}
try:
result = await func(*deserialized_args, **deserialized_kwargs)
return serialize_ok(result, root_module)
except Exception as e:
return serialize_err(e, root_module)
return wrapper
def _wrap_workload_calculator(
self,
root_module: str,
user_calc: Callable[..., float],
func_globals: dict,
) -> Callable[[dict], float]:
"""Deserialize the request payload and score it with the user's calculator.
Runs synchronously inside count_workload and returns a plain float,
unlike the async, result-serializing remote-function wrapper.
"""
def calculator(payload: dict) -> float:
deserialized_args = [
deserialize(a, root_module, func_globals)
for a in payload.get("args", [])
]
deserialized_kwargs = {
k: deserialize(v, root_module, func_globals)
for k, v in payload.get("kwargs", {}).items()
}
return float(user_calc(*deserialized_args, **deserialized_kwargs))
return calculator
def into_worker(self) -> Worker:
handlers: list[HandlerConfig] = []
if isinstance(self.root_module, str):
for key, entry in self.remote_funcs.items():
route = "/remote/" + "/".join(key)
benchmark_config = None
if (
entry.benchmark_dataset is not None
or entry.benchmark_generator is not None
):
# Serialize benchmark dataset values so they match the
# format a real client sends (the wrapper deserializes them).
dataset = entry.benchmark_dataset
if dataset is not None:
dataset = [
{
"kwargs": {
k: serialize(v, self.root_module)
for k, v in item.items()
}
for item in dataset
}
]
benchmark_generator = None
if entry.benchmark_generator is not None:
root_module = self.root_module # in case self.root_module changes, bind to current root_module
original_benchmark_generator = (
entry.benchmark_generator
) # see above
benchmark_generator = lambda: {
"kwargs": {
k: serialize(v, root_module)
for k, v in original_benchmark_generator().items()
}
}
benchmark_config = BenchmarkConfig(
dataset=dataset,
generator=benchmark_generator,
runs=entry.benchmark_runs,
)
wrapped = self._wrap_remote_func(
self.root_module, entry.func, entry.globals
)
workload_calculator = None
if entry.workload_calculator is not None:
workload_calculator = self._wrap_workload_calculator(
self.root_module, entry.workload_calculator, entry.globals
)
handlers.append(
HandlerConfig(
route=route,
remote_function=wrapped,
allow_parallel_requests=entry.allow_parallel_requests,
benchmark_config=benchmark_config,
max_queue_time=entry.max_queue_time,
workload_calculator=workload_calculator,
)
)
config = WorkerConfig(
handlers=handlers,
lifecycle=self,
)
return Worker(config)
raise TypeError("Refusing to create empty Worker!")
def remote(
self,
f: Callable[P, Awaitable[Any]] | None = None,
*,
allow_parallel_requests: bool = False,
max_queue_time: float = 30.0,
benchmark_dataset: list[dict] | None = None,
benchmark_generator: Callable[[], dict] | None = None,
benchmark_runs: int = 10,
workload_calculator: Callable[..., float] | None = None,
) -> (
Callable[P, Awaitable[Any]]
| Callable[[Callable[P, Awaitable[Any]]], Callable[P, Awaitable[Any]]]
):
def decorator(f: Callable[P, Awaitable[Any]]) -> Callable[P, Awaitable[Any]]:
key = self.relativize(f)
self.remote_funcs[key] = RemoteFunc(
func=f,
globals=f.__globals__,
allow_parallel_requests=allow_parallel_requests,
max_queue_time=max_queue_time,
benchmark_dataset=benchmark_dataset,
benchmark_generator=benchmark_generator,
benchmark_runs=benchmark_runs,
workload_calculator=workload_calculator,
)
return f
if f is not None:
return decorator(f)
return decorator
serverless/remote/utils.py
from __future__ import annotations
import dataclasses
import hashlib
import io
import json
import os
import re
import tarfile
import tempfile
from typing import TYPE_CHECKING, Tuple, BinaryIO
if TYPE_CHECKING:
from vastai.serverless.remote.base import Config
def _sanitize_info(
info: tarfile.TarInfo, arcname: str | None = None
) -> tarfile.TarInfo:
"""Set uid/gid to 0 on a TarInfo.
If arcname is provided, override info.name. This is needed because
gettarinfo strips leading '/' from absolute arcnames.
"""
info.uid = 0
info.gid = 0
info.uname = "root"
info.gname = "root"
if arcname is not None:
info.name = arcname
return info
def add_file(tf: tarfile.TarFile, src_path: str, arcname: str) -> None:
"""Add a single file to the tarball under arcname, with uid/gid 0."""
info = tf.gettarinfo(src_path, arcname=arcname)
_sanitize_info(info, arcname)
with open(src_path, "rb") as f:
tf.addfile(info, f)
def add_folder(tf: tarfile.TarFile, src_path: str, arcname: str) -> None:
"""Recursively add a folder to the tarball, renaming the root to arcname, with uid/gid 0."""
src_path = os.path.normpath(src_path)
for root, dirs, files in os.walk(src_path):
rel = os.path.relpath(root, src_path)
arc_root = arcname if rel == "." else os.path.join(arcname, rel)
# Add directory entry
info = tf.gettarinfo(root, arcname=arc_root)
_sanitize_info(info, arc_root)
tf.addfile(info)
for fname in files:
file_path = os.path.join(root, fname)
arc_path = os.path.join(arc_root, fname)
info = tf.gettarinfo(file_path, arcname=arc_path)
_sanitize_info(info, arc_path)
with open(file_path, "rb") as f:
tf.addfile(info, f)
def add_string(
tf: tarfile.TarFile,
content: str,
arcname: str,
executable: bool = False,
) -> None:
"""Add a string as a text file to the tarball with uid/gid 0.
If executable=True, the file gets mode 0o755; otherwise 0o644.
"""
data = content.encode("utf-8")
info = tarfile.TarInfo(name=arcname)
info.size = len(data)
info.mode = 0o755 if executable else 0o644
_sanitize_info(info)
tf.addfile(info, io.BytesIO(data))
def add_path(tf: tarfile.TarFile, src_path: str, arcname: str) -> None:
"""Add a file or directory to the tarball, dispatching to add_file or add_folder."""
if os.path.isdir(src_path):
add_folder(tf, src_path, arcname)
elif os.path.isfile(src_path):
add_file(tf, src_path, arcname)
else:
raise FileNotFoundError(f"Source path does not exist: {src_path}")
def serialize_config(config: Config) -> str:
"""JSON-serialize a Config dataclass."""
return json.dumps(dataclasses.asdict(config))
def is_python_package(path: str) -> bool:
"""Return True if path is a Python package directory (contains __init__.py)."""
return os.path.isdir(path) and os.path.isfile(os.path.join(path, "__init__.py"))
def is_python_module(path: str) -> bool:
"""Return True if path is a single .py file."""
return os.path.isfile(path) and path.endswith(".py")
def deployment_arcname(deployment_path: str) -> str:
"""Return the archive name for a deployment: './deployment' for packages, './deployment.py' for modules.
Raises ValueError if the path is neither a Python package nor a module.
"""
if is_python_package(deployment_path):
return "./deployment"
if is_python_module(deployment_path):
return "./deployment.py"
raise ValueError(
f"Deployment path is neither a Python package (directory with __init__.py) "
f"nor a Python module (.py file): {deployment_path}"
)
IGNORE_COMMENT_BODY = b"!VAST_IGNORE_CHANGES"
# Matches the closing delimiter (with escape handling) for a multiline string.
_CLOSE_MULTILINE_RE = {
b'"""': re.compile(rb'(?:[^\\]|\\.)*?"""'),
b"'''": re.compile(rb"(?:[^\\]|\\.)*?'''"),
}
# Tokenizer for Python source outside multiline strings.
# Matches (in priority order): closed triple-quoted strings, unclosed
# triple-quotes (groups 1/2), single-quoted strings, or comment start (group 3).
# Unclosed triple-quotes MUST precede single-quote patterns, otherwise
# "..." would consume "" from the start of """.
# Unmatched text (code, whitespace, etc.) is simply skipped by finditer.
_PYTHON_TOKEN_RE = re.compile(
rb'"""(?:[^\\]|\\.)*?"""|'
rb"'''(?:[^\\]|\\.)*?'''|"
rb'(""")|'
rb"(''')|"
rb'"(?:[^"\\\n]|\\.)*"|'
rb"'(?:[^'\\\n]|\\.)*'|"
rb"(#)"
)
def _scan_python_line(
line: bytes,
in_multiline: bytes | None,
) -> tuple[bool, bytes | None]:
"""Scan a Python source line left-to-right, tracking string context.
Returns (has_ignore_comment, new_multiline_state).
"""
i = 0
if in_multiline is not None:
m = _CLOSE_MULTILINE_RE[in_multiline].match(line)
if m is None:
return False, in_multiline
i = m.end()
for m in _PYTHON_TOKEN_RE.finditer(line, i):
if m.group(1):
return False, b'"""'
if m.group(2):
return False, b"'''"
if m.group(3):
comment_body = line[m.start() + 1 :].strip()
return comment_body == IGNORE_COMMENT_BODY, None
return False, None
def filter_ignored_lines(data: bytes) -> bytes:
"""Remove lines whose comment is exactly ``#!VAST_IGNORE_CHANGES``.
Tracks triple-quoted multiline strings so that occurrences inside string
literals are not treated as comments.
"""
lines = data.split(b"\n")
result: list[bytes] = []
in_multiline: bytes | None = None
for line in lines:
has_ignore, in_multiline = _scan_python_line(line, in_multiline)
if not has_ignore:
result.append(line)
return b"\n".join(result)
def read_file_for_hash(path: str, filter_comments: bool = False) -> bytes:
"""Read a file's bytes for hashing.
When filter_comments is True and the file is a .py file, lines whose
comment is exactly ``# !VAST_IGNORE_CHANGES`` are excluded.
"""
with open(path, "rb") as f:
data = f.read()
if filter_comments and path.endswith(".py"):
data = filter_ignored_lines(data)
return data
def hash_update_file(
hasher: hashlib._Hash,
src_path: str,
arcname: str,
filter_comments: bool = False,
) -> None:
"""Feed a single file into the hasher: its archive name followed by its content."""
hasher.update(arcname.encode("utf-8"))
hasher.update(read_file_for_hash(src_path, filter_comments=filter_comments))
def hash_update_directory(
hasher: hashlib._Hash,
src_path: str,
arcname: str,
filter_comments: bool = False,
) -> None:
"""Walk a directory in sorted order, feeding each file into the hasher."""
src_path = os.path.normpath(src_path)
for root, dirs, files in os.walk(src_path):
dirs.sort()
rel = os.path.relpath(root, src_path)
arc_root = arcname if rel == "." else os.path.join(arcname, rel)
for fname in sorted(files):
hash_update_file(
hasher,
os.path.join(root, fname),
os.path.join(arc_root, fname),
filter_comments=filter_comments,
)
def hash_update_path(
hasher: hashlib._Hash,
src_path: str,
arcname: str,
filter_comments: bool = False,
) -> None:
"""Feed a file or directory into the hasher, dispatching appropriately."""
if os.path.isdir(src_path):
hash_update_directory(
hasher, src_path, arcname, filter_comments=filter_comments
)
elif os.path.isfile(src_path):
hash_update_file(hasher, src_path, arcname, filter_comments=filter_comments)
else:
raise FileNotFoundError(f"Source path does not exist: {src_path}")
def compute_deployment_hash(
config: Config,
deployment_path: str,
extra_files: list[tuple[str, str]] | None = None,
) -> str:
"""Compute a SHA-256 hex digest over deployment data without creating a tarball.
Hashes (in order):
1. The JSON-serialized Config
2. The deployment files (with #!VAST_IGNORE_CHANGES comment lines excluded from .py files)
3. Any extra files (no comment filtering)
Args:
config: The deployment Config object.
deployment_path: Path to the deployment source (a .py file or a package directory).
extra_files: List of (source_path, dest_path) pairs.
"""
hasher = hashlib.sha256()
hasher.update(serialize_config(config).encode("utf-8"))
hash_update_path(
hasher,
deployment_path,
deployment_arcname(deployment_path),
filter_comments=True,
)
for src, dest in extra_files or []:
hash_update_path(hasher, src, dest)
return hasher.hexdigest()
def create_deployment_tarball(
tar_path: str,
config: Config,
deployment_path: str,
extra_files: list[tuple[str, str]] | None = None
):
"""Create a deployment tarball in /tmp and return its path.
The tarball contains:
- ./config.json — JSON-serialized Config
- ./deployment[.py] — the deployment package or module
- extra files at their specified destination paths
Args:
config: The deployment Config object.
deployment_path: Path to the deployment source (a .py file or a package directory).
extra_files: List of (source_path, dest_path) pairs. dest_path may be absolute.
"""
tf = tarfile.TarFile.open(tar_path, mode="w:gz")
try:
add_string(tf, serialize_config(config), "./config.json")
add_path(tf, deployment_path, deployment_arcname(deployment_path))
for src, dest in extra_files or []:
add_path(tf, src, dest)
finally:
tf.close()
serverless/server/__init__.py
serverless/server/lib/__init__.py
serverless/server/lib/backend.py
import os
import sys
import json
import time
import base64
import dataclasses
import logging
from asyncio import sleep, gather, Semaphore, create_task
from typing import (
Tuple,
Awaitable,
NoReturn,
List,
Union,
Callable,
Optional,
Any,
Dict,
AsyncContextManager,
)
from functools import cached_property
from distutils.util import strtobool
from collections import deque
from asyncio import sleep, CancelledError
from anyio import open_file
from aiohttp import (
web,
ClientResponse,
ClientSession,
ClientConnectorError,
ClientTimeout,
FormData,
TCPConnector,
)
import asyncio
import string
import random
from Crypto.Signature import pkcs1_15
from Crypto.Hash import SHA256
from Crypto.PublicKey import RSA
from .metrics import Metrics
from .data_types import (
AuthData,
EndpointHandler,
LogAction,
ApiPayload_T,
JsonDataException,
RequestMetrics,
BenchmarkResult,
Session,
)
VERSION = "1.1.0"
log = logging.getLogger(__file__)
# defines the minimum wait time between sending updates to autoscaler
LOG_POLL_INTERVAL = 0.1
# Defines waiting interval for session garbage collection
SESSION_GC_INTERVAL = 5.0
BENCHMARK_INDICATOR_FILE = ".has_benchmark"
MAX_PUBKEY_FETCH_ATTEMPTS = 5
@dataclasses.dataclass
class Backend:
"""
This class is responsible for:
1. Tailing logs and updating load time metrics
2. Taking an EndpointHandler alongside incoming payload, preparing a json to be sent to the model, and
sending the request. It also updates metrics as it makes those requests.
3. Running a benchmark from an EndpointHandler
"""
model_server_url: str
model_log_file: str
benchmark_handler: (
EndpointHandler # this endpoint handler will be used for benchmarking
)
log_actions: List[Tuple[LogAction, str]]
reqnum = -1
version = VERSION
sem: Semaphore = dataclasses.field(default_factory=Semaphore)
queue: deque = dataclasses.field(default_factory=deque, repr=False)
_queue_lock: asyncio.Lock = dataclasses.field(
default_factory=asyncio.Lock, repr=False
)
unsecured: bool = dataclasses.field(
default_factory=lambda: bool(strtobool(os.environ.get("UNSECURED", "false"))),
)
report_addr: str = dataclasses.field(
default_factory=lambda: os.environ.get("REPORT_ADDR", "https://run.vast.ai")
)
mtoken: str = dataclasses.field(
default_factory=lambda: os.environ.get("MASTER_TOKEN", "")
)
healthcheck_url: str = dataclasses.field(
default_factory=lambda: os.environ.get("MODEL_HEALTH_ENDPOINT", "")
)
_sessions_lock: asyncio.Lock = dataclasses.field(
default_factory=asyncio.Lock, repr=False
)
sessions: Dict[str, Session] = dataclasses.field(default_factory=dict)
session_metrics: Dict[str, RequestMetrics] = dataclasses.field(default_factory=dict)
max_sessions: int = dataclasses.field(default=-1)
lifecycle: Optional[AsyncContextManager] = dataclasses.field(default=None)
async def pyworker_update_handler(self, request: web.Request) -> web.Response:
# Verify authorization header matches mtoken
auth_header = request.headers.get("Authorization", "")
token = (
auth_header.removeprefix("Bearer ")
if auth_header.startswith("Bearer ")
else auth_header
)
if token != self.mtoken:
return web.json_response({"error": "unauthorized"}, status=401)
try:
async with await open_file("/.force_update", "w") as f:
await f.write("")
except Exception as e:
log.error(f"Failed to write /.force_update: {e}")
return web.json_response(
{"error": f"failed to write file: {e}"}, status=500
)
return web.json_response({"ok": True}, status=200)
async def session_health_handler(self, request: web.Request) -> web.Response:
try:
data = await request.json()
session_id = data.get("session_id")
session_auth = data.get("session_auth")
except json.JSONDecodeError:
return web.json_response({"error": "invalid JSON"}, status=422)
if not session_id:
return web.json_response({"error": "missing session_id"}, status=422)
async with self._sessions_lock:
session = self.sessions.get(session_id)
if session is None:
return web.json_response({"ok": False}, status=200)
if session_auth is None or session.auth_data != session_auth:
return web.json_response(
{"error": "session_auth is not valid"}, status=401
)
return web.json_response({"ok": True}, status=200)
async def session_get_handler(self, request: web.Request) -> web.Response:
try:
data = await request.json()
session_id = data.get("session_id")
session_auth = data.get("session_auth")
except json.JSONDecodeError:
return web.json_response({"error": "invalid JSON"}, status=422)
if not session_id:
return web.json_response({"error": "missing session_id"}, status=422)
session = None
async with self._sessions_lock:
session = self.sessions.get(session_id)
if session is None:
return web.json_response({"error": "session does not exist"}, status=400)
if session_auth is None or session.auth_data != session_auth:
return web.json_response({"error": "session_auth is not valid"}, status=401)
else:
return web.json_response(
{
"session_id": session_id,
"auth_data": session.auth_data,
"lifetime": session.lifetime,
"expiration": session.expiration,
"created_at": session.created_at,
"on_close_route": session.on_close_route,
"on_close_payload": session.on_close_payload,
}
)
async def __run_session_on_close(self, session: Session) -> None:
"""
Best-effort POST to session.on_close_route with session.on_close_payload as JSON body.
Intended to be non-fatal and bounded in time.
"""
on_close_route = getattr(session, "on_close_route", None)
if not on_close_route:
return
on_close_payload = getattr(session, "on_close_payload", None)
if isinstance(on_close_payload, dict):
body: Dict[str, Any] = dict(on_close_payload) # copy
elif on_close_payload is None:
body = {}
else:
# Allow non-dict payloads without breaking the call contract.
body = {"payload": on_close_payload}
# Make it easier for the receiver to correlate the callback.
body.setdefault("session_id", getattr(session, "session_id", None))
try:
timeout = ClientTimeout(total=10)
async with self.session.post(
url=on_close_route,
json=body,
timeout=timeout,
) as resp:
text = await resp.text()
if resp.status >= 400:
log.debug(
f"on_close POST failed: route={on_close_route} status={resp.status} body={text[:512]!r}"
)
except Exception as e:
log.debug(f"on_close POST exception: route={on_close_route} error={e!r}")
return
async def __close_session(self, session_id: str) -> bool:
"""
Close a session and cancel all in-flight tasks associated with it.
Returns True if a session was removed, False if it didn't exist.
"""
async with self._sessions_lock:
session = self.sessions.pop(session_id, None)
if session is None:
return False
# Cancel all in-flight request handler tasks
for req in list(session.requests):
try:
tr = getattr(req, "transport", None)
if tr is not None and not tr.is_closing():
tr.close()
except Exception:
pass
session.requests.clear()
request_metrics = self.session_metrics.pop(session_id, None)
# Run the on_close callback
try:
await self.__run_session_on_close(session)
except Exception:
pass
# Update metrics outside lock
if request_metrics is not None:
self.metrics._request_success(request_metrics)
self.metrics._request_end(request_metrics)
return True
async def session_end_handler(self, request: web.Request) -> web.Response:
try:
data = await request.json()
session_id = data.get("session_id")
session_auth = data.get("session_auth")
except json.JSONDecodeError:
return web.json_response({"error": "invalid JSON"}, status=422)
if not session_id:
return web.json_response({"error": "missing session_id"}, status=422)
session = None
async with self._sessions_lock:
session = self.sessions.get(session_id)
if session is None:
return web.json_response(
{"error": f"session with id {session_id} not found"},
status=400,
)
if session_auth is None or session.auth_data != session_auth:
return web.json_response(
{"error": "session_auth is not valid"}, status=401
)
closed = await self.__close_session(session_id)
if not closed:
return web.json_response({"error": "session already closed"}, status=410)
return web.json_response(
{"ended": True, "removed_session": session_id},
status=200,
)
def generate_session_id(self):
characters = string.ascii_letters + string.digits
random_string = "".join(random.choices(characters, k=13))
return random_string
async def session_create_handler(self, request: web.Request) -> web.Response:
try:
data = await request.json()
auth_data = data.get("auth_data")
payload = data.get("payload")
except json.JSONDecodeError:
return web.json_response(dict(error="invalid JSON"), status=422)
now = time.time()
session_request_metrics = RequestMetrics(
request_idx=auth_data.get("request_idx"),
reqnum=auth_data.get("reqnum"),
workload=auth_data.get("cost"),
status="SessionActive",
is_session=True,
entered_queue_at=now,
work_started_at=now,
)
async with self._sessions_lock:
if (
not (self.max_sessions is None or self.max_sessions == 0)
and len(self.sessions) >= self.max_sessions
):
self.metrics._request_reject(session_request_metrics)
return web.Response(status=429)
# Set the session expiration time, and the TTL extension per request/get
lifetime = payload.get("lifetime", 60.0)
now = time.time()
expiration = now + lifetime
session_id = self.generate_session_id()
on_close_route = None
on_close_payload = None
if payload.get("on_close_route") is not None:
on_close_route = payload.get("on_close_route")
if payload.get("on_close_payload") is not None:
on_close_payload = payload.get("on_close_payload")
session = Session(
session_id=session_id,
lifetime=lifetime,
expiration=expiration,
auth_data=auth_data,
on_close_route=on_close_route,
on_close_payload=on_close_payload,
request_idx=session_request_metrics.request_idx,
)
self.sessions[session_id] = session
self.session_metrics[session_id] = session_request_metrics
self.metrics._request_start(session_request_metrics)
return web.json_response(
{"session_id": session.session_id, "expiration": session.expiration},
status=201,
)
def __post_init__(self):
self.metrics = Metrics()
self.metrics._set_version(self.version)
self.metrics._set_mtoken(self.mtoken)
self._pubkey: Optional[RSA.RsaKey] = None
self.__pubkey_fetch_complete: asyncio.Event = asyncio.Event()
self.__pubkey_failed: bool = False
self.__start_healthcheck: asyncio.Event = asyncio.Event()
self.__healthcheck_ready: asyncio.Event = asyncio.Event()
self.__healthcheck_succeeded: bool = False
@cached_property
def session(self):
log.debug(f"Starting TCP session with model server at {self.model_server_url}")
connector = TCPConnector(
force_close=True, # Required for long running jobs
enable_cleanup_closed=True,
)
timeout = ClientTimeout(total=None)
return ClientSession(
self.model_server_url, timeout=timeout, connector=connector
)
def create_handler(
self,
handler: EndpointHandler[ApiPayload_T],
) -> Callable[[web.Request], Awaitable[Union[web.Response, web.StreamResponse]]]:
async def handler_fn(
request: web.Request,
) -> Union[web.Response, web.StreamResponse]:
return await self.__handle_request(handler=handler, request=request)
return handler_fn
#######################################Private#######################################
async def _fetch_pubkey(self) -> None:
"""Fetch the serverless public key with exponential backoff retries."""
report_addr = self.report_addr.rstrip("/")
url = f"{report_addr}/pubkey/"
timeout = ClientTimeout(total=10)
for attempt in range(1, MAX_PUBKEY_FETCH_ATTEMPTS + 1):
try:
async with ClientSession(timeout=timeout) as client:
async with client.get(url) as response:
response.raise_for_status()
text = await response.text()
log.debug("Serverless Public Key:")
log.debug(text)
key = RSA.import_key(text)
if key is not None:
self._pubkey = key
self.__pubkey_fetch_complete.set()
log.debug(
f"Successfully fetched pubkey on attempt {attempt}"
)
return
except (ValueError, ClientConnectorError, Exception) as e:
log.debug(
f"Error downloading key (attempt {attempt}/{MAX_PUBKEY_FETCH_ATTEMPTS}): {e}"
)
if attempt < MAX_PUBKEY_FETCH_ATTEMPTS:
backoff = 2 ** (attempt - 1) # 1, 2, 4, 8, 16 seconds
log.debug(f"Retrying in {backoff} seconds...")
await asyncio.sleep(backoff)
log.error(f"Failed to fetch pubkey after {MAX_PUBKEY_FETCH_ATTEMPTS} attempts")
self.__pubkey_failed = True
self.__pubkey_fetch_complete.set() # Signal that fetch is complete (even though it failed)
self.backend_errored(
"Failed to get Serverless public key after all retry attempts"
)
async def __handle_request(
self,
handler: EndpointHandler[ApiPayload_T],
request: web.Request,
) -> Union[web.Response, web.StreamResponse]:
"""use this function to forward requests to the model endpoint"""
try:
data = await request.json()
auth_data, payload, session_id = handler.get_data_from_request(data)
except JsonDataException as e:
return web.json_response(data=e.message, status=422)
except json.JSONDecodeError:
return web.json_response(dict(error="invalid JSON"), status=422)
workload = payload.count_workload()
request_metrics: RequestMetrics = RequestMetrics(
request_idx=auth_data.request_idx,
reqnum=auth_data.reqnum,
workload=workload,
status="Created",
entered_queue_at=time.time(),
)
event = asyncio.Event()
session = None
if session_id is not None:
async with self._sessions_lock:
session = self.sessions.get(session_id)
if session is None:
return web.json_response(dict(error="invalid session"), status=410)
session.expiration += session.lifetime
session.requests.append(request)
session.session_reqnum += 1
request_metrics.session = session
request_metrics.session_reqnum = session.session_reqnum
async def advance_queue_after_completion(event: asyncio.Event):
"""Pop current head and wake next waiter, if any."""
async with self._queue_lock:
# If this event is current head, wake next waiter
if self.queue and self.queue[0] is event:
self.queue.popleft()
if self.queue:
self.queue[0].set()
else:
# Else, remove it from the queue
try:
self.queue.remove(event)
except ValueError:
pass
async def make_request() -> Union[web.Response, web.StreamResponse]:
request_metrics.work_started_at = time.time()
try:
if handler.remote_function:
# Remote dispatch: return JSON directly, bypassing the
# fake ClientResponse → generate_client_response path
# to avoid double-serialization of large payloads.
remote_func_params = payload.generate_payload_json()
result = await handler.call_remote_dispatch_function(
params=remote_func_params
)
res = web.json_response({"result": result})
else:
response = await self.__call_backend(
handler=handler, payload=payload
)
res = await handler.generate_client_response(request, response)
self.metrics._request_success(request_metrics)
return res
except asyncio.CancelledError:
raise
except Exception as e:
self.metrics._request_errored(request_metrics, str(e))
return web.Response(status=500)
###########
if self.__check_signature(auth_data) is False:
self.metrics._request_reject(request_metrics)
return web.Response(status=401)
# Check queue time with lock to prevent race conditions on concurrent requests
if handler.max_queue_time is not None:
async with self._queue_lock:
if self.metrics.model_metrics.wait_time > handler.max_queue_time:
self.metrics._request_reject(request_metrics)
return web.Response(status=429)
work_task = None
self.metrics._request_start(request_metrics)
try:
if handler.allow_parallel_requests:
coro = make_request()
try:
work_task = create_task(coro)
except Exception:
coro.close()
raise
# Handler cancellation will raise CancelledError on client disconnect
return await work_task
# FIFO-queue branch
else:
# Insert a Event into the queue for this request (with lock)
# Event.set() == our request is up next
async with self._queue_lock:
self.queue.append(event)
if self.queue and self.queue[0] is event:
event.set()
# Wait for our turn - CancelledError raised if client disconnects
await event.wait()
# We are the next-up request in the queue
if request_metrics.session is None:
log.debug(f"Starting work on request {request_metrics.reqnum}")
else:
log.debug(
f"Starting work on request {request_metrics.session_reqnum} in session {request_metrics.session.request_idx}"
)
# Execute the work task
coro = make_request()
try:
work_task = create_task(coro)
except Exception:
coro.close()
raise
return await work_task
except asyncio.CancelledError:
# With handler_cancellation enabled, this indicates client disconnect
log.debug(f"Request {request_metrics.reqnum} cancelled (client disconnect)")
self.metrics._request_canceled(request_metrics)
return web.Response(status=499)
except Exception as e:
log.debug(f"Exception in main handler loop {e}")
return web.Response(status=500)
finally:
try:
# Remove request from session if present
if session is not None and session_id is not None:
async with self._sessions_lock:
s = self.sessions.get(session_id)
if s is not None:
try:
s.requests.remove(request)
except ValueError:
pass
if not handler.allow_parallel_requests:
await advance_queue_after_completion(event)
self.metrics._request_end(request_metrics)
# Cleanup work task if still pending
if work_task and not work_task.done():
work_task.cancel()
await asyncio.gather(work_task, return_exceptions=True)
except Exception as e:
log.error(f"Error during request cleanup: {e}")
async def __healthcheck(self) -> None:
"""
Periodically hit the healthcheck endpoint using the same session
configuration as normal API calls.
"""
health_check_url = self.healthcheck_url
if not health_check_url:
log.debug("No healthcheck endpoint defined, skipping healthcheck")
return
await self.__start_healthcheck.wait()
timeout = ClientTimeout(total=10)
while True:
try:
log.debug(f"Performing healthcheck on {health_check_url}")
async with self.session.get(
health_check_url,
timeout=timeout,
) as response:
status = response.status
if status == 200:
if not self.__healthcheck_succeeded:
self.__healthcheck_succeeded = True
self.__healthcheck_ready.set()
log.debug("First healthcheck succeeded - model is ready")
else:
log.debug("Healthcheck successful")
else:
msg = f"Healthcheck failed with status: {status}"
log.debug(msg)
if self.__healthcheck_succeeded:
self.backend_errored(msg)
except CancelledError:
log.debug("Healthcheck task cancelled; exiting loop")
break
except Exception as e:
log.debug(f"Healthcheck failed with exception: {e}")
if self.__healthcheck_succeeded:
self.backend_errored(str(e))
await sleep(10)
async def _start_tracking(self) -> None:
if self.lifecycle is not None:
await gather(
self._fetch_pubkey(),
self.__lifecycle_startup(),
self.metrics._send_metrics_loop(),
self.__healthcheck(),
self.metrics._send_delete_requests_loop(),
self.__session_gc_loop(),
)
else:
await gather(
self._fetch_pubkey(),
self.__read_logs(),
self.metrics._send_metrics_loop(),
self.__healthcheck(),
self.metrics._send_delete_requests_loop(),
self.__session_gc_loop(),
)
async def __lifecycle_startup(self) -> None:
"""
Alternative to __read_logs for deployments that use an AsyncContextManager
lifecycle instead of log-based ready detection.
Enters the lifecycle context (which replaces on_init / log-based readiness),
runs the benchmark, then marks the model as loaded. Guarantees __aexit__ runs
on cancellation or exception.
"""
try:
await self.lifecycle.__aenter__()
try:
max_throughput = await self.__run_benchmark()
self.__start_healthcheck.set()
if self.healthcheck_url:
log.debug("Lifecycle ready, waiting for healthcheck...")
try:
await asyncio.wait_for(
self.__healthcheck_ready.wait(), timeout=300.0
)
except asyncio.TimeoutError:
raise Exception(
"Timed out waiting for healthcheck after lifecycle ready (300s)"
)
else:
log.debug("Lifecycle ready, no healthcheck configured")
self.metrics._model_loaded(max_throughput=max_throughput)
log.debug("Model marked as loaded via lifecycle")
# Keep alive — this coroutine lives for the duration of the server
await asyncio.Event().wait()
except BaseException:
await self.lifecycle.__aexit__(*sys.exc_info())
raise
else:
await self.lifecycle.__aexit__(None, None, None)
except Exception as e:
log.error(f"Lifecycle startup failed: {e}")
self.backend_errored(f"Lifecycle startup failed: {e}")
def backend_errored(self, msg: str) -> None:
self.metrics._model_errored(msg)
async def __call_backend(
self, handler: EndpointHandler[ApiPayload_T], payload: ApiPayload_T
) -> ClientResponse:
if handler.remote_function:
return await self.__call_remote_dispatch_function(
handler=handler, payload=payload
)
else:
return await self.__call_api(handler=handler, payload=payload)
@staticmethod
def __build_form_data(fields: Dict[str, Any]) -> FormData:
"""A (filename, bytes, content_type) tuple is a file part, a list repeats the field,
None is omitted, a dict or nested list is sent as JSON, other scalars as text.
Bare bytes are refused: pass a file tuple."""
form = FormData(default_to_multipart=True) # multipart even with no file part
for name, value in fields.items():
values = value if isinstance(value, list) else [value]
for item in values:
if item is None:
continue
if isinstance(item, tuple):
filename, content, content_type = item
form.add_field(
name, content, filename=filename, content_type=content_type
)
elif isinstance(item, (bytes, bytearray)):
raise TypeError(
f"multipart field {name!r}: bytes must be a file part "
"(filename, bytes, content_type)"
)
elif isinstance(item, (dict, list)):
form.add_field(name, json.dumps(item))
elif isinstance(item, bool):
form.add_field(name, "true" if item else "false")
else:
form.add_field(name, str(item))
return form
async def __call_api(
self, handler: EndpointHandler[ApiPayload_T], payload: ApiPayload_T
) -> ClientResponse:
multipart_fields = payload.generate_payload_multipart()
if multipart_fields is not None:
log.debug(
f"posting multipart to endpoint: '{handler.endpoint}', "
f"fields: {list(multipart_fields)}"
)
return await self.session.post(
url=handler.endpoint, data=self.__build_form_data(multipart_fields)
)
api_payload = payload.generate_payload_json()
# Keys only: values carry prompts and inline media.
log.debug(f"posting to endpoint: '{handler.endpoint}', payload keys: {list(api_payload) if isinstance(api_payload, dict) else type(api_payload).__name__}")
return await self.session.post(url=handler.endpoint, json=api_payload)
async def __call_remote_dispatch_function(
self, handler: EndpointHandler[ApiPayload_T], payload: ApiPayload_T
) -> ClientResponse:
remote_func_params = payload.generate_payload_json()
log.debug(
f"Calling remote function on {handler.endpoint} "
f"with params {remote_func_params}"
)
result = await handler.call_remote_dispatch_function(params=remote_func_params)
# Wrap the result in a fake ClientResponse-like object
class RemoteFunctionClientResponse:
def __init__(self, data: Any, status: int = 200):
self._body = json.dumps({"result": data}).encode("utf-8")
self.status = status
self.content_type = "application/json"
self.headers = {"Content-Type": self.content_type}
async def read(self) -> bytes:
return self._body
return RemoteFunctionClientResponse(result)
def __check_signature(self, auth_data: AuthData) -> bool:
if self.unsecured is True:
return True
# Reject early if pubkey was never loaded
if self._pubkey is None:
log.error("Rejecting request: pubkey not loaded")
return False
def verify_signature(message, signature):
h = SHA256.new(message.encode())
try:
pkcs1_15.new(self._pubkey).verify(h, base64.b64decode(signature))
return True
except (ValueError, TypeError):
return False
message = {"url": auth_data.url}
if verify_signature(
json.dumps(message, indent=4, sort_keys=True), auth_data.signature
):
self.reqnum = max(auth_data.reqnum, self.reqnum)
return True
else:
log.error(
f"Signature error: signature verification failed, sig:{auth_data.signature}, message: {message}"
)
return False
async def __run_benchmark(self) -> float:
"""Run the benchmark against the benchmark handler. Returns max throughput."""
log.debug("Running benchmark")
try:
with open(BENCHMARK_INDICATOR_FILE, "r") as f:
perf = float(f.readline())
log.debug(f"Already ran benchmark for perf score of {perf}")
return perf
except FileNotFoundError:
pass
if self.benchmark_handler.do_warmup_benchmark:
log.debug(
f"Performing benchmark on endpoint {self.benchmark_handler.endpoint}"
)
log.debug("Initial run to trigger model loading...")
payload = self.benchmark_handler.make_benchmark_payload()
await self.__call_backend(handler=self.benchmark_handler, payload=payload)
max_throughput = 0
sum_throughput = 0
concurrent_requests = (
self.benchmark_handler.concurrency
if self.benchmark_handler.allow_parallel_requests
else 1
)
for run in range(1, self.benchmark_handler.benchmark_runs + 1):
start = time.time()
benchmark_requests = []
for i in range(concurrent_requests):
payload = self.benchmark_handler.make_benchmark_payload()
workload = payload.count_workload()
task = self.__call_backend(
handler=self.benchmark_handler, payload=payload
)
benchmark_requests.append(
BenchmarkResult(request_idx=i, workload=workload, task=task)
)
responses = await gather(*[br.task for br in benchmark_requests])
for br, response in zip(benchmark_requests, responses):
br.response = response
total_workload = sum(
br.workload for br in benchmark_requests if br.is_successful
)
time_elapsed = time.time() - start
successful_responses = sum(
[1 for br in benchmark_requests if br.is_successful]
)
if successful_responses == 0:
self.backend_errored("No successful responses from benchmark")
log.error(f"Benchmark Failed: No successful responses")
return 0.0
throughput = total_workload / time_elapsed
sum_throughput += throughput
max_throughput = max(max_throughput, throughput)
# Log results for debugging
log.debug(
"\n".join(
[
"#" * 60,
f"Run: {run}, concurrent_requests: {concurrent_requests}",
f"Total workload: {total_workload}, time_elapsed: {time_elapsed}s",
f"Throughput: {throughput} workload/s",
f"Successful responses: {successful_responses}/{concurrent_requests}",
"#" * 60,
]
)
)
average_throughput = sum_throughput / self.benchmark_handler.benchmark_runs
log.debug(
f"Benchmark complete: average perf is {average_throughput}, measured perf is {max_throughput}"
)
with open(BENCHMARK_INDICATOR_FILE, "w") as f:
f.write(str(max_throughput))
return max_throughput
async def __read_logs(self) -> Awaitable[NoReturn]:
async def handle_log_line(log_line: str) -> None:
"""
Implement this function to handle each log line for your model.
This function should mutate self.system_metrics and self.model_metrics
"""
for action, msg in self.log_actions:
match action:
case LogAction.ModelLoaded if msg in log_line:
log.debug(
f"Got log line indicating model is loaded: {log_line}"
)
try:
max_throughput = await self.__run_benchmark()
self.__start_healthcheck.set()
# Wait for the first successful healthcheck before marking model as loaded
if self.healthcheck_url:
log.debug(
"Benchmark succeeded, waiting for healthcheck to confirm model is ready..."
)
try:
await asyncio.wait_for(
self.__healthcheck_ready.wait(), timeout=300.0
)
log.debug(
"Healthcheck confirmed - marking model as loaded"
)
except asyncio.TimeoutError:
raise Exception(
"Timed out waiting for healthcheck after benchmark (waited 300s)"
)
else:
# No healthcheck endpoint defined, wait 10 seconds as fallback
log.debug(
"No healthcheck endpoint defined, waiting 10 seconds before marking model as loaded..."
)
await asyncio.sleep(10)
log.debug("Wait complete - marking model as loaded")
# Wait for pubkey to be fetched before proceeding (unless unsecured)
if not self.unsecured:
if not self.__pubkey_fetch_complete.is_set():
log.debug("Waiting for pubkey to be fetched...")
try:
await asyncio.wait_for(
self.__pubkey_fetch_complete.wait(),
timeout=60.0,
)
except asyncio.TimeoutError:
raise Exception(
"Timed out waiting for pubkey fetch (waited 60s)"
)
log.debug("Pubkey fetch complete")
if self.__pubkey_failed:
raise Exception(
"Cannot mark model as loaded: pubkey fetch failed"
)
# Mark worker ready!
self.metrics._model_loaded(
max_throughput=max_throughput,
)
except Exception as e:
log.debug(f"Benchmark failed with error: {e}")
self.backend_errored(f"Benchmark failed with error: {e}")
case LogAction.ModelError if msg in log_line:
log.debug(f"Got log line indicating error: {log_line}")
self.backend_errored(msg)
break
case LogAction.Info if msg in log_line:
log.debug(f"Info from model logs: {log_line}")
async def tail_log():
"""
Tail a log file with proper rotation handling.
Tracks inode to detect when file is rotated and reopens accordingly.
"""
log.debug(f"tailing file: {self.model_log_file}")
current_inode = None
f = None
try:
while True:
# Check if file exists and get its inode
try:
stat_info = os.stat(self.model_log_file)
file_inode = stat_info.st_ino
except FileNotFoundError:
# File doesn't exist, wait and retry
if f is not None:
await f.aclose()
f = None
current_inode = None
log.debug(
f"Log file {self.model_log_file} not found, waiting..."
)
await asyncio.sleep(1)
continue
# If inode changed (file was rotated) or we don't have a file handle, (re)open
if current_inode != file_inode:
if f is not None:
log.debug(
f"Log file rotation detected (inode changed from {current_inode} to {file_inode}), reopening..."
)
await f.aclose()
else:
log.debug(f"Opening log file (inode: {file_inode})")
f = await open_file(
self.model_log_file, encoding="utf-8", errors="ignore"
)
current_inode = file_inode
# Read a line
line = await f.readline()
if line:
await handle_log_line(line.rstrip())
else:
# No data available, sleep briefly
await asyncio.sleep(LOG_POLL_INTERVAL)
finally:
if f is not None:
await f.aclose()
# Wait for log file to exist before starting to tail
while True:
if os.path.isfile(self.model_log_file):
return await tail_log()
else:
await sleep(1)
async def __session_gc_loop(self) -> NoReturn:
while True:
try:
await sleep(SESSION_GC_INTERVAL)
now = time.time()
async with self._sessions_lock:
expired = [
sid
for sid, s in self.sessions.items()
if s.expiration is not None and s.expiration <= now
]
if expired:
log.debug(f"Removing {len(expired)} expired sessions")
for sid in expired:
await self.__close_session(sid)
except CancelledError:
log.debug("Session GC task cancelled; exiting loop")
raise
except Exception as e:
log.debug(f"Session GC loop error: {e}")
continue
serverless/server/lib/data_types.py
import time
import logging
from dataclasses import dataclass, field
from enum import Enum
from abc import ABC, abstractmethod
from typing import (
Dict,
Any,
Union,
Tuple,
Optional,
Set,
TypeVar,
Generic,
Type,
Callable,
Awaitable,
)
from aiohttp import web, ClientResponse
import inspect
import random
import json
import psutil
import os
import asyncio
"""
type variable representing an incoming payload to pyworker that will used to calculate load and will then
be forwarded to the model
"""
log = logging.getLogger(__file__)
class JsonDataException(Exception):
def __init__(self, json_msg: Dict[str, Any]):
self.message = json_msg
ApiPayload_T = TypeVar("ApiPayload_T", bound="ApiPayload")
@dataclass
class ApiPayload(ABC):
@classmethod
@abstractmethod
def for_test(cls: Type[ApiPayload_T]) -> ApiPayload_T:
"""defines how to create a payload for load testing"""
pass
@abstractmethod
def generate_payload_json(self) -> Dict[str, Any]:
"""defines how to convert an ApiPayload to JSON that will be sent to model API"""
pass
def generate_payload_multipart(self) -> Optional[Dict[str, Any]]:
"""form fields to POST as multipart, or None (default) to POST JSON"""
return None
@abstractmethod
def count_workload(self) -> float:
"""defines how to calculate workload for a payload"""
pass
@classmethod
@abstractmethod
def from_json_msg(
cls: Type[ApiPayload_T], json_msg: Dict[str, Any]
) -> ApiPayload_T:
"""
defines how to create an API payload from a JSON message,
it should throw an JsonDataException if there are issues with some fields
or they are missing in the format of
{
"field": "error msg"
}
"""
pass
@dataclass
class AuthData:
"""data used to authenticate requester"""
cost: str
endpoint: str
reqnum: int
request_idx: int
signature: str
url: str
@classmethod
def from_json_msg(cls, json_msg: Dict[str, Any]):
errors = {}
for param in inspect.signature(cls).parameters:
if param not in json_msg:
errors[param] = "missing parameter"
if errors:
raise JsonDataException(errors)
return cls(
**{
k: v
for k, v in json_msg.items()
if k in inspect.signature(cls).parameters
}
)
@dataclass
class EndpointHandler(ABC, Generic[ApiPayload_T]):
"""
Each model endpoint will have a handler responsible for counting workload from the incoming ApiPayload
and converting it to json to be forwarded to model API
"""
has_benchmark: bool = True
concurrency: int = 10
benchmark_runs: int = 8
allow_parallel_requests: bool = False
max_queue_time: float = None
remote_function: Callable = None
do_warmup_benchmark: bool = True
@property
@abstractmethod
def endpoint(self) -> str:
"""the endpoint on the model API"""
pass
@property
@abstractmethod
def healthcheck_endpoint(self) -> Optional[str]:
"""the endpoint on the model API that is used for healthchecks"""
pass
@classmethod
@abstractmethod
def payload_cls(cls) -> Type[ApiPayload_T]:
"""ApiPayload class"""
pass
@abstractmethod
def make_benchmark_payload(self) -> ApiPayload_T:
"""defines how to create an ApiPayload for benchmarking."""
pass
@abstractmethod
async def generate_client_response(
self, client_request: web.Request, model_response: ClientResponse
) -> Union[web.Response, web.StreamResponse]:
"""
defines how to convert a model API response to a response to PyWorker client
"""
pass
@abstractmethod
async def call_remote_dispatch_function(self, params: dict):
"""
define a remote function for this endpoint, return the result
"""
pass
@classmethod
def get_data_from_request(
cls, req_data: Dict[str, Any]
) -> Tuple[AuthData, ApiPayload_T]:
errors = {}
auth_data: Optional[AuthData] = None
session_id: Optional[str] = None
payload: Optional[ApiPayload_T] = None
try:
if "auth_data" in req_data:
auth_data = AuthData.from_json_msg(req_data["auth_data"])
else:
errors["auth_data"] = "field missing"
except JsonDataException as e:
errors["auth_data"] = e.message
try:
if "payload" in req_data:
payload_cls = cls.payload_cls()
payload = payload_cls.from_json_msg(req_data["payload"])
else:
errors["payload"] = "field missing"
except JsonDataException as e:
errors["payload"] = e.message
if "session_id" in req_data:
session_id = req_data["session_id"]
if errors:
raise JsonDataException(errors)
if auth_data and payload:
return (auth_data, payload, session_id)
else:
raise Exception("error deserializing request data")
@dataclass
class Session:
session_id: str
lifetime: float # extends TTL per-request
auth_data: dict
expiration: float # epoch seconds
on_close_route: str
on_close_payload: dict
requests: list[web.Request] = field(default_factory=list)
created_at: float = field(default_factory=time.time)
request_idx: int = 0 # Request idx associated with this session
session_reqnum: int = 0 # Internal per-session request count
@dataclass
class RequestMetrics:
"""Tracks metrics for an active request."""
request_idx: int
reqnum: int
workload: float
status: str
success: bool = False
is_session: bool = False
session: Session = None
session_reqnum: Optional[int] = None
entered_queue_at: float = 0.0
work_started_at: float = 0.0
work_completed_at: float = 0.0
@dataclass
class SystemMetrics:
"""General system metrics"""
model_loading_start: float
model_loading_time: Union[float, None]
last_disk_usage: float
additional_disk_usage: float
model_is_loaded: bool
disk_usage: float = 0.0
@staticmethod
def get_disk_usage_GB():
return psutil.disk_usage("/").used / (2**30) # want units of GB
@classmethod
def empty(cls):
disk_usage = SystemMetrics.get_disk_usage_GB()
return cls(
model_loading_start=time.time(),
model_loading_time=None,
last_disk_usage=disk_usage,
additional_disk_usage=0.0,
model_is_loaded=False,
disk_usage=disk_usage,
)
def update_disk_usage(self):
"""Refresh the absolute figure; the delta is measured against the last successful report."""
self.disk_usage = SystemMetrics.get_disk_usage_GB()
self.additional_disk_usage = self.disk_usage - self.last_disk_usage
def commit_disk_usage(self):
"""Call once a report carrying the current figures was delivered."""
self.last_disk_usage = self.disk_usage
@dataclass
class BenchmarkResult:
request_idx: int
workload: float
task: Awaitable[ClientResponse]
response: Optional[ClientResponse] = None
@property
def is_successful(self) -> bool:
return self.response is not None and self.response.status == 200
@dataclass
class ModelMetrics:
"""Model specific metrics"""
# these are reset after being sent to autoscaler
workload_served: float
workload_received: float
workload_cancelled: float
workload_errored: float
workload_rejected: float
# these are not
workload_pending: float
error_msg: Optional[str]
max_throughput: float
requests_recieved: Set[int] = field(default_factory=set)
requests_working: dict[int, RequestMetrics] = field(default_factory=dict)
requests_deleting: list[RequestMetrics] = field(default_factory=list)
last_update: float = field(default_factory=time.time)
@classmethod
def empty(cls):
return cls(
workload_pending=0.0,
workload_served=0.0,
workload_cancelled=0.0,
workload_errored=0.0,
workload_rejected=0.0,
workload_received=0.0,
error_msg=None,
max_throughput=0.0,
)
@property
def workload_processing(self) -> float:
return max(self.workload_received - self.workload_cancelled, 0.0)
@property
def wait_time(self) -> float:
if len(self.requests_working) == 0:
return 0.0
return sum(
[
request.workload
for request in self.requests_working.values()
if not request.is_session
]
) / max(self.max_throughput, 0.00001)
@property
def cur_load(self) -> float:
return sum([request.workload for request in self.requests_working.values()])
@property
def working_request_idxs(self) -> list[int]:
return [req.request_idx for req in self.requests_working.values()]
def set_errored(self, error_msg):
self.reset()
self.error_msg = error_msg
def reset(self):
self.workload_served = 0
self.workload_received = 0
self.workload_cancelled = 0
self.workload_errored = 0
self.workload_rejected = 0
self.last_update = time.time()
@dataclass
class WorkerStatusData:
"""Data that is reported to autoscaler's on_status endpoint"""
id: int
mtoken: str
version: str
loadtime: float
cur_load: float
rej_load: float
new_load: float
error_msg: str
max_perf: float
cur_perf: float
cur_capacity: float
max_capacity: float
num_requests_working: int
num_requests_recieved: int
additional_disk_usage: float
disk_usage: float
working_request_idxs: list[int]
url: str
class LogAction(Enum):
"""
These actions tell the backend what a log value means, for example:
actions [
# this marks the model server as loaded
(LogAction.ModelLoaded, "Starting server"),
# these mark the model server as errored
(LogAction.ModelError, "Exception loading model"),
(LogAction.ModelError, "Server failed to bind to port"),
# this tells the backend to print any logs containing the string into its own logs
# which are visible in the vast console instance logs
(LogAction.Info, "Starting model download"),
]
"""
ModelLoaded = 1
ModelError = 2
Info = 3
serverless/server/lib/metrics.py
import os
import time
import logging
import json
from asyncio import sleep
from dataclasses import dataclass, asdict, field
from functools import cache
import asyncio
from aiohttp import ClientSession, ClientTimeout, TCPConnector, ClientResponseError
from .data_types import WorkerStatusData, SystemMetrics, ModelMetrics, RequestMetrics
from typing import Awaitable, NoReturn, List
METRICS_UPDATE_INTERVAL = 1
DELETE_REQUESTS_INTERVAL = 1
log = logging.getLogger(__file__)
@cache
def get_url() -> str:
use_ssl = os.environ.get("USE_SSL", "false") == "true"
worker_port = os.environ[f"VAST_TCP_PORT_{os.environ['WORKER_PORT']}"]
public_ip = os.environ["PUBLIC_IPADDR"]
return f"http{'s' if use_ssl else ''}://{public_ip}:{worker_port}"
@dataclass
class Metrics:
version: str = "0"
mtoken: str = ""
last_metric_update: float = 0.0
last_request_served: float = 0.0
update_pending: bool = False
id: int = field(default_factory=lambda: int(os.environ["CONTAINER_ID"]))
report_addr: List[str] = field(
default_factory=lambda: os.environ["REPORT_ADDR"].split(",")
)
url: str = field(default_factory=get_url)
system_metrics: SystemMetrics = field(default_factory=SystemMetrics.empty)
model_metrics: ModelMetrics = field(default_factory=ModelMetrics.empty)
_session: ClientSession | None = field(default=None, init=False, repr=False)
async def http(self) -> ClientSession:
if self._session is None:
self._session = ClientSession(
timeout=ClientTimeout(total=10),
connector=TCPConnector(
limit=8,
limit_per_host=4,
force_close=True,
enable_cleanup_closed=True,
),
)
return self._session
async def aclose(self) -> None:
if self._session is not None:
await self._session.close()
self._session = None
def _request_id(self, request: RequestMetrics) -> str:
"""Return a formatted identifier for logging based on request type."""
if request.is_session:
return f"Session {request.request_idx}"
elif request.session:
return f"Request {request.session_reqnum} in Session {request.session.request_idx}"
else:
return f"Request {request.request_idx}"
def _request_start(self, request: RequestMetrics) -> None:
"""
this function is called prior to forwarding a request to a model API.
"""
log.debug(f"Starting {self._request_id(request)}")
request.status = "Started"
self.model_metrics.workload_pending += request.workload
self.model_metrics.workload_received += request.workload
if not request.session:
self.model_metrics.requests_recieved.add(request.reqnum)
self.model_metrics.requests_working[request.reqnum] = request
self.update_pending = True
def _request_end(self, request: RequestMetrics) -> None:
"""
this function is called after handling of a request ends, regardless of the outcome
"""
log.debug(f"Ending {self._request_id(request)}")
self.model_metrics.workload_pending -= request.workload
if not request.session:
self.model_metrics.requests_working.pop(request.reqnum, None)
self.model_metrics.requests_deleting.append(request)
self.last_request_served = time.time()
def _request_success(self, request: RequestMetrics) -> None:
"""
this function is called after a response from model API is received and forwarded.
"""
log.debug(f"{self._request_id(request)} succeeded")
request.work_completed_at = time.time()
self.model_metrics.workload_served += request.workload
request.status = "Success"
request.success = True
self.update_pending = True
def _request_errored(self, request: RequestMetrics, message: str) -> None:
"""
this function is called if model API returns an error
"""
log.error(f"{self._request_id(request)} errored: {message}")
request.work_completed_at = time.time()
self.model_metrics.workload_errored += request.workload
request.status = "Error"
request.success = False
self.update_pending = True
def _request_canceled(self, request: RequestMetrics) -> None:
"""
this function is called if client drops connection before model API has responded
"""
log.debug(f"Canceling {self._request_id(request)}")
self.model_metrics.workload_cancelled += request.workload
request.success = True
request.status = "Cancelled"
def _request_reject(self, request: RequestMetrics) -> None:
"""
this function is called if the current wait time for the model is above max_queue_time
"""
log.debug(f"Rejecting {self._request_id(request)}")
self.model_metrics.workload_rejected += request.workload
if not request.session:
self.model_metrics.requests_recieved.add(request.reqnum)
self.model_metrics.requests_deleting.append(request)
request.success = False
request.status = "Rejected"
self.update_pending = True
async def _send_delete_requests_loop(self) -> Awaitable[NoReturn]:
while True:
await sleep(DELETE_REQUESTS_INTERVAL)
if len(self.model_metrics.requests_deleting) > 0:
await self.__send_delete_requests_and_reset()
async def _send_metrics_loop(self) -> Awaitable[NoReturn]:
while True:
await sleep(METRICS_UPDATE_INTERVAL)
elapsed = time.time() - self.last_metric_update
if self.system_metrics.model_is_loaded is False and elapsed >= 10:
await self.__send_metrics_and_reset()
elif self.update_pending or elapsed > 10:
await self.__send_metrics_and_reset()
def _model_loaded(self, max_throughput: float) -> None:
self.system_metrics.model_loading_time = (
time.time() - self.system_metrics.model_loading_start
)
self.system_metrics.model_is_loaded = True
self.model_metrics.max_throughput = max_throughput
def _model_errored(self, error_msg: str) -> None:
self.model_metrics.set_errored(error_msg)
self.system_metrics.model_is_loaded = True
def _set_version(self, version: str) -> None:
self.version = version
def _set_mtoken(self, mtoken: str) -> None:
self.mtoken = mtoken
#######################################Private#######################################
async def __send_delete_requests_and_reset(self):
async def post(report_addr: str, requests: list[dict]) -> bool:
data = {
"worker_id": self.id,
"mtoken": self.mtoken,
"requests": requests,
}
log.debug(
f"Deleting requests: {[r['request_idx'] for r in requests]}"
)
full_path = report_addr.rstrip("/") + "/delete_requests/"
for attempt in range(1, 4):
try:
session = await self.http()
async with session.post(full_path, json=data) as res:
log.debug(f"delete_requests response: {res.status}")
res.raise_for_status()
return True
except asyncio.TimeoutError:
log.debug("delete_requests timed out")
except (ClientResponseError, Exception) as e:
log.debug(f"delete_requests failed with error: {e}")
await asyncio.sleep(2)
log.debug(f"Retrying delete_request, attempt: {attempt}")
return False
# Take a snapshot of what we plan to send this tick.
# New arrivals after this snapshot will remain in the queue for the next tick.
snapshot = list(self.model_metrics.requests_deleting)
if not snapshot:
return # nothing to do
requests_payload = [
{
"request_idx": r.request_idx,
"success": r.success,
"status": r.status,
"entered_queue_at": r.entered_queue_at,
"work_started_at": r.work_started_at,
"work_completed_at": r.work_completed_at,
}
for r in snapshot
]
for report_addr in self.report_addr:
# TODO: Add a Redis subscriber queue for delete_requests
if report_addr == "https://cloud.vast.ai/api/v0":
# Patch: ignore the Redis API report_addr
continue
sent = await post(report_addr, requests_payload)
if sent:
# Remove only the items we actually sent from the live queue.
sent_set = {r.request_idx for r in snapshot}
self.model_metrics.requests_deleting[:] = [
r for r in self.model_metrics.requests_deleting
if r.request_idx not in sent_set
]
break
async def __send_metrics_and_reset(self):
loadtime_snapshot = self.system_metrics.model_loading_time
def get_worker_status_data() -> WorkerStatusData:
return WorkerStatusData(
id=self.id,
mtoken=self.mtoken,
version=self.version,
loadtime=(loadtime_snapshot or 0.0),
new_load=self.model_metrics.workload_processing,
cur_load=self.model_metrics.cur_load,
rej_load=self.model_metrics.workload_rejected,
max_perf=self.model_metrics.max_throughput,
cur_perf=self.model_metrics.workload_served,
error_msg=self.model_metrics.error_msg or "",
num_requests_working=len(self.model_metrics.requests_working),
num_requests_recieved=len(self.model_metrics.requests_recieved),
additional_disk_usage=self.system_metrics.additional_disk_usage,
disk_usage=self.system_metrics.disk_usage,
working_request_idxs=self.model_metrics.working_request_idxs,
cur_capacity=0,
max_capacity=0,
url=self.url,
)
async def send_data(report_addr: str) -> bool:
data = get_worker_status_data()
log_data = asdict(data)
def obfuscate(secret: str) -> str:
if secret is None:
return ""
return secret[:7] + "..." if len(secret) > 12 else ("*" * len(secret))
log_data["mtoken"] = obfuscate(log_data.get("mtoken"))
log.debug(
"\n".join(
[
"\n" + "#" * 20 + " Worker Status " + "#" * 20,
f"{json.dumps(log_data, indent=2)}",
"#" * 55,
]
)
)
full_path = report_addr.rstrip("/") + "/worker_status/"
for attempt in range(1, 4):
try:
session = await self.http()
async with session.post(full_path, json=asdict(data)) as res:
res.raise_for_status()
return True
except asyncio.TimeoutError:
log.debug(f"Request to /worker_status/ timed out")
except (ClientResponseError, Exception) as e:
log.debug(f"Request to /worker_status/ failed with error: {e}")
await asyncio.sleep(2)
log.debug(f"Retrying /worker_status/ request, attempt: {attempt}")
log.debug(f"Failed to send worker status through {report_addr}")
return False
###########
self.system_metrics.update_disk_usage()
sent = False
for report_addr in self.report_addr:
if await send_data(report_addr):
sent = True
break
if sent:
self.update_pending = False
self.model_metrics.reset()
self.system_metrics.commit_disk_usage()
self.last_metric_update = time.time()
serverless/server/lib/server.py
import os
import logging
from typing import List
import ssl
from asyncio import run, gather
import asyncio
from .backend import Backend
from .metrics import Metrics
from aiohttp import web
log = logging.getLogger(__file__)
async def start_server_async(backend: Backend, routes: List[web.RouteDef], **kwargs):
try:
use_ssl = os.environ.get("USE_SSL", "false") == "true"
if use_ssl is True:
log.debug("Getting SSL Certificate from /etc/instance.crt")
try:
ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
ssl_context.load_cert_chain(
certfile="/etc/instance.crt",
keyfile="/etc/instance.key",
)
except Exception as ex:
raise Exception(f"Failed to get SSL Certificate: {ex}")
else:
ssl_context = None
log.debug("Starting Worker Server...")
app = web.Application(client_max_size=0)
app.add_routes(routes)
# Hardcoded session routes
app.router.add_post("/session/create", backend.session_create_handler)
app.router.add_post("/session/end", backend.session_end_handler)
app.router.add_post("/session/get", backend.session_get_handler)
app.router.add_post("/session/health", backend.session_health_handler)
# Pyworker management routes
app.router.add_post("/pyworker/update", backend.pyworker_update_handler)
runner = web.AppRunner(app, handler_cancellation=True)
await runner.setup()
site = web.TCPSite(
runner,
ssl_context=ssl_context,
port=int(os.environ["WORKER_PORT"]),
**kwargs
)
# Create HTTP-only app for internal webhooks to /session/end
http_app = web.Application()
http_app.router.add_post("/session/end", backend.session_end_handler)
http_runner = web.AppRunner(http_app)
await http_runner.setup()
http_port = int(os.environ.get("WORKER_HTTP_PORT", int(os.environ["WORKER_PORT"]) + 1))
http_site = web.TCPSite(
http_runner,
ssl_context=None, # No SSL for internal webhooks
port=http_port,
**kwargs
)
log.debug(f"Starting HTTP-only server for /session/end on port {http_port}")
await site.start()
await http_site.start()
log.debug("HTTP servers started, beginning log tracking")
await backend._start_tracking()
except Exception as e:
err_msg = f"Worker Server failed to launch: {e}"
log.error(err_msg)
async def beacon():
metrics = Metrics()
metrics._set_version(getattr(backend, "version", "0"))
metrics._set_mtoken(getattr(backend, "mtoken", ""))
try:
while True:
metrics._model_errored(err_msg)
await metrics._Metrics__send_metrics_and_reset()
await asyncio.sleep(10)
finally:
await metrics.aclose()
await beacon()
def start_server(backend: Backend, routes: List[web.RouteDef], **kwargs):
run(start_server_async(backend, routes, **kwargs))
serverless/server/worker.py
from vastai.serverless.server.lib import backend, server
from vastai.serverless.server.lib.data_types import (
ApiPayload,
EndpointHandler,
LogAction,
JsonDataException,
)
from dataclasses import dataclass, field
from aiohttp import web, ClientResponse
import logging
import json
import math
import random
import inspect
import logging
from typing import (
Optional,
Dict,
Callable,
Awaitable,
Union,
Any,
Type,
AsyncContextManager,
)
# Callable types
RequestPayloadParser = Callable[[Dict[str, Any]], Dict[str, Any]]
ClientResponseGenerator = Callable[
[web.Request, ClientResponse], Awaitable[Union[web.Response, web.StreamResponse]]
]
WorkloadCalculator = Callable[[Dict[str, Any]], float]
log = logging.getLogger(__name__)
DEFAULT_WORKLOAD = 100.0
@dataclass
class LogActionConfig:
"""Configuration for defining log actions"""
on_load: list[str] = field(default_factory=list)
on_error: list[str] = field(default_factory=list)
on_info: list[str] = field(default_factory=list)
@property
def log_actions(self) -> list[LogAction]:
log_actions_ = []
log_actions_.extend([(LogAction.ModelLoaded, log) for log in self.on_load])
log_actions_.extend([(LogAction.ModelError, log) for log in self.on_error])
log_actions_.extend([(LogAction.Info, log) for log in self.on_info])
return log_actions_
@dataclass
class BenchmarkConfig:
"""Configuration for defining a benchmark"""
dataset: list[dict] | None = None
generator: Callable[[], dict] | None = None # optional sample factory
runs: int = 8
concurrency: int | None = 10
do_warmup: bool = True
@dataclass
class HandlerConfig:
"""Configuration for defining handlers"""
route: str
healthcheck: Optional[str] = None
allow_parallel_requests: bool = False
max_queue_time: Optional[float] = 30.0
benchmark_config: Optional[BenchmarkConfig] = None
handler_class: Optional[Type[EndpointHandler]] = None
payload_class: Optional[Type[ApiPayload]] = None
request_parser: Optional[RequestPayloadParser] = None
response_generator: Optional[ClientResponseGenerator] = None
workload_calculator: Optional[WorkloadCalculator] = None
is_remote_dispatch: bool = False
remote_function: Optional[Callable] = None
@dataclass
class WorkerConfig:
model_server_url: str = None
model_server_port: int = None
model_log_file: str = None
benchmark_data: list[dict[str, Any]] = field(default_factory=list)
handlers: list[HandlerConfig] = field(default_factory=list)
model_healthcheck_url: str = None
benchmark_route: Optional[str] = None
log_action_config: LogActionConfig = field(default_factory=LogActionConfig)
max_sessions: Optional[int] = 10
lifecycle: Optional[AsyncContextManager] = None
class EndpointHandlerFactory:
"""Factory for creating endpoint handlers from WorkerConfig"""
def __init__(self, config: WorkerConfig):
self.config = config
self._handlers: Dict[str, EndpointHandler] = {}
self._build_handlers()
def _build_handlers(self) -> None:
"""Build endpoint handlers from config"""
# If no handlers, create a default handler at /
if not self.config.handlers:
default_handler_config = HandlerConfig(
route="/", healthcheck=self.config.model_healthcheck_url
)
handler = self._create_handler(default_handler_config)
self._handlers["/"] = handler
else:
# Build handlers from config
for handler_config in self.config.handlers:
# TODO: Override values from handler class with HandlerConfig values
if handler_config.handler_class is not None:
# Use custom handler class
handler = handler_config.handler_class()
self._handlers[handler_config.route] = handler
continue
handler = self._create_handler(handler_config)
self._handlers[handler_config.route] = handler
def _create_handler(self, handler_config: HandlerConfig) -> EndpointHandler:
"""Create a generic endpoint handler from HandlerConfig"""
# Extract config values with defaults
route_path = handler_config.route
healthcheck_path = handler_config.healthcheck
benchmark_config = handler_config.benchmark_config
user_payload_class = handler_config.payload_class
user_request_parser = handler_config.request_parser
user_response_generator = handler_config.response_generator
user_workload_calculator = handler_config.workload_calculator
# If user provided a custom payload class, use it
if user_payload_class:
PayloadClass = user_payload_class
else:
# Create a generic ApiPayload class
@dataclass
class GenericApiPayload(ApiPayload):
input: Dict[str, Any] = field(default_factory=dict)
@classmethod
def for_test(cls) -> "GenericApiPayload":
# Use random.choice from benchmark_data if available
if benchmark_config:
if benchmark_config.dataset:
test_data = random.choice(benchmark_config.dataset)
elif benchmark_config.generator:
try:
test_data = benchmark_config.generator()
except Exception as ex:
raise (
Exception(
f'Error generating benchmark data for "{handler_config.route}" handler: {ex}'
)
)
else:
raise (
Exception(
f'Error generating benchmark data for "{handler_config.route}" handler: Missing BenchmarkConfig!'
)
)
return cls(input=test_data.copy())
def generate_payload_json(self) -> Dict[str, Any]:
return self.input
@classmethod
def from_dict(cls, input: Dict[str, Any]) -> "GenericApiPayload":
return cls(input=input)
def count_workload(self) -> float:
if not user_workload_calculator:
return DEFAULT_WORKLOAD
# Fall back rather than let a bad calculator 500 requests or
# poison the autoscaler's load sums.
try:
workload = float(user_workload_calculator(self.input))
except Exception:
log.warning(
f'workload_calculator for "{route_path}" raised; '
f"falling back to {DEFAULT_WORKLOAD}",
exc_info=True,
)
return DEFAULT_WORKLOAD
if not math.isfinite(workload) or workload < 0:
log.warning(
f'workload_calculator for "{route_path}" returned '
f"{workload!r}; falling back to {DEFAULT_WORKLOAD}"
)
return DEFAULT_WORKLOAD
return workload
@classmethod
def from_json_msg(cls, json_msg: Dict[str, Any]) -> "GenericApiPayload":
if not isinstance(json_msg, dict):
raise JsonDataException(
{"data": "payload must be a dictionary"}
)
# Apply user's request parser if provided
if user_request_parser:
try:
json_msg = user_request_parser(json_msg)
except JsonDataException:
raise # a bad request (422), not a server error
except Exception as e:
raise Exception(f"Error in user response handler: {e}")
try:
return cls.from_dict(json_msg)
except (json.JSONDecodeError, JsonDataException) as e:
raise JsonDataException(f"Error in user response handler: {e}")
PayloadClass = GenericApiPayload
# Create a generic EndpointHandler class
@dataclass
class GenericEndpointHandler(EndpointHandler[PayloadClass]):
_route: str = field(default=route_path)
has_benchmark: bool = field(
default=(True if handler_config.benchmark_config else False)
)
allow_parallel_requests: bool = field(
default=handler_config.allow_parallel_requests
)
max_queue_time: float = field(default=handler_config.max_queue_time)
benchmark_runs: int = field(
default=(
handler_config.benchmark_config.runs
if handler_config.benchmark_config
else 8
)
)
concurrency: int = field(
default=(
handler_config.benchmark_config.concurrency
if handler_config.benchmark_config
and handler_config.benchmark_config.concurrency
else 10
)
)
do_warmup_benchmark: bool = field(
default=(
handler_config.benchmark_config.do_warmup
if handler_config.benchmark_config
and handler_config.benchmark_config.concurrency
else True
)
)
remote_function: Callable = field(
default=(
handler_config.remote_function
if handler_config.remote_function is not None
else None
)
)
@property
def endpoint(self) -> str:
"""The endpoint is the same as the route"""
return self._route
async def call_remote_dispatch_function(self, params: dict):
"""
define a remote function for this endpoint, return the result
"""
if self.remote_function is None:
raise RuntimeError(
f"remote_function is not configured for route {self._route}"
)
try:
return await self.remote_function(**params)
except Exception as ex:
raise RuntimeError(
f"Error calling remote function for route {self._route}: {ex}"
) from ex
@property
def healthcheck_endpoint(self) -> Optional[str]:
return None
@classmethod
def payload_cls(cls) -> Type[PayloadClass]:
return PayloadClass
def make_benchmark_payload(self) -> PayloadClass:
"""Just call the payload class's for_test() method"""
return PayloadClass.for_test()
async def generate_client_response(
self,
client_request: web.Request,
model_response: ClientResponse,
) -> Union[web.Response, web.StreamResponse]:
# User implemented override
if user_response_generator:
try:
return await user_response_generator(
client_request, model_response
)
except Exception as e:
raise Exception(f"Error in user response generator: {e}")
# Reasonable default: return response, handle streaming if present
# Detect streaming
content_type = model_response.content_type or ""
is_stream = (
content_type.startswith("text/event-stream")
or content_type in ("application/x-ndjson", "application/jsonl")
or "stream" in content_type.lower()
) or (model_response.headers.get("Transfer-Encoding") == "chunked")
if is_stream:
# Streaming passthrough
res = web.StreamResponse(
status=model_response.status,
)
if content_type:
res.content_type = content_type
await res.prepare(client_request)
async for chunk in model_response.content.iter_any():
if not chunk:
continue
await res.write(chunk)
await res.write_eof()
return res
else:
# Non-streaming
body = await model_response.read()
headers = model_response.headers.copy()
headers.pop("Content-Type", None)
return web.Response(
body=body,
status=model_response.status,
content_type=content_type or None,
headers=headers,
)
return GenericEndpointHandler()
def get_handler(self, route: str) -> Optional[EndpointHandler]:
"""Get handler for a specific route"""
return self._handlers.get(route)
def get_all_handlers(self) -> Dict[str, EndpointHandler]:
"""Get all registered handlers"""
return self._handlers.copy()
def get_benchmark_handler(self) -> Optional[EndpointHandler]:
"""
Get the benchmark handler.
Raises errors if there are too many / too little BenchmarkConfigs
Should be exactly one EndpointHandler with a BenchmarkConfig
"""
if not self._handlers:
return None
benchmark_handler: EndpointHandler = None
for handler in self._handlers.values():
if handler.has_benchmark:
if benchmark_handler is not None:
raise Exception(
"Cannot define BenchmarkConfig for more than one EndpointHandler!"
)
else:
benchmark_handler = handler
if not benchmark_handler:
raise Exception("Missing EndpointHandler with BenchmarkConfig")
return benchmark_handler
def has_handlers(self) -> bool:
"""Check if any handlers are registered"""
return len(self._handlers) > 0
@property
def model_server_base_url(self) -> str:
"""Get the full model server base URL"""
return f"{self.config.model_server_url}:{self.config.model_server_port}"
class Worker:
"""
This class provides a simple to use abstraction over the pyworker backend.
All custom implementations of pyworker can be created by configuring a Worker object.
The pyworker starts by calling Worker.run()
"""
def __init__(self, config: WorkerConfig):
# Configure logging for the pyworker internals.
# Force-override any prior basicConfig (e.g. from vastai CLI code) by
# setting the root logger level and adding a handler if none exist.
root_logger = logging.getLogger()
root_logger.setLevel(logging.DEBUG)
if not root_logger.handlers:
handler = logging.StreamHandler()
handler.setFormatter(
logging.Formatter(
fmt="%(asctime)s[%(levelname)-5s] %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
)
root_logger.addHandler(handler)
else:
for handler in root_logger.handlers:
handler.setLevel(logging.DEBUG)
handler_factory = EndpointHandlerFactory(config)
# Get all endpoint handlers
handlers = handler_factory.get_all_handlers()
benchmark_handler = handler_factory.get_benchmark_handler()
# Create backend
self.backend = backend.Backend(
model_server_url=f"{config.model_server_url}:{config.model_server_port}",
model_log_file=config.model_log_file,
benchmark_handler=benchmark_handler,
log_actions=config.log_action_config.log_actions,
healthcheck_url=config.model_healthcheck_url,
max_sessions=config.max_sessions,
lifecycle=config.lifecycle,
)
# Attach endpoint handlers to HTTP routes
self.routes = []
for route_path, handler in handlers.items():
self.routes.append(
web.post(route_path, self.backend.create_handler(handler))
)
async def run_async(self, **kwargs):
await server.start_server_async(self.backend, self.routes, **kwargs)
def run(self, **kwargs):
server.start_server(self.backend, self.routes, **kwargs)
SKILL.md
---
name: vastai
description: Vast.ai CLI to manage GPU instances, volumes, serverless endpoints, and billing.
allowed-tools: Bash(vastai:*)
compatibility: Linux, macOS
metadata:
author: vast-ai
---
# vastai
Manage GPU instances, templates, volumes, serverless endpoints, SSH keys, and billing on Vast.ai.
> Command is `vastai` (lowercase). Always use `--raw` for machine-readable JSON output.
## Install
```bash
# Linux/macOS (recommended, no Python required)
curl -fsSL https://vast.ai/install.sh | bash
# or via PyPI (Windows, or if you'd rather use pip)
pip install vastai
```
## Quick start
```bash
vastai set api-key <YOUR_API_KEY> # Authenticate (one-time); Create API Key in account at https://console.vast.ai/manage-keys/
vastai show user # Verify auth + check balance
vastai create ssh-key ~/.ssh/id_ed25519.pub # Register SSH key (do BEFORE create)
vastai search offers 'gpu_name=RTX_4090 num_gpus=1 verified=true direct_port_count>=1 rentable=true' -o 'dlperf_usd-'
# Note the offer ID from the output
vastai create instance <OFFER_ID> --image pytorch/pytorch:@vastai-automatic-tag --disk 20 --ssh --direct
# Automatically grab appropriate image tag; Response: {"success": true, "new_contract": <INSTANCE_ID>}
vastai show instance <INSTANCE_ID> # Poll until actual_status == "running" (see Instance status values below)
vastai ssh-url <INSTANCE_ID> # Get SSH connection string
vastai copy local:./data/ <INSTANCE_ID>:/workspace/ # Upload files
vastai destroy instance <INSTANCE_ID> -y # Clean up (stops all billing; -y skips confirmation)
```
API key: https://console.vast.ai/manage-keys/
## Global flags
Available on every command:
```
--api-key KEY Override stored API key
--raw Output machine-readable JSON (agents should always use this)
--full Print full results (don't page with less)
--explain Show underlying API calls (useful for debugging)
--curl Show equivalent curl command
--no-color Disable colored output
--url URL Override server REST API URL
--retry RETRY Set retry limit for API calls
--version Show CLI version
```
## Query syntax
Search commands accept filter expressions. Operators: `=`, `!=`, `>`, `>=`, `<`, `<=`, `in`, `notin`.
```bash
# Examples
'gpu_name=RTX_4090 num_gpus=1' # Exact match + numeric
'gpu_ram>=48 reliability>0.95' # Greater-than filters
'geolocation=EU dph_total<=2.0' # Region + price cap
```
Common filter fields: `num_gpus`, `gpu_name`, `gpu_ram`, `cpu_ram`, `disk_space`, `reliability`, `compute_cap`, `inet_up`, `inet_down`, `dph_total`, `geolocation`, `direct_port_count`, `verified`, `rentable`
Common sort fields: `score` (default — overall value), `dlperf_usd` (DL perf per dollar), `dph_total` (price), `num_gpus`, `reliability`
## Commands
### Instances
```bash
vastai show instances # List all your instances (fetches every page, no prompts)
vastai show instances --status running loading # Filter by status
vastai show instances --gpu-name 'RTX 4090' # Filter by GPU
vastai show instances --label training # Filter by label
vastai show instances --order-by start_date desc # Sort by column
vastai show instances --cols id,status,gpu,dph # Custom columns
vastai show instances --raw # Flat JSON list of all instances (scripting)
vastai show instance <id> # Poll single instance (use for status checks)
vastai create instance <offer-id> --image pytorch/pytorch:2.4.0-cuda12.4-cudnn9-runtime --disk 20 --ssh --direct
# Response includes "new_contract": <id> — that is your instance ID
vastai launch instance --gpu-name RTX_4090 --num-gpus 1 --image pytorch/pytorch
vastai start instance <id> # Start stopped instance
vastai stop instance <id> # Stop (preserves disk, no GPU charges)
vastai reboot instance <id> # Stop + start
vastai destroy instance <id> -y # Delete permanently (irreversible; -y required for non-interactive use)
vastai destroy instances <id1> <id2> -y # Batch delete (-y skips confirmation prompt)
vastai label instance <id> --label "training-run-1" # Tag instance
vastai update instance <id> # Recreate from updated template
vastai prepay instance <id> # Deposit credits into reserved instance
vastai recycle instance <id> # Destroy + recreate
```
**Recommended Vast.ai images** (use `@vastai-automatic-tag` to get the right tag for the machine automatically; browse pre-configured models at https://vast.ai/model-library):
<!-- Note: @vastai-automatic-tag is resolved server-side — CLI passes it unchanged -->
```bash
vastai/base-image:@vastai-automatic-tag # Minimal Ubuntu base
vastai/pytorch:@vastai-automatic-tag # PyTorch + CUDA
vastai/linux-desktop:@vastai-automatic-tag # Linux desktop (VNC/RDP)
# vLLM — set model via env vars with huggingface model example
vastai create instance <id> --image vastai/vllm:@vastai-automatic-tag --disk 40 --ssh --direct \
--env '-e MODEL_NAME=Qwen/Qwen2.5-3B-Instruct -e HF_TOKEN=hf_xxx'
# ComfyUI — set model checkpoint via env vars with huggingface model example
vastai create instance <id> --image vastai/comfy:@vastai-automatic-tag --disk 40 --ssh --direct \
--env '-e CHECKPOINT_MODEL=black-forest-labs/FLUX.1-schnell -e HF_TOKEN=hf_xxx'
```
**create instance flags:**
- `--image IMAGE` — Docker image
- `--disk DISK` — Local disk in GB
- `--ssh` / `--jupyter` — Connection type
- `--direct` — Faster direct connections
- `--label LABEL` — Instance label
- `--env ENV` — Env vars and port mappings, e.g. `'-e TZ=UTC -p 8080:8080'`
- `--onstart FILE` — Path to a startup script file
- `--onstart-cmd CMD` — Startup script as inline string (for longer scripts use `--onstart` or gzip+base64 encode)
- `--bid_price PRICE` — Interruptible (spot) pricing in $/hr
- `--template_hash HASH` — Create from template
- `--create-volume ID` — Attach new volume
- `--link-volume ID` — Attach existing volume
- `--cancel-unavail` — Fail if no machine available (vs. create stopped)
**Instance status values:**
| `actual_status` | Meaning |
|-----------------|---------|
| `null` | Provisioning |
| `created` | Instance created, not yet provisioned |
| `loading` | Image downloading / container starting |
| `running` | Active — GPU charges apply |
| `stopped` | Halted — disk charges only |
| `frozen` | Paused with memory — GPU charges apply |
| `exited` | Container process exited unexpectedly |
| `rebooting` | Restarting (transient) |
| `unknown` | No recent heartbeat from host |
| `offline` | Host disconnected from Vast servers |
> **Poll loop warning:** If `actual_status` becomes `exited`, `unknown`, or `offline` it will never reach `running`. Always add a timeout and error branch — otherwise your script loops forever while disk charges accrue. Destroy and retry with a different offer.
> **Charges:** Storage charges begin at creation. GPU charges begin when status reaches `running`.
### Interruptible (spot) rentals
Interruptible (spot) instances are priced below on-demand instances, but can be interrupted at any time by another user with a lower bid. Note: `search offers --type bid` exposes `min_bid`, but `create instance` defaults to **on-demand at `dph_total`** unless you pass `--bid_price <floor>`. Always pass `--bid_price` after a `--type bid` search, otherwise the instance will be rented as an on-demand instance/price instead of as an interruptible.
When outbid, the instance moves to `stopped` (not destroyed) and storage charges continue. Resume by raising the bid via `update instance --bid_price`.
### Search
```bash
vastai search offers # Default: verified, on-demand, sorted by score
vastai search offers 'gpu_name=RTX_4090 num_gpus=1 verified=true direct_port_count>=1' -o 'dlperf_usd-'
vastai search offers 'num_gpus>=4 reliability>0.99' -o 'num_gpus-'
vastai search offers --type bid # Interruptible (spot) pricing
vastai search offers --type reserved # Reserved pricing
vastai search offers -n 'gpu_name=H100_SXM' # No default filters
vastai search volumes # Search volume offers
vastai search templates "pytorch" # Search templates
vastai search benchmarks # Search benchmarks
vastai search invoices # Search invoice history
```
**search offers flags:** `--type on-demand|reserved|bid`, `--order/-o FIELD[-]`, `--limit`, `--storage GB`, `--no-default/-n`
### SSH & File Transfer
```bash
vastai ssh-url <id> # Get ssh:// connection URL
vastai scp-url <id> # Get scp:// URL
vastai attach ssh <id> "ssh-ed25519 AAAA..." # Attach SSH key to instance
vastai detach ssh <id> <ssh_key_id> # Remove SSH key (ssh_key_id from show ssh-keys)
vastai show ssh-keys # List account SSH keys
vastai create ssh-key ~/.ssh/id_ed25519.pub # Add SSH key from file (do BEFORE create instance)
vastai create ssh-key # Generate new key if you don't have one
vastai create ssh-key "ssh-ed25519 AAAA..." # Add SSH key inline
vastai delete ssh-key <id> # Remove SSH key from account
vastai update ssh-key <id> "ssh-ed25519 AAAA..." # Update SSH key value
```
**Note:** `ssh-url` returns a connection string — it does not open an interactive session. Use `$(vastai ssh-url <id>)` to extract the URL, or parse `--raw` JSON output.
### File Copy
```bash
vastai copy <src> <dst> # Copy between instance and local
vastai copy local:./data/ <id>:/workspace/data/ # Local → instance (preferred format)
vastai copy <id>:/workspace/results/ local:./results/ # Instance → local
vastai copy <id-a>:/workspace/ <id-b>:/workspace/ # Instance → instance
vastai copy s3.101:/data/ C.<id>:/workspace/ # Cloud service → instance (s3.<connection-id>)
vastai copy V.1234:/file C.<id>:/workspace/ # Volume → instance
vastai copy V.1234:/file s3.101:/workspace/ # Volume → cloud service
# Legacy format also works: vastai copy 12345:./data ./local-data
vastai cloud copy --src ./data --dst s3://bucket/path \
--instance 12345 --connection <conn-id> \
--transfer "Instance To Cloud" # To cloud storage (add connection in UI on settings page)
vastai cancel copy <dst-id> # Cancel in-progress copy
```
**cloud copy flags:** `--src`, `--dst`, `--instance`, `--connection`, `--transfer`
**Location formats:** `[instance_id:]path` (legacy), `C.instance_id:path`, `V.volume_id:path`, `cloud_service[.id]:path`, `local:path`. Volume copy is currently only supported for copying to other volumes, instances, or cloud services, not local. Never copy to `/root` or `/` — it breaks ssh permissions on the instance and future copies fail.
### Logs & Exec
```bash
vastai logs <id> # Container logs (last 1000 lines)
vastai logs <id> --tail 100 # Last 100 lines
vastai logs <id> --filter "error" # Grep filter
vastai execute <id> "nvidia-smi" # Run command on instance
vastai execute <id> "ls /workspace" --schedule DAILY # Scheduled execution
```
### Volumes
```bash
vastai search volumes # Search available volume offers
vastai show volumes # List your volumes
vastai create volume <offer_id> [-s SIZE] [-n NAME] # Create volume (offer_id from search volumes)
vastai clone volume <source_id> <dest_id> [-s SIZE] # Clone volume (dest_id from search volumes)
vastai delete volume <id> # Delete volume
vastai create network-volume ... # Create network volume
vastai list network-volume # List network volumes
vastai take snapshot <instance_id> --repo REPO --docker_login_user USER --docker_login_pass PASS # Snapshot instance to volume
```
### Serverless & Deployments
```bash
vastai show endpoints # List endpoint groups
vastai create endpoint --name "my-ep" ... # Create endpoint
vastai update endpoint <id> ... # Update endpoint
vastai delete endpoint <id> # Delete endpoint
vastai get endpt-logs <id> # Endpoint logs
vastai show workergroups # List worker groups
vastai create workergroup --name "wg" ... # Create worker group
vastai update workergroup <id> ... # Update worker group
vastai update workers <id> # Rolling update of workers
vastai delete workergroup <id> # Delete worker group
vastai get wrkgrp-logs <id> # Worker group logs
vastai show deployments # List deployments
vastai show deployment <id> # Deployment details
vastai show deployment-versions <id> # Version history
vastai delete deployment <id> # Delete deployment
vastai show scheduled-jobs # List scheduled jobs
vastai delete scheduled-job <id> # Delete scheduled job
```
### Templates
```bash
vastai search templates "pytorch" # Search templates
vastai create template --name "x" --image "img" # Create template
vastai update template <hash_id> ... # Update template (edits create a new version)
vastai delete template <id> # Delete template
```
### Account & API Keys
```bash
vastai set api-key <key> # Save API key locally
vastai show api-key <id> # Show a specific key
vastai show api-keys # List all your API keys
vastai create api-key --name "ci" --permissions '{...}' # Create restricted key
vastai delete api-key <id> # Delete key
vastai reset api-key # Reset main key (get new from console)
vastai show user # Account info, credit balance
vastai show audit-logs # Account action history
vastai show connections # Cloud storage connections
vastai show ipaddrs # IP address history
```
### Billing
```bash
vastai show invoices-v1 # Charges/invoices (paginated)
vastai show invoices-v1 --charges # Charges only
vastai show invoices-v1 --invoices # Invoices only
vastai show invoices-v1 --start-date 2026-01-01 --end-date 2026-02-01
vastai show invoices-v1 --limit 50 --latest-first
vastai show deposit <id> # Reserved instance deposit info
```
### Teams
```bash
vastai create team --team-name "myteam" # Create team; your key stays on your personal account
vastai show members # List team members
vastai invite member --email user@example.com # Invite member
vastai remove member <id> # Remove member
vastai create team-role --name "viewer" ... # Create role
vastai show team-role <id> # Role details
vastai show team-roles # List roles
vastai update team-role <id> ... # Update role
vastai remove team-role <id> # Remove role
vastai destroy team # Delete team
```
### Environment Variables
```bash
vastai show env-vars # List user env vars
vastai create env-var KEY val # Create env var
vastai update env-var KEY newval # Update env var
vastai delete env-var KEY # Delete env var
```
### Machine Management (hosts)
```bash
vastai show machines # List your machines
vastai list machine <id> # List/unlist machine on marketplace
vastai show machine <id> # Machine details
vastai cleanup machine <id> # Clean up machine state
vastai schedule maint <id> ... # Schedule maintenance window
vastai cancel maint <id> # Cancel scheduled maintenance
vastai show maints # List maintenance windows
vastai unlist machine <id> # Remove from marketplace
```
## Common errors
| Error | Cause | Fix |
|-------|-------|-----|
| `401 Unauthorized` | Invalid or expired API key | `vastai set api-key <new-key>` |
| `Insufficient credits` | Account balance too low | Add credits at https://cloud.vast.ai/billing/ |
| `No offers found` | Filters too restrictive | Relax filters, try `--no-default/-n` |
| `Permission denied` | SSH key not attached | `vastai create ssh-key` before `create instance` |
| `Connection refused` | Instance not yet running | Poll `show instance <id>` until `actual_status == "running"` |
| Hangs on `destroy instance` | Confirmation prompt waiting for input | Add `-y` flag: `vastai destroy instance <id> -y` |
## URLs
### Console
```
https://console.vast.ai/instances/ # Your instances
https://console.vast.ai/create/ # Search GPU offers
https://console.vast.ai/manage-keys/ # Create and manage API keys
https://cloud.vast.ai/billing/ # Billing
```
### API
```
https://console.vast.ai/api/v0/instances/ # Instances endpoint
https://console.vast.ai/api/v0/asks/ # Offers search
```
### Instance Ports
Access a port exposed on your instance:
```
ssh://root@<ssh_host>:<ssh_port> # SSH (from vastai ssh-url)
```
Direct connections (when `--direct` used at creation): `<direct_port_end>` field in instance JSON.
## Documentation
Fetch the complete documentation index at: https://docs.vast.ai/llms.txt
sync/__init__.py
from .client import SyncClient
from .results import SyncOffer, SyncInstance
sync/client.py
import json
import requests
from typing import Any, Optional, Union
from vastai._base import _BaseClient, _APIKEY_SENTINEL
from vastai.data.query import Query
from vastai.data.offer import Offer
from vastai.data.instance import Instance, InstanceConfig, CreateInstanceResponse
from vastai.sync.results import SyncOffer, SyncInstance
class SyncClient(_BaseClient):
"""
Synchronous Vast.ai API client.
api_key resolution order:
1. Explicit api_key argument
2. VAST_API_KEY environment variable
3. $XDG_CONFIG_HOME/vastai/vast_api_key
4. ~/.vast_api_key (legacy)
"""
def __init__(self, api_key: object = _APIKEY_SENTINEL, vast_server: str = "https://console.vast.ai"):
super().__init__(api_key, vast_server)
def search(
self,
query: Query,
type_: Optional[str] = None,
order: Optional[list] = [["score", "desc"]],
limit: Optional[int] = None,
allocated_storage: Optional[float] = None,
target_reliability: Optional[float] = None,
external: Optional[bool] = None,
extra_ids: Optional[list] = None,
has_avx: Optional[int] = None,
gpu_option: Optional[Any] = None,
show_incompatible: Optional[Any] = None,
sort_option: Optional[Any] = None,
template_id: Optional[Any] = None,
) -> list[SyncOffer]:
"""Search available offers and return a list of SyncOffer result wrappers."""
payload = dict(query.query)
for key, value in [
("type", type_),
("order", order),
("limit", limit),
("allocated_storage", allocated_storage),
("target_reliability", target_reliability),
("external", external),
("extra_ids", extra_ids),
("has_avx", has_avx),
("gpu_option", gpu_option),
("show_incompatible", show_incompatible),
("sort_option", sort_option),
("template_id", template_id),
]:
if value is not None:
payload[key] = value
response = requests.post(
self._url("/api/v0/bundles/"),
json=payload,
headers=self._headers(),
)
response.raise_for_status()
data = response.json()
return [SyncOffer(Offer.from_dict(o), self) for o in data.get("offers", [])]
def create_instance(
self,
offer_id: int,
config: InstanceConfig,
) -> SyncInstance:
"""Accept an offer and create an instance. Returns a SyncInstance wrapper."""
payload = config.to_dict()
response = requests.put(
self._url(f"/api/v0/asks/{offer_id}/"),
json=payload,
headers=self._headers(),
)
response.raise_for_status()
data = response.json()
if not data.get("success"):
raise Exception(f"Failed to create instance: {data.get('error', '')} - {data.get('msg', '')}")
resp = CreateInstanceResponse.from_dict(data)
# Fetch the newly created instance
instances = self.show_instances()
for inst in instances:
if inst.id == resp.new_contract:
return inst
# Fallback: return a minimal instance
return SyncInstance(Instance.from_dict({"id": resp.new_contract}), self)
def show_instances(self) -> list[SyncInstance]:
"""Return all instances owned by the authenticated user.
The legacy v0 ``/instances`` list endpoint is deprecated; this pages
through the v1 ``/api/v1/instances/`` endpoint (omitting ``select_cols``
so the backend returns full instance rows) and accumulates every page.
"""
instances = []
params = {
"select_filters": json.dumps({}),
"order_by": json.dumps([{"col": "id", "dir": "asc"}]),
"limit": 25,
}
while True:
response = requests.get(
self._url("/api/v1/instances/"),
params=params,
headers=self._headers(),
)
response.raise_for_status()
data = response.json()
instances.extend(data.get("instances") or [])
next_token = data.get("next_token")
if not next_token:
break
params["after_token"] = next_token
return [SyncInstance(Instance.from_dict(i), self) for i in instances]
def destroy_instance(self, instance_or_id: Union[SyncInstance, int]) -> None:
"""Destroy a running or stopped instance."""
instance_id = instance_or_id if isinstance(instance_or_id, int) else instance_or_id.id
response = requests.delete(
self._url(f"/api/v0/instances/{instance_id}/"),
headers=self._headers(),
)
response.raise_for_status()
sync/results.py
from __future__ import annotations
from typing import TYPE_CHECKING, Optional, Union
from vastai.data.offer import Offer
from vastai.data.instance import InstanceConfig
if TYPE_CHECKING:
from vastai.sync.client import SyncClient
from vastai.data.instance import Instance
class SyncOffer:
def __init__(self, data: Offer, client: SyncClient):
self._data = data
self._client = client
def create_instance(self, config: InstanceConfig) -> SyncInstance:
"""Create an instance from this offer. Synonym for SyncClient.create_instance(self.id, config)."""
return self._client.create_instance(self._data.id, config)
def __getattr__(self, name: str):
return getattr(self._data, name)
def __repr__(self) -> str:
return f"SyncOffer({self._data!r})"
class SyncInstance:
def __init__(self, data: Instance, client: SyncClient):
self._data = data
self._client = client
def destroy(self) -> None:
"""Synonym for SyncClient.destroy_instance(self)."""
return self._client.destroy_instance(self._data.id)
def __getattr__(self, name: str):
return getattr(self._data, name)
def __repr__(self) -> str:
return f"SyncInstance({self._data!r})"
utils.py
"""Shared utility functions for the Vast.ai SDK and CLI.
Generic helpers that are not specific to any single layer (CLI, API, SDK).
"""
import re
import os
import sys
import time
import math
import argparse
import subprocess
import importlib.metadata
from datetime import datetime, timedelta, timezone
from concurrent.futures import ThreadPoolExecutor
# ---------------------------------------------------------------------------
# Version helpers
# ---------------------------------------------------------------------------
def parse_version(version: str) -> tuple:
parts = version.split(".")
if len(parts) < 3:
print(f"Invalid version format: {version}", file=sys.stderr)
return tuple(int(part) for part in parts)
def _get_git_version():
try:
result = subprocess.run(
["git", "describe", "--tags", "--abbrev=0"],
capture_output=True,
text=True,
check=True,
cwd=os.path.dirname(__file__),
)
tag = result.stdout.strip()
return tag[1:] if tag.startswith("v") else tag
except Exception:
return "0.0.0"
def _get_local_version():
try:
return importlib.metadata.version("vastai")
except Exception:
return _get_git_version()
VERSION = _get_local_version()
# ---------------------------------------------------------------------------
# Validation helpers
# ---------------------------------------------------------------------------
def validate_seconds(value):
"""Validate that the input value is a valid number for seconds between yesterday and Jan 1, 2100."""
try:
val = int(value)
# Calculate min_seconds as the start of yesterday in seconds
yesterday = datetime.now() - timedelta(days=1)
min_seconds = int(yesterday.timestamp())
# Calculate max_seconds for Jan 1st, 2100 in seconds
max_date = datetime(2100, 1, 1, 0, 0, 0)
max_seconds = int(max_date.timestamp())
if not (min_seconds <= val <= max_seconds):
raise argparse.ArgumentTypeError(f"{value} is not a valid second timestamp.")
return val
except ValueError:
raise argparse.ArgumentTypeError(f"{value} is not a valid integer.")
# ---------------------------------------------------------------------------
# VRL (Vast Resource Locator) parsing
# ---------------------------------------------------------------------------
class VRLException(Exception):
pass
def parse_vast_url(url_str):
"""
Breaks up a vast-style url in the form instance_id:path and does
some basic sanity type-checking.
:param url_str:
:return:
"""
instance_id = None
path = url_str
#print(f'url_str: {url_str}')
if (":" in url_str):
url_parts = url_str.split(":", 2)
if len(url_parts) == 2:
(instance_id, path) = url_parts
else:
raise VRLException("Invalid VRL (Vast resource locator).")
else:
try:
instance_id = int(path)
path = "/"
except:
pass
valid_unix_path_regex = re.compile('^(/)?([^/\0]+(/)?)+$')
# Got this regex from https://stackoverflow.com/questions/537772/what-is-the-most-correct-regular-expression-for-a-unix-file-path
if (path != "/") and (valid_unix_path_regex.match(path) is None):
raise VRLException(f"Path component: {path} of VRL is not a valid Unix style path.")
#print(f'instance_id: {instance_id}')
#print(f'path: {path}')
return (instance_id, path)
# ---------------------------------------------------------------------------
# String / environment parsing
# ---------------------------------------------------------------------------
def smart_split(s, char):
in_double_quotes = False
in_single_quotes = False # note that isn't designed to work with nested quotes within the env
parts = []
current = []
for c in s:
if c == char and not (in_double_quotes or in_single_quotes):
parts.append(''.join(current))
current = []
elif c == '\'':
in_single_quotes = not in_single_quotes
current.append(c)
elif c == '\"':
in_double_quotes = not in_double_quotes
current.append(c)
else:
current.append(c)
parts.append(''.join(current)) # add last part
return parts
def parse_env(envs):
result = {}
if (envs is None):
return result
env = smart_split(envs, ' ')
prev = None
for e in env:
if (prev is None):
if (e in {"-e", "-p", "-h", "-v", "-n"}):
prev = e
else:
pass
else:
if (prev == "-p"):
if set(e).issubset(set("0123456789:tcp/udp")):
result["-p " + e] = "1"
else:
pass
elif (prev == "-e"):
kv = e.split('=')
if len(kv) >= 2:
val = kv[1]
if len(kv) > 2:
val = '='.join(kv[1:])
result[kv[0]] = val.strip("'\"")
else:
pass
elif (prev == "-v"):
if (set(e).issubset(set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789:./_"))):
result["-v " + e] = "1"
elif (prev == "-n"):
if (set(e).issubset(set("abcdefghijklmnopqrstuvwxyz0123456789-"))):
result["-n " + e] = "1"
else:
result[prev] = e
prev = None
return result
# ---------------------------------------------------------------------------
# List / threading helpers
# ---------------------------------------------------------------------------
def split_list(lst, k):
"""
Splits a list into sublists of maximum size k.
"""
return [lst[i:i + k] for i in range(0, len(lst), k)]
def exec_with_threads(f, args, nt=16, max_retries=5):
def worker(sub_args):
for arg in sub_args:
retries = 0
while retries <= max_retries:
try:
result = None
if isinstance(arg, tuple):
result = f(*arg)
else:
result = f(arg)
if result: # Assuming a truthy return value means success
break
except Exception as e:
print(str(e))
pass
retries += 1
stime = 0.25 * 1.3 ** retries
print(f"retrying in {stime}s")
time.sleep(stime) # Exponential backoff
# Split args into nt sublists
args_per_thread = math.ceil(len(args) / nt)
sublists = [args[i:i + args_per_thread] for i in range(0, len(args), args_per_thread)]
with ThreadPoolExecutor(max_workers=nt) as executor:
executor.map(worker, sublists)
# ---------------------------------------------------------------------------
# Date / scheduling helpers
# ---------------------------------------------------------------------------
def convert_timestamp_to_date(unix_timestamp):
utc_datetime = datetime.fromtimestamp(unix_timestamp, tz=timezone.utc)
return utc_datetime.strftime("%Y-%m-%d")
def parse_day_cron_style(value):
"""
Accepts an integer string 0-6 or '*' to indicate 'Every day'.
Returns 0-6 as int, or None if '*'.
"""
val = str(value).strip()
if val == "*":
return None
try:
day = int(val)
if 0 <= day <= 6:
return day
except ValueError:
pass
raise argparse.ArgumentTypeError("Day must be 0-6 (0=Sunday) or '*' for every day.")
def parse_hour_cron_style(value):
"""
Accepts an integer string 0-23 or '*' to indicate 'Every hour'.
Returns 0-23 as int, or None if '*'.
"""
val = str(value).strip()
if val == "*":
return None
try:
hour = int(val)
if 0 <= hour <= 23:
return hour
except ValueError:
pass
raise argparse.ArgumentTypeError("Hour must be 0-23 or '*' for every hour.")
# ---------------------------------------------------------------------------
# Georegion mapping and query/result hooks
# ---------------------------------------------------------------------------
_regions = {
'AF': ('DZ,AO,BJ,BW,BF,BI,CM,CV,CF,TD,KM,CD,CG,DJ,EG,GQ,ER,ET,GA,GM,GH,GN,'
'GW,KE,LS,LR,LY,MW,MA,ML,MR,MU,MZ,NA,NE,NG,RW,SH,ST,SN,SC,SL,SO,ZA,'
'SS,SD,SZ,TZ,TG,TN,UG,YE,ZM,ZW'), # Africa
'AS': ('AE,AM,AR,AU,AZ,BD,BH,BN,BT,MM,KH,KP,IN,ID,IR,IQ,IL,JP,JO,KZ,LV,'
'LI,MY,MV,MN,NP,KR,PK,PH,QA,SA,SG,LK,SY,TW,TJ,TH,TR,TM,VN,YE,HK,'
'CN,OM'), # Asia
'EU': ('AL,AD,AT,BY,BE,BA,BG,HR,CY,CZ,DK,EE,'
'FI,FR,GE,DE,GR,HU,IS,IT,KZ,LV,LI,LT,'
'LU,MT,MD,MC,ME,NL,NO,PL,PT,RO,RU,RS,'
'SK,SI,ES,SE,CH,UA,GB,VA,MK'), # Europe
'LC': ('AG,AR,BS,BB,BZ,BO,BR,CL,CO,CR,CU,DO,EC,SV,GY,HT,HN,JM,MX,NI,PA,PY,'
'PE,PR,RD,SUR,TT,UR,VZ'), # Latin America and the Caribbean
'NA': 'CA,US', # Northern America
'OC': ('AU,FJ,GU,KI,MH,FM,NR,NZ,PG,PW,SL,TO,TV,VU'), # Oceania
}
def _reverse_mapping(regions):
reversed_mapping = {}
for region, countries in regions.items():
for country in countries.split(','):
reversed_mapping[country] = region
return reversed_mapping
_regions_rev = _reverse_mapping(_regions)
def preprocess_search_query(query_str):
"""Pre-process a query string to expand georegion and chunked directives.
If the query contains ``georegion=true``, any ``geolocation = <REGION_CODE>``
clause is expanded into ``geolocation in [<country_list>]`` using :data:`_regions`.
Both ``georegion=true`` and ``chunked=true`` are stripped from the query
before it reaches the API.
Returns:
A ``(georegion_active, chunked_active, processed_query_str)`` tuple.
"""
if query_str is None:
return False, False, query_str
from pyparsing import Word, alphas, alphanums, one_of, Group, ZeroOrMore
key = Word(alphas + "_", alphanums + "_")
operator = one_of("= in != > < >= <=")
value = Word(alphanums + "_")
expr = Group(key + operator + value)
query = ZeroOrMore(expr)
parsed = query.parse_string(query_str)
directives = {'georegion', 'chunked'}
state = {}
for d in directives:
state[d] = any([d, '=', 'true'] == list(e) for e in parsed)
if not any(state.values()):
return False, False, query_str
parts = []
for e in parsed:
if e[0] in directives:
continue
elif e[0] == 'geolocation' and state['georegion'] and e[2] in _regions:
parts.append(f'geolocation in [{_regions[e[2]]}]')
else:
parts.append(' '.join(e))
return state['georegion'], state['chunked'], ' '.join(parts)
def postprocess_search_results(results, georegion_active=False, chunked=False):
"""Post-process search results with datacenter flag, georegion, and chunked filters.
Always adds a ``datacenter`` boolean derived from ``hosting_type``.
When *georegion_active*, appends continent codes to geolocation.
When *chunked*, filters and rounds values for skypilot compatibility.
"""
_chunked_cutoffs = {
'cpu_ram': 64 * 1024,
'cpu_cores': 32,
'min_bid': 0,
}
filtered = []
for res in results:
res['datacenter'] = (res.get('hosting_type') == 1)
if georegion_active:
geo = res.get('geolocation', '')
if geo:
country = geo[-2:]
region = _regions_rev.get(country)
if region:
res['geolocation'] = f'{geo}, {region}'
if chunked:
good = True
try:
for k, v in _chunked_cutoffs.items():
if res.get(k) is not None and res[k] < v:
good = False
else:
res[k] = v
except Exception:
good = False
if not good:
continue
res['gpu_ram'] = res.get('gpu_ram', 0) & 0xffffffffff0
res['disk_space'] = int(res.get('disk_space', 0)) & 0xffffffffffc0
filtered.append(res)
return filtered