vastai — Skillopedia

vastai Manage GPU instances, templates, volumes, serverless endpoints, SSH keys, and billing on Vast.ai. Command is (lowercase). Always use for machine-readable JSON output. Install Quick start API key: https://console.vast.ai/manage-keys/ Global flags Available on every command: Query syntax Search commands accept filter expressions. Operators: , , , , , , , . Common filter fields: , , , , , , , , , , , , , Common sort fields: (default — overall value), (DL perf per dollar), (price), , Commands Instances Recommended Vast.ai images (use to get the right tag for the machine automatically; brow…

)\n\n\ndef resolve_runtype(runtype: Optional[str] = None, *, ssh: bool = False,\n jupyter: bool = False, direct: bool = False,\n args=None, jupyter_lab: bool = False,\n jupyter_dir: Optional[str] = None,\n default: Optional[str] = None) -> tuple:\n \"\"\"Translate the friendly connection flags into a runtype string.\n\n Returns ``(runtype, args)``. ``args`` comes back because an empty ``--args``\n selects the args runtype and then has to be dropped from the payload.\n\n ``default`` is what an unflagged call resolves to: the CLI passes ``'ssh'``,\n the SDK passes ``None`` so the key stays out of the payload entirely.\n \"\"\"\n if args in ('', [''], []):\n args = None\n empty_args = True\n else:\n empty_args = False\n\n if runtype:\n if ssh or jupyter or direct:\n raise ValueError(\n \"Pass either runtype= or the ssh/jupyter/direct flags, not both.\"\n )\n return runtype, args\n\n resolved = default\n if args or empty_args:\n resolved = 'args'\n\n if not jupyter and (jupyter_dir or jupyter_lab):\n jupyter = True\n\n if jupyter and resolved == 'args':\n raise ValueError(\n \"Can't use jupyter and args together. \"\n \"Try onstart_cmd instead of args.\"\n )\n\n if jupyter:\n resolved = 'jupyter_direc ssh_direc ssh_proxy' if direct else 'jupyter_proxy ssh_proxy'\n elif ssh:\n resolved = 'ssh_direc ssh_proxy' if direct else 'ssh_proxy'\n elif direct and default is None:\n # The CLI reaches here with default='ssh' and keeps it, which is the\n # long-standing behaviour of a bare --direct. There is no such fallback\n # for the SDK, so say so instead of dropping the flag.\n raise ValueError(\"direct=True requires ssh=True or jupyter=True.\")\n\n return resolved, args\n\n\ndef build_volume_info(*, create_volume: Optional[int] = None,\n link_volume: Optional[int] = None,\n volume_size: Optional[float] = None,\n mount_path: Optional[str] = None,\n volume_label: Optional[str] = None) -> Optional[dict]:\n \"\"\"Build the volume_info blob, or None when no volume was requested.\"\"\"\n if volume_size is not None and not create_volume:\n raise ValueError(\"volume_size can only be used with create_volume.\")\n if not (create_volume or link_volume):\n if mount_path or volume_label:\n raise ValueError(\"mount_path and volume_label need create_volume or link_volume.\")\n return None\n if not mount_path:\n raise ValueError(\"mount_path is required when creating or linking a volume.\")\n if not VALID_MOUNT_PATH.match(mount_path):\n raise ValueError(f\"mount_path '{mount_path}' is not a valid Linux file path.\")\n\n volume_info = {\n \"mount_path\": mount_path,\n \"create_new\": True if create_volume else False,\n \"volume_id\": create_volume if create_volume else link_volume,\n }\n if volume_label:\n volume_info[\"name\"] = volume_label\n if volume_size:\n volume_info[\"size\"] = volume_size\n elif create_volume:\n volume_info[\"size\"] = 15\n return volume_info\n\n\ndef apply_portal_config(env: dict, runtype: Optional[str]) -> dict:\n \"\"\"Return env with jupyter portal entries stripped on non-jupyter runtypes.\"\"\"\n if runtype and 'jupyter' in runtype:\n return env\n if \"PORTAL_CONFIG\" not in env:\n return env\n\n filtered = [c for c in env[\"PORTAL_CONFIG\"].split(\"|\") if 'jupyter' not in c.lower()]\n if not filtered:\n raise ValueError(\n \"env variable PORTAL_CONFIG must contain at least one non-jupyter \"\n \"related config string if runtype is not jupyter\"\n )\n env = dict(env)\n env[\"PORTAL_CONFIG\"] = \"|\".join(filtered)\n return env\n\n\ndef build_create_instance_payload(\n *, image: Optional[str] = None, disk: float = 10, env=None,\n price: Optional[float] = None, bid_price: Optional[float] = None,\n label: Optional[str] = None, extra: Optional[str] = None,\n onstart_cmd: Optional[str] = None, login: Optional[str] = None,\n python_utf8: bool = False, lang_utf8: bool = False,\n jupyter_lab: bool = False, jupyter_dir: Optional[str] = None,\n force: bool = False, cancel_unavail: bool = False,\n template_hash: Optional[str] = None, user: Optional[str] = None,\n runtype: Optional[str] = None, ssh: bool = False, jupyter: bool = False,\n direct: bool = False, args=None, volume_info: Optional[dict] = None,\n create_volume: Optional[int] = None, link_volume: Optional[int] = None,\n volume_size: Optional[float] = None, mount_path: Optional[str] = None,\n volume_label: Optional[str] = None,\n) -> dict:\n \"\"\"Build the JSON body for a create-instance request.\n\n The one place the CLI, the SDK and the docs generator agree on what a\n create-instance call takes. Pure: no client, no argparse, no printing.\n \"\"\"\n if price is not None and bid_price is not None:\n raise ValueError(\"Pass either price= or bid_price=, not both.\")\n if price is None:\n price = bid_price\n\n if isinstance(env, str):\n env = parse_env(env)\n env = dict(env or {})\n\n if template_hash is None:\n runtype, args = resolve_runtype(\n runtype, ssh=ssh, jupyter=jupyter, direct=direct, args=args,\n jupyter_lab=jupyter_lab, jupyter_dir=jupyter_dir,\n )\n env = apply_portal_config(env, runtype)\n\n if volume_info is None:\n volume_info = build_volume_info(\n create_volume=create_volume, link_volume=link_volume,\n volume_size=volume_size, mount_path=mount_path,\n volume_label=volume_label,\n )\n elif create_volume or link_volume:\n raise ValueError(\"Pass either volume_info= or the create/link volume flags, not both.\")\n\n json_blob = {\n \"client_id\": \"me\",\n \"image\": image,\n \"env\": env,\n \"price\": price,\n \"disk\": disk,\n \"label\": label,\n \"extra\": extra,\n \"onstart\": onstart_cmd,\n \"image_login\": login,\n \"python_utf8\": python_utf8,\n \"lang_utf8\": lang_utf8,\n \"use_jupyter_lab\": jupyter_lab,\n \"jupyter_dir\": jupyter_dir,\n \"force\": force,\n \"cancel_unavail\": cancel_unavail,\n \"template_hash_id\": template_hash,\n \"user\": user,\n }\n if runtype:\n json_blob[\"runtype\"] = runtype\n if args is not None:\n json_blob[\"args\"] = args\n if volume_info:\n json_blob[\"volume_info\"] = volume_info\n return json_blob\n\n\ndef create_instance(client: VastClient, id, image=None, disk=10, env=None,\n price=None, label=None, extra=None, onstart_cmd=None,\n login=None, python_utf8=False, lang_utf8=False,\n jupyter_lab=False, jupyter_dir=None, force=False,\n cancel_unavail=False, template_hash=None, user=None,\n runtype=None, args=None, volume_info=None,\n # Appended, never inserted: the parameters above keep their\n # original positions so existing positional calls still bind\n # the same values.\n bid_price=None, ssh=False, jupyter=False, direct=False,\n create_volume=None, link_volume=None, volume_size=None,\n mount_path=None, volume_label=None) -> dict:\n \"\"\"Create one instance, or several when ``id`` is a list of offer ids.\"\"\"\n json_blob = build_create_instance_payload(\n image=image, disk=disk, env=env, price=price, bid_price=bid_price,\n label=label, extra=extra, onstart_cmd=onstart_cmd, login=login,\n python_utf8=python_utf8, lang_utf8=lang_utf8, jupyter_lab=jupyter_lab,\n jupyter_dir=jupyter_dir, force=force, cancel_unavail=cancel_unavail,\n template_hash=template_hash, user=user, runtype=runtype, ssh=ssh,\n jupyter=jupyter, direct=direct, args=args, volume_info=volume_info,\n create_volume=create_volume, link_volume=link_volume,\n volume_size=volume_size, mount_path=mount_path, volume_label=volume_label,\n )\n\n return create_instance_from_payload(client, id, json_blob)\n\n\ndef create_instance_from_payload(client: VastClient, id, json_blob: dict) -> dict:\n \"\"\"Send an already-built create-instance payload for one or many offers.\"\"\"\n if isinstance(id, list):\n json_blob[\"ids\"] = id\n r = client.post(\"/asks/bulk/\", json_data=json_blob)\n else:\n r = client.put(f\"/asks/{id}/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef destroy_instance(client: VastClient, id) -> dict:\n json_blob = {}\n if isinstance(id, list):\n json_blob[\"instance_ids\"] = id\n r = client.delete(\"/instances/\", json_data=json_blob)\n else:\n r = client.delete(f\"/instances/{id}/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef start_instance(client: VastClient, id) -> dict:\n json_blob = {\"state\": \"running\"}\n if isinstance(id, list):\n json_blob[\"ids\"] = id\n r = client.put(\"/instances/\", json_data=json_blob)\n else:\n r = client.put(f\"/instances/{id}/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef stop_instance(client: VastClient, id) -> dict:\n json_blob = {\"state\": \"stopped\"}\n if isinstance(id, list):\n json_blob[\"ids\"] = id\n r = client.put(\"/instances/\", json_data=json_blob)\n else:\n r = client.put(f\"/instances/{id}/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef reboot_instance(client: VastClient, id: int) -> dict:\n r = client.put(f\"/instances/reboot/{id}/\", json_data={})\n r.raise_for_status()\n return r.json()\n\n\ndef recycle_instance(client: VastClient, id: int) -> dict:\n r = client.put(f\"/instances/recycle/{id}/\", json_data={})\n r.raise_for_status()\n return r.json()\n\n\ndef label_instance(client: VastClient, id: int, label: str) -> dict:\n r = client.put(f\"/instances/{id}/\", json_data={\"label\": label})\n r.raise_for_status()\n return r.json()\n\n\ndef prepay_instance(client: VastClient, id: int, amount: float) -> dict:\n r = client.put(f\"/instances/prepay/{id}/\", json_data={\"amount\": amount})\n r.raise_for_status()\n return r.json()\n\n\ndef change_bid(client: VastClient, id: int, price: float = None) -> dict:\n r = client.put(f\"/instances/bid_price/{id}/\", json_data={\"client_id\": \"me\", \"price\": price})\n r.raise_for_status()\n return r.json()\n\n\ndef accept_price_increase(client: VastClient, id: int = None,\n instance_ids=None, host_id: int = None) -> dict:\n \"\"\"Deprecated shim that resolves to the per-row endpoint.\n\n Prefer :func:`vastai.api.price_increase.accept`. The old\n ``id`` / ``instance_ids`` / ``host_id`` selectors are no longer\n supported by the backend; ``id`` here is reinterpreted as an\n instance id, looked up against the pending list, and forwarded\n to the per-row accept. ``instance_ids`` and ``host_id`` raise.\n Removed in the release after this one.\n \"\"\"\n from vastai.api import price_increase as _pi\n if instance_ids is not None or host_id is not None:\n raise TypeError(\n \"accept_price_increase: instance_ids and host_id are no longer \"\n \"supported. Use vastai.api.price_increase.accept(client, pending_id) \"\n \"per row, or vastai.sdk.VastAI.accept_price_increase(instance_id=…).\")\n if id is None:\n raise TypeError(\"accept_price_increase: instance id is required\")\n try:\n match = _pi.resolve_instance_to_pending(client, id)\n except LookupError as err:\n raise LookupError(\n f\"accept_price_increase: no pending price increase for instance {id}\"\n ) from err\n return _pi.accept(client, match[\"pending_price_increase_id\"])\n\n\ndef execute(client: VastClient, id: int, command: str):\n \"\"\"Execute a command on an instance and return the output.\"\"\"\n r = client.put(f\"/instances/command/{id}/\", json_data={\"command\": command})\n r.raise_for_status()\n rj = r.json()\n result_url = rj.get(\"result_url\")\n if not result_url:\n return rj\n return _poll_result_url(result_url)\n\n\ndef logs(client: VastClient, instance_id: int, tail=None, filter=None, daemon_logs=False):\n \"\"\"Request logs for an instance and return the log text.\"\"\"\n json_blob = {}\n if filter:\n json_blob['filter'] = filter\n if tail:\n json_blob['tail'] = tail\n if daemon_logs:\n json_blob['daemon_logs'] = 'true'\n r = client.put(f\"/instances/request_logs/{instance_id}/\", json_data=json_blob)\n r.raise_for_status()\n rj = r.json()\n result_url = rj.get(\"result_url\")\n if not result_url:\n return rj\n return _poll_result_url(result_url)\n\n\ndef update_instance(client: VastClient, id: int, template_id=None, template_hash_id=None,\n image=None, args=None, env=None, onstart=None) -> dict:\n json_blob = {\"id\": id}\n if template_id is not None:\n json_blob[\"template_id\"] = template_id\n if template_hash_id is not None:\n json_blob[\"template_hash_id\"] = template_hash_id\n if image is not None:\n json_blob[\"image\"] = image\n if args is not None:\n json_blob[\"args\"] = args\n if env is not None:\n json_blob[\"env\"] = env\n if onstart is not None:\n json_blob[\"onstart\"] = onstart\n r = client.put(f\"/instances/update_template/{id}/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef take_snapshot(client: VastClient, instance_id, repo=None, container_registry=\"docker.io\",\n docker_login_user=None, docker_login_pass=None, pause=\"true\") -> dict:\n req_json = {\n \"id\": instance_id,\n \"container_registry\": container_registry,\n \"personal_repo\": repo,\n \"docker_login_user\": docker_login_user,\n \"docker_login_pass\": docker_login_pass,\n \"pause\": pause\n }\n r = client.post(f\"/instances/take_snapshot/{instance_id}/\", json_data=req_json)\n r.raise_for_status()\n return r.json()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":17462,"content_sha256":"0fe3dcddd971e38f5d48cf5d20a8704c9a4e5afbdb6888e7846042ea332fca77"},{"filename":"api/keys.py","content":"\"\"\"SSH key and API key operations.\"\"\"\nfrom vastai.api.client import VastClient\n\n\ndef create_ssh_key(client: VastClient, ssh_key: str) -> dict:\n \"\"\"Add an SSH public key to the account.\n\n Args:\n client: VastClient instance.\n ssh_key: SSH public key content string.\n\n Returns:\n Response dict with created key info.\n \"\"\"\n r = client.post(\"/ssh/\", json_data={\"ssh_key\": ssh_key})\n r.raise_for_status()\n return r.json()\n\n\ndef show_ssh_keys(client: VastClient) -> dict:\n \"\"\"List SSH keys associated with the account.\n\n Args:\n client: VastClient instance.\n\n Returns:\n Response dict with SSH key info.\n \"\"\"\n r = client.get(\"/ssh/\")\n r.raise_for_status()\n return r.json()\n\n\ndef update_ssh_key(client: VastClient, id: int, ssh_key: str) -> dict:\n \"\"\"Update an existing SSH key.\n\n Args:\n client: VastClient instance.\n id: SSH key ID.\n ssh_key: New SSH public key content string.\n\n Returns:\n Response dict.\n \"\"\"\n payload = {\n \"id\": id,\n \"ssh_key\": ssh_key,\n }\n r = client.put(f\"/ssh/{id}/\", json_data=payload)\n r.raise_for_status()\n return r.json()\n\n\ndef delete_ssh_key(client: VastClient, id: int) -> dict:\n \"\"\"Delete an SSH key from the account.\n\n Args:\n client: VastClient instance.\n id: SSH key ID to delete.\n\n Returns:\n Response dict.\n \"\"\"\n r = client.delete(f\"/ssh/{id}/\")\n r.raise_for_status()\n return r.json()\n\n\ndef attach_ssh(client: VastClient, instance_id: int, ssh_key: str) -> dict:\n \"\"\"Attach an SSH key to an instance.\n\n Args:\n client: VastClient instance.\n instance_id: Instance ID to attach the key to.\n ssh_key: SSH public key content string.\n\n Returns:\n Response dict.\n \"\"\"\n r = client.post(f\"/instances/{instance_id}/ssh/\", json_data={\"ssh_key\": ssh_key})\n r.raise_for_status()\n return r.json()\n\n\ndef detach_ssh(client: VastClient, instance_id: int, ssh_key_id: str) -> dict:\n \"\"\"Detach an SSH key from an instance.\n\n Args:\n client: VastClient instance.\n instance_id: Instance ID.\n ssh_key_id: SSH key ID to detach.\n\n Returns:\n Response dict.\n \"\"\"\n r = client.delete(f\"/instances/{instance_id}/ssh/{ssh_key_id}/\")\n r.raise_for_status()\n return r.json()\n\n\ndef create_api_key(client: VastClient, name: str, permissions: dict = None,\n key_params: str = None) -> dict:\n \"\"\"Create a new API key.\n\n Args:\n client: VastClient instance.\n name: Name for the new API key.\n permissions: Dict of permissions for the key.\n key_params: Optional key parameters string.\n\n Returns:\n Response dict with created API key info.\n \"\"\"\n json_blob = {\"name\": name}\n if permissions is not None:\n json_blob[\"permissions\"] = permissions\n if key_params is not None:\n json_blob[\"key_params\"] = key_params\n r = client.post(\"/auth/apikeys/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef show_api_key(client: VastClient, id: int) -> dict:\n \"\"\"Show details of a specific API key.\n\n Args:\n client: VastClient instance.\n id: API key ID.\n\n Returns:\n API key details dict.\n \"\"\"\n r = client.get(f\"/auth/apikeys/{id}/\")\n r.raise_for_status()\n return r.json()\n\n\ndef show_api_keys(client: VastClient) -> dict:\n \"\"\"List all API keys associated with the account.\n\n Args:\n client: VastClient instance.\n\n Returns:\n Envelope dict of the form ``{\"apikeys\": [...]}`` (what the backend\n sends). The high-level ``VastAI.show_api_keys`` wrapper unwraps this\n to a plain list.\n \"\"\"\n r = client.get(\"/auth/apikeys/\")\n r.raise_for_status()\n return r.json()\n\n\ndef delete_api_key(client: VastClient, id: int) -> dict:\n \"\"\"Delete an API key.\n\n Args:\n client: VastClient instance.\n id: API key ID to delete.\n\n Returns:\n Response dict.\n \"\"\"\n r = client.delete(f\"/auth/apikeys/{id}/\")\n r.raise_for_status()\n return r.json()\n\n\ndef reset_api_key(client: VastClient) -> dict:\n \"\"\"Reset the current API key (generates a new one).\n\n Args:\n client: VastClient instance.\n\n Returns:\n Response dict.\n \"\"\"\n json_blob = {\"client_id\": \"me\"}\n r = client.put(\"/commands/reset_apikey/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":4454,"content_sha256":"1f10251cefc98f92ab2a14024983e5aa6a35c0e4892492601dd98e7b0cfde546"},{"filename":"api/machines.py","content":"\"\"\"Machine management operations for hosts.\"\"\"\nfrom vastai.api.client import VastClient\n\n\ndef show_machine(client: VastClient, id: int) -> list:\n \"\"\"Show a single hosted machine.\n\n Args:\n client: VastClient instance.\n id: Machine ID.\n\n Returns:\n List of machine data dicts (API returns list even for single machine).\n \"\"\"\n r = client.get(f\"/machines/{id}\", query_args={\"owner\": \"me\"})\n r.raise_for_status()\n return r.json()\n\n\ndef show_machines(client: VastClient) -> list:\n \"\"\"Show all hosted machines for the current user.\n\n Args:\n client: VastClient instance.\n\n Returns:\n List of machine dicts.\n \"\"\"\n r = client.get(\"/machines\", query_args={\"owner\": \"me\"})\n r.raise_for_status()\n return r.json()[\"machines\"]\n\n\ndef show_maints(client: VastClient, machine_ids: list) -> list:\n \"\"\"Show maintenance information for host machines.\n\n Args:\n client: VastClient instance.\n machine_ids: List of machine ID integers.\n\n Returns:\n List of maintenance info dicts.\n \"\"\"\n r = client.get(\"/machines/maintenances\", query_args={\"owner\": \"me\", \"machine_ids\": machine_ids})\n r.raise_for_status()\n return r.json()\n\n\ndef list_machine(client: VastClient, id: int, price_gpu: float = None,\n price_disk: float = None, price_inetu: float = None,\n price_inetd: float = None, price_min_bid: float = None,\n min_chunk: int = None, end_date: float = None,\n discount_rate: float = None, duration: str = None,\n vol_size: int = None, vol_price: float = None) -> dict:\n \"\"\"List a machine for rent (create offers).\n\n Args:\n client: VastClient instance.\n id: Machine ID.\n price_gpu: Per GPU rental price in $/hour.\n price_disk: Storage price in $/GB/month.\n price_inetu: Price for upload bandwidth in $/GB.\n price_inetd: Price for download bandwidth in $/GB.\n price_min_bid: Per GPU minimum bid price floor in $/hour.\n min_chunk: Minimum number of GPUs.\n end_date: Contract offer expiration as unix epoch timestamp.\n discount_rate: Max long-term prepay discount rate fraction.\n duration: Duration string or seconds.\n vol_size: Volume contract offer size in GB.\n vol_price: Volume disk price.\n\n Returns:\n Response dict.\n \"\"\"\n json_blob = {\n \"machine\": id,\n \"price_gpu\": price_gpu,\n \"price_disk\": price_disk,\n \"price_inetu\": price_inetu,\n \"price_inetd\": price_inetd,\n \"price_min_bid\": price_min_bid,\n \"min_chunk\": min_chunk,\n \"end_date\": end_date,\n \"credit_discount_max\": discount_rate,\n \"duration\": duration,\n \"vol_size\": vol_size,\n \"vol_price\": vol_price,\n }\n r = client.put(\"/machines/create_asks/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef unlist_machine(client: VastClient, id: int) -> dict:\n \"\"\"Unlist a listed machine (remove all offers).\n\n Args:\n client: VastClient instance.\n id: Machine ID.\n\n Returns:\n Response dict.\n \"\"\"\n r = client.delete(f\"/machines/{id}/asks/\")\n r.raise_for_status()\n return r.json()\n\n\ndef cancel_maint(client: VastClient, id: int) -> dict:\n \"\"\"Cancel scheduled maintenance window(s) for a machine.\n\n Args:\n client: VastClient instance.\n id: Machine ID.\n\n Returns:\n Response dict.\n \"\"\"\n json_blob = {\"client_id\": \"me\", \"machine_id\": id}\n r = client.put(f\"/machines/{id}/cancel_maint/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef schedule_maint(client: VastClient, id: int, sdate: float = None,\n duration: float = None, maintenance_category: str = None) -> dict:\n \"\"\"Schedule a maintenance window for a machine.\n\n Args:\n client: VastClient instance.\n id: Machine ID.\n sdate: Start date as unix epoch timestamp.\n duration: Duration in hours.\n maintenance_category: Category of maintenance.\n\n Returns:\n Response dict.\n \"\"\"\n json_blob = {\n \"client_id\": \"me\",\n \"sdate\": sdate,\n \"duration\": duration,\n \"maintenance_category\": maintenance_category,\n }\n r = client.put(f\"/machines/{id}/dnotify/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef cleanup_machine(client: VastClient, id: int) -> dict:\n \"\"\"Remove all expired storage instances from a machine.\n\n Args:\n client: VastClient instance.\n id: Machine ID.\n\n Returns:\n Response dict.\n \"\"\"\n r = client.put(f\"/machines/{id}/cleanup/\", json_data={})\n r.raise_for_status()\n return r.json()\n\n\ndef defrag_machines(client: VastClient, machine_ids: list) -> dict:\n \"\"\"Defragment machines to rearrange GPU assignments.\n\n Args:\n client: VastClient instance.\n machine_ids: List of machine ID integers.\n\n Returns:\n Response dict with defragmentation results.\n \"\"\"\n json_blob = {\"machine_ids\": machine_ids}\n r = client.put(\"/machines/defrag_offers/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef delete_machine(client: VastClient, id: int) -> dict:\n \"\"\"Force delete a machine if not in use by clients.\n\n Args:\n client: VastClient instance.\n id: Machine ID.\n\n Returns:\n Response dict.\n \"\"\"\n r = client.post(f\"/machines/{id}/force_delete/\")\n r.raise_for_status()\n return r.json()\n\n\ndef reports(client: VastClient, id: int) -> dict:\n \"\"\"Get user reports for a given machine.\n\n Args:\n client: VastClient instance.\n id: Machine ID.\n\n Returns:\n Reports dict.\n \"\"\"\n json_blob = {\"machine_id\": id}\n r = client.get(f\"/machines/{id}/reports/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef set_defjob(client: VastClient, id: int, price_gpu: float = None,\n price_inetu: float = None, price_inetd: float = None,\n image: str = None, args: list = None) -> dict:\n \"\"\"Create default jobs for a machine.\n\n Args:\n client: VastClient instance.\n id: Machine ID.\n price_gpu: Per GPU rental price in $/hour.\n price_inetu: Price for upload bandwidth in $/GB.\n price_inetd: Price for download bandwidth in $/GB.\n image: Docker container image to launch.\n args: List of arguments passed to container launch.\n\n Returns:\n Response dict.\n \"\"\"\n json_blob = {\n \"machine\": id,\n \"price_gpu\": price_gpu,\n \"price_inetu\": price_inetu,\n \"price_inetd\": price_inetd,\n \"image\": image,\n \"args\": args,\n }\n r = client.put(\"/machines/create_bids/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef remove_defjob(client: VastClient, id: int) -> dict:\n \"\"\"Remove default jobs from a machine.\n\n Args:\n client: VastClient instance.\n id: Machine ID.\n\n Returns:\n Response dict.\n \"\"\"\n r = client.delete(f\"/machines/{id}/defjob/\")\n r.raise_for_status()\n return r.json()\n\n\ndef set_min_bid(client: VastClient, id: int, price: float = None) -> dict:\n \"\"\"Set the minimum bid/rental price for a machine.\n\n Args:\n client: VastClient instance.\n id: Machine ID.\n price: Per GPU min bid price in $/hour.\n\n Returns:\n Response dict.\n \"\"\"\n json_blob = {\"client_id\": \"me\", \"price\": price}\n r = client.put(f\"/machines/{id}/minbid/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\n","content_type":"text/x-python; charset=utf-8","language":"python","size":7613,"content_sha256":"6724080bf611369580e1328a17867383d1b3e72c2f459efcccb038e7109c2b92"},{"filename":"api/metrics.py","content":"\"\"\"Platform-wide GPU market metrics (available to admins and hosts).\"\"\"\nfrom typing import Optional\n\nfrom vastai.api.client import VastClient\n\n\ndef gpu_current(client: VastClient, verified: str = \"all\", hosting_type: str = \"all\",\n num_gpus: str = \"all\") -> dict:\n \"\"\"Current snapshot of GPU supply/demand, pricing, and perf per GPU type.\n\n Args:\n client: VastClient instance.\n verified: \"yes\", \"no\", or \"all\".\n hosting_type: \"all\", \"secure_cloud\", or \"community\".\n num_gpus: machine GPU-count bucket — \"all\", \"1\", \"2\", \"4\", or \"8\".\n\n Returns:\n Parsed JSON response. Shape: {\"success\": True, \"gpus\": [...]} or\n {\"success\": True, \"gpus\": [], \"needs_machine\": True} for hosts with no machines.\n \"\"\"\n r = client.get(\"/metrics/gpu/current/\", query_args={\"verified\": verified, \"hosting_type\": hosting_type, \"num_gpus\": num_gpus})\n r.raise_for_status()\n return r.json()\n\n\ndef gpu_history(client: VastClient, gpu_name: str, verified: str = \"all\", hosting_type: str = \"all\",\n start: Optional[int] = None, end: Optional[int] = None, step: Optional[int] = None,\n num_gpus: str = \"all\") -> dict:\n \"\"\"Time-series supply/demand, pricing, and stats per GPU type.\n\n Args:\n client: VastClient instance.\n gpu_name: GPU name, comma-separated list, or \"all\".\n verified: \"yes\", \"no\", or \"all\".\n hosting_type: \"all\", \"secure_cloud\", or \"community\".\n start: Start unix timestamp (defaults to end - 1 day server-side).\n end: End unix timestamp (defaults to now server-side).\n step: Step in seconds between data points.\n num_gpus: machine GPU-count bucket — \"all\", \"1\", \"2\", \"4\", or \"8\".\n\n Returns:\n Parsed JSON response. Shape: {\"success\": True, \"gpus\": {gpu_name: {supply_demand, pricing, stats}}}\n or {\"success\": True, \"gpus\": {}, \"needs_machine\": True}.\n \"\"\"\n params = {\"gpu_name\": gpu_name, \"verified\": verified, \"hosting_type\": hosting_type, \"num_gpus\": num_gpus}\n if start is not None:\n params[\"start\"] = str(start)\n if end is not None:\n params[\"end\"] = str(end)\n if step is not None:\n params[\"step\"] = str(step)\n r = client.get(\"/metrics/gpu/history/\", query_args=params)\n r.raise_for_status()\n return r.json()\n\n\ndef gpu_locations(client: VastClient) -> dict:\n \"\"\"Geographic locations of all GPUs on the platform.\n\n Returns:\n Parsed JSON response. Shape: {\"success\": True, \"locations\": [...]}\n or {\"success\": True, \"locations\": [], \"needs_machine\": True}.\n \"\"\"\n r = client.get(\"/metrics/gpu/locations/\")\n r.raise_for_status()\n return r.json()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":2694,"content_sha256":"5056fe04d7596cdaaaa953aa368fdf7d7cdfa4d00dfc607636f08d859fafd27e"},{"filename":"api/offers.py","content":"\"\"\"Search offers, templates, benchmarks, volumes, network volumes, and invoices.\"\"\"\nfrom typing import Optional\n\nfrom vastai.api.billing import show_user\nfrom vastai.api.client import VastClient\nfrom vastai.api.instances import resolve_runtype\nfrom vastai.api.query import (parse_order, parse_query, offers_alias,\n offers_fields, offers_mult)\nfrom vastai.utils import parse_env\n\n\ndef search_offers(client: VastClient, query: dict = None, offer_type: str = \"on-demand\",\n order: list = None, limit: int = None, storage: float = 5.0,\n no_default: bool = False, disable_bundling: bool = False) -> list:\n \"\"\"Search for instance offers using a query dict.\n\n Args:\n client: VastClient instance.\n query: Pre-parsed query dict of filters (e.g. {\"gpu_name\": {\"eq\": \"RTX 3090\"}}).\n offer_type: One of \"on-demand\", \"reserved\", or \"bid\".\n order: List of [field, direction] pairs, e.g. [[\"score\", \"desc\"]].\n limit: Max number of results.\n storage: Allocated storage in GiB for pricing (default 5.0).\n no_default: If True, skip default filters.\n disable_bundling: Deprecated bundling flag.\n\n Returns:\n List of offer dicts.\n \"\"\"\n if no_default:\n q = query or {}\n else:\n q = {\"verified\": {\"eq\": True}, \"external\": {\"eq\": False},\n \"rentable\": {\"eq\": True}, \"rented\": {\"eq\": False}}\n if query:\n q.update(query)\n\n if order is not None:\n q[\"order\"] = order\n else:\n q[\"order\"] = [[\"score\", \"desc\"]]\n\n q[\"type\"] = offer_type\n if offer_type == \"interruptible\":\n q[\"type\"] = \"bid\"\n\n if limit:\n q[\"limit\"] = int(limit)\n q[\"allocated_storage\"] = storage\n\n if disable_bundling:\n q[\"disable_bundling\"] = True\n\n r = client.post(\"/bundles/\", json_data=q)\n r.raise_for_status()\n return r.json()[\"offers\"]\n\n\ndef search_offers_new(client: VastClient, query: dict = None, offer_type: str = \"on-demand\",\n order: list = None, limit: int = None, storage: float = 5.0,\n no_default: bool = False, disable_bundling: bool = False) -> list:\n \"\"\"Search for instance offers using the new /search/asks/ endpoint.\n\n Args:\n client: VastClient instance.\n query: Pre-parsed query dict of filters (e.g. {\"gpu_name\": {\"eq\": \"RTX 3090\"}}).\n offer_type: One of \"on-demand\", \"reserved\", or \"bid\".\n order: List of [field, direction] pairs, e.g. [[\"score\", \"desc\"]].\n limit: Max number of results.\n storage: Allocated storage in GiB for pricing (default 5.0).\n no_default: If True, skip default filters.\n disable_bundling: Deprecated bundling flag.\n\n Returns:\n List of offer dicts.\n \"\"\"\n if no_default:\n q = query or {}\n else:\n q = {\"verified\": {\"eq\": True}, \"external\": {\"eq\": False},\n \"rentable\": {\"eq\": True}, \"rented\": {\"eq\": False}}\n if query:\n q.update(query)\n\n if order is not None:\n q[\"order\"] = order\n else:\n q[\"order\"] = [[\"score\", \"desc\"]]\n\n q[\"type\"] = offer_type\n if offer_type == \"interruptible\":\n q[\"type\"] = \"bid\"\n\n if limit:\n q[\"limit\"] = int(limit)\n q[\"allocated_storage\"] = storage\n\n if disable_bundling:\n q[\"disable_bundling\"] = True\n\n json_blob = {\"select_cols\": [\"*\"], \"q\": q}\n r = client.put(\"/search/asks/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()[\"offers\"]\n\n\ndef search_templates(client: VastClient, query: dict = None) -> list:\n \"\"\"Search for templates using a query dict.\n\n Args:\n client: VastClient instance.\n query: Pre-parsed query dict of select_filters.\n\n Returns:\n List of template dicts.\n \"\"\"\n query_args = {\"select_cols\": [\"*\"], \"select_filters\": query or {}}\n r = client.get(\"/template/\", query_args=query_args)\n r.raise_for_status()\n return r.json().get(\"templates\", [])\n\n\ndef benchmarks_query_args(query: dict = None, order: list = None,\n limit: int = None, after_token: str = None) -> dict:\n \"\"\"Build the query args for a ``/benchmarks`` request.\"\"\"\n query_args = {\"select_cols\": [\"*\"], \"select_filters\": query or {}}\n if order is not None:\n query_args[\"order_by\"] = order\n if limit is not None:\n query_args[\"limit\"] = int(limit)\n if after_token:\n query_args[\"after_token\"] = after_token\n return query_args\n\n\ndef search_benchmarks(client: VastClient, query: dict = None, order: list = None,\n limit: int = None, after_token: str = None) -> list:\n \"\"\"Search for benchmarks using a query dict, as a flat list.\n\n Pages through ``/benchmarks``, following ``next_token`` until it is\n exhausted, and concatenates every page. The backend caps a page at 200 rows\n by default, so callers that want one page at a time should use\n :func:`search_benchmarks_v1` instead.\n\n Args:\n client: VastClient instance.\n query: Pre-parsed query dict of select_filters.\n order: List of {\"col\": ..., \"dir\": ...} dicts, e.g. [{\"col\": \"last_update\", \"dir\": \"desc\"}].\n limit: Max number of results per page.\n after_token: Pagination token to resume from.\n\n Returns:\n List of benchmark dicts.\n \"\"\"\n rows = []\n params = benchmarks_query_args(query=query, order=order, limit=limit,\n after_token=after_token)\n while True:\n data = search_benchmarks_v1(client, params)\n rows.extend(data.get(\"benchmarks\") or [])\n next_token = data.get(\"next_token\")\n if not next_token:\n return rows\n params[\"after_token\"] = next_token\n\n\ndef search_benchmarks_v1(client: VastClient, params: dict) -> dict:\n \"\"\"Fetch one page of benchmarks using the paginated API.\n\n Args:\n client: VastClient instance.\n params: Dict with select_cols, select_filters, order_by, limit, after_token.\n\n Returns:\n Full response dict (benchmarks, benchmarks_found, next_token).\n \"\"\"\n r = client.get(\"/benchmarks\", query_args=params)\n r.raise_for_status()\n return r.json()\n\n\ndef search_volumes(client: VastClient, query: dict = None, order: list = None,\n limit: int = None, storage: float = 1.0,\n no_default: bool = False) -> list:\n \"\"\"Search for volume offers.\n\n Args:\n client: VastClient instance.\n query: Pre-parsed query dict of filters.\n order: List of [field, direction] pairs.\n limit: Max number of results.\n storage: Allocated storage in GiB for pricing (default 1.0).\n no_default: If True, skip default filters.\n\n Returns:\n List of volume offer dicts.\n \"\"\"\n if no_default:\n q = query or {}\n else:\n q = {\"verified\": {\"eq\": True}, \"external\": {\"eq\": False}, \"disk_space\": {\"gte\": 1}}\n if query:\n q.update(query)\n\n if order is not None:\n q[\"order\"] = order\n else:\n q[\"order\"] = [[\"score\", \"desc\"]]\n\n if limit:\n q[\"limit\"] = int(limit)\n q[\"allocated_storage\"] = storage\n\n r = client.post(\"/volumes/search/\", json_data=q)\n r.raise_for_status()\n return r.json()[\"offers\"]\n\n\ndef search_network_volumes(client: VastClient, query: dict = None, order: list = None,\n limit: int = None, storage: float = 1.0,\n no_default: bool = False) -> list:\n \"\"\"Search for network volume offers.\n\n Args:\n client: VastClient instance.\n query: Pre-parsed query dict of filters.\n order: List of [field, direction] pairs.\n limit: Max number of results.\n storage: Allocated storage in GiB for pricing (default 1.0).\n no_default: If True, skip default filters.\n\n Returns:\n List of network volume offer dicts.\n \"\"\"\n if no_default:\n q = query or {}\n else:\n q = {\"verified\": {\"eq\": True}, \"external\": {\"eq\": False}, \"disk_space\": {\"gte\": 1}}\n if query:\n q.update(query)\n\n if order is not None:\n q[\"order\"] = order\n else:\n q[\"order\"] = [[\"score\", \"desc\"]]\n\n if limit:\n q[\"limit\"] = int(limit)\n q[\"allocated_storage\"] = storage\n\n r = client.post(\"/network_volumes/search/\", json_data=q)\n r.raise_for_status()\n return r.json()[\"offers\"]\n\n\ndef search_invoices(client: VastClient, query: dict = None) -> list:\n \"\"\"Search for invoices using a query dict.\n\n Args:\n client: VastClient instance.\n query: Pre-parsed query dict of select_filters.\n\n Returns:\n List of invoice dicts.\n \"\"\"\n query_args = {\"select_cols\": [\"*\"], \"select_filters\": query or {}}\n r = client.get(\"/invoices\", query_args=query_args)\n r.raise_for_status()\n return r.json()\n\n\ndef template_fields_from_flags(*, ssh: bool = False, jupyter: bool = False,\n direct: bool = False, jupyter_lab: bool = False,\n login: Optional[str] = None,\n hide_readme: bool = False, public: bool = False,\n search_params: Optional[str] = None,\n no_default: bool = False,\n always_default: bool = False) -> dict:\n \"\"\"Translate the friendly template flags into api-layer template fields.\n\n Shared by the CLI commands and the SDK so both publish the same surface.\n\n ``always_default`` seeds the default offer filters even when no\n ``search_params`` were given, which is what the CLI has always done. The SDK\n passes False: it used to send no filters at all on a template created\n without search params, and quietly starting to attach three would change\n templates that existing scripts create.\n \"\"\"\n default_search_query = {}\n if not no_default and (always_default or search_params is not None):\n default_search_query = {\"verified\": {\"eq\": True}, \"external\": {\"eq\": False},\n \"rentable\": {\"eq\": True}}\n\n return {\n \"jup_direct\": jupyter and direct,\n \"ssh_direct\": ssh and direct,\n \"use_ssh\": ssh or jupyter,\n \"use_jupyter_lab\": jupyter_lab,\n \"runtype\": \"jupyter\" if jupyter else (\"ssh\" if ssh else \"args\"),\n \"docker_login_repo\": login.split(\" \")[0] if login else None,\n \"extra_filters\": parse_query(search_params, default_search_query,\n offers_fields, offers_alias, offers_mult),\n \"readme_visible\": not hide_readme,\n \"private\": not public,\n }\n\n\ndef create_template(client: VastClient, name: str = None, image: str = None,\n image_tag: str = None, href: str = None, repo: str = None,\n env: str = None, onstart_cmd: str = None,\n jup_direct: bool = False, ssh_direct: bool = False,\n use_jupyter_lab: bool = False, runtype: str = \"args\",\n use_ssh: bool = False, jupyter_dir: str = None,\n docker_login_repo: str = None, extra_filters: dict = None,\n disk_space: float = None, readme: str = None,\n readme_visible: bool = True, desc: str = None,\n private: bool = True) -> dict:\n \"\"\"Create a new template.\n\n Args:\n client: VastClient instance.\n name: Template name.\n image: Docker image.\n image_tag: Docker image tag.\n href: Link to provide.\n repo: Link to repository.\n env: Docker options env string.\n onstart_cmd: Onstart script contents.\n jup_direct: Supports jupyter direct.\n ssh_direct: Supports ssh direct.\n use_jupyter_lab: Launch with jupyter lab.\n runtype: Run type (jupyter, ssh, args).\n use_ssh: Supports ssh.\n jupyter_dir: Jupyter directory.\n docker_login_repo: Docker login repository.\n extra_filters: Search offer filters dict.\n disk_space: Recommended disk space.\n readme: Readme string.\n readme_visible: Whether readme is visible.\n desc: Description string.\n private: Whether template is private.\n\n Returns:\n Response dict with template info.\n \"\"\"\n template = {\n \"name\": name,\n \"image\": image,\n \"tag\": image_tag,\n \"href\": href,\n \"repo\": repo,\n \"env\": env,\n \"onstart\": onstart_cmd,\n \"jup_direct\": jup_direct,\n \"ssh_direct\": ssh_direct,\n \"use_jupyter_lab\": use_jupyter_lab,\n \"runtype\": runtype,\n \"use_ssh\": use_ssh,\n \"jupyter_dir\": jupyter_dir,\n \"docker_login_repo\": docker_login_repo,\n \"extra_filters\": extra_filters or {},\n \"recommended_disk_space\": disk_space,\n \"readme\": readme,\n \"readme_visible\": readme_visible,\n \"desc\": desc,\n \"private\": private,\n }\n r = client.post(\"/template/\", json_data=template)\n r.raise_for_status()\n return r.json()\n\n\ndef update_template(client: VastClient, hash_id: str, name: str = None,\n image: str = None, image_tag: str = None, href: str = None,\n repo: str = None, env: str = None, onstart_cmd: str = None,\n jup_direct: bool = False, ssh_direct: bool = False,\n use_jupyter_lab: bool = False, runtype: str = \"args\",\n use_ssh: bool = False, jupyter_dir: str = None,\n docker_login_repo: str = None, extra_filters: dict = None,\n disk_space: float = None, readme: str = None,\n readme_visible: bool = True, desc: str = None,\n private: bool = True) -> dict:\n \"\"\"Update an existing template.\n\n Resolves the template's numeric id from its hash_id first; the v1 update\n endpoint requires both. Edits always produce a new template version with a\n new hash_id, so send the full template rather than only the changed fields.\n\n Args:\n client: VastClient instance.\n hash_id: Hash ID of the template to update.\n (remaining args same as create_template)\n\n Returns:\n Response dict with updated template info.\n\n Raises:\n ValueError: No template with this hash_id is visible to the caller.\n \"\"\"\n template = {\n \"id\": resolve_template_id(client, hash_id),\n \"hash_id\": hash_id,\n \"name\": name,\n \"image\": image,\n \"tag\": image_tag,\n \"href\": href,\n \"repo\": repo,\n \"env\": env,\n \"onstart\": onstart_cmd,\n \"jup_direct\": jup_direct,\n \"ssh_direct\": ssh_direct,\n \"use_jupyter_lab\": use_jupyter_lab,\n \"runtype\": runtype,\n \"use_ssh\": use_ssh,\n \"jupyter_dir\": jupyter_dir,\n \"docker_login_repo\": docker_login_repo,\n \"extra_filters\": extra_filters or {},\n \"recommended_disk_space\": disk_space,\n \"readme\": readme,\n \"readme_visible\": readme_visible,\n \"desc\": desc,\n \"private\": private,\n }\n r = client.put(\"/api/v1/template/\", json_data=template)\n r.raise_for_status()\n return r.json()\n\n\ndef resolve_template_id(client: VastClient, hash_id: str) -> int:\n \"\"\"Return the numeric id of the template whose hash_id matches.\n\n The hash is content-based, so templates created by different users can share\n one. When several match, the caller's own template wins over a public one.\n\n Args:\n client: VastClient instance.\n hash_id: Hash ID of the template to look up.\n\n Returns:\n The template's numeric id.\n\n Raises:\n ValueError: No template with this hash_id is visible to the caller.\n \"\"\"\n matches = search_templates(client, {\"hash_id\": {\"eq\": hash_id}})\n if not matches:\n raise ValueError(f\"No template found with hash_id {hash_id}\")\n if len(matches) > 1:\n caller_id = show_user(client).get(\"id\")\n own = [match for match in matches if match.get(\"creator_id\") == caller_id]\n if own:\n return own[0][\"id\"]\n return matches[0][\"id\"]\n\n\ndef delete_template(client: VastClient, hash_id: str = None,\n template_id: int = None) -> dict:\n \"\"\"Delete a template by hash_id or template_id.\n\n Args:\n client: VastClient instance.\n hash_id: Hash ID of the template to delete.\n template_id: Numeric ID of the template to delete.\n\n Returns:\n Response dict.\n \"\"\"\n json_blob = {}\n if hash_id:\n json_blob[\"hash_id\"] = hash_id\n elif template_id:\n json_blob[\"template_id\"] = template_id\n r = client.delete(\"/template/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef launch_instance(client: VastClient, gpu_name: str, num_gpus: str, image: str,\n region: str = None, disk: float = 10, order: str = \"score-\",\n limit: int = None, env: dict = None, label: str = None,\n extra: str = None, onstart_cmd: str = None, login: str = None,\n python_utf8: bool = False, lang_utf8: bool = False,\n jupyter_lab: bool = False, jupyter_dir: str = None,\n cancel_unavail: bool = False,\n template_hash: str = None, runtype: str = None,\n args: str = None, query: dict = None,\n ssh: bool = False, jupyter: bool = False,\n direct: bool = False) -> dict:\n \"\"\"Launch the top instance from search offers matching the given criteria.\n\n Searches for offers and launches the best match in a single API call.\n\n Args:\n client: VastClient instance.\n gpu_name: GPU model name (e.g. \"RTX_4090\").\n num_gpus: Number of GPUs required.\n image: Docker image to launch.\n region: Region name or country code list (e.g. \"North_America\" or \"[US,CA]\").\n disk: Disk space in GB (default 10).\n order: Sort order for offers (default \"score-\").\n limit: Max number of offers to consider.\n env: Environment variables dict.\n label: Instance label.\n extra: Extra docker options.\n onstart_cmd: Onstart script contents.\n login: Docker login credentials.\n python_utf8: Enable Python UTF-8 mode.\n lang_utf8: Enable lang UTF-8 mode.\n jupyter_lab: Launch with Jupyter Lab.\n jupyter_dir: Jupyter directory.\n cancel_unavail: Cancel if unavailable.\n template_hash: Template hash ID.\n runtype: Run type (jupyter, ssh, args).\n ssh: Launch as an ssh instance type.\n jupyter: Launch as a jupyter instance instead of an ssh instance.\n direct: Use (faster) direct connections for jupyter & ssh.\n args: Container arguments.\n query: Pre-built query dict (overrides auto-built query from gpu_name/num_gpus).\n\n Returns:\n Response dict with launch result.\n \"\"\"\n if isinstance(env, str):\n env = parse_env(env)\n if template_hash is None:\n runtype, args = resolve_runtype(runtype, ssh=ssh, jupyter=jupyter,\n direct=direct, args=args,\n jupyter_lab=jupyter_lab,\n jupyter_dir=jupyter_dir)\n\n REGIONS = {\n \"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]\",\n \"South_America\": \"[AR, BO, BR, CL, CO, EC, FK, GF, GY, PY, PE, SR, UY, VE]\",\n \"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]\",\n \"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]\",\n \"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]\",\n \"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]\",\n }\n\n if query is None:\n args_query = f\"num_gpus={num_gpus} gpu_name={gpu_name}\"\n if region:\n region_query = REGIONS.get(region, region)\n args_query += f\" geolocation in {region_query}\"\n if disk:\n args_query += f\" disk_space>={disk}\"\n base_query = {\"verified\": {\"eq\": True}, \"external\": {\"eq\": False},\n \"rentable\": {\"eq\": True}, \"rented\": {\"eq\": False}}\n query = parse_query(args_query, base_query, offers_fields, offers_alias, offers_mult)\n\n # Parse order string\n order_list = parse_order(order) or []\n\n query[\"order\"] = order_list\n query[\"type\"] = \"on-demand\"\n if limit:\n query[\"limit\"] = int(limit)\n query[\"allocated_storage\"] = disk\n\n json_blob = {\n \"image\": image,\n \"disk\": disk,\n \"q\": query,\n \"env\": env or {},\n \"label\": label,\n \"extra\": extra,\n \"onstart\": onstart_cmd,\n \"image_login\": login,\n \"python_utf8\": python_utf8,\n \"lang_utf8\": lang_utf8,\n \"use_jupyter_lab\": jupyter_lab,\n \"jupyter_dir\": jupyter_dir,\n \"cancel_unavail\": cancel_unavail,\n \"template_hash_id\": template_hash,\n }\n if runtype:\n json_blob[\"runtype\"] = runtype\n if args is not None:\n json_blob[\"args\"] = args\n\n r = client.put(\"/launch_instance/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\n","content_type":"text/x-python; charset=utf-8","language":"python","size":21748,"content_sha256":"79c7ebc371412874cf75e12aa62292e1f75531b4695fda439893f7b50a5c5cf8"},{"filename":"api/price_increase.py","content":"\"\"\"Price-increase contract extension API.\n\nBackend pairs (vast/web/views/instance.py:257-312,\nvast/web/views/pydantic/models/instance.py):\n\n GET /api/v0/instances/pending-price-increases/\n PUT /api/v0/instances/accept-price-increase/\n PUT /api/v0/instances/reject-price-increase/\n\nBoth PUT endpoints require body exactly\n``{\"pending_price_increase_id\": \u003cint>}`` (``_Base.Config.extra='forbid'``).\nThe backend has no batch endpoint, no ``instance_ids``, no ``host_id``,\nand no ``snapshot`` field. The pending row id IS the identity.\n\nStale rows return ``HTTP 404`` with body\n``{\"success\": false, \"error\": \"no_pending_price_increase\"}``; older\nbackends may still return ``HTTP 409`` (the frontend recognises both).\n\"\"\"\n\nfrom vastai.api.client import VastClient\n\n\nNO_PENDING_PRICE_INCREASE = \"no_pending_price_increase\"\n\n\ndef list_pending(client: VastClient) -> dict:\n \"\"\"Return the pending-price-increase envelope.\n\n Shape (from web/price_increase_pending.py:_serialize_pending_row):\n\n {\n \"success\": True,\n \"count\": \u003cint>,\n \"truncated\": \u003cbool>,\n \"pending_price_increases\": [\n {\n \"pending_price_increase_id\": \u003cint>,\n \"contract_id\": \u003cint>,\n \"host_id\": \u003cint>,\n \"new_gpu_costpersec\": \u003cfloat|null>,\n \"new_disk_ram_costpersec\": \u003cfloat|null>,\n \"new_bwu_cost\": \u003cfloat|null>,\n \"new_bwd_cost\": \u003cfloat|null>,\n \"new_platform_fee\": \u003cfloat|null>,\n \"old_gpu_costpersec\": \u003cfloat|null>,\n \"old_disk_ram_costpersec\": \u003cfloat|null>,\n \"old_bwu_cost\": \u003cfloat|null>,\n \"old_bwd_cost\": \u003cfloat|null>,\n \"old_platform_fee\": \u003cfloat|null>,\n \"contract_end_date\": \u003cfloat|null>,\n \"ask_end_date\": \u003cfloat|null>,\n \"created_at\": \u003cfloat|null>\n },\n ...\n ]\n }\n \"\"\"\n r = client.get(\"/instances/pending-price-increases/\")\n r.raise_for_status()\n return r.json()\n\n\ndef pending_rows(envelope: dict) -> list[dict]:\n \"\"\"Return the pending rows from a pending-price-increase envelope.\"\"\"\n return envelope.get(\"pending_price_increases\", []) or []\n\n\ndef find_pending_for_instance(rows: list[dict], instance_id: int) -> dict | None:\n \"\"\"Find one pending row by contract/instance id.\"\"\"\n target = int(instance_id)\n return next((row for row in rows if row.get(\"contract_id\") == target), None)\n\n\ndef resolve_instance_to_pending(client: VastClient, instance_id: int) -> dict:\n \"\"\"Resolve an instance id to its pending-price-increase row.\n\n Raises ``LookupError`` when no pending row matches the instance id.\n \"\"\"\n envelope = list_pending(client)\n match = find_pending_for_instance(pending_rows(envelope), instance_id)\n if match is None:\n raise LookupError(\n f\"no pending price increase for instance {instance_id}\"\n )\n return match\n\n\ndef accept(client: VastClient, pending_id: int) -> dict:\n \"\"\"Accept one pending price-increase row.\n\n The body must be exactly ``{\"pending_price_increase_id\": int}``;\n backend ``extra='forbid'`` returns 400 for any extra key. Returns\n ``{\"success\": True, \"pending_price_increase_id\": int, \"contract_id\": int}``.\n \"\"\"\n r = client.put(\n \"/instances/accept-price-increase/\",\n json_data={\"pending_price_increase_id\": int(pending_id)},\n )\n r.raise_for_status()\n return r.json()\n\n\ndef reject(client: VastClient, pending_id: int) -> dict:\n \"\"\"Reject one pending price-increase row.\n\n Same body shape and response shape as :func:`accept`. The backend\n tombstones the row (status=rejected); no cutover follows.\n \"\"\"\n r = client.put(\n \"/instances/reject-price-increase/\",\n json_data={\"pending_price_increase_id\": int(pending_id)},\n )\n r.raise_for_status()\n return r.json()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":4102,"content_sha256":"a7eb16870e70b2bdd43c86620f1334c9468098670e02e696de4afe0b6617867f"},{"filename":"api/query.py","content":"\"\"\"Query parsing and field definitions for Vast.ai API.\n\nThis module is shared between the SDK and CLI layers.\n\"\"\"\n\nimport re\nimport sys\nimport time\nfrom datetime import datetime, timezone\nfrom typing import Dict, List\n\n\ndef numeric_version(version_str):\n try:\n # Split the version string by the period\n major, minor, patch = version_str.split('.')\n\n # Pad each part with leading zeros to make it 3 digits\n major = major.zfill(3)\n minor = minor.zfill(3)\n patch = patch.zfill(3)\n\n # Concatenate the padded parts\n numeric_version_str = f\"{major}{minor}{patch}\"\n\n # Convert the concatenated string to an integer\n result = int(numeric_version_str)\n #print(result)\n return result\n\n except ValueError:\n print(\"Invalid version string format. Expected format: X.X.X\")\n return None\n\n\ndef string_to_unix_epoch(date_string):\n if date_string is None:\n return None\n try:\n # Check if the input is a float or integer representing Unix time\n return float(date_string)\n except ValueError:\n # Parse MM/DD/YYYY input; interpret as UTC midnight so filtering\n # windows don't shift with the caller's local timezone.\n date_object = datetime.strptime(date_string, \"%m/%d/%Y\").replace(tzinfo=timezone.utc)\n return date_object.timestamp()\n\ndef fix_date_fields(query: Dict[str, Dict], date_fields: List[str]):\n \"\"\"Takes in a query and date fields to correct and returns query with appropriate epoch dates\"\"\"\n new_query: Dict[str, Dict] = {}\n for field, sub_query in query.items():\n # fix date values for given date fields\n if field in date_fields:\n new_sub_query = {k: string_to_unix_epoch(v) for k, v in sub_query.items()}\n new_query[field] = new_sub_query\n # else, use the original\n else: new_query[field] = sub_query\n\n return new_query\n\n\ndef version_string_sort(a, b) -> int:\n \"\"\"\n Accepts two version strings and decides whether a > b, a == b, or a \u003c b.\n This is meant as a sort function to be used for the driver versions in which only\n the == operator currently works correctly. Not quite finished...\n\n :param str a:\n :param str b:\n :return int:\n \"\"\"\n a_parts = a.split(\".\")\n b_parts = b.split(\".\")\n\n return 0\n\n\nvol_offers_fields = {\n \"cpu_arch\",\n \"cuda_vers\",\n \"cluster_id\",\n \"nw_disk_min_bw\",\n \"nw_disk_avg_bw\",\n \"nw_disk_max_bw\",\n \"datacenter\",\n \"disk_bw\",\n \"disk_space\",\n \"driver_version\",\n \"duration\",\n \"geolocation\",\n \"gpu_arch\",\n \"has_avx\",\n \"host_id\",\n \"id\",\n \"inet_down\",\n \"inet_up\",\n \"machine_id\",\n \"pci_gen\",\n \"pcie_bw\",\n \"reliability\",\n \"storage_cost\",\n \"static_ip\",\n \"total_flops\",\n \"ubuntu_version\",\n \"verified\",\n}\n\n\noffers_fields = {\n \"bw_nvlink\",\n \"compute_cap\",\n \"cpu_arch\",\n \"cpu_cores\",\n \"cpu_cores_effective\",\n \"cpu_ghz\",\n \"cpu_ram\",\n \"cuda_max_good\",\n \"datacenter\",\n \"direct_port_count\",\n \"driver_version\",\n \"disk_bw\",\n \"disk_space\",\n \"dlperf\",\n \"dlperf_per_dphtotal\",\n \"dph_total\",\n \"duration\",\n \"external\",\n \"flops_per_dphtotal\",\n \"gpu_arch\",\n \"gpu_display_active\",\n \"gpu_frac\",\n # \"gpu_ram_free_min\",\n \"gpu_mem_bw\",\n \"gpu_name\",\n \"gpu_ram\",\n \"gpu_total_ram\",\n \"gpu_display_active\",\n \"gpu_max_power\",\n \"gpu_max_temp\",\n \"has_avx\",\n \"host_id\",\n \"id\",\n \"inet_down\",\n \"inet_down_cost\",\n \"inet_up\",\n \"inet_up_cost\",\n \"machine_id\",\n \"min_bid\",\n \"mobo_name\",\n \"num_gpus\",\n \"pci_gen\",\n \"pcie_bw\",\n \"reliability\",\n #\"reliability2\",\n \"rentable\",\n \"rented\",\n \"storage_cost\",\n \"static_ip\",\n \"total_flops\",\n \"ubuntu_version\",\n \"verification\",\n \"verified\",\n \"vms_enabled\",\n \"geolocation\",\n \"cluster_id\"\n}\n\noffers_alias = {\n \"cuda_vers\": \"cuda_max_good\",\n \"display_active\": \"gpu_display_active\",\n #\"reliability\": \"reliability2\",\n \"dlperf_usd\": \"dlperf_per_dphtotal\",\n \"dph\": \"dph_total\",\n \"flops_usd\": \"flops_per_dphtotal\",\n}\n\noffers_mult = {\n \"cpu_ram\": 1000,\n \"gpu_ram\": 1000,\n \"gpu_total_ram\" : 1000,\n \"duration\": 24.0 * 60.0 * 60.0,\n}\n\n\nbenchmarks_fields = {\n \"contract_id\",# int ID of instance/contract reporting benchmark\n \"dph_base\",# float gpu $/hr at benchmark time\n \"gpu_name\",# string GPU model benchmarked (e.g. RTX_4090)\n \"id\",# int benchmark unique ID\n \"image\",# string image used for benchmark\n \"last_update\",# float date of benchmark\n \"launch_args\",# string launch args of the benchmarked worker\n \"machine_id\",# int id of machine benchmarked\n \"num_gpus\",# int number of gpus used in benchmark\n \"template_hash\",# string hash of the template benchmarked (formerly 'name')\n \"template_id\",# int id of the template benchmarked\n \"type\",# string benchmark type\n \"value\",# float benchmark result (formerly 'score')\n}\n\ninvoices_fields = {\n 'id',# int,\n 'user_id',# int,\n 'when',# float,\n 'paid_on',# float,\n 'payment_expected',# float,\n 'amount_cents',# int,\n 'is_credit',# bool,\n 'is_delayed',# bool,\n 'balance_before',# float,\n 'balance_after',# float,\n 'original_amount',# int,\n 'event_id',# string,\n 'cut_amount',# int,\n 'cut_percent',# float,\n 'extra',# json,\n 'service',# string,\n 'stripe_charge',# json,\n 'stripe_refund',# json,\n 'stripe_payout',# json,\n 'error',# json,\n 'paypal_email',# string,\n 'transfer_group',# string,\n 'failed',# bool,\n 'refunded',# bool,\n 'is_check',# bool,\n}\n\ntemplates_fields = {\n \"creator_id\",# int ID of creator\n \"created_at\",# float time of initial template creation (UTC epoch timestamp)\n \"count_created\",# int #instances created (popularity)\n \"default_tag\",# string image default tag\n \"docker_login_repo\",# string image docker repository\n \"id\",# int template unique ID\n \"image\",# string image used for benchmark\n \"jup_direct\",# bool supports jupyter direct\n \"hash_id\",# string unique hash ID of template\n \"private\",# bool true: only your templates, None: public templates\n \"name\",# string displayable name\n \"recent_create_date\",# float last time of instance creation (UTC epoch timestamp)\n \"recommended_disk_space\",# float min disk space required\n \"recommended\",# bool is templated on our recommended list\n \"ssh_direct\",# bool supports ssh direct\n \"tag\",# string image tag\n \"use_ssh\",# string supports ssh (direct or proxy)\n}\n\n\n\ndef parse_order(order, field_alias: Dict = None) -> List:\n \"\"\"Turn a comma-separated order string into [[field, direction], ...].\n\n A leading or trailing ``-`` means descending, ``+`` ascending. A list is\n returned unchanged so callers can pass either spelling.\n \"\"\"\n if order is None:\n return None\n if not isinstance(order, str):\n return order\n\n field_alias = offers_alias if field_alias is None else field_alias\n parsed = []\n for name in order.split(\",\"):\n name = name.strip()\n if not name:\n continue\n direction = \"asc\"\n field = name\n if name.strip(\"-\") != name:\n direction = \"desc\"\n field = name.strip(\"-\")\n if name.strip(\"+\") != name:\n direction = \"asc\"\n field = name.strip(\"+\")\n parsed.append([field_alias.get(field, field), direction])\n return parsed\n\ndef parse_query(query_str: str, res: Dict = None, fields = {}, field_alias = {}, field_multiplier = {}) -> Dict:\n \"\"\"\n Basically takes a query string (like the ones in the examples of commands for the search__offers function) and\n processes it into a dict of URL parameters to be sent to the server.\n\n :param str query_str:\n :param Dict res:\n :return Dict:\n \"\"\"\n if query_str is None:\n return res\n\n if res is None: res = {}\n if type(query_str) == list:\n query_str = \" \".join(query_str)\n query_str = query_str.strip()\n\n # Revised regex pattern to accurately capture quoted strings, bracketed lists, and single words/numbers\n #pattern = r\"([a-zA-Z0-9_]+)\\s*(=|!=|\u003c=|>=|\u003c|>| in | nin | eq | neq | not eq | not in )?\\s*(\\\"[^\\\"]*\\\"|\\[[^\\]]+\\]|[^ ]+)\"\n #pattern = \"([a-zA-Z0-9_]+)( *[=>\u003c!]+| +(?:[lg]te?|nin|neq|eq|not ?eq|not ?in|in) )?( *)(\\[[^\\]]+\\]|[^ ]+)?( *)\"\n pattern = r\"([a-zA-Z0-9_]+)( *[=>\u003c!]+| +(?:[lg]te?|nin|neq|eq|not ?eq|not ?in|in) )?( *)(\\[[^\\]]+\\]|\\\"[^\\\"]+\\\"|[^ ]+)?( *)\"\n opts = re.findall(pattern, query_str)\n\n #print(\"parse_query regex:\")\n #print(opts)\n\n #print(opts)\n # res = {}\n op_names = {\n \">=\": \"gte\",\n \">\": \"gt\",\n \"gt\": \"gt\",\n \"gte\": \"gte\",\n \"\u003c=\": \"lte\",\n \"\u003c\": \"lt\",\n \"lt\": \"lt\",\n \"lte\": \"lte\",\n \"!=\": \"neq\",\n \"==\": \"eq\",\n \"=\": \"eq\",\n \"eq\": \"eq\",\n \"neq\": \"neq\",\n \"noteq\": \"neq\",\n \"not eq\": \"neq\",\n \"notin\": \"notin\",\n \"not in\": \"notin\",\n \"nin\": \"notin\",\n \"in\": \"in\",\n }\n\n\n\n joined = \"\".join(\"\".join(x) for x in opts)\n if joined != query_str:\n raise ValueError(\n \"Unconsumed text. Did you forget to quote your query? \" + repr(joined) + \" != \" + repr(query_str))\n\n for field, op, _, value, _ in opts:\n value = value.strip(\",[]\")\n op = op.strip()\n op_name = op_names.get(op)\n\n if field in field_alias:\n old_field = field\n field = field_alias[field]\n if old_field in res:\n res[field] = res.pop(old_field)\n\n v = res.setdefault(field, {})\n\n if (field == \"driver_version\") and ('.' in value):\n value = numeric_version(value)\n\n if not field in fields:\n print(\"Warning: Unrecognized field: {}, see list of recognized fields.\".format(field), file=sys.stderr);\n if not op_name:\n raise ValueError(\"Unknown operator. Did you forget to quote your query? \" + repr(op).strip(\"u\"))\n if op_name in [\"in\", \"notin\"]:\n value = [x.strip() for x in value.split(\",\") if x.strip()]\n if not value:\n raise ValueError(\"Value cannot be blank. Did you forget to quote your query? \" + repr((field, op, value)))\n if not field:\n raise ValueError(\"Field cannot be blank. Did you forget to quote your query? \" + repr((field, op, value)))\n if value in [\"?\", \"*\", \"any\"]:\n if op_name != \"eq\":\n raise ValueError(\"Wildcard only makes sense with equals.\")\n if field in v:\n del v[field]\n if field in res:\n del res[field]\n continue\n\n if isinstance(value, str):\n value = value.replace('_', ' ')\n value = value.strip('\\\"')\n elif isinstance(value, list):\n value = [x.replace('_', ' ') for x in value]\n value = [x.strip('\\\"') for x in value]\n\n if field in field_multiplier:\n value = float(value) * field_multiplier[field]\n v[op_name] = value\n else:\n #print(value)\n if (value == 'true') or (value == 'True'):\n v[op_name] = True\n elif (value == 'false') or (value == 'False'):\n v[op_name] = False\n elif (value == 'None') or (value == 'null'):\n v[op_name] = None\n else:\n v[op_name] = value\n\n if field not in res:\n res[field] = v\n else:\n res[field].update(v)\n #print(res)\n return res\n","content_type":"text/x-python; charset=utf-8","language":"python","size":12475,"content_sha256":"d4aa31e600f67ff65b9db5d1b47e2e431cb23b42ec200813b1ec4bd5866158bd"},{"filename":"api/storage.py","content":"\"\"\"Storage, volume, and data transfer API functions for the Vast.ai SDK.\"\"\"\n\nimport time\n\n\ndef copy(client, src_id, dst_id, src_path, dst_path):\n \"\"\"Copy directories between instances, volumes, cloud services, and local.\n\n PUT /commands/copy_direct/ (remote-to-remote) or\n PUT /commands/rsync/ (when one side is local/None)\n\n Each of the source and destination locations can be either local or remote,\n subject to the read and write permissions required to carry out the action.\n Location IDs are the parsed form of a vast URL (see\n ``vastai.utils.parse_vast_url``):\n\n - instance_id (int) legacy format, still supported\n - \\\"C.instance_id\\\" container copy format\n - \\\"cloud_service\\\" cloud service, e.g. \\\"drive\\\"\n - \\\"cloud_service.id\\\" cloud service with ID, e.g. \\\"s3.101\\\" or \\\"hf.101\\\"\n - \\\"local\\\" or None local machine\n - \\\"V.volume_id\\\" volume copy, e.g. \\\"V.1234\\\"\n\n Volume copy is currently only supported for copying to other volumes,\n instances, or cloud services, not local.\n\n You should not copy to /root or / as a destination directory, as this can\n mess up the permissions on your instance ssh folder, breaking future copy\n operations (as they use ssh authentication). See\n https://vast.ai/docs/gpu-instances/data-movement#constraints for more\n information about constraints.\n\n Args:\n client: VastClient instance.\n src_id: Source location ID (or None/\\\"local\\\" for local).\n dst_id: Destination location ID (or None/\\\"local\\\" for local).\n src_path (str): Source path.\n dst_path (str): Destination path.\n\n Returns:\n dict: API response data. For local transfers this includes the\n ``src_addr``/``src_port`` or ``dst_addr``/``dst_port`` the caller\n uses to run rsync.\n\n Examples:\n copy(client, 6003036, 6003038, \\\"/workspace/\\\", \\\"/workspace/\\\")\n copy(client, \\\"C.11824\\\", \\\"local\\\", \\\"/data/test\\\", \\\"data/test\\\")\n copy(client, \\\"drive\\\", \\\"C.6003036\\\", \\\"/folder/file.txt\\\", \\\"/workspace/\\\")\n copy(client, \\\"s3.101\\\", \\\"C.6003036\\\", \\\"/data/\\\", \\\"/workspace/\\\")\n copy(client, \\\"V.1234\\\", \\\"C.5678\\\", \\\"/file\\\", \\\"/workspace/\\\")\n copy(client, \\\"V.1234\\\", \\\"s3.101\\\", \\\"/file\\\", \\\"/workspace/\\\")\n \"\"\"\n req_json = {\n \"client_id\": \"me\",\n \"src_id\": src_id,\n \"dst_id\": dst_id,\n \"src_path\": src_path,\n \"dst_path\": dst_path,\n }\n\n if src_id is None or dst_id is None:\n r = client.put(\"/commands/rsync/\", json_data=req_json)\n else:\n r = client.put(\"/commands/copy_direct/\", json_data=req_json)\n\n r.raise_for_status()\n return r.json()\n\n\ndef cancel_copy(client, dst_id):\n \"\"\"Cancel a remote copy in progress.\n\n DELETE /commands/copy_direct/\n\n Args:\n client: VastClient instance.\n dst_id: ID of copy destination to cancel.\n\n Returns:\n dict: API response data.\n \"\"\"\n req_json = {\"client_id\": \"me\", \"dst_id\": dst_id}\n r = client.delete(\"/commands/copy_direct/\", json_data=req_json)\n r.raise_for_status()\n return r.json()\n\n\ndef cancel_sync(client, dst_id):\n \"\"\"Cancel a remote cloud sync in progress.\n\n DELETE /commands/rclone/\n\n Args:\n client: VastClient instance.\n dst_id: ID of cloud sync destination to cancel.\n\n Returns:\n dict: API response data.\n \"\"\"\n req_json = {\"client_id\": \"me\", \"dst_id\": dst_id}\n r = client.delete(\"/commands/rclone/\", json_data=req_json)\n r.raise_for_status()\n return r.json()\n\n\ndef rclone_flags(*, dry_run: bool = False, size_only: bool = False,\n ignore_existing: bool = False, update: bool = False,\n delete_excluded: bool = False) -> list:\n \"\"\"Translate the friendly copy options into the rclone flag list.\"\"\"\n flags = []\n if dry_run:\n flags.append(\"--dry-run\")\n if size_only:\n flags.append(\"--size-only\")\n if ignore_existing:\n flags.append(\"--ignore-existing\")\n if update:\n flags.append(\"--update\")\n if delete_excluded:\n flags.append(\"--delete-excluded\")\n return flags\n\n\ndef cloud_copy(client, src, dst, instance, connection, transfer, flags=None):\n \"\"\"Copy files/folders to and from cloud providers.\n\n POST /commands/rclone/\n\n Supported cloud providers include \\\"drive\\\", \\\"s3\\\", \\\"b2\\\", \\\"dropbox\\\", and\n Hugging Face (\\\"hf\\\").\n\n Args:\n client: VastClient instance.\n src (str): Path to source of object to copy.\n dst (str): Path to target of copy operation. Default \\\"/workspace\\\".\n instance (str): ID of the instance.\n connection (str): ID of cloud connection on your account.\n transfer (str): Type of transfer (e.g. \\\"Instance to Cloud\\\",\n \\\"Cloud to Instance\\\").\n flags (list, optional): Additional rclone flags such as\n [\\\"--dry-run\\\", \\\"--size-only\\\", \\\"--ignore-existing\\\",\n \\\"--update\\\", \\\"--delete-excluded\\\"].\n\n Returns:\n dict: API response data.\n\n Hugging Face example (\\\"hf\\\" connection):\n cloud_copy(client, \\\"my-bucket/data\\\", \\\"/workspace\\\", \\\"6003036\\\",\n \\\"1234\\\", \\\"Cloud To Instance\\\")\n \"\"\"\n req_json = {\n \"src\": src,\n \"dst\": dst,\n \"instance_id\": instance,\n \"selected\": connection,\n \"transfer\": transfer,\n \"flags\": flags or [],\n }\n\n r = client.post(\"/commands/rclone/\", json_data=req_json)\n r.raise_for_status()\n return r.json()\n\n\ndef clone_volume(client, source, dest, size=None, disable_compression=False):\n \"\"\"Clone a volume to another volume.\n\n POST /volumes/copy/\n\n Args:\n client: VastClient instance.\n source: Source volume ID.\n dest: Destination volume ID.\n size (float, optional): Size in GB for destination volume.\n disable_compression (bool): Disable compression during clone.\n Default False.\n\n Returns:\n dict: API response data.\n \"\"\"\n json_blob = {\n \"src_id\": source,\n \"dst_id\": dest,\n }\n if size is not None:\n json_blob[\"size\"] = size\n if disable_compression:\n json_blob[\"disable_compression\"] = True\n\n r = client.post(\"/volumes/copy/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef show_volumes(client, type=\"all\"):\n \"\"\"Show stats on owned volumes.\n\n GET /volumes\n\n Args:\n client: VastClient instance.\n type (str): Volume type to display. Options: \\\"local\\\", \\\"network\\\",\n \\\"all\\\". Default \\\"all\\\".\n\n Returns:\n list: Volume data with computed duration field.\n \"\"\"\n types = {\n \"local\": \"local_volume\",\n \"network\": \"network_volume\",\n \"all\": \"all_volume\",\n }\n vol_type = types.get(type, \"all_volume\")\n r = client.get(\"/volumes\", query_args={\"owner\": \"me\", \"type\": vol_type})\n r.raise_for_status()\n rows = r.json()[\"volumes\"]\n processed = []\n for row in rows:\n row['duration'] = time.time() - row['start_date']\n processed.append(row)\n return processed\n\n\ndef show_instance_backups(client, contract_id=None):\n \"\"\"List the backups Vast has taken of the caller's instances and volumes.\n\n GET /instance_backups/\n\n Args:\n client: VastClient instance.\n contract_id (int): Restrict to one instance or volume contract id.\n\n Returns:\n list: Backup rows, newest first.\n \"\"\"\n query_args = {}\n if contract_id is not None:\n query_args[\"contract_id\"] = contract_id\n r = client.get(\"/instance_backups/\", query_args=query_args)\n r.raise_for_status()\n return r.json()[\"backups\"]\n\n\ndef instance_backup_files(client, contract_id):\n \"\"\"Objects stored for one backup, plus a short-lived token to fetch them.\n\n GET /instance_backups/\u003ccontract_id>/download/\n\n The returned URLs are not usable alone: each request must carry the\n authorization_token, either as an `Authorization` header or appended as an\n `?Authorization=\u003ctoken>` query parameter (what the CLI does, so the printed\n URLs are clickable). `truncated` is True when the backup holds more objects\n than were listed.\n\n Args:\n client: VastClient instance.\n contract_id (int): Instance or volume contract id.\n\n Returns:\n dict: authorization_token, expires_in_sec, truncated, files.\n \"\"\"\n r = client.get(f\"/instance_backups/{contract_id}/download/\")\n r.raise_for_status()\n return r.json()\n\n\ndef create_volume(client, id, size=15, name=None):\n \"\"\"Create a new volume.\n\n PUT /volumes/\n\n Args:\n client: VastClient instance.\n id (int): ID of volume offer.\n size (float): Size in GB. Default 15.\n name (str, optional): Optional name of volume.\n\n Returns:\n dict: API response data.\n \"\"\"\n json_blob = {\n \"size\": int(size),\n \"id\": int(id),\n }\n if name:\n json_blob[\"name\"] = name\n\n r = client.put(\"/volumes/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef delete_volume(client, id):\n \"\"\"Delete a volume.\n\n DELETE /volumes/\n\n All instances using the volume must be destroyed before deletion.\n\n Args:\n client: VastClient instance.\n id (int): ID of volume contract.\n\n Returns:\n dict: API response data.\n \"\"\"\n r = client.delete(\"/volumes/\", query_args={\"id\": id})\n r.raise_for_status()\n return r.json()\n\n\ndef list_volume(client, id, size=15, price_disk=0.10, end_date=None):\n \"\"\"[Host] List disk space for rent as a volume on a machine.\n\n POST /volumes/\n\n Args:\n client: VastClient instance.\n id (int): ID of machine to list.\n size (int): Size of disk space in GB. Default 15.\n price_disk (float): Storage price in $/GB/month. Default 0.10.\n end_date (str, optional): Contract offer expiration date\n (unix timestamp or MM/DD/YYYY format).\n\n Returns:\n dict: API response data.\n \"\"\"\n json_blob = {\n \"size\": int(size),\n \"machine\": int(id),\n \"price_disk\": float(price_disk),\n }\n if end_date is not None:\n json_blob[\"end_date\"] = end_date\n\n r = client.post(\"/volumes/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef list_volumes(client, ids, size=15, price_disk=0.10, end_date=None):\n \"\"\"[Host] List disk space for rent as volumes on multiple machines.\n\n POST /volumes/\n\n Args:\n client: VastClient instance.\n ids (list[int]): IDs of machines to list.\n size (int): Size of disk space in GB. Default 15.\n price_disk (float): Storage price in $/GB/month. Default 0.10.\n end_date (str, optional): Contract offer expiration date\n (unix timestamp or MM/DD/YYYY format).\n\n Returns:\n dict: API response data.\n \"\"\"\n json_blob = {\n \"size\": int(size),\n \"machine\": [int(mid) for mid in ids],\n \"price_disk\": float(price_disk),\n }\n if end_date is not None:\n json_blob[\"end_date\"] = end_date\n\n r = client.post(\"/volumes/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef unlist_volume(client, id):\n \"\"\"[Host] Unlist a volume offer.\n\n POST /volumes/unlist\n\n Args:\n client: VastClient instance.\n id (int): Volume ID to unlist.\n\n Returns:\n dict: API response data.\n \"\"\"\n json_blob = {\"id\": id}\n r = client.post(\"/volumes/unlist\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef create_network_volume(client, id, size=15, name=None):\n \"\"\"Create a new network volume.\n\n PUT /network_volumes/\n\n Args:\n client: VastClient instance.\n id (int): ID of network volume offer.\n size (float): Size in GB. Default 15.\n name (str, optional): Optional name of network volume.\n\n Returns:\n dict: API response data.\n \"\"\"\n json_blob = {\n \"size\": int(size),\n \"id\": int(id),\n }\n if name:\n json_blob[\"name\"] = name\n\n r = client.put(\"/network_volumes/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef list_network_volume(client, disk_id, price_disk=0.15, size=15, end_date=None):\n \"\"\"[Host] List disk space for rent as a network volume.\n\n POST /network_volumes/\n\n Args:\n client: VastClient instance.\n disk_id (int): ID of disk to list.\n price_disk (float): Storage price in $/GB/month. Default 0.15.\n size (int): Size of disk space in GB. Default 15.\n end_date (str, optional): Contract offer expiration date.\n\n Returns:\n dict: API response data.\n \"\"\"\n json_blob = {\n \"disk_id\": disk_id,\n \"price_disk\": price_disk,\n \"size\": size,\n }\n if end_date is not None:\n json_blob[\"end_date\"] = end_date\n\n r = client.post(\"/network_volumes/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef unlist_network_volume(client, id):\n \"\"\"[Host] Unlist a network volume offer.\n\n POST /network_volumes/unlist/\n\n Args:\n client: VastClient instance.\n id (int): ID of network volume offer to unlist.\n\n Returns:\n dict: API response data.\n \"\"\"\n json_blob = {\"id\": id}\n r = client.post(\"/network_volumes/unlist/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef show_network_disks(client):\n \"\"\"[Host] Show network disks associated with your account.\n\n GET /network_disk/\n\n Args:\n client: VastClient instance.\n\n Returns:\n dict: Network disk data including cluster and machine info.\n \"\"\"\n r = client.get(\"/network_disk/\")\n r.raise_for_status()\n return r.json()\n\n\ndef add_network_disk(client, machines, mount_point, disk_id=None):\n \"\"\"[Host] Add network disk to physical cluster.\n\n POST /network_disk/\n\n Args:\n client: VastClient instance.\n machines (list[int]): IDs of machines to add disk to.\n mount_point (str): Mount path of disk to add.\n disk_id (int, optional): ID of network disk to attach to machines.\n\n Returns:\n dict: API response data including disk_id.\n \"\"\"\n json_blob = {\n \"machines\": [int(mid) for mid in machines],\n \"mount_point\": mount_point,\n }\n if disk_id is not None:\n json_blob[\"disk_id\"] = disk_id\n\n r = client.post(\"/network_disk/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef show_connections(client):\n \"\"\"Display user's cloud connections/integrations.\n\n GET /users/cloud_integrations/\n\n Args:\n client: VastClient instance.\n\n Returns:\n dict/list: Cloud integration data.\n \"\"\"\n r = client.get(\"/users/cloud_integrations/\")\n r.raise_for_status()\n return r.json()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":14812,"content_sha256":"f24a2e69f4878f263b1c5c50df48a9960bd2f2310f1bbe9c16936068f3fd3f04"},{"filename":"api/teams.py","content":"\"\"\"Team CRUD operations, members, and roles.\"\"\"\nfrom vastai.api.client import VastClient\n\n\ndef create_team(client: VastClient, team_name: str, transfer_credit: float = 0) -> dict:\n \"\"\"Create a new team.\n\n Args:\n client: VastClient instance.\n team_name: Name of the team to create.\n transfer_credit: Amount of personal credit to transfer to the new team.\n\n Returns:\n Response dict with team info.\n \"\"\"\n json_data = {\"team_name\": team_name, \"transfer_credit\": transfer_credit}\n r = client.post(\"/team/\", json_data=json_data)\n r.raise_for_status()\n return r.json()\n\n\ndef destroy_team(client: VastClient) -> dict:\n \"\"\"Destroy the current team.\n\n Args:\n client: VastClient instance.\n\n Returns:\n Response dict.\n \"\"\"\n r = client.delete(\"/team/\")\n r.raise_for_status()\n return r.json()\n\n\ndef show_members(client: VastClient) -> dict:\n \"\"\"Show members of the current team.\n\n Args:\n client: VastClient instance.\n\n Returns:\n Response dict with member info.\n \"\"\"\n r = client.get(\"/team/members/\")\n r.raise_for_status()\n return r.json()\n\n\ndef invite_member(client: VastClient, email: str, role: str) -> dict:\n \"\"\"Invite a member to the current team.\n\n Args:\n client: VastClient instance.\n email: Email address of the member to invite.\n role: Role to assign to the invited member.\n\n Returns:\n Response dict.\n \"\"\"\n r = client.post(\"/team/invite/\", query_args={\"email\": email, \"role\": role})\n r.raise_for_status()\n return r.json()\n\n\ndef remove_member(client: VastClient, id: int) -> dict:\n \"\"\"Remove a member from the current team.\n\n Args:\n client: VastClient instance.\n id: Member ID to remove.\n\n Returns:\n Response dict.\n \"\"\"\n r = client.delete(f\"/team/members/{id}/\")\n r.raise_for_status()\n return r.json()\n\n\ndef create_team_role(client: VastClient, name: str, permissions: dict) -> dict:\n \"\"\"Add a new role to the current team.\n\n Args:\n client: VastClient instance.\n name: Name of the role.\n permissions: Dict of permissions for the role.\n\n Returns:\n Response dict.\n \"\"\"\n r = client.post(\"/team/roles/\", json_data={\"name\": name, \"permissions\": permissions})\n r.raise_for_status()\n return r.json()\n\n\ndef show_team_role(client: VastClient, name: str) -> dict:\n \"\"\"Show details of a specific team role.\n\n Args:\n client: VastClient instance.\n name: Name of the role.\n\n Returns:\n Role details dict.\n \"\"\"\n r = client.get(f\"/team/roles/{name}/\")\n r.raise_for_status()\n return r.json()\n\n\ndef show_team_roles(client: VastClient) -> dict:\n \"\"\"Show all roles for the current team.\n\n Args:\n client: VastClient instance.\n\n Returns:\n Response dict with roles info.\n \"\"\"\n r = client.get(\"/team/roles-full/\")\n r.raise_for_status()\n return r.json()\n\n\ndef update_team_role(client: VastClient, id: int, name: str = None,\n permissions: dict = None) -> dict:\n \"\"\"Update an existing team role.\n\n Args:\n client: VastClient instance.\n id: Role ID.\n name: New name for the role.\n permissions: Updated permissions dict.\n\n Returns:\n Response dict.\n \"\"\"\n json_blob = {}\n if name is not None:\n json_blob[\"name\"] = name\n if permissions is not None:\n json_blob[\"permissions\"] = permissions\n r = client.put(f\"/team/roles/{id}/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n\n\ndef remove_team_role(client: VastClient, name: str) -> dict:\n \"\"\"Remove a role from the current team.\n\n Args:\n client: VastClient instance.\n name: Name of the role to remove.\n\n Returns:\n Response dict.\n \"\"\"\n r = client.delete(f\"/team/roles/{name}/\")\n r.raise_for_status()\n return r.json()\n\n\ndef transfer_credit(client: VastClient, recipient: str, amount: float) -> dict:\n \"\"\"Transfer credit to another account.\n\n Args:\n client: VastClient instance.\n recipient: Recipient identifier (email or user ID).\n amount: Amount of credit to transfer.\n\n Returns:\n Response dict.\n \"\"\"\n json_blob = {\n \"sender\": \"me\",\n \"recipient\": recipient,\n \"amount\": amount,\n }\n r = client.put(\"/commands/transfer_credit/\", json_data=json_blob)\n r.raise_for_status()\n return r.json()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":4461,"content_sha256":"927e939dc69e53f968063944cb2b33ee0f867070ae065f2ec23e9b2c0af7405b"},{"filename":"async_/__init__.py","content":"from .client import AsyncClient\nfrom .results import AsyncOffer, AsyncInstance\n","content_type":"text/x-python; charset=utf-8","language":"python","size":79,"content_sha256":"4f9b2926101701c1ce047cf7506a8b451a58abeae81a67162382644ff3bc9886"},{"filename":"async_/client.py","content":"import aiohttp\nimport ssl\nimport os\nimport json\nimport tempfile\nimport asyncio\nfrom typing import Any, Optional, Union\n\nfrom vastai._base import _BaseClient, _APIKEY_SENTINEL\nfrom vastai.data.query import Query\nfrom vastai.data.offer import Offer\nfrom vastai.data.instance import Instance, InstanceConfig, CreateInstanceResponse\nfrom vastai.async_.results import AsyncOffer, AsyncInstance\nfrom vastai.serverless.client.client import CoroutineServerless\n\n\nclass AsyncClient(_BaseClient):\n \"\"\"\n Asynchronous Vast.ai API client.\n\n api_key resolution order:\n 1. Explicit api_key argument\n 2. VAST_API_KEY environment variable\n 3. $XDG_CONFIG_HOME/vastai/vast_api_key\n 4. ~/.vast_api_key (legacy)\n \"\"\"\n\n SSL_CERT_URL = \"https://console.vast.ai/static/jvastai_root.cer\"\n\n def __init__(\n self,\n api_key: object = _APIKEY_SENTINEL,\n vast_server: str = \"https://console.vast.ai\",\n connection_limit: int = 100,\n ):\n super().__init__(api_key, vast_server)\n self._connection_limit = connection_limit\n self._session: aiohttp.ClientSession | None = None\n self._ssl_context: ssl.SSLContext | None = None\n self._session_depth = 0\n\n async def get_ssl_context(self) -> ssl.SSLContext:\n \"\"\"Download Vast.ai root cert and build SSL context (cached).\n\n The custom cert is added on top of the system default CAs, so this\n context works for both console.vast.ai and serverless worker connections.\n \"\"\"\n if self._ssl_context is None:\n async with aiohttp.ClientSession() as s:\n async with s.get(self.SSL_CERT_URL) as resp:\n if resp.status != 200:\n raise Exception(f\"Failed to fetch SSL cert: {resp.status}\")\n cert_bytes = await resp.read()\n\n tmpfile = tempfile.NamedTemporaryFile(delete=False, suffix=\".cer\")\n tmpfile.write(cert_bytes)\n tmpfile.close()\n\n ctx = ssl.create_default_context()\n ctx.load_verify_locations(cafile=tmpfile.name)\n\n self._ssl_context = ctx\n os.unlink(tmpfile.name)\n\n return self._ssl_context\n\n async def _get_session(self, use_context=False) -> aiohttp.ClientSession:\n if self._session is None or self._session.closed:\n self._session = aiohttp.ClientSession(\n connector=aiohttp.TCPConnector(\n limit=self._connection_limit,\n ssl=await self.get_ssl_context(),\n )\n )\n self._session_depth += 1\n if not use_context:\n loop = asyncio.get_running_loop()\n original_close = loop.close\n\n def close():\n loop.run_until_complete(self.close(force=True))\n original_close()\n\n loop.close = close\n elif use_context:\n 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.\n return self._session\n\n def is_open(self):\n return self._session is not None and not self._session.closed\n\n async def close(self, force=False):\n if self._session and not self._session.closed:\n if self._session_depth > 0:\n self._session_depth -= 1\n if force or not self._session_depth:\n self._session_depth = 0\n await self._session.close()\n\n async def __aenter__(self):\n await self._get_session(use_context=True)\n return self\n\n async def __aexit__(self, exc_type, exc, tb):\n await self.close()\n\n # ── Search ──────────────────────────────────────────────────────────────\n\n async def search(\n self,\n query: Query,\n type_: Optional[str] = None,\n order: Optional[list] = [[\"score\", \"desc\"]],\n limit: Optional[int] = None,\n allocated_storage: Optional[float] = None,\n target_reliability: Optional[float] = None,\n external: Optional[bool] = None,\n extra_ids: Optional[list] = None,\n has_avx: Optional[int] = None,\n gpu_option: Optional[Any] = None,\n show_incompatible: Optional[Any] = None,\n sort_option: Optional[Any] = None,\n template_id: Optional[Any] = None,\n ) -> list[AsyncOffer]:\n \"\"\"Search available offers and return a list of AsyncOffer result wrappers.\"\"\"\n payload = dict(query.query)\n for key, value in [\n (\"type\", type_),\n (\"order\", order),\n (\"limit\", limit),\n (\"allocated_storage\", allocated_storage),\n (\"target_reliability\", target_reliability),\n (\"external\", external),\n (\"extra_ids\", extra_ids),\n (\"has_avx\", has_avx),\n (\"gpu_option\", gpu_option),\n (\"show_incompatible\", show_incompatible),\n (\"sort_option\", sort_option),\n (\"template_id\", template_id),\n ]:\n if value is not None:\n payload[key] = value\n\n session = await self._get_session()\n async with session.post(\n self._url(\"/api/v0/bundles/\"),\n json=payload,\n headers=self._headers(),\n ) as resp:\n resp.raise_for_status()\n data = await resp.json(content_type=None)\n\n return [AsyncOffer(Offer.from_dict(o), self) for o in data.get(\"offers\", [])]\n\n # ── Create Instance ─────────────────────────────────────────────────────\n\n async def create_instance(\n self,\n offer_id: int,\n config: InstanceConfig,\n ) -> AsyncInstance:\n \"\"\"Accept an offer and create an instance. Returns an AsyncInstance wrapper.\"\"\"\n payload = config.to_dict()\n session = await self._get_session()\n async with session.put(\n self._url(f\"/api/v0/asks/{offer_id}/\"),\n json=payload,\n headers=self._headers(),\n ) as resp:\n resp.raise_for_status()\n data = await resp.json(content_type=None)\n\n if not data.get(\"success\"):\n raise Exception(\n f\"Failed to create instance: {data.get('error', '')} - {data.get('msg', '')}\"\n )\n\n resp_data = CreateInstanceResponse.from_dict(data)\n # Fetch the newly created instance\n instances = await self.show_instances()\n for inst in instances:\n if inst.id == resp_data.new_contract:\n return inst\n # Fallback: return a minimal instance\n return AsyncInstance(Instance.from_dict({\"id\": resp_data.new_contract}), self)\n\n # ── Instances ───────────────────────────────────────────────────────────\n\n async def show_instances(self) -> list[AsyncInstance]:\n \"\"\"Return all instances owned by the authenticated user.\n\n The legacy v0 ``/instances`` list endpoint is deprecated; this pages\n through the v1 ``/api/v1/instances/`` endpoint (omitting ``select_cols``\n so the backend returns full instance rows) and accumulates every page.\n \"\"\"\n session = await self._get_session()\n instances = []\n params = {\n \"select_filters\": json.dumps({}),\n \"order_by\": json.dumps([{\"col\": \"id\", \"dir\": \"asc\"}]),\n \"limit\": 25,\n }\n while True:\n async with session.get(\n self._url(\"/api/v1/instances/\"),\n params=params,\n headers=self._headers(),\n ) as resp:\n resp.raise_for_status()\n data = await resp.json(content_type=None)\n instances.extend(data.get(\"instances\") or [])\n next_token = data.get(\"next_token\")\n if not next_token:\n break\n params[\"after_token\"] = next_token\n\n return [\n AsyncInstance(Instance.from_dict(i), self)\n for i in instances\n ]\n\n async def destroy_instance(self, instance_or_id: Union[AsyncInstance, int]) -> None:\n \"\"\"Destroy a running or stopped instance.\"\"\"\n instance_id = (\n instance_or_id if isinstance(instance_or_id, int) else instance_or_id.id\n )\n session = await self._get_session()\n async with session.delete(\n self._url(f\"/api/v0/instances/{instance_id}/\"),\n headers=self._headers(),\n ) as resp:\n resp.raise_for_status()\n\n # ── Serverless factory ──────────────────────────────────────────────────\n\n def serverless(\n self,\n *,\n instance: str = \"prod\",\n autoscaler_url: Optional[str] = None,\n debug: bool = False,\n default_request_timeout: float = 600.0,\n max_poll_interval: float = 5.0,\n ) -> CoroutineServerless:\n \"\"\"Create a CoroutineServerless that shares this client's aiohttp session and SSL context.\"\"\"\n sl = CoroutineServerless(\n api_key=self._api_key,\n instance=instance,\n autoscaler_url=autoscaler_url,\n webserver_url=self._vast_server,\n debug=debug,\n connection_limit=self._connection_limit,\n default_request_timeout=default_request_timeout,\n max_poll_interval=max_poll_interval,\n )\n sl._transport_owner = self\n return sl\n","content_type":"text/x-python; charset=utf-8","language":"python","size":9873,"content_sha256":"8327e28bf5ab6fe1967f477172b935ad2408a0fed2a739d8077a8f2d171964a3"},{"filename":"async_/results.py","content":"from __future__ import annotations\nfrom typing import TYPE_CHECKING, Optional, Union\n\nfrom vastai.data.offer import Offer\nfrom vastai.data.instance import InstanceConfig\n\nif TYPE_CHECKING:\n from vastai.async_.client import AsyncClient\n from vastai.data.instance import Instance\n\n\nclass AsyncOffer:\n def __init__(self, data: Offer, client: AsyncClient):\n self._data = data\n self._client = client\n\n async def create_instance(self, config: InstanceConfig) -> AsyncInstance:\n \"\"\"Create an instance from this offer. Synonym for AsyncClient.create_instance(self.id, config).\"\"\"\n return await self._client.create_instance(self._data.id, config)\n\n def __getattr__(self, name: str):\n return getattr(self._data, name)\n\n def __repr__(self) -> str:\n return f\"AsyncOffer({self._data!r})\"\n\n\nclass AsyncInstance:\n def __init__(self, data: Instance, client: AsyncClient):\n self._data = data\n self._client = client\n\n async def destroy(self) -> None:\n \"\"\"Synonym for AsyncClient.destroy_instance(self).\"\"\"\n return await self._client.destroy_instance(self._data.id)\n\n def __getattr__(self, name: str):\n return getattr(self._data, name)\n\n def __repr__(self) -> str:\n return f\"AsyncInstance({self._data!r})\"\n","content_type":"text/x-python; charset=utf-8","language":"python","size":1301,"content_sha256":"bb40f8ebd9128ac89befb59b12a47e49fc0ba36c5091b83a41feb3f34b2c7713"},{"filename":"cli/__init__.py","content":"\"\"\"Vast.ai CLI layer.\"\"\"\n","content_type":"text/x-python; charset=utf-8","language":"python","size":25,"content_sha256":"4f177f45e6e2c9a4053cbf9274ccb12cb41f757db0463135d5fd120fa9ce412c"},{"filename":"cli/commands/__init__.py","content":"\"\"\"CLI command modules.\n\nEach module registers its commands with the global parser at import time\nvia the ``@parser.command(...)`` decorator pattern. Importing a module\nis sufficient to register all of its commands.\n\"\"\"\n\n\ndef register_all_commands(parser):\n \"\"\"Import all command modules to register their commands with the parser.\n\n The imports themselves trigger the decorator registrations -- no\n explicit ``register()`` call is needed per module.\n \"\"\"\n from vastai.cli.commands import ( # noqa: F401\n instances, offers, machines, teams, keys, endpoints,\n billing, storage, auth, misc, deployments, metrics,\n benchmarks,\n price_increase,\n # clusters, # cluster/overlay commands disabled for now\n )\n","content_type":"text/x-python; charset=utf-8","language":"python","size":757,"content_sha256":"7d7180ed5f310c849e39f5fd48b71496114bcff2abc16ddc9ffaada89849f27f"},{"filename":"cli/commands/auth.py","content":"\"\"\"CLI commands for authentication and 2FA.\"\"\"\n\nimport json\nimport os\nimport sys\nfrom datetime import datetime, timezone\n\nfrom vastai.cli.parser import argument\nfrom vastai.cli.display import deindent, display_table\nfrom vastai.api import auth as auth_api\nfrom vastai.cli.util import SUCCESS, WARN, FAIL, format_key_suffix\n\n\n# ---------------------------------------------------------------------------\n# TFA helper functions (ported from old/vast.py)\n# ---------------------------------------------------------------------------\n\nTFA_METHOD_FIELDS = (\n (\"id\", \"ID\", \"{}\", None, True),\n (\"user_id\", \"User ID\", \"{}\", None, True),\n (\"is_primary\", \"Primary\", \"{}\", None, True),\n (\"method\", \"Method\", \"{}\", None, True),\n (\"label\", \"Label\", \"{}\", None, True),\n (\"phone_number\", \"Phone Number\", \"{}\", None, False),\n (\"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),\n (\"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),\n (\"fail_count\", \"Failures\", \"{}\", None, True),\n (\"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),\n)\n\n\ndef display_tfa_methods(methods):\n \"\"\"Helper function to display 2FA methods in a table.\"\"\"\n method_fields = TFA_METHOD_FIELDS\n has_sms = any(m['method'] == 'sms' for m in methods)\n if not has_sms: # Don't show Phone Number column if the user has no SMS methods\n method_fields = tuple(field for field in TFA_METHOD_FIELDS if field[0] != 'phone_number')\n display_table(methods, method_fields, replace_spaces=False)\n\n\ndef confirm_destructive_action(prompt=\"Are you sure? (y/n): \"):\n \"\"\"Prompt user for confirmation of destructive actions.\"\"\"\n try:\n response = input(f\" {prompt}\").strip().lower()\n return 'y' in response\n except (EOFError, KeyboardInterrupt):\n print(\"\\nOperation cancelled.\")\n raise\n\n\ndef handle_failed_tfa_verification(args, e):\n \"\"\"Parse TFA error responses and provide contextual error messages.\"\"\"\n try:\n error_data = e.response.json()\n except Exception:\n print(f\"\\n{FAIL} Error: {e}\")\n return\n\n error_msg = error_data.get(\"msg\", str(e))\n error_code = error_data.get(\"error\", \"\")\n\n if args.raw:\n print(json.dumps(error_data, indent=2))\n\n print(f\"\\n{FAIL} Error: {error_msg}\")\n\n # Provide helpful context for common errors\n if error_code in {\"tfa_locked\", \"2fa_verification_failed\"}:\n fail_count = error_data.get(\"fail_count\", 0)\n locked_until = error_data.get(\"locked_until\")\n\n if fail_count > 0:\n print(f\" Failed attempts: {fail_count}\")\n if locked_until:\n lock_time_sec = (datetime.fromtimestamp(locked_until) - datetime.now()).seconds\n minutes, seconds = divmod(lock_time_sec, 60)\n print(f\" Time Remaining for 2FA Lock: {minutes} minutes and {seconds} seconds...\")\n\n elif error_code == \"2fa_expired\":\n # Note: Only SMS uses tfa challenges that expire when verifying\n print(\"\\n The SMS code and secret have expired. Please start over:\")\n print(\" vastai tfa send-sms\")\n\n\ndef format_backup_codes(backup_codes):\n \"\"\"Format backup codes for display or file output.\"\"\"\n output_lines = [\n \"=\" * 60, \" VAST.AI 2FA BACKUP CODES\", \"=\" * 60,\n f\"\\nGenerated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\",\n f\"\\n{WARN} WARNING: All previous backup codes are now invalid!\",\n \"\\nYour New Backup Codes (one-time use only):\",\n \"-\" * 40,\n ]\n for i, code in enumerate(backup_codes, 1):\n output_lines.append(f\" {i:2d}. {code}\")\n output_lines.extend([\n \"-\" * 40,\n \"\\nIMPORTANT:\",\n \" \\u2022 Each code can only be used once\",\n \" \\u2022 Store them in a secure location\",\n \" \\u2022 Use these codes to log in if you lose access to your 2FA device\",\n \"\\n\" + \"=\" * 60,\n ])\n return \"\\n\".join(output_lines)\n\n\ndef _save_to_file(content, filepath):\n \"\"\"Save content to file, creating parent directories if needed.\"\"\"\n try:\n filepath = os.path.abspath(os.path.expanduser(filepath))\n parent_dir = os.path.dirname(filepath)\n if parent_dir:\n os.makedirs(parent_dir, exist_ok=True)\n with open(filepath, \"w\") as f:\n f.write(content)\n return True\n except (IOError, OSError) as e:\n print(f\"\\n{FAIL} Error saving file: {e}\")\n return False\n\n\ndef _get_backup_codes_filename():\n \"\"\"Generate a timestamped filename for backup codes.\"\"\"\n timestamp = datetime.now().strftime(\"%Y-%m-%d_%H-%M-%S\")\n return f\"vastai_backup_codes_{timestamp}.txt\"\n\n\ndef save_backup_codes(backup_codes):\n \"\"\"Save or display 2FA backup codes based on user choice.\"\"\"\n print(f\"\\nBackup codes regenerated successfully! {SUCCESS}\")\n print(f\"\\n{WARN} WARNING: All previous backup codes are now invalid!\")\n\n formatted_content = format_backup_codes(backup_codes)\n filename = _get_backup_codes_filename()\n\n while True:\n print(\"\\nHow would you like to save your new backup codes?\")\n print(f\" 1. Save to default location (~/Downloads/{filename})\")\n print(f\" 2. Save to a custom path\")\n print(f\" 3. Print to screen ({WARN} potentially unsafe - visible to onlookers)\")\n\n try:\n choice = input(\"\\nEnter choice (1-3): \").strip()\n\n if choice in {'1', '2'}:\n if choice == '1':\n downloads_dir = os.path.expanduser(\"~/Downloads\")\n filepath = os.path.join(downloads_dir, filename)\n else: # choice == '2'\n custom_path = input(\"\\nEnter full path for backup codes file: \").strip()\n if not custom_path:\n print(\"Error: Path cannot be empty\")\n continue\n\n filepath = os.path.abspath(os.path.expanduser(custom_path))\n if os.path.isdir(filepath):\n filepath = os.path.join(filepath, filename)\n\n if _save_to_file(formatted_content, filepath):\n print(f\"\\n{SUCCESS} Backup codes saved to: {filepath}\")\n print(f\"\\nIMPORTANT:\")\n print(f\" \\u2022 The file contains {len(backup_codes)} one-time use backup codes\")\n if choice == '1':\n print(f\" \\u2022 Move this file to a secure location\")\n return\n else:\n print(\"Please try again with a different path.\")\n continue\n\n elif choice == '3':\n print(f\"\\n{WARN} WARNING: Printing sensitive codes to screen!\")\n confirm = input(\"\\nAre you sure? Anyone nearby can see these codes. (yes/no): \").strip().lower()\n if confirm in {'yes', 'y'}:\n print(\"\\n\" + formatted_content + \"\\n\")\n return\n else:\n print(\"Cancelled. Please choose another option.\")\n continue\n\n else:\n print(\"Invalid choice. Please enter 1, 2, or 3.\")\n\n except (EOFError, KeyboardInterrupt):\n print(\"\\n\\nOperation cancelled. Your backup codes were generated but not saved.\")\n print(\"You will need to regenerate them to get new codes.\")\n raise\n\n\nfrom vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401\n\n\nparser = _get_parser()\n\n\n# ---------------------------------------------------------------------------\n# set api-key\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"new_api_key\", help=\"Api key to set as currently logged in user\"),\n usage=\"vastai set api-key APIKEY\",\n help=\"Set api-key (get your api-key from the console/CLI)\",\n)\ndef set__api_key(args):\n \"\"\"Set the api-key.\"\"\"\n from vastai.cli.util import APIKEY_FILE, APIKEY_FILE_HOME\n\n with open(APIKEY_FILE, \"w\") as writer:\n writer.write(args.new_api_key)\n print(\"Your api key has been saved in {}\".format(APIKEY_FILE))\n\n if os.path.exists(APIKEY_FILE_HOME):\n os.remove(APIKEY_FILE_HOME)\n if getattr(args, \"explain\", False):\n print(\"Removed legacy api key file at {} (no longer used).\".format(APIKEY_FILE_HOME))\n\n env_key = os.environ.get(\"VAST_API_KEY\")\n if env_key:\n 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.\")\n print(\"Unset VAST_API_KEY to use the saved key.\")\n\n\n# ---------------------------------------------------------------------------\n# set role\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"role\", help=\"'host' or 'client'\", choices=[\"host\", \"client\"]),\n usage=\"vastai set role host|client\",\n help=\"Set host role to view host-only commands\",\n epilog=deindent(\"\"\"\n This only changes what --help and tab completion display — every command\n still runs regardless of role, since the server enforces real permissions.\n \"\"\"),\n)\ndef set__role(args):\n \"\"\"Set the client/host CLI display role.\"\"\"\n from vastai.cli.util import set_role_file\n\n set_role_file(args.role)\n print(f\"CLI role set to '{args.role}'.\")\n if args.role == \"client\":\n print(\"Host-only commands are now hidden from --help and tab completion.\")\n else:\n print(\"All commands (client + host) are now shown in --help and tab completion.\")\n\n\n# ---------------------------------------------------------------------------\n# Note: set__user and show__user are in billing.py.\n# Note: show__ipaddrs is in billing.py.\n# ---------------------------------------------------------------------------\n\n\n# ---------------------------------------------------------------------------\n# show audit-logs\n# ---------------------------------------------------------------------------\n\[email protected](\n usage=\"vastai show audit-logs [--api-key API_KEY] [--raw]\",\n help=\"Display account's history of important actions\",\n)\ndef show__audit_logs(args):\n \"\"\"Show the history of important actions.\"\"\"\n from vastai.cli.display import display_table, audit_log_fields\n client = get_client(args)\n rows = auth_api.show_audit_logs(client)\n if args.raw:\n return rows\n else:\n display_table(rows, audit_log_fields)\n\n\n# ---------------------------------------------------------------------------\n# environment variables\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"-s\", \"--show-values\", action=\"store_true\", help=\"Show the values of environment variables\"),\n usage=\"vastai show env-vars [-s]\",\n help=\"Show user environment variables\",\n)\ndef show__env_vars(args):\n \"\"\"Show the environment variables for the current user.\"\"\"\n client = get_client(args)\n env_vars = auth_api.show_env_vars(client)\n\n if args.raw:\n if not args.show_values:\n masked_env_vars = {k: \"*****\" for k, v in env_vars.items()}\n return masked_env_vars\n else:\n return env_vars\n else:\n if not env_vars:\n print(\"No environment variables found.\")\n else:\n for key, value in env_vars.items():\n print(f\"Name: {key}\")\n if args.show_values:\n print(f\"Value: {value}\")\n else:\n print(\"Value: *****\")\n print(\"---\")\n\n if not args.show_values:\n print(\"\\nNote: Values are hidden. Use --show-values or -s option to display them.\")\n\n\[email protected](\n argument(\"name\", help=\"Environment variable name\", type=str),\n argument(\"value\", help=\"Environment variable value\", type=str),\n usage=\"vastai create env-var \u003cname> \u003cvalue>\",\n help=\"Create a new user environment variable\",\n)\ndef create__env_var(args):\n \"\"\"Create a new environment variable for the current user.\"\"\"\n client = get_client(args)\n result = auth_api.create_env_var(client, name=args.name, value=args.value)\n\n if result.get(\"success\"):\n print(result.get(\"msg\", \"Environment variable created successfully.\"))\n else:\n print(f\"Failed to create environment variable: {result.get('msg', 'Unknown error')}\")\n\n\[email protected](\n argument(\"name\", help=\"Environment variable name to update\", type=str),\n argument(\"value\", help=\"New environment variable value\", type=str),\n usage=\"vastai update env-var \u003cname> \u003cvalue>\",\n help=\"Update an existing user environment variable\",\n)\ndef update__env_var(args):\n \"\"\"Update an existing environment variable for the current user.\"\"\"\n client = get_client(args)\n result = auth_api.update_env_var(client, name=args.name, value=args.value)\n\n if result.get(\"success\"):\n print(result.get(\"msg\", \"Environment variable updated successfully.\"))\n else:\n print(f\"Failed to update environment variable: {result.get('msg', 'Unknown error')}\")\n\n\[email protected](\n argument(\"name\", help=\"Environment variable name to delete\", type=str),\n usage=\"vastai delete env-var \u003cname>\",\n help=\"Delete a user environment variable\",\n)\ndef delete__env_var(args):\n \"\"\"Delete an environment variable for the current user.\"\"\"\n client = get_client(args)\n result = auth_api.delete_env_var(client, name=args.name)\n\n if result.get(\"success\"):\n print(result.get(\"msg\", \"Environment variable deleted successfully.\"))\n else:\n print(f\"Failed to delete environment variable: {result.get('msg', 'Unknown error')}\")\n\n\n# ---------------------------------------------------------------------------\n# Note: show__scheduled_jobs and delete__scheduled_job are in billing.py.\n# ---------------------------------------------------------------------------\n\n\n# ---------------------------------------------------------------------------\n# 2FA commands\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"code\", help=\"6-digit verification code from SMS or Authenticator app\", type=str),\n argument(\"-t\", \"--method-type\", choices=[\"sms\", \"totp\"], help=\"New 2FA Method type to activate\", type=str, default=\"totp\"),\n argument(\"--secret\", help=\"Secret token from setup process (required)\", type=str, required=True),\n argument(\"--phone-number\", help=\"Phone number for SMS method (E.164 format)\", type=str, default=None),\n argument(\"-l\", \"--label\", help=\"Label for the new 2FA method\", type=str, default=None),\n usage=\"vastai tfa activate CODE --secret SECRET [--method-type METHOD_TYPE] [--phone-number PHONE_NUMBER] [--label LABEL]\",\n help=\"Activate a new 2FA method by verifying the code\",\n epilog=deindent(f\"\"\"\n Complete the 2FA setup process by verifying your code.\n\n {'*'*120}\n 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.\n This is required to ensure that only you can add new 2FA methods to your account.\n {'*'*120}\n\n For TOTP (Authenticator app):\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 this command with the 6-digit code from your app and the secret token from step 1\n\n For SMS:\n 1. Run 'vastai tfa send-sms --phone-number \u003cPHONE_NUMBER>' to receive SMS and get secret token\n 2. Run this command with the code you received via SMS and the phone number it was sent to\n\n If this is your first 2FA method, backup codes will be generated and displayed.\n Save these backup codes in a secure location!\n\n Examples:\n vastai tfa activate --method-type totp --secret abc123def456 123456\n vastai tfa activate --method-type sms --secret abc123def456 --phone-number +12345678901 123456\n vastai tfa activate --method-type sms --secret abc123def456 --phone-number +12345678901 --label \"Work Phone\" 123456\n \"\"\"),\n)\ndef tfa__activate(args):\n \"\"\"Activate a new 2FA method by confirming the verification code.\"\"\"\n client = get_client(args)\n response_data = auth_api.tfa_activate(\n client, code=args.code, secret=args.secret, method_type=args.method_type,\n phone_number=args.phone_number, label=args.label,\n )\n method_name = \"SMS\" if args.phone_number or args.method_type == \"sms\" else \"TOTP (Authenticator App)\"\n print(f\"\\n{SUCCESS} {method_name} 2FA method activated successfully!\")\n if \"backup_codes\" in response_data:\n save_backup_codes(response_data[\"backup_codes\"])\n\n\[email protected](\n argument(\"-id\", \"--id-to-delete\", help=\"ID of the 2FA method to delete (see `vastai tfa status`)\", type=int, default=None),\n argument(\"-c\", \"--code\", mutex_group='code_grp', required=True, help=\"2FA code from your Authenticator app, SMS, or Email to authorize deletion\", type=str),\n argument(\"-t\", \"--method-type\", mutex_group=\"type_grp\", choices=[\"email\", \"sms\", \"totp\"],\n help=\"2FA Method type. Only use when you only have one method of this type\", type=str, default=None),\n argument(\"-s\", \"--secret\", help=\"Secret token (required for SMS or Email 2FA)\", type=str, default=None),\n argument(\"-bc\", \"--backup-code\", mutex_group='code_grp', required=True, help=\"One-time backup code (alternative to regular 2FA code)\", type=str, default=None),\n 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),\n usage=\"vastai tfa delete [--id-to-delete ID] [--code CODE] [--method-type METHOD_TYPE] [--secret SECRET] [--backup-code BACKUP_CODE] [--method-id ID]\",\n help=\"Remove a 2FA method from your account\",\n epilog=deindent(f\"\"\"\n Remove a 2FA method from your account.\n\n This action requires 2FA verification to prevent unauthorized removals.\n\n {'*'*120}\n NOTE: If you do not specify --id-to-delete, the system will attempt to delete the method you are using to authenticate.\n However please be advised, it is much safer to specify the ID to avoid confusion if you have multiple methods.\n {'*'*120}\n\n Use `vastai tfa status` to see your active methods and their IDs.\n\n Examples:\n # Delete method #123, authorize with email code and secret from `tfa send-email`\n vastai tfa delete --id-to-delete 123 --method-type email --secret abc123def456 -c 456789\n\n # Delete method #123, authorize with TOTP/Authenticator code\n vastai tfa delete --id-to-delete 123 --method-type totp --code 456789\n\n # Delete method #123, authorize with SMS and secret from `tfa send-sms`\n vastai tfa delete -id 123 --method-type sms --secret abc123def456 -c 456789\n\n # Delete method #123, authorize with backup code\n vastai tfa delete --id-to-delete 123 --backup-code ABCD-EFGH-IJKL\n\n # Delete method #123, specify which TOTP method to use if you have multiple\n vastai tfa delete -id 123 --method-id 456 -c 456789\n\n # Delete the TOTP method you are using to authenticate (use with caution)\n vastai tfa delete -c 456789\n \"\"\"),\n)\ndef tfa__delete(args):\n \"\"\"Remove a 2FA method from the user's account.\"\"\"\n if args.method_type in (\"sms\", \"email\") and not args.secret:\n print(f\"\\n{FAIL} Error: --secret is required for deletion authorization when using --method-type {args.method_type}.\")\n print(\"\\nPlease use: `vastai tfa send-sms` or `vastai tfa send-email` to get the missing secret and try again.\")\n return 1\n\n prompt = \"\\nAre you sure you want to delete this 2FA method? (y|n): \"\n if confirm_destructive_action(prompt) == False:\n print(\"Operation cancelled.\")\n return\n\n client = get_client(args)\n try:\n response_data = auth_api.tfa_delete(\n client, code=args.code, backup_code=args.backup_code,\n method_type=args.method_type, secret=args.secret, method_id=args.method_id,\n id_to_delete=args.id_to_delete,\n )\n print(f\"\\n{SUCCESS} 2FA method deleted successfully.\")\n\n if \"remaining_methods\" in response_data:\n remaining = response_data[\"remaining_methods\"]\n print(f\"\\nYou have {remaining} 2FA method{'s' if remaining != 1 else ''} remaining.\")\n else:\n print(f\"\\n{WARN} WARNING: You have removed all 2FA methods from your account.\")\n print(\"Your backup codes have been invalidated and 2FA is now fully disabled.\")\n\n except Exception as e:\n if hasattr(e, 'response'):\n handle_failed_tfa_verification(args, e)\n else:\n print(f\"Failed to delete 2FA method: {e}\")\n return 1\n\n\[email protected](\n argument(\"-c\", \"--code\", mutex_group='code_grp', required=True, help=\"2FA code from Authenticator app, SMS, or Email\", type=str),\n argument(\"-t\", \"--method-type\", mutex_group=\"type_grp\", choices=[\"email\", \"sms\", \"totp\"],\n help=\"2FA Method type. Only use when you only have one method of this type\", type=str, default=None),\n argument(\"-s\", \"--secret\", help=\"Secret token from previous login step (required for SMS or Email 2FA)\", type=str, default=None),\n argument(\"-bc\", \"--backup-code\", mutex_group='code_grp', required=True, help=\"One-time backup code (alternative to regular 2FA code)\", type=str, default=None),\n 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),\n usage=\"vastai tfa login [--code CODE] [--method-type METHOD_TYPE] [--secret SECRET] [--backup-code BACKUP_CODE]\",\n help=\"Complete 2FA login by verifying code\",\n epilog=deindent(\"\"\"\n Complete Two-Factor Authentication login by providing the 2FA code.\n\n For Email: Include the --method-type email flag and provide -s/--secret from the `tfa send-email` command response\n For TOTP: Include the --method-type totp and provide the 6-digit code from your Authenticator app\n For SMS: Include the --method-type sms flag and provide -s/--secret from the `tfa send-sms` command response\n For backup code: Use --backup-code instead of code (codes may only be used once)\n\n Examples:\n vastai tfa login --method-type totp -c 123456\n vastai tfa login --method-type sms --code 123456 --secret abc123def456\n vastai tfa login -t email -c 123456 -s abc123def456\n vastai tfa login --backup-code ABCD-EFGH-IJKL\n \"\"\"),\n)\ndef tfa__login(args):\n \"\"\"Complete 2FA login and store the session key.\"\"\"\n from vastai.cli.util import TFAKEY_FILE\n\n client = get_client(args)\n try:\n response_data = auth_api.tfa_login(\n client, code=args.code, backup_code=args.backup_code,\n method_type=args.method_type, secret=args.secret, method_id=args.method_id,\n )\n\n if \"session_key\" in response_data:\n session_key = response_data[\"session_key\"]\n if session_key != args.api_key:\n with open(TFAKEY_FILE, \"w\") as f:\n f.write(session_key)\n print(f\"{SUCCESS} 2FA login successful! Session key saved to {TFAKEY_FILE}\")\n else:\n print(f\"{SUCCESS} 2FA login successful! Your session key has been refreshed.\")\n\n # Display remaining backup codes if present\n if \"backup_codes_remaining\" in response_data:\n remaining = response_data[\"backup_codes_remaining\"]\n if remaining == 0:\n 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.\")\n elif remaining \u003c= 3:\n print(f\"{WARN} Warning: You only have {remaining} backup codes remaining. Consider regenerating them.\")\n else:\n print(f\"Backup codes remaining: {remaining}\")\n else:\n print(\"2FA login successful but a session key was not returned. Please check that you have an API Key that's properly set up\")\n except Exception as e:\n if hasattr(e, 'response'):\n handle_failed_tfa_verification(args, e)\n else:\n print(f\"2FA login failed: {e}\")\n return 1\n\n\[email protected](\n argument(\"-p\", \"--phone-number\", help=\"Phone number to receive SMS code (E.164 format, e.g., +1234567890)\", type=str, default=None),\n argument(\"-s\", \"--secret\", help=\"Secret token from the original 2FA login attempt\", type=str, required=True),\n usage=\"vastai tfa resend-sms --secret SECRET [--phone-number PHONE_NUMBER]\",\n help=\"Resend SMS 2FA code\",\n epilog=deindent(\"\"\"\n Resend the SMS verification code to your phone.\n\n This is useful if:\n \\u2022 You didn't receive the original SMS\n \\u2022 The code expired before you could use it\n \\u2022 You accidentally deleted the message\n\n You must provide the same secret token from the original request.\n\n Example:\n vastai tfa resend-sms --secret abc123def456\n \"\"\"),\n)\ndef tfa__resend_sms(args):\n \"\"\"Resend SMS 2FA code to the user's phone.\"\"\"\n client = get_client(args)\n response_data = auth_api.tfa_resend_sms(client, secret=args.secret, phone_number=args.phone_number)\n print(\"SMS code resent successfully!\")\n print(f\"\\n{response_data['msg']}\")\n\n\[email protected](\n argument(\"-c\", \"--code\", mutex_group='code_grp', required=True, help=\"2FA code from Authenticator app, SMS, or Email\", type=str),\n argument(\"-t\", \"--method-type\", mutex_group=\"type_grp\", choices=[\"email\", \"sms\", \"totp\"],\n help=\"2FA Method type. Only use when you only have one method of this type\", type=str, default=None),\n argument(\"-s\", \"--secret\", help=\"Secret token from previous login step (required for SMS or Email 2FA)\", type=str, default=None),\n argument(\"-bc\", \"--backup-code\", mutex_group='code_grp', required=True, help=\"One-time backup code (alternative to regular 2FA code)\", type=str, default=None),\n 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),\n usage=\"vastai tfa regen-codes [--code CODE] [--method-type METHOD_TYPE] [--secret SECRET] [--backup-code BACKUP_CODE] [--method-id ID]\",\n help=\"Regenerate backup codes for 2FA\",\n epilog=deindent(\"\"\"\n Generate a new set of backup codes for your account.\n\n This action requires 2FA verification to prevent unauthorized regeneration.\n\n WARNING: This will invalidate all existing backup codes!\n Any previously generated codes will no longer work.\n\n Backup codes are one-time use codes that allow you to log in\n if you lose access to your primary 2FA method (lost phone, etc).\n\n You should regenerate your backup codes if:\n \\u2022 You've used several codes and are running low\n \\u2022 You think your codes may have been compromised\n \\u2022 You lost your saved codes and need new ones\n\n Important: Save the new codes in a secure location immediately!\n They will not be shown again.\n\n Examples:\n vastai tfa regen-codes --code --method-type totp 123456\n vastai tfa regen-codes -c 123456 -t sms --secret abc123def456\n vastai tfa regen-codes --backup-code ABCD-EFGH-IJKL\n \"\"\"),\n)\ndef tfa__regen_codes(args):\n \"\"\"Regenerate backup codes for 2FA recovery.\"\"\"\n prompt = \"\\nThis will invalidate all existing backup codes. Continue? (y|n): \"\n if confirm_destructive_action(prompt) == False:\n print(\"Operation cancelled.\")\n return\n\n client = get_client(args)\n try:\n response_data = auth_api.tfa_regen_codes(\n client, code=args.code, backup_code=args.backup_code,\n method_type=args.method_type, secret=args.secret, method_id=args.method_id,\n )\n if \"backup_codes\" in response_data:\n save_backup_codes(response_data[\"backup_codes\"])\n else:\n print(f\"\\n{SUCCESS} Backup codes regenerated successfully!\")\n print(\"(No codes returned in response - this may be an error)\")\n except Exception as e:\n if hasattr(e, 'response'):\n handle_failed_tfa_verification(args, e)\n else:\n print(f\"Failed to regenerate backup codes: {e}\")\n return 1\n\n\[email protected](\n argument(\"-p\", \"--phone-number\", help=\"Phone number to receive SMS code (E.164 format, e.g., +1234567890)\", type=str, default=None),\n usage=\"vastai tfa send-sms [--phone-number PHONE_NUMBER]\",\n help=\"Request a 2FA SMS verification code\",\n epilog=deindent(\"\"\"\n Request a two-factor authentication code to be sent via SMS.\n\n If --phone-number is not provided, uses the phone number on your account.\n The secret token will be returned and must be used with 'vastai tfa activate'.\n\n Examples:\n vastai tfa send-sms\n vastai tfa send-sms --phone-number +12345678901\n \"\"\"),\n)\ndef tfa__send_sms(args):\n \"\"\"Request a 2FA SMS code to be sent to the user's phone.\"\"\"\n client = get_client(args)\n response_data = auth_api.tfa_send_sms(client, phone_number=args.phone_number)\n\n secret = response_data.get(\"secret\", \"\")\n print(\"SMS code sent successfully!\")\n print(f\" Secret token: {secret}\")\n print(f\"\\nOnce you receive the SMS code:\")\n phone_num = f\"--phone-number {args.phone_number}\" if args.phone_number else \"[--phone-number \u003cPHONE_NUMBER>]\"\n print(f\"\\n If you are setting up SMS 2FA for the first time, run:\")\n print(f\" vastai tfa activate --method-type sms --secret {secret} {phone_num} [--label \u003cLABEL>] \u003cCODE>\")\n print(f\"\\n Otherwise you can complete your 2FA log in with:\")\n print(f\" vastai tfa login --method-type sms --secret {secret} -c \u003cCODE>\\n\")\n\n\[email protected](\n usage=\"vastai tfa send-email\",\n help=\"Request a 2FA Email verification code\",\n epilog=deindent(\"\"\"\n Request a two-factor authentication code to be sent via Email.\n\n The secret token will be returned and must be used with 'vastai tfa activate'.\n\n Examples:\n vastai tfa send-email\n \"\"\"),\n)\ndef tfa__send_email(args):\n \"\"\"Request a 2FA email code to be sent to the user's email.\"\"\"\n client = get_client(args)\n response_data = auth_api.tfa_send_email(client)\n\n secret = response_data.get(\"secret\", \"\")\n print(f\"{SUCCESS} Email code sent successfully!\")\n print(f\" Secret token: {secret}\")\n print(f\"\\nOnce you receive the Email code:\")\n print(f\"\\n You can complete your 2FA log in with:\")\n print(f\" vastai tfa login --method-type email --secret {secret} -c \u003cCODE>\\n\")\n\n\ndef print_next_steps_after_new_method_auth():\n print(f\"\\nNext Steps:\"\n \"\\n To add a new SMS 2FA method:\"\n \"\\n 1. Run `vastai tfa send-sms --phone-number \u003cPHONE_NUMBER>` to receive SMS and get secret token\"\n \"\\n 2. Run `vastai tfa activate --method-type sms --secret \u003cSECRET> --phone-number \u003cPHONE_NUMBER> CODE` to activate the new method with the code you received via SMS\\n\"\n \"\\n To add a new TOTP (Authenticator app) 2FA method:\"\n \"\\n 1. Run `vastai tfa totp-setup` to get the manual key/QR code and secret\"\n \"\\n 2. Enter the manual key or scan the QR code with your Authenticator app\"\n \"\\n 3. Run `vastai tfa activate --method-type totp --secret \u003cSECRET> CODE` to activate the new method with the 6-digit code from your app\")\n\n\[email protected](\n argument(\"-c\", \"--code\", help=\"2FA code from Authenticator app, SMS, or Email\", type=str),\n argument(\"-s\", \"--secret\", help=\"Secret token from previous auth step\", type=str, default=None),\n argument(\"-t\", \"--method-type\", mutex_group=\"type_grp\", choices=[\"email\", \"sms\", \"totp\"],\n help=\"2FA Method type. Only use when you only have one method of this type\", type=str, default=\"email\"),\n argument(\"-bc\", \"--backup-code\", mutex_group='type_grp',\n help=\"One-time backup code (alternative to regular 2FA code)\", type=str, default=None),\n argument(\"-id\", \"--method-id\", mutex_group=\"type_grp\",\n help=\"2FA Method ID if you have more than one of the same type ('id' from `tfa status`)\", type=str, default=None),\n usage=\"vastai tfa auth-new {[--method-type METHOD_TYPE | --method-id ID | --backup-code BACKUP_CODE] | [--secret SECRET --code CODE]}\",\n help=\"Authorize your account to add a new 2FA method\",\n epilog=deindent(\"\"\"\n Authorize your account to add a new 2FA method by verifying via email or an existing method.\n\n This is a required step to ensure that only you can add new 2FA methods to your account.\n\n Step 1. Run command with your chosen verification method:\n - If you have an existing 2FA method set up, you can use '--backup-code BACKUP_CODE' to immediately authorize (skip Step 2)\n - 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)\n - 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\n\n Step 2. When prompted, enter the 2FA code from your 2FA method of choice to confirm authorization.\n\n Note:\n If you exit the command before being able to enter the code, you can run this command again\n with --secret SECRET and --code CODE to complete the authorization step as long as the code has not expired.\n\n Examples:\n # Initiating New Method Authorization\n vastai tfa auth-new (method type is email by default)\n vastai tfa auth-new --method-type totp\n vastai tfa auth-new -t sms\n vastai tfa auth-new --method-id 456\n vastai tfa auth-new --backup-code ABCD-EFGH-IJKL\n\n # Completing Authorization with code and secret if not completed in previous run\n vastai tfa auth-new --secret abc123def456 --code 123456\n \"\"\"),\n)\ndef tfa__auth_new(args):\n \"\"\"Authorize the user to add a new 2FA method by verifying with an existing method.\"\"\"\n client = get_client(args)\n\n secret, code = args.secret, args.code\n if not secret and not code:\n try:\n response_data = auth_api.tfa_auth_new(\n client, backup_code=args.backup_code,\n method_type=args.method_type, method_id=args.method_id,\n )\n except Exception as e:\n if hasattr(e, 'response'):\n handle_failed_tfa_verification(args, e)\n else:\n print(f\"Authorization failed: {e}\")\n return 1\n\n if args.backup_code and response_data.get(\"msg\") == \"Authorization successful.\":\n print(f\"\\n{SUCCESS} Successfully authorized account for adding new 2FA method using backup code\")\n print_next_steps_after_new_method_auth()\n return 0\n\n secret = response_data.get(\"secret\")\n if not secret:\n print(f\"\\n{FAIL} Error: No secret token received for authorization. Please try again.\")\n return 1\n\n print(f\"\\n{SUCCESS} Authorization initiated successfully.\")\n print(f\"2FA Secret: {secret}\")\n try:\n code = input(\"Enter 2FA code to complete authorization: \").strip()\n except (EOFError, KeyboardInterrupt):\n print(\"\\nOperation cancelled.\")\n print(\"You can still complete this authorization later by running:\"\n f\"\\n vastai tfa auth-new --secret {secret} --code \u003cCODE>\")\n return 1\n\n try:\n auth_api.tfa_auth_new(client, code=code, secret=secret)\n print(f\"\\n{SUCCESS} Successfully authorized account for adding new 2FA method!\")\n print_next_steps_after_new_method_auth()\n except Exception as e:\n if hasattr(e, 'response'):\n handle_failed_tfa_verification(args, e)\n else:\n print(f\"Authorization failed: {e}\")\n print(f\"\\n{FAIL} Authorization failed. Please try again.\")\n return 1\n\n\[email protected](\n help=\"Shows the current 2FA status and configured methods\",\n epilog=deindent(\"\"\"\n Show the current 2FA status for your account, including:\n \\u2022 Whether or not 2FA is enabled\n \\u2022 A list of active 2FA methods\n \\u2022 The number of backup codes remaining (if 2FA is enabled)\n \"\"\")\n)\ndef tfa__status(args):\n \"\"\"Show the current 2FA status for the user.\"\"\"\n client = get_client(args)\n response_data = auth_api.tfa_status(client)\n\n if args.raw:\n return response_data\n\n tfa_enabled = response_data.get(\"tfa_enabled\", False)\n methods = response_data.get(\"methods\", [])\n backup_codes_remaining = response_data.get(\"backup_codes_remaining\", 0)\n\n if not tfa_enabled or not methods:\n print(f\"{WARN} No active 2FA methods found\")\n else:\n print(f\"2FA Status: Enabled {SUCCESS}\")\n print(f\"\\nActive 2FA Methods:\")\n display_tfa_methods(methods)\n print(f\"\\nBackup codes remaining: {backup_codes_remaining}\")\n\n\[email protected](\n usage=\"vastai tfa totp-setup\",\n help=\"Generate TOTP secret and QR code for Authenticator app setup\",\n epilog=deindent(\"\"\"\n Set up TOTP (Time-based One-Time Password) 2FA using an Authenticator app.\n\n This command generates a new TOTP secret and displays:\n \\u2022 A QR code (for scanning with your app)\n \\u2022 A manual entry key (for typing into your app)\n \\u2022 A secret token (needed for the next step)\n\n Workflow:\n 1. Run this command to generate the TOTP secret\n 2. Add the account to your Authenticator app by either:\n \\u2022 Scanning the displayed QR code, OR\n \\u2022 Manually entering the key shown\n 3. Once added, your app will display a 6-digit code\n 4. Complete setup by running:\n vastai tfa activate --method-type totp --secret \u003cSECRET> \u003cCODE>\n\n Supported Authenticator Apps:\n \\u2022 Google Authenticator\n \\u2022 Microsoft Authenticator\n \\u2022 Authy\n \\u2022 1Password\n \\u2022 Any TOTP-compatible app\n\n Example:\n vastai tfa totp-setup\n \"\"\"),\n)\ndef tfa__totp_setup(args):\n \"\"\"Generate a TOTP secret and QR code for setting up Authenticator app 2FA.\"\"\"\n client = get_client(args)\n response_data = auth_api.tfa_totp_setup(client)\n\n if args.raw:\n return response_data\n\n secret = response_data.get(\"secret\", \"\")\n provisioning_uri = response_data.get(\"provisioning_uri\", \"\")\n\n print(\"\\n\" + \"=\" * 60)\n print(\"TOTP (Authenticator App) 2FA Setup\")\n print(\"=\" * 60)\n\n print(\"\\nScan this QR code with your Authenticator app:\\n\")\n\n try:\n import qrcode\n qr = qrcode.QRCode(border=2)\n qr.add_data(provisioning_uri)\n qr.make()\n qr.print_ascii(tty=True)\n except ImportError:\n print(\" [QR code display requires 'qrcode' package]\")\n print(f\" Install with: pip install qrcode\")\n print(f\"\\n Or manually enter this URI in your app:\")\n print(f\" {provisioning_uri}\")\n\n print(\"\\nOR Manual Entry Key (type this into your Authenticator app):\")\n print(f\" {secret}\")\n\n print(\"\\nNext Steps:\")\n print(\" 1. Your Authenticator app should now display a 6-digit code\")\n print(\" 2. Complete setup by running:\")\n print(f\" vastai tfa activate --method-type totp --secret {secret} \u003cCODE>\")\n print(\"\\n\" + \"=\" * 60 + \"\\n\")\n\n\[email protected](\n argument(\"method_id\", metavar=\"METHOD_ID\", help=\"ID of the 2FA method to update (see `vastai tfa status`)\", type=int),\n argument(\"-l\", \"--label\", help=\"New label/name for this 2FA method\", type=str, default=None),\n argument(\"-p\", \"--set-primary\", help=\"Set this method as the primary/default 2FA method\", default=None),\n usage=\"vastai tfa update METHOD_ID [--label LABEL] [--set-primary]\",\n help=\"Update a 2FA method's settings\",\n epilog=deindent(\"\"\"\n Update the label or primary status of a 2FA method.\n\n The label is a friendly name to help you identify different methods\n (e.g. \"Work Phone\", \"Personal Authenticator\").\n\n The primary method is your preferred/default 2FA method.\n\n Examples:\n vastai tfa update 123 --label \"Work Phone\"\n vastai tfa update 456 --set-primary\n vastai tfa update 789 --label \"Backup Authenticator\" --set-primary\n \"\"\"),\n)\ndef tfa__update(args):\n \"\"\"Update settings for an existing 2FA method.\"\"\"\n label = args.label\n set_primary = args.set_primary\n\n if set_primary is not None:\n if set_primary.lower() in {'true', 't'}:\n set_primary = True\n elif set_primary.lower() in {'false', 'f'}:\n set_primary = False\n else:\n print(\"Error: --set-primary must be \u003ct|true> or \u003cf|false>\")\n return\n\n if label is None and set_primary is None:\n print(\"Error: You must specify at least one field to update (--label or --set-primary)\")\n return 1\n\n client = get_client(args)\n try:\n response_data = auth_api.tfa_update(client, method_id=args.method_id, label=label, set_primary=set_primary)\n except ValueError as e:\n print(f\"Error: {e}\")\n return 1\n\n if args.raw:\n return response_data\n\n print(\"\\n2FA method updated successfully!\")\n if args.label:\n print(f\" New label: {args.label}\")\n if set_primary is not None:\n print(f\" Set as primary method = {set_primary}\")\n","content_type":"text/x-python; charset=utf-8","language":"python","size":42156,"content_sha256":"52c494f657c10899384433aa027f5ad914e4890ac1d74e53c685a796a37d4fcf"},{"filename":"cli/commands/benchmarks.py","content":"\"\"\"\nA CLI command to benchmark a template against multiple GPUs in parallel.\nThe template's pyworker benchmark determines the measured perf (ex. tokens/$).\n\nvastai run benchmarks\n\"\"\"\n\nimport argparse\nimport atexit\nimport functools\nimport json\nimport math\nimport random\nimport re\nimport signal\nimport statistics\nimport sys\nimport threading\nimport time\nimport uuid\nfrom concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait\n\nfrom rich.console import Console\nfrom rich.live import Live\nfrom rich.table import Table\n\nfrom vastai import VastAI\nfrom vastai.cli.parser import argument\nfrom vastai.cli.display import deindent\nfrom vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401 (used by test conftest)\nfrom vastai.cli.util import _get_gpu_types\nfrom vastai.data.query import OP_TO_STR\nfrom vastai.data import query as _query_consts\n\n\nparser = _get_parser()\n\n\n# Default GPU sweep set + per-card VRAM fallback (MB). Catalog gpu_ram_mb\n# overrides these values when populated; the keys stay the editorial sweep set.\nDEFAULT_GPUS = {\n \"RTX 5090\": 32607,\n \"RTX 4090\": 24564,\n \"RTX 3090\": 24576,\n \"RTX A6000\": 49140,\n}\n\n# One worker per GPU, kept alive while idle so we can read its measured perf.\nENDPOINT_CONFIG = {\"cold_workers\": 1, \"max_workers\": 1, \"min_load\": 1.0}\n\nRENTAL_TIMEOUT = 120 \nDEFAULT_BENCHMARK_TIMEOUT = 30 * 60 \nWORKER_POLL_INTERVAL = 10.0 \nPARALLEL_SUBMIT_DELAY = 4.0 # spacing between parallel endpoint submits\nMAX_ABANDONS = 3 # bail after this many workers vanish without one reaching idle\n\n# states in which the autoscaler shouldnt retry the worker \nTERMINAL_STATES = {\"stopped\", \"destroying\", \"unavail\"} \nTERMINAL_GRACE_SECONDS = 30\nTERMINAL_STATUSES = {\"done\", \"skipped\", \"timeout\", \"failed\", \"no_worker\"}\n\n\ndef format_time_elapsed(seconds):\n m, s = divmod(int(seconds), 60)\n return f\"{m}:{s:02d}\"\n\n\[email protected]_cache(maxsize=1)\ndef _canonical_gpu_names():\n \"\"\"Map of lower-cased GPU name -> canonical name, for resolving user input\n like \"rtx_4090\"/\"RTX 4090\" to the canonical \"RTX 4090\".\n\n Union of the hardcoded query.py constants (fallback floor) with the\n /api/v0/gpu_types/ catalog (picks up newly-onboarded SKUs). On a catalog\n fetch failure the constants alone are used, so resolution never gets worse.\n \"\"\"\n canonical = {\n v.lower(): v for k, v in vars(_query_consts).items()\n if isinstance(v, str) and not k.startswith(\"_\")\n }\n for gpu_type in (_get_gpu_types() or []):\n name = gpu_type.get(\"canonical_name\")\n if name:\n canonical[name.lower()] = name\n return canonical\n\n\[email protected]_cache(maxsize=1)\ndef _catalog_vram_map():\n \"\"\"Map canonical_name -> gpu_ram_mb for entries the catalog has populated.\n Empty when the catalog is unreachable or the column is unbackfilled.\n \"\"\"\n return {\n gpu_type[\"canonical_name\"]: gpu_type[\"gpu_ram_mb\"]\n for gpu_type in (_get_gpu_types() or [])\n if gpu_type.get(\"canonical_name\") and gpu_type.get(\"gpu_ram_mb\")\n }\n\n\ndef _per_card_vram_mb(gpu_name):\n \"\"\"Per-card VRAM in MB: catalog gpu_ram_mb if populated, else the hardcoded\n DEFAULT_GPUS value. None when neither source knows the GPU.\n \"\"\"\n return _catalog_vram_map().get(gpu_name) or DEFAULT_GPUS.get(gpu_name)\n\n\ndef parse_gpu_spec(token, default_num_gpus):\n \"\"\"Parse a GPU token into (canonical_gpu_name, num_gpus).\n Example: 2x RTX_4090\" -> (\"RTX 4090\", 2). \n \"\"\"\n\n # so \"rtx_4090\", \"RTX 4090\", \"Rtx_4090\" all map to \"RTX 4090\"\n canonical = _canonical_gpu_names()\n\n token = token.strip()\n # string to tuple. \"2x RTX_4090\" to (\"RTX 4090\", 2)\n m = re.match(r\"^(\\d+)[\\s_]*x[\\s_]*(.+)$\", token, re.IGNORECASE)\n if m:\n raw_name, n = m.group(2).strip().replace(\"_\", \" \"), int(m.group(1))\n else:\n raw_name, n = token.replace(\"_\", \" \"), default_num_gpus\n return (canonical.get(raw_name.lower(), raw_name), n)\n\n\ndef format_filter_query(filters):\n \"\"\"Render extra_filters dict into comparison string.\n e.g. {\"gpu_ram\": {\"gte\": 16000}} -> \"gpu_ram>=16000\"\n \"\"\"\n parts = []\n for field, ops in (filters or {}).items():\n if not isinstance(ops, dict):\n continue\n for op, value in ops.items():\n sym = OP_TO_STR.get(op, op)\n parts.append(f\"{field}{sym}{value}\")\n return \", \".join(parts)\n\n\n# ------------------------------------------------------------------------------------\n# Pre-flight helpers: check whether each (template, gpu_name, num_gpus) is rentable\n# before we benchmark it, so a run can't fail late due to 0 matching offers. On a\n# miss they build a user-friendly explanation of which filter blocked it.\n# ------------------------------------------------------------------------------------\n\n\ndef has_matching_offer(vast, *, gpu_name, num_gpus, extra_filters):\n \"\"\"True if any verified+rentable offer matches the template's filters\"\"\"\n query = dict(extra_filters or {})\n query[\"gpu_name\"] = {\"eq\": gpu_name}\n query[\"num_gpus\"] = {\"eq\": num_gpus}\n return bool(vast.search_offers(query=query, limit=1))\n\n\ndef format_skip_message(vast, *, gpu_name, num_gpus, extra_filters):\n \"\"\"Build a user-friendly explanation for why a GPU has 0 matching offers.\n Calls search_offers to find the blocking filter.\n \"\"\"\n if not extra_filters:\n return (f\"no offers for {gpu_name} (num_gpus={num_gpus})\")\n\n diag = find_blockers(\n vast, gpu_name=gpu_name, num_gpus=num_gpus,\n extra_filters=extra_filters,\n )\n if diag[\"base_count\"] == 0:\n return (f\"no offers for {gpu_name} (num_gpus={num_gpus}), \"\n f\"independent of template filters\")\n\n blockers = diag[\"single_blockers\"]\n if blockers:\n sample = diag.get(\"sample_offer\") or {}\n blocker_lines = []\n for key in blockers:\n ops = extra_filters.get(key)\n if not isinstance(ops, dict):\n blocker_lines.append(f\" {key}\")\n continue\n gpu_value = sample.get(key)\n for op, threshold in ops.items():\n if check_template_filter(gpu_value, op, threshold) is True:\n continue # search-engine quirk, skip\n sym = OP_TO_STR.get(op, op)\n if gpu_value is None:\n detail = \"no sample available to check\"\n else:\n detail = f\"{gpu_name} has {gpu_value}\"\n blocker_lines.append(f\" {key}{sym}{threshold}: {detail}\")\n\n if blocker_lines:\n plural = \"s\" if len(blocker_lines) > 1 else \"\"\n lines = [f\"blocked by template filter{plural}:\"] + blocker_lines\n # Only hint num_gpus when gpu_total_ram is the sole blocker; otherwise bumping it won't fix the run.\n if blockers == [\"gpu_total_ram\"]:\n ops = extra_filters.get(\"gpu_total_ram\") or {}\n if ops:\n op, value = next(iter(ops.items()))\n raw = min_gpus_for_ram(\n value, op, diag.get(\"per_card_gpu_ram\"))\n if raw and raw > num_gpus:\n viable = get_gpu_chunk_size(raw)\n lines.append(\n f\" hint: try {viable}x {gpu_name} \"\n f\"(host total then satisfies {value})\")\n return \"\\n\".join(lines)\n # All \"blockers\" turned out to be quirks; fall through to generic failure message\n\n return (\"0 offers match. No single filter is the culprit; the combination \"\n f\"of template filters excludes {gpu_name} \"\n f\"({format_filter_query(extra_filters)})\")\n\n\ndef find_blockers(vast, *, gpu_name, num_gpus, extra_filters):\n \"\"\"Find which template filter blocks this GPU from having offers.\n\n Bails on the first real blocker. Returns a dict\n {base_count, single_blockers, per_card_gpu_ram, sample_offer}.\n \"\"\"\n base_query = {\"gpu_name\": {\"eq\": gpu_name},\n \"num_gpus\": {\"eq\": num_gpus}}\n base_offers = vast.search_offers(query=base_query, limit=1)\n base = len(base_offers)\n sample = base_offers[0] if base_offers else None\n per_card = sample.get(\"gpu_ram\") if sample else None\n blockers = []\n if base > 0:\n for key, ops in (extra_filters or {}).items():\n # if GPU satisfies the given filter, its not a blocker so skip\n if isinstance(ops, dict) and sample is not None:\n gpu_value = sample.get(key)\n if gpu_value is not None and all(\n check_template_filter(gpu_value, op, threshold) is True\n for op, threshold in ops.items()\n ):\n continue\n if not has_matching_offer(\n vast, gpu_name=gpu_name, num_gpus=num_gpus,\n extra_filters={key: ops},\n ):\n blockers.append(key)\n break\n return {\"base_count\": base, \"single_blockers\": blockers,\n \"per_card_gpu_ram\": per_card, \"sample_offer\": sample}\n\n\ndef check_template_filter(gpu_value, op, threshold):\n \"\"\"Checks if gpu_value matches the template filters (e.g. gpu_total_ram>=16000)\"\"\"\n if gpu_value is None or threshold is None:\n return None\n try:\n a, t = float(gpu_value), float(threshold)\n except (TypeError, ValueError):\n return None\n if op == \"gt\": return a > t\n if op == \"gte\": return a >= t\n if op == \"lt\": return a \u003c t\n if op == \"lte\": return a \u003c= t\n if op == \"eq\": return a == t\n if op == \"neq\": return a != t\n return None\n\n\ndef min_gpus_for_ram(threshold, op, per_card_gpu_ram):\n \"\"\"Smallest ``num_gpus`` such that ``num_gpus * per_card_gpu_ram``\n satisfies the template's ``gpu_total_ram`` filter. Returns None if any\n input is missing or the operator isn't a comparison we can solve.\n \"\"\"\n if not per_card_gpu_ram or not threshold:\n return None\n if op == \"gt\":\n return int(threshold // per_card_gpu_ram) + 1\n if op == \"gte\":\n return int(math.ceil(threshold / per_card_gpu_ram))\n return None\n\n\ndef get_gpu_chunk_size(n):\n \"\"\"Round n up to the next common GPU chunk size (1, 2, 4, 8). Returns n\n unchanged for >8 since that's rare enough we'd rather be honest than guess.\n \"\"\"\n for c in (1, 2, 4, 8):\n if n \u003c= c:\n return c\n return n\n\n\ndef pick_num_gpus(gpu_name, extra_filters):\n \"\"\"Smallest num_gpus where ``num_gpus * per_card_ram`` satisfies the\n template's ``gpu_total_ram`` filter, rounded up to a chassis size.\n Falls back to 1 when the template has no such filter or one card\n already satisfies it.\n \"\"\"\n per_card = _per_card_vram_mb(gpu_name)\n if per_card is None:\n return 1\n ops = (extra_filters or {}).get(\"gpu_total_ram\") or {}\n if not ops:\n return 1\n op, threshold = next(iter(ops.items()))\n raw_min = min_gpus_for_ram(threshold, op, per_card)\n if raw_min is None or raw_min \u003c= 1:\n return 1\n return get_gpu_chunk_size(raw_min)\n\n\nDEFAULT_CACHE_MAX_AGE_DAYS = 30\nMAX_CACHE_ROWS = 100 # newest rows per template column; caps popular template+GPU combos\n\n\ndef lookup_cached_benchmark(vast, *, gpu_name, num_gpus, template_hash,\n template_id, max_age_days):\n \"\"\"Recent reported benchmarks for this exact spec (median + spread), or None.\n\n Queried once per template column rather than filtered client-side: a\n workergroup is created with either template_id or template_hash, so a row\n can carry either one. Both columns are indexed with last_update on the\n benchmarks table, so each query is index-backed instead of pulling every\n perf row for the GPU and discarding the non-matching ones.\n \"\"\"\n base = {\n \"type\": {\"eq\": \"perf\"},\n \"gpu_name\": {\"eq\": gpu_name},\n \"num_gpus\": {\"eq\": num_gpus},\n \"last_update\": {\"gte\": time.time() - max_age_days * 86400},\n }\n rows = []\n seen_ids = set()\n for col, value in ((\"template_hash\", template_hash), (\"template_id\", template_id)):\n if not value:\n continue\n found = vast.search_benchmarks(\n query=dict(base, **{col: {\"eq\": value}}),\n order=[{\"col\": \"last_update\", \"dir\": \"desc\"}],\n limit=MAX_CACHE_ROWS,\n )\n for r in found if isinstance(found, list) else []:\n if not isinstance(r, dict):\n continue\n # A row carrying both columns comes back from both queries; dedupe on\n # id, but keep rows without one rather than collapsing them together.\n row_id = r.get(\"id\")\n if row_id is not None:\n if row_id in seen_ids:\n continue\n seen_ids.add(row_id)\n rows.append(r)\n matched = [\n r for r in rows\n if isinstance(r.get(\"value\"), (int, float)) and r[\"value\"] > 0\n ]\n if not matched:\n return None\n values = [r[\"value\"] for r in matched]\n newest = max((r.get(\"last_update\") or 0) for r in matched)\n # skip null rows (from before we tracked dph_base in benchmarks)\n dphs = [\n r[\"dph_base\"]\n for r in matched\n if isinstance(r.get(\"dph_base\"), (int, float)) and r[\"dph_base\"] > 0\n ]\n # use p5/p95 for the range\n if len(values) >= 2:\n qs = statistics.quantiles(values, n=20, method=\"inclusive\")\n low, high = qs[0], qs[-1]\n else:\n low = high = values[0]\n return {\n \"median\": statistics.median(values),\n \"low\": low,\n \"high\": high,\n \"n\": len(matched),\n \"age_days\": max(0.0, (time.time() - newest) / 86400),\n \"median_dph_base\": statistics.median(dphs) if dphs else None,\n }\n\n\ndef update_worker_states(worker_states, current_workers, gpu_name):\n \"\"\"Update worker_states with current poll, and print only on worker\n rotation (the live rich table already shows current status / elapsed).\n Tracks state_started per worker for the terminal-state grace check.\n Returns the count of workers abandoned during this poll.\n \"\"\"\n now = time.monotonic()\n abandoned = 0\n waiting_key = \"__waiting__\" # non-int sentinel: \"no worker rented yet\"\n\n if not current_workers:\n worker_states.setdefault(waiting_key, {\"started\": now})\n return abandoned\n\n worker_states.pop(waiting_key, None)\n\n seen = set()\n for w in current_workers:\n wid = w.get(\"id\")\n if wid is None:\n continue\n seen.add(wid)\n status = str(w.get(\"status\", \"\")).lower()\n prev = worker_states.get(wid)\n if prev is None:\n worker_states[wid] = {\"status\": status, \"state_started\": now}\n continue\n if prev[\"status\"] != status:\n prev[\"status\"] = status\n prev[\"state_started\"] = now\n\n for wid in list(worker_states):\n if not isinstance(wid, int):\n continue # skip the non-int waiting sentinel\n if wid not in seen:\n # Worker abandoned; live table only shows current worker, so rotation is otherwise invisible.\n prev = worker_states[wid]\n elapsed = format_time_elapsed(now - prev[\"state_started\"])\n print(f\"[{gpu_name}] worker {wid} abandoned \"\n f\"(last state={prev['status']} for {elapsed})\",\n file=sys.stderr)\n del worker_states[wid]\n abandoned += 1\n return abandoned\n\n\ndef update_row(class_states, row_id, **fields):\n \"\"\"Update one GPU's row in the live-table state dict. Sets\n ``run_started`` on first non-``queued`` status and ``run_ended`` on\n terminal status (so elapsed freezes); each thread only writes its own\n key so no lock needed.\n \"\"\"\n if class_states is None:\n return\n cur = class_states.setdefault(row_id, {})\n new_status = fields.get(\"status\")\n now_ts = time.monotonic()\n if new_status and new_status != cur.get(\"status\"):\n if \"run_started\" not in cur and new_status not in (\"queued\", \"cached\"):\n cur[\"run_started\"] = now_ts\n if new_status in TERMINAL_STATUSES:\n cur[\"run_ended\"] = now_ts\n else:\n cur.pop(\"run_ended\", None) # recovered from a transient terminal status; unfreeze elapsed\n cur.update(fields)\n\n\ndef render_table(class_states):\n \"\"\"Render the per-GPU progress as a rich.Table for the live display.\"\"\"\n table = Table(show_header=True, header_style=\"bold cyan\")\n table.add_column(\"GPU\", no_wrap=True)\n table.add_column(\"Status\", no_wrap=True)\n table.add_column(\"Endpoint\", justify=\"right\", no_wrap=True)\n table.add_column(\"Worker\", justify=\"right\", no_wrap=True)\n table.add_column(\"Elapsed\", justify=\"right\", no_wrap=True)\n table.add_column(\"Perf\", justify=\"right\", no_wrap=True)\n table.add_column(\"$/hr\", justify=\"right\", no_wrap=True)\n table.add_column(\"Perf/$/hr\", justify=\"right\", no_wrap=True)\n now = time.monotonic()\n for gpu in sorted(class_states):\n s = class_states[gpu]\n status = s.get(\"status\") or \"-\"\n # Skipped rows are noise during the run; reason already printed at pre-flight, final summary still includes them.\n if status == \"skipped\":\n continue\n run_started = s.get(\"run_started\")\n run_ended = s.get(\"run_ended\")\n # Blank for not-yet-running, frozen for terminal, live otherwise.\n if run_started is None:\n elapsed_str = \"-\"\n elif run_ended is not None:\n elapsed_str = format_time_elapsed(run_ended - run_started)\n else:\n elapsed_str = format_time_elapsed(now - run_started)\n perf = s.get(\"perf\")\n dph = s.get(\"dph\")\n pps = (perf / dph) if (perf and dph) else None\n table.add_row(\n gpu,\n status,\n str(s.get(\"endpoint_id\") or \"-\"),\n str(s.get(\"worker_id\") or \"-\"),\n elapsed_str,\n f\"{perf:.1f}\" if perf else \"-\",\n f\"${dph:.3f}\" if dph else \"-\",\n f\"{pps:.1f}\" if pps else \"-\",\n )\n return table\n\n\ndef benchmark_gpu(vast, *, gpu_name, num_gpus, timeout,\n workergroups, endpoints,\n template_hash=None, template_id=None,\n auto_instance=None, autoscaler_url=None,\n class_states=None, stop_event=None):\n \"\"\"Rent one instance for ``gpu_name``, poll until idle + measured_perf,\n tear down. Each call owns its own endpoint + workergroup so it can run\n in parallel safely. Returns ``(gpu_name, num_gpus, status, perf, err, dph_base)``.\n \"\"\"\n endpoint_id = None\n workergroup_id = None\n # uuid suffix keeps the name unique when benchmarking the same GPU in parallel;\n # no parens/shell chars, the backend rejects them in endpoint_name.\n endpoint_name = f\"benchmark {num_gpus}x {gpu_name} {uuid.uuid4().hex[:8]}\"\n 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\n start = time.monotonic()\n update_row(class_states, row_id, status=\"provisioning\")\n search = f\"gpu_name={gpu_name.replace(' ', '_')} num_gpus={num_gpus}\"\n if autoscaler_url:\n status_vast = VastAI(api_key=vast.client.api_key, server_url=autoscaler_url, retry=vast.client.retry)\n else:\n status_vast = vast\n\n try:\n endpoint_kwargs = dict(ENDPOINT_CONFIG, endpoint_name=endpoint_name)\n if auto_instance is not None:\n endpoint_kwargs[\"auto_instance\"] = auto_instance\n ep_resp = vast.create_endpoint(**endpoint_kwargs)\n # autoscaler returns the new id under either \"result\" or \"id\" depending on resource\n endpoint_id = ep_resp.get(\"result\", ep_resp.get(\"id\")) if isinstance(ep_resp, dict) else None\n if not isinstance(endpoint_id, int):\n return (gpu_name, num_gpus, \"error\", None,\n f\"create_endpoint returned no id: {ep_resp!r}\", None)\n endpoints.add(endpoint_id)\n update_row(class_states, row_id, endpoint_id=endpoint_id)\n\n # cold_workers and min_load are set on ENDPOINT_CONFIG, not here; the autoscaler reads\n # both from the endpoint group, not the workergroup (verified via autoscaler.cpp).\n workergroup_kwargs = dict(\n endpoint_id=endpoint_id,\n endpoint_name=endpoint_name,\n search_params=search,\n )\n if template_id is not None:\n workergroup_kwargs[\"template_id\"] = template_id\n elif template_hash is not None:\n workergroup_kwargs[\"template_hash\"] = template_hash\n if auto_instance is not None:\n workergroup_kwargs[\"auto_instance\"] = auto_instance\n resp = vast.create_workergroup(**workergroup_kwargs)\n workergroup_id = resp.get(\"result\", resp.get(\"id\")) if isinstance(resp, dict) else None\n if not isinstance(workergroup_id, int):\n return (gpu_name, num_gpus, \"error\", None,\n f\"create_workergroup returned no id: {resp!r}\", None)\n workergroups.add(workergroup_id)\n update_row(class_states, row_id, status=\"waiting_for_worker\")\n worker_states = {}\n dph_by_worker = {} # cache so we don't re-fetch dph on every poll\n abandoned_count = 0 # bail when this hits MAX_ABANDONS; autoscaler is renting + giving up repeatedly, no progress\n\n while time.monotonic() - start \u003c timeout:\n if stop_event is not None and stop_event.is_set():\n update_row(class_states, row_id, status=\"aborted\")\n return (gpu_name, num_gpus, \"aborted\", None, \"user aborted\", None)\n # calls get endpoint workers which includes measured perf(get workergroup workers doesnt)\n resp = status_vast.get_endpoint_workers(endpoint_id)\n if isinstance(resp, list):\n workers = resp\n elif isinstance(resp, dict) and isinstance(resp.get(\"workers\"), list):\n workers = resp[\"workers\"]\n else:\n workers = []\n abandoned_count += update_worker_states(worker_states, workers, gpu_name)\n if abandoned_count >= MAX_ABANDONS:\n update_row(class_states, row_id, status=\"failed\")\n return (gpu_name, num_gpus, \"failed\", None,\n f\"autoscaler abandoned {abandoned_count} workers without \"\n f\"any reaching idle (likely template + GPU config not loadable \"\n f\"on these hosts)\", None)\n if workers:\n primary = workers[0]\n primary_id = primary.get(\"id\")\n fields = {\n \"status\": str(primary.get(\"status\") or \"?\").lower(),\n \"worker_id\": primary_id,\n }\n # Fetch dph once per new worker_id; surfaces $/hr in the live table as soon as the autoscaler rents (before model loads).\n if primary_id and primary_id not in dph_by_worker:\n try:\n inst = vast.show_instance(id=primary_id)\n dph_by_worker[primary_id] = inst.get(\"dph_base\")\n except Exception as e:\n # Cache None so we don't re-attempt every poll, but log so a real bug isn't silently swallowed.\n dph_by_worker[primary_id] = None\n print(f\"[warn] failed to fetch dph for worker {primary_id}: {e}\",\n file=sys.stderr)\n if primary_id and dph_by_worker.get(primary_id) is not None:\n fields[\"dph\"] = dph_by_worker[primary_id]\n update_row(class_states, row_id, **fields)\n # measured_perf is only real once status==idle; before that it's a dlperf placeholder.\n ready = [w for w in workers\n if str(w.get(\"status\", \"\")).lower() == \"idle\"\n and (w.get(\"measured_perf\") or 0) > 0]\n if ready:\n # dph was already fetched and cached when this worker_id first appeared in polling.\n worker_id = ready[0].get(\"id\")\n dph = dph_by_worker.get(worker_id) if worker_id else None\n update_row(class_states, row_id, status=\"done\",\n perf=ready[0][\"measured_perf\"], dph=dph)\n return (gpu_name, num_gpus, \"ok\", ready[0][\"measured_perf\"], None, dph)\n # Fail fast: every worker stuck terminal past the grace period, none reached idle.\n now_ts = time.monotonic()\n terminal_workers_long = workers and all(\n str(w.get(\"status\", \"\")).lower() in TERMINAL_STATES\n and isinstance(w.get(\"id\"), int)\n and w[\"id\"] in worker_states\n and now_ts - worker_states[w[\"id\"]][\"state_started\"]\n > TERMINAL_GRACE_SECONDS\n for w in workers\n )\n if terminal_workers_long:\n states = sorted({str(w.get(\"status\", \"\")).lower()\n for w in workers})\n update_row(class_states, row_id, status=\"failed\")\n return (gpu_name, num_gpus, \"failed\", None,\n f\"all workers terminal ({', '.join(states)}) for \"\n f\">{TERMINAL_GRACE_SECONDS}s without reaching idle; \"\n f\"autoscaler not rotating\", None)\n # Fail fast if no real worker (int id) has ever appeared.\n if (not any(isinstance(k, int) for k in worker_states)\n and time.monotonic() - start > RENTAL_TIMEOUT):\n update_row(class_states, row_id, status=\"no_worker\")\n return (gpu_name, num_gpus, \"no_worker\", None,\n f\"autoscaler did not rent in {RENTAL_TIMEOUT}s \"\n f\"(possible causes: insufficient credit, scoring issue, \"\n f\"all candidates failed silently, or template+GPU mismatch \"\n f\"missed by pre-flight)\", None)\n # Jitter so N parallel threads don't all hit the autoscaler at the same instant.\n delay = WORKER_POLL_INTERVAL + random.uniform(0, 2.0)\n if stop_event is not None:\n if stop_event.wait(timeout=delay):\n update_row(class_states, row_id, status=\"aborted\")\n return (gpu_name, num_gpus, \"aborted\", None, \"user aborted\", None)\n else:\n time.sleep(delay)\n\n update_row(class_states, row_id, status=\"timeout\")\n return (gpu_name, num_gpus, \"timeout\", None,\n f\"no measured_perf in {timeout}s\", None)\n finally:\n # Workergroup first (stops new workers), then endpoint. atexit sweeps anything we miss.\n # 404s are silently absorbed since \"already gone\" is the desired end state.\n if workergroup_id is not None:\n try:\n vast.delete_workergroup(id=workergroup_id)\n workergroups.discard(workergroup_id)\n except Exception as e:\n if getattr(getattr(e, \"response\", None), \"status_code\", None) != 404:\n print(f\"[cleanup] failed to delete workergroup {workergroup_id}: {e}\",\n file=sys.stderr)\n else:\n workergroups.discard(workergroup_id)\n if endpoint_id is not None:\n try:\n vast.delete_endpoint(id=endpoint_id)\n endpoints.discard(endpoint_id)\n except Exception as e:\n if getattr(getattr(e, \"response\", None), \"status_code\", None) != 404:\n print(f\"[cleanup] failed to delete endpoint {endpoint_id}: {e}\",\n file=sys.stderr)\n else:\n endpoints.discard(endpoint_id)\n\n\[email protected](\n argument(\"--template_hash\", type=str, default=None,\n help=\"the template to benchmark; provide either --template_hash or --template_id\"),\n argument(\"--template_id\", type=int, default=None,\n help=\"the template to benchmark; provide either --template_hash or --template_id\"),\n argument(\"--gpus\", type=str,\n 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)\"),\n argument(\"--num_gpus\", type=int, default=None,\n help=\"GPUs per instance for tokens without an Nx prefix (default 1); overridden by inline Nx in --gpus\"),\n argument(\"--timeout\", type=int, default=DEFAULT_BENCHMARK_TIMEOUT,\n help=f\"max seconds to wait for a benchmark before giving up (default {DEFAULT_BENCHMARK_TIMEOUT})\"),\n argument(\"--no-cache\", dest=\"no_cache\", action=\"store_true\",\n help=\"re-measure instead of reusing recent cached benchmark results\"),\n argument(\"-y\", \"--yes\", action=\"store_true\",\n help=\"Skip confirmation prompt\"),\n argument(\"--auto_instance\", type=str, default=None, help=argparse.SUPPRESS),\n argument(\"--autoscaler_url\", type=str, default=None, help=argparse.SUPPRESS),\n usage=\"vastai run benchmarks [OPTIONS]\",\n help=\"Benchmark a template against one or more GPUs\",\n epilog=deindent(\"\"\"\n Rents one instance per GPU in parallel, measures perf, tears down.\n Each rental runs for up to --timeout seconds and costs real money.\n\n Specs benchmarked recently (same template, GPU, and count, by any\n user) are served from the benchmarks table as an estimated perf, and\n you're asked whether to reuse that or run a fresh benchmark. Pass\n --no-cache to always re-measure, or -y to always reuse the cache.\n\n Perf/$ uses the GPU $/hr, so cached and freshly measured rows are\n comparable. Cached rows recorded before that price was tracked show\n a perf with no price.\n\n Examples:\n # auto-sweep the default GPUs against TGI\n vastai run benchmarks --template_hash 79ebdd2ebfb9d42cedf7a221c42d37a5\n\n # specific GPUs against vLLM\n vastai run benchmarks --template_hash 393fa8572e6c73c927c8275fe4dffd53 --gpus RTX_4090,RTX_3090\n\n # multi-GPU configurations via inline Nx prefix against ComfyUI\n vastai run benchmarks --template_hash 40ef49becc953aa910ee05bd4653b9b3 --gpus \"2x RTX_4090, 2x RTX_3090\"\n\n # default count for tokens without an Nx prefix\n vastai run benchmarks --template_hash 79ebdd2ebfb9d42cedf7a221c42d37a5 --gpus RTX_4090,RTX_3090 --num_gpus 2\n\n # shorter timeout (10 min), skipping the cost prompt\n vastai run benchmarks --template_hash 393fa8572e6c73c927c8275fe4dffd53 --timeout 600 -y\n\n # ignore cached results and re-measure everything\n vastai run benchmarks --template_hash 79ebdd2ebfb9d42cedf7a221c42d37a5 --no-cache\n\n # raw JSON output for piping into another tool\n vastai run benchmarks --template_hash 40ef49becc953aa910ee05bd4653b9b3 --raw\n \"\"\"),\n)\ndef run__benchmarks(args):\n if args.template_id is None and args.template_hash is None:\n print(\"error: one of --template_id or --template_hash is required \"\n \"(run `vastai run benchmarks --help` for usage)\",\n file=sys.stderr)\n return 1\n if args.template_id is not None:\n query = {\"id\": {\"eq\": args.template_id}}\n ident = f\"id={args.template_id}\"\n else:\n query = {\"hash_id\": {\"eq\": args.template_hash}}\n ident = f\"hash={args.template_hash}\"\n\n # bump retry budget so parallel polls don't exhaust VastClient's tiny default (3 attempts, ~0.7s) under autoscaler rate limits.\n retry = max(args.retry, 8)\n vast = VastAI(api_key=args.api_key, server_url=args.url, retry=retry,\n explain=getattr(args, 'explain', False),\n curl=getattr(args, 'curl', False))\n\n templates = vast.search_templates(query=query)\n if not templates:\n print(f\"error: template not found ({ident})\", file=sys.stderr)\n return 1\n template = templates[0]\n try:\n extra_filters = json.loads(template.get(\"extra_filters\") or \"{}\")\n except (ValueError, TypeError):\n extra_filters = {}\n\n template_name = template.get(\"name\") or template.get(\"title\") or \"?\"\n filter_summary = format_filter_query(extra_filters) or \"none\"\n print(f\"\\nTemplate id={template.get('id')} {template_name}, filters: {filter_summary}\")\n\n # if user provides specific GPU count like \"2x RTX 4090\", that takes precedence over --num_gpus\n if args.gpus:\n gpu_specs = [parse_gpu_spec(t, args.num_gpus or 1)\n for t in args.gpus.split(\",\") if t.strip()]\n elif args.num_gpus is not None:\n gpu_specs = [(g, args.num_gpus) for g in DEFAULT_GPUS]\n else:\n gpu_specs = [\n (g, pick_num_gpus(g, extra_filters))\n for g in DEFAULT_GPUS\n ]\n\n console = Console(stderr=True)\n\n # deduplicates (gpu, num_gpus) so we dont benchmark the same GPU config twice\n seen = set()\n deduped = []\n for g, n in gpu_specs:\n if (g, n) in seen:\n console.print(f\"[yellow]skipping duplicate spec {n}x {g}[/yellow]\",\n highlight=False)\n continue\n seen.add((g, n))\n deduped.append((g, n))\n gpu_specs = deduped\n\n # Pre-flight: for each spec, serve a recent cached benchmark if one exists\n # (unless --no-cache), else skip specs with 0 matching offers, so the user\n # sees cache hits and skip reasons before approving any rentals.\n compatible_specs = []\n skipped_results = []\n cached_results = []\n interactive = not (args.yes or args.raw)\n for g, n in gpu_specs:\n if not args.no_cache:\n hit = lookup_cached_benchmark(\n vast, gpu_name=g, num_gpus=n,\n template_hash=template.get(\"hash_id\"),\n template_id=template.get(\"id\"),\n max_age_days=DEFAULT_CACHE_MAX_AGE_DAYS,\n )\n if hit:\n rentable = has_matching_offer(\n vast, gpu_name=g, num_gpus=n, extra_filters=extra_filters)\n age = (f\"{hit['age_days']:.0f}d\" if hit[\"age_days\"] >= 1\n else \"\u003c1d\")\n note = None if rentable else \"no offers available to rent right now\"\n median_dph_base = hit.get(\"median_dph_base\")\n pps = f\", {hit['median'] / median_dph_base:.1f} perf/$\" if median_dph_base else \"\"\n msg = (f\"[green][{n}x {g}] cached:[/green] median perf \"\n f\"{hit['median']:.1f}{pps} \"\n f\"(n={hit['n']}, range {hit['low']:.1f}-{hit['high']:.1f}, \"\n f\"newest {age} ago)\")\n if note:\n msg += f\"; [yellow]{note}[/yellow]\"\n console.print(msg, highlight=False)\n # Default is to reuse the cached result; only an interactive user\n # who opts in pays to re-measure (and only when offers exist).\n run_fresh = (\n interactive and rentable\n and input(\n f\" [{n}x {g}] run a fresh benchmark anyway? this rents a \"\n f\"real GPU and charges your account [y/N] \"\n ).strip().lower() in (\"y\", \"yes\")\n )\n if run_fresh:\n compatible_specs.append((g, n)) # offers already confirmed\n else:\n cached_results.append((g, n, \"cached\", hit[\"median\"], None, hit.get(\"median_dph_base\")))\n continue\n if not has_matching_offer(\n vast, gpu_name=g, num_gpus=n,\n extra_filters=extra_filters,\n ):\n msg = format_skip_message(\n vast, gpu_name=g, num_gpus=n,\n extra_filters=extra_filters,\n )\n console.print(f\"[yellow][{n}x {g}] skipping:[/yellow] {msg}\",\n highlight=False)\n skipped_results.append((g, n, \"skipped\", None, msg, None))\n else:\n compatible_specs.append((g, n))\n\n # Live-table state. Pre-populate every GPU so the table is complete.\n class_states = {}\n for g, n in compatible_specs:\n update_row(class_states, f\"{n}x {g}\", status=\"queued\")\n for sr in skipped_results:\n update_row(class_states, f\"{sr[1]}x {sr[0]}\", status=\"skipped\")\n for cr in cached_results:\n update_row(class_states, f\"{cr[1]}x {cr[0]}\", status=\"cached\",\n perf=cr[3], dph=cr[5])\n\n timeout_minutes = args.timeout / 60.0\n n = len(compatible_specs)\n if n == 0:\n if cached_results:\n console.print()\n console.print(render_table(class_states))\n else:\n console.print(\n \"\\nNo compatible GPUs to benchmark for this template.\",\n style=\"bold red\")\n return print_results(args, skipped_results + cached_results)\n\n spec_strs = [f\"{cnt}x {gpu}\" for gpu, cnt in compatible_specs]\n spec_summary = \", \".join(spec_strs)\n if n == 1:\n print(f\"\\nWill rent 1 GPU configuration ({spec_summary}). \"\n f\"Each runs ~2 to {timeout_minutes:.0f} min.\")\n else:\n print(f\"\\nWill rent {n} GPU configurations in parallel \"\n f\"({spec_summary}). Each runs ~2 to {timeout_minutes:.0f} min.\")\n print(\"This rents real GPUs and charges your account for usage.\")\n if not args.yes:\n if input(\"Continue? [y/N] \").strip().lower() not in (\"y\", \"yes\"):\n print(\"Aborted.\", file=sys.stderr)\n sys.exit(130)\n print()\n\n # Backing sets for the atexit sweep; per-GPU teardown discards on success.\n workergroups = set()\n endpoints = set()\n stop_event = threading.Event()\n\n def _cleanup():\n # Mask SIGINT + catch BaseException so a second Ctrl+C can't leak records mid-iteration.\n prev_handler = None\n try:\n prev_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)\n except (ValueError, OSError):\n pass\n try:\n for workergroup_id in list(workergroups):\n try:\n vast.delete_workergroup(id=workergroup_id)\n workergroups.discard(workergroup_id)\n except BaseException:\n pass\n for ep_id in list(endpoints):\n try:\n vast.delete_endpoint(id=ep_id)\n endpoints.discard(ep_id)\n except BaseException:\n pass\n finally:\n if prev_handler is not None:\n try:\n signal.signal(signal.SIGINT, prev_handler)\n except (ValueError, OSError):\n pass\n atexit.register(_cleanup)\n\n def _run_one_gpu(g, n):\n try:\n return benchmark_gpu(\n vast,\n gpu_name=g,\n num_gpus=n,\n timeout=args.timeout,\n workergroups=workergroups,\n endpoints=endpoints,\n template_hash=args.template_hash,\n template_id=args.template_id,\n auto_instance=args.auto_instance,\n autoscaler_url=args.autoscaler_url,\n class_states=class_states,\n stop_event=stop_event,\n )\n except Exception as e:\n update_row(class_states, f\"{n}x {g}\", status=\"error\")\n # Surface response body for HTTPError so backend rate-limit / validation\n # messages aren't swallowed (default __str__ on HTTPError just shows the URL).\n body = \"\"\n resp = getattr(e, \"response\", None)\n if resp is not None:\n try:\n body = f\" | body: {resp.text[:500]}\"\n except Exception:\n pass\n return (g, n, \"error\", None, f\"{type(e).__name__}: {e}{body}\", None)\n\n run_results = []\n executor = ThreadPoolExecutor(max_workers=len(compatible_specs))\n try:\n # Stagger submits to keep cold-start POSTs from tripping rate limits.\n futures = []\n for i, (g, n_for_class) in enumerate(compatible_specs):\n futures.append(executor.submit(_run_one_gpu, g, n_for_class))\n if i \u003c len(compatible_specs) - 1:\n time.sleep(PARALLEL_SUBMIT_DELAY)\n try:\n with Live(render_table(class_states), console=console,\n refresh_per_second=2, transient=False) as live:\n pending = set(futures)\n while pending:\n done, pending = wait(pending, timeout=0.5,\n return_when=FIRST_COMPLETED)\n for fut in done:\n result = fut.result()\n run_results.append(result)\n # Surface failures live above the table.\n g, n, status, _perf, err, _price = result\n if status not in (\"ok\", \"skipped\") and err:\n console.print(\n f\"[red][{n}x {g}] {status}:[/red] {err}\",\n highlight=False)\n live.update(render_table(class_states))\n except KeyboardInterrupt:\n # Set stop_event so threads exit their poll/sleep immediately, then cleanup synchronously.\n console.print(\n \"\\n[yellow]Deleting endpoints, please wait... \"\n \"If you exit now, you'll need to delete them manually with `vastai delete endpoint`.[/yellow]\",\n highlight=False)\n stop_event.set()\n executor.shutdown(wait=False, cancel_futures=True)\n _cleanup()\n sys.exit(130)\n finally:\n executor.shutdown(wait=True)\n _cleanup()\n\n results = skipped_results + cached_results + run_results\n return print_results(args, results)\n\n\ndef print_results(args, results):\n rows = []\n for gpu, num_gpus, status, perf, err, price in results:\n pps = (perf / price) if (perf and price) else None\n rows.append({\n \"gpu_name\": gpu,\n \"num_gpus\": num_gpus,\n \"rental_dph\": price,\n \"measured_perf\": perf,\n \"status\": status,\n \"perf_per_dollar\": pps,\n })\n rows.sort(key=lambda r: (r[\"perf_per_dollar\"] or -1), reverse=True)\n\n if args.raw:\n return rows\n\n n_ok = sum(1 for r in rows if r[\"status\"] == \"ok\")\n n_cached = sum(1 for r in rows if r[\"status\"] == \"cached\")\n line = f\"\\nBenchmark complete: {n_ok + n_cached}/{len(rows)} GPUs measured\"\n print(line + (f\" ({n_cached} from cache).\" if n_cached else \".\"))\n","content_type":"text/x-python; charset=utf-8","language":"python","size":43055,"content_sha256":"331b0a06f32ba03cb5093c2dd7162d80c80684990b6a714ddad396228c64b069"},{"filename":"cli/commands/billing.py","content":"\"\"\"CLI commands for billing, invoices, earnings, and user account management.\"\"\"\n\nimport json\nimport sys\nimport time\nfrom copy import deepcopy\nfrom datetime import datetime, timezone\nfrom io import StringIO\n\nfrom vastai.cli.parser import argument\nfrom vastai.cli.display import (\n display_table, invoice_fields, user_fields, ipaddr_fields,\n scheduled_jobs_fields, deindent, print_or_page,\n)\nfrom vastai.api import billing as billing_api\nfrom vastai.api import auth as auth_api\nfrom vastai.api import teams as teams_api\n\n\nfrom vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401\n\n\nparser = _get_parser()\n\n\n# ---------------------------------------------------------------------------\n# Helper utilities\n# ---------------------------------------------------------------------------\n\ndef convert_timestamp_to_date(unix_timestamp):\n utc_datetime = datetime.fromtimestamp(unix_timestamp, tz=timezone.utc)\n return utc_datetime.strftime(\"%Y-%m-%d\")\n\n\ndef to_timestamp_(val):\n \"\"\"Convert a date string or int to a UNIX timestamp.\"\"\"\n if isinstance(val, int):\n return val\n if isinstance(val, str):\n if val.isdigit():\n return int(val)\n return int(datetime.strptime(val + \"+0000\", '%Y-%m-%d%z').timestamp())\n raise ValueError(\"Invalid date format\")\n\n\ndef _sum(X, k):\n y = 0\n for x in X:\n a = float(x.get(k, 0))\n y += a\n return y\n\n\ndef _select(X, k):\n Y = set()\n for x in X:\n v = x.get(k, None)\n if v is not None:\n Y.add(v)\n return Y\n\n\ndef filter_invoice_items(args, rows):\n \"\"\"Filter invoice items by date range and charge/credit type.\"\"\"\n from datetime import date\n\n try:\n from dateutil import parser as dateutil_parser\n except ImportError:\n print(\"\\nWARNING: Missing dateutil, can't parse time format\")\n\n end_timestamp = 9999999999\n start_timestamp = 0\n start_date_txt = \"\"\n end_date_txt = \"\"\n\n if args.end_date:\n try:\n end_date = dateutil_parser.parse(str(args.end_date))\n end_date_txt = end_date.isoformat()\n if end_date.tzinfo is None:\n end_date = end_date.replace(tzinfo=timezone.utc)\n end_timestamp = end_date.timestamp()\n except ValueError:\n print(\"Warning: Invalid end date format! Ignoring end date!\")\n if args.start_date:\n try:\n start_date = dateutil_parser.parse(str(args.start_date))\n start_date_txt = start_date.isoformat()\n if start_date.tzinfo is None:\n start_date = start_date.replace(tzinfo=timezone.utc)\n start_timestamp = start_date.timestamp()\n except ValueError:\n print(\"Warning: Invalid start date format! Ignoring start date!\")\n\n selector_flag = \"\"\n if args.only_charges:\n type_txt = \"Only showing charges.\"\n selector_flag = \"only_charges\"\n def type_filter_fn(row):\n return row[\"type\"] == \"charge\"\n elif args.only_credits:\n type_txt = \"Only showing credits.\"\n selector_flag = \"only_credits\"\n def type_filter_fn(row):\n return row[\"type\"] == \"payment\"\n else:\n type_txt = \"\"\n def type_filter_fn(row):\n return True\n\n if args.end_date:\n if args.start_date:\n header_text = f'Invoice items after {start_date_txt} and before {end_date_txt}.'\n else:\n header_text = f'Invoice items before {end_date_txt}.'\n elif args.start_date:\n header_text = f'Invoice items after {start_date_txt}.'\n else:\n header_text = \" \"\n\n header_text = header_text + \" \" + type_txt\n\n rows = list(filter(\n lambda row: end_timestamp >= (row[\"timestamp\"] or 0.0) >= start_timestamp\n and type_filter_fn(row) and float(row[\"amount\"]) != 0,\n rows\n ))\n\n if start_date_txt:\n start_date_txt = \"S:\" + start_date_txt\n if end_date_txt:\n end_date_txt = \"E:\" + end_date_txt\n\n now = date.today()\n invoice_number = now.year * 12 + now.month - 1\n\n pdf_filename_fields = list(filter(\n lambda fld: fld != \"\",\n [str(invoice_number), start_date_txt, end_date_txt, selector_flag]\n ))\n filename = \"invoice_\" + \"-\".join(pdf_filename_fields) + \".pdf\"\n return {\"rows\": rows, \"header_text\": header_text, \"pdf_filename\": filename}\n\n\n# ---------------------------------------------------------------------------\n# Invoices-v1 helpers (rich formatting)\n# ---------------------------------------------------------------------------\n\ncharge_types = ['instance', 'volume', 'serverless', 'i', 'v', 's']\ninvoice_types = {\n \"transfers\": \"transfer\",\n \"stripe\": \"stripe_payments\",\n \"bitpay\": \"bitpay\",\n \"coinbase\": \"coinbase\",\n \"crypto.com\": \"crypto.com\",\n \"reserved\": \"instance_prepay\",\n \"payout_paypal\": \"paypal_manual\",\n \"payout_wise\": \"wise_manual\",\n}\n\n\ndef format_invoices_charges_results(args, results):\n indices_to_remove = []\n for i, item in enumerate(results):\n item['start'] = convert_timestamp_to_date(item['start']) if item['start'] else None\n item['end'] = convert_timestamp_to_date(item['end']) if item['end'] else None\n if item['amount'] == 0:\n indices_to_remove.append(i)\n elif args.invoices:\n if item['type'] not in {'transfer', 'payout'}:\n item['amount'] *= -1\n item['amount_str'] = f\"${item['amount']:.2f}\" if item['amount'] > 0 else f\"-${abs(item['amount']):.2f}\"\n else:\n item['amount'] = f\"${item['amount']:.3f}\"\n\n if args.charges:\n if item['type'] in {'instance', 'volume'} and not args.verbose:\n item['items'] = []\n if item['source'] and '-' in item['source']:\n item['type'], item['source'] = item['source'].capitalize().split('-')\n\n item['items'] = format_invoices_charges_results(args, item['items'])\n\n for i in reversed(indices_to_remove):\n del results[i]\n\n return results\n\n\ndef rich_object_to_string(rich_obj, no_color=True):\n \"\"\"Render a Rich object (Table or Tree) to a string.\"\"\"\n from rich.console import Console\n buffer = StringIO()\n console = Console(record=True, file=buffer)\n console.print(rich_obj)\n return console.export_text(clear=True, styles=not no_color)\n\n\ndef create_charges_tree(results, parent=None, title=\"Charges Breakdown\"):\n \"\"\"Build and return a Rich Tree from nested charge results.\"\"\"\n from rich.text import Text\n from rich.tree import Tree\n from rich.panel import Panel\n if parent is None:\n root = Tree(Text(title, style=\"bold red\"))\n create_charges_tree(results, root)\n return Panel(root, style=\"white on #000000\", expand=False)\n\n top_level = (parent.label.plain == title)\n for item in results:\n end_date = f\" -> {item['end']}\" if item['start'] != item['end'] else \"\"\n label = Text.assemble(\n (item[\"type\"], \"bold cyan\"),\n (f\" {item['source']}\" if item.get('source') else \"\", \"gold1\"), \" -> \",\n (f\"{item['amount']}\", 'bold green1' if top_level else 'green1'),\n (f\" -- {item['description']}\", \"bright_white\" if top_level else \"dim white\"),\n (f\" ({item['start']}{end_date})\", \"bold bright_white\" if top_level else \"white\")\n )\n node = parent.add(label, guide_style=\"blue3\")\n if item.get(\"items\"):\n create_charges_tree(item[\"items\"], node)\n return parent\n\n\ndef create_rich_table_for_charges(args, results):\n \"\"\"Build and return a Rich Table from charge results.\"\"\"\n from rich.table import Table\n from rich.text import Text\n from rich import box\n from rich.padding import Padding\n table = Table(style=\"white\", header_style=\"bold bright_yellow\", box=box.DOUBLE_EDGE, row_styles=[\"on grey11\", \"none\"])\n table.add_column(Text(\"Type\", justify=\"center\"), style=\"bold steel_blue1\", justify=\"center\")\n table.add_column(Text(\"ID\", justify=\"center\"), style=\"gold1\", justify=\"center\")\n table.add_column(Text(\"Amount\", justify=\"center\"), style=\"sea_green2\", justify=\"right\")\n table.add_column(Text(\"Start\", justify=\"center\"), style=\"bright_white\", justify=\"center\")\n table.add_column(Text(\"End\", justify=\"center\"), style=\"bright_white\", justify=\"center\")\n if not args.charge_type or 'serverless' in args.charge_type:\n table.add_column(Text(\"Endpoint\", justify=\"center\"), style=\"bright_red\", justify=\"center\")\n table.add_column(Text(\"Workergroup\", justify=\"center\"), style=\"orchid\", justify=\"center\")\n for item in results:\n row = [item['type'].capitalize(), item['source'], item['amount'], item['start'], item['end']]\n if not args.charge_type or 'serverless' in args.charge_type:\n row.append(str(item['metadata'].get('endpoint_id', '')))\n row.append(str(item['metadata'].get('workergroup_id', '')))\n table.add_row(*row)\n return Padding(table, (1, 2), style=\"on #000000\", expand=False)\n\n\ndef create_rich_table_for_invoices(results):\n \"\"\"Build and return a Rich Table from invoice results.\"\"\"\n from rich.table import Table\n from rich.text import Text\n from rich import box\n from rich.padding import Padding\n invoice_type_to_color = {\n \"credit\": \"green1\",\n \"transfer\": \"gold1\",\n \"payout\": \"orchid\",\n \"reserved\": \"sky_blue1\",\n \"refund\": \"bright_red\",\n }\n table = Table(style=\"white\", header_style=\"bold bright_yellow\", box=box.DOUBLE_EDGE, row_styles=[\"on grey11\", \"none\"])\n table.add_column(Text(\"ID\", justify=\"center\"), style=\"bright_white\", justify=\"center\")\n table.add_column(Text(\"Created\", justify=\"center\"), style=\"yellow3\", justify=\"center\")\n table.add_column(Text(\"Paid\", justify=\"center\"), style=\"yellow3\", justify=\"center\")\n table.add_column(Text(\"Type\", justify=\"center\"), justify=\"center\")\n table.add_column(Text(\"Result\", justify=\"center\"), justify=\"right\")\n table.add_column(Text(\"Source\", justify=\"center\"), style=\"bright_cyan\", justify=\"center\")\n table.add_column(Text(\"Description\", justify=\"center\"), style=\"bright_white\", justify=\"left\")\n for item in results:\n table.add_row(\n str(item['metadata']['invoice_id']),\n item['start'],\n item['end'] if item['end'] else 'N/A',\n Text(item['type'].capitalize(), style=invoice_type_to_color.get(item['type'], \"white\")),\n Text(item['amount_str'], style=\"sea_green2\" if item['amount'] > 0 else \"bright_red\"),\n item['source'].capitalize() if item['type'] != 'transfer' else item['source'],\n item['description'],\n )\n return Padding(table, (1, 2), style=\"on #000000\", expand=False)\n\n\n# ---------------------------------------------------------------------------\n# show invoices (deprecated)\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"-q\", \"--quiet\", action=\"store_true\", help=\"only display numeric ids\"),\n argument(\"-s\", \"--start_date\", help=\"start date and time for report. Many formats accepted (optional)\", type=str),\n argument(\"-e\", \"--end_date\", help=\"end date and time for report. Many formats accepted (optional)\", type=str),\n argument(\"-c\", \"--only_charges\", action=\"store_true\", help=\"Show only charge items\"),\n argument(\"-p\", \"--only_credits\", action=\"store_true\", help=\"Show only credit items\"),\n argument(\"--instance_label\", help=\"Filter charges on a particular instance label (useful for autoscaler groups)\"),\n usage=\"(DEPRECATED) vastai show invoices [OPTIONS]\",\n help=\"(DEPRECATED) Get billing history reports. Use `vastai show invoices-v1` instead.\",\n)\ndef show__invoices(args):\n \"\"\"Show current payments and charges (deprecated).\"\"\"\n if not args.quiet:\n print(\"DEPRECATED: `vastai show invoices` will be removed in a future release. \"\n \"Use `vastai show invoices-v1` for the new paginated command.\", file=sys.stderr)\n client = get_client(args)\n\n result = billing_api.show_invoices(client, start_date=args.start_date, end_date=args.end_date,\n only_charges=args.only_charges, only_credits=args.only_credits)\n rows = result[\"invoices\"]\n current_charges = result[\"current\"]\n\n invoice_filter_data = filter_invoice_items(args, rows)\n rows = invoice_filter_data[\"rows\"]\n filter_header = invoice_filter_data[\"header_text\"]\n\n if args.instance_label:\n contract_ids = _select(rows, 'instance_id')\n req_json = {\n \"label\": args.instance_label,\n \"contract_ids\": list(contract_ids)\n }\n if args.explain:\n print(\"request json: \")\n print(req_json)\n result2 = client.post(\"/contracts/fetch/\", json_data=req_json)\n result2.raise_for_status()\n filtered_rows = result2.json()[\"contracts\"]\n contract_ids = _select(filtered_rows, 'id')\n rows = [row for row in rows if row.get(\"instance_id\") in contract_ids]\n\n if args.quiet:\n for row in rows:\n id = row.get(\"id\", None)\n if id is not None:\n print(id)\n elif args.raw:\n return rows\n else:\n print(filter_header)\n display_table(rows, invoice_fields)\n print(f\"Total: ${_sum(rows, 'amount')}\")\n print(\"Current: \", current_charges)\n\n\n# ---------------------------------------------------------------------------\n# show invoices-v1 (advanced with pagination and rich formatting)\n# ---------------------------------------------------------------------------\n\[email protected](\n argument('-i', '--invoices', mutex_group='grp', action='store_true', required=True, help='Show invoices instead of charges'),\n argument('-it', '--invoice-type', choices=invoice_types.keys(), nargs='+', metavar='type',\n help=f'Filter which types of invoices to show: {{{\", \".join(invoice_types.keys())}}}'),\n argument('-c', '--charges', mutex_group='grp', action='store_true', required=True, help='Show charges instead of invoices'),\n argument('-ct', '--charge-type', choices=charge_types, nargs='+', metavar='type',\n help='Filter which types of charges to show: {i|instance, v|volume, s|serverless}'),\n argument('-s', '--start-date', help='Start date (YYYY-MM-DD or timestamp)'),\n argument('-e', '--end-date', help='End date (YYYY-MM-DD or timestamp)'),\n argument('-l', '--limit', type=int, default=20, help='Number of results per page (default: 20, max: 100)'),\n argument('-t', '--next-token', help='Pagination token for next page'),\n argument('-f', '--format', choices=['table', 'tree'], default='table', help='Output format for charges (default: table)'),\n argument('-v', '--verbose', action='store_true', help='Include full Instance Charge details and Invoice Metadata (tree view only)'),\n argument('--latest-first', action='store_true', help='Sort by latest first'),\n usage=\"vastai show invoices-v1 [OPTIONS]\",\n help=\"Get billing (invoices/charges) history reports with advanced filtering and pagination\",\n epilog=deindent(\"\"\"\n This command supports colored output and rich formatting if the 'rich' python module is installed!\n\n Examples:\n # Show the first 20 invoices in the last week (note: default window is a 7 day period ending today)\n vastai show invoices-v1 --invoices\n\n # Show the first 50 charges over a 7 day period starting from 2025-11-30 in tree format\n vastai show invoices-v1 --charges -s 2025-11-30 -f tree -l 50\n\n # Show the first 20 invoices of specific types for the month of November 2025\n vastai show invoices-v1 -i -it stripe bitpay transfers --start-date 2025-11-01 --end-date 2025-11-30\n\n # Show the first 20 charges for only volumes and serverless instances between two dates, including all details and metadata\n vastai show invoices-v1 -c --charge-type v s -s 2025-11-01 -e 2025-11-05 --format tree --verbose\n\n # Get the next page of paginated invoices, limit to 50 per page (note: type/date filters MUST match previous request for pagination to work)\n vastai show invoices-v1 --invoices --limit 50 --next-token eyJ2YWx1ZXMiOiB7ImlkIjogMjUwNzgyMzR9LCAib3NfcGFnZSI6IDB9\n\n # Show the last 10 instance (only) charges over a 7 day period ending in 2025-12-25, sorted by latest charges first\n vastai show invoices-v1 --charges -ct instance --end-date 2025-12-25 -l 10 --latest-first\n \"\"\"),\n)\ndef show__invoices_v1(args):\n \"\"\"Get billing (invoices/charges) history with advanced filtering and pagination.\"\"\"\n output_lines = []\n try:\n from rich.prompt import Confirm\n has_rich = True\n except ImportError:\n output_lines.append(\"NOTE: To view results in color and table/tree format please install the 'rich' python module with 'pip install rich'\\n\")\n has_rich = False\n\n # Handle default start and end date values\n if not args.start_date and not args.end_date:\n args.end_date = int(time.time())\n if not args.start_date:\n args.start_date = args.end_date - 7 * 24 * 60 * 60\n elif not args.end_date:\n args.end_date = args.start_date + 7 * 24 * 60 * 60\n\n try:\n start_timestamp = to_timestamp_(args.start_date)\n end_timestamp = to_timestamp_(args.end_date)\n except Exception as e:\n print(f\"Error parsing dates: {e}\")\n print(\"Use format YYYY-MM-DD or UNIX timestamp\")\n return\n\n if has_rich and not args.no_color and not args.raw:\n print(\"(use --no-color to disable colored output)\\n\")\n\n start_date = convert_timestamp_to_date(start_timestamp)\n end_date = convert_timestamp_to_date(end_timestamp)\n data_type = \"Instance Charges\" if args.charges else \"Invoices\"\n output_lines.append(f\"Fetching {data_type} from {start_date} to {end_date}...\")\n\n # Build SDK params\n params = {\n \"charges\": args.charges,\n \"start_date\": args.start_date,\n \"end_date\": args.end_date,\n \"limit\": min(args.limit, 100) if args.limit > 0 else 20,\n \"latest_first\": args.latest_first,\n \"format\": args.format if args.charges else \"table\",\n \"charge_type\": args.charge_type or [],\n \"invoice_type\": args.invoice_type or [],\n \"next_token\": args.next_token,\n }\n\n client = get_client(args)\n\n if args.raw:\n # Single-page raw response; dispatcher JSON-prints. Callers paginate\n # by re-running with --next-token.\n return billing_api.show_invoices_v1(client, params)\n\n found_results, found_count = [], 0\n looping = True\n while looping:\n response = billing_api.show_invoices_v1(client, params)\n\n found_results += response.get('results', [])\n found_count += response.get('count', 0)\n total = response.get('total', 0)\n next_token = response.get('next_token')\n\n if has_rich is False:\n output_lines.append(\"Raw response:\\n\" + json.dumps(response, indent=2))\n if next_token:\n print(f\"Next page token: {next_token}\\n\")\n elif not found_results:\n output_lines.append(\"No results found\")\n else:\n formatted_results = format_invoices_charges_results(args, deepcopy(found_results))\n if args.invoices:\n rich_obj = create_rich_table_for_invoices(formatted_results)\n elif args.format == 'tree':\n rich_obj = create_charges_tree(formatted_results)\n else:\n rich_obj = create_rich_table_for_charges(args, formatted_results)\n\n output_lines.append(rich_object_to_string(rich_obj, no_color=args.no_color))\n output_lines.append(f\"Showing {found_count} of {total} results\")\n if next_token:\n output_lines.append(f\"Next page token: {next_token}\\n\")\n\n paging = print_or_page(args, '\\n'.join(output_lines))\n\n if next_token and not paging:\n if has_rich:\n ans = Confirm.ask(\"Fetch next page?\", show_default=False, default=False)\n else:\n ans = input(\"Fetch next page? (y/N): \").strip().lower() == 'y'\n if ans:\n params['next_token'] = next_token\n output_lines.clear()\n args.full = True\n else:\n looping = False\n else:\n looping = False\n\n\n# ---------------------------------------------------------------------------\n# show earnings\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"-q\", \"--quiet\", action=\"store_true\", help=\"only display numeric ids\"),\n argument(\"-s\", \"--start_date\", help=\"start date and time for report. Many formats accepted\", type=str),\n argument(\"-e\", \"--end_date\", help=\"end date and time for report. Many formats accepted \", type=str),\n argument(\"-m\", \"--machine_id\", help=\"Machine id (optional)\", type=int),\n usage=\"vastai show earnings [OPTIONS]\",\n help=\"Get machine earning history reports\",\n)\ndef show__earnings(args):\n \"\"\"Show earnings history for a time range, optionally per machine.\"\"\"\n client = get_client(args)\n rows = billing_api.show_earnings(client, start_date=args.start_date, end_date=args.end_date,\n machine_id=args.machine_id)\n\n if args.raw:\n return rows\n print(json.dumps(rows, indent=1, sort_keys=True))\n\n\n# ---------------------------------------------------------------------------\n# show deposit\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of instance to get info for\", type=int),\n usage=\"vastai show deposit ID [options]\",\n help=\"Display reserve deposit info for an instance\",\n)\ndef show__deposit(args):\n \"\"\"Show reserve deposit info for an instance.\"\"\"\n client = get_client(args)\n result = billing_api.show_deposit(client, id=args.id)\n print(json.dumps(result, indent=1, sort_keys=True))\n\n\n# ---------------------------------------------------------------------------\n# show / set user\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"-q\", \"--quiet\", action=\"store_true\", help=\"display information about user\"),\n usage=\"vastai show user [OPTIONS]\",\n help=\"Get current user data\",\n epilog=deindent(\"\"\"\n Shows stats for logged-in user. These include user balance, email, and ssh key. Does not show API key.\n \"\"\"),\n)\ndef show__user(args):\n \"\"\"Show stats for logged-in user.\"\"\"\n client = get_client(args)\n user_blob = billing_api.show_user(client)\n\n if args.raw:\n return user_blob\n else:\n display_table([user_blob], user_fields)\n\n\[email protected](\n argument(\"--file\", help=\"file path for params in json format\", type=str),\n usage=\"vastai set user --file FILE\",\n help=\"Update user data from json file\",\n epilog=deindent(\"\"\"\n\n Available fields:\n\n Name Type Description\n\n ssh_key string\n paypal_email string\n wise_email string\n email string\n normalized_email string\n username string\n fullname string\n billaddress_line1 string\n billaddress_line2 string\n billaddress_city string\n billaddress_zip string\n billaddress_country string\n billaddress_taxinfo string\n balance_threshold_enabled string\n balance_threshold string\n autobill_threshold string\n phone_number string\n \"\"\"),\n)\ndef set__user(args):\n \"\"\"Update user data from a json file.\"\"\"\n params = None\n with open(args.file, 'r') as file:\n params = json.load(file)\n client = get_client(args)\n result = billing_api.set_user(client, params=params)\n print(f\"{result}\")\n\n\n# ---------------------------------------------------------------------------\n# show subaccounts / create subaccount\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"-q\", \"--quiet\", action=\"store_true\", help=\"display subaccounts from current user\"),\n usage=\"vastai show subaccounts [OPTIONS]\",\n help=\"Get current subaccounts\",\n)\ndef show__subaccounts(args):\n \"\"\"Show subaccounts.\"\"\"\n client = get_client(args)\n rows = billing_api.show_subaccounts(client)\n if args.raw:\n return rows\n else:\n display_table(rows, user_fields)\n\n\[email protected](\n argument(\"--email\", help=\"email address to use for login\", type=str),\n argument(\"--username\", help=\"username to use for login\", type=str),\n argument(\"--password\", help=\"password to use for login\", type=str),\n argument(\"--type\", help=\"host/client\", type=str),\n usage=\"vastai create subaccount --email EMAIL --username USERNAME --password PASSWORD --type TYPE\",\n help=\"Create a subaccount\",\n epilog=deindent(\"\"\"\n Creates a new account that is considered a child of your current account as defined via the API key.\n\n vastai create subaccount --email [email protected] --username bob --password password --type host\n\n vastai create subaccount --email [email protected] --username vast --password password --type host\n \"\"\"),\n)\ndef create__subaccount(args):\n \"\"\"Create a subaccount under your current account.\"\"\"\n host_only = False\n if args.type:\n host_only = args.type.lower() == \"host\"\n\n if getattr(args, 'explain', False):\n print(\"Request JSON would be: \")\n print({\n \"email\": args.email, \"username\": args.username,\n \"password\": args.password, \"host_only\": host_only, \"parent_id\": \"me\",\n })\n return\n\n client = get_client(args)\n try:\n rj = billing_api.create_subaccount(client, email=args.email, username=args.username,\n password=args.password, host_only=host_only)\n print(rj)\n except Exception as e:\n print(f\"Failed with error: {e}\")\n\n\n# ---------------------------------------------------------------------------\n# show ipaddrs\n# ---------------------------------------------------------------------------\n\[email protected](\n usage=\"vastai show ipaddrs [--api-key API_KEY] [--raw]\",\n help=\"Display user's history of ip addresses\",\n)\ndef show__ipaddrs(args):\n \"\"\"Show the history of ip address accesses.\"\"\"\n client = get_client(args)\n rows = billing_api.show_ipaddrs(client)\n if args.raw:\n return rows\n else:\n display_table(rows, ipaddr_fields)\n\n\n# ---------------------------------------------------------------------------\n# transfer credit\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"recipient_pos\", nargs=\"?\", default=None, metavar=\"RECIPIENT\",\n help=\"email (or id) of recipient account (positional; use --recipient to name)\"),\n argument(\"amount_pos\", nargs=\"?\", default=None, type=float, metavar=\"AMOUNT\",\n help=\"dollars of credit to transfer (positional; use --amount to name)\"),\n argument(\"-r\", \"--recipient\", help=\"email (or id) of recipient account\", type=str, default=None),\n argument(\"-a\", \"--amount\", help=\"dollars of credit to transfer\", type=float, default=None),\n argument(\"--skip\", help=\"skip confirmation\", action=\"store_true\", default=False),\n usage=\"vastai transfer credit [--recipient EMAIL] [--amount DOLLARS] [RECIPIENT AMOUNT]\",\n help=\"Transfer credits to another account\",\n epilog=deindent(\"\"\"\n Transfer credits to another account. This action is irreversible.\n\n Supports both named flags (recommended) and legacy positional arguments:\n vastai transfer credit --recipient [email protected] --amount 25.50\n vastai transfer credit -r [email protected] -a 25.50\n vastai transfer credit [email protected] 25.50\n \"\"\"),\n)\ndef transfer__credit(args):\n \"\"\"Transfer credits to another account.\"\"\"\n recipient = args.recipient if args.recipient is not None else args.recipient_pos\n amount = args.amount if args.amount is not None else args.amount_pos\n\n if recipient is None or amount is None:\n print(\"Error: both recipient and amount are required (as flags or positional args).\")\n return\n\n if not args.skip:\n print(f\"Transfer ${amount} credit to account {recipient}? This is irreversible.\")\n ok = input(\"Continue? [y/n] \")\n if ok.strip().lower() != \"y\":\n return\n\n if args.explain:\n print(\"request json: \")\n print({\"sender\": \"me\", \"recipient\": recipient, \"amount\": amount})\n\n client = get_client(args)\n try:\n rj = teams_api.transfer_credit(client, recipient=recipient, amount=amount)\n if rj[\"success\"]:\n print(f\"Sent {amount} to {recipient}\")\n else:\n print(rj[\"msg\"])\n except Exception as e:\n print(f\"Failed with error: {e}\")\n\n\n# ---------------------------------------------------------------------------\n# scheduled jobs\n# ---------------------------------------------------------------------------\n\ndef normalize_schedule_fields(job):\n \"\"\"Mutates the job dict to replace None values with readable scheduling labels.\"\"\"\n if job.get(\"day_of_the_week\") is None:\n job[\"day_of_the_week\"] = \"Everyday\"\n else:\n days = [\"Sunday\", \"Monday\", \"Tuesday\", \"Wednesday\", \"Thursday\", \"Friday\", \"Saturday\"]\n job[\"day_of_the_week\"] = days[int(job[\"day_of_the_week\"])]\n\n if job.get(\"hour_of_the_day\") is None:\n job[\"hour_of_the_day\"] = \"Every hour\"\n else:\n hour = int(job[\"hour_of_the_day\"])\n suffix = \"AM\" if hour \u003c 12 else \"PM\"\n hour_12 = hour % 12\n hour_12 = 12 if hour_12 == 0 else hour_12\n job[\"hour_of_the_day\"] = f\"{hour_12}_{suffix}\"\n\n if job.get(\"min_of_the_hour\") is None:\n job[\"min_of_the_hour\"] = \"Every minute\"\n else:\n job[\"min_of_the_hour\"] = f\"{int(job['min_of_the_hour']):02d}\"\n\n return job\n\n\[email protected](\n usage=\"vastai show scheduled-jobs [--api-key API_KEY] [--raw]\",\n help=\"Display the list of scheduled jobs\",\n)\ndef show__scheduled_jobs(args):\n \"\"\"Show the list of scheduled jobs for the account.\"\"\"\n client = get_client(args)\n rows = auth_api.show_scheduled_jobs(client)\n if args.raw:\n return rows\n else:\n rows = [normalize_schedule_fields(job) for job in rows]\n display_table(rows, scheduled_jobs_fields)\n\n\[email protected](\n argument(\"id\", help=\"id of scheduled job to remove\", type=int),\n usage=\"vastai delete scheduled-job ID\",\n help=\"Delete a scheduled job\",\n)\ndef delete__scheduled_job(args):\n \"\"\"Delete a scheduled job.\"\"\"\n client = get_client(args)\n result = auth_api.delete_scheduled_job(client, id=args.id)\n print(result)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":31121,"content_sha256":"92d66d4b2247f847aec3e7866c557e192109aaf03ac2746ea2cdbdb644f695f9"},{"filename":"cli/commands/clusters.py","content":"\"\"\"CLI commands for managing clusters and overlays.\"\"\"\n\nimport json\n\nfrom vastai.cli.parser import argument\nfrom vastai.cli.display import (\n display_table, cluster_fields, overlay_fields, deindent,\n)\nfrom vastai.api import clusters as clusters_api\n\n\nfrom vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401\n\n\nparser = _get_parser()\n\n\n# ---------------------------------------------------------------------------\n# clusters\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"subnet\", help=\"local subnet for cluster, ex: '0.0.0.0/24'\", type=str),\n argument(\"manager_id\", help=\"Machine ID of manager node in cluster. Must exist already.\", type=int),\n usage=\"vastai create cluster SUBNET MANAGER_ID\",\n help=\"Create Vast cluster\",\n epilog=deindent(\"\"\"\n Create Vast Cluster by defining a local subnet and manager id.\"\"\")\n)\ndef create__cluster(args):\n \"\"\"Create a Vast Cluster.\"\"\"\n if args.explain:\n print(\"request json:\", {\"subnet\": args.subnet, \"manager_id\": args.manager_id})\n\n client = get_client(args)\n result = clusters_api.create_cluster(client, subnet=args.subnet, manager_id=args.manager_id)\n\n if args.raw:\n return result\n\n print(result.get(\"msg\", result))\n\n\[email protected](\n usage=\"vastai show clusters\",\n help=\"Show clusters associated with your account.\",\n epilog=deindent(\"\"\"\n Show clusters associated with your account.\n \"\"\")\n)\ndef show__clusters(args):\n \"\"\"Show clusters associated with your account.\"\"\"\n client = get_client(args)\n response_data = clusters_api.show_clusters(client)\n\n if args.raw:\n return response_data\n\n rows = []\n for cluster_id, cluster_data in response_data['clusters'].items():\n machine_ids = [node[\"machine_id\"] for node in cluster_data[\"nodes\"]]\n manager_node = next(node for node in cluster_data['nodes'] if node['is_cluster_manager'])\n row_data = {\n 'id': cluster_id,\n 'subnet': cluster_data['subnet'],\n 'node_count': len(cluster_data['nodes']),\n 'machine_ids': str(machine_ids),\n 'manager_id': str(manager_node['machine_id']),\n 'manager_ip': manager_node['local_ip'],\n }\n rows.append(row_data)\n\n display_table(rows, cluster_fields, replace_spaces=False)\n\n\[email protected](\n argument(\"cluster_id\", help=\"ID of cluster to delete\", type=int),\n usage=\"vastai delete cluster CLUSTER_ID\",\n help=\"Delete Cluster\",\n epilog=deindent(\"\"\"\n Delete Vast Cluster\"\"\")\n)\ndef delete__cluster(args):\n \"\"\"Delete a Vast Cluster.\"\"\"\n if args.explain:\n print(\"request json:\", {\"cluster_id\": args.cluster_id})\n\n client = get_client(args)\n result = clusters_api.delete_cluster(client, cluster_id=args.cluster_id)\n\n if args.raw:\n return result\n print(result[\"msg\"])\n\n\[email protected](\n argument(\"cluster_id\", help=\"ID of cluster to add machine to\", type=int),\n argument(\"machine_ids\", help=\"machine id(s) to join cluster\", type=int, nargs=\"+\"),\n usage=\"vastai join cluster CLUSTER_ID MACHINE_IDS\",\n help=\"Join Machine to Cluster\",\n epilog=deindent(\"\"\"\n Join's Machine to Vast Cluster\n \"\"\")\n)\ndef join__cluster(args):\n \"\"\"Join machine(s) to a Vast Cluster.\"\"\"\n if args.explain:\n print(\"request json:\", {\"cluster_id\": args.cluster_id, \"machine_ids\": args.machine_ids})\n\n client = get_client(args)\n result = clusters_api.join_cluster(client, cluster_id=args.cluster_id, machine_ids=args.machine_ids)\n\n if args.raw:\n return result\n print(result[\"msg\"])\n\n\[email protected](\n argument(\"cluster_id\", help=\"ID of cluster you want to remove machine from.\", type=int),\n argument(\"machine_id\", help=\"ID of machine to remove from cluster.\", type=int),\n argument(\"new_manager_id\", help=\"ID of machine to promote to manager. Must already be in cluster\", type=int, nargs=\"?\"),\n usage=\"vastai remove-machine-from-cluster CLUSTER_ID MACHINE_ID NEW_MANAGER_ID\",\n help=\"Removes machine from cluster\",\n epilog=deindent(\"\"\"Removes machine from cluster and also reassigns manager ID,\n if we're removing the manager node\"\"\")\n)\ndef remove_machine_from_cluster(args):\n \"\"\"Remove a machine from a cluster.\"\"\"\n if args.explain:\n print(\"request json:\", {\"cluster_id\": args.cluster_id, \"machine_id\": args.machine_id,\n \"new_manager_id\": args.new_manager_id})\n\n client = get_client(args)\n result = clusters_api.remove_machine_from_cluster(\n client, cluster_id=args.cluster_id, machine_id=args.machine_id,\n new_manager_id=args.new_manager_id,\n )\n\n if args.raw:\n return result\n print(result.get(\"msg\", result))\n\n\n# ---------------------------------------------------------------------------\n# overlays\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"cluster_id\", help=\"ID of cluster to create overlay on top of\", type=int),\n argument(\"name\", help=\"overlay network name\"),\n usage=\"vastai create overlay CLUSTER_ID OVERLAY_NAME\",\n help=\"Creates overlay network on top of a physical cluster\",\n epilog=deindent(\"\"\"\n Creates an overlay network to allow local networking between instances on a physical cluster\"\"\"),\n)\ndef create__overlay(args):\n \"\"\"Create an overlay network on a physical cluster.\"\"\"\n if args.explain:\n print(\"request json:\", {\"cluster_id\": args.cluster_id, \"name\": args.name})\n\n client = get_client(args)\n result = clusters_api.create_overlay(client, cluster_id=args.cluster_id, name=args.name)\n\n if args.raw:\n return result\n print(result[\"msg\"])\n\n\[email protected](\n usage=\"vastai show overlays\",\n help=\"Show overlays associated with your account.\",\n epilog=deindent(\"\"\"\n Show overlays associated with your account.\n \"\"\"),\n)\ndef show__overlays(args):\n \"\"\"Show overlays associated with your account.\"\"\"\n client = get_client(args)\n response_data = clusters_api.show_overlays(client)\n if args.raw:\n return response_data\n\n rows = []\n for overlay in response_data:\n row_data = {\n 'overlay_id': overlay['overlay_id'],\n 'name': overlay['name'],\n 'subnet': overlay['internal_subnet'] if overlay['internal_subnet'] else 'N/A',\n 'cluster_id': overlay['cluster_id'],\n 'instance_count': len(overlay['instances']),\n 'instances': str(overlay['instances']),\n }\n rows.append(row_data)\n display_table(rows, overlay_fields, replace_spaces=False)\n\n\[email protected](\n argument(\"overlay_identifier\", help=\"ID (int) or name (str) of overlay to delete\", nargs=\"?\"),\n usage=\"vastai delete overlay OVERLAY_IDENTIFIER\",\n help=\"Deletes overlay and removes all of its associated instances\",\n)\ndef delete__overlay(args):\n \"\"\"Delete an overlay network.\"\"\"\n if args.explain:\n print(\"request json:\", {\"overlay_identifier\": args.overlay_identifier})\n\n client = get_client(args)\n result = clusters_api.delete_overlay(client, overlay_identifier=args.overlay_identifier)\n\n if args.raw:\n return result\n print(result[\"msg\"])\n\n\[email protected](\n argument(\"name\", help=\"Overlay network name to join instance to.\", type=str),\n argument(\"instance_id\", help=\"Instance ID to add to overlay.\", type=int),\n usage=\"vastai join overlay OVERLAY_NAME INSTANCE_ID\",\n help=\"Adds instance to an overlay network\",\n epilog=deindent(\"\"\"\n Adds an instance to a compatible overlay network.\"\"\"),\n)\ndef join__overlay(args):\n \"\"\"Add an instance to an overlay network.\"\"\"\n if args.explain:\n print(\"request json:\", {\"name\": args.name, \"instance_id\": args.instance_id})\n\n client = get_client(args)\n result = clusters_api.join_overlay(client, name=args.name, instance_id=args.instance_id)\n\n if args.raw:\n return result\n print(result[\"msg\"])\n","content_type":"text/x-python; charset=utf-8","language":"python","size":7994,"content_sha256":"f94411ceaeef83cfc8ad17f3eef81a7ca0289c588568b01a67fd2f8203141490"},{"filename":"cli/commands/deployments.py","content":"\"\"\"CLI commands for managing deployments.\"\"\"\n\nfrom vastai.cli.parser import argument\nfrom vastai.cli.display import display_table, deindent\nfrom vastai.api import deployments as deployments_api\n\nfrom vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401\n\nparser = _get_parser()\n\ndeployment_fields = (\n (\"id\", \"ID\", \"{}\", None, True),\n (\"name\", \"Name\", \"{}\", None, True),\n (\"tag\", \"Tag\", \"{}\", None, True),\n (\"image\", \"Image\", \"{}\", None, True),\n (\"endpoint_id\", \"Endpoint\", \"{}\", None, True),\n (\"file_hash\", \"Hash\", \"{}\", None, True),\n (\"current_version_id\", \"Version\", \"{}\", None, True),\n (\"ttl\", \"TTL\", \"{}\", None, True),\n (\"created_at\", \"Created\", \"{}\", None, True),\n (\"updated_at\", \"Updated\", \"{}\", None, True),\n)\n\ndeployment_detail_fields = (\n (\"id\", \"ID\", \"{}\", None, True),\n (\"name\", \"Name\", \"{}\", None, True),\n (\"tag\", \"Tag\", \"{}\", None, True),\n (\"image\", \"Image\", \"{}\", None, True),\n (\"endpoint_id\", \"Endpoint\", \"{}\", None, True),\n (\"endpoint_state\", \"State\", \"{}\", None, True),\n (\"worker_count\", \"Workers\", \"{}\", None, True),\n (\"file_hash\", \"Hash\", \"{}\", None, True),\n (\"s3_key\", \"S3 Key\", \"{}\", None, True),\n (\"current_version_id\", \"Version\", \"{}\", None, True),\n (\"last_healthy_version_id\", \"Last Healthy\", \"{}\", None, True),\n (\"storage\", \"Storage\", \"{}\", None, True),\n (\"ttl\", \"TTL\", \"{}\", None, True),\n (\"created_at\", \"Created\", \"{}\", None, True),\n (\"updated_at\", \"Updated\", \"{}\", None, True),\n)\n\n\n# ---------------------------------------------------------------------------\n# show deployments\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"-q\", \"--quiet\", action=\"store_true\", help=\"only display numeric ids\"),\n usage=\"vastai show deployments [OPTIONS]\",\n help=\"Display user's current deployments\",\n epilog=deindent(\"\"\"\n Shows the user's current deployments.\n\n Examples:\n vastai show deployments\n vastai show deployments --raw\n vastai show deployments -q\n \"\"\"),\n)\ndef show__deployments(args):\n \"\"\"Show the user's current deployments.\"\"\"\n client = get_client(args)\n rows = deployments_api.show_deployments(client)\n\n if args.quiet:\n for row in rows:\n id = row.get(\"id\", None)\n if id is not None:\n print(id)\n elif args.raw:\n return rows\n else:\n display_table(rows, deployment_fields)\n\n\[email protected](\n argument(\"id\", help=\"id of deployment to show info for\", type=int),\n argument(\"-q\", \"--quiet\", action=\"store_true\", help=\"only display numeric id\"),\n usage=\"vastai show deployment ID [OPTIONS]\",\n help=\"Display a single deployment\",\n epilog=deindent(\"\"\"\n Shows detailed info for a single deployment, including endpoint state and worker count.\n\n Examples:\n vastai show deployment 123\n vastai show deployment 123 --raw\n \"\"\"),\n)\ndef show__deployment(args):\n \"\"\"Show details of a single deployment.\"\"\"\n client = get_client(args)\n result = deployments_api.show_deployment(client, id=args.id)\n if args.raw:\n return result\n if args.quiet:\n print(result.get(\"id\", \"\"))\n else:\n display_table([result], deployment_detail_fields)\n\n\n# ---------------------------------------------------------------------------\n# show deployment-versions\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of deployment to show versions for\", type=int),\n usage=\"vastai show deployment-versions ID\",\n help=\"Display versions for a deployment\",\n epilog=deindent(\"\"\"\n Example: vastai show deployment-versions 1234\n \"\"\"),\n)\ndef show__deployment_versions(args):\n \"\"\"Show version history for a deployment.\"\"\"\n client = get_client(args)\n rows = deployments_api.show_deployment_versions(client, id=args.id)\n if args.raw:\n return rows\n else:\n import json\n print(json.dumps(rows, indent=1, sort_keys=True))\n\n\n# ---------------------------------------------------------------------------\n# stop deployment\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of deployment to stop\", type=int),\n usage=\"vastai stop deployment ID [OPTIONS]\",\n help=\"Stop a running deployment\",\n epilog=deindent(\"\"\"\n Stops the endpoint associated with a deployment. The deployment is not\n deleted and can be restarted with 'vastai start deployment'.\n\n Examples:\n vastai stop deployment 1234\n vastai stop deployment 1234 --raw\n \"\"\"),\n)\ndef stop__deployment(args):\n \"\"\"Stop a running deployment.\"\"\"\n client = get_client(args)\n rj = deployments_api.stop_deployment(client, id=args.id)\n if args.raw:\n return rj\n elif rj.get(\"success\"):\n print(f\"Deployment {args.id} stopped.\")\n else:\n print(rj.get(\"msg\", rj))\n\n\n# ---------------------------------------------------------------------------\n# start deployment\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of deployment to start\", type=int),\n usage=\"vastai start deployment ID [OPTIONS]\",\n help=\"Start a stopped deployment\",\n epilog=deindent(\"\"\"\n Starts (or restarts) the endpoint associated with a deployment.\n\n Examples:\n vastai start deployment 1234\n vastai start deployment 1234 --raw\n \"\"\"),\n)\ndef start__deployment(args):\n \"\"\"Start a stopped deployment.\"\"\"\n client = get_client(args)\n rj = deployments_api.start_deployment(client, id=args.id)\n if args.raw:\n return rj\n elif rj.get(\"success\"):\n print(f\"Deployment {args.id} started.\")\n else:\n print(rj.get(\"msg\", rj))\n\n\n# ---------------------------------------------------------------------------\n# delete deployment\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of deployment to delete\", type=int, nargs=\"?\", default=None),\n argument(\"--name\", help=\"name of deployment to delete (deletes all tags unless --tag is specified)\", type=str, default=None),\n argument(\"--tag\", help=\"tag to filter by when deleting by name\", type=str, default=None),\n usage=\"vastai delete deployment [ID | --name NAME [--tag TAG]]\",\n help=\"Delete a deployment by id, or by name and optional tag\",\n epilog=deindent(\"\"\"\n Examples:\n vastai delete deployment 1234\n vastai delete deployment --name my-deployment\n vastai delete deployment --name my-deployment --tag prod\n \"\"\"),\n)\ndef delete__deployment(args):\n \"\"\"Delete a deployment.\"\"\"\n if args.id is not None and args.name is not None:\n print(\"Error: specify either an id or --name, not both\")\n return\n if args.tag is not None and args.name is None:\n print(\"Error: --tag can only be used with --name\")\n return\n\n client = get_client(args)\n\n if args.id is not None:\n rj = deployments_api.delete_deployment(client, id=args.id)\n elif args.name is not None:\n if args.explain:\n json_blob = {\"name\": args.name}\n if args.tag is not None:\n json_blob[\"tag\"] = args.tag\n print(\"request json: \")\n print(json_blob)\n rj = deployments_api.delete_deployment_by_name(client, name=args.name, tag=args.tag)\n else:\n print(\"Error: must specify either an id or --name\")\n return\n\n if args.raw:\n return rj\n elif rj.get(\"success\"):\n if \"count\" in rj:\n print(f\"Deleted {rj['count']} deployment(s)\")\n elif args.id is not None:\n print(f\"Deleted deployment {args.id}.\")\n else:\n print(\"Deployment deleted successfully\")\n else:\n print(rj.get(\"msg\", rj))\n","content_type":"text/x-python; charset=utf-8","language":"python","size":8002,"content_sha256":"7866f0414a4be0b29aef6b8f56b22bbf8a336b4004eb956c0a17911d994cc73c"},{"filename":"cli/commands/endpoints.py","content":"\"\"\"CLI commands for managing endpoints and workergroups.\"\"\"\n\nimport json\nimport argparse\n\nfrom vastai.cli.parser import argument\nfrom vastai.cli.display import deindent\nfrom vastai.api import endpoints as endpoints_api\n\n\nfrom vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401\n\n\nparser = _get_parser()\n\n\n# ---------------------------------------------------------------------------\n# endpoints\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"--min_load\", help=\"minimum floor load in perf units/s (token/s for LLms)\", type=float, default=0.0),\n 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),\n argument(\"--target_util\", help=\"target capacity utilization (fraction, max 1.0, default 0.9)\", type=float, default=0.9),\n argument(\"--cold_mult\", help=\"cold/stopped instance capacity target as multiple of hot capacity target (default 2.5)\", type=float, default=2.5),\n argument(\"--cold_workers\", help=\"min number of workers to keep 'cold' when you have no load (default 5)\", type=int, default=5),\n argument(\"--max_workers\", help=\"max number of workers your endpoint group can have (default 20)\", type=int, default=20),\n argument(\"--endpoint_name\", help=\"deployment endpoint name (allows multiple autoscale groups to share same deployment endpoint)\", type=str),\n argument(\"--max_queue_time\", help=\"maximum seconds requests may be queued on each worker (default 30.0)\", type=float),\n argument(\"--target_queue_time\", help=\"target seconds for the queue to be cleared (default 10.0)\", type=float),\n argument(\"--inactivity_timeout\", help=\"seconds of no traffic before the endpoint can scale to zero active workers\", type=int),\n argument(\"--auto_instance\", help=argparse.SUPPRESS, type=str, default=\"prod\"),\n usage=\"vastai create endpoint [OPTIONS]\",\n help=\"Create a new endpoint group\",\n epilog=deindent(\"\"\"\n Create a new endpoint group to manage many autoscaling groups\n\n Example: vastai create endpoint --target_util 0.9 --cold_mult 2.0 --endpoint_name \"LLama\"\n \"\"\"),\n)\ndef create__endpoint(args):\n \"\"\"Create a new endpoint group.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\n \"client_id\": \"me\", \"min_load\": args.min_load, \"min_cold_load\": args.min_cold_load,\n \"target_util\": args.target_util, \"cold_mult\": args.cold_mult,\n \"cold_workers\": args.cold_workers, \"max_workers\": args.max_workers,\n \"endpoint_name\": args.endpoint_name, \"autoscaler_instance\": args.auto_instance,\n })\n\n client = get_client(args)\n result = endpoints_api.create_endpoint(\n client, min_load=args.min_load, min_cold_load=args.min_cold_load,\n target_util=args.target_util, cold_mult=args.cold_mult,\n cold_workers=args.cold_workers, max_workers=args.max_workers,\n endpoint_name=args.endpoint_name, auto_instance=args.auto_instance,\n max_queue_time=args.max_queue_time, target_queue_time=args.target_queue_time,\n inactivity_timeout=args.inactivity_timeout,\n )\n if args.raw:\n return result\n print(\"create endpoint {}\".format(result))\n\n\[email protected](\n usage=\"vastai show endpoints [--api-key API_KEY]\",\n help=\"Display user's current endpoint groups\",\n epilog=deindent(\"\"\"\n Example: vastai show endpoints\n \"\"\"),\n)\ndef show__endpoints(args):\n \"\"\"Display user's current endpoint groups.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"client_id\": \"me\", \"api_key\": args.api_key})\n\n client = get_client(args)\n result = endpoints_api.show_endpoints(client)\n\n if isinstance(result, dict) and \"error\" in result:\n print(result[\"error\"])\n return\n\n if args.raw:\n return result\n else:\n print(json.dumps(result, indent=1, sort_keys=True))\n\n\[email protected](\n argument(\"id\", help=\"id of endpoint group to update\", type=int),\n argument(\"--min_load\", help=\"minimum floor load in perf units/s (token/s for LLms)\", type=float),\n argument(\"--min_cold_load\", help=\"minimum floor load in perf units/s (token/s for LLms), but allow handling with cold workers\", type=float),\n argument(\"--endpoint_state\", help=\"active, suspended, or stopped\", type=str),\n argument(\"--auto_instance\", help=argparse.SUPPRESS, type=str, default=\"prod\"),\n argument(\"--target_util\", help=\"target capacity utilization (fraction, max 1.0, default 0.9)\", type=float),\n argument(\"--cold_mult\", help=\"cold/stopped instance capacity target as multiple of hot capacity target (default 2.5)\", type=float),\n argument(\"--cold_workers\", help=\"min number of workers to keep 'cold' when you have no load (default 5)\", type=int),\n argument(\"--max_workers\", help=\"max number of workers your endpoint group can have (default 20)\", type=int),\n argument(\"--endpoint_name\", help=\"deployment endpoint name (allows multiple workergroups to share same deployment endpoint)\", type=str),\n argument(\"--max_queue_time\", help=\"maximum seconds requests may be queued on each worker (default 30.0)\", type=float),\n argument(\"--target_queue_time\", help=\"target seconds for the queue to be cleared (default 10.0)\", type=float),\n argument(\"--inactivity_timeout\", help=\"seconds of no traffic before the endpoint can scale to zero active workers\", type=int),\n usage=\"vastai update endpoint ID [OPTIONS]\",\n help=\"Update an existing endpoint group\",\n epilog=deindent(\"\"\"\n Example: vastai update endpoint 4242 --min_load 100 --target_util 0.9 --cold_mult 2.0 --endpoint_name \"LLama\"\n \"\"\"),\n)\ndef update__endpoint(args):\n \"\"\"Update an existing endpoint group.\"\"\"\n client = get_client(args)\n result = endpoints_api.update_endpoint(\n client, id=args.id,\n min_load=args.min_load, min_cold_load=args.min_cold_load,\n target_util=args.target_util, cold_mult=args.cold_mult,\n cold_workers=args.cold_workers, max_workers=args.max_workers,\n endpoint_name=args.endpoint_name, endpoint_state=args.endpoint_state,\n auto_instance=args.auto_instance,\n max_queue_time=args.max_queue_time, target_queue_time=args.target_queue_time,\n inactivity_timeout=args.inactivity_timeout,\n )\n if args.raw:\n return result\n print(\"update endpoint {}\".format(result))\n\n\[email protected](\n argument(\"id\", help=\"id of endpoint group to delete\", type=int),\n usage=\"vastai delete endpoint ID \",\n help=\"Delete an endpoint group\",\n epilog=deindent(\"\"\"\n Example: vastai delete endpoint 4242\n \"\"\"),\n)\ndef delete__endpoint(args):\n \"\"\"Delete an endpoint group.\"\"\"\n id = args.id\n if args.explain:\n print(\"request json: \")\n print({\"client_id\": \"me\", \"endptjob_id\": args.id})\n\n client = get_client(args)\n result = endpoints_api.delete_endpoint(client, id=id)\n print(\"delete endpoint {}\".format(result))\n\n\[email protected](\n argument(\"id\", help=\"id of endpoint group to fetch logs from\", type=int),\n argument(\"--level\", help=\"log detail level (0 to 3)\", type=int, default=1),\n argument(\"--tail\", help=\"\", type=int, default=None),\n usage=\"vastai get endpt-logs ID [--api-key API_KEY]\",\n help=\"Fetch logs for a specific serverless endpoint group\",\n epilog=deindent(\"\"\"\n Example: vastai get endpt-logs 382\n \"\"\"),\n)\ndef get__endpt_logs(args):\n \"\"\"Fetch logs for a specific serverless endpoint group.\"\"\"\n if args.explain:\n print(f\"Fetching endpoint logs for id={args.id}\")\n\n client = get_client(args)\n rj = endpoints_api.get_endpt_logs(client, id=args.id, level=args.level, tail=args.tail)\n\n levels = {0: \"info0\", 1: \"info1\", 2: \"trace\", 3: \"debug\"}\n\n if isinstance(rj, dict) and \"error\" in rj:\n print(rj[\"error\"])\n return\n\n if args.raw:\n return rj\n else:\n dbg_lvl = levels[args.level]\n if rj and dbg_lvl:\n print(rj[dbg_lvl])\n\n\[email protected](\n argument(\"id\", help=\"id of endpoint group whose workers to list\", type=int),\n usage=\"vastai get endpt-workers ID [--api-key API_KEY]\",\n help=\"List workers on an endpoint group with status and measured_perf\",\n epilog=deindent(\"\"\"\n Example: vastai get endpt-workers 21979\n \"\"\"),\n)\ndef get__endpt_workers(args):\n \"\"\"List workers on a specific serverless endpoint group.\"\"\"\n client = get_client(args)\n rj = endpoints_api.get_endpoint_workers(client, id=args.id)\n\n if isinstance(rj, dict) and \"error\" in rj:\n print(rj[\"error\"])\n return\n\n workers = rj if isinstance(rj, list) else rj.get(\"workers\", [])\n if args.raw:\n return workers\n if not workers:\n print(\"No workers found.\")\n return\n for w in workers:\n print(f\"id={w.get('id')} status={w.get('status')} \"\n f\"gpu_name={w.get('gpu_name')} measured_perf={w.get('measured_perf')}\")\n\n\n# ---------------------------------------------------------------------------\n# workergroups\n# ---------------------------------------------------------------------------\n\[email protected](\n 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),\n argument(\"--template_id\", help=\"template id (optional)\", type=int),\n argument(\"-n\", \"--no-default\", action=\"store_true\", help=\"Disable default search param query args\"),\n 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),\n argument(\"--endpoint_name\", help=\"deployment endpoint name (allows multiple workergroups to share same deployment endpoint)\", type=str),\n argument(\"--endpoint_id\", help=\"deployment endpoint id (allows multiple workergroups to share same deployment endpoint)\", type=int),\n argument(\"--gpu_ram\", help=\"estimated GPU RAM req (independent of search string)\", type=float),\n 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),\n argument(\"--auto_instance\", help=argparse.SUPPRESS, type=str, default=\"prod\"),\n usage=\"vastai create workergroup [OPTIONS]\",\n help=\"Create a new autoscale group\",\n epilog=deindent(\"\"\"\n Create a new autoscaling group to manage a pool of worker instances.\n\n Example: vastai create workergroup --template_hash HASH --endpoint_name \"LLama\"\n \"\"\"),\n)\ndef create__workergroup(args):\n \"\"\"Create a new workergroup.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"template_hash\": args.template_hash, \"search_params\": args.search_params})\n\n client = get_client(args)\n try:\n result = endpoints_api.create_workergroup(\n client, template_hash=args.template_hash, template_id=args.template_id,\n no_default=args.no_default, launch_args=args.launch_args,\n endpoint_name=args.endpoint_name, endpoint_id=args.endpoint_id,\n gpu_ram=args.gpu_ram, search_params=args.search_params,\n auto_instance=args.auto_instance,\n )\n print(\"workergroup create {}\".format(result))\n except Exception as e:\n print(f\"Error creating workergroup: {e}\")\n\n\[email protected](\n usage=\"vastai show workergroups [--api-key API_KEY]\",\n help=\"Display user's current workergroups\",\n epilog=deindent(\"\"\"\n Example: vastai show workergroups\n \"\"\"),\n)\ndef show__workergroups(args):\n \"\"\"Display user's current workergroups.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"client_id\": \"me\", \"api_key\": args.api_key})\n\n client = get_client(args)\n result = endpoints_api.show_workergroups(client)\n\n if isinstance(result, dict) and \"error\" in result:\n print(result[\"error\"])\n return\n\n if args.raw:\n return result\n else:\n print(json.dumps(result, indent=1, sort_keys=True))\n\n\[email protected](\n argument(\"id\", help=\"id of autoscale group to update\", type=int),\n argument(\"--gpu_ram\", help=\"estimated GPU RAM req (independent of search string)\", type=float),\n 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),\n argument(\"--template_id\", help=\"template id\", type=int),\n 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),\n argument(\"-n\", \"--no-default\", action=\"store_true\", help=\"Disable default search param query args\"),\n 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),\n argument(\"--endpoint_name\", help=\"deployment endpoint name (allows multiple workergroups to share same deployment endpoint)\", type=str),\n argument(\"--endpoint_id\", help=\"deployment endpoint id (allows multiple workergroups to share same deployment endpoint)\", type=int),\n usage=\"vastai update workergroup WORKERGROUP_ID --endpoint_id ENDPOINT_ID [options]\",\n help=\"Update an existing autoscale group\",\n epilog=deindent(\"\"\"\n 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\n \"\"\"),\n)\ndef update__workergroup(args):\n \"\"\"Update an existing workergroup.\"\"\"\n client = get_client(args)\n result = endpoints_api.update_workergroup(\n client, id=args.id,\n gpu_ram=args.gpu_ram,\n template_hash=args.template_hash, template_id=args.template_id,\n search_params=args.search_params, no_default=args.no_default,\n launch_args=args.launch_args, endpoint_name=args.endpoint_name,\n endpoint_id=args.endpoint_id,\n )\n if args.raw:\n return result\n print(\"workergroup update {}\".format(result))\n\n\[email protected](\n argument(\"id\", help=\"id of workergroup to update workers for\", type=int),\n argument(\"--cancel\", action=\"store_true\", help=\"cancel an in-progress update for the workergroup\"),\n usage=\"vastai update workers WORKERGROUP_ID [--cancel]\",\n help=\"Trigger a rolling update of all workers in a workergroup, or cancel an in-progress update\",\n epilog=deindent(\"\"\"\n Starts a rolling update of all workers in the specified workergroup. The autoscaler\n will cycle through workers, updating them while maintaining capacity.\n\n Use --cancel to cancel an update that is currently in progress.\n\n Examples:\n vastai update workers 4242\n vastai update workers 4242 --cancel\n \"\"\"),\n)\ndef update__workers(args):\n \"\"\"Trigger a rolling update of workers in a workergroup.\"\"\"\n client = get_client(args)\n result = endpoints_api.update_workers(client, id=args.id, cancel=args.cancel)\n if args.raw:\n return result\n if result.get(\"success\"):\n if result.get(\"cancelled\"):\n print(f\"Update cancelled for workergroup {args.id}\")\n else:\n print(f\"Update started for workergroup {args.id} ({result.get('workers_to_update', 0)} workers)\")\n else:\n print(f\"Error: {result.get('error_msg', 'unknown error')}\")\n\n\[email protected](\n argument(\"id\", help=\"id of group to delete\", type=int),\n usage=\"vastai delete workergroup ID \",\n help=\"Delete a workergroup group\",\n epilog=deindent(\"\"\"\n Note that deleting a workergroup doesn't automatically destroy all the instances that are associated with your workergroup.\n Example: vastai delete workergroup 4242\n \"\"\"),\n)\ndef delete__workergroup(args):\n \"\"\"Delete a workergroup.\"\"\"\n id = args.id\n if args.explain:\n print(\"request json: \")\n print({\"client_id\": \"me\", \"autojob_id\": args.id})\n\n client = get_client(args)\n result = endpoints_api.delete_workergroup(client, id=id)\n print(\"workergroup delete {}\".format(result))\n\n\[email protected](\n argument(\"id\", help=\"id of workergroup to fetch logs from\", type=int),\n argument(\"--level\", help=\"log detail level (0 to 3)\", type=int, default=1),\n argument(\"--tail\", help=\"\", type=int, default=None),\n usage=\"vastai get wrkgrp-logs ID [--api-key API_KEY]\",\n help=\"Fetch logs for a specific serverless worker group\",\n epilog=deindent(\"\"\"\n Example: vastai get endpt-logs 382\n \"\"\"),\n)\ndef get__wrkgrp_logs(args):\n \"\"\"Fetch logs for a specific serverless worker group.\"\"\"\n if args.explain:\n print(f\"Fetching workergroup logs for id={args.id}\")\n\n client = get_client(args)\n rj = endpoints_api.get_wrkgrp_logs(client, id=args.id, level=args.level, tail=args.tail)\n\n levels = {0: \"info0\", 1: \"info1\", 2: \"trace\", 3: \"debug\"}\n\n if isinstance(rj, dict) and \"error\" in rj:\n print(rj[\"error\"])\n return\n\n if args.raw:\n return rj\n else:\n dbg_lvl = levels[args.level]\n if rj and dbg_lvl:\n print(rj[dbg_lvl])\n","content_type":"text/x-python; charset=utf-8","language":"python","size":17716,"content_sha256":"b3ce7219835e22f010c9a1b22b95fc314f8fb373ee8769dc0e262db86b6e2acd"},{"filename":"cli/commands/instances.py","content":"\"\"\"CLI commands for managing instances.\"\"\"\n\nimport json\nimport sys\nimport time\nimport argparse\nimport subprocess\n\nfrom vastai.cli.parser import argument, hidden_aliases, MyWideHelpFormatter\nfrom vastai.cli.util import (\n default_start_date, default_end_date,\n parse_day_cron_style, parse_hour_cron_style,\n validate_frequency_values, add_scheduled_job,\n parse_env, split_list, exec_with_threads,\n _get_gpu_names,\n)\nfrom vastai.api import instances as instances_api\nfrom vastai.api import offers as offers_api\n\n\nfrom vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401\n\n\n# ---------------------------------------------------------------------------\n# Instance display\n# ---------------------------------------------------------------------------\n\nfrom vastai.cli.display import display_table, print_or_page, instance_fields, deindent\n\nparser = _get_parser()\n\n\n# ---------------------------------------------------------------------------\n# show instances (defined further down, after the display helpers it uses)\n# ---------------------------------------------------------------------------\n\n# ---------------------------------------------------------------------------\n# show instance (single)\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of instance to get\", type=int),\n argument(\"-q\", \"--quiet\", action=\"store_true\", help=\"only display numeric id\"),\n usage=\"vastai show instance ID [options]\",\n help=\"Display user's current instances\",\n)\ndef show__instance(args):\n \"\"\"Shows stats for a single instance.\"\"\"\n client = get_client(args)\n result = instances_api.show_instance(client, id=args.id)\n if result is None:\n if args.raw:\n return {\"instances\": None}\n if not args.quiet:\n print(f\"Instance {args.id} not found or no longer exists.\", file=sys.stderr)\n return 1\n if args.raw:\n return result if result is not None else {\"instances\": None}\n if result is None:\n print(f\"Instance {args.id} not found or no longer exists.\")\n return 1\n if args.quiet:\n print(result.get(\"id\", \"\"))\n else:\n display_table([result], instance_fields)\n\n\n# ---------------------------------------------------------------------------\n# create instance\n# ---------------------------------------------------------------------------\n\ndef cli_runtype(args):\n \"\"\"Resolve the runtype from the CLI flags, keeping the CLI's 'ssh' default.\"\"\"\n if args.template_hash is not None:\n return None\n runtype, args.args = instances_api.resolve_runtype(\n ssh=args.ssh, jupyter=args.jupyter, direct=args.direct, args=args.args,\n jupyter_lab=args.jupyter_lab, jupyter_dir=args.jupyter_dir, default='ssh',\n )\n return runtype\n\n\ndef create_instance_impl(id, args):\n if args.onstart:\n with open(args.onstart, \"r\") as reader:\n args.onstart_cmd = reader.read()\n if args.onstart_cmd is None:\n args.onstart_cmd = args.entrypoint\n\n try:\n runtype = cli_runtype(args)\n except ValueError as e:\n print(f\"Error: {e}\", file=sys.stderr)\n return 1\n\n json_blob = instances_api.build_create_instance_payload(\n image=args.image, disk=args.disk, env=args.env, bid_price=args.bid_price,\n label=args.label, extra=args.extra, onstart_cmd=args.onstart_cmd,\n login=args.login, python_utf8=args.python_utf8, lang_utf8=args.lang_utf8,\n jupyter_lab=args.jupyter_lab, jupyter_dir=args.jupyter_dir,\n force=args.force, cancel_unavail=args.cancel_unavail,\n template_hash=args.template_hash, user=args.user, runtype=runtype,\n args=args.args, create_volume=args.create_volume,\n link_volume=args.link_volume, volume_size=args.volume_size,\n mount_path=args.mount_path, volume_label=args.volume_label,\n )\n\n if args.explain:\n print(\"request json: \")\n print(json_blob)\n\n client = get_client(args)\n rj = instances_api.create_instance_from_payload(client, id, json_blob)\n\n if args.raw:\n return rj\n else:\n print(\"Started. {}\".format(rj))\n return True\n\n\n_create_instance_args = [\n argument(\"--template_hash\", help=\"Create instance from template info\", type=str),\n 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),\n argument(\"--disk\", help=\"size of local disk partition in GB\", type=float, default=10),\n argument(\"--image\", help=\"docker container image to launch\", type=str),\n argument(\"--login\", help=\"docker login arguments for private repo authentication, surround with '' \", type=str),\n argument(\"--label\", help=\"label to set on the instance\", type=str),\n argument(\"--onstart\", help=\"filename to use as onstart script\", type=str),\n argument(\"--onstart-cmd\", help=\"contents of onstart script as single argument\", type=str),\n argument(\"--entrypoint\", help=\"override entrypoint for args launch instance\", type=str),\n argument(\"--ssh\", help=\"Launch as an ssh instance type\", action=\"store_true\"),\n argument(\"--jupyter\", help=\"Launch as a jupyter instance instead of an ssh instance\", action=\"store_true\"),\n argument(\"--direct\", help=\"Use (faster) direct connections for jupyter & ssh\", action=\"store_true\"),\n argument(\"--jupyter-dir\", help=\"For runtype 'jupyter', directory in instance to use to launch jupyter. Defaults to image's working directory\", type=str),\n argument(\"--jupyter-lab\", help=\"For runtype 'jupyter', Launch instance with jupyter lab\", action=\"store_true\"),\n 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\"),\n argument(\"--python-utf8\", help=\"Workaround for images with locale problems: set python's locale to C.UTF-8\", action=\"store_true\"),\n argument(\"--extra\", help=argparse.SUPPRESS),\n argument(\"--env\", help=\"env variables and port mapping options, surround with '' \", type=str),\n argument(\"--args\", nargs=argparse.REMAINDER, help=\"list of arguments passed to container ENTRYPOINT. Onstart is recommended for this purpose. (must be last argument)\"),\n argument(\"--force\", help=\"Skip sanity checks when creating from an existing instance\", action=\"store_true\"),\n argument(\"--cancel-unavail\", help=\"Return error if scheduling fails (rather than creating a stopped instance)\", action=\"store_true\"),\n argument(\"--bid_price\", help=\"(OPTIONAL) create an INTERRUPTIBLE instance with per machine bid price in $/hour\", type=float),\n 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),\n 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),\n argument(\"--volume-size\", help=\"Size of the volume to create in GB. Only usable with --create-volume (default 15GB)\", type=int),\n argument(\"--mount-path\", help=\"The path to the volume from within the new instance container. e.g. /root/volume\", type=str),\n argument(\"--volume-label\", help=\"(optional) A name to give the new volume. Only usable with --create-volume\", type=str),\n]\n\[email protected](\n argument(\"id\", help=\"id of instance type to launch (returned from search offers)\", type=int),\n *_create_instance_args,\n usage=\"vastai create instance ID [OPTIONS] [--args ...]\",\n help=\"Create a new instance\",\n epilog=deindent(\"\"\"\n Performs the same action as pressing the \"RENT\" button on the website at https://console.vast.ai/create/\n Creates an instance from an offer ID (which is returned from \"search offers\"). Each offer ID can only be used to create one instance.\n Besides the offer ID, you must pass in an '--image' argument as a minimum.\n\n If you use args/entrypoint launch mode, we create a container from your image as is, without attempting to inject ssh and or jupyter.\n If you use the args launch mode, you can override the entrypoint with --entrypoint, and pass arguments to the entrypoint with --args.\n If you use --args, that must be the last argument, as any following tokens are consumed into the args string.\n For ssh/jupyter launch types, use --onstart-cmd to pass in startup script, instead of --entrypoint and --args.\n\n Examples:\n\n # create an on-demand instance with the PyTorch (cuDNN Devel) template and 64GB of disk\n vastai create instance 384826 --template_hash 661d064bbda1f2a133816b6d55da07c3 --disk 64\n\n # 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\n 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'\";\n\n # create an on-demand instance with the bobsrepo/pytorch:latest image, 20GB of disk, open 22, 8080, jupyter ssh, and set some env variables\n 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\n\n # 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)\n vastai create instance 5801802 --image pytorch/pytorch --disk 40 --onstart-cmd 'bash' --args -c 'echo hello; sleep infinity;'\n\n # create an interruptible (spot) instance with the PyTorch (cuDNN Devel) template, 64GB of disk, and a bid price of $0.10/hr\n vastai create instance 384826 --template_hash 661d064bbda1f2a133816b6d55da07c3 --disk 64 --bid_price 0.1\n\n Return value:\n Returns a json reporting the instance ID of the newly created instance:\n {'success': True, 'new_contract': 7835610}\n \"\"\"),\n)\ndef create__instance(args):\n \"\"\"Create an instance from an offer ID.\"\"\"\n rj = create_instance_impl(args.id, args)\n if args.raw:\n return rj\n\n\[email protected](\n argument(\"ids\", help=\"ids of instance types to launch (returned from search offers)\", type=int, nargs='+'),\n *_create_instance_args,\n usage=\"vastai create instances [OPTIONS] ID0 ID1 ID2... [--args ...]\",\n help=\"Create instances from a list of offers\",\n)\ndef create__instances(args):\n \"\"\"Bulk create instances.\"\"\"\n idlist = split_list(args.ids, 64)\n exec_with_threads(lambda ids: create_instance_impl(ids, args), idlist, nt=8)\n\n\n# ---------------------------------------------------------------------------\n# destroy instance\n# ---------------------------------------------------------------------------\n\ndef destroy_instance_impl(id, args):\n client = get_client(args)\n rj = instances_api.destroy_instance(client, id=id)\n\n if args.raw:\n return rj\n elif rj.get(\"success\"):\n print(\"destroying instance {id}.\".format(**(locals())))\n return True\n else:\n print(rj.get(\"msg\", rj))\n return False\n\n\[email protected](\n argument(\"id\", help=\"id of instance to delete\", type=int),\n argument(\"-y\", \"--yes\", help=\"Skip confirmation prompt\", action=\"store_true\"),\n usage=\"vastai destroy instance id [-h] [-y] [--api-key API_KEY] [--raw]\",\n help=\"Destroy an instance (irreversible, deletes data)\",\n epilog=deindent(\"\"\"\n Perfoms the same action as pressing the \"DESTROY\" button on the website at https://console.vast.ai/instances/\n Example: vastai destroy instance 4242\n Use -y or --yes to skip the confirmation prompt.\n \"\"\"),\n)\ndef destroy__instance(args):\n \"\"\"Destroy a single instance.\"\"\"\n if not args.yes:\n try:\n 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()\n except (EOFError, KeyboardInterrupt):\n print(\"\\nAborted.\")\n return\n if confirm != \"y\":\n print(\"Aborted.\")\n return\n destroy_instance_impl(args.id, args)\n\n\[email protected](\n argument(\"ids\", help=\"ids of instance to destroy\", type=int, nargs='+'),\n argument(\"-y\", \"--yes\", help=\"Skip confirmation prompt\", action=\"store_true\"),\n usage=\"vastai destroy instances [--raw] [-y] \u003cid> [\u003cid> ...]\",\n help=\"Destroy a list of instances (irreversible, deletes data)\",\n)\ndef destroy__instances(args):\n \"\"\"Bulk destroy instances.\"\"\"\n if not args.yes:\n id_list_str = \", \".join(str(i) for i in args.ids)\n try:\n 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()\n except (EOFError, KeyboardInterrupt):\n print(\"\\nAborted.\")\n return\n if confirm != \"y\":\n print(\"Aborted.\")\n return\n idlist = split_list(args.ids, 64)\n exec_with_threads(lambda ids: destroy_instance_impl(ids, args), idlist, nt=8)\n\n\n# ---------------------------------------------------------------------------\n# start/stop/reboot/recycle instance\n# ---------------------------------------------------------------------------\n\ndef start_instance_impl(id, args):\n client = get_client(args)\n rj = instances_api.start_instance(client, id=id)\n\n if args.explain:\n print(\"request json: \")\n print({\"state\": \"running\"})\n if rj.get(\"success\"):\n print(\"starting instance {id}.\".format(**(locals())))\n return True\n else:\n print(rj.get(\"msg\", rj))\n return False\n\n\[email protected](\n argument(\"id\", help=\"ID of instance to start/restart\", type=int),\n usage=\"vastai start instance ID [OPTIONS]\",\n help=\"Start a stopped instance\",\n epilog=deindent(\"\"\"\n 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.\n 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.\n Examples:\n vastai start instances $(vastai show instances -q)\n vastai start instance 329838\n \"\"\"),\n)\ndef start__instance(args):\n \"\"\"Start a stopped instance.\"\"\"\n start_instance_impl(args.id, args)\n\n\[email protected](\n argument(\"ids\", help=\"ids of instance to start\", type=int, nargs='+'),\n usage=\"vastai start instances [OPTIONS] ID0 ID1 ID2...\",\n help=\"Start a list of instances\",\n)\ndef start__instances(args):\n \"\"\"Bulk start instances.\"\"\"\n idlist = split_list(args.ids, 64)\n exec_with_threads(lambda ids: start_instance_impl(ids, args), idlist, nt=8)\n\n\ndef stop_instance_impl(id, args):\n client = get_client(args)\n rj = instances_api.stop_instance(client, id=id)\n\n if args.explain:\n print(\"request json: \")\n print({\"state\": \"stopped\"})\n if rj.get(\"success\"):\n print(\"stopping instance {id}.\".format(**(locals())))\n return True\n else:\n print(rj.get(\"msg\", rj))\n return False\n\n\[email protected](\n argument(\"id\", help=\"id of instance to stop\", type=int),\n usage=\"vastai stop instance ID [OPTIONS]\",\n help=\"Stop a running instance\",\n epilog=deindent(\"\"\"\n 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,\n 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.\n There are ways to move data off of a stopped instance, which are described here: https://vast.ai/docs/gpu-instances/data-movement\n \"\"\"),\n)\ndef stop__instance(args):\n \"\"\"Stop a running instance.\"\"\"\n stop_instance_impl(args.id, args)\n\n\[email protected](\n argument(\"ids\", help=\"ids of instance to stop\", type=int, nargs='+'),\n usage=\"vastai stop instances [OPTIONS] ID0 ID1 ID2...\",\n help=\"Stop a list of instances\",\n epilog=deindent(\"\"\"\n Examples:\n vastai stop instances $(vastai show instances -q)\n vastai stop instances 329838 984849\n \"\"\"),\n)\ndef stop__instances(args):\n \"\"\"Bulk stop instances.\"\"\"\n idlist = split_list(args.ids, 64)\n exec_with_threads(lambda ids: stop_instance_impl(ids, args), idlist, nt=8)\n\n\[email protected](\n argument(\"id\", help=\"id of instance to reboot\", type=int),\n 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\"),\n 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)\"),\n 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)\"),\n 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),\n 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),\n usage=\"vastai reboot instance ID [OPTIONS]\",\n help=\"Reboot (stop/start) an instance\",\n epilog=deindent(\"\"\"\n Stops and starts container without any risk of losing GPU priority.\n \"\"\"),\n)\ndef reboot__instance(args):\n \"\"\"Reboot an instance.\"\"\"\n client = get_client(args)\n rj = instances_api.reboot_instance(client, id=args.id)\n\n if args.schedule:\n validate_frequency_values(args.day, args.hour, args.schedule)\n cli_command = \"reboot instance\"\n api_endpoint = \"/api/v0/instances/reboot/{id}/\".format(id=args.id)\n json_blob = {\"instance_id\": args.id}\n add_scheduled_job(client, args, json_blob, cli_command, api_endpoint, \"PUT\", instance_id=args.id)\n return\n\n if rj.get(\"success\"):\n print(\"Rebooting instance {args.id}.\".format(**(locals())))\n else:\n print(rj.get(\"msg\", rj))\n\n\[email protected](\n argument(\"id\", help=\"id of instance to recycle\", type=int),\n usage=\"vastai recycle instance ID [OPTIONS]\",\n help=\"Recycle (destroy/create) an instance\",\n epilog=deindent(\"\"\"\n Destroys and recreates container in place (from newly pulled image) without any risk of losing GPU priority.\n \"\"\"),\n)\ndef recycle__instance(args):\n \"\"\"Recycle an instance.\"\"\"\n client = get_client(args)\n rj = instances_api.recycle_instance(client, id=args.id)\n if rj.get(\"success\"):\n print(\"Recycling instance {args.id}.\".format(**(locals())))\n else:\n print(rj.get(\"msg\", rj))\n\n\n# ---------------------------------------------------------------------------\n# update instance\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of instance to update\", type=int),\n argument(\"--template_id\", help=\"new template ID to associate with the instance\", type=int),\n argument(\"--template_hash_id\", help=\"new template hash ID to associate with the instance\", type=str),\n argument(\"--image\", help=\"new image UUID for the instance\", type=str),\n argument(\"--args\", help=\"new arguments for the instance\", type=str),\n argument(\"--env\", help=\"new environment variables for the instance\", type=json.loads),\n argument(\"--onstart\", help=\"new onstart script for the instance\", type=str),\n usage=\"vastai update instance ID [OPTIONS]\",\n help=\"Update recreate an instance from a new/updated template\",\n epilog=deindent(\"\"\"\n Example: vastai update instance 1234 --template_hash_id 661d064bbda1f2a133816b6d55da07c3\n \"\"\"),\n)\ndef update__instance(args):\n \"\"\"Update/recreate an instance from a new or updated template.\"\"\"\n client = get_client(args)\n rj = instances_api.update_instance(\n client, id=args.id,\n template_id=args.template_id,\n template_hash_id=args.template_hash_id,\n image=args.image,\n args=args.args,\n env=args.env,\n onstart=args.onstart,\n )\n\n if args.raw:\n return rj\n if rj.get(\"success\"):\n print(f\"Instance {args.id} updated successfully.\")\n if rj.get(\"updated_instance\"):\n print(\"Updated instance details:\")\n print(rj.get(\"updated_instance\"))\n else:\n print(f\"Failed to update instance {args.id}: {rj.get('msg')}\")\n\n\n# ---------------------------------------------------------------------------\n# label / prepay / change bid\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of instance to label\", type=int),\n argument(\"label\", help=\"label to set\", type=str),\n usage=\"vastai label instance \u003cid> \u003clabel>\",\n help=\"Assign a string label to an instance\",\n)\ndef label__instance(args):\n \"\"\"Set a label on an instance.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"label\": args.label})\n\n client = get_client(args)\n rj = instances_api.label_instance(client, id=args.id, label=args.label)\n if rj.get(\"success\"):\n print(\"label for {args.id} set to {args.label}.\".format(**(locals())))\n else:\n print(rj.get(\"msg\", rj))\n\n\[email protected](\n argument(\"id\", help=\"id of instance to prepay for\", type=int),\n argument(\"amount\", help=\"amount of instance credit prepayment (default discount func of 0.2 for 1 month, 0.3 for 3 months)\", type=float),\n usage=\"vastai prepay instance ID AMOUNT\",\n help=\"Deposit credits into reserved instance\",\n)\ndef prepay__instance(args):\n \"\"\"Prepay for an instance.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"amount\": args.amount})\n\n client = get_client(args)\n rj = instances_api.prepay_instance(client, id=args.id, amount=args.amount)\n if rj.get(\"success\"):\n timescale = round(rj[\"timescale\"], 3)\n discount_rate = 100.0 * round(rj[\"discount_rate\"], 3)\n print(\"prepaid for {timescale} months of instance {args.id} applying ${args.amount} credits for a discount of {discount_rate}%\".format(**(locals())))\n else:\n print(rj.get(\"msg\", rj))\n\n\[email protected](\n argument(\"id\", help=\"id of instance type to change bid\", type=int),\n argument(\"--price\", help=\"per machine bid price in $/hour\", type=float),\n 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\"),\n 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)\"),\n 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)\"),\n 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),\n 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),\n usage=\"vastai change bid id [--price PRICE]\",\n help=\"Change the bid price for a spot/interruptible instance\",\n epilog=deindent(\"\"\"\n Change the current bid price of instance id to PRICE.\n If PRICE is not specified, then a winning bid price is used as the default.\n \"\"\"),\n)\ndef change__bid(args):\n \"\"\"Change the bid price for a spot instance.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"client_id\": \"me\", \"price\": args.price})\n\n client = get_client(args)\n\n if args.schedule:\n validate_frequency_values(args.day, args.hour, args.schedule)\n cli_command = \"change bid\"\n api_endpoint = \"/api/v0/instances/bid_price/{id}/\".format(id=args.id)\n json_blob = {\"client_id\": \"me\", \"price\": args.price, \"instance_id\": args.id}\n add_scheduled_job(client, args, json_blob, cli_command, api_endpoint, \"PUT\", instance_id=args.id)\n return\n\n instances_api.change_bid(client, id=args.id, price=args.price)\n print(\"Per gpu bid price changed\")\n\n\n# The accept/reject/show subcommands for price-increase live in\n# vastai/cli/commands/price_increase.py to keep this module focused on\n# instance CRUD/bid changes.\n\n\n# ---------------------------------------------------------------------------\n# launch instance\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"-g\", \"--gpu-name\", type=str, required=True, choices=_get_gpu_names(), help=\"Name of the GPU model, replace spaces with underscores\"),\n argument(\"-n\", \"--num-gpus\", type=str, required=True, choices=[\"1\", \"2\", \"4\", \"8\", \"12\", \"14\"], help=\"Number of GPUs required\"),\n argument(\"-r\", \"--region\", type=str, help=\"Geographical location of the instance\"),\n argument(\"-i\", \"--image\", required=True, help=\"Name of the image to use for instance\"),\n argument(\"-d\", \"--disk\", type=float, default=16.0, help=\"Disk space required in GB\"),\n argument(\"--limit\", default=3, type=int, help=\"\"),\n 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-'),\n argument(\"--login\", help=\"docker login arguments for private repo authentication, surround with '' \", type=str),\n argument(\"--label\", help=\"label to set on the instance\", type=str),\n argument(\"--onstart\", help=\"filename to use as onstart script\", type=str),\n argument(\"--onstart-cmd\", help=\"contents of onstart script as single argument\", type=str),\n argument(\"--entrypoint\", help=\"override entrypoint for args launch instance\", type=str),\n argument(\"--ssh\", help=\"Launch as an ssh instance type\", action=\"store_true\"),\n argument(\"--jupyter\", help=\"Launch as a jupyter instance instead of an ssh instance\", action=\"store_true\"),\n argument(\"--direct\", help=\"Use (faster) direct connections for jupyter & ssh\", action=\"store_true\"),\n argument(\"--jupyter-dir\", help=\"For runtype 'jupyter', directory in instance to use to launch jupyter. Defaults to image's working directory\", type=str),\n argument(\"--jupyter-lab\", help=\"For runtype 'jupyter', Launch instance with jupyter lab\", action=\"store_true\"),\n 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\"),\n argument(\"--python-utf8\", help=\"Workaround for images with locale problems: set python's locale to C.UTF-8\", action=\"store_true\"),\n argument(\"--extra\", help=argparse.SUPPRESS),\n argument(\"--env\", help=\"env variables and port mapping options, surround with '' \", type=str),\n argument(\"--args\", nargs=argparse.REMAINDER, help=\"list of arguments passed to container ENTRYPOINT. Onstart is recommended for this purpose. (must be last argument)\"),\n argument(\"--cancel-unavail\", help=\"Return error if scheduling fails (rather than creating a stopped instance)\", action=\"store_true\"),\n 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),\n usage=\"vastai launch instance [--help] [--api-key API_KEY] \u003cgpu_name> \u003cnum_gpus> \u003cimage> [geolocation] [disk_space]\",\n help=\"Launch the top instance from the search offers based on the given parameters\",\n epilog=deindent(\"\"\"\n Launches an instance based on the given parameters. The instance will be created with the top offer from the search results.\n Besides the gpu_name and num_gpus, you must pass in an '--image' argument as a minimum.\n\n If you use args/entrypoint launch mode, we create a container from your image as is, without attempting to inject ssh and or jupyter.\n If you use the args launch mode, you can override the entrypoint with --entrypoint, and pass arguments to the entrypoint with --args.\n If you use --args, that must be the last argument, as any following tokens are consumed into the args string.\n For ssh/jupyter launch types, use --onstart-cmd to pass in startup script, instead of --entrypoint and --args.\n\n Examples:\n\n # launch a single RTX 3090 instance with the pytorch image and 16 GB of disk space located anywhere\n python vast.py launch instance -g RTX_3090 -n 1 -i pytorch/pytorch\n\n # launch a 4x RTX 3090 instance with the pytorch image and 32 GB of disk space located in North America\n python vast.py launch instance -g RTX_3090 -n 4 -i pytorch/pytorch -d 32.0 -r North_America\n\n Available fields:\n\n Name Type Description\n\n num_gpus: int # of GPUs\n gpu_name: string GPU model name\n region: string Region of the instance\n image: string Docker image name\n disk_space: float Disk space in GB\n ssh, jupyter, direct: bool Flags to specify the instance type and connection method.\n env: str Environment variables and port mappings, encapsulated in single quotes.\n args: list Arguments passed to the container's ENTRYPOINT, used only if '--args' is specified.\n \"\"\"),\n)\ndef launch__instance(args):\n \"\"\"Launch the top instance from search offers.\"\"\"\n if args.onstart:\n with open(args.onstart, \"r\") as reader:\n args.onstart_cmd = reader.read()\n if args.onstart_cmd is None:\n args.onstart_cmd = args.entrypoint\n\n try:\n runtype = cli_runtype(args)\n except ValueError as e:\n print(f\"Error: {e}\", file=sys.stderr)\n return 1\n\n client = get_client(args)\n try:\n response_data = offers_api.launch_instance(\n client,\n gpu_name=args.gpu_name,\n num_gpus=args.num_gpus,\n image=args.image,\n region=args.region,\n disk=args.disk,\n order=args.order,\n limit=args.limit,\n env=parse_env(args.env),\n label=args.label,\n extra=args.extra,\n onstart_cmd=args.onstart_cmd,\n login=args.login,\n python_utf8=args.python_utf8,\n lang_utf8=args.lang_utf8,\n jupyter_lab=args.jupyter_lab,\n jupyter_dir=args.jupyter_dir,\n cancel_unavail=args.cancel_unavail,\n template_hash=args.template_hash,\n runtype=runtype,\n args=args.args,\n )\n if args.raw:\n return response_data\n else:\n print(\"Started. {}\".format(response_data))\n if response_data.get('success'):\n print(f\"Instance launched successfully: {response_data.get('new_contract')}\")\n else:\n print(f\"Failed to launch instance: {response_data.get('error')}, {response_data.get('message')}\")\n except Exception as err:\n print(f\"An error occurred: {err}\")\n\n\n# ---------------------------------------------------------------------------\n# show instances -- display helpers (Rich tables, filtering, pagination)\n# ---------------------------------------------------------------------------\n\n_DEFAULT_INSTANCE_SELECT_COLS = [\n \"id\", \"actual_status\", \"label\",\n \"num_gpus\", \"gpu_name\", \"gpu_util\",\n \"disk_space\", \"disk_usage\", \"disk_util\",\n \"volume_info\",\n \"dph_total\", \"image_uuid\",\n \"start_date\", \"verification\",\n]\n\n_VERBOSE_INSTANCE_SELECT_COLS = _DEFAULT_INSTANCE_SELECT_COLS + [\n \"machine_id\", \"template_id\", \"template_name\",\n \"geolocation\", \"inet_up\", \"inet_down\",\n \"ssh_host\", \"ssh_port\", \"status_msg\",\n]\n\n\ndef _fmt_age(start_date):\n \"\"\"Format seconds elapsed since start_date as e.g. '2d 3h' or '4h 15m'.\"\"\"\n if not start_date:\n return \"—\"\n secs = max(0, time.time() - start_date)\n d, rem = divmod(int(secs), 86400)\n h, rem = divmod(rem, 3600)\n m, _ = divmod(rem, 60)\n if d: return f\"{d}d {h}h\"\n if h: return f\"{h}h {m}m\"\n return f\"{m}m\"\n\n\ndef _fmt_disk(disk_usage, disk_space, disk_util):\n \"\"\"Format disk as 'used/total GB (X%)' or '?/total GB'.\"\"\"\n total = f\"{disk_space:.0f}\" if disk_space is not None else \"?\"\n if disk_usage is None or disk_usage \u003c 0:\n return f\"?/{total} GB\"\n used = f\"{disk_usage:.1f}\"\n if disk_util is not None and disk_util >= 0:\n pct = disk_util * 100\n return f\"{used}/{total} GB ({pct:.0f}%)\"\n return f\"{used}/{total} GB\"\n\n\ndef _fmt_volumes(volume_info):\n \"\"\"Format volume_info list as a compact string showing IDs and usage.\"\"\"\n if not volume_info:\n return \"—\"\n if len(volume_info) == 1:\n v = volume_info[0]\n vid = v.get(\"id\", \"?\")\n avail = v.get(\"avail_space\")\n total = v.get(\"total_space\")\n if avail is not None and total is not None:\n used = total - avail\n return f\"#{vid} {used:.0f}/{total:.0f} GB\"\n return f\"#{vid}\"\n return \", \".join(f\"#{v.get('id', '?')}\" for v in volume_info)\n\n\ndef _fmt_gpu(num_gpus, gpu_name, gpu_util):\n \"\"\"Format as '4x RTX 3090' or '4x RTX 3090 (72%)'.\"\"\"\n base = f\"{int(num_gpus)}x {gpu_name}\" if num_gpus and gpu_name else (gpu_name or \"—\")\n if gpu_util is not None and gpu_util >= 0:\n return f\"{base} ({gpu_util:.0f}%)\"\n return base\n\n\n_STATUS_COLORS = {\"running\": \"bold green\", \"loading\": \"bold yellow\", \"exited\": \"bright_red\", \"created\": \"bright_white\"}\n_VERIF_COLORS = {\"verified\": \"sea_green2\", \"unverified\": \"gold1\", \"deverified\": \"bright_red\"}\n\n\ndef _status_style(status):\n return _STATUS_COLORS.get(status, \"white\")\n\ndef _verif_style(v):\n return _VERIF_COLORS.get(v, \"white\")\n\n\n_INSTANCE_COL_MAX_WIDTHS = {\n \"gpu\": 20,\n \"image\": 30,\n \"age\": 8,\n \"volumes\": 17,\n \"location\": 22,\n \"net\": 11,\n \"ssh\": 22,\n \"template\": 32,\n \"msg\": 30,\n}\n\n# Column spec: (name, header, style, justify, min_width, drop_order, verbose_only)\n_INSTANCE_COL_SPECS = [\n (\"id\", \"ID\", \"bright_white\", \"right\", 4, 0, False),\n (\"status\", \"Status\", None, \"center\", 7, 0, False),\n (\"label\", \"Label\", \"bright_white\", \"left\", 7, 0, False),\n (\"gpu\", \"GPU\", \"steel_blue1\", \"left\", 13, 0, False),\n (\"disk\", \"Disk\", \"bright_white\", \"right\", 8, 5, False),\n (\"volumes\", \"Volumes\", \"bright_white\", \"left\", 10, 6, False),\n (\"dph\", \"$/hr\", \"sea_green2\", \"right\", 7, 1, False),\n (\"image\", \"Image\", \"orchid\", \"left\", 30, 2, False),\n (\"age\", \"Age\", \"bright_white\", \"left\", 8, 3, False),\n (\"verified\", \"Verified\", None, \"center\", 10, 0, False),\n (\"machine\", \"Machine\", \"gold1\", \"center\", 5, 8, True),\n (\"net\", \"Net Mbps\", \"bright_white\", \"left\", 9, 9, True),\n (\"location\", \"Location\", \"bright_white\", \"center\", 10, 10, True),\n (\"template\", \"Template\", \"bright_white\", \"center\", 20, 11, True),\n (\"ssh\", \"SSH\", \"cyan\", \"left\", 21, 7, True),\n (\"msg\", \"Msg\", \"dim white\", \"left\", 15, 12, True),\n]\n\n_INSTANCE_COL_SPEC_BY_NAME = {s[0]: s for s in _INSTANCE_COL_SPECS}\n\ntry:\n from rich.text import Text as _RichText\nexcept ImportError:\n _RichText = None\n\n\ndef _render_instance_col(name, inst):\n \"\"\"Render a single cell value for the given column name.\"\"\"\n if name == \"id\":\n return str(inst.get(\"id\", \"—\"))\n if name == \"status\":\n s = inst.get(\"actual_status\") or \"—\"\n return _RichText(s, style=_status_style(s))\n if name == \"label\":\n return inst.get(\"label\") or \"—\"\n if name == \"gpu\":\n return _fmt_gpu(inst.get(\"num_gpus\"), inst.get(\"gpu_name\"), inst.get(\"gpu_util\"))\n if name == \"disk\":\n return _fmt_disk(inst.get(\"disk_usage\"), inst.get(\"disk_space\"), inst.get(\"disk_util\"))\n if name == \"volumes\":\n return _fmt_volumes(inst.get(\"volume_info\") or [])\n if name == \"dph\":\n dph = inst.get(\"dph_total\")\n return f\"${dph:.4f}\" if dph is not None else \"—\"\n if name == \"image\":\n return (inst.get(\"image_uuid\") or \"—\")[:50]\n if name == \"age\":\n return _fmt_age(inst.get(\"start_date\"))\n if name == \"verified\":\n v = inst.get(\"verification\") or \"—\"\n return _RichText(v, style=_verif_style(v))\n if name == \"ssh\":\n return f\"{inst.get('ssh_host')}:{inst.get('ssh_port', '')}\" if inst.get(\"ssh_host\") else \"—\"\n if name == \"machine\":\n return str(inst.get(\"machine_id\", \"—\"))\n if name == \"net\":\n up, down = inst.get(\"inet_up\"), inst.get(\"inet_down\")\n return f\"↑{up:.0f} ↓{down:.0f}\" if (up is not None and down is not None) else \"—\"\n if name == \"location\":\n return inst.get(\"geolocation\") or \"—\"\n if name == \"template\":\n tid, tname = inst.get(\"template_id\"), inst.get(\"template_name\") or \"\"\n return (f\"{tname[:28]} ({tid})\" if tid else tname[:30] or \"—\")\n if name == \"msg\":\n return (inst.get(\"status_msg\") or \"—\")[:40]\n return \"—\"\n\n\ndef _estimate_table_width(specs):\n \"\"\"Estimate rendered table width for a list of col specs.\"\"\"\n n = len(specs)\n content = sum(\n _INSTANCE_COL_MAX_WIDTHS.get(s[0]) or max(len(s[1]), s[4])\n for s in specs\n )\n return 4 + 2 * n + (n - 1) + content\n\n\ndef _build_instances_table(instances, verbose=False, cols=None):\n \"\"\"Build the Rich table for instances.\"\"\"\n import shutil\n from rich.table import Table\n from rich import box\n\n show_volumes = any(inst.get(\"volume_info\") for inst in instances)\n term_width = shutil.get_terminal_size((120, 24)).columns\n\n if cols is not None:\n active = [_INSTANCE_COL_SPEC_BY_NAME[c] for c in cols if c in _INSTANCE_COL_SPEC_BY_NAME]\n hidden = []\n else:\n candidate = [\n s for s in _INSTANCE_COL_SPECS\n if (not s[6] or verbose) and not (s[0] == \"volumes\" and not show_volumes)\n ]\n droppable = sorted((s for s in candidate if s[5] > 0), key=lambda s: s[5], reverse=True)\n active = list(candidate)\n for drop_spec in droppable:\n if _estimate_table_width(active) \u003c= term_width:\n break\n active.remove(drop_spec)\n hidden = [s[1] for s in candidate if s not in active]\n\n tbl = Table(\n style=\"white\",\n header_style=\"bold bright_yellow\",\n box=box.DOUBLE_EDGE,\n row_styles=[\"on grey11\", \"none\"],\n )\n for name, header, style, justify, min_width, *_ in active:\n kwargs = dict(justify=justify, no_wrap=True, min_width=min_width)\n if name in _INSTANCE_COL_MAX_WIDTHS:\n kwargs[\"max_width\"] = _INSTANCE_COL_MAX_WIDTHS[name]\n if style:\n kwargs[\"style\"] = style\n tbl.add_column(_RichText(header, justify=\"center\"), **kwargs)\n\n for inst in instances:\n tbl.add_row(*[_render_instance_col(name, inst) for name, *_ in active])\n\n return tbl, hidden\n\n\n_FILTER_VALUE_COLORS = {\n \"status\": _STATUS_COLORS,\n \"verification\": _VERIF_COLORS,\n}\n\n\ndef _render_filter_values(values, colors=None, bold=False, line_sep=False):\n \"\"\"Render a sequence of filter values as a Rich Text, dot-separated.\"\"\"\n t = _RichText()\n for i, v in enumerate(values):\n if i: t.append(\"|\" if line_sep else \" · \", style=\"bright_white\")\n style = (colors or {}).get(v, \"bright_white\")\n t.append(v, style=(\"bold \" + style) if bold else style)\n return t\n\n\ndef _build_summary_panel(total, label_counts, active_filters=None, order_by=None):\n \"\"\"Build a Rich Panel summarising the instance query.\"\"\"\n from rich.panel import Panel\n from rich.table import Table\n from rich.text import Text\n\n lines = []\n lines.append(Text.assemble((\"Total: \", \"bold bright_yellow\"), (f\"{total} instances\", \"bold bright_white\")))\n\n if label_counts:\n parts = []\n for lbl, cnt in sorted(label_counts.items(), key=lambda x: -x[1]):\n display = lbl if lbl else \"(unlabeled)\"\n parts.append(f\"{display}: {cnt}\")\n lines.append(Text.assemble((\"Labels: \", \"bold bright_yellow\"), (\" · \".join(parts), \"bright_white\")))\n\n if active_filters:\n filter_line = Text.assemble((\"Filters: \", \"bold bright_yellow\"))\n for i, (k, vals) in enumerate(active_filters.items()):\n if i: filter_line.append(\" \", style=\"dim\")\n filter_line.append(f\"{k}=\", style=\"bold bright_white\")\n filter_line.append_text(_render_filter_values(vals, _FILTER_VALUE_COLORS.get(k), bold=True, line_sep=True))\n lines.append(filter_line)\n\n if order_by:\n order_line = Text.assemble((\"Order by: \", \"bold bright_yellow\"))\n for i, key in enumerate(order_by):\n if i: order_line.append(\" > \", style=\"bright_white\")\n order_line.append(key[\"col\"], style=\"bold bright_white\")\n order_line.append(f\" ({key['dir']})\", style=\"bright_white\")\n lines.append(order_line)\n\n grid = Table.grid(padding=(0, 0))\n grid.add_column()\n for line in lines:\n grid.add_row(line)\n\n return Panel(grid, title=\"[bold bright_yellow]Results Summary[/bold bright_yellow]\", style=\"on #000000\", border_style=\"bright_yellow\", expand=False)\n\n\ndef _build_filters_panel(filters):\n \"\"\"Build a Rich Panel showing the distinct filterable values.\"\"\"\n from rich.panel import Panel\n from rich.table import Table\n from rich.text import Text\n\n statuses = sorted({f[\"actual_status\"] for f in filters if f.get(\"actual_status\")})\n verifs = sorted({f[\"verification\"] for f in filters if f.get(\"verification\")})\n gpus = sorted({f[\"gpu_name\"] for f in filters if f.get(\"gpu_name\")})\n\n lines = [\n Text.assemble((\"--status: \", \"bold bright_yellow\"), _render_filter_values(statuses, _STATUS_COLORS)),\n Text.assemble((\"--verification: \", \"bold bright_yellow\"), _render_filter_values(verifs, _VERIF_COLORS)),\n Text.assemble((\"--gpu-name: \", \"bold bright_yellow\"), _render_filter_values(gpus)),\n ]\n\n grid = Table.grid(padding=(0, 0))\n grid.add_column()\n for line in lines:\n grid.add_row(line)\n\n return Panel(grid, title=\"[bright_white]Filterable Values[/bright_white]\", style=\"on #000000\", border_style=\"bright_white\", expand=False)\n\n\ndef _rich_object_to_string(rich_obj, no_color=True):\n \"\"\"Render a Rich object to a string.\"\"\"\n from io import StringIO\n from rich.console import Console\n buffer = StringIO()\n console = Console(record=True, file=buffer)\n console.print(rich_obj)\n return console.export_text(clear=True, styles=not no_color)\n\n\[email protected](\n argument(\"-q\", \"--quiet\", action=\"store_true\", help=\"only display numeric ids\"),\n argument(\"-v\", \"--verbose\", action=\"store_true\", help=\"show additional columns (SSH, location, template, etc.)\"),\n 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)\"),\n argument(\"-s\", \"--status\", metavar=\"STATUS\", nargs=\"+\", help=\"filter by container status: running loading exited\"),\n argument(\"--label\", metavar=\"LABEL\", nargs=\"+\", help=\"filter by instance label; pass empty string '' for unlabeled\"),\n argument(\"--gpu-name\", metavar=\"GPU\", nargs=\"+\", dest=\"gpu_name\", help=\"filter by GPU model name\"),\n argument(\"--verification\", metavar=\"VERIF\", nargs=\"+\", choices=[\"verified\", \"unverified\", \"deverified\"], help=\"filter by verification status\"),\n 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\"),\n argument(\"-t\", \"--next-token\", dest=\"next_token\", help=\"resume from a pagination token (implies single-page mode)\"),\n argument(\"--order-by\", dest=\"order_by\", metavar=\"COL [asc|desc]\", action=\"append\", help=\"sort by column; repeat for multiple keys\"),\n argument(\"--cols\", metavar=\"COLS\", help=f\"override displayed columns (available: {','.join(s[0] for s in _INSTANCE_COL_SPECS)})\"),\n usage=\"vastai show instances [OPTIONS]\",\n help=\"Display user's current instances\",\n aliases=hidden_aliases([\"show instances-v1\"]),\n epilog=deindent(\"\"\"\n Shows the instances the user is currently renting. By default this fetches\n every page and prints one combined result, which is safe for scripts and\n cron jobs. Passing --limit or --next-token switches to single-page mode for\n manual pagination (add --all to still fetch every page in that mode).\n\n Examples:\n vastai show instances\n vastai show instances --raw\n vastai show instances -q\n vastai show instances -v\n vastai show instances --status running loading\n vastai show instances --gpu-name 'RTX A5000' 'GTX 1070'\n vastai show instances --label training --order-by start_date desc\n vastai show instances --cols id,status,gpu,dph\n vastai show instances --next-token eyJ2YWx1ZXMi...\n \"\"\"),\n)\ndef show__instances(args, extra_filters=None):\n \"\"\"Show the user's current instances.\"\"\"\n if extra_filters and extra_filters.get('internal'):\n client = get_client(args)\n rows = instances_api.show_instances(client)\n field = extra_filters.get('field', 'id')\n return [str(r.get(field, '')) for r in rows]\n\n try:\n from rich.prompt import Confirm\n from rich.text import Text\n from rich.padding import Padding\n has_rich = True\n except ImportError:\n has_rich = False\n\n client = get_client(args)\n\n # build select_filters\n select_filters = {}\n active_display_filters = {}\n\n if args.status:\n invalid_statuses = [s for s in args.status if s not in _STATUS_COLORS]\n if invalid_statuses:\n valid = \", \".join(sorted(_STATUS_COLORS))\n print(f\"Warning: unknown status value(s): {', '.join(invalid_statuses)}. Valid: {valid}\", file=sys.stderr)\n select_filters[\"actual_status\"] = {\"in\": args.status}\n active_display_filters[\"status\"] = args.status\n\n if args.label is not None:\n vals = [None if l == \"\" else l for l in args.label]\n select_filters[\"label\"] = {\"in\": vals}\n active_display_filters[\"label\"] = [l or \"(unlabeled)\" for l in vals]\n\n if args.gpu_name:\n select_filters[\"gpu_name\"] = {\"in\": args.gpu_name}\n active_display_filters[\"gpu_name\"] = args.gpu_name\n\n if args.verification:\n select_filters[\"verification\"] = {\"in\": args.verification}\n active_display_filters[\"verification\"] = args.verification\n\n # order_by\n order_by = [{\"col\": \"id\", \"dir\": \"asc\"}]\n if args.order_by:\n order_by = []\n seen_cols = set()\n for entry in args.order_by:\n parts = entry.split()\n col = parts[0]\n dirn = parts[1].lower() if len(parts) > 1 and parts[1].lower() in (\"asc\", \"desc\") else \"asc\"\n order_by.append({\"col\": col, \"dir\": dirn})\n seen_cols.add(col)\n if \"id\" not in seen_cols:\n order_by.append({\"col\": \"id\", \"dir\": \"asc\"})\n\n # Legacy contract: fetch every page and print/return one combined result by\n # default (safe for scripts). Only an explicit --limit or --next-token opts\n # into single-page mode, matching the old paginated `instances-v1` behavior.\n explicit_pagination = args.next_token is not None or args.limit is not None\n fetch_all = args.all or not explicit_pagination\n limit = max(1, min(args.limit, 25)) if args.limit is not None else 25\n\n if args.raw and fetch_all and not args.quiet:\n # Flat list of every matching instance, full unrestricted row shape --\n # this is the contract scripts have always depended on for --raw.\n return instances_api.show_instances(client, select_filters=select_filters, order_by=order_by)\n\n params = {\n \"select_filters\": json.dumps(select_filters),\n \"order_by\": json.dumps(order_by),\n \"limit\": limit,\n }\n if not args.raw:\n select_cols = _VERBOSE_INSTANCE_SELECT_COLS if args.verbose else _DEFAULT_INSTANCE_SELECT_COLS\n if not has_rich:\n select_cols = [s[0] for s in instance_fields]\n params[\"select_cols\"] = json.dumps(select_cols)\n if args.next_token:\n params[\"after_token\"] = args.next_token\n\n if args.quiet and fetch_all:\n while True:\n data = instances_api.show_instances_v1(client, params)\n for inst in data.get(\"instances\") or []:\n instance_id = inst.get(\"id\")\n if instance_id is not None:\n print(instance_id)\n next_token = data.get(\"next_token\")\n if not next_token:\n return\n params[\"after_token\"] = next_token\n\n # fetch filter breakdown\n filter_combos = None\n if not args.quiet and not args.raw:\n try:\n filter_combos = instances_api.show_instance_filters(client)\n except Exception:\n pass\n\n user_cols = [c.strip() for c in args.cols.split(\",\")] if args.cols and has_rich else None\n\n page = 0\n offset = 0\n looping = True\n all_instances = []\n while looping:\n data = instances_api.show_instances_v1(client, params)\n\n instances = data.get(\"instances\", [])\n next_token = data.get(\"next_token\")\n total = data.get(\"total_instances\", 0)\n label_cnts = data.get(\"label_counts\", {})\n page += 1\n\n if args.quiet:\n for inst in instances:\n instance_id = inst.get(\"id\")\n if instance_id is not None:\n print(instance_id)\n break\n\n if args.raw:\n return data\n\n if fetch_all:\n all_instances.extend(instances)\n if next_token:\n sys.stderr.write(f\"\\rLoading more... (page {page + 1})\")\n sys.stderr.flush()\n params[\"after_token\"] = next_token\n continue\n sys.stderr.write(\"\\r\\033[K\")\n sys.stderr.flush()\n instances = all_instances\n\n output_parts = []\n\n if has_rich:\n if page == 1 or fetch_all:\n if filter_combos:\n output_parts.append(_rich_object_to_string(_build_filters_panel(filter_combos), no_color=args.no_color))\n output_parts.append(_rich_object_to_string(_build_summary_panel(\n total, label_cnts,\n active_filters=active_display_filters,\n order_by=order_by if args.order_by else None,\n ), no_color=args.no_color).rstrip(\"\\n\"))\n else:\n output_parts.append('')\n\n if not instances:\n empty_msg = \"No instances matched your filters.\" if active_display_filters else \"No instances found.\"\n output_parts.append(_rich_object_to_string(Text(empty_msg, style=\"bright_white\"), no_color=args.no_color).rstrip())\n else:\n tbl, hidden = _build_instances_table(instances, verbose=args.verbose, cols=user_cols)\n\n caption = Text()\n if not fetch_all:\n if not args.next_token:\n caption.append(f\"[Page {page}]\", style=\"bright_white\")\n caption.append(\" · \", style=\"bright_white\")\n caption.append(\"Fetched Results: \", style=\"bright_white\")\n caption.append(f\"{offset + 1} – {offset + len(instances)}\", style=\"bright_white\")\n caption.append(f\" of {total}\", style=\"bright_white\")\n if hidden and (page == 1 or fetch_all):\n caption.append(\n f\"\\nColumns hidden to fit terminal width: {', '.join(hidden)}\"\n f\" · use --cols to customize (see --help)\",\n style=\"dim\",\n )\n if caption:\n tbl.caption = caption\n tbl.caption_justify = \"left\"\n\n padded = Padding(tbl, (0, 1, 0, 1), style=\"on #000000\", expand=False)\n output_parts.append(_rich_object_to_string(padded, no_color=args.no_color))\n if next_token:\n output_parts.append(f\"Next page token: {next_token}\\n\")\n else:\n if page == 1 or fetch_all:\n if filter_combos:\n statuses = sorted({f[\"actual_status\"] for f in filter_combos if f.get(\"actual_status\")})\n verifs = sorted({f[\"verification\"] for f in filter_combos if f.get(\"verification\")})\n gpus = sorted({f[\"gpu_name\"] for f in filter_combos if f.get(\"gpu_name\")})\n print(\"Filterable Values:\")\n print(f\" --status: {' | '.join(statuses) if statuses else '(none)'}\")\n print(f\" --verification: {' | '.join(verifs) if verifs else '(none)'}\")\n print(f\" --gpu-name: {' | '.join(gpus) if gpus else '(none)'}\")\n print()\n summary_lines = [f\"Total: {total} instances\"]\n if label_cnts:\n lbl_parts = [f\"{(lbl or '(unlabeled)')}: {cnt}\" for lbl, cnt in sorted(label_cnts.items(), key=lambda x: -x[1])]\n summary_lines.append(f\"Labels: {' · '.join(lbl_parts)}\")\n if active_display_filters:\n filter_parts = [f\"{k}={' | '.join(str(v) for v in vals)}\" for k, vals in active_display_filters.items()]\n summary_lines.append(f\"Filters: {' '.join(filter_parts)}\")\n if args.order_by:\n order_parts = [f\"{o['col']} ({o['dir']})\" for o in order_by]\n summary_lines.append(f\"Order by: {' > '.join(order_parts)}\")\n print(\"Results Summary:\")\n for line in summary_lines:\n print(f\" {line}\")\n print()\n\n if not instances:\n print(\"No instances matched your filters.\" if active_display_filters else \"No instances found.\")\n else:\n display_table(instances, instance_fields)\n if not fetch_all:\n print(f\"[Page {page}] Fetched Results: {offset + 1} – {offset + len(instances)} of {total}\")\n if next_token:\n print(f\"Next page token: {next_token}\")\n if page == 1:\n print(\"\\nNOTE: install the 'rich' module for colored output (pip install rich)\")\n\n if not fetch_all:\n offset += len(instances)\n\n if fetch_all or not next_token:\n print_or_page(args, \"\\n\".join(output_parts))\n looping = False\n elif sys.stdin.isatty() and sys.stdout.isatty():\n print(\"\\n\".join(output_parts))\n try:\n if has_rich:\n ans = Confirm.ask(f\"Fetch next page? (page {page + 1})\", default=False)\n else:\n ans = input(f\"Fetch next page? (page {page + 1}) (y/N): \").strip().lower() == \"y\"\n except (EOFError, KeyboardInterrupt):\n ans = False\n if ans:\n params[\"after_token\"] = next_token\n else:\n looping = False\n else:\n # Non-interactive invocation (no TTY): never block on a prompt.\n print_or_page(args, \"\\n\".join(output_parts))\n looping = False\n","content_type":"text/x-python; charset=utf-8","language":"python","size":56763,"content_sha256":"0ea41f90d38f861212c893cf1fcfd25a203061344e4177d26bd4b5cb8fd52072"},{"filename":"cli/commands/keys.py","content":"\"\"\"CLI commands for managing API keys and SSH keys.\"\"\"\n\nimport json\nimport os\n\nfrom vastai.cli.parser import argument\nfrom vastai.cli.display import deindent\nfrom vastai.api import keys as keys_api\n\n\nfrom vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401\n\n\nparser = _get_parser()\n\n\n# ---------------------------------------------------------------------------\n# API keys\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"--name\", help=\"name of the api-key\", type=str),\n argument(\"--permission_file\", help=\"file path for json encoded permissions, see https://vast.ai/docs/cli/roles-and-permissions for more information\", type=str),\n argument(\"--key_params\", help=\"optional wildcard key params for advanced keys\", type=str),\n usage=\"vastai create api-key --name NAME --permission_file PERMISSIONS\",\n help=\"Create a new api-key with restricted permissions. Can be sent to other users and teammates\",\n epilog=deindent(\"\"\"\n In order to create api keys you must understand how permissions must be sent via json format.\n You can find more information about permissions here: https://vast.ai/docs/cli/roles-and-permissions\n \"\"\")\n)\ndef create__api_key(args):\n \"\"\"Create a new api-key with restricted permissions.\"\"\"\n from vastai.cli.util import load_permissions_from_file\n try:\n client = get_client(args)\n permissions = load_permissions_from_file(args.permission_file)\n result = keys_api.create_api_key(client, name=args.name, permissions=permissions, key_params=args.key_params)\n print(\"api-key created {}\".format(result))\n except FileNotFoundError:\n print(\"Error: Permission file '{}' not found.\".format(args.permission_file))\n\n\[email protected](\n argument(\"id\", help=\"id of apikey to get\", type=int),\n usage=\"vastai show api-key\",\n help=\"Show an api-key\",\n)\ndef show__api_key(args):\n \"\"\"Show an api-key.\"\"\"\n client = get_client(args)\n result = keys_api.show_api_key(client, id=args.id)\n print(result)\n\n\[email protected](\n usage=\"vastai show api-keys\",\n help=\"List your api-keys associated with your account\",\n)\ndef show__api_keys(args):\n \"\"\"List api-keys associated with your account.\"\"\"\n client = get_client(args)\n result = keys_api.show_api_keys(client)\n if args.raw:\n return result\n else:\n print(result)\n\n\[email protected](\n argument(\"id\", help=\"id of apikey to remove\", type=int),\n usage=\"vastai delete api-key ID\",\n help=\"Remove an api-key\",\n)\ndef delete__api_key(args):\n \"\"\"Remove an api-key.\"\"\"\n client = get_client(args)\n result = keys_api.delete_api_key(client, id=args.id)\n print(result)\n\n\[email protected](\n usage=\"vastai reset api-key\",\n help=\"Reset your api-key (get new key from website)\",\n)\ndef reset__api_key(args):\n \"\"\"Reset your api-key.\"\"\"\n client = get_client(args)\n rj = keys_api.reset_api_key(client)\n if rj.get(\"success\"):\n print(\"New api-key: {}\".format(rj.get(\"new_api_key\", \"\")))\n else:\n print(rj.get(\"msg\", \"Failed to reset api-key\"))\n\n\n# ---------------------------------------------------------------------------\n# SSH keys\n# ---------------------------------------------------------------------------\n\[email protected](\n 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='?'),\n argument(\"-y\", \"--yes\", help=\"automatically answer yes to prompts\", action=\"store_true\"),\n usage=\"vastai create ssh-key [ssh_public_key] [-y]\",\n help=\"Create a new ssh-key\",\n epilog=deindent(\"\"\"\n 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.\n\n 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.\n\n Example: $vastai create ssh-key 'ssh_public_key.pub'\n\n If you don't provide an ssh_public_key.pub argument, a new Ed25519 key pair will be generated.\n\n Example: $vastai create ssh-key\n\n 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_\u003ctimestamp>.\n The public key will be added to your Vast account.\n\n All ssh public keys are stored in your Vast account and can be used to connect to instances they've been added to.\n \"\"\")\n)\ndef create__ssh_key(args):\n \"\"\"Create or add an SSH key to your account.\"\"\"\n from vastai.cli.util import get_ssh_key, generate_ssh_key\n\n ssh_key_content = args.ssh_key\n\n if not ssh_key_content:\n ssh_key_content = generate_ssh_key(args.yes)\n else:\n print(\"Adding provided SSH public key to account...\")\n\n client = get_client(args)\n result = keys_api.create_ssh_key(client, ssh_key=ssh_key_content)\n print(\"ssh-key created {}\\nNote: You may need to add the new public key to any pre-existing instances\".format(result))\n\n\[email protected](\n usage=\"vastai show ssh-keys\",\n help=\"List your ssh keys associated with your account\",\n)\ndef show__ssh_keys(args):\n \"\"\"List ssh keys associated with your account.\"\"\"\n client = get_client(args)\n result = keys_api.show_ssh_keys(client)\n if args.raw:\n return result\n else:\n print(result)\n\n\[email protected](\n argument(\"id\", help=\"id ssh key to delete\", type=int),\n usage=\"vastai delete ssh-key ID\",\n help=\"Remove an ssh-key\",\n)\ndef delete__ssh_key(args):\n \"\"\"Remove an ssh-key.\"\"\"\n client = get_client(args)\n result = keys_api.delete_ssh_key(client, id=args.id)\n print(result)\n\n\[email protected](\n argument(\"id\", help=\"id of the ssh key to update\", type=int),\n argument(\"ssh_key\", help=\"new public key value\", type=str),\n usage=\"vastai update ssh-key ID SSH_KEY\",\n help=\"Update an existing SSH key\",\n)\ndef update__ssh_key(args):\n \"\"\"Update an existing SSH key.\"\"\"\n from vastai.cli.util import get_ssh_key\n\n ssh_key = get_ssh_key(args.ssh_key)\n client = get_client(args)\n result = keys_api.update_ssh_key(client, id=args.id, ssh_key=ssh_key)\n print(result)\n\n\n# ---------------------------------------------------------------------------\n# attach / detach ssh\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"instance_id\", help=\"id of instance to attach to\", type=int),\n argument(\"ssh_key\", help=\"ssh key to attach to instance\", type=str),\n usage=\"vastai attach ssh instance_id ssh_key\",\n help=\"Attach an ssh key to an instance. This will allow you to connect to the instance with the ssh key\",\n epilog=deindent(\"\"\"\n Attach an ssh key to an instance. This will allow you to connect to the instance with the ssh key.\n\n Examples:\n vastai attach ssh 12371 AAAAB3NzaC1yc2EAAA...\n vastai attach ssh 12371 $(cat ~/.ssh/id_rsa.pub)\n vastai attach ssh 12371 ~/.ssh/id_rsa.pub\n\n All examples attaches the ssh key to instance 12371\n \"\"\"),\n)\ndef attach__ssh(args):\n \"\"\"Attach an ssh key to an instance.\"\"\"\n from vastai.cli.util import get_ssh_key\n\n ssh_key = get_ssh_key(args.ssh_key)\n client = get_client(args)\n result = keys_api.attach_ssh(client, instance_id=args.instance_id, ssh_key=ssh_key)\n print(result)\n\n\[email protected](\n argument(\"instance_id\", help=\"id of the instance\", type=int),\n argument(\"ssh_key_id\", help=\"id of the key to detach from the instance\", type=str),\n usage=\"vastai detach instance_id ssh_key_id\",\n help=\"Detach an ssh key from an instance\",\n epilog=deindent(\"\"\"\n Example: vastai detach 99999 12345\n \"\"\")\n)\ndef detach__ssh(args):\n \"\"\"Detach an ssh key from an instance.\"\"\"\n client = get_client(args)\n result = keys_api.detach_ssh(client, instance_id=args.instance_id, ssh_key_id=args.ssh_key_id)\n print(result)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":8057,"content_sha256":"b05300ac613a66982ce621f65d5334233bef1422fcb0dbb9ee88517e54118559"},{"filename":"cli/commands/machines.py","content":"\"\"\"CLI commands for managing host machines.\"\"\"\n\nimport json\nimport os\nimport sys\nimport time\nimport warnings\nimport argparse\nfrom contextlib import redirect_stdout, redirect_stderr\n\nimport requests\nimport urllib3\n\nfrom vastai.cli.parser import argument\nfrom vastai.cli.display import (\n display_table, machine_fields, maintenance_fields,\n network_disk_fields, network_disk_machine_fields, deindent,\n)\nfrom vastai.api import machines as machines_api\nfrom vastai.api import instances as instances_api\nfrom vastai.api import offers as offers_api\nfrom vastai.api import storage as storage_api\n\n\nfrom vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401\nfrom vastai.cli.self_test.machine_diagnostics import (\n base_result,\n compact_offer_metadata,\n failed_checks,\n no_offer_failure,\n preflight_requirement_checks,\n render_preflight_advisories,\n render_preflight_failure,\n requirement_failure,\n)\nfrom vastai.cli.self_test.runtime_diagnostics import (\n DAEMON_STARTUP_FAILED,\n INSTANCE_CREATE_FAILED,\n INSTANCE_CREATE_MISSING_CONTRACT,\n INSTANCE_OFFLINE_BEFORE_TEST,\n INSTANCE_START_TIMEOUT,\n INTERRUPTED,\n LegacyProgressParser,\n MISSING_PUBLIC_IP,\n PROGRESS_CONTAINER_PORT,\n PROGRESS_EMPTY_TIMEOUT,\n PROGRESS_ENDPOINT_LOST,\n PROGRESS_ENDPOINT_UNREACHABLE,\n PROGRESS_PORT_NOT_MAPPED,\n RUNTIME_TEST_TIMEOUT,\n classify_status_msg,\n make_progress_endpoint_diagnostic,\n make_failure,\n redact_secret_text,\n)\nfrom vastai.cli.self_test.support_bundle import (\n create_support_bundle,\n format_bundle_summary,\n support_bundles_enabled,\n)\n\n\nparser = _get_parser()\nSELF_TEST_INSTANCE_LABEL_PREFIX = \"vast-self-test-machine\"\nINSTANCE_LOG_TAIL_LINES = 1000\n\n\n# ---------------------------------------------------------------------------\n# show machine / machines\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"Machine\", help=\"id of machine to display\", type=int),\n argument(\"-q\", \"--quiet\", action=\"store_true\", help=\"only display numeric ids\"),\n usage=\"vastai show machine ID [OPTIONS]\",\n help=\"[Host] Show hosted machines\",\n)\ndef show__machine(args):\n \"\"\"Show a machine the host is offering for rent.\"\"\"\n client = get_client(args)\n rows = machines_api.show_machine(client, id=args.Machine)\n if args.raw:\n return rows\n else:\n if args.quiet:\n ids = [f\"{row['id']}\" for row in rows]\n print(\" \".join(id for id in ids))\n else:\n display_table(rows, machine_fields)\n\n\[email protected](\n argument(\"-q\", \"--quiet\", action=\"store_true\", help=\"only display numeric ids\"),\n usage=\"vastai show machines [OPTIONS]\",\n help=\"[Host] Show hosted machines\",\n)\ndef show__machines(args):\n \"\"\"Show the machines user is offering for rent.\"\"\"\n client = get_client(args)\n rows = machines_api.show_machines(client)\n if args.raw:\n return {\"machines\": rows}\n else:\n if args.quiet:\n ids = [f\"{row['id']}\" for row in rows]\n print(\" \".join(id for id in ids))\n else:\n display_table(rows, machine_fields)\n\n\n# ---------------------------------------------------------------------------\n# show maints\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"-i\", \"--ids\", help=\"comma separated string of machine_ids for which to get maintenance information\", type=str),\n argument(\"-q\", \"--quiet\", action=\"store_true\", help=\"only display numeric ids of the machines in maintenance\"),\n usage=\"vastai show maints --ids 'machine_id_1[,machine_id_2,...]' [OPTIONS]\",\n help=\"[Host] Show maintenance information for host machines\",\n)\ndef show__maints(args):\n \"\"\"Show the maintenance information for the machines.\"\"\"\n machine_ids = args.ids.split(',')\n machine_ids = list(map(int, machine_ids))\n\n client = get_client(args)\n rows = machines_api.show_maints(client, machine_ids=machine_ids)\n if args.raw:\n return rows\n else:\n if args.quiet:\n ids = [f\"{row['machine_id']}\" for row in rows]\n print(\" \".join(id for id in ids))\n else:\n display_table(rows, maintenance_fields)\n\n\n# ---------------------------------------------------------------------------\n# show network-disks\n# ---------------------------------------------------------------------------\n\[email protected](\n usage=\"vastai show network-disks\",\n help=\"[Host] Show network disks associated with your account.\",\n epilog=deindent(\"\"\"\n Show network disks associated with your account.\n \"\"\")\n)\ndef show__network_disks(args):\n \"\"\"Show network disks associated with your account.\"\"\"\n client = get_client(args)\n response_data = storage_api.show_network_disks(client)\n\n if args.raw:\n return response_data\n\n for cluster_data in response_data['data']:\n print(f\"Cluster ID: {cluster_data['cluster_id']}\")\n display_table(cluster_data['network_disks'], network_disk_fields, replace_spaces=False)\n\n machine_rows = []\n for machine_id in cluster_data['machine_ids']:\n machine_rows.append({\n \"machine_id\": machine_id,\n \"mount_point\": cluster_data['mounts'].get(str(machine_id), \"N/A\"),\n })\n print()\n display_table(machine_rows, network_disk_machine_fields, replace_spaces=False)\n print(\"\\n\")\n\n\n# ---------------------------------------------------------------------------\n# list machine / machines\n# ---------------------------------------------------------------------------\n\ndef list_machine_impl(args, id):\n \"\"\"Shared logic for listing a single machine.\"\"\"\n from vastai.cli.util import string_to_unix_epoch\n\n client = get_client(args)\n end_date = string_to_unix_epoch(args.end_date) if args.end_date else None\n\n json_blob = {\n 'machine': id,\n 'price_gpu': args.price_gpu,\n 'price_disk': args.price_disk,\n 'price_inetu': args.price_inetu,\n 'price_inetd': args.price_inetd,\n 'price_min_bid': args.price_min_bid,\n 'min_chunk': args.min_chunk,\n 'end_date': end_date,\n 'credit_discount_max': args.discount_rate,\n 'duration': args.duration,\n 'vol_size': args.vol_size,\n 'vol_price': args.vol_price,\n }\n if args.explain:\n print(\"request json: \")\n print(json_blob)\n\n rj = machines_api.list_machine(\n client, id=id, price_gpu=args.price_gpu, price_disk=args.price_disk,\n price_inetu=args.price_inetu, price_inetd=args.price_inetd,\n price_min_bid=args.price_min_bid, min_chunk=args.min_chunk,\n end_date=end_date, discount_rate=args.discount_rate,\n duration=args.duration, vol_size=args.vol_size, vol_price=args.vol_price,\n )\n\n if rj.get(\"success\"):\n price_gpu_ = str(args.price_gpu) if args.price_gpu is not None else \"def\"\n price_inetu_ = str(args.price_inetu)\n price_inetd_ = str(args.price_inetd)\n min_chunk_ = str(args.min_chunk)\n discount_rate_ = str(args.discount_rate)\n duration_ = str(args.duration)\n if args.raw:\n return rj\n else:\n 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_}\")\n num_extended = rj.get(\"extended\", 0)\n if num_extended > 0:\n print(f\"extended {num_extended} client contracts to {args.end_date}\")\n else:\n if args.raw:\n return rj\n else:\n print(rj.get(\"msg\", rj))\n\n\[email protected](\n argument(\"id\", help=\"id of machine to list\", type=int),\n argument(\"-g\", \"--price_gpu\", help=\"per gpu rental price in $/hour (price for active instances)\", type=float),\n argument(\"-s\", \"--price_disk\",\n help=\"storage price in $/GB/month (price for inactive instances), default: $0.10/GB/month\", type=float),\n argument(\"-u\", \"--price_inetu\", help=\"price for internet upload bandwidth in $/GB\", type=float),\n argument(\"-d\", \"--price_inetd\", help=\"price for internet download bandwidth in $/GB\", type=float),\n argument(\"-b\", \"--price_min_bid\", help=\"per gpu minimum bid price floor in $/hour\", type=float),\n argument(\"-r\", \"--discount_rate\", help=\"Max long term prepay discount rate fraction, default: 0.4 \", type=float),\n argument(\"-m\", \"--min_chunk\", help=\"minimum amount of gpus (default: 1)\", type=int, default=1),\n argument(\"-e\", \"--end_date\", help=\"contract offer expiration - the available until date (optional, in unix float timestamp or MM/DD/YYYY format)\", type=str),\n 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.\"),\n 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),\n argument(\"-z\", \"--vol_price\", help=\"Price for disk on volume contract offer. Defaults to price_disk. Invalid if vol_size is 0.\", type=float),\n usage=\"vastai list machine ID [options]\",\n help=\"[Host] list a machine for rent\",\n epilog=deindent(\"\"\"\n Performs the same action as pressing the \"LIST\" button on the site https://cloud.vast.ai/host/machines.\n 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.\n Once you list your machine and it is rented, it is extremely important that you don't interfere with the machine in any way.\n If your machine has an active client job and then goes offline, crashes, or has performance problems, this could permanently lower your reliability rating.\n We strongly recommend you test the machine first and only list when ready.\n\n Raising any resource price above the current contract price writes a\n pending row to `pending_price_increases` for each affected client.\n Clients review and accept those rows via the console or\n `vastai show pending-price-increases`; auto-extend stops at the old\n price until they accept.\n \"\"\")\n)\ndef list__machine(args):\n \"\"\"List a machine for rent.\"\"\"\n return list_machine_impl(args, args.id)\n\n\[email protected](\n argument(\"ids\", help=\"ids of machines to list\", type=int, nargs='+'),\n argument(\"-g\", \"--price_gpu\", help=\"per gpu on-demand rental price in $/hour (base price for active instances)\", type=float),\n argument(\"-s\", \"--price_disk\",\n help=\"storage price in $/GB/month (price for inactive instances), default: $0.10/GB/month\", type=float),\n argument(\"-u\", \"--price_inetu\", help=\"price for internet upload bandwidth in $/GB\", type=float),\n argument(\"-d\", \"--price_inetd\", help=\"price for internet download bandwidth in $/GB\", type=float),\n argument(\"-b\", \"--price_min_bid\", help=\"per gpu minimum bid price floor in $/hour\", type=float),\n argument(\"-r\", \"--discount_rate\", help=\"Max long term prepay discount rate fraction, default: 0.4 \", type=float),\n argument(\"-m\", \"--min_chunk\", help=\"minimum amount of gpus (default: 1)\", type=int, default=1),\n argument(\"-e\", \"--end_date\", help=\"contract offer expiration - the available until date (optional, in unix float timestamp or MM/DD/YYYY format)\", type=str),\n 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.\"),\n 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),\n argument(\"-z\", \"--vol_price\", help=\"Price for disk on volume contract offer. Defaults to price_disk. Invalid if vol_size is 0.\", type=float),\n usage=\"vastai list machines IDs [options]\",\n help=\"[Host] list machines for rent\",\n epilog=deindent(\"\"\"\n This variant can be used to list or update the listings for multiple machines at once with the same args.\n You could extend the end dates of all your machines using a command combo like this:\n ./vast.py list machines $(./vast.py show machines -q) -e 12/31/2024 --retry 6\n \"\"\")\n)\ndef list__machines(args):\n \"\"\"List multiple machines for rent.\"\"\"\n return [list_machine_impl(args, id) for id in args.ids]\n\n\n# ---------------------------------------------------------------------------\n# unlist machine\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of machine to unlist\", type=int),\n usage=\"vastai unlist machine \u003cid>\",\n help=\"[Host] Unlist a listed machine\",\n)\ndef unlist__machine(args):\n \"\"\"Remove machine from list of machines for rent.\"\"\"\n client = get_client(args)\n rj = machines_api.unlist_machine(client, id=args.id)\n if rj.get(\"success\"):\n print(\"all offers for machine {machine_id} removed, machine delisted.\".format(machine_id=args.id))\n else:\n print(rj.get(\"msg\", rj))\n\n\n# ---------------------------------------------------------------------------\n# delete machine\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of machine to delete\", type=int),\n usage=\"vastai delete machine \u003cid>\",\n 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.\",\n)\ndef delete__machine(args):\n \"\"\"Delete machine if the machine is not being used by clients.\"\"\"\n client = get_client(args)\n rj = machines_api.delete_machine(client, id=args.id)\n if rj.get(\"success\"):\n print(\"deleted machine_id ({machine_id}) and all related contracts.\".format(machine_id=args.id))\n else:\n print(rj.get(\"msg\", rj))\n\n\n# ---------------------------------------------------------------------------\n# cleanup machine\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of machine to cleanup\", type=int),\n usage=\"vastai cleanup machine ID [options]\",\n help=\"[Host] Remove all expired storage instances from the machine, freeing up space\",\n epilog=deindent(\"\"\"\n Instances expire on their end date. Expired instances still pay storage fees, but can not start.\n Since hosts are still paid storage fees for expired instances, we do not auto delete them.\n Instead you can use this CLI/API function to delete all expired storage instances for a machine.\n This is useful if you are running low on storage, want to do maintenance, or are subsidizing storage, etc.\n \"\"\")\n)\ndef cleanup__machine(args):\n \"\"\"Remove expired storage instances from a machine.\"\"\"\n client = get_client(args)\n rj = machines_api.cleanup_machine(client, id=args.id)\n\n if args.raw:\n return rj\n if rj.get(\"success\"):\n print(json.dumps(rj, indent=1))\n else:\n print(rj.get(\"msg\", rj))\n\n\n# ---------------------------------------------------------------------------\n# defrag machines\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"IDs\", help=\"ids of machines\", type=int, nargs='+'),\n usage=\"vastai defragment machines IDs \",\n help=\"[Host] Defragment machines\",\n epilog=deindent(\"\"\"\n Defragment some of your machines. This will rearrange GPU assignments to try and make more multi-gpu offers available.\n \"\"\"),\n)\ndef defrag__machines(args):\n \"\"\"Defragment machines to make more multi-gpu offers available.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"machine_ids\": args.IDs})\n\n client = get_client(args)\n try:\n result = machines_api.defrag_machines(client, machine_ids=args.IDs)\n print(f\"defragment result: {result}\")\n except Exception as e:\n print(f\"Error: {e}\")\n\n\n# ---------------------------------------------------------------------------\n# set min_bid\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of machine to set min bid price for\", type=int),\n argument(\"--price\", help=\"per gpu min bid price in $/hour\", type=float),\n usage=\"vastai set min_bid id [--price PRICE]\",\n help=\"[Host] Set the minimum bid/rental price for a machine\",\n epilog=deindent(\"\"\"\n Change the current min bid price of machine id to PRICE.\n \"\"\"),\n)\ndef set__min_bid(args):\n \"\"\"Set the minimum bid/rental price for a machine.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"client_id\": \"me\", \"price\": args.price})\n\n client = get_client(args)\n machines_api.set_min_bid(client, id=args.id, price=args.price)\n print(\"Per gpu min bid price changed\")\n\n\n# ---------------------------------------------------------------------------\n# set defjob\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of machine to launch default instance on\", type=int),\n argument(\"--price_gpu\", help=\"per gpu rental price in $/hour\", type=float),\n argument(\"--price_inetu\", help=\"price for internet upload bandwidth in $/GB\", type=float),\n argument(\"--price_inetd\", help=\"price for internet download bandwidth in $/GB\", type=float),\n argument(\"--image\", help=\"docker container image to launch\", type=str),\n argument(\"--args\", nargs=argparse.REMAINDER, help=\"list of arguments passed to container launch\"),\n usage=\"vastai set defjob id [--api-key API_KEY] [--price_gpu PRICE_GPU] [--price_inetu PRICE_INETU] [--price_inetd PRICE_INETD] [--image IMAGE] [--args ...]\",\n help=\"[Host] Create default jobs for a machine\",\n epilog=deindent(\"\"\"\n Performs the same action as creating a background job at https://cloud.vast.ai/host/create.\n \"\"\")\n)\ndef set__defjob(args):\n \"\"\"Create default jobs for a machine.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({'machine': args.id, 'price_gpu': args.price_gpu, 'price_inetu': args.price_inetu,\n 'price_inetd': args.price_inetd, 'image': args.image, 'args': args.args})\n\n client = get_client(args)\n rj = machines_api.set_defjob(client, id=args.id, price_gpu=args.price_gpu,\n price_inetu=args.price_inetu, price_inetd=args.price_inetd,\n image=args.image, args=args.args)\n if rj.get(\"success\"):\n print(\"bids created for machine {args.id}, @ ${args.price_gpu}/gpu/day, ${args.price_inetu}/GB up, ${args.price_inetd}/GB down\".format(**locals()))\n else:\n print(rj.get(\"msg\", rj))\n\n\n# ---------------------------------------------------------------------------\n# remove defjob\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of machine to remove default instance from\", type=int),\n usage=\"vastai remove defjob id\",\n help=\"[Host] Delete default jobs\",\n)\ndef remove__defjob(args):\n \"\"\"Delete default jobs for a machine.\"\"\"\n client = get_client(args)\n rj = machines_api.remove_defjob(client, id=args.id)\n\n if rj.get(\"success\"):\n print(\"default instance for machine {machine_id} removed.\".format(machine_id=args.id))\n else:\n print(rj.get(\"msg\", rj))\n\n\n# ---------------------------------------------------------------------------\n# schedule / cancel maint\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of machine to schedule maintenance for\", type=int),\n argument(\"--sdate\", help=\"maintenance start date in unix epoch time (UTC seconds)\", type=float),\n argument(\"--duration\", help=\"maintenance duration in hours\", type=float),\n argument(\"--maintenance_category\", help=\"(optional) can be one of [power, internet, disk, gpu, software, other]\", type=str, default=\"not provided\"),\n usage=\"vastai schedule maintenance id [--sdate START_DATE --duration DURATION --maintenance_category MAINTENANCE_CATEGORY]\",\n help=\"[Host] Schedule upcoming maint window\",\n epilog=deindent(\"\"\"\n The proper way to perform maintenance on your machine is to wait until all active contracts have expired or the machine is vacant.\n 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.\n You can specify a date, duration, reason and category for the maintenance.\n\n Example: vastai schedule maint 8207 --sdate 1677562671 --duration 0.5 --maintenance_category \"power\"\n \"\"\"),\n )\ndef schedule__maint(args):\n \"\"\"Schedule upcoming maintenance window.\"\"\"\n from datetime import datetime, timezone\n from vastai.cli.util import string_to_unix_epoch\n\n dt = datetime.fromtimestamp(args.sdate, tz=timezone.utc)\n print(f\"Scheduling maintenance window starting {dt} lasting {args.duration} hours\")\n print(f\"This will notify all clients of this machine.\")\n ok = input(\"Continue? [y/n] \")\n if ok.strip().lower() != \"y\":\n return\n\n if args.explain:\n print(\"request json: \")\n print({\"client_id\": \"me\", \"sdate\": args.sdate, \"duration\": args.duration,\n \"maintenance_category\": args.maintenance_category})\n\n client = get_client(args)\n machines_api.schedule_maint(client, id=args.id, sdate=string_to_unix_epoch(args.sdate),\n duration=args.duration, maintenance_category=args.maintenance_category)\n print(f\"Maintenance window scheduled for {dt} success\")\n\n\[email protected](\n argument(\"id\", help=\"id of machine to cancel maintenance(s) for\", type=int),\n usage=\"vastai cancel maint id\",\n help=\"[Host] Cancel maint window\",\n epilog=deindent(\"\"\"\n For deleting a machine's scheduled maintenance window(s), use this cancel maint command.\n Example: vastai cancel maint 8207\n \"\"\"),\n )\ndef cancel__maint(args):\n \"\"\"Cancel scheduled maintenance window(s).\"\"\"\n print(f\"Cancelling scheduled maintenance window(s) for machine {args.id}.\")\n ok = input(\"Continue? [y/n] \")\n if ok.strip().lower() != \"y\":\n return\n\n if args.explain:\n print(\"request json: \")\n print({\"client_id\": \"me\", \"machine_id\": args.id})\n\n client = get_client(args)\n machines_api.cancel_maint(client, id=args.id)\n print(f\"Cancel maintenance window(s) scheduled for machine {args.id} success\")\n\n\n# ---------------------------------------------------------------------------\n# add network-disk\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"machines\", help=\"ids of machines to add disk to, that is networked to be on the same LAN as machine\", type=int, nargs='+'),\n argument(\"mount_point\", help=\"mount path of disk to add\", type=str),\n argument(\"-d\", \"--disk_id\", help=\"id of network disk to attach to machines in the cluster\", type=int, nargs='?'),\n usage=\"vastai add network-disk MACHINES MOUNT_PATH [options]\",\n help=\"[Host] Add Network Disk to Physical Cluster.\",\n epilog=deindent(\"\"\"\n This variant can be used to add a network disk to a physical cluster.\n When you add a network disk for the first time, you just need to specify the machine(s) and mount_path.\n When you add a network disk for the second time, you need to specify the disk_id.\n Example:\n vastai add network-disk 1 /mnt/disk1\n vastai add network-disk 1 /mnt/disk1 -d 12345\n \"\"\")\n)\ndef add__network_disk(args):\n \"\"\"Add network disk to a physical cluster.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"machines\": [int(id) for id in args.machines], \"mount_point\": args.mount_point, \"disk_id\": args.disk_id})\n\n client = get_client(args)\n result = storage_api.add_network_disk(client, machines=args.machines, mount_point=args.mount_point,\n disk_id=args.disk_id)\n\n if args.raw:\n return result\n\n print(\"Attached network disk to machines. Disk id: \" + str(result[\"disk_id\"]))\n\n\n# ---------------------------------------------------------------------------\n# dump-logs\n# ---------------------------------------------------------------------------\n\ndef _diagnostic_error(error):\n return redact_secret_text(error) or \"\"\n\n\ndef _artifact_text(value):\n if isinstance(value, str):\n return value\n return json.dumps(value, indent=2, sort_keys=True, default=str)\n\n\ndef collect_instance_log_artifacts(client, instance_id):\n files = {}\n errors = []\n\n try:\n instance_info = instances_api.show_instance(client, id=instance_id)\n files[\"instance/show-instance.json\"] = _artifact_text(instance_info)\n except Exception as e:\n errors.append({\n \"artifact\": \"instance/show-instance.json\",\n \"error\": _diagnostic_error(e),\n })\n\n for archive_name, daemon_logs in (\n (\"instance/container.log\", False),\n (\"instance/daemon.log\", True),\n ):\n try:\n logs = instances_api.logs(\n client,\n instance_id=instance_id,\n tail=INSTANCE_LOG_TAIL_LINES,\n daemon_logs=daemon_logs,\n )\n files[archive_name] = _artifact_text(logs)\n except Exception as e:\n errors.append({\n \"artifact\": archive_name,\n \"error\": _diagnostic_error(e),\n })\n\n return files, errors\n\n\[email protected](\n argument(\"machine_id\", help=\"Machine ID\", type=str),\n argument(\"--instance-id\", help=\"Instance ID to pull Vast instance logs from\", type=int),\n argument(\"--output-dir\", help=\"Directory for the diagnostic bundle (default: /tmp)\", type=str),\n argument(\n \"--include-local-host-artifacts\",\n action=\"store_true\",\n help=\"Include local OS/kaalia artifacts; only use when running on the actual Vast host\",\n ),\n usage=(\n \"vastai dump-logs \u003cmachine_id> [--instance-id INSTANCE_ID] \"\n \"[--output-dir DIR] [--include-local-host-artifacts]\"\n ),\n help=\"[Host] Bundle self-test diagnostics for support\",\n epilog=deindent(\"\"\"\n Creates a redacted diagnostic tarball containing CLI-visible self-test\n evidence. If --instance-id is provided, the command also pulls container\n and daemon logs from the Vast instance logs API.\n\n Local OS/kaalia artifacts are only collected with\n --include-local-host-artifacts. Use that option only when running this\n command on the actual host machine; from a laptop, it would collect the\n laptop's logs instead.\n\n Example:\n vastai dump-logs 12345 --instance-id 67890\n \"\"\"),\n)\ndef dump_logs(args):\n run_started_at = time.strftime(\"%Y%m%dT%H%M%SZ\", time.gmtime())\n cli_output = [f\"Manual diagnostic bundle requested for machine {args.machine_id}.\"]\n extra_files = {}\n extra_errors = []\n if args.instance_id:\n client = get_client(args)\n cli_output.append(f\"Requested Vast instance logs for instance {args.instance_id}.\")\n instance_files, instance_errors = collect_instance_log_artifacts(client, args.instance_id)\n extra_files.update(instance_files)\n extra_errors.extend(instance_errors)\n else:\n cli_output.append(\"No --instance-id provided; Vast instance logs were not requested.\")\n\n if not args.include_local_host_artifacts:\n cli_output.append(\n \"Local host OS/kaalia artifacts were not collected. Use \"\n \"--include-local-host-artifacts only when running on the actual host.\"\n )\n\n try:\n bundle = create_support_bundle(\n machine_id=args.machine_id,\n output_dir=args.output_dir,\n result={\n \"machine_id\": args.machine_id,\n \"stage\": \"manual_dump_logs\",\n \"reason\": \"Manual diagnostic bundle requested with vastai dump-logs.\",\n \"instance_id\": args.instance_id,\n \"includes_local_host_artifacts\": args.include_local_host_artifacts,\n },\n cli_output=cli_output,\n extra_files=extra_files,\n extra_errors=extra_errors,\n run_started_at=run_started_at,\n command=sys.argv,\n secrets=[getattr(args, \"api_key\", None), os.environ.get(\"VAST_API_KEY\")],\n include_local_host_artifacts=args.include_local_host_artifacts,\n )\n except Exception as e:\n error = _diagnostic_error(e)\n if getattr(args, \"raw\", False):\n return {\n \"success\": False,\n \"machine_id\": args.machine_id,\n \"error\": f\"Failed to create diagnostic bundle: {error}\",\n }\n print(f\"WARNING: failed to create diagnostic bundle: {error}\")\n sys.exit(1)\n if getattr(args, \"raw\", False):\n return bundle\n for line in format_bundle_summary(bundle):\n print(line)\n\n\n# ---------------------------------------------------------------------------\n# self-test machine\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"machine_id\", help=\"Machine ID\", type=str),\n argument(\"--debugging\", action=\"store_true\", help=\"Enable debugging output\"),\n argument(\"--ignore-requirements\", action=\"store_true\", help=\"Ignore the minimum system requirements and run the self test regardless\"),\n 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),\n argument(\"--support-bundle-dir\", help=\"Directory for failure diagnostic bundles (default: /tmp)\", type=str),\n argument(\"--no-support-bundle\", action=\"store_true\", help=\"Do not create a diagnostic tarball when the self-test fails\"),\n usage=\"vastai self-test machine \u003cmachine_id> [--debugging] [--ignore-requirements] [--test-image IMAGE]\",\n help=\"[Host] Perform a self-test on the specified machine\",\n epilog=deindent(\"\"\"\n This command tests if a machine meets specific requirements and\n runs a series of tests to ensure it's functioning correctly.\n\n Examples:\n vastai self-test machine 12345\n vastai self-test machine 12345 --debugging\n \"\"\"),\n)\ndef self_test__machine(args):\n \"\"\"\n Performs a self-test on the specified machine to verify its compliance with\n required specifications and functionality.\n \"\"\"\n instance_id = None\n result = base_result(args.machine_id)\n run_started_at = time.strftime(\"%Y%m%dT%H%M%SZ\", time.gmtime())\n cli_output = []\n instance_log_files = {}\n instance_log_errors = []\n instance_log_collected = set()\n ignore_requirements_warning = (\n \"WARNING: --ignore-requirements is set. Requirement checks are skipped as a \"\n \"pass/fail gate, and passing this self-test does not qualify this machine for verification.\"\n )\n\n if not hasattr(args, 'debugging'):\n args.debugging = False\n if not hasattr(args, 'test_image'):\n args.test_image = None\n if not hasattr(args, 'support_bundle_dir'):\n args.support_bundle_dir = None\n if not hasattr(args, 'no_support_bundle'):\n args.no_support_bundle = False\n if getattr(args, \"ignore_requirements\", False):\n result[\"warning\"] = ignore_requirements_warning\n result[\"diagnostics\"][\"requirements_ignored\"] = True\n\n def output_line(*args_to_print):\n return \" \".join(str(item) for item in args_to_print)\n\n def progress_print(*args_to_print):\n cli_output.append(output_line(*args_to_print))\n if not args.raw:\n print(*args_to_print)\n\n def debug_print(*args_to_print):\n if args.debugging:\n cli_output.append(f\"DEBUG: {output_line(*args_to_print)}\")\n if args.debugging and not args.raw:\n print(*args_to_print)\n\n def collect_instance_logs(inst_id):\n if args.no_support_bundle or not support_bundles_enabled():\n return\n if not inst_id or inst_id in instance_log_collected:\n return\n instance_log_collected.add(inst_id)\n files, errors = collect_instance_log_artifacts(client, inst_id)\n instance_log_files.update(files)\n instance_log_errors.extend(errors)\n result[\"diagnostics\"][\"instance_log_collection\"] = {\n \"instance_id\": inst_id,\n \"files\": sorted(instance_log_files.keys()),\n \"collection_errors\": instance_log_errors,\n }\n\n def ensure_support_bundle():\n if result.get(\"success\"):\n return None\n if result.get(\"diagnostics\", {}).get(\"support_bundle\"):\n return result[\"diagnostics\"][\"support_bundle\"]\n if args.no_support_bundle or not support_bundles_enabled():\n return None\n if instance_id:\n collect_instance_logs(instance_id)\n try:\n bundle = create_support_bundle(\n machine_id=args.machine_id,\n output_dir=args.support_bundle_dir,\n result=result,\n cli_output=cli_output,\n extra_files=instance_log_files,\n extra_errors=instance_log_errors,\n run_started_at=run_started_at,\n command=sys.argv,\n secrets=[getattr(args, \"api_key\", None), os.environ.get(\"VAST_API_KEY\")],\n include_local_host_artifacts=False,\n )\n except Exception as e:\n error = redact_secret_text(e) or \"\"\n result[\"diagnostics\"][\"support_bundle_error\"] = error\n progress_print(f\"WARNING: failed to create self-test diagnostic bundle: {error}\")\n return None\n result[\"diagnostics\"][\"support_bundle\"] = bundle\n return bundle\n\n def finish_failure():\n if args.raw:\n ensure_support_bundle()\n return result\n if result.get(\"warning\"):\n print(result[\"warning\"])\n render_runtime_failure()\n bundle = ensure_support_bundle()\n if bundle:\n for line in format_bundle_summary(bundle):\n print(line)\n print(f\"Test failed: {result['reason']}\")\n sys.exit(1)\n\n def set_runtime_failure(diagnostic, fallback_reason=None):\n result[\"failure\"] = diagnostic\n result[\"failure_code\"] = diagnostic[\"code\"]\n result[\"stage\"] = diagnostic.get(\"stage\") or result.get(\"stage\")\n result[\"reason\"] = fallback_reason or diagnostic.get(\"summary\") or \"\"\n result[\"diagnostics\"][\"runtime_failure\"] = diagnostic\n\n def safe_error(error):\n return redact_secret_text(error) or \"\"\n\n def render_runtime_failure():\n diagnostic = result.get(\"diagnostics\", {}).get(\"runtime_failure\")\n if not diagnostic:\n return\n progress_print(\"Runtime failure diagnostics:\")\n progress_print(f\"- code: {diagnostic.get('code')}\")\n if diagnostic.get(\"summary\"):\n progress_print(f\"- summary: {diagnostic['summary']}\")\n if diagnostic.get(\"underlying_error\"):\n progress_print(f\"- underlying error: {diagnostic['underlying_error']}\")\n endpoint = diagnostic.get(\"progress_endpoint\") or result.get(\"diagnostics\", {}).get(\"progress_endpoint\")\n if endpoint:\n if endpoint.get(\"url\"):\n progress_print(f\"- tried: {endpoint['url']}\")\n external_port = endpoint.get(\"external_port\") or endpoint.get(\"host_port\")\n if external_port:\n progress_print(f\"- external port tested: {external_port}\")\n last_bits = []\n if endpoint.get(\"last_status_code\") is not None:\n last_bits.append(f\"HTTP {endpoint['last_status_code']}\")\n if endpoint.get(\"last_error_type\"):\n last_bits.append(str(endpoint[\"last_error_type\"]))\n if endpoint.get(\"last_error\"):\n last_bits.append(str(endpoint[\"last_error\"]))\n if last_bits:\n progress_print(f\"- last result: {' - '.join(last_bits)}\")\n if endpoint.get(\"mapped_ports\"):\n progress_print(f\"- mapped container ports: {', '.join(endpoint['mapped_ports'])}\")\n if diagnostic.get(\"remediation\"):\n progress_print(f\"- remediation: {diagnostic['remediation']}\")\n steps = diagnostic.get(\"suggested_steps\") or []\n if steps:\n progress_print(\"- suggested steps:\")\n for step in steps:\n progress_print(f\" - {step}\")\n\n client = get_client(args)\n\n try:\n def http_status_code(error):\n response = getattr(error, \"response\", None)\n return getattr(response, \"status_code\", None)\n\n def offer_search_error_summary(error):\n return {\n \"status_code\": http_status_code(error),\n \"error\": safe_error(error),\n }\n\n def lookup_machine_for_offer_failure(machine_id):\n try:\n rows = machines_api.show_machine(client, id=machine_id)\n if isinstance(rows, dict):\n rows_for_status = [rows] if rows else []\n else:\n rows_for_status = rows or []\n return {\n \"status\": \"visible\" if rows_for_status else \"empty\",\n \"rows\": rows_for_status,\n \"row_count\": len(rows_for_status),\n }\n except requests.exceptions.HTTPError as e:\n status_code = http_status_code(e)\n status = \"permission_denied\" if status_code in (401, 403) else \"not_found\" if status_code == 404 else \"error\"\n return {\n \"status\": status,\n \"status_code\": status_code,\n \"error\": safe_error(e),\n \"rows\": [],\n \"row_count\": 0,\n }\n except requests.exceptions.RequestException as e:\n return {\n \"status\": \"lookup_error\",\n \"status_code\": http_status_code(e),\n \"error\": safe_error(e),\n \"rows\": [],\n \"row_count\": 0,\n }\n\n def compact_machine_lookup(machine_lookup):\n if not machine_lookup:\n return None\n return {\n \"status\": machine_lookup.get(\"status\"),\n \"status_code\": machine_lookup.get(\"status_code\"),\n \"row_count\": machine_lookup.get(\"row_count\", len(machine_lookup.get(\"rows\") or [])),\n \"error\": machine_lookup.get(\"error\"),\n }\n\n def selected_offer_for_self_test(machine_id):\n strict_query = {\n \"machine_id\": {\"eq\": machine_id},\n \"verified\": {\"eq\": \"any\"},\n \"rentable\": {\"eq\": True},\n \"rented\": {\"eq\": \"any\"},\n }\n diagnostics = {\n \"strict_offer_count\": None,\n \"broader_offer_count\": None,\n \"broader_offers\": [],\n \"search_error\": None,\n \"machine_lookup\": None,\n }\n try:\n strict_offers = offers_api.search_offers(\n client, query=strict_query, offer_type=\"on-demand\",\n order=[[\"score\", \"desc\"]], storage=5.0, no_default=True,\n )\n except requests.exceptions.HTTPError as e:\n if http_status_code(e) in (401, 403):\n diagnostics[\"search_error\"] = offer_search_error_summary(e)\n check, failure = no_offer_failure(machine_id, [], search_error=e)\n return None, (check, failure), diagnostics\n raise\n debug_print(\"Captured strict offers from search_offers:\", strict_offers)\n diagnostics[\"strict_offer_count\"] = len(strict_offers or [])\n if strict_offers:\n sorted_offers = sorted(strict_offers, key=lambda x: x.get(\"dlperf\", 0), reverse=True)\n selected = dict(sorted_offers[0])\n selected[\"machine_id\"] = selected.get(\"machine_id\") or machine_id\n debug_print(\"Selected offer found:\", selected)\n return selected, None, diagnostics\n\n broader_query = {\n \"machine_id\": {\"eq\": machine_id},\n \"verified\": {\"eq\": \"any\"},\n \"rentable\": {\"eq\": \"any\"},\n \"rented\": {\"eq\": \"any\"},\n }\n try:\n broader_offers = offers_api.search_offers(\n client, query=broader_query, offer_type=\"on-demand\",\n order=[[\"score\", \"desc\"]], storage=5.0, no_default=True,\n )\n except requests.exceptions.HTTPError as e:\n if http_status_code(e) in (401, 403):\n diagnostics[\"search_error\"] = offer_search_error_summary(e)\n check, failure = no_offer_failure(machine_id, [], search_error=e)\n return None, (check, failure), diagnostics\n raise\n diagnostics[\"broader_offer_count\"] = len(broader_offers or [])\n diagnostics[\"broader_offers\"] = [\n compact_offer_metadata(dict(offer, machine_id=offer.get(\"machine_id\", machine_id)))\n for offer in (broader_offers or [])[:5]\n ]\n machine_lookup = None\n if not broader_offers:\n machine_lookup = lookup_machine_for_offer_failure(machine_id)\n diagnostics[\"machine_lookup\"] = compact_machine_lookup(machine_lookup)\n check, failure = no_offer_failure(machine_id, broader_offers, machine_lookup=machine_lookup)\n return None, (check, failure), diagnostics\n\n selected_offer, offer_failure, offer_diagnostics = selected_offer_for_self_test(args.machine_id)\n result[\"diagnostics\"][\"offer_search\"] = offer_diagnostics\n if offer_failure:\n check, failure = offer_failure\n result[\"checks\"] = [check]\n result[\"failure\"] = failure\n result[\"failure_code\"] = failure[\"code\"]\n result[\"stage\"] = \"select_offer\"\n result[\"reason\"] = failure[\"summary\"]\n render_preflight_failure(args.machine_id, result[\"checks\"], failure, progress_print)\n return finish_failure()\n\n result[\"offer\"] = compact_offer_metadata(selected_offer)\n checks = preflight_requirement_checks(selected_offer)\n result[\"checks\"] = checks\n unmet_checks = failed_checks(checks)\n if unmet_checks:\n failure = requirement_failure(checks)\n result[\"diagnostics\"][\"preflight_failure\"] = failure\n render_preflight_failure(args.machine_id, checks, failure, progress_print)\n render_preflight_advisories(args.machine_id, checks, progress_print)\n if not args.ignore_requirements:\n result[\"failure\"] = failure\n result[\"failure_code\"] = failure[\"code\"]\n result[\"stage\"] = \"preflight_requirements\"\n result[\"reason\"] = failure[\"summary\"]\n return finish_failure()\n progress_print(\"Continuing despite unmet requirements because --ignore-requirements is set.\")\n else:\n progress_print(f\"Machine ID {args.machine_id} meets all the requirements.\")\n render_preflight_advisories(args.machine_id, checks, progress_print)\n if args.ignore_requirements:\n progress_print(ignore_requirements_warning)\n\n # ----- CUDA version to docker image mapping -----\n def cuda_map_to_image(cuda_version, compute_cap=None):\n \"\"\"Return (image, reason). Reason explains why this image was picked.\"\"\"\n docker_repo = \"vastai/test\"\n image_tag_prefix = \"self-test-v2-cuda\"\n\n def image_for(version):\n return f\"{docker_repo}:{image_tag_prefix}-{version}\"\n\n if isinstance(cuda_version, float):\n cuda_version = str(cuda_version)\n original_cuda = cuda_version\n\n # cuda-12.8 (torch 2.10) still ships sm_70 (Volta); cuda-13.0\n # (torch 2.11) never did. Neither builds sm_50/sm_60 kernels.\n # Anything pre-Volta (compute_cap \u003c 700) must use the cuda-11.8\n # legacy image.\n if compute_cap is not None and compute_cap \u003c 700:\n return (\n image_for(\"11.8\"),\n f\"compute_cap={compute_cap} below sm_70 -> forced {image_tag_prefix}-11.8\",\n )\n\n # Volta sm_70/sm_72 hosts: cuda-12.8 wheels include sm_70 but\n # cuda-13.0 wheels never did. Cap the driver-reported CUDA version\n # at 12.8 so the map below resolves to cuda-12.8 even if the\n # operator has installed a CUDA 13 driver on a V100.\n clamped_for_volta = False\n if compute_cap is not None and compute_cap \u003c 750:\n if float(cuda_version) > 12.8:\n cuda_version = \"12.8\"\n clamped_for_volta = True\n\n docker_tag_map = {\n \"11.8\": image_for(\"11.8\"),\n \"12.8\": image_for(\"12.8\"),\n \"13.0\": image_for(\"13.0\"),\n \"13.3\": image_for(\"13.3\"),\n }\n\n cap_hint = f\"compute_cap={compute_cap}\" if compute_cap is not None else \"compute_cap=unknown\"\n\n cuda_float = float(cuda_version)\n compatible_versions = sorted(float(version) for version in docker_tag_map)\n selected_version = max(\n (version for version in compatible_versions if version \u003c= cuda_float),\n default=None,\n )\n\n if selected_version is not None:\n selected_version_str = f\"{selected_version:.1f}\"\n image = docker_tag_map[selected_version_str]\n if clamped_for_volta:\n reason = f\"{cap_hint} (Volta) + cuda_max_good={original_cuda} -> clamped to {cuda_version} -> {image}\"\n elif selected_version_str == cuda_version:\n reason = f\"{cap_hint}, cuda_max_good={cuda_version} -> exact match -> {image}\"\n else:\n reason = (\n f\"{cap_hint}, cuda_max_good={original_cuda} -> \"\n f\"selected newest image \u003c= host CUDA ({selected_version_str}) -> {image}\"\n )\n return image, reason\n\n raise KeyError(f\"No CUDA version found for {cuda_version} or any lower version\")\n\n top_offer = selected_offer\n if not top_offer:\n progress_print(f\"No valid offers found for Machine ID {args.machine_id}\")\n result[\"reason\"] = \"No valid offers found.\"\n else:\n ask_contract_id = top_offer[\"id\"]\n cuda_version = top_offer[\"cuda_max_good\"]\n compute_cap = top_offer.get(\"compute_cap\")\n image_override = args.test_image or os.environ.get(\"VAST_SELF_TEST_IMAGE\")\n if image_override:\n docker_image = image_override\n image_reason = \"custom self-test image override\"\n else:\n docker_image, image_reason = cuda_map_to_image(cuda_version, compute_cap)\n result[\"diagnostics\"][\"image\"] = {\n \"image\": docker_image,\n \"reason\": image_reason,\n \"override\": bool(image_override),\n }\n\n # ----- create the test instance -----\n try:\n result[\"phase\"] = \"rental\"\n result[\"stage\"] = \"create_instance\"\n from vastai.cli.util import parse_env\n env = parse_env(\"-e TZ=PDT -e XNAME=XX4 -p 5000:5000 -p 1234:1234\")\n runtype = \"ssh_direc ssh_proxy\"\n self_test_label = f\"{SELF_TEST_INSTANCE_LABEL_PREFIX}-{args.machine_id}\"\n result[\"diagnostics\"][\"launch\"] = {\n \"runtype\": runtype,\n \"jupyter_lab\": False,\n \"ports\": [\"5000/tcp\", \"1234/tcp\"],\n \"label\": self_test_label,\n }\n\n progress_print(f\"Starting test with {docker_image} ({image_reason})\")\n rj = instances_api.create_instance(\n client,\n id=ask_contract_id,\n image=docker_image,\n disk=40,\n env=env,\n price=None,\n label=self_test_label,\n extra=None,\n onstart_cmd=\"/verification/remote.sh\",\n login=None,\n python_utf8=False,\n lang_utf8=False,\n jupyter_lab=False,\n jupyter_dir=None,\n force=False,\n cancel_unavail=False,\n template_hash=None,\n user=None,\n runtype=runtype,\n args=None,\n )\n debug_print(\"Captured instance_info from create_instance:\", rj)\n except Exception as e:\n error = safe_error(e)\n progress_print(f\"Error creating instance: {error}\")\n set_runtime_failure(\n make_failure(\n INSTANCE_CREATE_FAILED,\n stage=\"create_instance\",\n summary=\"Failed to create instance. Check the docker configuration.\",\n error=error,\n underlying_error=error,\n )\n )\n result[\"error\"] = error\n return finish_failure()\n\n instance_id = rj.get(\"new_contract\")\n if not instance_id:\n progress_print(\"Instance creation response did not contain 'new_contract'.\")\n set_runtime_failure(\n make_failure(\n INSTANCE_CREATE_MISSING_CONTRACT,\n stage=\"create_instance\",\n details=f\"Create-instance response: {rj}\",\n ),\n \"Instance creation failed.\",\n )\n else:\n # ----- helper: check if instance exists -----\n def instance_exist(inst_id):\n try:\n info = instances_api.show_instance(client, id=inst_id)\n if not info:\n return False\n status = info.get('intended_status') or info.get('actual_status')\n if status in ['destroyed', 'terminated', 'offline']:\n return False\n return True\n except requests.exceptions.HTTPError as e:\n if e.response.status_code == 404:\n return False\n debug_print(f\"HTTPError when checking instance existence: {safe_error(e)}\")\n return False\n except Exception as e:\n debug_print(f\"No instance found or Unexpected error checking instance existence: {safe_error(e)}\")\n return False\n\n # ----- helper: destroy instance silently with retries -----\n def destroy_instance_silent(inst_id, collect_logs=False):\n if collect_logs:\n collect_instance_logs(inst_id)\n max_retries = 10\n for attempt in range(1, max_retries + 1):\n try:\n if args.raw:\n with open(os.devnull, 'w') as devnull, redirect_stdout(devnull), redirect_stderr(devnull):\n instances_api.destroy_instance(client, id=inst_id)\n else:\n instances_api.destroy_instance(client, id=inst_id)\n if not args.raw:\n print(f\"Instance {inst_id} destroyed successfully on attempt {attempt}.\")\n return {\"success\": True}\n except Exception as e:\n if not args.raw:\n print(f\"Error destroying instance {inst_id}: {safe_error(e)}\")\n if attempt \u003c max_retries:\n if not args.raw:\n print(f\"Retrying in 10 seconds... (Attempt {attempt}/{max_retries})\")\n time.sleep(10)\n else:\n if not args.raw:\n print(f\"Failed to destroy instance {inst_id} after {max_retries} attempts.\")\n return {\"success\": False, \"error\": \"Max retries exceeded\"}\n\n # ----- wait for instance to start -----\n def wait_for_instance(inst_id, timeout=900, interval=10):\n start_time = time.time()\n debug_print(\"Starting wait_for_instance with ID:\", inst_id)\n\n while time.time() - start_time \u003c timeout:\n try:\n instance_info = instances_api.show_instance(client, id=inst_id)\n if not instance_info:\n progress_print(f\"No information returned for instance {inst_id}. Retrying...\")\n time.sleep(interval)\n continue\n\n status_msg = instance_info.get('status_msg', '')\n status_msg_clean = status_msg.strip() if isinstance(status_msg, str) else \"\"\n status_msg_lower = status_msg_clean.lower()\n status_msg_is_error = any(\n token in status_msg_lower\n for token in (\n \"error\",\n \"failed\",\n \"failure\",\n \"exception\",\n \"traceback\",\n \"oci runtime\",\n \"permission denied\",\n )\n )\n if status_msg_clean and status_msg_is_error:\n diagnostic = classify_status_msg(status_msg_clean) or make_failure(\n DAEMON_STARTUP_FAILED,\n stage=\"startup\",\n error=status_msg_clean,\n underlying_error=status_msg_clean,\n )\n reason = f\"Instance {inst_id} encountered an error: {status_msg_clean}\"\n progress_print(reason)\n if instance_exist(inst_id):\n destroy_instance_silent(inst_id, collect_logs=True)\n progress_print(f\"Instance {inst_id} has been destroyed due to error.\")\n else:\n progress_print(f\"Instance {inst_id} could not be destroyed or does not exist.\")\n return False, reason, diagnostic\n\n actual_status = instance_info.get('actual_status', 'unknown')\n intended_status = instance_info.get('intended_status', 'unknown')\n if actual_status == 'offline':\n reason = \"Instance offline during testing\"\n diagnostic = make_failure(INSTANCE_OFFLINE_BEFORE_TEST, stage=\"startup\")\n progress_print(reason)\n if instance_exist(inst_id):\n destroy_instance_silent(inst_id, collect_logs=True)\n progress_print(f\"Instance {inst_id} has been destroyed due to being offline.\")\n else:\n progress_print(f\"Instance {inst_id} could not be destroyed or does not exist.\")\n return False, reason, diagnostic\n\n if intended_status in ('stopped', 'exited') or actual_status in ('stopped', 'exited'):\n reason = f\"Instance {inst_id} stopped before reaching running status\"\n if status_msg_clean:\n reason = f\"{reason}: {status_msg_clean}\"\n diagnostic = classify_status_msg(status_msg_clean) if status_msg_clean else None\n if diagnostic is None:\n diagnostic = make_failure(\n DAEMON_STARTUP_FAILED,\n stage=\"startup\",\n summary=\"Instance stopped before the self-test runtime started.\",\n details=reason,\n error=status_msg_clean or None,\n underlying_error=status_msg_clean or None,\n )\n progress_print(reason)\n if instance_exist(inst_id):\n destroy_instance_silent(inst_id, collect_logs=True)\n progress_print(f\"Instance {inst_id} has been destroyed due to startup failure.\")\n else:\n progress_print(f\"Instance {inst_id} could not be destroyed or does not exist.\")\n return False, reason, diagnostic\n\n if intended_status == 'running' and actual_status == 'running':\n debug_print(f\"Instance {inst_id} is now running.\")\n return instance_info, None, None\n\n progress_print(f\"Instance {inst_id} status: {actual_status}... waiting for 'running' status.\")\n time.sleep(interval)\n\n except Exception as e:\n error = safe_error(e)\n progress_print(f\"Error retrieving instance info for {inst_id}: {error}. Retrying...\")\n debug_print(f\"Exception details: {error}\")\n time.sleep(interval)\n\n reason = f\"Instance did not become running within {timeout} seconds. Verify network configuration. Use the self-test machine function in the vastai CLI\"\n progress_print(reason)\n return False, reason, make_failure(\n INSTANCE_START_TIMEOUT,\n stage=\"startup\",\n details=reason,\n )\n\n # ----- run machine tester -----\n def run_machinetester(ip_address, port, inst_id, machine_id, delay, mapped_ports=None):\n urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)\n delay = int(delay)\n message = ''\n legacy_parser = LegacyProgressParser()\n progress_url = f\"https://{ip_address}:{port}/progress\"\n request_timeout = 10\n attempt_count = 0\n last_error_type = None\n last_error = None\n last_status_code = None\n first_connection_established = False\n\n def update_progress_endpoint():\n endpoint = make_progress_endpoint_diagnostic(\n url=progress_url,\n public_ip=ip_address,\n container_port=PROGRESS_CONTAINER_PORT,\n host_port=port,\n timeout_seconds=request_timeout,\n attempt_count=attempt_count,\n first_connection_established=first_connection_established,\n last_error_type=last_error_type,\n last_error=last_error,\n last_status_code=last_status_code,\n mapped_ports=mapped_ports,\n )\n result[\"diagnostics\"][\"progress_endpoint\"] = endpoint\n return endpoint\n\n update_progress_endpoint()\n\n def is_instance(iid):\n try:\n info = instances_api.show_instance(client, id=iid)\n debug_print(f\"is_instance(): Output from show instance: {info}\")\n if not info or not isinstance(info, dict):\n debug_print(\"is_instance(): No valid instance information received.\")\n return 'unknown'\n actual_status = info.get('actual_status', 'unknown')\n return actual_status if actual_status in ['running', 'offline', 'exited', 'created'] else 'unknown'\n except Exception as e:\n debug_print(f\"is_instance(): Error: {safe_error(e)}\")\n return 'unknown'\n\n if delay > 0:\n debug_print(f\"Sleeping for {delay} seconds before starting tests.\")\n time.sleep(delay)\n\n start_time = time.time()\n no_response_seconds = 0\n printed_lines = set()\n instance_destroyed = False\n try:\n while time.time() - start_time \u003c 600:\n status = is_instance(inst_id)\n debug_print(f\"Instance {inst_id} status: {status}\")\n\n if status == 'offline':\n reason = \"Instance offline during testing\"\n progress_print(f\"Instance {inst_id} went offline. {reason}\")\n destroy_instance_silent(inst_id, collect_logs=True)\n instance_destroyed = True\n return False, reason, make_failure(INSTANCE_OFFLINE_BEFORE_TEST, stage=\"runtime\")\n\n try:\n debug_print(f\"Sending GET request to {progress_url}\")\n attempt_count += 1\n response = requests.get(progress_url, verify=False, timeout=request_timeout)\n last_status_code = response.status_code\n last_error_type = None\n last_error = None\n\n if response.status_code == 200 and not first_connection_established:\n progress_print(\"Successfully established HTTPS connection to the server.\")\n first_connection_established = True\n\n if response.status_code == 200:\n message = response.text.strip()\n else:\n last_error_type = \"HTTPStatus\"\n last_error = f\"HTTP {response.status_code} from progress endpoint\"\n message = ''\n debug_print(f\"Received message: '{message}'\")\n except requests.exceptions.RequestException as e:\n last_status_code = None\n last_error_type = e.__class__.__name__\n last_error = safe_error(e)\n debug_print(f\"Error making HTTPS request: {last_error}\")\n message = ''\n update_progress_endpoint()\n\n if message:\n lines = message.split('\\n')\n new_lines = [line for line in lines if line not in printed_lines]\n for line in new_lines:\n diagnostic = legacy_parser.process_line(line)\n if line == 'DONE':\n progress_print(\"Test completed successfully.\")\n progress_print(\"Test passed.\")\n destroy_instance_silent(inst_id)\n instance_destroyed = True\n return True, \"\", None\n elif line.startswith('ERROR'):\n progress_print(line)\n progress_print(f\"Test failed with error: {line}.\")\n destroy_instance_silent(inst_id, collect_logs=True)\n instance_destroyed = True\n return False, line, diagnostic\n else:\n progress_print(line)\n printed_lines.add(line)\n no_response_seconds = 0\n else:\n no_response_seconds += 20\n debug_print(f\"No message received. Incremented no_response_seconds to {no_response_seconds}.\")\n\n if status == 'running' and no_response_seconds >= 120:\n endpoint = update_progress_endpoint()\n if not first_connection_established:\n progress_print(\"No response for 120s — port was never reachable.\")\n progress_print(f\"Tried: {endpoint.get('url')}\")\n if endpoint.get(\"external_port\"):\n progress_print(f\"External port tested: {endpoint.get('external_port')}\")\n if endpoint.get(\"last_error_type\") or endpoint.get(\"last_status_code\") is not None:\n last_result = endpoint.get(\"last_error_type\") or f\"HTTP {endpoint.get('last_status_code')}\"\n if endpoint.get(\"last_error\"):\n last_result = f\"{last_result}: {endpoint.get('last_error')}\"\n progress_print(f\"Last result: {last_result}\")\n progress_print(\"Possible causes:\")\n progress_print(\" 1. TCP firewall/NAT forwarding is blocking the mapped public port\")\n progress_print(\" 2. Container did not start or did not bind the progress server\")\n progress_print(\" 3. NAT loopback/hairpinning may fail when testing from the same LAN as the host\")\n 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'\")\n return_reason = \"Port never reachable within 120 seconds\"\n diagnostic = make_failure(\n PROGRESS_ENDPOINT_UNREACHABLE,\n stage=\"runtime_connect\",\n details=return_reason,\n progress_endpoint=endpoint,\n )\n elif endpoint.get(\"last_status_code\") == 200 and not endpoint.get(\"last_error_type\"):\n progress_print(\"Progress endpoint was reachable but returned no output for 120s.\")\n progress_print(f\"Tried: {endpoint.get('url')}\")\n if endpoint.get(\"external_port\"):\n progress_print(f\"External port tested: {endpoint.get('external_port')}\")\n progress_print(\"Possible causes:\")\n progress_print(\" 1. Runtime script stalled before writing progress\")\n progress_print(\" 2. Container process is alive but the test worker is hung\")\n progress_print(\" 3. Host stall, GPU hang, or slow I/O prevented progress updates\")\n return_reason = \"Progress endpoint returned no output for 120 seconds\"\n diagnostic = make_failure(\n PROGRESS_EMPTY_TIMEOUT,\n stage=\"runtime_progress\",\n details=return_reason,\n progress_endpoint=endpoint,\n )\n else:\n progress_print(\"Connection lost after initial success — instance may have crashed.\")\n progress_print(f\"Tried: {endpoint.get('url')}\")\n if endpoint.get(\"external_port\"):\n progress_print(f\"External port tested: {endpoint.get('external_port')}\")\n if endpoint.get(\"last_error_type\") or endpoint.get(\"last_status_code\") is not None:\n last_result = endpoint.get(\"last_error_type\") or f\"HTTP {endpoint.get('last_status_code')}\"\n if endpoint.get(\"last_error\"):\n last_result = f\"{last_result}: {endpoint.get('last_error')}\"\n progress_print(f\"Last result: {last_result}\")\n progress_print(\"Possible causes:\")\n progress_print(\" 1. Container crash, OOM, or runtime process exit — check docker logs\")\n progress_print(\" 2. GPU errors or reset — check dmesg for Xid errors\")\n progress_print(\" 3. Host stall or network loss after the progress server started\")\n return_reason = \"Connection lost after 120 seconds — possible crash during stress test\"\n diagnostic = make_failure(\n PROGRESS_ENDPOINT_LOST,\n stage=\"runtime_connect\",\n details=return_reason,\n progress_endpoint=endpoint,\n )\n destroy_instance_silent(inst_id, collect_logs=True)\n instance_destroyed = True\n return False, return_reason, diagnostic\n\n debug_print(\"Waiting for 20 seconds before the next check.\")\n time.sleep(20)\n\n debug_print(f\"Time limit reached. Destroying instance {inst_id}.\")\n return False, \"Test did not complete within the time limit\", make_failure(\n RUNTIME_TEST_TIMEOUT,\n stage=\"runtime\",\n )\n finally:\n if not instance_destroyed and inst_id and instance_exist(inst_id):\n destroy_instance_silent(inst_id, collect_logs=True)\n progress_print(f\"Machine: {machine_id} Done with testing remote.py results {message}\")\n warnings.simplefilter('default')\n\n # ----- main orchestration: wait then test -----\n result[\"phase\"] = \"wait\"\n result[\"stage\"] = \"wait_for_instance\"\n instance_info, wait_reason, wait_diagnostic = wait_for_instance(instance_id)\n if not instance_info:\n set_runtime_failure(wait_diagnostic, wait_reason)\n else:\n ip_address = instance_info.get(\"public_ipaddr\")\n if not ip_address:\n set_runtime_failure(\n make_failure(MISSING_PUBLIC_IP, stage=\"instance_network\"),\n \"Failed to retrieve public IP address.\",\n )\n else:\n all_ports = instance_info.get(\"ports\", {})\n port_mappings = all_ports.get(\"5000/tcp\", [])\n port = port_mappings[0].get(\"HostPort\") if port_mappings else None\n if not port:\n endpoint = make_progress_endpoint_diagnostic(\n public_ip=ip_address,\n container_port=PROGRESS_CONTAINER_PORT,\n host_port=None,\n timeout_seconds=10,\n mapped_ports=all_ports,\n )\n result[\"diagnostics\"][\"progress_endpoint\"] = endpoint\n progress_print(f\"Port 5000/tcp not found in mapped ports. Available ports: {list(all_ports.keys())}\")\n progress_print(\"Possible causes:\")\n progress_print(\" 1. The instance launch did not map the self-test progress port\")\n 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'\")\n progress_print(\" 3. Container is not exposing port 5000/tcp\")\n set_runtime_failure(\n make_failure(\n PROGRESS_PORT_NOT_MAPPED,\n stage=\"instance_network\",\n details=f\"Available ports: {list(all_ports.keys())}\",\n progress_endpoint=endpoint,\n ),\n \"Failed to retrieve mapped port.\",\n )\n else:\n delay = \"15\"\n result[\"phase\"] = \"test\"\n result[\"stage\"] = \"run_machinetester\"\n success, reason, runtime_diagnostic = run_machinetester(\n ip_address, port, instance_id, args.machine_id, delay, mapped_ports=all_ports,\n )\n result[\"success\"] = success\n result[\"reason\"] = reason\n if success:\n result[\"phase\"] = \"complete\"\n result[\"stage\"] = \"complete\"\n else:\n set_runtime_failure(runtime_diagnostic, reason)\n\n except KeyboardInterrupt:\n result[\"success\"] = False\n set_runtime_failure(\n make_failure(INTERRUPTED, stage=result.get(\"stage\")),\n \"Interrupted by user (Ctrl+C)\",\n )\n result[\"error\"] = result[\"reason\"]\n progress_print(\"\\nInterrupted — cleaning up test instance...\")\n except Exception as e:\n error = safe_error(e)\n result[\"success\"] = False\n result[\"reason\"] = error\n result[\"failure_code\"] = \"unexpected_error\"\n result[\"failure\"] = {\n \"code\": \"unexpected_error\",\n \"summary\": error,\n \"remediation\": \"Retry with --debugging and inspect the error details.\",\n }\n result[\"error\"] = error\n\n finally:\n # Always attempt to destroy the test instance, including on Ctrl+C.\n # KeyboardInterrupt is BaseException (not Exception) so it skips the\n # typed except above and lands here. Surface failures loudly: a\n # silently-leaked instance keeps billing the host.\n if instance_id:\n try:\n info = instances_api.show_instance(client, id=instance_id)\n if not info:\n debug_print(f\"Test instance {instance_id} already gone during cleanup.\")\n info = {}\n status = (info or {}).get('intended_status') or (info or {}).get('actual_status')\n if info and status not in ('destroyed', 'terminated', 'offline'):\n if not result.get(\"success\"):\n collect_instance_logs(instance_id)\n progress_print(f\"Destroying test instance {instance_id} (status: {status})...\")\n instances_api.destroy_instance(client, id=instance_id)\n progress_print(f\"Test instance {instance_id} destroyed.\")\n except KeyboardInterrupt:\n progress_print(\n f\"\\nSecond interrupt during cleanup — instance {instance_id} may still be running.\\n\"\n f\" Destroy it manually: vastai destroy instance {instance_id}\"\n )\n raise\n except requests.exceptions.HTTPError as e:\n if e.response is not None and e.response.status_code == 404:\n debug_print(f\"Test instance {instance_id} already gone during cleanup.\")\n else:\n progress_print(\n f\"WARNING: failed to destroy test instance {instance_id}: {safe_error(e)}\\n\"\n f\" Destroy it manually: vastai destroy instance {instance_id}\"\n )\n except Exception as e:\n progress_print(\n f\"WARNING: failed to destroy test instance {instance_id}: {safe_error(e)}\\n\"\n f\" Destroy it manually: vastai destroy instance {instance_id}\"\n )\n\n if args.raw:\n if not result.get(\"success\"):\n ensure_support_bundle()\n return result\n else:\n if result.get(\"warning\"):\n print(result[\"warning\"])\n if result[\"success\"]:\n if result.get(\"diagnostics\", {}).get(\"preflight_failure\"):\n print(\"Runtime checks passed, but minimum requirement checks were ignored.\")\n print(\"This run does not qualify the machine for verification.\")\n else:\n print(\"Test completed successfully.\")\n sys.exit(0)\n else:\n render_runtime_failure()\n bundle = ensure_support_bundle()\n if bundle:\n for line in format_bundle_summary(bundle):\n print(line)\n print(f\"Test failed: {result['reason']}\")\n sys.exit(1)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":80326,"content_sha256":"1c58876ada8c3912550396a6b2862e5de565c68b53c84554eb368e2140dc66d3"},{"filename":"cli/commands/metrics.py","content":"\"\"\"CLI commands for platform-wide GPU market metrics (host/admin access).\"\"\"\n\nimport json\nimport sys\nfrom datetime import datetime, timezone\n\nfrom vastai.cli.parser import argument, hidden_aliases\nfrom vastai.cli.display import display_table, deindent\nfrom vastai.cli.utils import get_parser as _get_parser, get_client\nfrom vastai.api import metrics as metrics_api\n\n\nparser = _get_parser()\n\n\n_VERIFIED_MAP = {\"true\": \"yes\", \"false\": \"no\", \"all\": \"all\"}\n_HOSTING_MAP = {\"true\": \"secure_cloud\", \"false\": \"community\", \"all\": \"all\"}\n\n_NEEDS_MACHINE_MSG = \"No metrics available. This endpoint is for hosts with active machines.\"\n\n\ndef notice_if_old_name(args, old, new):\n if getattr(args, \"command\", None) == old:\n print(f\"NOTICE: `vastai {old}` is now `vastai {new}`. The old name will be removed in the next release.\",\n file=sys.stderr)\n\n\n_GPU_CURRENT_FIELDS = (\n (\"gpu_name\", \"GPU\", \"{}\", None, True),\n (\"total\", \"Total\", \"{}\", None, False),\n (\"available\", \"Avail\", \"{}\", None, False),\n (\"unavailable\", \"Unrent\", \"{}\", None, False),\n (\"usage\", \"Usage%\", \"{:.1f}\", None, False),\n (\"rented_verified\", \"Rnt/Ver\", \"{}\", None, False),\n (\"avail_verified\", \"Avl/Ver\", \"{}\", None, False),\n (\"rented_unverified\", \"Rnt/Unv\", \"{}\", None, False),\n (\"avail_unverified\", \"Avl/Unv\", \"{}\", None, False),\n (\"dlperf\", \"DLPerf\", \"{:.1f}\", None, False),\n (\"tflops\", \"TFLOPS\", \"{:.1f}\", None, False),\n (\"tflops_per_dollar\", \"TFLOPS/$\", \"{:.1f}\", None, False),\n (\"price_p10\", \"P10 $/hr\", \"{:.4f}\", None, False),\n (\"price_median\", \"Med $/hr\", \"{:.4f}\", None, False),\n (\"price_p90\", \"P90 $/hr\", \"{:.4f}\", None, False),\n)\n\n\n_GPU_TRENDS_FIELDS = (\n (\"date\", \"Date\", \"{}\", None, True),\n (\"rented_verified\", \"Rnt/Ver\", \"{}\", None, False),\n (\"avail_verified\", \"Avl/Ver\", \"{}\", None, False),\n (\"rented_unverified\", \"Rnt/Unv\", \"{}\", None, False),\n (\"avail_unverified\", \"Avl/Unv\", \"{}\", None, False),\n (\"unavail_verified\", \"Unr/Ver\", \"{}\", None, False),\n (\"unavail_unverified\",\"Unr/Unv\", \"{}\", None, False),\n (\"total\", \"Total\", \"{}\", None, False),\n (\"rented_p10\", \"R.p10\", \"{:.4f}\", None, False),\n (\"rented_median\", \"R.med\", \"{:.4f}\", None, False),\n (\"rented_p90\", \"R.p90\", \"{:.4f}\", None, False),\n (\"avail_p10\", \"A.p10\", \"{:.4f}\", None, False),\n (\"avail_median\", \"A.med\", \"{:.4f}\", None, False),\n (\"avail_p90\", \"A.p90\", \"{:.4f}\", None, False),\n)\n\n\n_GPU_LOCATION_FIELDS = (\n (\"gpu_name\", \"GPU\", \"{}\", None, True),\n (\"city\", \"City\", \"{}\", None, True),\n (\"country_code\", \"CC\", \"{}\", None, True),\n (\"num_gpus\", \"GPUs\", \"{}\", None, False),\n (\"state\", \"State\", \"{}\", None, False),\n (\"verified\", \"Verified\", \"{}\", None, False),\n (\"latitude\", \"Lat\", \"{:.4f}\", None, False),\n (\"longitude\", \"Lon\", \"{:.4f}\", None, False),\n)\n\n\[email protected](\n argument(\"--verified\", type=str, choices=[\"true\", \"false\", \"all\"], default=\"all\",\n help=\"Filter GPUs by verification status\"),\n argument(\"--datacenter\", type=str, choices=[\"true\", \"false\", \"all\"], default=\"all\",\n help=\"Filter GPUs by datacenter hosting type\"),\n argument(\"--num-gpus\", type=str, choices=[\"all\", \"1\", \"2\", \"4\", \"8\", \"other\"], default=\"all\",\n help=\"Filter by machine GPU-count bucket (1x, 2x, 4x, 8x, or other)\"),\n usage=\"vastai show gpu-metrics [OPTIONS]\",\n aliases=hidden_aliases([\"metrics gpu\"]),\n help=\"[Host] Get current GPU market metrics\",\n epilog=deindent(\"\"\"\n Get current GPU metrics with counts, usage, performance, and pricing.\n For historical metrics, see the `show gpu-trends` command.\n Requires host or admin access.\n\n Examples:\n vastai show gpu-metrics\n vastai show gpu-metrics --verified true --datacenter true\n vastai show gpu-metrics --num-gpus 8\n vastai show gpu-metrics --raw\n \"\"\"),\n)\ndef show__gpu_metrics(args):\n \"\"\"Get current GPU metrics.\"\"\"\n notice_if_old_name(args, \"metrics gpu\", \"show gpu-metrics\")\n client = get_client(args)\n resp = metrics_api.gpu_current(\n client,\n verified=_VERIFIED_MAP[args.verified],\n hosting_type=_HOSTING_MAP[args.datacenter],\n num_gpus=args.num_gpus,\n )\n if resp.get(\"needs_machine\"):\n print(_NEEDS_MACHINE_MSG)\n return\n if args.raw:\n print(json.dumps(resp, indent=1))\n return\n display_table(resp.get(\"gpus\", []), _GPU_CURRENT_FIELDS)\n\n\[email protected](\n argument(\"name\", type=str, nargs=\"?\", default=\"RTX 5090,RTX 4090,RTX 3090\",\n help=\"GPU name, comma-separated list, or 'all'. Underscores are accepted in place of spaces.\"),\n argument(\"--verified\", type=str, choices=[\"true\", \"false\", \"all\"], default=\"all\",\n help=\"Filter by verified status\"),\n argument(\"--datacenter\", type=str, choices=[\"true\", \"false\", \"all\"], default=\"all\",\n help=\"Filter by datacenter hosting type\"),\n argument(\"--num-gpus\", type=str, choices=[\"all\", \"1\", \"2\", \"4\", \"8\", \"other\"], default=\"all\",\n help=\"Filter by machine GPU-count bucket (1x, 2x, 4x, 8x, or other)\"),\n argument(\"--start\", type=int, default=None, help=\"Start unix timestamp\"),\n argument(\"--end\", type=int, default=None, help=\"End unix timestamp\"),\n argument(\"--step\", type=int, default=None,\n help=\"Time between data points in seconds (e.g. 3600 for hourly). Minimum 60s; step may be raised server-side to cap points returned.\"),\n argument(\"--full\", action=\"store_true\", default=False,\n help=\"Show all data points instead of sampling ~20\"),\n usage=\"vastai show gpu-trends [NAME] [OPTIONS]\",\n aliases=hidden_aliases([\"metrics gpu-trends\"]),\n help=\"[Host] Get GPU market history\",\n epilog=deindent(\"\"\"\n Show GPU supply/demand and pricing trends over time. Defaults to RTX 5090, 4090, 3090\n for the last 24 hours. Requires host or admin access.\n\n Examples:\n vastai show gpu-trends\n vastai show gpu-trends \"RTX 4090\"\n vastai show gpu-trends \"RTX 4090\" --full\n vastai show gpu-trends \"RTX 4090\" --raw\n vastai show gpu-trends all --verified true --datacenter true\n vastai show gpu-trends \"RTX 5090\" --num-gpus 8\n vastai show gpu-trends \"RTX 4090,H100_SXM\" --start 1773298800 --end 1773817200 --step 3600\n \"\"\"),\n)\ndef show__gpu_trends(args):\n \"\"\"Get GPU metrics history.\"\"\"\n notice_if_old_name(args, \"metrics gpu-trends\", \"show gpu-trends\")\n # Accept underscores as space aliases: \"H100_SXM\" -> \"H100 SXM\"\n args.name = args.name.replace(\"_\", \" \")\n client = get_client(args)\n resp = metrics_api.gpu_history(\n client,\n gpu_name=args.name,\n verified=_VERIFIED_MAP[args.verified],\n hosting_type=_HOSTING_MAP[args.datacenter],\n start=args.start,\n end=args.end,\n step=args.step,\n num_gpus=args.num_gpus,\n )\n if resp.get(\"needs_machine\"):\n print(_NEEDS_MACHINE_MSG)\n return\n\n if \"gpus\" in resp:\n requested = [g.strip() for g in args.name.split(\",\")]\n gpu_items = [(name, resp[\"gpus\"][name]) for name in requested if name in resp[\"gpus\"]]\n seen = set(requested)\n for name, data in resp[\"gpus\"].items():\n if name not in seen:\n gpu_items.append((name, data))\n else:\n gpu_items = [(args.name, resp)]\n\n sd_keys = [\"rented_verified\", \"avail_verified\", \"rented_unverified\", \"avail_unverified\",\n \"unavail_verified\", \"unavail_unverified\", \"total\"]\n pr_keys = [\"rented_p10\", \"rented_median\", \"rented_p90\", \"avail_p10\", \"avail_median\", \"avail_p90\"]\n\n if args.raw:\n raw_out = {}\n for gpu_name, gpu_data in gpu_items:\n sd = gpu_data.get(\"supply_demand\", {})\n pr = gpu_data.get(\"pricing\", {})\n timestamps = sd.get(\"timestamps\", [])\n rows = []\n for i, ts in enumerate(timestamps):\n row = {\"date\": datetime.fromtimestamp(ts, tz=timezone.utc).strftime(\"%Y-%m-%d %H:%M:%S\"), \"timestamp\": ts}\n for key in sd_keys:\n arr = sd.get(key, [])\n row[key] = arr[i] if i \u003c len(arr) else None\n for key in pr_keys:\n arr = pr.get(key, [])\n row[key] = arr[i] if i \u003c len(arr) else None\n rows.append(row)\n raw_out[gpu_name] = {\"stats\": gpu_data.get(\"stats\", {}), \"rows\": rows}\n print(json.dumps(raw_out, indent=1))\n return\n\n for gpu_name, gpu_data in gpu_items:\n stats = gpu_data.get(\"stats\", {})\n print(f\"\\n=== {gpu_name} ===\")\n print(f\" tflops: {stats.get('tflops', '-')} dlperf: {stats.get('dlperf', '-')} tflops/$: {stats.get('tflops_per_dollar', '-')}\")\n sd = gpu_data.get(\"supply_demand\", {})\n pr = gpu_data.get(\"pricing\", {})\n timestamps = sd.get(\"timestamps\", [])\n if not timestamps:\n print(\" (no data)\")\n continue\n\n n = len(timestamps)\n sample_step = 1 if args.full else max(1, n // 20)\n rows = []\n for i in range(0, n, sample_step):\n row = {\"date\": datetime.fromtimestamp(timestamps[i], tz=timezone.utc).strftime(\"%Y-%m-%d %H:%M\")}\n for key in sd_keys:\n arr = sd.get(key, [])\n row[key] = arr[i] if i \u003c len(arr) else None\n for key in pr_keys:\n arr = pr.get(key, [])\n row[key] = arr[i] if i \u003c len(arr) else None\n rows.append(row)\n\n display_table(rows, _GPU_TRENDS_FIELDS)\n if sample_step > 1:\n print(f\"\\n ({n} data points, showing every {sample_step}th. Use --full for all.)\")\n\n\[email protected](\n argument(\"--verified\", type=str, choices=[\"true\", \"false\", \"all\"], default=\"all\",\n help=\"Filter by verification status\"),\n argument(\"--datacenter\", type=str, choices=[\"true\", \"false\", \"all\"], default=\"all\",\n help=\"Filter by datacenter hosting type\"),\n argument(\"--rented\", type=str, choices=[\"true\", \"false\", \"all\"], default=\"all\",\n help=\"Filter by rented status. 'false' covers both available and unrentable GPUs; use --state to tell them apart\"),\n argument(\"--state\", type=str, choices=[\"all\", \"rented\", \"available\", \"unavailable\"], default=\"all\",\n help=\"Filter by listing state. 'unavailable' means listed but neither rentable nor rented\"),\n argument(\"--gpu\", type=str, default=None,\n help=\"Filter by GPU name (comma-separated list). Underscores are accepted in place of spaces.\"),\n usage=\"vastai show gpu-locations [OPTIONS]\",\n aliases=hidden_aliases([\"metrics gpu-locations\"]),\n help=\"[Host] Get GPU location metrics\",\n epilog=deindent(\"\"\"\n Show geographic locations of GPUs on the platform. Filtering is applied\n client-side — the endpoint returns one shared dataset and the CLI narrows\n the rows locally. Requires host or admin access.\n\n Examples:\n vastai show gpu-locations\n vastai show gpu-locations --verified true --datacenter true\n vastai show gpu-locations --gpu \"RTX 4090,H100_SXM\"\n vastai show gpu-locations --rented false --raw\n \"\"\"),\n)\ndef show__gpu_locations(args):\n \"\"\"Get GPU location metrics.\"\"\"\n notice_if_old_name(args, \"metrics gpu-locations\", \"show gpu-locations\")\n client = get_client(args)\n resp = metrics_api.gpu_locations(client)\n if resp.get(\"needs_machine\"):\n print(_NEEDS_MACHINE_MSG)\n return\n\n locations = resp.get(\"locations\", [])\n\n for field in (\"verified\", \"datacenter\"):\n choice = getattr(args, field)\n if choice != \"all\":\n want = choice == \"true\"\n locations = [loc for loc in locations if bool(loc.get(field)) == want]\n\n # current servers send a three-valued state; older ones only the rented boolean\n def state_of(loc):\n return loc.get(\"state\") or (\"rented\" if loc.get(\"rented\") else \"available\")\n\n if args.rented != \"all\":\n locations = [loc for loc in locations if (state_of(loc) == \"rented\") == (args.rented == \"true\")]\n if args.state != \"all\":\n locations = [loc for loc in locations if state_of(loc) == args.state]\n if args.gpu:\n wanted_gpus = {g.strip().replace(\"_\", \" \") for g in args.gpu.split(\",\") if g.strip()}\n locations = [loc for loc in locations if loc.get(\"gpu_name\") in wanted_gpus]\n\n if args.raw:\n print(json.dumps({\"success\": True, \"locations\": locations}, indent=1))\n return\n display_table(locations, _GPU_LOCATION_FIELDS)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":13039,"content_sha256":"571f1df187e00e9d08ac82857abf84049fd1e7cfff43d8fd8b58edc025a6e0fb"},{"filename":"cli/commands/misc.py","content":"\"\"\"CLI commands for miscellaneous operations (execute, logs, reports, ssh/scp urls, snapshots, templates, self-test).\n\nNote: create/update/delete template commands are registered in offers.py.\n self_test__machine is registered in machines.py.\n\"\"\"\n\nimport json\nimport re\nimport time\nimport subprocess\nimport argparse\nimport os\nimport requests as req_lib\n\nfrom vastai.cli.parser import argument\nfrom vastai.cli.display import deindent\nfrom vastai.cli.util import (\n default_start_date, default_end_date,\n parse_day_cron_style, parse_hour_cron_style,\n validate_frequency_values, add_scheduled_job,\n)\nfrom vastai.api import instances as instances_api\nfrom vastai.api import machines as machines_api\n\n\nfrom vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401\n\n\nparser = _get_parser()\n\n\n# ---------------------------------------------------------------------------\n# execute\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of instance to execute on\", type=int),\n argument(\"COMMAND\", help=\"bash command surrounded by single quotes\", type=str),\n 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\"),\n 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)\"),\n 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)\"),\n 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),\n 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),\n usage=\"vastai execute id COMMAND\",\n help=\"Execute a (constrained) remote command on a machine\",\n epilog=deindent(\"\"\"\n Examples:\n vastai execute 99999 'ls -l -o -r'\n vastai execute 99999 'rm -r home/delete_this.txt'\n vastai execute 99999 'du -d2 -h'\n\n available commands:\n ls List directory contents\n rm Remote files or directories\n du Summarize device usage for a set of files\n\n Return value:\n Returns the output of the command which was executed on the instance, if successful. May take a few seconds to retrieve the results.\n\n \"\"\"),\n)\ndef execute(args):\n \"\"\"Execute a (constrained) remote command on a machine.\"\"\"\n json_blob = {\"command\": args.COMMAND}\n if args.explain:\n print(\"request json: \")\n print(json_blob)\n\n client = get_client(args)\n rj = instances_api.execute(client, id=args.id, command=args.COMMAND)\n\n if args.schedule:\n validate_frequency_values(args.day, args.hour, args.schedule)\n cli_command = \"execute\"\n api_endpoint = \"/api/v0/instances/command/{id}/\".format(id=args.id)\n json_blob[\"instance_id\"] = args.id\n add_scheduled_job(client, args, json_blob, cli_command, api_endpoint, \"PUT\", instance_id=args.id)\n return\n\n if rj.get(\"success\"):\n for i in range(0, 30):\n time.sleep(0.3)\n url = rj[\"result_url\"]\n r2 = req_lib.get(url)\n if r2.status_code == 200:\n filtered_text = r2.text.replace(rj[\"writeable_path\"], '')\n print(filtered_text)\n break\n else:\n print(rj)\n\n\n# ---------------------------------------------------------------------------\n# logs\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"INSTANCE_ID\", help=\"id of instance\", type=int),\n argument(\"--tail\", help=\"Number of lines to show from the end of the logs (default '1000')\", type=str),\n argument(\"--filter\", help=\"Grep filter for log entries\", type=str),\n argument(\"--daemon-logs\", help=\"Fetch daemon system logs instead of container logs\", action=\"store_true\"),\n usage=\"vastai logs INSTANCE_ID [OPTIONS] \",\n help=\"Get the logs for an instance\",\n)\ndef logs(args):\n \"\"\"Get the logs for an instance.\"\"\"\n if args.explain:\n print(\"request json: \")\n json_blob = {}\n if args.filter:\n json_blob['filter'] = args.filter\n if args.tail:\n json_blob['tail'] = args.tail\n if args.daemon_logs:\n json_blob['daemon_logs'] = 'true'\n print(json_blob)\n\n client = get_client(args)\n rj = instances_api.logs(client, instance_id=args.INSTANCE_ID, tail=args.tail,\n filter=args.filter, daemon_logs=args.daemon_logs)\n\n if isinstance(rj, str):\n cleaned_text = re.sub(r'\\n\\s*\\n', '\\n', rj)\n print(cleaned_text)\n else:\n print(rj.get(\"msg\", rj))\n\n\n# ---------------------------------------------------------------------------\n# reports\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"machine id\", type=int),\n usage=\"vastai reports ID\",\n help=\"Get the user reports for a given machine\",\n)\ndef reports(args):\n \"\"\"Get the user reports for a given machine.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"machine_id\": args.id})\n\n client = get_client(args)\n result = machines_api.reports(client, id=args.id)\n print(f\"reports: {json.dumps(result, indent=2)}\")\n\n\n# ---------------------------------------------------------------------------\n# ssh-url / scp-url\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of instance\", type=int),\n usage=\"vastai ssh-url ID\",\n help=\"ssh url helper\",\n)\ndef ssh_url(args):\n \"\"\"Get the SSH URL for an instance.\"\"\"\n return _ssh_url(args, \"ssh://\")\n\n\[email protected](\n argument(\"id\", help=\"id\", type=int),\n usage=\"vastai scp-url ID\",\n help=\"scp url helper\",\n)\ndef scp_url(args):\n \"\"\"Get the SCP URL for an instance.\"\"\"\n return _ssh_url(args, \"scp://\")\n\n\ndef _ssh_url(args, protocol):\n \"\"\"Internal helper for ssh/scp URL generation.\"\"\"\n from vastai.cli.util import DIRS\n\n json_object = None\n\n # Try reading cached SSH info\n try:\n with open(f\"{DIRS['temp']}/ssh_{args.id}.json\", 'r') as openfile:\n json_object = json.load(openfile)\n except Exception:\n pass\n\n port = None\n ipaddr = None\n\n if json_object is not None:\n ipaddr = json_object[\"ipaddr\"]\n port = json_object[\"port\"]\n\n if ipaddr is None or ipaddr.endswith('.vast.ai'):\n client = get_client(args)\n rows = instances_api.show_instances(client)\n\n if args.id:\n matches = [row for row in rows if row['id'] == args.id]\n if not matches:\n print(f\"error: no instance found with id {args.id}\")\n return 1\n instance = matches[0]\n elif len(rows) > 1:\n print(\"Found multiple running instances\")\n return 1\n else:\n instance = rows[0]\n\n ports = instance.get(\"ports\", {})\n port_22d = ports.get(\"22/tcp\", None)\n port = -1\n try:\n if port_22d is not None:\n ipaddr = instance[\"public_ipaddr\"]\n port = int(port_22d[0][\"HostPort\"])\n else:\n ipaddr = instance[\"ssh_host\"]\n port = int(instance[\"ssh_port\"]) + 1 if \"jupyter\" in instance[\"image_runtype\"] else int(instance[\"ssh_port\"])\n except Exception:\n port = -1\n\n if port > 0:\n print(f'{protocol}root@{ipaddr}:{port}')\n else:\n print('error: ssh port not found')\n\n # Cache SSH info\n try:\n with open(f\"{DIRS['temp']}/ssh_{args.id}.json\", \"w\") as outfile:\n json.dump({\"ipaddr\": ipaddr, \"port\": port}, outfile)\n except Exception:\n pass\n\n\n# ---------------------------------------------------------------------------\n# take snapshot\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"instance_id\", help=\"instance_id of the container instance to snapshot\", type=str),\n argument(\"--container_registry\", help=\"Container registry to push the snapshot to. Default will be docker.io\", type=str, default=\"docker.io\"),\n argument(\"--repo\", help=\"repo to push the snapshot to\", type=str),\n argument(\"--docker_login_user\", help=\"Username for container registry with repo\", type=str),\n argument(\"--docker_login_pass\", help=\"Password or token for container registry with repo\", type=str),\n 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\"),\n usage=\"vastai take snapshot INSTANCE_ID \"\n \"--repo REPO --docker_login_user USER --docker_login_pass PASS\"\n \"[--container_registry REGISTRY] [--pause true|false]\",\n help=\"Schedule a snapshot of a running container and push it to your repo in a container registry\",\n epilog=deindent(\"\"\"\n Takes a snapshot of a running container instance and pushes snapshot to the specified repository in container registry.\n\n Use pause=true to pause the container during commit (safer but slower),\n or pause=false to leave it running (faster but may produce a less safe snapshot).\n \"\"\"),\n)\ndef take__snapshot(args):\n \"\"\"Take a container snapshot and push to registry.\"\"\"\n instance_id = args.instance_id\n repo = args.repo\n container_registry = args.container_registry\n user = args.docker_login_user\n password = args.docker_login_pass\n pause_flag = args.pause\n\n print(f\"Taking snapshot for instance {instance_id} and pushing to repo {repo} in container registry {container_registry}\")\n\n if args.explain:\n print(\"Request JSON:\")\n print(json.dumps({\n \"id\": instance_id, \"container_registry\": container_registry,\n \"personal_repo\": repo, \"docker_login_user\": user,\n \"docker_login_pass\": password, \"pause\": pause_flag,\n }, indent=2))\n\n client = get_client(args)\n data = instances_api.take_snapshot(\n client, instance_id=instance_id, repo=repo,\n container_registry=container_registry,\n docker_login_user=user, docker_login_pass=password,\n pause=pause_flag,\n )\n\n if data.get(\"success\"):\n 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.\")\n else:\n print(data.get(\"msg\", \"Unknown error with snapshot request\"))\n\n\n\n# ---------------------------------------------------------------------------\n# Note: create__template, update__template, delete__template are in offers.py.\n# Note: self_test__machine is in machines.py.\n# ---------------------------------------------------------------------------\n","content_type":"text/x-python; charset=utf-8","language":"python","size":11300,"content_sha256":"8366a3f9b34be3c7c248d80388b3aaa23eccce6bc64397af9c9697fa28b125e8"},{"filename":"cli/commands/offers.py","content":"\"\"\"CLI commands for searching offers, templates, and benchmarks.\"\"\"\n\nimport json\nimport sys\nimport time\n\nfrom vastai.cli.parser import argument, hidden_aliases\nfrom vastai.cli.display import display_table, displayable_fields, displayable_fields_reserved, benchmarks_displayable_fields, deindent\nfrom vastai.api import offers as offers_api\n\n\nfrom vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401\n\n\nparser = _get_parser()\n\n\n# ---------------------------------------------------------------------------\n# search offers\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"-t\", \"--type\", default=\"on-demand\", help=\"Show 'on-demand', 'reserved', or 'bid'(interruptible) pricing. default: on-demand\"),\n argument(\"-i\", \"--interruptible\", dest=\"type\", const=\"bid\", action=\"store_const\", help=\"Alias for --type=bid\"),\n argument(\"-b\", \"--bid\", dest=\"type\", const=\"bid\", action=\"store_const\", help=\"Alias for --type=bid\"),\n argument(\"-r\", \"--reserved\", dest=\"type\", const=\"reserved\", action=\"store_const\", help=\"Alias for --type=reserved\"),\n argument(\"-d\", \"--on-demand\", dest=\"type\", const=\"on-demand\", action=\"store_const\", help=\"Alias for --type=on-demand\"),\n argument(\"-n\", \"--no-default\", action=\"store_true\", help=\"Disable default query\"),\n argument(\"--new\", action=\"store_true\", help=\"New search exp\"),\n argument(\"--limit\", type=int, help=\"\"),\n argument(\"--storage\", type=float, default=5.0, help=\"Amount of storage to use for pricing, in GiB. default=5.0GiB\"),\n 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-'),\n 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),\n usage=\"vastai search offers [--help] [--api-key API_KEY] [--raw] \u003cquery>\",\n help=\"Search for instance types using custom query\",\n epilog=deindent(\"\"\"\n Query syntax:\n\n query = comparison comparison...\n comparison = field op value\n field = \u003cname of a field>\n op = one of: \u003c, \u003c=, ==, !=, >=, >, in, notin\n value = \u003cbool, int, float, string> | 'any' | [value0, value1, ...]\n bool: True, False\n\n note: to pass '>' and '\u003c' on the command line, make sure to use quotes\n note: to encode a string query value (ie for gpu_name), replace any spaces ' ' with underscore '_'\n\n Examples:\n\n # search for somewhat reliable single RTX 3090 instances\n vastai search offers 'reliability > 0.98 num_gpus=1 gpu_name=RTX_3090'\n\n # search for datacenter gpus with minimal compute_cap and total_flops\n vastai search offers 'compute_cap > 610 total_flops > 5 datacenter=True'\n\n # search for reliable 4 gpu offers in Taiwan or Sweden\n vastai search offers 'reliability>0.99 num_gpus=4 geolocation in [TW,SE]'\n\n # search for reliable RTX 3090 or 4090 gpus NOT in China or Vietnam\n vastai search offers 'reliability>0.99 gpu_name in [\"RTX 4090\", \"RTX 3090\"] geolocation notin [CN,VN]'\n\n # search for machines with nvidia drivers 535.86.05 or greater (and various other options)\n vastai search offers 'disk_space>146 duration>24 gpu_ram>10 cuda_vers>=12.1 direct_port_count>=2 driver_version >= 535.86.05'\n\n # search for reliable machines with at least 4 gpus, unverified, order by num_gpus\n vastai search offers 'reliability > 0.99 num_gpus>=4 verified=False' -o 'num_gpus-'\n\n # search for arm64 cpu architecture\n vastai search offers 'cpu_arch=arm64'\n\n Available fields:\n\n Name Type Description\n\n bw_nvlink float bandwidth NVLink\n compute_cap: int cuda compute capability*100 (ie: 650 for 6.5, 700 for 7.0)\n cpu_arch string host machine cpu architecture (e.g. amd64, arm64)\n cpu_cores: int # virtual cpus\n cpu_ghz: Float # cpu clock speed GHZ\n cpu_cores_effective: float # virtual cpus you get\n cpu_ram: float system RAM in gigabytes\n cuda_vers: float machine max supported cuda version (based on driver version)\n datacenter: bool show only datacenter offers\n direct_port_count int open ports on host's router\n disk_bw: float disk read bandwidth, in MB/s\n disk_space: float disk storage space, in GB\n dlperf: float DL-perf score (see FAQ for explanation)\n dlperf_usd: float DL-perf/$\n dph: float $/hour rental cost\n driver_version: string machine's nvidia/amd driver version as 3 digit string ex. \"535.86.05,\"\n duration: float max rental duration in days\n external: bool show external offers in addition to datacenter offers\n flops_usd: float TFLOPs/$\n geolocation: string Two letter country code. Works with operators =, !=, in, notin (e.g. geolocation not in ['XV','XZ'])\n gpu_arch string host machine gpu architecture (e.g. nvidia, amd)\n gpu_max_power float GPU power limit (watts)\n gpu_max_temp float GPU temp limit (C)\n gpu_mem_bw: float GPU memory bandwidth in GB/s\n gpu_name: string GPU model name (no quotes, replace spaces with underscores, ie: RTX_3090 rather than 'RTX 3090')\n gpu_ram: float per GPU RAM in GB\n gpu_total_ram: float total GPU RAM in GB\n gpu_frac: float Ratio of GPUs in the offer to gpus in the system\n gpu_display_active: bool True if the GPU has a display attached\n has_avx: bool CPU supports AVX instruction set.\n id: int instance unique ID\n inet_down: float internet download speed in Mb/s\n inet_down_cost: float internet download bandwidth cost in $/GB\n inet_up: float internet upload speed in Mb/s\n inet_up_cost: float internet upload bandwidth cost in $/GB\n machine_id int machine id of instance\n min_bid: float current minimum bid price in $/hr for interruptible\n num_gpus: int # of GPUs\n pci_gen: float PCIE generation\n pcie_bw: float PCIE bandwidth (CPU to GPU)\n reliability: float machine reliability score (see FAQ for explanation)\n rentable: bool is the instance currently rentable\n storage_cost: float storage cost in $/GB/month\n static_ip: bool is the IP addr static/stable\n total_flops: float total TFLOPs from all GPUs\n ubuntu_version string host machine ubuntu OS version\n verified: bool is the machine verified\n vms_enabled: bool is the machine a VM instance\n \"\"\"),\n aliases=hidden_aliases([\"search instances\"]),\n)\ndef search__offers(args):\n \"\"\"Creates a query based on search parameters as in the examples above.\n\n :param argparse.Namespace args: should supply all the command-line options\n \"\"\"\n from vastai.api.query import parse_query, offers_fields, offers_alias, offers_mult\n\n try:\n if args.no_default:\n query = {}\n else:\n query = {\"verified\": {\"eq\": True}, \"external\": {\"eq\": False}, \"rentable\": {\"eq\": True}}\n\n if args.query is not None:\n query = parse_query(args.query, query, offers_fields, offers_alias, offers_mult)\n if not args.no_default:\n # Request every cpu arch like the console (CLN-3461); explicit cpu_arch wins.\n query.setdefault(\"cpu_arch\", {\"in\": [\"amd64\", \"arm64\"]})\n\n order = []\n for name in args.order.split(\",\"):\n name = name.strip()\n if not name:\n continue\n direction = \"asc\"\n field = name\n if name.strip(\"-\") != name:\n direction = \"desc\"\n field = name.strip(\"-\")\n if name.strip(\"+\") != name:\n direction = \"asc\"\n field = name.strip(\"+\")\n if field in offers_alias:\n field = offers_alias[field]\n order.append([field, direction])\n\n query[\"order\"] = order\n query[\"type\"] = args.type\n if args.limit:\n query[\"limit\"] = int(args.limit)\n query[\"allocated_storage\"] = args.storage\n if query[\"type\"] == 'interruptible':\n query[\"type\"] = 'bid'\n except ValueError as e:\n print(\"Error: \", e)\n return 1\n\n json_blob = query\n client = get_client(args)\n\n if args.new:\n json_blob = {\"select_cols\": ['*'], \"q\": query}\n if args.explain:\n print(\"request json: \")\n print(json_blob)\n r = client.put(\"/search/asks/\", json_data=json_blob)\n else:\n if args.explain:\n print(\"request json: \")\n print(json_blob)\n r = client.post(\"/bundles/\", json_data=json_blob)\n\n r.raise_for_status()\n\n if r.headers.get('Content-Type') != 'application/json':\n print(f\"invalid return Content-Type: {r.headers.get('Content-Type')}\")\n return\n\n rows = r.json()[\"offers\"]\n\n if args.raw:\n return rows\n else:\n if args.type == \"reserved\":\n display_table(rows, displayable_fields_reserved)\n else:\n display_table(rows, displayable_fields)\n\n\n# ---------------------------------------------------------------------------\n# search benchmarks\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"query\", help=\"Search query in simple query syntax (see below)\", nargs=\"*\", default=None),\n argument(\"-a\", \"--all\", action=\"store_true\", help=\"force fetching every page, even with --limit/--next-token (fetching everything is already the default otherwise)\"),\n argument(\"-l\", \"--limit\", type=int, default=None, help=\"max results per page; passing this switches to single-page mode instead of fetching everything\"),\n argument(\"-t\", \"--next-token\", dest=\"next_token\", help=\"resume from a pagination token (implies single-page mode)\"),\n usage=\"vastai search benchmarks [--help] [--api-key API_KEY] [--raw] [--limit N] [--next-token TOKEN] \u003cquery>\",\n help=\"Search for benchmark results using custom query\",\n epilog=deindent(\"\"\"\n By default this fetches every page and prints one combined result. Passing\n --limit or --next-token switches to single-page mode, which prints a page\n and asks whether to fetch the next one (add --all to still fetch every page\n in that mode). The prompt is skipped when output is piped or --raw is used;\n those print the next page token so you can resume with --next-token.\n\n --raw prints a flat list when fetching every page, and the full response\n (benchmarks plus next_token) in single-page mode.\n\n Query syntax:\n\n query = comparison comparison...\n comparison = field op value\n field = \u003cname of a field>\n op = one of: \u003c, \u003c=, ==, !=, >=, >, in, notin\n value = \u003cbool, int, float, string> | 'any' | [value0, value1, ...]\n bool: True, False\n\n note: to pass '>' and '\u003c' on the command line, make sure to use quotes\n note: to encode a string query value (ie for gpu_name), replace any spaces ' ' with underscore '_'\n\n Examples:\n\n # search for benchmarks with value > 100 on RTX 4090s for 2 specific machines\n vastai search benchmarks 'value > 100.0 gpu_name=RTX_4090 machine_id in [302,402]'\n\n # a page at a time, prompting between pages\n vastai search benchmarks 'gpu_name=RTX_5090' --limit 50\n\n Available fields:\n\n Name Type Description\n\n contract_id int ID of instance/contract reporting benchmark\n dph_base float gpu $/hr at the time of the benchmark\n gpu_name string GPU model benchmarked (e.g. RTX_4090)\n id int benchmark unique ID\n image string image used for benchmark\n last_update float date of benchmark\n launch_args string launch args of the benchmarked worker\n machine_id int id of machine benchmarked\n num_gpus int number of gpus used in benchmark\n template_hash string hash of the template benchmarked (formerly 'name')\n template_id int id of the template benchmarked\n type string benchmark type\n value float benchmark result (formerly 'score')\n \"\"\"),\n aliases=hidden_aliases([\"search benchmarks\"]),\n)\ndef search__benchmarks(args):\n \"\"\"Creates a query based on search parameters as in the examples above.\n :param argparse.Namespace args: should supply all the command-line options\n \"\"\"\n from vastai.api.query import parse_query, benchmarks_fields, fix_date_fields\n\n try:\n query = {}\n if args.query is not None:\n query = parse_query(args.query, query, benchmarks_fields)\n query = fix_date_fields(query, ['last_update'])\n except ValueError as e:\n print(\"Error: \", e)\n return 1\n\n client = get_client(args)\n # Fetch every page by default; --limit/--next-token opt into manual pagination.\n fetch_all = args.all or (args.limit is None and args.next_token is None)\n if fetch_all:\n rows = offers_api.search_benchmarks(\n client, query=query, limit=args.limit, after_token=args.next_token)\n if args.raw:\n # Flat list of every matching row -- the contract scripts have\n # always depended on for --raw. Single-page mode returns the\n # envelope instead, since only it has a token worth handing back.\n return rows\n display_table(rows, benchmarks_displayable_fields)\n return\n\n params = offers_api.benchmarks_query_args(\n query=query, limit=args.limit, after_token=args.next_token)\n page = 0\n while True:\n page += 1\n data = offers_api.search_benchmarks_v1(client, params)\n rows = data.get(\"benchmarks\") or []\n next_token = data.get(\"next_token\")\n if args.raw:\n return {\"success\": True, \"benchmarks_found\": len(rows),\n \"benchmarks\": rows, \"next_token\": next_token}\n display_table(rows, benchmarks_displayable_fields)\n if not next_token:\n return\n print(f\"Next page token: {next_token}\")\n # Never block on a prompt when piped or scripted.\n if not (sys.stdin.isatty() and sys.stdout.isatty()):\n return\n try:\n if input(f\"Fetch next page? (page {page + 1}) (y/N): \").strip().lower() != \"y\":\n return\n except (EOFError, KeyboardInterrupt):\n return\n params[\"after_token\"] = next_token\n\n\n# ---------------------------------------------------------------------------\n# search templates\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"query\", help=\"Search query in simple query syntax (see below)\", nargs=\"*\", default=None),\n usage=\"vastai search templates [--help] [--api-key API_KEY] [--raw] \u003cquery>\",\n help=\"Search for template results using custom query\",\n epilog=deindent(\"\"\"\n Query syntax:\n\n query = comparison comparison...\n comparison = field op value\n field = \u003cname of a field>\n op = one of: \u003c, \u003c=, ==, !=, >=, >, in, notin\n value = \u003cbool, int, float, string> | 'any' | [value0, value1, ...]\n bool: True, False\n\n note: to pass '>' and '\u003c' on the command line, make sure to use quotes\n note: to encode a string query value (ie for gpu_name), replace any spaces ' ' with underscore '_'\n\n Examples:\n\n # search for somewhat reliable single RTX 3090 instances, filter out any duplicates or offers that conflict with our existing stopped instances\n vastai search templates 'count_created > 100 creator_id in [38382,48982]'\n\n Available fields:\n\n Name Type Description\n\n creator_id int ID of creator\n created_at float time of initial template creation (UTC epoch timestamp)\n count_created int #instances created (popularity)\n default_tag string image default tag\n docker_login_repo string image docker repository\n id int template unique ID\n image string image used for template\n jup_direct bool supports jupyter direct\n hash_id string unique hash ID of template\n name string displayable name\n recent_create_date float last time of instance creation (UTC epoch timestamp)\n recommended_disk_space float min disk space required\n recommended bool is templated on our recommended list\n ssh_direct bool supports ssh direct\n tag string image tag\n use_ssh bool supports ssh (direct or proxy) \"\"\"),\n aliases=hidden_aliases([\"search templates\"]),\n)\ndef search__templates(args):\n \"\"\"Creates a query based on search parameters as in the examples above.\n :param argparse.Namespace args: should supply all the command-line options\n \"\"\"\n from vastai.api.query import parse_query, templates_fields, fix_date_fields\n\n try:\n query = {}\n if args.query is not None:\n query = parse_query(args.query, query, templates_fields)\n query = fix_date_fields(query, ['created_at', 'recent_create_date'])\n except ValueError as e:\n print(\"Error: \", e)\n return 1\n\n client = get_client(args)\n try:\n rows = offers_api.search_templates(client, query=query)\n print(json.dumps(rows, indent=1, sort_keys=True))\n except Exception as e:\n print(f\"Error: {e}\")\n\n\n# ---------------------------------------------------------------------------\n# search invoices\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"query\", help=\"Search query in simple query syntax (see below)\", nargs=\"*\", default=None),\n usage=\"vastai search invoices [--help] [--api-key API_KEY] [--raw] \u003cquery>\",\n help=\"Search for invoices using custom query\",\n epilog=deindent(\"\"\"\n Query syntax:\n\n query = comparison comparison...\n comparison = field op value\n field = \u003cname of a field>\n op = one of: \u003c, \u003c=, ==, !=, >=, >, in, notin\n value = \u003cbool, int, float, string> | 'any' | [value0, value1, ...]\n bool: True, False\n\n note: to pass '>' and '\u003c' on the command line, make sure to use quotes\n note: to encode a string query value (ie for gpu_name), replace any spaces ' ' with underscore '_'\n\n Examples:\n\n # search for somewhat reliable single RTX 3090 instances, filter out any duplicates or offers that conflict with our existing stopped instances\n vastai search invoices 'amount_cents>3000 '\n\n Available fields:\n\n Name Type Description\n\n id int,\n user_id int,\n when float, utc epoch timestamp of initial invoice creation\n paid_on float, actual payment date (utc epoch timestamp )\n payment_expected float, expected payment date (utc epoch timestamp )\n amount_cents int, amount of payment in cents\n is_credit bool, is a credit purchase\n is_delayed bool, is not yet paid\n balance_before float, balance before\n balance_after float, balance after\n original_amount int, original amount of payment\n event_id string,\n cut_amount int,\n cut_percent float,\n extra json,\n service string, type of payment\n stripe_charge json,\n stripe_refund json,\n stripe_payout json,\n error json,\n paypal_email string, email for paypal/wise payments\n transfer_group string,\n failed bool,\n refunded bool,\n is_check bool,\n \"\"\"),\n aliases=hidden_aliases([\"search invoices\"]),\n)\ndef search__invoices(args):\n \"\"\"Search for invoices using custom query.\"\"\"\n from vastai.api.query import parse_query, invoices_fields, fix_date_fields\n\n try:\n query = {}\n if args.query is not None:\n query = parse_query(args.query, query, invoices_fields)\n query = fix_date_fields(query, ['when', 'paid_on', 'payment_expected'])\n except ValueError as e:\n print(\"Error: \", e)\n return 1\n\n client = get_client(args)\n rows = offers_api.search_invoices(client, query=query)\n if args.raw:\n return rows\n print(json.dumps(rows, indent=1, sort_keys=True))\n\n\n# ---------------------------------------------------------------------------\n# create / update / delete template\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"--name\", help=\"name of the template\", type=str),\n argument(\"--image\", help=\"docker container image to launch\", type=str),\n argument(\"--image_tag\", help=\"docker image tag\", type=str),\n argument(\"--href\", help=\"link you want to provide\", type=str),\n argument(\"--repo\", help=\"link to repository\", type=str),\n argument(\"--login\", help=\"docker login arguments for private repo authentication, surround with ''\", type=str),\n argument(\"--env\", help=\"Contents of the 'Docker options' field\", type=str),\n argument(\"--ssh\", help=\"Launch as an ssh instance type\", action=\"store_true\"),\n argument(\"--jupyter\", help=\"Launch as a jupyter instance instead of an ssh instance\", action=\"store_true\"),\n argument(\"--direct\", help=\"Use (faster) direct connections for jupyter & ssh\", action=\"store_true\"),\n argument(\"--jupyter-dir\", help=\"For runtype 'jupyter', directory in instance to use to launch jupyter\", type=str),\n argument(\"--jupyter-lab\", help=\"For runtype 'jupyter', Launch instance with jupyter lab\", action=\"store_true\"),\n argument(\"--onstart-cmd\", help=\"contents of onstart script as single argument\", type=str),\n argument(\"--search_params\", help=\"search offers filters\", type=str),\n argument(\"-n\", \"--no-default\", action=\"store_true\", help=\"Disable default search param query args\"),\n argument(\"--disk_space\", help=\"disk storage space, in GB\", type=str),\n argument(\"--readme\", help=\"readme string\", type=str),\n argument(\"--hide-readme\", help=\"hide the readme from users\", action=\"store_true\"),\n argument(\"--desc\", help=\"description string\", type=str),\n argument(\"--public\", help=\"make template available to public\", action=\"store_true\"),\n usage=\"vastai create template\",\n help=\"Create a new template\",\n)\ndef create__template(args):\n \"\"\"Create a new template.\"\"\"\n fields = offers_api.template_fields_from_flags(\n ssh=args.ssh, jupyter=args.jupyter, direct=args.direct,\n jupyter_lab=args.jupyter_lab, login=args.login,\n hide_readme=args.hide_readme, public=args.public,\n search_params=args.search_params, no_default=args.no_default,\n always_default=True,\n )\n extra_filters = fields[\"extra_filters\"]\n\n if args.explain:\n print(\"request json: \")\n print({\"name\": args.name, \"image\": args.image, \"extra_filters\": extra_filters})\n\n client = get_client(args)\n try:\n rj = offers_api.create_template(\n client, name=args.name, image=args.image, image_tag=args.image_tag,\n href=args.href, repo=args.repo, env=args.env, onstart_cmd=args.onstart_cmd,\n jupyter_dir=args.jupyter_dir, disk_space=args.disk_space,\n readme=args.readme, desc=args.desc, **fields,\n )\n if rj.get(\"success\"):\n print(f\"New Template: {rj['template']}\")\n else:\n print(rj.get('msg', rj))\n except Exception:\n print(\"The response is not valid JSON.\")\n\n\[email protected](\n argument(\"HASH_ID\", help=\"hash id of the template\", type=str),\n argument(\"--name\", help=\"name of the template\", type=str),\n argument(\"--image\", help=\"docker container image to launch\", type=str),\n argument(\"--image_tag\", help=\"docker image tag\", type=str),\n argument(\"--href\", help=\"link you want to provide\", type=str),\n argument(\"--repo\", help=\"link to repository\", type=str),\n argument(\"--login\", help=\"docker login arguments for private repo authentication, surround with ''\", type=str),\n argument(\"--env\", help=\"Contents of the 'Docker options' field\", type=str),\n argument(\"--ssh\", help=\"Launch as an ssh instance type\", action=\"store_true\"),\n argument(\"--jupyter\", help=\"Launch as a jupyter instance instead of an ssh instance\", action=\"store_true\"),\n argument(\"--direct\", help=\"Use (faster) direct connections for jupyter & ssh\", action=\"store_true\"),\n argument(\"--jupyter-dir\", help=\"For runtype 'jupyter', directory in instance to use to launch jupyter\", type=str),\n argument(\"--jupyter-lab\", help=\"For runtype 'jupyter', Launch instance with jupyter lab\", action=\"store_true\"),\n argument(\"--onstart-cmd\", help=\"contents of onstart script as single argument\", type=str),\n argument(\"--search_params\", help=\"search offers filters\", type=str),\n argument(\"-n\", \"--no-default\", action=\"store_true\", help=\"Disable default search param query args\"),\n argument(\"--disk_space\", help=\"disk storage space, in GB\", type=str),\n argument(\"--readme\", help=\"readme string\", type=str),\n argument(\"--hide-readme\", help=\"hide the readme from users\", action=\"store_true\"),\n argument(\"--desc\", help=\"description string\", type=str),\n argument(\"--public\", help=\"make template available to public\", action=\"store_true\"),\n usage=\"vastai update template HASH_ID\",\n help=\"Update an existing template\",\n epilog=deindent(\"\"\"\n Edits create a new version of the template with a new hash id; the original stays\n valid for anyone already using it. Send the full template, not only the changed fields.\n\n Examples:\n\n 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\n \"\"\"),\n)\ndef update__template(args):\n \"\"\"Update an existing template.\"\"\"\n fields = offers_api.template_fields_from_flags(\n ssh=args.ssh, jupyter=args.jupyter, direct=args.direct,\n jupyter_lab=args.jupyter_lab, login=args.login,\n hide_readme=args.hide_readme, public=args.public,\n search_params=args.search_params, no_default=args.no_default,\n always_default=True,\n )\n extra_filters = fields[\"extra_filters\"]\n\n if args.explain:\n print(\"request json: \")\n print({\"hash_id\": args.HASH_ID, \"name\": args.name, \"image\": args.image})\n\n client = get_client(args)\n try:\n rj = offers_api.update_template(\n client, hash_id=args.HASH_ID, name=args.name, image=args.image,\n image_tag=args.image_tag, href=args.href, repo=args.repo, env=args.env,\n onstart_cmd=args.onstart_cmd, jupyter_dir=args.jupyter_dir,\n disk_space=args.disk_space, readme=args.readme, desc=args.desc,\n **fields,\n )\n if rj.get(\"success\"):\n print(f\"updated template: {json.dumps(rj['template'], indent=1)}\")\n else:\n print(rj.get(\"msg\", \"template update failed\"))\n return 1\n except Exception as e:\n print(_response_msg(getattr(e, \"response\", None)) or str(e))\n return 1\n\n\ndef _response_msg(response):\n \"\"\"Return the API's msg from an error response body, or None when there is none.\"\"\"\n if response is None:\n return None\n try:\n body = response.json()\n except ValueError:\n return None\n return body.get(\"msg\") if isinstance(body, dict) else None\n\n\[email protected](\n argument(\"--template-id\", help=\"Template ID of Template to Delete\", type=int),\n argument(\"--hash-id\", help=\"Hash ID of Template to Delete\", type=str),\n usage=\"vastai delete template [--template-id \u003cid> | --hash-id \u003chash_id>]\",\n help=\"Delete a Template\",\n)\ndef delete__template(args):\n \"\"\"Delete a template.\"\"\"\n if not args.hash_id and not args.template_id:\n print('ERROR: Must Specify either Template ID or Hash ID to delete a template')\n return\n\n if args.explain:\n print(\"request json: \")\n print({\"hash_id\": args.hash_id, \"template_id\": args.template_id})\n\n client = get_client(args)\n try:\n rj = offers_api.delete_template(client, hash_id=args.hash_id, template_id=args.template_id)\n print(rj.get('msg', rj))\n except Exception as e:\n print(f\"Error: {e}\")\n","content_type":"text/x-python; charset=utf-8","language":"python","size":30712,"content_sha256":"f884f85405096aa78ca085a64c9c2c680a34edace0b51d8c9371eb7667c9d6fb"},{"filename":"cli/commands/price_increase.py","content":"\"\"\"CLI commands for the per-row price-increase contract extension flow.\n\nBackend source of truth: ``vast/web/views/instance.py:257-312`` and\n``vast/web/price_increase_pending.py``. All three endpoints are per-row\nkeyed on ``pending_price_increase_id``; the body is exactly\n``{\"pending_price_increase_id\": int}`` with ``extra='forbid'``.\n\nThe CLI argument stays a positional **instance ID** to preserve UX\nparity with ``show instances``; this module resolves it to the row\nid with a single ``GET /pending-price-increases/`` per command\ninvocation, mirroring what the frontend modal does implicitly.\n\nStale signal is ``HTTP 404`` with body\n``{\"success\": false, \"error\": \"no_pending_price_increase\"}`` (legacy\n``HTTP 409`` is treated identically — the FE still recognises it).\n\"\"\"\n\nimport json\nimport sys\nfrom datetime import datetime, timezone\n\nfrom requests.exceptions import HTTPError\n\nfrom vastai.api import price_increase as price_increase_api\nfrom vastai.cli.display import (\n display_table,\n deindent,\n pending_price_increase_fields,\n)\nfrom vastai.cli.parser import argument\nfrom vastai.cli.utils import get_parser as _get_parser, get_client\n\nparser = _get_parser()\n\n\n# Display rescale: backend writes per-second / fractional values; the human\n# table renders the conventions the frontend uses on the in-app modal.\n_SEC_PER_HOUR = 3600\n_SEC_PER_MONTH = 3600 * 24 * 30\n\n\n# ---------------------------------------------------------------------------\n# Formatters\n# ---------------------------------------------------------------------------\n\n\ndef _fmt_epoch(ts):\n \"\"\"Render an epoch float as ``YYYY-MM-DD HH:MM UTC``; ``-`` if absent.\"\"\"\n if ts is None:\n return \"-\"\n return datetime.fromtimestamp(float(ts), tz=timezone.utc).strftime(\"%Y-%m-%d %H:%M UTC\")\n\n\ndef _fmt_delta(old, new, scale, precision):\n \"\"\"Render ``old → new`` for a resource; ``-`` when the resource did not change.\n\n A ``null`` ``new`` means the backend did not change the resource on this\n row; the column collapses to ``-`` so the table does not invent a delta.\n \"\"\"\n if new is None:\n return \"-\"\n old_part = f\"{old * scale:.{precision}f}\" if old is not None else \"?\"\n new_part = f\"{new * scale:.{precision}f}\"\n return f\"{old_part} -> {new_part}\"\n\n\ndef _prepare_row(row):\n \"\"\"Project a backend row into the table dict ``pending_price_increase_fields`` expects.\"\"\"\n return {\n \"pending_id\": row.get(\"pending_price_increase_id\"),\n \"instance_id\": row.get(\"contract_id\"),\n \"host_id\": row.get(\"host_id\"),\n \"current_end\": _fmt_epoch(row.get(\"contract_end_date\")),\n \"new_end\": _fmt_epoch(row.get(\"ask_end_date\")),\n \"gpu\": _fmt_delta(\n row.get(\"old_gpu_costpersec\"), row.get(\"new_gpu_costpersec\"),\n _SEC_PER_HOUR, 4,\n ),\n \"storage\": _fmt_delta(\n row.get(\"old_disk_ram_costpersec\"), row.get(\"new_disk_ram_costpersec\"),\n _SEC_PER_MONTH, 4,\n ),\n \"bw_up\": _fmt_delta(\n row.get(\"old_bwu_cost\"), row.get(\"new_bwu_cost\"), 1, 4,\n ),\n \"bw_down\": _fmt_delta(\n row.get(\"old_bwd_cost\"), row.get(\"new_bwd_cost\"), 1, 4,\n ),\n \"fee\": _fmt_delta(\n row.get(\"old_platform_fee\"), row.get(\"new_platform_fee\"), 100, 2,\n ),\n }\n\n\n# ---------------------------------------------------------------------------\n# Stale-row detection\n# ---------------------------------------------------------------------------\n\n\ndef _is_stale_error(err):\n \"\"\"Detect the backend's stale-row signal across the documented variants.\n\n Modern backend: ``HTTP 404`` with body\n ``{\"error\": \"no_pending_price_increase\"}``.\n Legacy: ``HTTP 409`` (the frontend still recognises it, so do we).\n \"\"\"\n response = getattr(err, \"response\", None)\n if response is None:\n return False\n status = getattr(response, \"status_code\", None)\n if status == 409:\n return True\n if status != 404:\n return False\n try:\n body = response.json() or {}\n except Exception:\n return False\n return body.get(\"error\") == price_increase_api.NO_PENDING_PRICE_INCREASE\n\n\n# ---------------------------------------------------------------------------\n# show pending-price-increases\n# ---------------------------------------------------------------------------\n\n\[email protected](\n argument(\"-q\", \"--quiet\", action=\"store_true\",\n help=\"print only pending_price_increase_id values, one per line\"),\n usage=\"vastai show pending-price-increases [OPTIONS]\",\n help=\"List pending price-increase challenges for the authenticated user\",\n epilog=deindent(\"\"\"\n Lists every pending price-increase challenge the backend has open for\n you. Each row shows the per-resource old -> new prices and the cutover\n boundary: the new rate applies only after each contract's current\n end_date — your remaining time on the current term is billed at the\n original price.\n\n Pipe the pending IDs to `vastai accept price-increase` or\n `vastai reject price-increase` (use --quiet for one ID per line):\n\n vastai show pending-price-increases --quiet\n \"\"\"),\n)\ndef show__pending_price_increases(args):\n \"\"\"List pending price-increase challenges in a human or raw form.\"\"\"\n if args.explain:\n print(\"GET /instances/pending-price-increases/\")\n client = get_client(args)\n envelope = price_increase_api.list_pending(client)\n if args.raw:\n return envelope\n rows = price_increase_api.pending_rows(envelope)\n if args.quiet:\n for row in rows:\n print(row.get(\"pending_price_increase_id\"))\n return\n if not rows:\n print(\"No pending price increases.\")\n return\n display_table([_prepare_row(r) for r in rows], pending_price_increase_fields)\n if envelope.get(\"truncated\"):\n print(\"\\nResult truncated by the server; refine with `vastai show \"\n \"pending-price-increases --raw` to retrieve the full payload.\")\n\n\n# ---------------------------------------------------------------------------\n# Shared accept/reject confirmation + fan-out\n# ---------------------------------------------------------------------------\n\n\ndef _unique_instance_ids(ids):\n \"\"\"Return integer ids in first-seen order.\"\"\"\n unique = []\n seen = set()\n for raw_id in ids:\n instance_id = int(raw_id)\n if instance_id in seen:\n continue\n seen.add(instance_id)\n unique.append(instance_id)\n return unique\n\n\ndef _resolve_instance_ids_to_pending(rows, instance_ids):\n \"\"\"Return ``[(instance_id, pending_id_or_None)]`` preserving CLI argument order.\"\"\"\n return [\n (\n iid,\n (\n price_increase_api.find_pending_for_instance(rows, iid) or {}\n ).get(\"pending_price_increase_id\"),\n )\n for iid in instance_ids\n ]\n\n\ndef _prompt_or_exit(present_imperative, instance_ids, rows, *, yes, raw=False):\n \"\"\"Confirm before mutating; honours the spec's TTY + --yes rules.\n\n Returns ``True`` to proceed, ``False`` to abort cleanly with exit 0.\n Non-TTY without ``--yes`` exits 1 with a clear message (matches the\n Two-Step Confirmation requirement).\n \"\"\"\n if yes:\n return True\n if raw:\n print(\n \"--yes is required with --raw so stdout remains machine-readable.\",\n file=sys.stderr,\n )\n sys.exit(1)\n if not sys.stdin.isatty():\n print(\n '--yes is required when stdin is not a TTY. Run \"vastai show '\n 'pending-price-increases\" to review first.',\n file=sys.stderr,\n )\n sys.exit(1)\n instance_id_set = set(instance_ids)\n matching = [r for r in rows if r.get(\"contract_id\") in instance_id_set]\n if matching:\n display_table(\n [_prepare_row(r) for r in matching], pending_price_increase_fields,\n )\n answer = input(\n f\"{present_imperative} these price increases? [y/N]: \"\n ).strip().lower()\n return answer == \"y\"\n\n\ndef _run_per_row(args, *, route_verb, api_call, present_imperative, past_tense, cutover_note):\n \"\"\"Shared accept/reject driver: resolve, confirm, fan-out, summarise.\"\"\"\n instance_ids = _unique_instance_ids(args.ids)\n if args.explain:\n for iid in instance_ids:\n print(f\"PUT /instances/{route_verb}-price-increase/ body for instance {iid}\")\n client = get_client(args)\n envelope = price_increase_api.list_pending(client)\n rows = price_increase_api.pending_rows(envelope)\n pairs = _resolve_instance_ids_to_pending(rows, instance_ids)\n has_matching_rows = any(pid is not None for _, pid in pairs)\n if has_matching_rows and not _prompt_or_exit(\n present_imperative, instance_ids, rows, yes=args.yes, raw=args.raw,\n ):\n print(\"Aborted.\")\n return\n\n acted_ids, stale, failed = [], 0, 0\n results = []\n for iid, pid in pairs:\n if pid is None:\n if not args.raw:\n print(f\"pending price increase no longer available for instance {iid}\")\n results.append({\n \"instance_id\": iid,\n \"pending_price_increase_id\": None,\n \"status\": \"stale\",\n \"error\": price_increase_api.NO_PENDING_PRICE_INCREASE,\n })\n stale += 1\n continue\n try:\n result = api_call(client, pid)\n except HTTPError as err:\n if _is_stale_error(err):\n if not args.raw:\n print(\n \"pending price increase no longer available — re-run \"\n \"vastai show pending-price-increases\"\n )\n results.append({\n \"instance_id\": iid,\n \"pending_price_increase_id\": pid,\n \"status\": \"stale\",\n \"error\": price_increase_api.NO_PENDING_PRICE_INCREASE,\n })\n stale += 1\n continue\n if not args.raw:\n print(f\"failed for instance {iid}: {err}\", file=sys.stderr)\n results.append({\n \"instance_id\": iid,\n \"pending_price_increase_id\": pid,\n \"status\": \"failed\",\n \"error\": str(err),\n })\n failed += 1\n continue\n if not args.raw:\n print(\n f\"{past_tense} pending_id={result.get('pending_price_increase_id')} \"\n f\"contract_id={result.get('contract_id')}\"\n )\n results.append({\n \"instance_id\": iid,\n \"pending_price_increase_id\": result.get(\"pending_price_increase_id\"),\n \"contract_id\": result.get(\"contract_id\"),\n \"status\": past_tense.lower(),\n \"response\": result,\n })\n acted_ids.append(result.get(\"contract_id\"))\n\n total = len(instance_ids)\n n_ok = len(acted_ids)\n if args.raw:\n exit_code = 1 if failed else 2 if stale else 0\n payload = {\n \"success\": exit_code == 0,\n \"exit_code\": exit_code,\n \"requested\": total,\n \"acted\": n_ok,\n \"stale\": stale,\n \"failed\": failed,\n \"results\": results,\n }\n print(json.dumps(payload, indent=1, sort_keys=True))\n sys.exit(exit_code)\n\n print(\n f\"\\n{past_tense} {n_ok} / Stale {stale} / Failed {failed} \"\n f\"of {total} requested.\"\n )\n if n_ok:\n ids_str = \", \".join(str(cid) for cid in acted_ids)\n print(f\"{past_tense} price increase for {n_ok} instance(s): {ids_str}\")\n if cutover_note:\n print(cutover_note)\n\n # Spec: non-stale failures take precedence over stale (so a script can\n # distinguish \"all stale\" from \"mixed failures\").\n if failed:\n sys.exit(1)\n if stale:\n sys.exit(2)\n\n\n# ---------------------------------------------------------------------------\n# accept price-increase\n# ---------------------------------------------------------------------------\n\n\[email protected](\n argument(\"ids\", help=\"instance IDs to accept (one or more).\", type=int, nargs=\"+\"),\n argument(\"--yes\", \"-y\", action=\"store_true\",\n help=\"skip the interactive prompt; required when stdin is not a TTY\"),\n usage=\"vastai accept price-increase ID [ID ...] [--yes]\",\n help=\"Accept one or more pending host price increases\",\n epilog=deindent(\"\"\"\n Review pending price increases with `vastai show pending-price-increases`\n before accepting. The CLI fans out one PUT per instance ID (no batch\n endpoint exists on the backend), and the new rate applies only after\n each contract's current end_date — your remaining time on the current\n term is billed at the original price.\n\n Examples:\n vastai accept price-increase 123 --yes\n vastai accept price-increase 1 2 3 --yes\n\n Exit codes:\n 0 — every requested row was accepted\n 1 — at least one row failed (non-stale)\n 2 — at least one row was stale (re-run `show pending-price-increases`)\n \"\"\"),\n)\ndef accept__price_increase(args):\n \"\"\"Per-row accept for pending price-increase challenges.\"\"\"\n return _run_per_row(\n args,\n route_verb=\"accept\",\n api_call=price_increase_api.accept,\n present_imperative=\"Accept\",\n past_tense=\"Accepted\",\n cutover_note=\"New rate applies after each contract's current end_date.\",\n )\n\n\n# ---------------------------------------------------------------------------\n# reject price-increase\n# ---------------------------------------------------------------------------\n\n\[email protected](\n argument(\"ids\", help=\"instance IDs to reject (one or more).\", type=int, nargs=\"+\"),\n argument(\"--yes\", \"-y\", action=\"store_true\",\n help=\"skip the interactive prompt; required when stdin is not a TTY\"),\n usage=\"vastai reject price-increase ID [ID ...] [--yes]\",\n help=\"Reject one or more pending host price increases\",\n epilog=deindent(\"\"\"\n Tombstones the matching pending rows on the backend (no cutover follows).\n Use `vastai show pending-price-increases` to review the rows first; the\n CLI fans out one PUT per instance ID.\n\n Exit codes mirror `accept price-increase`:\n 0 — every requested row was rejected\n 1 — at least one row failed (non-stale)\n 2 — at least one row was stale\n \"\"\"),\n)\ndef reject__price_increase(args):\n \"\"\"Per-row reject for pending price-increase challenges.\"\"\"\n return _run_per_row(\n args,\n route_verb=\"reject\",\n api_call=price_increase_api.reject,\n present_imperative=\"Reject\",\n past_tense=\"Rejected\",\n cutover_note=None,\n )\n","content_type":"text/x-python; charset=utf-8","language":"python","size":14896,"content_sha256":"4618c7764dc6ef9d5a91b07ebfc945359af2830c721104a57ba04f84445604b7"},{"filename":"cli/commands/repl.py","content":"\"\"\"`vastai repl` — an interactive shell over the CLI's own commands.\"\"\"\nfrom vastai.cli.display import deindent\nfrom vastai.cli.utils import get_parser as _get_parser\n\nparser = _get_parser()\n\n\[email protected](\n usage=\"vastai repl\",\n help=\"Start an interactive vastai shell\",\n epilog=deindent(\"\"\"\n Runs vastai commands from one long-lived prompt, with Tab completion\n over commands, flags and your live instance ids:\n\n vast> show instances\n vast> destroy instance \u003ctab>\n\n Commands behave exactly as they do from the shell — same auth, same\n flags, same output.\n \"\"\"),\n)\ndef repl(args):\n from vastai.cli.repl.session import run_repl\n run_repl(args)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":719,"content_sha256":"e9db5c6fbf0d6a9b40d627c7684a7110c8d6f073f76fc951306ba6af7a98969f"},{"filename":"cli/commands/storage.py","content":"\"\"\"CLI commands for managing volumes, storage, copies, and cloud sync.\"\"\"\n\nimport json\nimport os\nimport time\nimport subprocess\nfrom urllib.parse import quote, urlparse\n\nfrom vastai.cli.parser import argument\nfrom vastai.cli.display import (\n display_table, vol_displayable_fields, nw_vol_displayable_fields,\n volume_fields, connection_fields, deindent,\n instance_backup_fields, instance_backup_file_fields,\n)\nfrom vastai.cli.display import strip_strings\nfrom vastai.cli.util import (\n default_start_date,\n parse_day_cron_style, parse_hour_cron_style,\n validate_frequency_values, add_scheduled_job,\n)\nfrom vastai.api import storage as storage_api\nfrom vastai.api import offers as offers_api\n\n\nfrom vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401\n\n\nparser = _get_parser()\n\n\n# ---------------------------------------------------------------------------\n# copy / cancel copy / cancel sync\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"src\", help=\"Source location for copy operation (supports multiple formats)\", type=str),\n argument(\"dst\", help=\"Target location for copy operation (supports multiple formats)\", type=str),\n argument(\"-i\", \"--identity\", help=\"Location of ssh private key\", type=str),\n usage=\"vastai copy SRC DST\",\n help=\"Copy directories between instances and/or local\",\n epilog=deindent(\"\"\"\n Copies a directory from a source location to a target location. Each of source and destination\n directories can be either local or remote, subject to appropriate read and write\n permissions required to carry out the action.\n\n Supported location formats:\n - [instance_id:]path (legacy format, still supported)\n - C.instance_id:path (container copy format)\n - cloud_service:path (cloud service format)\n - cloud_service.cloud_service_id:path (cloud service with ID)\n - local:path (explicit local path)\n - V.volume_id:path (volume copy)\n\n 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)\n You can see more information about constraints here: https://vast.ai/docs/gpu-instances/data-movement#constraints\n Volume copy is currently only supported for copying to other volumes, instances, or cloud services, not local.\n\n Examples:\n vastai copy 6003036:/workspace/ 6003038:/workspace/\n vastai copy C.11824:/data/test local:data/test\n vastai copy local:data/test C.11824:/data/test\n vastai copy drive:/folder/file.txt C.6003036:/workspace/\n vastai copy s3.101:/data/ C.6003036:/workspace/\n vastai copy hf.101:/my-bucket/data/ C.6003036:/workspace/\n vastai copy V.1234:/file C.5678:/workspace/\n vastai copy V.1234:/file s3.101:/workspace/\n\n The first example copy syncs all files from the absolute directory '/workspace' on instance 6003036 to the directory '/workspace' on instance 6003038.\n The second example copy syncs files from container 11824 to the local machine using structured syntax.\n The third example copy syncs files from local to container 11824 using structured syntax.\n The fourth example copy syncs files from Google Drive to an instance.\n The fifth example copy syncs files from S3 bucket with id 101 to an instance.\n The sixth example copy syncs files from a Hugging Face bucket using connection id 101 to an instance.\n The seventh example copy syncs files from volume 1234 to instance 5678.\n The eighth example copy syncs files from volume 1234 to S3 bucket with id 101.\n \"\"\"),\n)\ndef copy(args):\n \"\"\"Transfer data from one instance to another.\"\"\"\n from vastai.utils import parse_vast_url\n client = get_client(args)\n\n (src_id, src_path) = parse_vast_url(args.src)\n (dst_id, dst_path) = parse_vast_url(args.dst)\n\n print(f\"copying {str(src_id)+':' if src_id else ''}{src_path} {str(dst_id)+':' if dst_id else ''}{dst_path}\")\n\n if args.explain:\n print(\"request json: \")\n print({\"client_id\": \"me\", \"src_id\": src_id, \"dst_id\": dst_id, \"src_path\": src_path, \"dst_path\": dst_path})\n\n rj = storage_api.copy(client, src_id=src_id, dst_id=dst_id, src_path=src_path, dst_path=dst_path)\n\n if (rj.get(\"success\")) and ((src_id is None or src_id == \"local\") or (dst_id is None or dst_id == \"local\")):\n ssh_args = [\"ssh\"]\n if args.identity is not None:\n ssh_args.extend([\"-i\", args.identity])\n if (src_id is None or src_id == \"local\"):\n remote_port = rj[\"dst_port\"]\n remote_addr = rj[\"dst_addr\"]\n ssh_cmd = \" \".join(ssh_args + [\"-p\", str(remote_port), \"-o\", \"StrictHostKeyChecking=no\"])\n cmd = [\n \"rsync\", \"-arz\", \"-v\", \"--progress\", \"--rsh=ssh\", \"-e\", ssh_cmd,\n src_path, f\"vastai_kaalia@{remote_addr}::{dst_id}/{dst_path}\",\n ]\n print(\" \".join(cmd))\n subprocess.run(cmd)\n elif (dst_id is None or dst_id == \"local\"):\n os.makedirs(dst_path, exist_ok=True)\n remote_port = rj[\"src_port\"]\n remote_addr = rj[\"src_addr\"]\n ssh_cmd = \" \".join(ssh_args + [\"-p\", str(remote_port), \"-o\", \"StrictHostKeyChecking=no\"])\n cmd = [\n \"rsync\", \"-arz\", \"-v\", \"--progress\", \"--rsh=ssh\", \"-e\", ssh_cmd,\n f\"vastai_kaalia@{remote_addr}::{src_id}/{src_path}\", dst_path,\n ]\n print(\" \".join(cmd))\n subprocess.run(cmd)\n else:\n if rj.get(\"success\"):\n print(\"Remote to Remote copy initiated - check instance status bar for progress updates (~30 seconds delayed).\")\n else:\n if rj.get(\"msg\") == \"src_path not supported VMs.\":\n print(\"copy between VM instances does not currently support subpaths (only full disk copy)\")\n elif rj.get(\"msg\") == \"dst_path not supported for VMs.\":\n print(\"copy between VM instances does not currently support subpaths (only full disk copy)\")\n else:\n print(rj.get(\"msg\", rj))\n\n\[email protected](\n argument(\"dst\", help=\"instance_id:/path to target of copy operation\", type=str),\n usage=\"vastai cancel copy DST\",\n help=\"Cancel a remote copy in progress, specified by DST id\",\n epilog=deindent(\"\"\"\n Use this command to cancel any/all current remote copy operations copying to a specific named instance, given by DST.\n\n Examples:\n vastai cancel copy 12371\n\n The first example cancels all copy operations currently copying data into instance 12371\n\n \"\"\"),\n)\ndef cancel__copy(args):\n \"\"\"Cancel a remote copy in progress, specified by DST id.\"\"\"\n client = get_client(args)\n dst_id = args.dst\n if dst_id is None:\n print(\"invalid arguments\")\n return\n\n print(f\"canceling remote copies to {dst_id} \")\n\n rj = storage_api.cancel_copy(client, dst_id=dst_id)\n if rj.get(\"success\"):\n print(\"Remote copy canceled - check instance status bar for progress updates (~30 seconds delayed).\")\n else:\n print(rj.get(\"msg\", rj))\n\n\[email protected](\n argument(\"dst\", help=\"instance_id:/path to target of sync operation\", type=str),\n usage=\"vastai cancel sync DST\",\n help=\"Cancel a remote copy in progress, specified by DST id\",\n epilog=deindent(\"\"\"\n Use this command to cancel any/all current remote cloud sync operations copying to a specific named instance, given by DST.\n\n Examples:\n vastai cancel sync 12371\n\n The first example cancels all copy operations currently copying data into instance 12371\n\n \"\"\"),\n)\ndef cancel__sync(args):\n \"\"\"Cancel a remote cloud sync in progress, specified by DST id.\"\"\"\n client = get_client(args)\n dst_id = args.dst\n if dst_id is None:\n print(\"invalid arguments\")\n return\n\n print(f\"canceling remote copies to {dst_id} \")\n\n rj = storage_api.cancel_sync(client, dst_id=dst_id)\n if rj.get(\"success\"):\n print(\"Remote copy canceled - check instance status bar for progress updates (~30 seconds delayed).\")\n else:\n print(rj.get(\"msg\", rj))\n\n\n# ---------------------------------------------------------------------------\n# cloud copy\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"--src\", help=\"path to source of object to copy\", type=str),\n argument(\"--dst\", help=\"path to target of copy operation\", type=str, default=\"/workspace\"),\n argument(\"--instance\", help=\"id of the instance\", type=str),\n argument(\"--connection\", help=\"id of cloud connection on your account (get from calling 'vastai show connections')\", type=str),\n argument(\"--transfer\", help=\"type of transfer, possible options include Instance To Cloud and Cloud To Instance\", type=str, default=\"Instance to Cloud\"),\n argument(\"--dry-run\", help=\"show what would have been transferred\", action=\"store_true\"),\n argument(\"--size-only\", help=\"skip based on size only, not mod-time or checksum\", action=\"store_true\"),\n argument(\"--ignore-existing\", help=\"skip all files that exist on destination\", action=\"store_true\"),\n argument(\"--update\", help=\"skip files that are newer on the destination\", action=\"store_true\"),\n argument(\"--delete-excluded\", help=\"delete files on dest excluded from transfer\", action=\"store_true\"),\n 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\"),\n 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)\"),\n 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)\"),\n 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),\n 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),\n usage=\"vastai cloud copy --src SRC --dst DST --instance INSTANCE_ID --connection CONNECTION_ID --transfer TRANSFER_TYPE\",\n help=\"Copy files/folders to and from cloud providers\",\n epilog=deindent(\"\"\"\n Copies a directory from a source location to a target location. Each of source and destination\n directories can be either local or remote, subject to appropriate read and write\n permissions required to carry out the action. The format for both src and dst is [instance_id:]path.\n You can find more information about the cloud copy operation here: https://vast.ai/docs/gpu-instances/cloud-sync\n\n Examples:\n vastai show connections\n ID NAME Cloud Type\n 1001 test_dir drive\n 1003 data_dir drive\n 1004 hf_data hf\n\n vastai cloud copy --src /folder --dst /workspace --instance 6003036 --connection 1001 --transfer \"Instance To Cloud\"\n vastai cloud copy --src my-bucket/data --dst /workspace --instance 6003036 --connection 1004 --transfer \"Cloud To Instance\"\n\n The first cloud copy example copies all contents of /folder on instance 6003036 into /workspace in gdrive connection 'test_dir'.\n The second cloud copy example copies my-bucket/data from Hugging Face connection 'hf_data' into /workspace on the instance.\n\n Note: Hugging Face cloud destinations must already exist as bucket paths; directories are not created automatically.\n \"\"\"),\n)\ndef cloud__copy(args):\n \"\"\"Transfer data to/from cloud providers.\"\"\"\n client = get_client(args)\n\n if (args.src is None) and (args.dst is None):\n print(\"invalid arguments\")\n return\n\n flags = storage_api.rclone_flags(\n dry_run=args.dry_run, size_only=args.size_only,\n ignore_existing=args.ignore_existing, update=args.update,\n delete_excluded=args.delete_excluded,\n )\n\n print(f\"copying {args.src} {args.dst} {args.instance} {args.connection} {args.transfer}\")\n\n req_json = {\n \"src\": args.src,\n \"dst\": args.dst,\n \"instance_id\": args.instance,\n \"selected\": args.connection,\n \"transfer\": args.transfer,\n \"flags\": flags\n }\n\n if args.explain:\n print(\"request json: \")\n print(req_json)\n\n if args.schedule:\n validate_frequency_values(args.day, args.hour, args.schedule)\n r = client.get(f\"/instances/{args.instance}/\", query_args={\"owner\": \"me\"})\n r.raise_for_status()\n row = r.json()[\"instances\"]\n\n if args.transfer.lower() == \"instance to cloud\":\n if row:\n up_cost = row.get(\"internet_up_cost_per_tb\", None)\n if up_cost is not None:\n confirm = input(\n f\"Internet upload cost is ${up_cost} per TB. \"\n \"Are you sure you want to schedule a cloud backup? (y/n): \"\n ).strip().lower()\n if confirm != \"y\":\n print(\"Cloud backup scheduling aborted.\")\n return\n else:\n 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.\")\n\n cli_command = \"cloud copy\"\n api_endpoint = \"/api/v0/commands/rclone/\"\n contract_end_date = row.get(\"end_date\", None)\n add_scheduled_job(client, args, req_json, cli_command, api_endpoint, \"POST\", instance_id=args.instance, contract_end_date=contract_end_date)\n return\n else:\n print(\"Instance not found. Please check the instance ID.\")\n return\n\n rj = storage_api.cloud_copy(client, src=args.src, dst=args.dst, instance=args.instance,\n connection=args.connection, transfer=args.transfer, flags=flags)\n if rj:\n print(\"Cloud Copy Started - check instance status bar for progress updates (~30 seconds delayed).\")\n print(\"When the operation is finished you should see 'Cloud Copy Operation Finished' in the instance status bar.\")\n\n\n# ---------------------------------------------------------------------------\n# show connections\n# ---------------------------------------------------------------------------\n\[email protected](\n usage=\"vastai show connections [--api-key API_KEY] [--raw]\",\n help=\"Display user's cloud connections\",\n)\ndef show__connections(args):\n \"\"\"Show the user's cloud connections.\"\"\"\n client = get_client(args)\n rows = storage_api.show_connections(client)\n\n if args.raw:\n return rows\n else:\n display_table(rows, connection_fields)\n\n\n# ---------------------------------------------------------------------------\n# Note: show__network_disks and add__network_disk are in machines.py.\n# ---------------------------------------------------------------------------\n\n\n# ---------------------------------------------------------------------------\n# search volumes\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"-n\", \"--no-default\", action=\"store_true\", help=\"Disable default query\"),\n argument(\"--limit\", type=int, help=\"\"),\n argument(\"--storage\", type=float, default=1.0, help=\"Amount of storage to use for pricing, in GiB. default=1.0GiB\"),\n 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-'),\n argument(\"query\", help=\"Query to search for. default: 'external=false verified=true disk_space>=1', pass -n to ignore default\", nargs=\"*\", default=None),\n usage=\"vastai search volumes [--help] [--api-key API_KEY] [--raw] \u003cquery>\",\n help=\"Search for volume offers using custom query\",\n epilog=deindent(\"\"\"\n Query syntax:\n\n query = comparison comparison...\n comparison = field op value\n field = \u003cname of a field>\n op = one of: \u003c, \u003c=, ==, !=, >=, >, in, notin\n value = \u003cbool, int, float, string> | 'any' | [value0, value1, ...]\n bool: True, False\n\n note: to pass '>' and '\u003c' on the command line, make sure to use quotes\n note: to encode a string query value (ie for gpu_name), replace any spaces ' ' with underscore '_'\n\n Examples:\n\n # search for volumes with greater than 50GB of available storage and greater than 500 Mb/s upload and download speed\n vastai search volumes \"disk_space>50 inet_up>500 inet_down>500\"\n\n Available fields:\n\n Name Type Description\n\n cpu_arch: string host machine cpu architecture (e.g. amd64, arm64)\n cuda_vers: float machine max supported cuda version (based on driver version)\n datacenter: bool show only datacenter offers\n disk_bw: float disk read bandwidth, in MB/s\n disk_space: float disk storage space, in GB\n driver_version: string machine's nvidia/amd driver version as 3 digit string ex. \"535.86.05\"\n duration: float max rental duration in days\n geolocation: string Two letter country code. Works with operators =, !=, in, notin (e.g. geolocation not in ['XV','XZ'])\n gpu_arch: string host machine gpu architecture (e.g. nvidia, amd)\n gpu_name: string GPU model name (no quotes, replace spaces with underscores, ie: RTX_3090 rather than 'RTX 3090')\n has_avx: bool CPU supports AVX instruction set.\n id: int volume offer unique ID\n inet_down: float internet download speed in Mb/s\n inet_up: float internet upload speed in Mb/s\n machine_id: int machine id of volume offer\n pci_gen: float PCIE generation\n pcie_bw: float PCIE bandwidth (CPU to GPU)\n reliability: float machine reliability score (see FAQ for explanation)\n storage_cost: float storage cost in $/GB/month\n static_ip: bool is the IP addr static/stable\n total_flops: float total TFLOPs from all GPUs\n ubuntu_version: string host machine ubuntu OS version\n verified: bool is the machine verified\n \"\"\"),\n)\ndef search__volumes(args):\n \"\"\"Search for volume offers using custom query.\"\"\"\n from vastai.api.query import parse_query, vol_offers_fields, offers_alias, offers_mult\n\n try:\n if args.no_default:\n query = {}\n else:\n query = {\"verified\": {\"eq\": True}, \"external\": {\"eq\": False}, \"disk_space\": {\"gte\": 1}}\n\n if args.query is not None:\n query = parse_query(args.query, query, vol_offers_fields, {}, offers_mult)\n\n order = []\n for name in args.order.split(\",\"):\n name = name.strip()\n if not name:\n continue\n direction = \"asc\"\n field = name\n if name.strip(\"-\") != name:\n direction = \"desc\"\n field = name.strip(\"-\")\n if name.strip(\"+\") != name:\n direction = \"asc\"\n field = name.strip(\"+\")\n if field in offers_alias:\n field = offers_alias[field]\n order.append([field, direction])\n except ValueError as e:\n print(\"Error: \", e)\n return 1\n\n client = get_client(args)\n\n if args.explain:\n print(\"request json: \")\n print(query)\n\n rows = offers_api.search_volumes(\n client, query=query, order=order,\n limit=args.limit, storage=args.storage,\n no_default=True, # defaults already applied above\n )\n\n if args.raw:\n return rows\n else:\n display_table(rows, vol_displayable_fields)\n\n\n# ---------------------------------------------------------------------------\n# search network volumes\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"-n\", \"--no-default\", action=\"store_true\", help=\"Disable default query\"),\n argument(\"--limit\", type=int, help=\"\"),\n argument(\"--storage\", type=float, default=1.0, help=\"Amount of storage to use for pricing, in GiB. default=1.0GiB\"),\n 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-'),\n argument(\"query\", help=\"Query to search for. default: 'external=false verified=true disk_space>=1', pass -n to ignore default\", nargs=\"*\", default=None),\n usage=\"vastai search network volumes [--help] [--api-key API_KEY] [--raw] \u003cquery>\",\n help=\"Search for network volume offers using custom query\",\n epilog=deindent(\"\"\"\n Query syntax:\n\n query = comparison comparison...\n comparison = field op value\n field = \u003cname of a field>\n op = one of: \u003c, \u003c=, ==, !=, >=, >, in, notin\n value = \u003cbool, int, float, string> | 'any' | [value0, value1, ...]\n bool: True, False\n\n note: to pass '>' and '\u003c' on the command line, make sure to use quotes\n note: to encode a string query value (ie for gpu_name), replace any spaces ' ' with underscore '_'\n\n Examples:\n\n # search for volumes with greater than 50GB of available storage and greater than 500 Mb/s upload and download speed\n vastai search volumes \"disk_space>50 inet_up>500 inet_down>500\"\n\n Available fields:\n\n Name Type Description\n duration: float max rental duration in days\n geolocation: string Two letter country code. Works with operators =, !=, in, notin (e.g. geolocation not in ['XV','XZ'])\n id: int volume offer unique ID\n inet_down: float internet download speed in Mb/s\n inet_up: float internet upload speed in Mb/s\n reliability: float machine reliability score (see FAQ for explanation)\n storage_cost: float storage cost in $/GB/month\n verified: bool is the machine verified\n \"\"\"),\n)\ndef search__network_volumes(args):\n \"\"\"Search for network volume offers using custom query.\"\"\"\n from vastai.api.query import parse_query, vol_offers_fields, offers_alias, offers_mult\n\n try:\n if args.no_default:\n query = {}\n else:\n query = {\"verified\": {\"eq\": True}, \"external\": {\"eq\": False}, \"disk_space\": {\"gte\": 1}}\n\n if args.query is not None:\n query = parse_query(args.query, query, vol_offers_fields, {}, offers_mult)\n\n order = []\n for name in args.order.split(\",\"):\n name = name.strip()\n if not name:\n continue\n direction = \"asc\"\n field = name\n if name.strip(\"-\") != name:\n direction = \"desc\"\n field = name.strip(\"-\")\n if name.strip(\"+\") != name:\n direction = \"asc\"\n field = name.strip(\"+\")\n if field in offers_alias:\n field = offers_alias[field]\n order.append([field, direction])\n except ValueError as e:\n print(\"Error: \", e)\n return 1\n\n client = get_client(args)\n\n if args.explain:\n print(\"request json: \")\n print(query)\n\n rows = offers_api.search_network_volumes(\n client, query=query, order=order,\n limit=args.limit, storage=args.storage,\n no_default=True, # defaults already applied above\n )\n\n if args.raw:\n return rows\n else:\n display_table(rows, nw_vol_displayable_fields)\n\n\n# ---------------------------------------------------------------------------\n# create / delete / clone volume\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of volume offer\", type=int),\n argument(\"-s\", \"--size\",\n help=\"size in GB of volume. Default %(default)s GB.\", default=15, type=float),\n argument(\"-n\", \"--name\", help=\"Optional name of volume.\", type=str),\n usage=\"vastai create volume ID [options]\",\n help=\"Create a new volume\",\n epilog=deindent(\"\"\"\n Creates a volume from an offer ID (which is returned from \"search volumes\"). Each offer ID can be used to create multiple volumes,\n provided the size of all volumes does not exceed the size of the offer.\n \"\"\"),\n)\ndef create__volume(args):\n \"\"\"Create a new volume from an offer ID.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"size\": int(args.size), \"id\": int(args.id)})\n\n client = get_client(args)\n result = storage_api.create_volume(client, id=args.id, size=args.size, name=args.name)\n if args.raw:\n return result\n else:\n print(\"Created. {}\".format(result))\n\n\[email protected](\n argument(\"id\", help=\"id of network volume offer\", type=int),\n argument(\"-s\", \"--size\",\n help=\"size in GB of network volume. Default %(default)s GB.\", default=15, type=float),\n argument(\"-n\", \"--name\", help=\"Optional name of network volume.\", type=str),\n usage=\"vastai create network volume ID [options]\",\n help=\"Create a new network volume\",\n epilog=deindent(\"\"\"\n 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,\n provided the size of all volumes does not exceed the size of the offer.\n \"\"\"),\n)\ndef create__network_volume(args):\n \"\"\"Create a new network volume from an offer ID.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"size\": int(args.size), \"id\": int(args.id)})\n\n client = get_client(args)\n result = storage_api.create_network_volume(client, id=args.id, size=args.size, name=args.name)\n if args.raw:\n return result\n else:\n print(\"Created. {}\".format(result))\n\n\[email protected](\n argument(\"id\", help=\"id of volume contract\", type=int),\n usage=\"vastai delete volume ID\",\n help=\"Delete a volume\",\n epilog=deindent(\"\"\"\n Deletes volume with the given ID. All instances using the volume must be destroyed before the volume can be deleted.\n \"\"\"),\n)\ndef delete__volume(args):\n \"\"\"Delete a volume.\"\"\"\n client = get_client(args)\n result = storage_api.delete_volume(client, id=args.id)\n if args.raw:\n return result\n else:\n print(\"Deleted. {}\".format(result))\n\n\[email protected](\n argument(\"source\", help=\"id of volume contract being cloned\", type=int),\n argument(\"dest\", help=\"id of volume offer volume is being copied to\", type=int),\n 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),\n argument(\"-d\", \"--disable_compression\", action=\"store_true\", help=\"Do not compress volume data before copying.\"),\n usage=\"vastai clone volume \u003csource_id> \u003cdest_id> [options]\",\n help=\"Clone an existing volume\",\n epilog=deindent(\"\"\"\n Create a new volume with the given offer, by copying the existing volume.\n Size defaults to the size of the existing volume, but can be increased if there is available space.\n \"\"\"),\n)\ndef clone__volume(args):\n \"\"\"Clone an existing volume.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"src_id\": args.source, \"dst_id\": args.dest})\n\n client = get_client(args)\n result = storage_api.clone_volume(client, source=args.source, dest=args.dest,\n size=args.size, disable_compression=args.disable_compression)\n if args.raw:\n return result\n else:\n print(\"Created. {}\".format(result))\n\n\n# ---------------------------------------------------------------------------\n# show volumes\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"-t\", \"--type\", help=\"volume type to display. Default to all. Possible values are \\\"local\\\", \\\"all\\\", \\\"network\\\"\", type=str, default=\"all\"),\n usage=\"vastai show volumes [OPTIONS]\",\n help=\"Show stats on owned volumes.\",\n epilog=deindent(\"\"\"\n Show stats on owned volumes\n \"\"\"),\n)\ndef show__volumes(args):\n \"\"\"Show stats on owned volumes.\"\"\"\n client = get_client(args)\n processed = storage_api.show_volumes(client, type=args.type)\n for row in processed:\n for k, v in row.items():\n row[k] = strip_strings(v)\n if args.raw:\n return processed\n else:\n display_table(processed, volume_fields, replace_spaces=False)\n\n\n# ---------------------------------------------------------------------------\n# show instance-backups / show instance-backup-files\n# ---------------------------------------------------------------------------\n\ndef _authorized_url(url, token):\n \"\"\"Fold the download token into `url` as a query parameter.\n\n B2 takes the token either in the Authorization header or as an `Authorization`\n query parameter, and the second form is what makes a URL usable on its own -- paste\n it into a browser, or curl it without -H. Returns the url untouched when either\n piece is missing, so a malformed payload still lists its files.\n \"\"\"\n if not url or not token:\n return url\n sep = \"&\" if urlparse(url).query else \"?\"\n return \"{}{}Authorization={}\".format(url, sep, quote(token, safe=\"\"))\n\n\[email protected](\n argument(\"--contract-id\", help=\"only show backups of this instance or volume contract id\", type=int),\n usage=\"vastai show instance-backups [OPTIONS]\",\n help=\"Show backups Vast has taken of your instances.\",\n epilog=deindent(\"\"\"\n Lists the backups Vast has taken of your instances and volumes, newest first. These\n are created by Vast (for example before a machine is decommissioned), not by you.\n\n Status is pending while the copy runs, then success or error. A backup outlives the\n instance it came from, so entries stay listed after the instance is gone.\n\n Restore one into a running instance with `vastai copy`, naming vast.CONTRACT_ID as\n the source. An instance's own filesystem sits under upper/:\n\n vastai show instance-backups\n vastai copy vast.12345678:/upper/workspace 99887766:/workspace\n\n A backed-up volume has no upper/ -- its contents sit at the root of its own\n contract id:\n\n vastai copy vast.12345679:/ 99887766:/volumes/data\n\n To pull the objects down directly instead of restoring into an instance, see\n `vastai show instance-backup-files CONTRACT_ID`.\n\n Example: vastai show instance-backups --contract-id 12345678\n \"\"\"),\n)\ndef show__instance_backups(args):\n \"\"\"Show backups Vast has taken of your instances.\"\"\"\n client = get_client(args)\n rows = storage_api.show_instance_backups(client, contract_id=args.contract_id)\n if args.raw:\n return rows\n now = time.time()\n for row in rows:\n row[\"age_hours\"] = (now - (row.get(\"created_at\") or now)) / 3600.0\n # A backup whose contract has been archived out of the live tables reports no type.\n row[\"type\"] = row.get(\"type\") or \"unknown\"\n row[\"info\"] = strip_strings(row.get(\"info\") or \"\")\n display_table(rows, instance_backup_fields, replace_spaces=False)\n\n\[email protected](\n argument(\"contract_id\", help=\"instance or volume contract id whose backup to list\", type=int),\n usage=\"vastai show instance-backup-files CONTRACT_ID [OPTIONS]\",\n help=\"List the stored objects of one backup, with direct download URLs.\",\n epilog=deindent(\"\"\"\n Lists the objects stored for one backup. Each URL already carries a download token\n scoped to this backup, so it works on its own -- open it in a browser, or:\n\n curl -o model.safetensors \"\u003curl>\"\n\n The token is also printed on its own for the header form, `curl -H \"Authorization:\n \u003ctoken>\"`, which keeps it out of shell history and proxy logs. Either way it expires\n (the lifetime is printed), so re-run this command rather than saving the URLs.\n\n If the listing reports itself truncated, the backup holds more objects than are\n shown; restore with `vastai copy vast.CONTRACT_ID:/ INSTANCE_ID:/path` instead of\n fetching file by file.\n\n Example: vastai show instance-backup-files 12345678\n \"\"\"),\n)\ndef show__instance_backup_files(args):\n \"\"\"List the stored objects of one backup, with direct download URLs.\"\"\"\n client = get_client(args)\n blob = storage_api.instance_backup_files(client, args.contract_id)\n if args.raw:\n return blob\n files = blob.get(\"files\") or []\n # --raw stays a passthrough of the API payload; only the table gets ready-to-click URLs.\n token = blob.get(\"authorization_token\")\n for row in files:\n row[\"size_mb\"] = (row.get(\"size\") or 0) / 1e6\n row[\"url\"] = _authorized_url(row.get(\"url\"), token)\n print(\"authorization token: {} (already embedded in the URLs below)\".format(token))\n print(\"expires in: {} seconds\".format(blob.get(\"expires_in_sec\")))\n if blob.get(\"truncated\"):\n print(\"WARNING: listing truncated -- this backup holds more objects than shown below.\")\n if not files:\n print(\"No objects stored for this backup yet.\")\n return\n display_table(files, instance_backup_file_fields, replace_spaces=False)\n\n\n# ---------------------------------------------------------------------------\n# list / unlist volumes (host)\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"id\", help=\"id of machine to list\", type=int),\n argument(\"-p\", \"--price_disk\",\n help=\"storage price in $/GB/month, default: $%(default).2f/GB/month\", default=.10, type=float),\n 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),\n argument(\"-s\", \"--size\", help=\"size of disk space allocated to offer in GB, default %(default)s GB\", default=15),\n usage=\"vastai list volume ID [options]\",\n help=\"[Host] list disk space for rent as a volume on a machine\",\n epilog=deindent(\"\"\"\n Allocates a section of disk on a machine to be used for volumes.\n \"\"\"),\n)\ndef list__volume(args):\n \"\"\"List disk space for rent as a volume on a machine.\"\"\"\n from vastai.cli.util import string_to_unix_epoch\n\n end_date = string_to_unix_epoch(args.end_date) if args.end_date else None\n\n if args.explain:\n print(\"request json: \")\n print({\"size\": int(args.size), \"machine\": int(args.id), \"price_disk\": float(args.price_disk)})\n\n client = get_client(args)\n result = storage_api.list_volume(client, id=args.id, size=int(args.size),\n price_disk=float(args.price_disk), end_date=end_date)\n if args.raw:\n return result\n else:\n print(\"Created. {}\".format(result))\n\n\[email protected](\n argument(\"ids\", help=\"id of machines list\", type=int, nargs='+'),\n argument(\"-p\", \"--price_disk\",\n help=\"storage price in $/GB/month, default: $%(default).2f/GB/month\", default=.10, type=float),\n 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),\n argument(\"-s\", \"--size\", help=\"size of disk space allocated to offer in GB, default %(default)s GB\", default=15),\n usage=\"vastai list volume IDs [options]\",\n help=\"[Host] list disk space for rent as a volume on machines\",\n epilog=deindent(\"\"\"\n Allocates a section of disk on machines to be used for volumes.\n \"\"\"),\n)\ndef list__volumes(args):\n \"\"\"List disk space for rent as a volume on multiple machines.\"\"\"\n from vastai.cli.util import string_to_unix_epoch\n\n end_date = string_to_unix_epoch(args.end_date) if args.end_date else None\n\n if args.explain:\n print(\"request json: \")\n print({\"size\": int(args.size), \"machine\": [int(id) for id in args.ids], \"price_disk\": float(args.price_disk)})\n\n client = get_client(args)\n result = storage_api.list_volumes(client, ids=args.ids, size=int(args.size),\n price_disk=float(args.price_disk), end_date=end_date)\n if args.raw:\n return result\n else:\n print(\"Created. {}\".format(result))\n\n\[email protected](\n argument(\"disk_id\", help=\"id of network disk to list\", type=int),\n argument(\"-p\", \"--price_disk\", help=\"storage price in $/GB/month, default: $%(default).2f/GB/month\", default=.15, type=float),\n 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),\n argument(\"-s\", \"--size\", help=\"size of disk space allocated to offer in GB, default %(default)s GB\", default=15, type=int),\n usage=\"vastai list network volume DISK_ID [options]\",\n help=\"[Host] list disk space for rent as a network volume\",\n)\ndef list__network_volume(args):\n \"\"\"List disk space for rent as a network volume.\"\"\"\n from vastai.cli.util import string_to_unix_epoch\n\n end_date = string_to_unix_epoch(args.end_date) if args.end_date else None\n\n if args.explain:\n print(\"request json: \")\n print({\"disk_id\": args.disk_id, \"price_disk\": args.price_disk, \"size\": args.size})\n\n client = get_client(args)\n result = storage_api.list_network_volume(client, disk_id=args.disk_id, price_disk=args.price_disk,\n size=args.size, end_date=end_date)\n if args.raw:\n return result\n print(result[\"msg\"])\n\n\[email protected](\n argument(\"id\", help=\"volume ID you want to unlist\", type=int),\n usage=\"vastai unlist volume ID\",\n help=\"[Host] unlist volume offer\",\n)\ndef unlist__volume(args):\n \"\"\"Unlist a volume offer.\"\"\"\n if args.explain:\n print(\"request json:\", {\"id\": args.id})\n\n client = get_client(args)\n result = storage_api.unlist_volume(client, id=args.id)\n if args.raw:\n return result\n else:\n print(result[\"msg\"])\n\n\[email protected](\n argument(\"id\", help=\"id of network volume offer to unlist\", type=int),\n usage=\"vastai unlist network volume OFFER_ID\",\n help=\"[Host] Unlists network volume offer\",\n)\ndef unlist__network_volume(args):\n \"\"\"Unlist a network volume offer.\"\"\"\n if args.explain:\n print(\"request json: \")\n print({\"id\": args.id})\n\n client = get_client(args)\n result = storage_api.unlist_network_volume(client, id=args.id)\n if args.raw:\n return result\n print(result[\"msg\"])\n","content_type":"text/x-python; charset=utf-8","language":"python","size":39875,"content_sha256":"47f3008e07ca8a5c5f6953951e9ef2cabf76bdf9964d992d42a2092d18d3c611"},{"filename":"cli/commands/teams.py","content":"\"\"\"CLI commands for managing teams.\"\"\"\n\nimport json\n\nfrom vastai.cli.parser import argument\nfrom vastai.cli.display import deindent\nfrom vastai.api import teams as teams_api\n\n\nfrom vastai.cli.utils import get_parser as _get_parser, get_client # noqa: F401\n\n\nparser = _get_parser()\n\n\n# ---------------------------------------------------------------------------\n# create / destroy team\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"--team-name\", help=\"name of the team\", type=str),\n argument(\"--transfer-credit\", help=\"amount of personal credit to transfer to the new team\", type=float, default=0),\n usage=\"vastai create-team --team-name TEAM_NAME [--transfer-credit AMOUNT]\",\n help=\"Create a new team\",\n epilog=deindent(\"\"\"\n Creates a new team under your account.\n\n Unlike legacy teams, this command does NOT convert your personal account into a team.\n Each team is created as a separate account, and you can be a member of multiple teams.\n\n When you create a team:\n - You become the team owner.\n - The team starts as an independent account with its own billing, credits, and resources.\n - Default roles (owner, manager, member) are automatically created.\n - You can invite others, assign roles, and manage resources within the team.\n\n Optional:\n You can transfer a portion of your existing personal credits to the team by using\n the `--transfer-credit` flag. Example:\n vastai create-team --team-name myteam --transfer-credit 25\n\n Notes:\n - You cannot create a team from within another team account.\n - Your API key keeps acting on your personal account. To work in the team,\n switch to it in the console and create an API key there.\n\n For more details, see:\n https://vast.ai/docs/teams-quickstart\n \"\"\"),\n)\ndef create__team(args):\n \"\"\"Create a new team.\"\"\"\n client = get_client(args)\n result = teams_api.create_team(client, team_name=args.team_name, transfer_credit=args.transfer_credit)\n print(result)\n\n\[email protected](\n usage=\"vastai destroy team\",\n help=\"Destroy your team\",\n)\ndef destroy__team(args):\n \"\"\"Destroy your team.\"\"\"\n client = get_client(args)\n result = teams_api.destroy_team(client)\n print(result)\n\n\n# ---------------------------------------------------------------------------\n# team roles\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"--name\", help=\"name of the role\", type=str),\n argument(\"--permissions\", help=\"file path for json encoded permissions, look in the docs for more information\", type=str),\n usage=\"vastai create team-role --name NAME --permissions PERMISSIONS\",\n help=\"Add a new role to your team\",\n epilog=deindent(\"\"\"\n Creating a new team role involves understanding how permissions must be sent via json format.\n You can find more information about permissions here: https://vast.ai/docs/cli/roles-and-permissions\n \"\"\")\n)\ndef create__team_role(args):\n \"\"\"Create a new team role.\"\"\"\n from vastai.cli.util import load_permissions_from_file\n client = get_client(args)\n permissions = load_permissions_from_file(args.permissions)\n result = teams_api.create_team_role(client, name=args.name, permissions=permissions)\n print(result)\n\n\[email protected](\n argument(\"NAME\", help=\"name of the role\", type=str),\n usage=\"vastai show team-role NAME\",\n help=\"Show your team role\",\n)\ndef show__team_role(args):\n \"\"\"Show a specific team role.\"\"\"\n client = get_client(args)\n result = teams_api.show_team_role(client, name=args.NAME)\n print(json.dumps(result, indent=1, sort_keys=True))\n\n\[email protected](\n usage=\"vastai show team-roles\",\n help=\"Show roles for a team\",\n)\ndef show__team_roles(args):\n \"\"\"Show all team roles.\"\"\"\n client = get_client(args)\n result = teams_api.show_team_roles(client)\n if args.raw:\n return result\n else:\n print(result)\n\n\[email protected](\n argument(\"id\", help=\"id of the role\", type=int),\n argument(\"--name\", help=\"name of the role\", type=str),\n argument(\"--permissions\", help=\"file path for json encoded permissions, look in the docs for more information\", type=str),\n usage=\"vastai update team-role ID --name NAME --permissions PERMISSIONS\",\n help=\"Update an existing team role\",\n)\ndef update__team_role(args):\n \"\"\"Update an existing team role.\"\"\"\n from vastai.cli.util import load_permissions_from_file\n client = get_client(args)\n permissions = load_permissions_from_file(args.permissions)\n result = teams_api.update_team_role(client, id=args.id, name=args.name, permissions=permissions)\n if args.raw:\n return result\n else:\n print(json.dumps(result, indent=1))\n\n\[email protected](\n argument(\"NAME\", help=\"name of the role\", type=str),\n usage=\"vastai remove team-role NAME\",\n help=\"Remove a role from your team\",\n)\ndef remove__team_role(args):\n \"\"\"Remove a team role.\"\"\"\n client = get_client(args)\n result = teams_api.remove_team_role(client, name=args.NAME)\n print(result)\n\n\n# ---------------------------------------------------------------------------\n# team members\n# ---------------------------------------------------------------------------\n\[email protected](\n argument(\"--email\", help=\"email of user to be invited\", type=str),\n argument(\"--role\", help=\"role of user to be invited\", type=str),\n usage=\"vastai invite member --email EMAIL --role ROLE\",\n help=\"Invite a team member\",\n)\ndef invite__member(args):\n \"\"\"Invite a team member.\"\"\"\n client = get_client(args)\n result = teams_api.invite_member(client, email=args.email, role=args.role)\n if result:\n print(f\"successfully invited {args.email} to your current team\")\n\n\[email protected](\n usage=\"vastai show members\",\n help=\"Show your team members\",\n)\ndef show__members(args):\n \"\"\"Show team members.\"\"\"\n client = get_client(args)\n result = teams_api.show_members(client)\n if args.raw:\n return result\n else:\n print(result)\n\n\[email protected](\n argument(\"id\", help=\"id of user to remove\", type=int),\n usage=\"vastai remove member ID\",\n help=\"Remove a team member\",\n)\ndef remove__member(args):\n \"\"\"Remove a team member.\"\"\"\n client = get_client(args)\n result = teams_api.remove_member(client, id=args.id)\n print(result)\n\n","content_type":"text/x-python; charset=utf-8","language":"python","size":6505,"content_sha256":"07b7e4f4dd9f8a316e7fcec932435ff4456254a5d673df34f7fdaf159026d2ec"},{"filename":"cli/commands/uninstall.py","content":"\"\"\"`vastai uninstall` — remove a managed (curl | bash installer) install.\n\npip installs aren't managed by the installer; for those this command only\nprints the correct pip uninstall instructions and exits non-zero — it never\nshells out to pip, matching `update`'s same-method discipline.\n\"\"\"\n\nimport sys\n\nfrom vastai.cli.parser import argument\nfrom vastai.cli.display import deindent\nfrom vastai.cli.util import DIRS, VERSION\nfrom vastai.cli.selfupdate import (\n INSTALL_SH_HINT, UninstallError, is_managed_install, install_root, perform_uninstall,\n)\nfrom vastai.cli.utils import get_parser as _get_parser\n\nparser = _get_parser()\n\nPIP_UNINSTALL_HINT = \"pip uninstall vastai\"\n\n\[email protected](\n argument(\"-y\", \"--yes\", action=\"store_true\", help=\"skip the confirmation prompt\"),\n usage=\"vastai uninstall [-y]\",\n help=\"Remove the managed CLI install\",\n hidden=True, # still under testing — remove once ready to announce\n epilog=deindent(f\"\"\"\n Removes a CLI installed with the managed installer\n ({INSTALL_SH_HINT}): the install root and its symlinks in\n ~/.local/bin. Config in {DIRS['config']} (your API key) is left\n untouched, so re-running the installer keeps you logged in.\n For pip installs, uninstall with: {PIP_UNINSTALL_HINT}\n \"\"\"),\n)\ndef uninstall(args):\n if not is_managed_install():\n print(\n \"This CLI was not installed with the managed installer, so this \"\n f\"command can't remove it.\\n Installed via pip? Run: {PIP_UNINSTALL_HINT}\",\n file=sys.stderr,\n )\n return 1\n\n root = install_root()\n if not args.yes:\n if not sys.stdin.isatty():\n print(\n \"Refusing to uninstall without confirmation in a non-interactive \"\n \"shell. Re-run with --yes to skip the prompt.\",\n file=sys.stderr,\n )\n return 1\n try:\n answer = input(f\"Remove vastai {VERSION} from {root}? [y/N] \").strip().lower()\n except EOFError:\n answer = \"\"\n if answer not in (\"y\", \"yes\"):\n print(\"Aborted.\")\n return 1\n\n try:\n removed = perform_uninstall()\n except UninstallError as e:\n print(str(e), file=sys.stderr)\n return 1\n\n print(f\"vastai {VERSION} uninstalled from {removed}.\")\n print(f\"Config in {DIRS['config']} was left untouched.\")\n return 0\n","content_type":"text/x-python; charset=utf-8","language":"python","size":2427,"content_sha256":"be68563de3ec532a431704fc22da138b308d2921e5c55b2ab8359be8788c60b9"},{"filename":"cli/commands/update.py","content":"\"\"\"`vastai update` — self-update for managed (curl | bash installer) installs.\n\npip installs aren't managed by the installer; for those this command only\nprints the correct pip upgrade instructions and exits non-zero — it never\nshells out to pip, since the CLI cannot know which pip/venv owns its\nenvironment.\n\"\"\"\n\nimport sys\n\nfrom vastai.cli.parser import argument\nfrom vastai.cli.display import deindent\nfrom vastai.cli.util import VERSION\nfrom vastai.cli.selfupdate import (\n INSTALL_SH_HINT, PIP_UPGRADE_HINT, UpdateError,\n fetch_manifest, is_managed_install, is_newer, perform_update,\n)\nfrom vastai.cli.utils import get_parser as _get_parser\n\nparser = _get_parser()\n\nEXIT_STALE = 10 # `update --check` exit code when a newer version exists\n\n\[email protected](\n argument(\"--check\", action=\"store_true\", help=\"only check for a newer version; exit 0 if current, 10 if an update is available\"),\n argument(\"--version\", dest=\"target_version\", metavar=\"VERSION\", help=\"install a specific version instead of the latest (also how you pin or roll back)\"),\n usage=\"vastai update [--check | --version VERSION]\",\n help=\"Update the CLI to the latest version\",\n epilog=deindent(\"\"\"\n Updates a CLI installed with the managed installer\n (curl -fsSL https://vast.ai/install.sh | bash). The new version is\n installed into a fresh environment and verified before it is swapped\n in, so an interrupted update never leaves a broken CLI. To roll back a\n bad release, pin the previous one: vastai update --version 1.2.3\n For pip installs, update with: pip install --upgrade vastai\n \"\"\"),\n)\ndef update(args):\n try:\n if args.check:\n latest = fetch_manifest()[\"latest\"]\n if is_newer(latest, VERSION):\n print(f\"Update available: {latest} (you have {VERSION})\")\n return EXIT_STALE\n print(f\"vastai {VERSION} is up to date (latest: {latest})\")\n return 0\n\n if not is_managed_install():\n print(\n \"This CLI was not installed with the managed installer, so it \"\n f\"cannot self-update.\\n Installed via pip? Run: {PIP_UPGRADE_HINT}\\n\"\n f\" Or switch to the managed install: {INSTALL_SH_HINT}\",\n file=sys.stderr,\n )\n return 1\n\n manifest = fetch_manifest()\n target = args.target_version or manifest[\"latest\"]\n current = VERSION\n if target == current:\n latest = manifest[\"latest\"]\n if args.target_version:\n print(f\"vastai {current} is already installed (latest available: {latest}).\")\n else:\n print(f\"vastai {current} is up to date (latest: {latest}).\")\n return 0\n print(f\"Updating vastai {current} -> {target} ...\")\n perform_update(target, manifest)\n print(f\"Done. vastai is now {target}.\")\n return 0\n\n except UpdateError as e:\n print(str(e), file=sys.stderr)\n if not args.check:\n print(\"The current install was left untouched.\", file=sys.stderr)\n return 1\n","content_type":"text/x-python; charset=utf-8","language":"python","size":3142,"content_sha256":"6442a47886483f1e7aca55929ced5ac576321eb64db0364093290daaeb63853e"},{"filename":"cli/display.py","content":"\"\"\"Display formatting and field definitions for the Vast.ai CLI.\"\"\"\n\nimport re\nimport shutil\nimport subprocess\nfrom datetime import datetime, timezone\nfrom typing import Dict, Tuple\n\n\ndef strip_strings(value):\n if isinstance(value, str):\n return value.strip()\n elif isinstance(value, dict):\n return {k: strip_strings(v) for k, v in value.items()}\n elif isinstance(value, list):\n return [strip_strings(item) for item in value]\n return value # Return as is if not a string, list, or dict\n\ndef translate_null_strings_to_blanks(d: Dict) -> Dict:\n \"\"\"Map over a dict and translate any null string values into ' '.\n Leave everything else as is. This is needed because you cannot add TableCell\n objects with only a null string or the client crashes.\n\n :param Dict d: dict of item values.\n :rtype Dict:\n \"\"\"\n\n # Beware: locally defined function.\n def translate_nulls(s):\n if s == \"\":\n return \" \"\n return s\n\n new_d = {k: translate_nulls(v) for k, v in d.items()}\n return new_d\n\n\ndef deindent(message: str, add_separator: bool = True) -> str:\n \"\"\"\n Deindent a quoted string. Scans message and finds the smallest number of whitespace characters in any line and\n removes that many from the start of every line.\n\n :param str message: Message to deindent.\n :rtype str:\n \"\"\"\n message = re.sub(r\" *$\", \"\", message, flags=re.MULTILINE)\n indents = [len(x) for x in re.findall(\"^ *(?=[^ ])\", message, re.MULTILINE) if len(x)]\n a = min(indents)\n message = re.sub(r\"^ {,\" + str(a) + \"}\", \"\", message, flags=re.MULTILINE)\n if add_separator:\n # For help epilogs - cleanly separating extra help from options\n line_width = min(150, shutil.get_terminal_size((80, 20)).columns)\n message = \"_\"*line_width + \"\\n\"*2 + message.strip() + \"\\n\" + \"_\"*line_width\n return message.strip()\n\n\n# These are the fields that are displayed when a search is run\ndisplayable_fields = (\n # (\"bw_nvlink\", \"Bandwidth NVLink\", \"{}\", None, True),\n (\"id\", \"ID\", \"{}\", None, True),\n (\"cuda_max_good\", \"CUDA\", \"{:0.1f}\", None, True),\n (\"num_gpus\", \"N\", \"{}x\", None, False),\n (\"gpu_name\", \"Model\", \"{}\", None, True),\n (\"cpu_arch\", \"arch\", \"{}\", None, True),\n (\"pcie_bw\", \"PCIE\", \"{:0.1f}\", None, True),\n (\"cpu_ghz\", \"cpu_ghz\", \"{:0.1f}\", None, True),\n (\"cpu_cores_effective\", \"vCPUs\", \"{:0.1f}\", None, True),\n (\"cpu_ram\", \"RAM\", \"{:0.1f}\", lambda x: x / 1000, False),\n (\"gpu_ram\", \"VRAM\", \"{:0.1f}\", lambda x: x / 1000, False),\n (\"disk_space\", \"Disk\", \"{:.0f}\", None, True),\n (\"dph_total\", \"$/hr\", \"{:0.4f}\", None, True),\n (\"dlperf\", \"DLP\", \"{:0.1f}\", None, True),\n (\"dlperf_per_dphtotal\", \"DLP/$\", \"{:0.2f}\", None, True),\n (\"score\", \"score\", \"{:0.1f}\", None, True),\n (\"driver_version\", \"NV Driver\", \"{}\", None, True),\n (\"inet_up\", \"Net_up\", \"{:0.1f}\", None, True),\n (\"inet_down\", \"Net_down\", \"{:0.1f}\", None, True),\n (\"reliability\", \"R\", \"{:0.1f}\", lambda x: x * 100, True),\n (\"duration\", \"Max_Days\", \"{:0.1f}\", lambda x: x / (24.0 * 60.0 * 60.0), True),\n (\"machine_id\", \"mach_id\", \"{}\", None, True),\n (\"verification\", \"status\", \"{}\", None, True),\n (\"host_id\", \"host_id\", \"{}\", None, True),\n (\"direct_port_count\", \"ports\", \"{}\", None, True),\n (\"geolocation\", \"country\", \"{}\", None, True),\n # (\"direct_port_count\", \"Direct Port Count\", \"{}\", None, True),\n)\n\ndisplayable_fields_reserved = (\n # (\"bw_nvlink\", \"Bandwidth NVLink\", \"{}\", None, True),\n (\"id\", \"ID\", \"{}\", None, True),\n (\"cuda_max_good\", \"CUDA\", \"{:0.1f}\", None, True),\n (\"num_gpus\", \"N\", \"{}x\", None, False),\n (\"gpu_name\", \"Model\", \"{}\", None, True),\n (\"cpu_arch\", \"arch\", \"{}\", None, True),\n (\"pcie_bw\", \"PCIE\", \"{:0.1f}\", None, True),\n (\"cpu_ghz\", \"cpu_ghz\", \"{:0.1f}\", None, True),\n (\"cpu_cores_effective\", \"vCPUs\", \"{:0.1f}\", None, True),\n (\"cpu_ram\", \"RAM\", \"{:0.1f}\", lambda x: x / 1000, False),\n (\"disk_space\", \"Disk\", \"{:.0f}\", None, True),\n (\"discounted_dph_total\", \"$/hr\", \"{:0.4f}\", None, True),\n (\"dlperf\", \"DLP\", \"{:0.1f}\", None, True),\n (\"dlperf_per_dphtotal\", \"DLP/$\", \"{:0.2f}\", None, True),\n (\"driver_version\", \"NV Driver\", \"{}\", None, True),\n (\"inet_up\", \"Net_up\", \"{:0.1f}\", None, True),\n (\"inet_down\", \"Net_down\", \"{:0.1f}\", None, True),\n (\"reliability\", \"R\", \"{:0.1f}\", lambda x: x * 100, True),\n (\"duration\", \"Max_Days\", \"{:0.1f}\", lambda x: x / (24.0 * 60.0 * 60.0), True),\n (\"machine_id\", \"mach_id\", \"{}\", None, True),\n (\"verification\", \"status\", \"{}\", None, True),\n (\"host_id\", \"host_id\", \"{}\", None, True),\n (\"direct_port_count\", \"ports\", \"{}\", None, True),\n (\"geolocation\", \"country\", \"{}\", None, True),\n # (\"direct_port_count\", \"Direct Port Count\", \"{}\", None, True),\n)\n\n\nvol_displayable_fields = (\n (\"id\", \"ID\", \"{}\", None, True),\n (\"cuda_max_good\", \"CUDA\", \"{:0.1f}\", None, True),\n (\"cpu_ghz\", \"cpu_ghz\", \"{:0.1f}\", None, True),\n (\"disk_bw\", \"Disk B/W\", \"{:0.1f}\", None, True),\n (\"disk_space\", \"Disk\", \"{:.0f}\", None, True),\n (\"disk_name\", \"Disk Name\", \"{}\", None, True),\n (\"storage_cost\", \"$/Gb/Month\", \"{:.2f}\", None, True),\n (\"driver_version\", \"NV Driver\", \"{}\", None, True),\n (\"inet_up\", \"Net_up\", \"{:0.1f}\", None, True),\n (\"inet_down\", \"Net_down\", \"{:0.1f}\", None, True),\n (\"reliability\", \"R\", \"{:0.1f}\", lambda x: x * 100, True),\n (\"duration\", \"Max_Days\", \"{:0.1f}\", lambda x: x / (24.0 * 60.0 * 60.0), True),\n (\"machine_id\", \"mach_id\", \"{}\", None, True),\n (\"verification\", \"status\", \"{}\", None, True),\n (\"host_id\", \"host_id\", \"{}\", None, True),\n (\"geolocation\", \"country\", \"{}\", None, True),\n)\n\nnw_vol_displayable_fields = (\n (\"id\", \"ID\", \"{}\", None, True),\n (\"disk_space\", \"Disk\", \"{:.0f}\", None, True),\n (\"storage_cost\", \"$/Gb/Month\", \"{:.2f}\", None, True),\n (\"inet_up\", \"Net_up\", \"{:0.1f}\", None, True),\n (\"inet_down\", \"Net_down\", \"{:0.1f}\", None, True),\n (\"reliability\", \"R\", \"{:0.1f}\", lambda x: x * 100, True),\n (\"duration\", \"Max_Days\", \"{:0.1f}\", lambda x: x / (24.0 * 60.0 * 60.0), True),\n (\"verification\", \"status\", \"{}\", None, True),\n (\"host_id\", \"host_id\", \"{}\", None, True),\n (\"cluster_id\", \"cluster_id\", \"{}\", None, True),\n (\"geolocation\", \"country\", \"{}\", None, True),\n (\"nw_disk_min_bw\", \"Min BW MiB/s\", \"{}\", None, True),\n (\"nw_disk_max_bw\", \"Max BW MiB/s\", \"{}\", None, True),\n (\"nw_disk_avg_bw\", \"Avg BW MiB/s\", \"{}\", None, True),\n\n)\n\n# These fields are displayed for 'search benchmarks' (benchmarks table, AUTO-1418 schema)\nbenchmarks_displayable_fields = (\n (\"id\", \"ID\", \"{}\", None, True),\n (\"machine_id\", \"mach_id\", \"{}\", None, True),\n (\"gpu_name\", \"Model\", \"{}\", None, True),\n (\"num_gpus\", \"N\", \"{}x\", None, False),\n (\"value\", \"Value\", \"{:0.2f}\", None, True),\n (\"dph_base\", \"$/hr\", \"{:0.4f}\", None, True),\n (\"type\", \"Type\", \"{}\", None, True),\n (\"template_hash\", \"Template\", \"{}\", None, True),\n (\"template_id\", \"Templ_ID\", \"{}\", None, True),\n (\"contract_id\", \"contract\", \"{}\", None, True),\n (\"image\", \"Image\", \"{}\", None, True),\n (\"last_update\", \"Date\", \"{}\", lambda x: datetime.fromtimestamp(x, timezone.utc).strftime(\"%Y-%m-%d %H:%M\"), True),\n)\n\n# These fields are displayed when you do 'show instances'\ninstance_fields = (\n (\"id\", \"ID\", \"{}\", None, True),\n (\"machine_id\", \"Machine\", \"{}\", None, True),\n (\"actual_status\", \"Status\", \"{}\", None, True),\n (\"num_gpus\", \"Num\", \"{}x\", None, False),\n (\"gpu_name\", \"Model\", \"{}\", None, True),\n (\"gpu_util\", \"Util. %\", \"{:0.1f}\", None, True),\n (\"cpu_cores_effective\", \"vCPUs\", \"{:0.1f}\", None, True),\n (\"cpu_ram\", \"RAM\", \"{:0.1f}\", lambda x: x / 1000, False),\n (\"disk_space\", \"Storage\", \"{:.0f}\", None, True),\n (\"ssh_host\", \"SSH Addr\", \"{}\", None, True),\n (\"ssh_port\", \"SSH Port\", \"{}\", None, True),\n (\"dph_total\", \"$/hr\", \"{:0.4f}\", None, True),\n (\"image_uuid\", \"Image\", \"{}\", None, True),\n # (\"dlperf\", \"DLPerf\", \"{:0.1f}\", None, True),\n # (\"dlperf_per_dphtotal\", \"DLP/$\", \"{:0.1f}\", None, True),\n (\"inet_up\", \"Net up\", \"{:0.1f}\", None, True),\n (\"inet_down\", \"Net down\", \"{:0.1f}\", None, True),\n (\"reliability2\", \"R\", \"{:0.1f}\", lambda x: x * 100, True),\n (\"label\", \"Label\", \"{}\", None, True),\n (\"duration\", \"age(hours)\", \"{:0.2f}\", lambda x: x/(3600.0), True),\n (\"uptime_mins\", \"uptime(mins)\", \"{:0.2f}\", None, True),\n)\n\ncluster_fields = (\n (\"id\", \"ID\", \"{}\", None, True),\n (\"subnet\", \"Subnet\", \"{}\", None, True),\n (\"node_count\", \"Nodes\", \"{}\", None, True),\n (\"manager_id\", \"Manager ID\", \"{}\", None, True),\n (\"manager_ip\", \"Manager IP\", \"{}\", None, True),\n (\"machine_ids\", \"Machine ID's\", \"{}\", None, True)\n)\n\nnetwork_disk_fields = (\n (\"network_disk_id\", \"Network Disk ID\", \"{}\", None, True),\n (\"free_space\", \"Free Space (GB)\", \"{}\", None, True),\n (\"total_space\", \"Total Space (GB)\", \"{}\", None, True),\n)\n\nnetwork_disk_machine_fields = (\n (\"machine_id\", \"Machine ID\", \"{}\", None, True),\n (\"mount_point\", \"Mount Point\", \"{}\", None, True),\n)\n\noverlay_fields = (\n (\"overlay_id\", \"Overlay ID\", \"{}\", None, True),\n (\"name\", \"Name\", \"{}\", None, True),\n (\"subnet\", \"Subnet\", \"{}\", None, True),\n (\"cluster_id\", \"Cluster ID\", \"{}\", None, True),\n (\"instance_count\", \"Instances\", \"{}\", None, True),\n (\"instances\", \"Instance IDs\", \"{}\", None, True),\n)\ninstance_backup_fields = (\n (\"id\", \"ID\", \"{}\", None, True),\n (\"contract_id\", \"Contract\", \"{}\", None, True),\n (\"type\", \"Type\", \"{}\", None, True),\n (\"machine_id\", \"Machine\", \"{}\", None, True),\n (\"status\", \"Status\", \"{}\", None, True),\n (\"age_hours\", \"Age(hours)\", \"{:0.2f}\", None, True),\n (\"info\", \"Info\", \"{}\", None, True),\n)\n\ninstance_backup_file_fields = (\n (\"name\", \"Name\", \"{}\", None, True),\n (\"size_mb\", \"Size(MB)\", \"{:0.2f}\", None, False),\n (\"url\", \"URL\", \"{}\", None, True),\n)\n\nvolume_fields = (\n (\"id\", \"ID\", \"{}\", None, True),\n (\"cluster_id\", \"Cluster ID\", \"{}\", None, True),\n (\"label\", \"Name\", \"{}\", None, True),\n (\"disk_space\", \"Disk\", \"{:.0f}\", None, True),\n (\"status\", \"status\", \"{}\", None, True),\n (\"disk_name\", \"Disk Name\", \"{}\", None, True),\n (\"driver_version\", \"NV Driver\", \"{}\", None, True),\n (\"inet_up\", \"Net_up\", \"{:0.1f}\", None, True),\n (\"inet_down\", \"Net_down\", \"{:0.1f}\", None, True),\n (\"reliability2\", \"R\", \"{:0.1f}\", lambda x: x * 100, True),\n (\"duration\", \"age(hours)\", \"{:0.2f}\", lambda x: x/(3600.0), True),\n (\"machine_id\", \"mach_id\", \"{}\", None, True),\n (\"verification\", \"Verification\", \"{}\", None, True),\n (\"host_id\", \"host_id\", \"{}\", None, True),\n (\"geolocation\", \"country\", \"{}\", None, True),\n (\"instances\", \"instances\",\"{}\", None, True)\n)\n\n# These fields are displayed when you do 'show machines'\nmachine_fields = (\n (\"id\", \"ID\", \"{}\", None, True),\n (\"num_gpus\", \"#gpus\", \"{}\", None, True),\n (\"gpu_name\", \"gpu_name\", \"{}\", None, True),\n (\"disk_space\", \"disk\", \"{}\", None, True),\n (\"hostname\", \"hostname\", \"{}\", lambda x: x[:16], True),\n (\"driver_version\", \"driver\", \"{}\", None, True),\n (\"reliability2\", \"reliab\", \"{:0.4f}\", None, True),\n (\"verification\", \"veri\", \"{}\", None, True),\n (\"public_ipaddr\", \"ip\", \"{}\", None, True),\n (\"geolocation\", \"geoloc\", \"{}\", None, True),\n (\"num_reports\", \"reports\", \"{}\", None, True),\n (\"listed_gpu_cost\", \"gpuD_$/h\", \"{:0.2f}\", None, True),\n (\"min_bid_price\", \"gpuI$/h\", \"{:0.2f}\", None, True),\n (\"credit_discount_max\", \"rdisc\", \"{:0.2f}\", None, True),\n (\"listed_inet_up_cost\", \"netu_$/TB\", \"{:0.2f}\", lambda x: x * 1024, True),\n (\"listed_inet_down_cost\", \"netd_$/TB\", \"{:0.2f}\", lambda x: x * 1024, True),\n (\"gpu_occupancy\", \"occup\", \"{}\", None, True),\n)\n\n# These fields are displayed when you do 'show maints'\nmaintenance_fields = (\n (\"machine_id\", \"Machine ID\", \"{}\", None, True),\n (\"start_time\", \"Start (Date/Time in UTC)\", \"{}\", lambda x: datetime.fromtimestamp(x, tz=timezone.utc).strftime('%Y-%m-%d/%H:%M'), True),\n (\"end_time\", \"End (Date/Time in UTC)\", \"{}\", lambda x: datetime.fromtimestamp(x, tz=timezone.utc).strftime('%Y-%m-%d/%H:%M'), True),\n (\"duration_hours\", \"Duration (Hrs)\", \"{}\", None, True),\n (\"maintenance_category\", \"Category\", \"{}\", None, True),\n)\n\n\nipaddr_fields = (\n (\"ip\", \"ip\", \"{}\", None, True),\n (\"first_seen\", \"first_seen\", \"{}\", None, True),\n (\"first_location\", \"first_location\", \"{}\", None, True),\n)\n\naudit_log_fields = (\n (\"ip_address\", \"ip_address\", \"{}\", None, True),\n (\"api_key_id\", \"api_key_id\", \"{}\", None, True),\n (\"created_at\", \"created_at\", \"{}\", None, True),\n (\"api_route\", \"api_route\", \"{}\", None, True),\n (\"args\", \"args\", \"{}\", None, True),\n)\n\n\n# Backend writes per-second costs (web/price_increase_pending.py:_serialize_pending_row).\n# Display matches the frontend convention: GPU per hour, storage per month, bandwidth per\n# GB, platform fee as a percentage. Pre-format `old → new` strings in the command before\n# rendering so a null `new_*` collapses to \"-\" via the empty-row fallback in display_table.\npending_price_increase_fields = (\n (\"pending_id\", \"Pending ID\", \"{}\", None, True),\n (\"instance_id\", \"Instance\", \"{}\", None, True),\n (\"host_id\", \"Host\", \"{}\", None, True),\n (\"current_end\", \"Current End (UTC)\", \"{}\", None, True),\n (\"new_end\", \"New End (UTC)\", \"{}\", None, True),\n (\"gpu\", \"GPU ($/GPU/hr)\", \"{}\", None, True),\n (\"storage\", \"Storage ($/GB/mo)\", \"{}\", None, True),\n (\"bw_up\", \"BW Up ($/GB)\", \"{}\", None, True),\n (\"bw_down\", \"BW Down ($/GB)\", \"{}\", None, True),\n (\"fee\", \"Platform Fee (%)\", \"{}\", None, True),\n)\n\n\nscheduled_jobs_fields = (\n (\"id\", \"Scheduled Job ID\", \"{}\", None, True),\n (\"instance_id\", \"Instance ID\", \"{}\", None, True),\n (\"api_endpoint\", \"API Endpoint\", \"{}\", None, True),\n (\"start_time\", \"Start (Date/Time in UTC)\", \"{}\", lambda x: datetime.fromtimestamp(x, tz=timezone.utc).strftime('%Y-%m-%d/%H:%M'), True),\n (\"end_time\", \"End (Date/Time in UTC)\", \"{}\", lambda x: datetime.fromtimestamp(x, tz=timezone.utc).strftime('%Y-%m-%d/%H:%M'), True),\n (\"day_of_the_week\", \"Day of the Week\", \"{}\", None, True),\n (\"hour_of_the_day\", \"Hour of the Day in UTC\", \"{}\", None, True),\n (\"min_of_the_hour\", \"Minute of the Hour\", \"{}\", None, True),\n (\"frequency\", \"Frequency\", \"{}\", None, True),\n)\n\ninvoice_fields = (\n (\"description\", \"Description\", \"{}\", None, True),\n (\"quantity\", \"Quantity\", \"{}\", None, True),\n (\"rate\", \"Rate\", \"{}\", None, True),\n (\"amount\", \"Amount\", \"{}\", None, True),\n (\"timestamp\", \"Timestamp\", \"{:0.1f}\", None, True),\n (\"type\", \"Type\", \"{}\", None, True)\n)\n\nuser_fields = (\n # (\"api_key\", \"api_key\", \"{}\", None, True),\n (\"balance\", \"Balance\", \"{}\", None, True),\n (\"balance_threshold\", \"Bal. Thld\", \"{}\", None, True),\n (\"balance_threshold_enabled\", \"Bal. Thld Enabled\", \"{}\", None, True),\n (\"billaddress_city\", \"City\", \"{}\", None, True),\n (\"billaddress_country\", \"Country\", \"{}\", None, True),\n (\"billaddress_line1\", \"Addr Line 1\", \"{}\", None, True),\n (\"billaddress_line2\", \"Addr line 2\", \"{}\", None, True),\n (\"billaddress_zip\", \"Zip\", \"{}\", None, True),\n (\"billed_expected\", \"Billed Expected\", \"{}\", None, True),\n (\"billed_verified\", \"Billed Vfy\", \"{}\", None, True),\n (\"billing_creditonly\", \"Billing Creditonly\", \"{}\", None, True),\n (\"can_pay\", \"Can Pay\", \"{}\", None, True),\n (\"credit\", \"Credit\", \"{:0.2f}\", None, True),\n (\"email\", \"Email\", \"{}\", None, True),\n (\"email_verified\", \"Email Vfy\", \"{}\", None, True),\n (\"fullname\", \"Full Name\", \"{}\", None, True),\n (\"got_signup_credit\", \"Got Signup Credit\", \"{}\", None, True),\n (\"has_billing\", \"Has Billing\", \"{}\", None, True),\n (\"has_payout\", \"Has Payout\", \"{}\", None, True),\n (\"id\", \"Id\", \"{}\", None, True),\n (\"last4\", \"Last4\", \"{}\", None, True),\n (\"paid_expected\", \"Paid Expected\", \"{}\", None, True),\n (\"paid_verified\", \"Paid Vfy\", \"{}\", None, True),\n (\"password_resettable\", \"Pwd Resettable\", \"{}\", None, True),\n (\"paypal_email\", \"Paypal Email\", \"{}\", None, True),\n (\"ssh_key\", \"Ssh Key\", \"{}\", None, True),\n (\"user\", \"User\", \"{}\", None, True),\n (\"username\", \"Username\", \"{}\", None, True)\n)\n\nconnection_fields = (\n (\"id\", \"ID\", \"{}\", None, True),\n (\"name\", \"NAME\", \"{}\", None, True),\n (\"cloud_type\", \"Cloud Type\", \"{}\", None, True),\n)\n\n\n# ANSI escape codes for background/foreground colors\nBG_DARK_GRAY = '\\033[40m' # Dark gray background\nBG_LIGHT_GRAY = '\\033[48;5;240m' # Light gray background\nFG_WHITE = '\\033[97m' # Bright white text\nBG_RESET = '\\033[0m' # Reset all formatting\n\ndef display_table(rows: list, fields: Tuple, replace_spaces: bool = True, auto_width: bool = True) -> None:\n \"\"\"Basically takes a set of field names and rows containing the corresponding data and prints a nice tidy table\n of it.\n\n :param list rows: Each row is a dict with keys corresponding to the field names (first element) in the fields tuple.\n\n :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.\n\n :rtype None:\n\n Example of 5-tuple: (\"cpu_ram\", \"RAM\", \"{:0.1f}\", lambda x: x / 1000, False)\n \"\"\"\n header = [name for _, name, _, _, _ in fields]\n out_rows = [header]\n lengths = [len(x) for x in header]\n for instance in rows:\n row = []\n out_rows.append(row)\n for key, name, fmt, conv, _ in fields:\n conv = conv or (lambda x: x)\n val = instance.get(key, None)\n if val is None:\n s = \"-\"\n else:\n val = conv(val)\n s = fmt.format(val)\n if replace_spaces:\n s = s.replace(' ', '_')\n idx = len(row)\n lengths[idx] = max(len(s), lengths[idx])\n row.append(s)\n\n if auto_width:\n width = shutil.get_terminal_size((80, 20)).columns\n start_col_idxs = [0]\n total_len = 4 # +6ch for row label and -2ch for missing last sep in \" \".join()\n for i, l in enumerate(lengths):\n total_len += l + 2\n if total_len > width:\n start_col_idxs.append(i) # index for the start of the next group\n total_len = l + 6 # l + 2 + the 4 from the initial length\n\n groups = {}\n for row in out_rows:\n grp_num = 0\n for i in range(len(start_col_idxs)):\n start = start_col_idxs[i]\n end = start_col_idxs[i+1] if i+1 \u003c len(start_col_idxs) else len(lengths)\n groups.setdefault(grp_num, []).append(row[start:end])\n grp_num += 1\n\n for i, group in groups.items():\n idx = start_col_idxs[i]\n group_lengths = lengths[idx:idx+len(group[0])]\n for row_num, row in enumerate(group):\n bg_color = BG_DARK_GRAY if (row_num - 1) % 2 else BG_LIGHT_GRAY\n row_label = \" #\" if row_num == 0 else f\"{row_num:3d}\"\n out = [row_label]\n for l, s, f in zip(group_lengths, row, fields[idx:idx+len(row)]):\n _, _, _, _, ljust = f\n if ljust: s = s.ljust(l)\n else: s = s.rjust(l)\n out.append(s)\n print(bg_color + FG_WHITE + \" \".join(out) + BG_RESET)\n print()\n else:\n for row in out_rows:\n out = []\n for l, s, f in zip(lengths, row, fields):\n _, _, _, _, ljust = f\n if ljust:\n s = s.ljust(l)\n else:\n s = s.rjust(l)\n out.append(s)\n print(\" \".join(out))\n\n\ndef print_or_page(args, text):\n \"\"\" Print text to terminal, or pipe to pager_cmd if too long. \"\"\"\n line_threshold = shutil.get_terminal_size(fallback=(80, 24)).lines\n lines = text.splitlines()\n if not args.full and len(lines) > line_threshold:\n pager_cmd = ['less', '-R'] if shutil.which('less') else None\n if pager_cmd:\n proc = subprocess.Popen(pager_cmd, stdin=subprocess.PIPE)\n proc.communicate(input=text.encode())\n return True\n else:\n print(text)\n return False\n else:\n print(text)\n return False\n","content_type":"text/x-python; charset=utf-8","language":"python","size":20569,"content_sha256":"f7fa3c41387683862c85c2d5e6ebb605b12e2de7bef45d936133bce65bf5d6c5"},{"filename":"cli/main.py","content":"\"\"\"CLI entry point for vastai.\"\"\"\nimport sys\nimport os\nimport json\nimport requests\nfrom datetime import datetime, timezone\n\nfrom vastai.api.client import is_search_quota_exceeded\nfrom vastai.cli.parser import apwrap, argument, MyWideHelpFormatter, set_completers\nfrom vastai.cli.util import (\n APIKEY_FILE, TFAKEY_FILE, VERSION, server_url_default, api_key_guard,\n format_key_suffix,\n)\n\ntry:\n JSONDecodeError = json.JSONDecodeError\nexcept AttributeError:\n JSONDecodeError = ValueError\n\n\ndef _is_tfa_session_expired(status_code, errmsg):\n \"\"\"True if this error signals an expired 2FA session key.\n\n 401 \"Invalid user key\" is sent when a deleted API key is used; an expired\n TFA session instead returns 404 \"Session expired. Please log in again.\"\n (matches legacy vast.py's handling of both signals).\n \"\"\"\n return (\n (status_code == 401 and errmsg == \"Invalid user key\")\n or (status_code == 404 and errmsg == \"Session expired. Please log in again.\")\n )\n\n\ndef with_quota_reset(response, errmsg):\n \"\"\"Append the reset time from ``X-RateLimit-Reset`` to a search quota denial's message.\"\"\"\n if not is_search_quota_exceeded(response):\n return errmsg\n try:\n reset = datetime.fromtimestamp(int(response.headers[\"X-RateLimit-Reset\"]), timezone.utc)\n except (KeyError, ValueError):\n return errmsg\n return f\"{errmsg}. Resets at {reset:%Y-%m-%d %H:%M} UTC.\"\n\n\ndef read_key(path):\n \"\"\"The key stored at a path, or None if it is absent or unreadable.\"\"\"\n try:\n with open(path) as f:\n return f.read().strip() or None\n except OSError:\n return None\n\n\ndef _emit_error(args, status_code, message):\n \"\"\"Emit a command error in the appropriate format.\n\n In ``--raw`` mode, prints a JSON error object to stderr so scripts can\n parse it; otherwise prints a human-readable line. Always goes to stderr\n so stdout stays clean for scripting consumers.\n \"\"\"\n if getattr(args, \"raw\", False):\n payload = {\"error\": True, \"status_code\": status_code, \"msg\": message}\n print(json.dumps(payload), file=sys.stderr)\n return\n\n if status_code:\n print(f\"Failed with error {status_code}: {message}\", file=sys.stderr)\n else:\n print(message, file=sys.stderr)\n\n if status_code == 401:\n if \"Two Factor Authentication\" in message:\n print(\n \"\\nThis endpoint requires a 2FA session. Authenticate with the method you have set up:\\n\"\n \" vastai tfa login --method-type totp -c \u003cCODE>\\n\"\n \" vastai tfa login --method-type sms --secret \u003cSECRET> -c \u003cCODE>\\n\"\n \" vastai tfa login --method-type email --secret \u003cSECRET> -c \u003cCODE>\\n\"\n \" vastai tfa login --backup-code ABCD-EFGH-IJKL\\n\"\n \"For SMS or email, first run 'vastai tfa send-sms' or 'vastai tfa send-email' to get \u003cSECRET>.\",\n file=sys.stderr,\n )\n return\n\n env = os.environ.get(\"VAST_API_KEY\")\n file_key = read_key(APIKEY_FILE)\n tfa_key = read_key(TFAKEY_FILE)\n\n # A --api-key outranks the rest, so a value matching none came from there.\n sent = getattr(args, \"api_key\", None)\n key_missing = not sent and not env and not file_key\n\n if sent and sent == tfa_key:\n print(f\" Sent your 2FA session key from {TFAKEY_FILE} (ends in {format_key_suffix(sent)}).\", file=sys.stderr)\n print(\" If it has expired, start a new one: vastai tfa login\", file=sys.stderr)\n elif sent and sent != env and sent != file_key:\n server = getattr(args, \"url\", None) or server_url_default\n print(f\" Sent key from --api-key (ends in {format_key_suffix(sent)}).\", file=sys.stderr)\n print(f\" Check that key is valid for {server}.\", file=sys.stderr)\n elif key_missing:\n print(\" No API key is configured.\", file=sys.stderr)\n print(\" Run: vastai set api-key \u003cKEY>\", file=sys.stderr)\n print(\" Create a key at https://console.vast.ai/manage-keys/\", file=sys.stderr)\n elif env and file_key:\n print(f\" Sent key from $VAST_API_KEY (ends in {format_key_suffix(env)}). Env var overrides the file.\", file=sys.stderr)\n 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)\n elif env:\n 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)\n elif file_key:\n print(f\" Sent key from {APIKEY_FILE} (ends in {format_key_suffix(file_key)}). Update with: vastai set api-key \u003cKEY>\", file=sys.stderr)\n\n\n# Create the global parser instance\nparser = apwrap(\n 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\",\n formatter_class=MyWideHelpFormatter\n)\n\n\ndef main():\n # Import all command modules - the import itself triggers decorator\n # registrations on the global parser via _get_parser().\n from vastai.cli.commands import ( # noqa: F401\n instances, offers, machines, teams, keys, endpoints,\n billing, storage, auth, misc, deployments, metrics,\n benchmarks,\n price_increase,\n repl,\n update,\n uninstall,\n # clusters, # cluster/overlay commands disabled for now\n )\n\n # Wire up tab completers now that command modules are loaded\n set_completers(\n instance_machine_fn=lambda **kw: instances.show__instances(\n type('Args', (), {'api_key': os.getenv('VAST_API_KEY'), 'url': server_url_default,\n 'retry': 3, 'explain': False, 'curl': False,\n 'quiet': False, 'raw': False})(),\n {'internal': True, 'field': 'machine_id'}),\n instance_fn=lambda **kw: instances.show__instances(\n type('Args', (), {'api_key': os.getenv('VAST_API_KEY'), 'url': server_url_default,\n 'retry': 3, 'explain': False, 'curl': False,\n 'quiet': False, 'raw': False})(),\n {'internal': True, 'field': 'id'}),\n )\n\n # Add global arguments\n parser.add_argument(\"--url\", help=\"Server REST API URL\", default=server_url_default)\n parser.add_argument(\"--retry\", help=\"Retry limit\", type=int, default=3)\n parser.add_argument(\"--explain\", action=\"store_true\", help=\"Output verbose explanation of mapping of CLI calls to HTTPS API endpoints\")\n parser.add_argument(\"--raw\", action=\"store_true\", help=\"Output machine-readable json\")\n parser.add_argument(\"--full\", action=\"store_true\", help=\"Print full results instead of paging with `less` for commands that support it\")\n parser.add_argument(\"--curl\", action=\"store_true\", help=\"Show a curl equivalency to the call\")\n 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))\n parser.add_argument(\"--version\", help=\"Show CLI version\", action=\"version\", version=VERSION)\n parser.add_argument(\"--no-color\", action=\"store_true\", help=\"Disable colored output for commands that support it\")\n\n # Tab completion\n try:\n import argcomplete\n import argparse\n\n from typing import List, Optional\n from vastai.cli.parser import build_command_maps, two_stage_command_completions\n class MyAutocomplete(argcomplete.CompletionFinder):\n # Two-stage (verb -> object) completion over the flat \"verb object\"\n # subparser names. Logic in parser.py (pure + unit-tested).\n def _command_maps(self):\n cached = getattr(self, \"_cmd_maps\", None)\n if cached is None:\n from vastai.cli.util import get_role\n cached = self._cmd_maps = build_command_maps(self._parser, role=get_role())\n return cached\n\n def _get_completions(self, comp_words, cword_prefix, cword_prequote, last_wordbreak_pos):\n verbs, verb_objs, singles = self._command_maps()\n cands, merged = two_stage_command_completions(\n comp_words, cword_prefix, verbs, verb_objs, singles)\n if cands is not None:\n return cands\n return super()._get_completions(merged, cword_prefix, cword_prequote, last_wordbreak_pos)\n\n def quote_completions(self, completions: List[str], cword_prequote: str, last_wordbreak_pos: Optional[int]) -> List[str]:\n pre = super().quote_completions(completions, cword_prequote, last_wordbreak_pos)\n return sorted(pre, key=lambda x: x.startswith('-'))\n\n myautocc = MyAutocomplete()\n myautocc(parser.parser)\n except ImportError:\n pass\n\n args = parser.parse_args()\n\n # Passive upgrade nudge: best-effort, ≤1 manifest GET + ≤1 stderr line per\n # 24h, silent when offline/piped/CI. Never raises. Opt out with\n # VASTAI_NO_UPDATE_CHECK=1. Runs first so it's visible before any command\n # output, not buried after it. See selfupdate.py and §7.\n from vastai.cli.selfupdate import notify_update\n notify_update(args)\n\n # API key resolution\n if args.api_key is api_key_guard:\n key_file = TFAKEY_FILE if os.path.exists(TFAKEY_FILE) else APIKEY_FILE\n if os.path.exists(key_file):\n with open(key_file, \"r\") as reader:\n args.api_key = reader.read().strip()\n else:\n args.api_key = None\n\n run_command(args)\n\n\ndef run_command(args):\n \"\"\"Execute ``args.func(args)``, retrying once on an expired 2FA session key.\"\"\"\n while True:\n try:\n res = args.func(args)\n\n if args.raw and res is not None:\n try:\n print(json.dumps(res, indent=1, sort_keys=True))\n except (TypeError, ValueError):\n print(json.dumps(res.json(), indent=1, sort_keys=True))\n sys.exit(0)\n sys.exit(res)\n\n except requests.exceptions.HTTPError as e:\n try:\n errmsg = e.response.json().get(\"msg\")\n except JSONDecodeError:\n if e.response.status_code == 401:\n errmsg = \"Please log in or sign up\"\n else:\n errmsg = \"(no detail message supplied)\"\n\n # 2FA session key expired -> retry with long-lived API key.\n if _is_tfa_session_expired(e.response.status_code, errmsg) and os.path.exists(TFAKEY_FILE):\n print(f\"Failed with error {e.response.status_code}: Your 2FA session has expired.\")\n os.remove(TFAKEY_FILE)\n if os.path.exists(APIKEY_FILE):\n with open(APIKEY_FILE, \"r\") as reader:\n args.api_key = reader.read().strip()\n print(f\"Trying again with your normal API Key from {APIKEY_FILE}...\")\n print(\"To start a new 2FA session, run: vastai tfa login\")\n continue\n else:\n print(\"Run `vastai tfa login` to start a new 2FA session and try again.\")\n break\n\n _emit_error(args, e.response.status_code, with_quota_reset(e.response, errmsg))\n break\n\n except ValueError as e:\n _emit_error(args, 0, str(e))\n break\n\n\nif __name__ == \"__main__\":\n try:\n main()\n except (KeyboardInterrupt, BrokenPipeError):\n pass\n","content_type":"text/x-python; charset=utf-8","language":"python","size":11741,"content_sha256":"d2d1a9cd1bfa3c11af8babd80975dcd01f222e09a727df4b1d087525da4656e5"},{"filename":"cli/parser.py","content":"\"\"\"\nArgument parser infrastructure for the Vast.ai CLI.\n\nExtracted from vast.py - contains the core parser wrapper classes\nthat handle command registration, argument parsing, and tab completion.\n\"\"\"\n\nfrom __future__ import unicode_literals, print_function\n\nimport sys\nimport argparse\nfrom pathlib import Path\n\nfrom vastai.cli.util import get_role, is_client_view\n\n\n# ---------------------------------------------------------------------------\n# Tab-completion helpers\n#\n# These reference show__instances which lives in the CLI commands layer.\n# We keep them as lazy references (initially None) so this module can be\n# imported without pulling in the whole command tree. The CLI entry-point\n# populates them once everything is wired up.\n# ---------------------------------------------------------------------------\n\n# Populated later by the CLI entry point\n_complete_instance_machine = None\n_complete_instance = None\n\ndef complete_instance_machine(prefix=None, action=None, parser=None, parsed_args=None):\n if _complete_instance_machine is not None:\n return _complete_instance_machine(prefix=prefix, action=action, parser=parser, parsed_args=parsed_args)\n return []\n\ndef complete_instance(prefix=None, action=None, parser=None, parsed_args=None):\n if _complete_instance is not None:\n return _complete_instance(prefix=prefix, action=action, parser=parser, parsed_args=parsed_args)\n return []\n\ndef complete_sshkeys(prefix=None, action=None, parser=None, parsed_args=None):\n return [str(m) for m in Path.home().joinpath('.ssh').glob('*.pub')]\n\n\n# ---------------------------------------------------------------------------\n# Two-stage (verb -> object) command-name completion\n#\n# Pure, unit-testable helpers that turn the flat \"verb object\" subparser names\n# into staged completion. The argcomplete adapter in cli/main.py calls them.\n# ---------------------------------------------------------------------------\n\ndef build_command_maps(parser, role=None):\n \"\"\"Derive (verbs, verb_objs, singles) from a parser's subparser names, excluding hidden and (unless role='host') host-only commands.\"\"\"\n names = []\n for action in parser._actions:\n if isinstance(action, argparse._SubParsersAction):\n for name, sp in action.choices.items():\n if getattr(sp, 'hidden', False) or name != getattr(sp, 'command_name', name):\n continue\n if is_client_view(role) and getattr(sp, 'host_only', False):\n continue\n names.append(name)\n break\n verbs, verb_objs, singles = set(), {}, set()\n for name in names:\n verb, sep, obj = name.partition(\" \")\n if sep:\n verbs.add(verb)\n verb_objs.setdefault(verb, set()).add(obj)\n else:\n singles.add(verb)\n return verbs, verb_objs, singles\n\n\ndef two_stage_command_completions(comp_words, cword_prefix, verbs, verb_objs, singles):\n \"\"\"Stage command-name completion. Returns (candidates, merged): candidates is\n the verb/object list, or None to delegate flag completion using merged.\"\"\"\n n = len(comp_words)\n if n == 1: # completing the verb / bare command\n return sorted(c for c in (verbs | singles) if c.startswith(cword_prefix)), None\n if n == 2 and comp_words[1] in verbs: # completing the object for a verb\n objs = verb_objs.get(comp_words[1], set())\n return sorted(o for o in objs if o.startswith(cword_prefix)), None\n # Past the command name: fuse \"verb object\" so the right subparser is found.\n # Never fuse across a flag, mirroring apwrap.parse_args.\n merged = comp_words\n if n >= 3 and comp_words[1] in verbs and not comp_words[2].startswith(\"-\"):\n merged = [comp_words[0], comp_words[1] + \" \" + comp_words[2]] + list(comp_words[3:])\n return None, merged\n\n\ndef set_completers(instance_machine_fn=None, instance_fn=None):\n \"\"\"\n Wire up the tab-completion functions. Called once from the CLI\n entry-point after all command modules have been loaded.\n \"\"\"\n global _complete_instance_machine, _complete_instance\n if instance_machine_fn is not None:\n _complete_instance_machine = instance_machine_fn\n if instance_fn is not None:\n _complete_instance = instance_fn\n\n\n# ---------------------------------------------------------------------------\n# Argument wrapper\n# ---------------------------------------------------------------------------\n\nclass argument(object):\n \"\"\"Thin wrapper that stores positional and keyword args for later\n application to an argparse subparser.\"\"\"\n def __init__(self, *args, mutex_group=None, **kwargs):\n self.args = args\n self.kwargs = kwargs\n self.mutex_group = mutex_group # Name of the mutually exclusive group this arg belongs to\n\n\nclass hidden_aliases(object):\n \"\"\"A list-like that is falsy so argparse won't display aliases in help.\"\"\"\n def __init__(self, l):\n self.l = l\n\n def __iter__(self):\n return iter(self.l)\n\n def __bool__(self):\n return False\n\n def __nonzero__(self):\n return False\n\n def append(self, x):\n self.l.append(x)\n\n\nclass MyWideHelpFormatter(argparse.RawTextHelpFormatter):\n def __init__(self, prog):\n super().__init__(prog, width=128, max_help_position=50, indent_increment=1)\n\n\nCOMMAND_GROUPS = {\n 'vastai.cli.commands.offers': 'Search',\n 'vastai.cli.commands.instances': 'Instances',\n 'vastai.cli.commands.benchmarks': 'Instances',\n 'vastai.cli.commands.machines': 'Host machines',\n 'vastai.cli.commands.teams': 'Teams',\n 'vastai.cli.commands.keys': 'Auth & keys',\n 'vastai.cli.commands.auth': 'Auth & keys',\n 'vastai.cli.commands.billing': 'Billing & account',\n 'vastai.cli.commands.endpoints': 'Serverless',\n 'vastai.cli.commands.deployments': 'Serverless',\n 'vastai.cli.commands.metrics': 'Metrics',\n 'vastai.cli.commands.storage': 'Storage volumes',\n 'vastai.cli.commands.misc': 'Other',\n 'vastai.cli.commands.repl': 'Other',\n 'vastai.cli.commands.update': 'Other',\n}\n\n# Some misc.py commands belong with their target resource section, not Other.\nCOMMAND_OVERRIDES = {\n 'execute': 'Instances',\n 'logs': 'Instances',\n 'ssh-url': 'Instances',\n 'scp-url': 'Instances',\n 'take snapshot': 'Instances',\n 'reports': 'Host machines',\n}\n\n# Host-only command classification (CLN-3582): display-only, hides host-only commands from --help/completion/hints for the client role; unclassified defaults to visible.\n\n# Modules where every command is host-only.\nHOST_ONLY_MODULES = {\n 'vastai.cli.commands.machines',\n 'vastai.cli.commands.metrics',\n}\n\n# Individual host-only commands in modules that are mostly client-facing (same shape as COMMAND_OVERRIDES above).\nHOST_ONLY_COMMANDS = {\n 'reports', # misc.py: machine reports, not instance logs\n 'show earnings', # billing.py: machine earning history\n 'list network-volume', # storage.py: list disk space for rent\n 'list volume',\n 'list volumes',\n 'unlist network-volume',\n 'unlist volume',\n}\n\n\ndef is_host_only_command(name, module, explicit=None):\n \"\"\"Whether a command is hidden from the client CLI view: explicit host_only= kwarg > per-command override > per-module default > False (visible).\"\"\"\n if explicit is not None:\n return bool(explicit)\n if name in HOST_ONLY_COMMANDS:\n return True\n return module in HOST_ONLY_MODULES\n\n\nGROUP_ORDER = [\n 'Search', 'Instances', 'Host machines', 'Teams', 'Auth & keys',\n 'Billing & account', 'Serverless', 'Metrics', 'Storage volumes', 'Other',\n]\n\n# ---------------------------------------------------------------------------\n# Unreleased/feature-flagged commands\n#\n# Hidden from --help and tab completion regardless of anything else — a\n# discoverability gate, not an access gate: the command still runs if typed\n# directly, so internal testing keeps working. Remove an entry once the\n# feature is ready to announce.\n# ---------------------------------------------------------------------------\n\nHIDDEN_COMMANDS = {\n 'repl', # interactive shell: internal testing before announcing\n 'search network-volumes', # network volumes are not yet released\n 'create network-volume',\n 'list network-volume',\n 'unlist network-volume',\n 'show network-disks', # network disks are not yet released\n 'add network-disk',\n}\n\n\ndef is_hidden_command(name, explicit=None):\n \"\"\"Whether a registered command should be hidden from --help/completion.\n\n Priority: explicit ``hidden=`` kwarg on the decorator > ``HIDDEN_COMMANDS``.\n Defaults to ``False`` (visible).\n \"\"\"\n if explicit is not None:\n return bool(explicit)\n return name in HIDDEN_COMMANDS\n\n\nclass GroupedArgumentParser(argparse.ArgumentParser):\n \"\"\"ArgumentParser that renders subcommands grouped by registering module.\"\"\"\n\n def format_help(self):\n sub_action = self._subparsers_action()\n if sub_action is None or not sub_action.choices:\n return super().format_help()\n\n formatter = self._get_formatter()\n formatter.add_usage(self.usage, self._actions, self._mutually_exclusive_groups)\n formatter.add_text(self.description)\n\n # Skip the subparsers action; we render commands grouped below.\n for action_group in self._action_groups:\n non_sub_actions = [\n a for a in action_group._group_actions\n if not isinstance(a, argparse._SubParsersAction)\n ]\n if not non_sub_actions:\n continue\n formatter.start_section(action_group.title)\n formatter.add_text(action_group.description)\n formatter.add_arguments(non_sub_actions)\n formatter.end_section()\n\n self._add_grouped_commands(formatter, sub_action)\n formatter.add_text(self.epilog)\n return formatter.format_help()\n\n def _subparsers_action(self):\n for a in self._actions:\n if isinstance(a, argparse._SubParsersAction):\n return a\n return None\n\n def _add_grouped_commands(self, formatter, sub_action):\n \"\"\"Render grouped commands.\"\"\"\n role = get_role()\n grouped = {}\n for pseudo in sub_action._choices_actions:\n sp = sub_action.choices.get(pseudo.dest)\n if sp is None:\n continue\n if getattr(sp, 'hidden', False):\n continue\n if is_client_view(role) and getattr(sp, 'host_only', False):\n continue\n label = COMMAND_OVERRIDES.get(pseudo.dest)\n if label is None:\n func = sp.get_default('func')\n module = getattr(func, '__module__', None)\n label = COMMAND_GROUPS.get(module, 'Other')\n grouped.setdefault(label, []).append(pseudo)\n\n seen = set()\n for label in GROUP_ORDER:\n if label in grouped:\n formatter.start_section(label)\n formatter.add_arguments(grouped[label])\n formatter.end_section()\n seen.add(label)\n for label, actions in grouped.items():\n if label in seen:\n continue\n formatter.start_section(label)\n formatter.add_arguments(actions)\n formatter.end_section()\n\n\n# ---------------------------------------------------------------------------\n# Main parser wrapper\n# ---------------------------------------------------------------------------\n\nclass apwrap(object):\n \"\"\"Wraps :class:`argparse.ArgumentParser` with convenience methods for\n two-word command registration (``verb object``) and tab-completion.\"\"\"\n\n def __init__(self, *args, **kwargs):\n if \"formatter_class\" not in kwargs:\n kwargs[\"formatter_class\"] = MyWideHelpFormatter\n self.parser = GroupedArgumentParser(*args, **kwargs)\n self.parser.set_defaults(func=self.fail_with_help)\n self.subparsers_ = None\n self.subparser_objs = []\n self.added_help_cmd = False\n self.post_setup = []\n self.verbs = set()\n self.objs = set()\n\n def fail_with_help(self, *a, **kw):\n self.parser.print_help(sys.stderr)\n raise SystemExit\n\n def add_argument(self, *a, **kw):\n if not kw.get(\"parent_only\"):\n for x in self.subparser_objs:\n try:\n # Create a global options group for better visual separation\n if not hasattr(x, '_global_options_group'):\n x._global_options_group = x.add_argument_group('Global options (available for all commands)')\n # Use SUPPRESS as default for subparsers so they don't overwrite\n # values already set by the main parser when the argument is placed\n # before the subcommand (e.g., `vastai --url \u003curl> get wrkgrp-logs`)\n subparser_kw = kw.copy()\n subparser_kw['default'] = argparse.SUPPRESS\n x._global_options_group.add_argument(*a, **subparser_kw)\n except argparse.ArgumentError:\n # duplicate - or maybe other things, hopefully not\n pass\n return self.parser.add_argument(*a, **kw)\n\n def subparsers(self, *a, **kw):\n if self.subparsers_ is None:\n kw[\"metavar\"] = \"command\"\n kw[\"help\"] = \"command to run. one of:\"\n kw[\"dest\"] = \"command\" # the name as typed, so a command can tell when its alias was used\n self.subparsers_ = self.parser.add_subparsers(*a, **kw)\n return self.subparsers_\n\n def get_name(self, verb, obj):\n if obj:\n self.verbs.add(verb)\n self.objs.add(obj)\n name = verb + ' ' + obj\n else:\n self.objs.add(verb)\n name = verb\n return name\n\n def command(self, *arguments, aliases=(), help=None, hidden=None, host_only=None, **kwargs):\n help_ = help\n if not self.added_help_cmd:\n self.added_help_cmd = True\n\n @self.command(argument(\"subcommand\", default=None, nargs=\"?\"), help=\"print this help message\")\n def help(*a, **kw):\n self.fail_with_help()\n\n def inner(func):\n dashed_name = func.__name__.replace(\"_\", \"-\")\n verb, _, obj = dashed_name.partition(\"--\")\n name = self.get_name(verb, obj)\n aliases_transformed = [] if aliases else hidden_aliases([])\n for x in aliases:\n verb, _, obj = x.partition(\" \")\n aliases_transformed.append(self.get_name(verb, obj))\n if \"formatter_class\" not in kwargs:\n kwargs[\"formatter_class\"] = MyWideHelpFormatter\n\n sp = self.subparsers().add_parser(name, aliases=aliases_transformed, help=help_, **kwargs)\n sp.hidden = is_hidden_command(name, explicit=hidden)\n sp.host_only = is_host_only_command(name, func.__module__, explicit=host_only)\n sp.command_name = name\n\n # TODO: Sometimes the parser.command has a help parameter. Ideally\n # I'd extract this during the sdk phase but for the life of me\n # I can't find it.\n setattr(func, \"mysignature\", sp)\n setattr(func, \"mysignature_help\", help_)\n\n self.subparser_objs.append(sp)\n\n self._process_arguments_with_groups(sp, arguments)\n\n sp.set_defaults(func=func)\n return func\n\n if len(arguments) == 1 and type(arguments[0]) != argument:\n func = arguments[0]\n arguments = []\n return inner(func)\n return inner\n\n def parse_args(self, argv=None, *a, **kw):\n if argv is None:\n argv = sys.argv[1:]\n argv_ = []\n for x in argv:\n # Merge \"verb object\" into one command token, but never swallow a\n # flag: a verb that is also a bare command (e.g. `update`) must\n # accept options (`vastai update --check`).\n if argv_ and argv_[-1] in self.verbs and not x.startswith(\"-\"):\n argv_[-1] += \" \" + x\n else:\n argv_.append(x)\n args = self.parser.parse_args(argv_, *a, **kw)\n for func in self.post_setup:\n func(args)\n return args\n\n def _process_arguments_with_groups(self, parser_obj, arguments):\n \"\"\"Process arguments and handle mutually exclusive groups\"\"\"\n mutex_groups_to_required = {}\n arg_to_group = {}\n\n # Determine if any mutex groups are required\n for arg in arguments:\n key = arg.args[0]\n if arg.mutex_group:\n is_required = arg.kwargs.pop('required', False)\n group_name = arg.mutex_group\n arg_to_group[key] = group_name\n if mutex_groups_to_required.get(group_name):\n continue # if marked as required then it stays required\n else:\n mutex_groups_to_required[group_name] = is_required\n\n name_to_group_parser = {} # Create mutually exclusive group parsers\n for group_name, is_required in mutex_groups_to_required.items():\n mutex_group = parser_obj.add_mutually_exclusive_group(required=is_required)\n name_to_group_parser[group_name] = mutex_group\n\n for arg in arguments: # Add args via the appropriate parser\n key = arg.args[0]\n if arg_to_group.get(key):\n group_parser = name_to_group_parser[arg_to_group[key]]\n tsp = group_parser.add_argument(*arg.args, **arg.kwargs)\n else:\n tsp = parser_obj.add_argument(*arg.args, **arg.kwargs)\n self._add_completer(tsp, arg)\n\n\n def _add_completer(self, tsp, arg):\n \"\"\"Helper function to add completers based on argument names\"\"\"\n myCompleter = None\n comparator = arg.args[0].lower()\n if comparator.startswith('machine'):\n myCompleter = complete_instance_machine\n elif comparator.startswith('id') or comparator.endswith('id'):\n myCompleter = complete_instance\n elif comparator.startswith('ssh'):\n myCompleter = complete_sshkeys\n\n if myCompleter:\n setattr(tsp, 'completer', myCompleter)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":18558,"content_sha256":"c96573e01c1ed601004655d2c086d4bc9db505070b6cd5e2cf30e9d9202f6807"},{"filename":"cli/repl/__init__.py","content":"\"\"\"Interactive shell for the vast CLI (`vastai repl`).\n\nOne process, the same commands: `session.py` owns the loop, `bridge.py` runs a\nline through the real parser, `completion.py` stages Tab completion and\n`catalog.py` indexes the command tree for both.\n\"\"\"\nfrom vastai.cli.repl.session import run_repl # noqa: F401\n","content_type":"text/x-python; charset=utf-8","language":"python","size":318,"content_sha256":"c9b6cb2ac0822ae568b432cd8b50f8a5fdd6ccebe3f721a37b3440237940f7c2"},{"filename":"cli/repl/bridge.py","content":"\"\"\"Run one CLI line inside the REPL process.\n\nThe REPL reuses the exact parser and command functions that ``vastai ...``\nuses, so ``show instances`` at the prompt behaves identically to\n``vastai show instances`` — same auth, same output, same flags. The only\ndifference is control flow: a failing command ends the line, never the session.\n\"\"\"\nimport argparse\nimport json\nimport os\nimport shlex\nimport sys\n\nimport requests\n\nfrom vastai.cli import main as cli_main\nfrom vastai.cli.repl.catalog import option_action\n\n# Globals the session carries onto every line, so a bare `show instances`\n# doesn't have to repeat --api-key/--url/--raw. A flag typed on the line wins.\nSESSION_GLOBALS = (\n \"api_key\", \"url\", \"retry\", \"explain\", \"curl\", \"raw\", \"full\", \"no_color\",\n)\n\n# Outcomes of an expired-2FA-session recovery attempt.\nRETRY = \"retry\" # a fresh key was loaded; run the command again\nHANDLED = \"handled\" # the expiry was reported; don't report it twice\n\n\ndef run_line(parser, line, session_args):\n \"\"\"Parse and execute one command line. Returns the command's return value,\n or None if the line failed to parse or the command errored.\"\"\"\n try:\n argv = tokenize(line)\n except ValueError as exc:\n print(f\"parse error: {exc}\", file=sys.stderr)\n return None\n if not argv:\n return None\n\n try:\n args = parser.parse_args(argv)\n except SystemExit:\n return None # argparse already printed usage or an error\n\n apply_session_globals(parser, args, session_args, argv)\n func = getattr(args, \"func\", None)\n if func is None:\n return None\n return _invoke(args, func, session_args)\n\n\ndef tokenize(line):\n \"\"\"Split a line into argv the way a shell would.\n\n POSIX rules everywhere except Windows, where a backslash is a path\n separator rather than an escape: POSIX splitting would quietly turn\n ``--onstart C:\\\\tmp\\\\go.sh`` into ``C:tmpgo.sh``.\n \"\"\"\n if os.name != \"nt\":\n return shlex.split(line)\n lexer = shlex.shlex(line, posix=False)\n lexer.whitespace_split = True\n return [_unquote(token) for token in lexer]\n\n\ndef _unquote(token):\n if len(token) > 1 and token[0] == token[-1] and token[0] in \"\\\"'\":\n return token[1:-1]\n return token\n\n\ndef apply_session_globals(parser, args, session_args, argv=()):\n \"\"\"Carry the session's global flags onto one line's args.\n\n Presence is read from the line itself rather than by comparing parsed\n values against defaults: an option typed with exactly its default value\n (``--retry 3`` while the session runs with ``--retry 10``) must still win.\n \"\"\"\n inner = getattr(parser, \"parser\", parser)\n sub = getattr(getattr(args, \"func\", None), \"mysignature\", None)\n typed = explicit_dests(inner, sub, argv)\n for name in SESSION_GLOBALS:\n if name in typed or not hasattr(args, name) or not hasattr(session_args, name):\n continue\n setattr(args, name, getattr(session_args, name))\n\n\ndef explicit_dests(root, sub, argv):\n \"\"\"The dests of the options a line actually names, in any form argparse\n accepts (``--flag``, ``--flag=value``, a short alias, an abbreviation).\n\n Each option is read against the parser argparse itself would use: the root\n parser before the command name, the command's own parser after it. An\n abbreviation can be unambiguous in one and ambiguous in the other — ``--c``\n is ``--curl`` up front, but could be ``--cols`` on ``show instances``.\n\n Scanning stops where argparse stops treating words as options: at ``--`` or\n at a REMAINDER option, so ``create instance --args --raw`` passes ``--raw``\n to the container rather than counting it as a global.\n \"\"\"\n dests = set()\n parser, skip = root, 0\n for token in argv:\n if skip:\n skip -= 1\n continue\n if token == \"--\":\n break\n if not token.startswith(\"-\") or token == \"-\":\n parser = sub or root # past the command name from here on\n continue\n action = option_action(parser, token)\n if action is None:\n continue\n if action.nargs == argparse.REMAINDER:\n break\n dests.add(action.dest)\n if parser is root and action.nargs != 0 and \"=\" not in token:\n skip = 1 # its value is not the command name\n return dests\n\n\ndef _invoke(args, func, session_args=None):\n try:\n res = func(args)\n except SystemExit:\n # Commands (and --help) exit the process in one-shot mode; here that\n # just ends the line.\n return None\n except requests.exceptions.HTTPError as exc:\n outcome = _recover_expired_tfa_session(args, exc, session_args)\n if outcome == RETRY:\n return _invoke(args, func) # retry once, as the one-shot CLI does\n if outcome != HANDLED:\n _emit_http_error(args, exc)\n return None\n except ValueError as exc:\n cli_main._emit_error(args, 0, str(exc))\n return None\n except KeyboardInterrupt:\n print(\"^C\", file=sys.stderr)\n return None\n except Exception as exc: # a broken command must not kill the session\n print(f\"{type(exc).__name__}: {exc}\", file=sys.stderr)\n return None\n\n if getattr(args, \"raw\", False) and res is not None:\n _print_raw(res)\n return res\n\n\ndef _recover_expired_tfa_session(args, exc, session_args):\n \"\"\"Fall back to the saved API key when a 2FA session expires, as\n ``main.run_command`` does — and keep the new key on the session, so the\n rest of the REPL's lines work too rather than failing one by one.\n\n Returns RETRY, HANDLED (reported, nothing left to say) or None (not a 2FA\n expiry, so the caller reports it).\n \"\"\"\n status, msg = _error_detail(exc)\n if not cli_main._is_tfa_session_expired(status, msg):\n return None\n if not os.path.exists(cli_main.TFAKEY_FILE):\n return None\n\n print(f\"Failed with error {status}: Your 2FA session has expired.\")\n os.remove(cli_main.TFAKEY_FILE)\n if not os.path.exists(cli_main.APIKEY_FILE):\n print(\"Run `vastai tfa login` to start a new 2FA session and try again.\")\n return HANDLED\n\n with open(cli_main.APIKEY_FILE, \"r\") as reader:\n key = reader.read().strip()\n args.api_key = key\n if session_args is not None:\n session_args.api_key = key\n print(f\"Trying again with your normal API Key from {cli_main.APIKEY_FILE}...\")\n print(\"To start a new 2FA session, run: vastai tfa login\")\n return RETRY\n\n\ndef _error_detail(exc):\n \"\"\"(status, message) from an API error, however malformed the response.\"\"\"\n resp = exc.response\n status = getattr(resp, \"status_code\", 0)\n try:\n msg = resp.json().get(\"msg\")\n except (ValueError, AttributeError):\n msg = None\n if not isinstance(msg, str):\n # A body that is valid JSON but carries no \"msg\" would otherwise hand\n # None to _emit_error, which tests it with `in` and raises TypeError.\n msg = \"Please log in or sign up\" if status == 401 else \"(no detail message supplied)\"\n return status, msg\n\n\ndef _emit_http_error(args, exc):\n \"\"\"Format an API error the same way the one-shot CLI's main loop does.\"\"\"\n status, msg = _error_detail(exc)\n cli_main._emit_error(args, status, msg)\n\n\ndef _print_raw(res):\n try:\n print(json.dumps(res, indent=1, sort_keys=True))\n except (TypeError, ValueError):\n print(json.dumps(res.json(), indent=1, sort_keys=True))\n","content_type":"text/x-python; charset=utf-8","language":"python","size":7516,"content_sha256":"96bacfdf3195ac8d508427e38fc799b3c125aa16e7d6b84b0fe0c3dfe2e87b05"},{"filename":"cli/repl/catalog.py","content":"\"\"\"What the REPL knows about the CLI's command tree.\n\nCommands are registered on one argparse parser as flat ``\"verb object\"``\nsubparser names (\"show instances\"). The catalog derives the two-stage view of\nthat — verbs, the objects each verb takes, bare commands, per-command options —\nonce at startup, so completion and command resolution stay dict lookups instead\nof a walk over ~150 subparsers on every keystroke.\n\"\"\"\nimport argparse\nimport difflib\n\nfrom vastai.cli.parser import build_command_maps\n\n\ndef option_action(parser, token):\n \"\"\"The action a flag token names on a parser: exact match, ``--flag=value``,\n or the unambiguous abbreviation argparse itself would accept.\"\"\"\n options = parser._option_string_actions\n name = token.split(\"=\", 1)[0]\n action = options.get(name)\n if action is not None:\n return action\n if name.startswith(\"--\"):\n matches = {id(a): a for opt, a in options.items() if opt.startswith(name)}\n if len(matches) == 1:\n return next(iter(matches.values()))\n return None\n\n\nclass CommandCatalog:\n \"\"\"Two-stage (verb -> object) index over a live ``apwrap`` parser.\"\"\"\n\n def __init__(self, parser):\n self._parser = getattr(parser, \"parser\", parser)\n self.verbs, self.verb_objects, self.singles = build_command_maps(self._parser)\n self._choices = self._subparser_choices()\n # Hidden commands are left out of the maps above — they are gated from\n # discovery, not from use — but a line naming one must still resolve,\n # or the REPL would reject a command the one-shot CLI runs.\n self._runnable_verbs = {name.split(\" \")[0] for name in self._choices if \" \" in name}\n self._flag_cache = {}\n self.first_words = sorted(self.verbs | self.singles)\n self.names = sorted(self.singles | {\n f\"{verb} {obj}\" for verb, objs in self.verb_objects.items() for obj in objs})\n\n def _subparser_choices(self):\n for action in self._parser._actions:\n if isinstance(action, argparse._SubParsersAction):\n return action.choices\n return {}\n\n def objects(self, verb):\n \"\"\"The objects registered for a verb, e.g. show -> instances, user, ...\"\"\"\n return sorted(self.verb_objects.get(verb, ()))\n\n def strip_options(self, tokens):\n \"\"\"Drop the global options a line leads with (``--raw show user``), so\n what's left starts at the command name.\"\"\"\n i = 0\n while i \u003c len(tokens) and tokens[i].startswith(\"-\") and tokens[i] != \"-\":\n token = tokens[i]\n action = option_action(self._parser, token)\n i += 1\n if action is not None and action.nargs != 0 and \"=\" not in token:\n i += 1 # the option's value\n return tokens[i:]\n\n def resolve(self, tokens):\n \"\"\"The command name a tokenised line names, or None if it names none.\n\n Mirrors ``apwrap.parse_args``: a verb absorbs the next token as its\n object, but never a flag, so ``update --check`` stays the bare command.\n \"\"\"\n if not tokens:\n return None\n if tokens[0] in self._runnable_verbs and len(tokens) > 1 and not tokens[1].startswith(\"-\"):\n fused = f\"{tokens[0]} {tokens[1]}\"\n if fused in self._choices:\n return fused\n # `update bogus` is not the bare `update` command. Only a bare\n # command that takes an argument can legitimately own a second word.\n if not self.takes_positional(tokens[0]):\n return None\n return tokens[0] if tokens[0] in self._choices else None\n\n def suggest(self, tokens):\n \"\"\"Close command names for a line that resolves to nothing — 'did you mean'.\"\"\"\n if not tokens:\n return []\n head = tokens[0]\n if head in self.verbs: # known verb, missing or misspelled object\n objs = self.objects(head)\n near = difflib.get_close_matches(tokens[1], objs, n=3, cutoff=0.5) if len(tokens) > 1 else []\n return [f\"{head} {obj}\" for obj in (near or objs)]\n return (difflib.get_close_matches(\" \".join(tokens[:2]), self.names, n=3, cutoff=0.5)\n or difflib.get_close_matches(head, self.first_words, n=3, cutoff=0.5))\n\n def flags(self, name):\n \"\"\"``{\"--flag\": action}`` for one command, cached after the first lookup.\n\n Long forms only: they are what completion offers. Use ``option`` to look\n up a token the user actually typed, which may be a short alias.\n \"\"\"\n if name not in self._flag_cache:\n sub = self._choices.get(name)\n flags = {}\n if sub is not None:\n for action in sub._actions:\n for opt in action.option_strings:\n if opt.startswith(\"--\"):\n flags[opt] = action\n self._flag_cache[name] = flags\n return self._flag_cache[name]\n\n def option(self, name, token):\n \"\"\"The action a flag token names on one command, short form included.\"\"\"\n sub = self._choices.get(name)\n return option_action(sub, token) if sub is not None else None\n\n def global_option(self, token):\n \"\"\"The action a flag token names among the global options.\"\"\"\n return option_action(self._parser, token)\n\n def positionals(self, name):\n sub = self._choices.get(name)\n return [a for a in sub._actions if not a.option_strings] if sub is not None else []\n\n def takes_positional(self, name):\n return bool(self.positionals(name))\n\n def positional_completer(self, name, index=0):\n \"\"\"The completer for a command's index-th positional, if it has one.\n\n ``apwrap._add_completer`` tags id/machine/ssh positionals with a\n completer function; that is what turns ``destroy instance \u003ctab>`` into a\n list of live instance ids, and the second argument of\n ``update ssh-key \u003cid> \u003ctab>`` into local .pub paths instead.\n \"\"\"\n positionals = self.positionals(name)\n if not positionals:\n return None\n if index >= len(positionals):\n # A trailing nargs='+'/'*' positional keeps accepting values.\n last = positionals[-1]\n return getattr(last, \"completer\", None) if last.nargs in (\"+\", \"*\") else None\n return getattr(positionals[index], \"completer\", None)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":6453,"content_sha256":"c8f184de02554cf14b21cffe0f47354823832b259f3856578d96e0c53902bff6"},{"filename":"cli/repl/completion.py","content":"\"\"\"Tab completion for the REPL.\n\nInside our own readline loop the completer sees the whole line, so completion\ncan be staged the way the command tree actually is:\n\n sh\u003ctab> -> show, search, set, ssh-url, ... (verbs)\n show \u003ctab> -> instances, user, invoices, ... (objects)\n show instances --\u003ctab> -> --raw, --quiet, ... (flags)\n destroy instance \u003ctab> -> live instance ids (values)\n\nIt is also the fast path: `vastai` under argcomplete forks a fresh interpreter\nper Tab press, while here the parser is already in memory and ids are cached,\nso completion is a dict lookup.\n\"\"\"\nimport argparse\nimport threading\nimport time\n\nfrom vastai.cli.repl.catalog import CommandCatalog\n\nVALUE_CACHE_TTL = 30.0 # seconds a fetched id list stays warm\nCOMPLETION_BUDGET = 2.0 # seconds a Tab press will wait for that fetch\n\n\nclass LiveValues:\n \"\"\"Cache over the id completers the parser attaches to positionals.\n\n Each lookup would otherwise be an API round trip on a keystroke, so the\n full list is fetched once per completer and filtered locally. Failures are\n cached too — a down API or a missing key must not hang every Tab press.\n \"\"\"\n\n def __init__(self, ttl=VALUE_CACHE_TTL, clock=time.monotonic,\n budget=COMPLETION_BUDGET):\n self._ttl = ttl\n self._clock = clock\n self._budget = budget\n self._cache = {}\n self._pending = set()\n self._generation = 0\n\n def clear(self):\n \"\"\"Forget every cached list — called when the session's credentials\n change, so Tab can't offer the previous account's ids.\"\"\"\n self._cache.clear()\n self._pending.clear()\n self._generation += 1 # discard whatever is still in flight\n\n def matching(self, completer, prefix):\n cached = self._cache.get(completer)\n if cached is None or self._clock() - cached[0] > self._ttl:\n self._fetch(completer)\n cached = self._cache.get(completer)\n if cached is None:\n return [] # still in flight; a later Tab press will have it\n return [v for v in cached[1] if v.startswith(prefix)]\n\n def _fetch(self, completer):\n \"\"\"Fetch in a worker, and wait only a moment for it.\n\n The API client allows 120s per request and retries, so fetching on the\n input thread would let one Tab press freeze the prompt for minutes\n against an unreachable endpoint. Whatever the worker eventually returns\n lands in the cache for the next press.\n \"\"\"\n if completer in self._pending:\n return\n self._pending.add(completer)\n generation = self._generation\n\n def work():\n try:\n values = [str(v) for v in (completer(prefix=\"\") or [])]\n except Exception:\n values = []\n if generation != self._generation:\n return # the account changed while we were fetching\n self._cache[completer] = (self._clock(), values)\n self._pending.discard(completer)\n\n worker = threading.Thread(target=work, daemon=True)\n worker.start()\n worker.join(self._budget)\n\n\nclass ReplCompleter:\n \"\"\"A readline completer over the live command tree.\n\n ``suggestions`` is the whole of the logic and is pure with respect to\n readline (it takes the line, not the terminal state); ``complete`` is just\n the readline protocol adapter around it.\n \"\"\"\n\n def __init__(self, catalog, values=None):\n self.catalog = catalog if isinstance(catalog, CommandCatalog) else CommandCatalog(catalog)\n self.values = values if values is not None else LiveValues()\n self.matches = []\n\n # -- readline protocol -------------------------------------------------\n def complete(self, text, state):\n if state == 0:\n try:\n import readline\n buf = readline.get_line_buffer()[: readline.get_endidx()]\n except Exception:\n buf = text\n try:\n self.matches = self.suggestions(buf, text)\n except Exception:\n self.matches = [] # a broken completion must never break typing\n try:\n return self.matches[state]\n except IndexError:\n return None\n\n def install(self):\n \"\"\"Bind this completer to readline. No-op where readline is missing.\"\"\"\n try:\n import readline\n except ImportError:\n return False\n readline.set_completer(self.complete)\n # Split on whitespace only, so the completer sees '--flag' and\n # '--flag=value' as single tokens and decides what they mean.\n readline.set_completer_delims(\" \\t\\n\")\n if \"libedit\" in (getattr(readline, \"__doc__\", \"\") or \"\"):\n readline.parse_and_bind(\"bind ^I rl_complete\") # macOS libedit\n else:\n readline.parse_and_bind(\"tab: complete\")\n return True\n\n # -- the actual logic --------------------------------------------------\n def suggestions(self, buf, text=None):\n \"\"\"Candidates for the token being typed at the end of ``buf``.\"\"\"\n if text is None:\n text = \"\" if (not buf or buf[-1].isspace()) else buf.split()[-1]\n stripped = buf.lstrip()\n tokens = stripped.split()\n typed = tokens if (not stripped or buf[-1:].isspace()) else tokens[:-1]\n typed = self.catalog.strip_options(typed) # `--raw show \u003ctab>` completes too\n\n if not typed:\n return self._starting_with(self.catalog.first_words, text)\n\n if len(typed) == 1 and typed[0] in self.catalog.verbs and not text.startswith(\"-\"):\n return self._starting_with(self.catalog.objects(typed[0]), text)\n\n name = self.catalog.resolve(typed)\n if name is None:\n return []\n if text.startswith(\"-\"):\n return self._starting_with(self.catalog.flags(name), text)\n return self._values(name, typed[len(name.split()):], text)\n\n def _values(self, name, args, text):\n \"\"\"Complete an argument value: the choices or ids of whichever option or\n positional the cursor sits on.\"\"\"\n option, index = self._position(name, args)\n if option is not None:\n if option.choices:\n return self._starting_with([str(c) for c in option.choices], text)\n return self._live(getattr(option, \"completer\", None), text)\n return self._live(self.catalog.positional_completer(name, index), text)\n\n def _position(self, name, args):\n \"\"\"Walk a command's arguments as argparse would.\n\n Returns the option still awaiting values (the cursor is typing one of\n them) and the index of the positional being typed — so a second\n positional gets its own completer, and a multi-value option keeps\n offering its choices.\n \"\"\"\n option, index, i = None, 0, 0\n while i \u003c len(args):\n token = args[i]\n i += 1\n if token.startswith(\"-\") and token != \"-\":\n option = None\n action = self.catalog.option(name, token)\n if action is None or action.nargs == 0 or \"=\" in token:\n continue\n taken = 0\n while (i \u003c len(args) and _accepts_more(action, taken)\n and not args[i].startswith(\"-\")):\n taken += 1\n i += 1\n if _accepts_more(action, taken):\n option = action # still hungry: the cursor's word is its value\n continue\n option = None\n index += 1\n return option, index\n\n def _live(self, completer, text):\n return self.values.matching(completer, text) if completer else []\n\n @staticmethod\n def _starting_with(candidates, text):\n return sorted(c for c in candidates if c.startswith(text))\n\n\ndef _accepts_more(action, taken):\n \"\"\"Whether an option can still absorb another value after `taken` of them.\"\"\"\n nargs = action.nargs\n if nargs in (\"+\", \"*\", argparse.REMAINDER):\n return True\n if nargs == \"?\":\n return taken \u003c 1\n if isinstance(nargs, int):\n return taken \u003c nargs\n return taken \u003c 1 # nargs None: exactly one value\n","content_type":"text/x-python; charset=utf-8","language":"python","size":8389,"content_sha256":"c99464aed920e1399738620cf1fd4448fa967d9d7e109a2d6efcf83b7b7d8aa1"},{"filename":"cli/repl/session.py","content":"\"\"\"The `vastai repl` loop: the same CLI, one live process.\n\n vast> show instances\n vast> search offers 'gpu_name=RTX_4090 num_gpus=1' -o dph\n vast> destroy instance \u003ctab> # completes from your live instances\n\nEvery line is handed to the same parser `vastai` uses, so behaviour is\nidentical to the one-shot CLI — only faster, because the interpreter and the\ncommand tree load once. Type `help` for the commands, `clear` to clear the\nscreen, `exit` or Ctrl-D to leave.\n\"\"\"\nimport copy\nimport os\nimport shutil\nimport sys\nimport textwrap\n\nfrom vastai.cli.repl.bridge import run_line\nfrom vastai.cli.repl.catalog import CommandCatalog\nfrom vastai.cli.repl.completion import ReplCompleter\nfrom vastai.cli.util import APIKEY_FILE, TFAKEY_FILE\n\nPROMPT = \"vast> \"\nEXIT_WORDS = (\"exit\", \"quit\", \"q\")\nCLEAR_WORDS = (\"clear\", \"cls\")\nCLEAR_SCREEN = \"\\033[H\\033[2J\\033[3J\" # cursor home, clear screen, clear scrollback\n\n\nclass Repl:\n def __init__(self, parser, session_args, stdout=None):\n self.parser = parser\n # A copy: the session's own state is not the caller's.\n self.args = copy.copy(session_args)\n self.out = stdout or sys.stdout\n self._stored_key = _stored_api_key()\n # Whether this session's key came from the config file. A key from\n # $VAST_API_KEY or --api-key outranks the file (the CLI says so, and\n # `set api-key` warns about it), so it must not be replaced by one.\n self._from_disk = getattr(session_args, \"api_key\", None) in (None, self._stored_key)\n self.catalog = CommandCatalog(parser)\n self.completer = ReplCompleter(self.catalog)\n\n # -- one line ----------------------------------------------------------\n def handle(self, line):\n \"\"\"Run one input line. Returns False when the session should end.\"\"\"\n s = line.strip()\n if not s:\n return True\n if s in EXIT_WORDS:\n return False\n if s in CLEAR_WORDS:\n self._clear_screen()\n return True\n\n # Leading global flags (`--raw show user`) are the parser's business,\n # not part of the command name.\n tokens = self.catalog.strip_options(s.split())\n if tokens and tokens[0] == \"repl\":\n self._print(\"Already in the REPL. Use exit or Ctrl-D to leave.\")\n return True\n # Resolve first so a typo gets a short 'did you mean' instead of\n # argparse dumping all ~150 command names. A line that is only flags\n # (`--version`) names no command and goes straight to the parser.\n if tokens and self.catalog.resolve(tokens) is None:\n self._unknown(tokens)\n return True\n run_line(self.parser, s, self.args)\n self._refresh_credentials()\n return True\n\n def _unknown(self, tokens):\n \"\"\"Report an unrecognised line. argparse would answer with all ~150\n command names; a verb's own objects, or a few near misses, is the\n answer worth reading.\"\"\"\n if len(tokens) == 1 and tokens[0] in self.catalog.verbs:\n self._print_wrapped(f\"{tokens[0]} takes an object: \",\n self.catalog.objects(tokens[0]))\n return\n print(f\"unknown command: {' '.join(tokens[:2])}\", file=sys.stderr)\n suggestions = self.catalog.suggest(tokens)[:6]\n if suggestions:\n self._print_wrapped(\"did you mean: \", suggestions)\n print(\"(help lists every command)\", file=sys.stderr)\n\n def _print_wrapped(self, lead, items):\n width = shutil.get_terminal_size(fallback=(80, 24)).columns\n print(textwrap.fill(\", \".join(items), width=max(40, width),\n initial_indent=lead, subsequent_indent=\" \" * len(lead)),\n file=sys.stderr)\n\n def _clear_screen(self):\n # Legacy Windows consoles don't honour ANSI escapes; cls always works there.\n if os.name == \"nt\":\n os.system(\"cls\")\n else:\n self.out.write(CLEAR_SCREEN)\n self.out.flush()\n\n def _print(self, text):\n print(text, file=self.out)\n\n def _refresh_credentials(self):\n \"\"\"Pick up a key written mid-session by `set api-key` or `tfa login`.\n\n Those commands write the config file but not our namespace, so without\n this the session would keep sending the key it started with — often\n none at all, leaving every later line to fail on auth. Only a session\n that took its key from the file follows the file; one running on\n $VAST_API_KEY or --api-key keeps what it was given, as the one-shot CLI\n would. A key that disappears is dropped rather than reused: an expired\n 2FA session with no fallback should report \"no API key\", not repeat the\n same failure on every line.\n \"\"\"\n if not self._from_disk:\n return\n key = _stored_api_key()\n if key == self._stored_key:\n return\n self.args.api_key = key\n self._stored_key = key\n self.completer.values.clear() # ids belong to the old account\n\n # -- the loop ----------------------------------------------------------\n def run(self):\n self._setup_terminal()\n while True:\n try:\n line = input(PROMPT)\n except KeyboardInterrupt:\n self._print(\"^C\")\n continue\n except EOFError:\n self._print(\"\")\n break\n if not self.handle(line):\n break\n\n def _setup_terminal(self):\n \"\"\"Wire completion. History lives in readline for the session only —\n persisting it would put whatever was typed, credentials included, in\n plain text on disk, which wants a redaction design of its own.\"\"\"\n self._wire_live_completions()\n self.completer.install()\n\n def _wire_live_completions(self):\n \"\"\"Point the parser's id completers at this session.\n\n The one-shot CLI wires these for argcomplete using only $VAST_API_KEY;\n rebinding them here means `destroy instance \u003ctab>` works for a key that\n came from the config file or from `--api-key` on the repl command.\n \"\"\"\n from vastai.cli.commands import instances\n from vastai.cli.parser import set_completers\n\n def ids(field):\n return lambda **kw: instances.show__instances(\n self._completion_args(), {\"internal\": True, \"field\": field})\n\n set_completers(instance_machine_fn=ids(\"machine_id\"), instance_fn=ids(\"id\"))\n\n def _completion_args(self):\n \"\"\"Args for completion's own API calls: this session's auth, endpoint\n and retries, with the output modes off. `--curl` in particular makes\n the client print a curl command and exit, which would otherwise land in\n the middle of the prompt on every Tab press.\n \"\"\"\n args = copy.copy(self.args)\n for flag in (\"explain\", \"curl\", \"raw\", \"full\"):\n setattr(args, flag, False)\n args.retry = 1 # a Tab press must not queue minutes of retries\n return args\n\n\ndef _stored_api_key():\n \"\"\"The key on disk right now, resolved as the one-shot CLI resolves it.\"\"\"\n path = TFAKEY_FILE if os.path.exists(TFAKEY_FILE) else APIKEY_FILE\n try:\n with open(path, \"r\") as reader:\n return reader.read().strip() or None\n except OSError:\n return None\n\n\ndef run_repl(args):\n from vastai.cli.main import parser\n Repl(parser, args).run()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":7524,"content_sha256":"b047f859be2b1e1b6ebaa8c8289bb26af159c7e460c732c60a5ef31a1aeed646"},{"filename":"cli/self_test/__init__.py","content":"\"\"\"Helpers for CLI self-test diagnostics.\"\"\"\n\n","content_type":"text/x-python; charset=utf-8","language":"python","size":46,"content_sha256":"2712a98324d097071d87a9b7b79b38bb317eea3d2dcc07389fc7d1b98fddb9ec"},{"filename":"cli/self_test/machine_diagnostics.py","content":"\"\"\"Structured diagnostics for ``vastai self-test machine``.\"\"\"\n\nfrom vastai.cli.util import required_inet_mbps\n\n\nDIAGNOSTICS_VERSION = \"1\"\nROOT_CURRENTLY_RENTED = \"currently_rented\"\nROOT_DEVERIFIED_OR_BELOW_THRESHOLD = \"deverified_or_below_threshold\"\nROOT_API_PERMISSION_FAILED = \"api_permission_failed\"\nROOT_ZERO_ACTIVE_OFFERS = \"zero_active_offers\"\nROOT_OFFLINE_OR_NOT_LISTED = \"offline_or_not_listed\"\nROOT_UNKNOWN_NO_RENTABLE_OFFER = \"unknown_no_rentable_offer\"\n\nNO_OFFER_ROOT_STATES = (\n ROOT_CURRENTLY_RENTED,\n ROOT_DEVERIFIED_OR_BELOW_THRESHOLD,\n ROOT_API_PERMISSION_FAILED,\n ROOT_ZERO_ACTIVE_OFFERS,\n ROOT_OFFLINE_OR_NOT_LISTED,\n ROOT_UNKNOWN_NO_RENTABLE_OFFER,\n)\n\n# Very large GPU hosts, such as 8x B300 systems, can exceed 2 TB of total VRAM.\n# Once a host has about 2 TB of system RAM, do not disqualify it only because it\n# is slightly below 95% of total VRAM.\nSYSTEM_RAM_REQUIREMENT_CAP_MIB = 2_000_000\n\n\ndef safe_float(value):\n if value is None:\n return 0.0\n try:\n return float(value)\n except (ValueError, TypeError):\n return 0.0\n\n\ndef per_gpu_vram_gib(offer):\n \"\"\"Return per-GPU VRAM in GiB from canonical offer fields when possible.\"\"\"\n gpu_total_ram = safe_float(offer.get(\"gpu_total_ram\"))\n num_gpus = safe_float(offer.get(\"num_gpus\"))\n if gpu_total_ram > 0 and num_gpus > 0:\n return gpu_total_ram / num_gpus / 1024\n\n gpu_ram = safe_float(offer.get(\"gpu_ram\"))\n if gpu_ram >= 1024:\n return gpu_ram / 1024\n return gpu_ram\n\n\ndef compact_offer_metadata(offer):\n if not offer:\n return None\n fields = (\n \"id\",\n \"machine_id\",\n \"gpu_name\",\n \"num_gpus\",\n \"dph_total\",\n \"dlperf\",\n \"cuda_max_good\",\n \"compute_cap\",\n \"reliability\",\n \"direct_port_count\",\n \"pcie_bw\",\n \"inet_down\",\n \"inet_up\",\n \"gpu_ram\",\n \"gpu_total_ram\",\n \"cpu_ram\",\n \"cpu_cores\",\n \"rentable\",\n \"rented\",\n \"vericode\",\n \"verified\",\n \"verification\",\n \"error_description\",\n )\n return {field: offer.get(field) for field in fields if field in offer}\n\n\ndef base_result(machine_id):\n return {\n \"success\": False,\n \"reason\": \"\",\n \"warning\": None,\n \"phase\": \"preflight\",\n \"diagnostics_version\": DIAGNOSTICS_VERSION,\n \"machine_id\": machine_id,\n \"offer\": None,\n \"checks\": [],\n \"failure\": None,\n \"failure_code\": None,\n \"stage\": None,\n \"error\": None,\n \"diagnostics\": {},\n }\n\n\ndef _check(\n check_id,\n title,\n actual,\n required,\n operator,\n unit,\n passed,\n purpose,\n remediation,\n):\n status = \"pass\" if passed else \"fail\"\n return {\n \"id\": check_id,\n \"title\": title,\n \"status\": status,\n \"actual\": actual,\n \"required\": required,\n \"operator\": operator,\n \"unit\": unit,\n \"summary\": f\"{title}: actual {actual} {unit}, required {operator} {required} {unit}\",\n \"purpose\": purpose,\n \"remediation\": remediation,\n }\n\n\ndef _info_check(\n check_id,\n title,\n actual,\n recommended,\n operator,\n unit,\n summary,\n purpose,\n remediation,\n):\n return {\n \"id\": check_id,\n \"title\": title,\n \"status\": \"info\",\n \"actual\": actual,\n \"required\": recommended,\n \"operator\": operator,\n \"unit\": unit,\n \"summary\": summary,\n \"purpose\": purpose,\n \"remediation\": remediation,\n }\n\n\ndef preflight_requirement_checks(offer):\n gpu_total_ram = safe_float(offer.get(\"gpu_total_ram\"))\n per_gpu_ram_gib = per_gpu_vram_gib(offer)\n required_mbps = required_inet_mbps(gpu_total_ram)\n cpu_ram = safe_float(offer.get(\"cpu_ram\"))\n cpu_cores = int(safe_float(offer.get(\"cpu_cores\")))\n num_gpus = int(safe_float(offer.get(\"num_gpus\")))\n listed_gpus = num_gpus if num_gpus > 0 else 1\n direct_port_count = safe_float(offer.get(\"direct_port_count\"))\n required_min_ports = 3 * listed_gpus\n recommended_max_ports = 64 * listed_gpus\n uncapped_required_cpu_ram = 0.95 * gpu_total_ram\n required_cpu_ram = min(uncapped_required_cpu_ram, SYSTEM_RAM_REQUIREMENT_CAP_MIB)\n required_cpu_cores = 2 * num_gpus\n\n checks = [\n _check(\n \"cuda.version\",\n \"CUDA version\",\n safe_float(offer.get(\"cuda_max_good\")),\n 11.8,\n \">=\",\n \"CUDA\",\n safe_float(offer.get(\"cuda_max_good\")) >= 11.8,\n \"The self-test image needs a host driver stack compatible with CUDA 11.8 or newer.\",\n \"Update the NVIDIA driver/CUDA stack, then confirm with: vastai search offers \"\n f\"'machine_id={offer.get('machine_id') or 'unknown'} rentable=any rented=any'\",\n ),\n _check(\n \"reliability\",\n \"Reliability\",\n safe_float(offer.get(\"reliability\")),\n 0.90,\n \">\",\n \"ratio\",\n safe_float(offer.get(\"reliability\")) > 0.90,\n \"Self-test uses reliability as a guardrail before launching a temporary instance.\",\n \"Let the host stabilize and resolve recent failures before retrying the self-test.\",\n ),\n _check(\n \"network.direct_ports\",\n \"Direct port count\",\n direct_port_count,\n required_min_ports,\n \">=\",\n \"ports\",\n direct_port_count >= required_min_ports,\n (\n \"The tester needs at least 3 directly mapped ports per listed GPU for remote \"\n \"progress, SSH checks, and runtime port allocation.\"\n ),\n (\n f\"Open at least {required_min_ports} direct ports for {listed_gpus} listed GPU(s), \"\n \"check firewall/NAT forwarding, and make TCP/UDP forwarding symmetric where required.\"\n ),\n ),\n _check(\n \"pcie.bandwidth\",\n \"PCIe bandwidth\",\n safe_float(offer.get(\"pcie_bw\")),\n 2.85,\n \">\",\n \"GB/s\",\n safe_float(offer.get(\"pcie_bw\")) > 2.85,\n \"Low PCIe bandwidth can make GPU stress and transfer checks fail or time out.\",\n \"Check BIOS PCIe generation/lane settings and confirm GPUs are seated in full-speed slots.\",\n ),\n _check(\n \"network.download\",\n \"Download speed\",\n safe_float(offer.get(\"inet_down\")),\n round(required_mbps, 2),\n \">=\",\n \"Mb/s\",\n safe_float(offer.get(\"inet_down\")) >= required_mbps,\n \"The bandwidth floor scales with total VRAM so large GPU hosts can complete data movement tests.\",\n \"Improve host download bandwidth or reduce contention, then rerun the Vast host verification.\",\n ),\n _check(\n \"network.upload\",\n \"Upload speed\",\n safe_float(offer.get(\"inet_up\")),\n round(required_mbps, 2),\n \">=\",\n \"Mb/s\",\n safe_float(offer.get(\"inet_up\")) >= required_mbps,\n \"The tester needs enough upload bandwidth to report progress and complete network checks.\",\n \"Improve host upload bandwidth or reduce contention, then rerun the Vast host verification.\",\n ),\n _check(\n \"gpu.ram\",\n \"GPU RAM\",\n round(per_gpu_ram_gib, 2),\n 7,\n \">\",\n \"GiB\",\n per_gpu_ram_gib > 7,\n \"The verification workload requires more than 7 GB of VRAM per GPU.\",\n \"Use a GPU with more VRAM for this self-test.\",\n ),\n _check(\n \"system.ram\",\n \"System RAM\",\n cpu_ram,\n round(required_cpu_ram, 2),\n \">=\",\n \"MiB\",\n cpu_ram >= required_cpu_ram,\n (\n \"System RAM must be close to total VRAM so CPU-side staging does not starve the tests. \"\n \"For very large GPU hosts, this requirement is capped at about 2 TB.\"\n ),\n (\n \"Add system RAM or reduce the listed GPU set so system RAM is at least 95% of \"\n \"total VRAM, up to the 2 TB cap.\"\n ),\n ),\n _check(\n \"cpu.cores\",\n \"CPU cores\",\n cpu_cores,\n required_cpu_cores,\n \">=\",\n \"cores\",\n cpu_cores >= required_cpu_cores,\n \"The tester expects at least two CPU cores per GPU for stable orchestration.\",\n \"Expose more CPU cores to the host or reduce the GPU count for this offer.\",\n ),\n ]\n if direct_port_count > recommended_max_ports:\n checks.append(\n _info_check(\n \"network.direct_ports.recommended_max\",\n \"Direct port count advisory\",\n direct_port_count,\n recommended_max_ports,\n \"\u003c=\",\n \"ports\",\n (\n f\"Direct port count: actual {direct_port_count} ports, recommended \u003c= \"\n f\"{recommended_max_ports} ports for {listed_gpus} GPU(s)\"\n ),\n (\n \"Vast instances can use at most 64 open ports each. Mapping more than \"\n \"64 ports per listed GPU is usually wasted effort.\"\n ),\n (\n \"This is advisory only, not a self-test gate. Keep enough direct ports for \"\n \"self-test and normal workloads, but avoid mapping very large unused ranges.\"\n ),\n )\n )\n return checks\n\n\ndef failed_checks(checks):\n return [check for check in checks if check.get(\"status\") == \"fail\"]\n\n\ndef informational_checks(checks):\n return [check for check in checks if check.get(\"status\") == \"info\"]\n\n\ndef _http_status(error):\n if not error:\n return None\n if isinstance(error, dict):\n return error.get(\"status_code\")\n response = getattr(error, \"response\", None)\n return getattr(response, \"status_code\", None)\n\n\ndef _machine_lookup_status(machine_lookup):\n if not machine_lookup:\n return None, []\n if isinstance(machine_lookup, dict) and \"status\" in machine_lookup:\n status = machine_lookup.get(\"status\")\n rows = machine_lookup.get(\"rows\") or []\n if isinstance(rows, dict):\n rows = [rows] if rows else []\n return status, rows\n if isinstance(machine_lookup, list):\n return (\"visible\", machine_lookup) if machine_lookup else (\"empty\", [])\n if isinstance(machine_lookup, dict):\n return (\"visible\", [machine_lookup]) if machine_lookup else (\"empty\", [])\n return None, []\n\n\ndef diagnose_no_offer_state(machine_id, broader_offers, machine_lookup=None, search_error=None):\n \"\"\"Infer the most likely no-rentable root state from CLI-visible data.\n\n This is intentionally heuristic. Public offer search alone cannot prove\n whether a machine is offline, hidden, or has zero active offers, so the\n return value includes a confidence field for UI/host copy.\n \"\"\"\n offers = broader_offers or []\n evidence = []\n suggested_steps = [\n f\"Inspect all visible offers with: vastai search offers 'machine_id={machine_id} rentable=any rented=any'.\"\n ]\n\n search_status = _http_status(search_error)\n lookup_status, lookup_rows = _machine_lookup_status(machine_lookup)\n lookup_status_code = machine_lookup.get(\"status_code\") if isinstance(machine_lookup, dict) else None\n\n if search_status in (401, 403) or lookup_status in (\"permission_denied\", \"unauthorized\") or lookup_status_code in (401, 403):\n if search_status:\n evidence.append(f\"Offer search failed with HTTP {search_status}.\")\n if lookup_status_code:\n evidence.append(f\"Machine lookup failed with HTTP {lookup_status_code}.\")\n evidence.append(\"The API key or account could not read the required machine/offer state.\")\n suggested_steps = [\n \"Check that the API key is valid and has permission to read offers and hosted machines.\",\n \"Retry after logging in with an account that owns or can inspect this host.\",\n ]\n return {\n \"root_state\": ROOT_API_PERMISSION_FAILED,\n \"confidence\": \"high\",\n \"evidence\": evidence,\n \"suggested_steps\": suggested_steps,\n }\n\n if offers:\n evidence.append(f\"Found {len(offers)} visible offer(s), but none were currently rentable.\")\n rented_offers = [offer for offer in offers if offer.get(\"rented\") is True]\n if rented_offers:\n evidence.append(f\"{len(rented_offers)} visible offer(s) are marked rented=true.\")\n suggested_steps = [\n \"Wait for the current rental to end, then retry the self-test.\",\n suggested_steps[0],\n ]\n return {\n \"root_state\": ROOT_CURRENTLY_RENTED,\n \"confidence\": \"medium\",\n \"evidence\": evidence,\n \"suggested_steps\": suggested_steps,\n }\n\n state_evidence = []\n for offer in offers:\n reliability = safe_float(offer.get(\"reliability\"))\n vericode = offer.get(\"vericode\")\n verified = offer.get(\"verified\")\n verification = offer.get(\"verification\")\n error_description = offer.get(\"error_description\")\n if reliability and reliability \u003c= 0.90:\n state_evidence.append(f\"Offer reliability is {reliability}, which is at or below the 0.90 self-test threshold.\")\n if vericode not in (None, \"\", 0, \"0\"):\n state_evidence.append(f\"Offer reports vericode={vericode}.\")\n if verified is False:\n state_evidence.append(\"Offer is marked verified=false.\")\n if isinstance(verification, str) and verification.lower() not in (\"\", \"verified\", \"true\"):\n state_evidence.append(f\"Offer verification state is {verification}.\")\n if error_description:\n state_evidence.append(f\"Offer error_description is {error_description}.\")\n\n if state_evidence:\n evidence.extend(state_evidence)\n suggested_steps = [\n \"Resolve the host verification or reliability issue, then wait for the offer state to refresh.\",\n suggested_steps[0],\n ]\n return {\n \"root_state\": ROOT_DEVERIFIED_OR_BELOW_THRESHOLD,\n \"confidence\": \"medium\",\n \"evidence\": evidence,\n \"suggested_steps\": suggested_steps,\n }\n\n evidence.append(\"Visible offers exist, but their payload did not expose a specific non-rentable reason.\")\n suggested_steps = [\n \"Refresh the marketplace/host state and retry shortly.\",\n suggested_steps[0],\n ]\n return {\n \"root_state\": ROOT_UNKNOWN_NO_RENTABLE_OFFER,\n \"confidence\": \"low\",\n \"evidence\": evidence,\n \"suggested_steps\": suggested_steps,\n }\n\n evidence.append(\"No visible on-demand offers were returned for this machine.\")\n if lookup_status == \"visible\" and lookup_rows:\n evidence.append(\"Machine lookup returned a visible machine record.\")\n suggested_steps = [\n \"List or relist an on-demand offer for this machine, then retry self-test.\",\n \"Check host listing state with: vastai show machines.\",\n ]\n return {\n \"root_state\": ROOT_ZERO_ACTIVE_OFFERS,\n \"confidence\": \"medium\",\n \"evidence\": evidence,\n \"suggested_steps\": suggested_steps,\n }\n\n if lookup_status in (\"empty\", \"not_found\") or lookup_status_code == 404:\n if lookup_status_code == 404:\n evidence.append(\"Machine lookup returned HTTP 404.\")\n else:\n evidence.append(\"Machine lookup returned no visible machine record.\")\n elif lookup_status:\n evidence.append(f\"Machine lookup status was {lookup_status}.\")\n else:\n evidence.append(\"No authoritative machine lookup result was available.\")\n\n suggested_steps = [\n \"Confirm the machine ID is correct and the host is online.\",\n \"If this is your host, check that it is listed and visible with: vastai show machines.\",\n suggested_steps[0],\n ]\n return {\n \"root_state\": ROOT_OFFLINE_OR_NOT_LISTED,\n \"confidence\": \"low\",\n \"evidence\": evidence,\n \"suggested_steps\": suggested_steps,\n }\n\n\ndef no_offer_failure(machine_id, broader_offers, machine_lookup=None, search_error=None):\n diagnosis = diagnose_no_offer_state(\n machine_id,\n broader_offers,\n machine_lookup=machine_lookup,\n search_error=search_error,\n )\n broad_count = len(broader_offers or [])\n if diagnosis[\"root_state\"] == ROOT_API_PERMISSION_FAILED:\n summary = f\"Could not inspect rentable offers for machine {machine_id} because API authorization failed.\"\n causes = [\n \"The API key may be invalid, expired, or missing permission to read offer/machine state.\",\n ]\n code = \"api_permission_failed\"\n remediation = \"Check the API key permissions, then retry with an account that can inspect this host.\"\n elif broad_count:\n summary = f\"No currently rentable on-demand offer found for machine {machine_id}.\"\n causes = [\n \"The machine may be listed but not rentable right now.\",\n \"The offer may be unavailable to new rentals because host state, verification, or pricing changed.\",\n \"The machine may already be occupied or temporarily unavailable.\",\n ]\n code = \"no_rentable_offer\"\n remediation = (\n \"Check host state with: vastai show machines. Then inspect offers with: \"\n f\"vastai search offers 'machine_id={machine_id} rentable=any rented=any'.\"\n )\n else:\n summary = f\"No on-demand offer found for machine {machine_id}.\"\n causes = [\n \"The machine may be offline.\",\n \"The machine may not have an active listed offer.\",\n \"The machine ID may be incorrect or not visible to this account.\",\n ]\n code = \"no_offer\"\n remediation = (\n \"Check host state with: vastai show machines. Then inspect offers with: \"\n f\"vastai search offers 'machine_id={machine_id} rentable=any rented=any'.\"\n )\n check = {\n \"id\": \"offer.available\",\n \"title\": \"Rentable offer available\",\n \"status\": \"fail\",\n \"actual\": 0,\n \"required\": 1,\n \"operator\": \">=\",\n \"unit\": \"offers\",\n \"summary\": summary,\n \"purpose\": \"The self-test needs a selected offer so it can rent one temporary diagnostic instance.\",\n \"remediation\": remediation,\n }\n failure = {\n \"code\": code,\n \"summary\": summary,\n \"likely_causes\": causes,\n \"remediation\": remediation,\n \"root_state\": diagnosis[\"root_state\"],\n \"confidence\": diagnosis[\"confidence\"],\n \"evidence\": diagnosis[\"evidence\"],\n \"suggested_steps\": diagnosis[\"suggested_steps\"],\n }\n return check, failure\n\n\ndef requirement_failure(checks):\n failures = failed_checks(checks)\n summary = f\"{len(failures)} preflight requirement check(s) failed.\"\n return {\n \"code\": \"preflight_requirements_failed\",\n \"summary\": summary,\n \"failed_check_ids\": [check[\"id\"] for check in failures],\n \"remediation\": \"Resolve the failed checks below, or rerun with --ignore-requirements to dogfood anyway.\",\n }\n\n\ndef render_preflight_failure(machine_id, checks, failure=None, print_fn=print):\n print_fn(f\"Preflight diagnostics for machine {machine_id} failed:\")\n for check in failed_checks(checks):\n print_fn(f\"- {check['title']}\")\n print_fn(f\" actual: {check['actual']} {check['unit']}\")\n print_fn(f\" required: {check['operator']} {check['required']} {check['unit']}\")\n print_fn(f\" purpose: {check['purpose']}\")\n print_fn(f\" remediation: {check['remediation']}\")\n if failure and failure.get(\"likely_causes\"):\n print_fn(\"Likely causes:\")\n for cause in failure[\"likely_causes\"]:\n print_fn(f\"- {cause}\")\n if failure and failure.get(\"root_state\"):\n print_fn(f\"Root state: {failure['root_state']} (confidence: {failure.get('confidence', 'unknown')})\")\n evidence = failure.get(\"evidence\") if failure else None\n if evidence:\n print_fn(\"Evidence:\")\n for item in evidence:\n print_fn(f\"- {item}\")\n steps = failure.get(\"suggested_steps\") if failure else None\n if steps:\n print_fn(\"Suggested steps:\")\n for step in steps:\n print_fn(f\"- {step}\")\n\n\ndef render_preflight_advisories(machine_id, checks, print_fn=print):\n info_checks = informational_checks(checks)\n if not info_checks:\n return\n print_fn(f\"Preflight advisory for machine {machine_id}:\")\n for check in info_checks:\n print_fn(f\"- {check['title']}\")\n print_fn(f\" actual: {check['actual']} {check['unit']}\")\n print_fn(f\" recommended: {check['operator']} {check['required']} {check['unit']}\")\n print_fn(f\" purpose: {check['purpose']}\")\n print_fn(f\" note: {check['remediation']}\")\n","content_type":"text/x-python; charset=utf-8","language":"python","size":21506,"content_sha256":"9cd1ddbdffbc97d86e0ecf00b3e166bcbffc8d0f5d9daf97cc561eda61e72152"},{"filename":"cli/self_test/runtime_diagnostics.py","content":"\"\"\"Runtime diagnostic helpers for CLI self-test machine flows.\n\nThe functions in this module are intentionally pure so the CLI can use them for\nraw output shaping without coupling parser/classifier behavior to live Vast API\ncalls or the existing self-test orchestration.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass\nimport re\nfrom typing import Iterable\n\n\nINSTANCE_CREATE_FAILED = \"instance_create_failed\"\nINSTANCE_CREATE_MISSING_CONTRACT = \"instance_create_missing_contract\"\nINSTANCE_STATUS_ERROR = \"instance_status_error\"\nINSTANCE_STATUS_POLL_FAILED = \"instance_status_poll_failed\"\nINSTANCE_START_TIMEOUT = \"instance_start_timeout\"\nINSTANCE_OFFLINE_BEFORE_TEST = \"instance_offline_before_test\"\nMISSING_PUBLIC_IP = \"missing_public_ip\"\nPROGRESS_PORT_NOT_MAPPED = \"progress_port_not_mapped\"\nPROGRESS_ENDPOINT_UNREACHABLE = \"progress_endpoint_unreachable\"\nPROGRESS_ENDPOINT_LOST = \"progress_endpoint_lost\"\nPROGRESS_EMPTY_TIMEOUT = \"progress_empty_timeout\"\nRUNTIME_TEST_TIMEOUT = \"runtime_test_timeout\"\nLEGACY_PROGRESS_ERROR = \"legacy_progress_error\"\nDOCKER_PULL_FAILED = \"docker_pull_failed\"\nDAEMON_STARTUP_FAILED = \"daemon_startup_failed\"\nNVML_FAILED = \"nvml_failed\"\nRESNET_FAILED = \"resnet_failed\"\nECC_FAILED = \"ecc_failed\"\nNCCL_FAILED = \"nccl_failed\"\nSTRESS_GPU_BURN_FAILED = \"stress_gpu_burn_failed\"\nINTERRUPTED = \"interrupted\"\nCLEANUP_FAILED = \"cleanup_failed\"\nPROGRESS_CONTAINER_PORT = \"5000/tcp\"\n\n\nRUNTIME_FAILURE_CODES = (\n INSTANCE_CREATE_FAILED,\n INSTANCE_CREATE_MISSING_CONTRACT,\n INSTANCE_STATUS_ERROR,\n INSTANCE_STATUS_POLL_FAILED,\n INSTANCE_START_TIMEOUT,\n INSTANCE_OFFLINE_BEFORE_TEST,\n MISSING_PUBLIC_IP,\n PROGRESS_PORT_NOT_MAPPED,\n PROGRESS_ENDPOINT_UNREACHABLE,\n PROGRESS_ENDPOINT_LOST,\n PROGRESS_EMPTY_TIMEOUT,\n RUNTIME_TEST_TIMEOUT,\n LEGACY_PROGRESS_ERROR,\n DOCKER_PULL_FAILED,\n DAEMON_STARTUP_FAILED,\n NVML_FAILED,\n RESNET_FAILED,\n ECC_FAILED,\n NCCL_FAILED,\n STRESS_GPU_BURN_FAILED,\n INTERRUPTED,\n CLEANUP_FAILED,\n)\n\n\n@dataclass(frozen=True)\nclass FailureCatalogEntry:\n code: str\n summary: str\n remediation: str\n suggested_steps: tuple[str, ...]\n\n\nFAILURE_CATALOG: dict[str, FailureCatalogEntry] = {\n INSTANCE_CREATE_FAILED: FailureCatalogEntry(\n INSTANCE_CREATE_FAILED,\n \"Failed to create the runtime test instance.\",\n \"Check the offer, docker image, and instance creation response.\",\n (\"Retry with --debugging enabled.\", \"Inspect the create-instance API error.\"),\n ),\n INSTANCE_CREATE_MISSING_CONTRACT: FailureCatalogEntry(\n INSTANCE_CREATE_MISSING_CONTRACT,\n \"Instance creation did not return a new contract id.\",\n \"Treat the create response as malformed or incomplete.\",\n (\"Inspect the raw create-instance response.\", \"Retry after confirming the offer is still rentable.\"),\n ),\n INSTANCE_STATUS_ERROR: FailureCatalogEntry(\n INSTANCE_STATUS_ERROR,\n \"The instance reported an error while starting.\",\n \"Inspect the instance status message and host/container logs.\",\n (\"Run show instance for the contract.\", \"Check docker logs on the host if available.\"),\n ),\n INSTANCE_STATUS_POLL_FAILED: FailureCatalogEntry(\n INSTANCE_STATUS_POLL_FAILED,\n \"Failed to poll instance status.\",\n \"Confirm API connectivity and retry the status check.\",\n (\"Retry the CLI command.\", \"Check network/API errors from the status request.\"),\n ),\n INSTANCE_START_TIMEOUT: FailureCatalogEntry(\n INSTANCE_START_TIMEOUT,\n \"The instance did not reach running state before timeout.\",\n \"Check host capacity, docker startup, and network configuration.\",\n (\"Inspect instance status_msg.\", \"Try the test again after confirming the host is healthy.\"),\n ),\n INSTANCE_OFFLINE_BEFORE_TEST: FailureCatalogEntry(\n INSTANCE_OFFLINE_BEFORE_TEST,\n \"The instance went offline before the runtime test could run.\",\n \"Investigate host availability and instance lifecycle events.\",\n (\"Check machine status.\", \"Review host daemon and container logs.\"),\n ),\n MISSING_PUBLIC_IP: FailureCatalogEntry(\n MISSING_PUBLIC_IP,\n \"The running instance did not expose a public IP address.\",\n \"Confirm the instance network configuration and public IP assignment.\",\n (\"Inspect show instance output.\", \"Retry on a machine with public networking available.\"),\n ),\n PROGRESS_PORT_NOT_MAPPED: FailureCatalogEntry(\n PROGRESS_PORT_NOT_MAPPED,\n \"The runtime progress port was not mapped.\",\n \"Confirm port 5000/tcp is exposed and direct ports are available.\",\n (\"Check the available mapped ports in the diagnostic output.\", \"Verify the machine has enough direct ports.\"),\n ),\n PROGRESS_ENDPOINT_UNREACHABLE: FailureCatalogEntry(\n PROGRESS_ENDPOINT_UNREACHABLE,\n \"The runtime progress endpoint was never reachable.\",\n \"Check TCP firewall/NAT forwarding, direct port mapping, container startup, and NAT hairpinning.\",\n (\n \"Confirm the mapped public TCP port forwards to the host LAN IP.\",\n \"Inspect docker logs to confirm the progress server bound port 5000/tcp.\",\n \"If testing from the same LAN as the host, retry from an outside network to rule out NAT loopback/hairpinning.\",\n ),\n ),\n PROGRESS_ENDPOINT_LOST: FailureCatalogEntry(\n PROGRESS_ENDPOINT_LOST,\n \"The runtime progress endpoint became unreachable after connecting.\",\n \"Look for container crashes, OOM, GPU errors, or host instability.\",\n (\n \"Inspect docker logs for a crash or missing progress server.\",\n \"Check dmesg for Xid, GPU reset, OOM, or host stall messages.\",\n \"Check for network loss between the CLI and host public endpoint.\",\n ),\n ),\n PROGRESS_EMPTY_TIMEOUT: FailureCatalogEntry(\n PROGRESS_EMPTY_TIMEOUT,\n \"The progress endpoint returned no new output before timeout.\",\n \"Check whether the runtime script stalled or stopped writing progress.\",\n (\"Inspect runtime logs.\", \"Retry with debugging enabled.\"),\n ),\n RUNTIME_TEST_TIMEOUT: FailureCatalogEntry(\n RUNTIME_TEST_TIMEOUT,\n \"The runtime test did not complete before timeout.\",\n \"Investigate long-running or stalled test stages.\",\n (\"Check the last reported progress stage.\", \"Run individual tests to isolate the stall.\"),\n ),\n LEGACY_PROGRESS_ERROR: FailureCatalogEntry(\n LEGACY_PROGRESS_ERROR,\n \"The legacy runtime progress stream reported an unclassified error.\",\n \"Use the raw error text and active stage to decide the next action.\",\n (\"Inspect the original ERROR line.\", \"Retry with debugging enabled if the cause is unclear.\"),\n ),\n DOCKER_PULL_FAILED: FailureCatalogEntry(\n DOCKER_PULL_FAILED,\n \"The test image could not be pulled.\",\n \"Check image name, tag availability, registry access, and credentials.\",\n (\"Verify the docker image tag exists.\", \"Check for registry unauthorized or denied errors.\"),\n ),\n DAEMON_STARTUP_FAILED: FailureCatalogEntry(\n DAEMON_STARTUP_FAILED,\n \"The container or daemon failed during startup.\",\n \"Inspect docker daemon, OCI runtime, and container startup logs.\",\n (\"Check docker logs.\", \"Verify NVIDIA container runtime and host daemon health.\"),\n ),\n NVML_FAILED: FailureCatalogEntry(\n NVML_FAILED,\n \"NVML or nvidia-smi failed during system checks.\",\n \"Check NVIDIA driver, NVML, and GPU visibility on the host.\",\n (\"Run nvidia-smi on the host.\", \"Check driver/library mismatch or GPU reset errors.\"),\n ),\n RESNET_FAILED: FailureCatalogEntry(\n RESNET_FAILED,\n \"The ResNet runtime test failed.\",\n \"Check PyTorch/CUDA health, available VRAM, and GPU compute stability.\",\n (\"Look for CUDA OOM, cuDNN, or torch runtime errors.\", \"Run a smaller isolated torch workload.\"),\n ),\n ECC_FAILED: FailureCatalogEntry(\n ECC_FAILED,\n \"The ECC runtime test failed.\",\n \"Check GPU ECC counters and hardware health.\",\n (\"Inspect ECC error counters.\", \"Review dmesg and nvidia-smi health output.\"),\n ),\n NCCL_FAILED: FailureCatalogEntry(\n NCCL_FAILED,\n \"The NCCL distributed runtime test failed.\",\n \"Check multi-GPU connectivity, NCCL transport, and network fabric.\",\n (\"Inspect NCCL error output.\", \"Verify peer-to-peer and multi-GPU communication.\"),\n ),\n STRESS_GPU_BURN_FAILED: FailureCatalogEntry(\n STRESS_GPU_BURN_FAILED,\n \"The stress-ng or gpu-burn runtime test failed.\",\n \"Check thermals, power stability, GPU Xid errors, and host stress logs.\",\n (\"Inspect dmesg for Xid errors.\", \"Review gpu-burn and stress-ng output.\"),\n ),\n INTERRUPTED: FailureCatalogEntry(\n INTERRUPTED,\n \"The runtime test was interrupted.\",\n \"Ensure cleanup completed or destroy the test instance manually.\",\n (\"Check whether the test instance still exists.\", \"Destroy leaked instances if needed.\"),\n ),\n CLEANUP_FAILED: FailureCatalogEntry(\n CLEANUP_FAILED,\n \"Runtime test cleanup failed.\",\n \"Destroy the temporary test instance manually to avoid continued billing.\",\n (\"Run destroy instance for the temporary contract.\", \"Retry cleanup after checking API connectivity.\"),\n ),\n}\n\n\nSTAGE_SYSTEM_REQUIREMENTS = \"system_requirements\"\nSTAGE_RESNET = \"resnet\"\nSTAGE_ECC = \"ecc\"\nSTAGE_NCCL = \"nccl\"\nSTAGE_STRESS_GPU_BURN = \"stress_gpu_burn\"\nSTAGE_STARTUP = \"startup\"\n\n\n_STAGE_PATTERNS: tuple[tuple[re.Pattern[str], str], ...] = (\n (re.compile(r\"^\\s*Running system requirements test\\.\\.\\.\\s*$\", re.IGNORECASE), STAGE_SYSTEM_REQUIREMENTS),\n (re.compile(r\"^\\s*Running ResNet50/ResNet18 test\\.\\.\\.\\s*$\", re.IGNORECASE), STAGE_RESNET),\n (re.compile(r\"^\\s*Running ECC test\\.\\.\\.\\s*$\", re.IGNORECASE), STAGE_ECC),\n (re.compile(r\"^\\s*Running NCCL distributed test\\.\\.\\.\\s*$\", re.IGNORECASE), STAGE_NCCL),\n (re.compile(r\"^\\s*Running stress-ng and gpu-burn\\.\\.\\.\\s*$\", re.IGNORECASE), STAGE_STRESS_GPU_BURN),\n)\n\n_NVML_RE = re.compile(\n r\"nvml|nvidia-smi|driver/library version mismatch|failed to initialize.*nvidia\",\n re.IGNORECASE,\n)\n_RESNET_RE = re.compile(\n r\"resnet|torch|pytorch|cuda out of memory|outofmemory|cudnn|cublas|cuda error|runtimeerror\",\n re.IGNORECASE,\n)\n_ECC_RE = re.compile(r\"\\becc\\b|volatile double bit|aggregate single bit\", re.IGNORECASE)\n_NCCL_RE = re.compile(r\"\\bnccl\\b|unhandled system error|connection timed out|allreduce|peer access\", re.IGNORECASE)\n_STRESS_RE = re.compile(r\"stress-ng|gpu-burn|xid|thermal|power limit|burn-in|hardware error\", re.IGNORECASE)\n\n_DOCKER_PULL_RE = re.compile(\n r\"pull|manifest|not found|unauthorized|denied|repository does not exist|no such image|\"\n r\"pull access denied|requested access to the resource is denied\",\n re.IGNORECASE,\n)\n_STARTUP_RE = re.compile(\n r\"daemon|container|startup|start container|failed to start|oci runtime|runc|nvidia-container|\"\n r\"docker|exited|exec format error|permission denied|mount|entrypoint\",\n re.IGNORECASE,\n)\n\n_API_KEY_RE = re.compile(r\"([?&]api_key=)[^&\\s]+\")\n\n\ndef redact_secret_text(value: object) -> str | None:\n \"\"\"Return a host-visible string with obvious API-key query values redacted.\"\"\"\n if value is None:\n return None\n return _API_KEY_RE.sub(r\"\\1REDACTED\", str(value))\n\n\ndef _mapped_port_names(mapped_ports) -> list[str]:\n if not mapped_ports:\n return []\n if isinstance(mapped_ports, dict):\n return sorted(str(key) for key in mapped_ports.keys())\n return sorted(str(port) for port in mapped_ports)\n\n\ndef make_progress_endpoint_diagnostic(\n *,\n url: str | None = None,\n public_ip: str | None = None,\n container_port: str = PROGRESS_CONTAINER_PORT,\n host_port: str | int | None = None,\n timeout_seconds: int | float | None = None,\n attempt_count: int = 0,\n first_connection_established: bool = False,\n last_error_type: str | None = None,\n last_error: object = None,\n last_status_code: int | None = None,\n mapped_ports=None,\n) -> dict[str, object]:\n \"\"\"Shape progress endpoint state for raw output and UI consumption.\"\"\"\n if url is None and public_ip and host_port:\n url = f\"https://{public_ip}:{host_port}/progress\"\n diagnostic: dict[str, object] = {\n \"url\": url,\n \"public_ip\": public_ip,\n \"container_port\": container_port,\n \"external_port\": str(host_port) if host_port is not None else None,\n \"host_port\": str(host_port) if host_port is not None else None,\n \"timeout_seconds\": timeout_seconds,\n \"attempt_count\": int(attempt_count or 0),\n \"first_connection_established\": bool(first_connection_established),\n \"last_error_type\": last_error_type,\n \"last_error\": redact_secret_text(last_error),\n \"last_status_code\": last_status_code,\n \"mapped_ports\": _mapped_port_names(mapped_ports),\n }\n return diagnostic\n\n\ndef failure_catalog() -> dict[str, dict[str, object]]:\n \"\"\"Return a JSON-serializable copy of the runtime failure catalog.\"\"\"\n return {\n code: {\n \"code\": entry.code,\n \"summary\": entry.summary,\n \"remediation\": entry.remediation,\n \"suggested_steps\": list(entry.suggested_steps),\n }\n for code, entry in FAILURE_CATALOG.items()\n }\n\n\ndef get_failure_entry(code: str) -> FailureCatalogEntry:\n \"\"\"Return the catalog entry for a stable runtime failure code.\"\"\"\n try:\n return FAILURE_CATALOG[code]\n except KeyError as exc:\n raise ValueError(f\"Unknown runtime failure code: {code}\") from exc\n\n\ndef make_failure(\n code: str,\n *,\n stage: str | None = None,\n summary: str | None = None,\n details: str | None = None,\n error: str | None = None,\n remediation: str | None = None,\n suggested_steps: Iterable[str] | None = None,\n underlying_error: str | None = None,\n progress_endpoint: dict[str, object] | None = None,\n) -> dict[str, object]:\n \"\"\"Build a raw-output-friendly diagnostic dictionary.\"\"\"\n entry = get_failure_entry(code)\n diagnostic: dict[str, object] = {\n \"code\": code,\n \"stage\": stage,\n \"summary\": summary or entry.summary,\n \"remediation\": remediation or entry.remediation,\n \"suggested_steps\": list(suggested_steps) if suggested_steps is not None else list(entry.suggested_steps),\n }\n if details:\n diagnostic[\"details\"] = details\n if error:\n diagnostic[\"error\"] = error\n if underlying_error:\n diagnostic[\"underlying_error\"] = underlying_error\n if progress_endpoint:\n diagnostic[\"progress_endpoint\"] = progress_endpoint\n return diagnostic\n\n\ndef stage_from_progress_line(line: str) -> str | None:\n \"\"\"Return the legacy runtime stage introduced by a progress line.\"\"\"\n for pattern, stage in _STAGE_PATTERNS:\n if pattern.search(line):\n return stage\n return None\n\n\ndef classify_legacy_error_line(line: str, stage: str | None = None) -> dict[str, object]:\n \"\"\"Classify a legacy progress ``ERROR`` line into a runtime diagnostic.\"\"\"\n stripped = line.strip()\n lowered_stage = stage.lower() if stage else None\n\n code = LEGACY_PROGRESS_ERROR\n if _NVML_RE.search(stripped):\n code = NVML_FAILED\n elif _NCCL_RE.search(stripped) or lowered_stage == STAGE_NCCL:\n code = NCCL_FAILED\n elif _ECC_RE.search(stripped) or lowered_stage == STAGE_ECC:\n code = ECC_FAILED\n elif _STRESS_RE.search(stripped) or lowered_stage == STAGE_STRESS_GPU_BURN:\n code = STRESS_GPU_BURN_FAILED\n elif _RESNET_RE.search(stripped) or lowered_stage == STAGE_RESNET:\n code = RESNET_FAILED\n\n return make_failure(\n code,\n stage=stage,\n error=stripped,\n details=f\"Legacy progress stream reported: {stripped}\",\n underlying_error=stripped,\n )\n\n\nclass LegacyProgressParser:\n \"\"\"Stateful parser for the legacy runtime progress text stream.\"\"\"\n\n def __init__(self) -> None:\n self.stage: str | None = None\n\n def process_line(self, line: str) -> dict[str, object] | None:\n \"\"\"Update parser state and return a diagnostic for ``ERROR`` lines.\"\"\"\n stage = stage_from_progress_line(line)\n if stage:\n self.stage = stage\n return None\n\n if line.strip().upper().startswith(\"ERROR\"):\n return classify_legacy_error_line(line, self.stage)\n\n return None\n\n def parse(self, text: str) -> list[dict[str, object]]:\n \"\"\"Parse progress text and return diagnostics for any ``ERROR`` lines.\"\"\"\n diagnostics: list[dict[str, object]] = []\n for line in text.splitlines():\n diagnostic = self.process_line(line)\n if diagnostic is not None:\n diagnostics.append(diagnostic)\n return diagnostics\n\n\ndef parse_legacy_progress(text: str) -> list[dict[str, object]]:\n \"\"\"Parse a legacy progress payload with a fresh parser instance.\"\"\"\n return LegacyProgressParser().parse(text)\n\n\ndef classify_status_msg(status_msg: str | None) -> dict[str, object] | None:\n \"\"\"Classify startup/status messages reported while the instance starts.\"\"\"\n if not status_msg:\n return None\n\n msg = status_msg.strip()\n if not msg:\n return None\n\n code = INSTANCE_STATUS_ERROR\n if _DOCKER_PULL_RE.search(msg):\n code = DOCKER_PULL_FAILED\n elif _STARTUP_RE.search(msg):\n code = DAEMON_STARTUP_FAILED\n\n return make_failure(\n code,\n stage=STAGE_STARTUP,\n error=msg,\n details=f\"Instance status message reported: {msg}\",\n underlying_error=msg,\n )\n\n\n__all__ = [\n \"CLEANUP_FAILED\",\n \"DAEMON_STARTUP_FAILED\",\n \"DOCKER_PULL_FAILED\",\n \"ECC_FAILED\",\n \"FAILURE_CATALOG\",\n \"FailureCatalogEntry\",\n \"INSTANCE_CREATE_FAILED\",\n \"INSTANCE_CREATE_MISSING_CONTRACT\",\n \"INSTANCE_OFFLINE_BEFORE_TEST\",\n \"INSTANCE_START_TIMEOUT\",\n \"INSTANCE_STATUS_ERROR\",\n \"INSTANCE_STATUS_POLL_FAILED\",\n \"INTERRUPTED\",\n \"LEGACY_PROGRESS_ERROR\",\n \"LegacyProgressParser\",\n \"MISSING_PUBLIC_IP\",\n \"NCCL_FAILED\",\n \"NVML_FAILED\",\n \"PROGRESS_EMPTY_TIMEOUT\",\n \"PROGRESS_CONTAINER_PORT\",\n \"PROGRESS_ENDPOINT_LOST\",\n \"PROGRESS_ENDPOINT_UNREACHABLE\",\n \"PROGRESS_PORT_NOT_MAPPED\",\n \"RESNET_FAILED\",\n \"RUNTIME_FAILURE_CODES\",\n \"RUNTIME_TEST_TIMEOUT\",\n \"STAGE_ECC\",\n \"STAGE_NCCL\",\n \"STAGE_RESNET\",\n \"STAGE_STARTUP\",\n \"STAGE_STRESS_GPU_BURN\",\n \"STAGE_SYSTEM_REQUIREMENTS\",\n \"STRESS_GPU_BURN_FAILED\",\n \"classify_legacy_error_line\",\n \"classify_status_msg\",\n \"failure_catalog\",\n \"get_failure_entry\",\n \"make_failure\",\n \"make_progress_endpoint_diagnostic\",\n \"parse_legacy_progress\",\n \"redact_secret_text\",\n \"stage_from_progress_line\",\n]\n","content_type":"text/x-python; charset=utf-8","language":"python","size":19010,"content_sha256":"bf32ebc9079b9666fc848e1a58d804910c52c3bd884b5f118fc153d2a302a7d9"},{"filename":"cli/self_test/support_bundle.py","content":"\"\"\"Support bundle helpers for host self-test failures.\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nimport io\nimport os\nimport re\nimport subprocess\nimport tarfile\nimport time\nfrom datetime import datetime, timezone\nfrom pathlib import Path\nfrom typing import Any\n\nfrom vastai.cli.self_test.runtime_diagnostics import redact_secret_text\n\n\nDEFAULT_BUNDLE_DIR = \"/tmp\"\nMAX_TEXT_BYTES = 256 * 1024\nMAX_LOG_BYTES = 256 * 1024\nMAX_COMMAND_SECONDS = 10\nJSON_ARTIFACT_RE = re.compile(r\"(^|/)[^/]+\\.json$\")\nSENSITIVE_KEY_RE = re.compile(\n r\"(api[_-]?key|token|secret|password|passwd|authorization|jupyter_token|instance_api_key|host_port_range)\",\n re.IGNORECASE,\n)\nSENSITIVE_ASSIGNMENT_RE = re.compile(\n r\"(?i)\\b(api[_-]?key|token|secret|password|passwd|authorization|jupyter_token|instance_api_key|host_port_range)\"\n r\"\\b\\s*[:=]\\s*([^\\s,;]+)\"\n)\nDMESG_KEYWORDS_RE = re.compile(r\"\\b(warn|warning|err|error|fail|fault|crit|panic|oops|bug)\\b\", re.IGNORECASE)\n\n\ndef utc_timestamp() -> str:\n return datetime.now(timezone.utc).strftime(\"%Y%m%dT%H%M%SZ\")\n\n\ndef support_bundles_enabled() -> bool:\n value = os.environ.get(\"VAST_SELF_TEST_SUPPORT_BUNDLE\", \"\")\n return value.strip().lower() not in (\"0\", \"false\", \"no\", \"off\")\n\n\ndef _safe_machine_id(machine_id: object) -> str:\n return re.sub(r\"[^A-Za-z0-9_.-]+\", \"_\", str(machine_id or \"unknown\"))[:80]\n\n\ndef _redact_text(text: object, secrets: list[str] | None = None) -> str:\n redacted = redact_secret_text(\"\" if text is None else str(text)) or \"\"\n for secret in secrets or []:\n if secret:\n redacted = redacted.replace(str(secret), \"REDACTED\")\n return SENSITIVE_ASSIGNMENT_RE.sub(lambda match: f\"{match.group(1)}=REDACTED\", redacted)\n\n\ndef _redact_json(value: Any, secrets: list[str] | None = None) -> Any:\n if isinstance(value, dict):\n return {\n str(key): \"REDACTED\" if SENSITIVE_KEY_RE.search(str(key)) else _redact_json(item, secrets)\n for key, item in value.items()\n }\n if isinstance(value, list):\n return [_redact_json(item, secrets) for item in value]\n if isinstance(value, str):\n return _redact_text(value, secrets)\n return value\n\n\ndef _truncate_text(text: str, max_bytes: int) -> str:\n encoded = text.encode(\"utf-8\", errors=\"replace\")\n if len(encoded) \u003c= max_bytes:\n return text\n tail = encoded[-max_bytes:].decode(\"utf-8\", errors=\"replace\")\n return f\"[truncated to last {max_bytes} bytes]\\n{tail}\"\n\n\ndef _decode_bytes(data: bytes) -> str:\n return data.decode(\"utf-8\", errors=\"replace\")\n\n\ndef _read_file_tail(path: Path, max_bytes: int) -> str:\n size = path.stat().st_size\n with path.open(\"rb\") as handle:\n if size > max_bytes:\n handle.seek(-max_bytes, os.SEEK_END)\n data = handle.read()\n text = _decode_bytes(data)\n if size > max_bytes:\n return f\"[truncated from {size} bytes to last {max_bytes} bytes]\\n{text}\"\n return text\n\n\ndef _run_command(args: list[str], *, timeout: int = MAX_COMMAND_SECONDS) -> tuple[str, str | None]:\n try:\n proc = subprocess.run(args, capture_output=True, text=True, timeout=timeout, check=False)\n except FileNotFoundError:\n return \"\", f\"command not found: {args[0]}\"\n except subprocess.TimeoutExpired:\n return \"\", f\"command timed out after {timeout}s: {' '.join(args)}\"\n output = proc.stdout or \"\"\n if proc.stderr:\n output = f\"{output}\\n[stderr]\\n{proc.stderr}\".strip()\n error = None\n if proc.returncode != 0:\n error = f\"command exited {proc.returncode}: {' '.join(args)}\"\n return output, error\n\n\ndef _safe_archive_name(name: str) -> str:\n return re.sub(r\"(^/+|\\.\\.|\\\\)\", \"_\", name)\n\n\ndef _summarize_ss_output(text: str) -> str:\n lines = [line for line in text.splitlines() if line.strip()]\n if not lines:\n return \"\"\n summarized = []\n for idx, line in enumerate(lines):\n parts = line.split()\n if idx == 0 and parts:\n summarized.append(\" \".join(parts[:5]))\n continue\n if len(parts) >= 5:\n summarized.append(\" \".join(parts[:5]))\n else:\n summarized.append(line)\n return \"\\n\".join(summarized)\n\n\ndef _collect_host_artifacts(secrets: list[str] | None = None) -> tuple[dict[str, str], list[dict[str, str]]]:\n files: dict[str, str] = {}\n errors: list[dict[str, str]] = []\n\n for path in sorted(Path(\"/var/lib/vastai_kaalia\").glob(\"kaalia.log*\")):\n if not path.is_file():\n continue\n archive_name = f\"host/kaalia/{path.name}\"\n try:\n files[archive_name] = _redact_text(_read_file_tail(path, MAX_LOG_BYTES), secrets)\n except Exception as exc:\n errors.append({\"artifact\": archive_name, \"error\": _redact_text(exc, secrets)})\n\n daemon_json = Path(\"/etc/docker/daemon.json\")\n if daemon_json.exists():\n try:\n files[\"host/etc-docker-daemon.json\"] = _redact_text(_read_file_tail(daemon_json, MAX_TEXT_BYTES), secrets)\n except Exception as exc:\n errors.append({\"artifact\": \"host/etc-docker-daemon.json\", \"error\": _redact_text(exc, secrets)})\n else:\n errors.append({\"artifact\": \"host/etc-docker-daemon.json\", \"error\": \"file not found\"})\n\n command_specs = [\n (\"host/dmesg-filtered.log\", [\"dmesg\", \"-T\"], \"dmesg\"),\n (\n \"host/journalctl-docker-vastai_kaalia.log\",\n [\"journalctl\", \"-b\", \"-p\", \"err..alert\", \"-u\", \"docker\", \"-u\", \"vastai_kaalia\", \"--since\", \"-24h\", \"--no-pager\"],\n \"journalctl\",\n ),\n (\"host/nvidia-smi.txt\", [\"nvidia-smi\"], \"nvidia-smi\"),\n (\"host/nvidia-smi-L.txt\", [\"nvidia-smi\", \"-L\"], \"nvidia-smi -L\"),\n (\"host/docker-info.txt\", [\"docker\", \"info\"], \"docker info\"),\n (\"host/ip-addr.txt\", [\"ip\", \"addr\"], \"ip addr\"),\n (\"host/ss-tlnp-summary.txt\", [\"ss\", \"-tlnp\"], \"ss -tlnp summary\"),\n ]\n\n for archive_name, command, label in command_specs:\n output, error = _run_command(command)\n if archive_name == \"host/dmesg-filtered.log\":\n output = \"\\n\".join(line for line in output.splitlines() if DMESG_KEYWORDS_RE.search(line))\n elif archive_name == \"host/ss-tlnp-summary.txt\":\n output = _summarize_ss_output(output)\n files[archive_name] = _truncate_text(_redact_text(output, secrets), MAX_TEXT_BYTES)\n if error:\n errors.append({\"artifact\": label, \"error\": _redact_text(error, secrets)})\n\n mounts = Path(\"/proc/mounts\")\n if mounts.exists():\n try:\n mount_lines = [\n line for line in _read_file_tail(mounts, MAX_TEXT_BYTES).splitlines()\n if \"/var/lib/docker\" in line\n ]\n files[\"host/docker-mounts.txt\"] = _redact_text(\"\\n\".join(mount_lines), secrets)\n except Exception as exc:\n errors.append({\"artifact\": \"host/docker-mounts.txt\", \"error\": _redact_text(exc, secrets)})\n else:\n errors.append({\"artifact\": \"host/docker-mounts.txt\", \"error\": \"/proc/mounts not found\"})\n\n return files, errors\n\n\ndef create_support_bundle(\n *,\n machine_id: object,\n output_dir: str | None = None,\n result: dict[str, Any] | None = None,\n cli_output: list[str] | None = None,\n extra_files: dict[str, str] | None = None,\n extra_errors: list[dict[str, str]] | None = None,\n run_started_at: str | None = None,\n command: list[str] | None = None,\n secrets: list[str] | None = None,\n include_local_host_artifacts: bool = False,\n include_host_logs: bool | None = None,\n) -> dict[str, Any]:\n \"\"\"Create a redacted self-test support tarball and return metadata.\"\"\"\n\n if include_host_logs is not None:\n include_local_host_artifacts = include_host_logs\n\n created_at = utc_timestamp()\n safe_id = _safe_machine_id(machine_id)\n bundle_dir = Path(output_dir or DEFAULT_BUNDLE_DIR).expanduser()\n bundle_dir.mkdir(parents=True, exist_ok=True)\n bundle_path = bundle_dir / f\"vast_selftest_{safe_id}_{created_at}.tar.gz\"\n\n files: dict[str, str] = {}\n errors: list[dict[str, str]] = []\n if cli_output:\n files[\"self-test-output.log\"] = _redact_text(\"\\n\".join(cli_output), secrets)\n if result is not None:\n files[\"self-test-result.json\"] = json.dumps(_redact_json(result, secrets), indent=2, sort_keys=True)\n if extra_files:\n for name, content in extra_files.items():\n files[name] = _redact_text(content, secrets)\n if extra_errors:\n errors.extend(_redact_json(extra_errors, secrets))\n if include_local_host_artifacts:\n host_files, host_errors = _collect_host_artifacts(secrets)\n files.update(host_files)\n errors.extend(host_errors)\n\n manifest = {\n \"bundle_version\": 1,\n \"machine_id\": str(machine_id),\n \"created_at_utc\": created_at,\n \"run_started_at_utc\": run_started_at,\n \"command\": _redact_json(command or [], secrets),\n \"max_text_bytes_per_artifact\": MAX_TEXT_BYTES,\n \"max_log_bytes_per_artifact\": MAX_LOG_BYTES,\n \"includes_local_host_artifacts\": include_local_host_artifacts,\n \"files\": sorted(files.keys()) + [\"manifest.json\", \"collection-errors.json\"],\n \"collection_errors\": errors,\n \"note\": (\n \"This bundle contains CLI-visible self-test evidence. Host OS artifacts \"\n \"are included only when includes_local_host_artifacts=true; otherwise \"\n \"host daemon/kaalia/system logs need host-side daemon/backend collection \"\n \"or an explicit dump run on the actual host.\"\n ),\n }\n files[\"manifest.json\"] = json.dumps(manifest, indent=2, sort_keys=True)\n files[\"collection-errors.json\"] = json.dumps(errors, indent=2, sort_keys=True)\n\n with tarfile.open(bundle_path, \"w:gz\") as tar:\n for name, content in sorted(files.items()):\n safe_name = _safe_archive_name(name)\n artifact_text = content if JSON_ARTIFACT_RE.search(name) else _truncate_text(content, MAX_LOG_BYTES)\n data = artifact_text.encode(\"utf-8\", errors=\"replace\")\n info = tarfile.TarInfo(safe_name)\n info.size = len(data)\n info.mtime = int(time.time())\n tar.addfile(info, fileobj=io.BytesIO(data))\n os.chmod(bundle_path, 0o600)\n\n size_bytes = bundle_path.stat().st_size\n return {\n \"path\": str(bundle_path),\n \"created_at_utc\": created_at,\n \"size_bytes\": size_bytes,\n \"files\": sorted(files.keys()),\n \"collection_errors\": errors,\n }\n\n\ndef format_bundle_summary(bundle: dict[str, Any]) -> list[str]:\n lines = [\n \"Self-test diagnostic bundle saved to:\",\n f\" {bundle.get('path')}\",\n \"Bundle contents:\",\n ]\n for name in bundle.get(\"files\", []):\n lines.append(f\" - {name}\")\n errors = bundle.get(\"collection_errors\") or []\n if errors:\n lines.append(\"Some artifacts could not be collected:\")\n for item in errors[:10]:\n lines.append(f\" - {item.get('artifact')}: {item.get('error')}\")\n if len(errors) > 10:\n lines.append(f\" - ... {len(errors) - 10} more\")\n lines.append(\"Review this tarball before sharing it with support.\")\n return lines\n","content_type":"text/x-python; charset=utf-8","language":"python","size":11238,"content_sha256":"88db065f17e32ecbdcd0e284064ea590cb1f66fe28302bd85611f2506bcbd89f"},{"filename":"cli/selfupdate.py","content":"\"\"\"Self-update engine for the managed (curl | bash) CLI install.\n\nLayout (docs/install-design.md, XDG Base Directory Specification): the\ninstall root defaults to $XDG_DATA_HOME/vastai (~/.local/share/vastai),\noverride with $VASTAI_INSTALL_DIR — a single active venv at current/, with\nbin/vastai a fixed symlink into it. Updating rebuilds current/ in place\n(build temp → verify → swap); there is no version retention.\n\n\"Managed install?\" is detected structurally (``is_managed_install``): a managed\nCLI runs from \u003croot>/current with a sibling \u003croot>/bin/uv. That's ground\ntruth — no on-disk marker to drift or hand-copy. pip installs fail the check\nand the updater never touches them.\n\nThe passive nudge is throttled to one manifest GET and one notice per week,\ncapped at 1s, and swallows every failure — the CLI must never get slower or\nnoisier because the manifest endpoint is unreachable.\n\"\"\"\n\nimport json\nimport os\nimport re\nimport shutil\nimport subprocess\nimport sys\nimport time\nfrom pathlib import Path\n\nimport requests\n\nfrom vastai.cli.util import DATA_HOME, DIRS, VERSION\n\nDEFAULT_MANIFEST_URL = \"https://vast.ai/cli/manifest.json\"\nINSTALL_SH_HINT = \"curl -fsSL https://vast.ai/install.sh | bash\"\nPIP_UPGRADE_HINT = \"pip install --upgrade vastai\"\n\n# Oldest version whose bundled selfupdate.py understands this layout\n# (\u003croot>/current + \u003croot>/bin/uv, XDG-based root). Versions before this pin\n# assumed ~/.vastai/env and would install fine but then see themselves as unmanaged,\n# permanently losing self-update until the install is rebuilt from scratch.\nMIN_DOWNGRADE_VERSION = \"1.4.0\"\n\nUPDATE_CHECK_FILE = os.path.join(DIRS['state'], \"update_check.json\")\nCHECK_INTERVAL_S = 7 * 24 * 60 * 60\nNUDGE_TIMEOUT_S = 1.0\n\n# Console scripts shipped in the wheel; each gets a swapped symlink in bin/.\nMANAGED_BINARIES = (\"vastai\", \"serve-vast-deployment\", \"register-python-argcomplete\")\n\n# Fixed, not XDG-overridable — every other XDG-aware CLI lands its symlink here too.\nLOCAL_BIN = Path.home() / \".local\" / \"bin\"\n\n\nclass UpdateError(Exception):\n \"\"\"A self-update step failed; the current install is untouched.\"\"\"\n\n\nclass UninstallError(Exception):\n \"\"\"A step of vastai uninstall failed; nothing further was attempted.\"\"\"\n\n\n# ---------------------------------------------------------------------------\n# Install detection / manifest / versions\n# ---------------------------------------------------------------------------\n\ndef install_root() -> Path:\n \"\"\"Where the managed install lives.\n\n Prefers ground truth from where this interpreter is actually running: if\n we're already inside a managed install (``\u003croot>/current`` next to\n ``\u003croot>/bin/uv``), that's authoritative and immune to\n $VASTAI_INSTALL_DIR/$XDG_DATA_HOME differing between install time and\n now — e.g. a custom install dir that isn't set in the current shell would\n otherwise make a real managed install look unmanaged. Only when we're\n *not* currently running from one (a pip install checking\n ``is_managed_install``) do we fall back to computing where a managed\n install would live.\n \"\"\"\n prefix = Path(sys.prefix).resolve()\n if prefix.name == \"current\" and (prefix.parent / \"bin\" / \"uv\").exists():\n return prefix.parent\n if os.environ.get(\"VASTAI_INSTALL_DIR\"):\n return Path(os.environ[\"VASTAI_INSTALL_DIR\"]).expanduser()\n return Path(DATA_HOME) / \"vastai\"\n\n\ndef is_managed_install() -> bool:\n \"\"\"True iff this CLI was installed by the managed installer.\n\n Detected from where the interpreter actually runs: a managed CLI lives in\n ``\u003croot>/current`` next to ``\u003croot>/bin/uv``. Ground truth, no marker file.\n pip installs run from elsewhere and fail the check.\n \"\"\"\n try:\n root = install_root()\n return (\n Path(sys.prefix).resolve() == (root / \"current\").resolve()\n and (root / \"bin\" / \"uv\").exists()\n )\n except Exception:\n return False\n\n\ndef fetch_manifest(timeout: float = 10.0) -> dict:\n url = os.environ.get(\"VASTAI_MANIFEST_URL\") or DEFAULT_MANIFEST_URL\n try:\n r = requests.get(url, timeout=timeout)\n r.raise_for_status()\n data = r.json()\n except (requests.exceptions.RequestException, ValueError) as e:\n raise UpdateError(f\"Could not fetch release manifest from {url}: {e}\")\n if not isinstance(data, dict) or \"latest\" not in data:\n raise UpdateError(f\"Malformed release manifest at {url}\")\n if data.get(\"schema\") != 1:\n raise UpdateError(\n f\"Release manifest schema {data.get('schema')!r} is not supported \"\n f\"by this CLI version.\\nRe-run the installer: {INSTALL_SH_HINT}\"\n )\n return data\n\n\ndef version_key(v: str):\n \"\"\"Sort key tolerant of dev/rc (below release), post (above), +local (ignored).\"\"\"\n v = (v or \"\").strip().lstrip(\"v\").split(\"+\", 1)[0]\n nums, extra = [], 0\n for part in v.split(\".\"):\n if part.isdigit():\n nums.append(int(part))\n continue\n m = re.match(r\"\\d+\", part)\n if m:\n nums.append(int(m.group()))\n extra = 1 if part.startswith(\"post\") else -1\n break\n while len(nums) \u003c 3:\n nums.append(0)\n return tuple(nums), extra\n\n\ndef is_newer(candidate: str, current: str) -> bool:\n return version_key(candidate) > version_key(current)\n\n\n# ---------------------------------------------------------------------------\n# Passive nudge\n# ---------------------------------------------------------------------------\n\ndef _load_check_state() -> dict:\n try:\n with open(UPDATE_CHECK_FILE) as f:\n data = json.load(f)\n return data if isinstance(data, dict) else {}\n except (OSError, ValueError):\n return {}\n\n\ndef _save_check_state(state: dict) -> None:\n try:\n with open(UPDATE_CHECK_FILE, \"w\") as f:\n json.dump(state, f)\n except OSError:\n pass\n\n\ndef _stderr_is_tty() -> bool:\n try:\n return sys.stderr.isatty()\n except Exception:\n return False\n\n\ndef notify_update(args=None) -> None:\n \"\"\"Pre-command, best-effort upgrade nudge. Must never raise or block.\"\"\"\n try:\n _notify_update(args)\n except Exception:\n pass\n\n\ndef _notify_update(args) -> None:\n suppressed = (\n os.environ.get(\"VASTAI_NO_UPDATE_CHECK\")\n or os.environ.get(\"CI\")\n or getattr(args, \"raw\", False)\n or not _stderr_is_tty()\n # the updater itself doesn't need a nudge\n or getattr(getattr(args, \"func\", None), \"__module__\", \"\") == \"vastai.cli.commands.update\"\n )\n if suppressed:\n return\n\n now = time.time()\n state = _load_check_state()\n if now - state.get(\"checked_at\", 0) > CHECK_INTERVAL_S:\n state[\"checked_at\"] = now # stamped even on failure, so errors back off too\n try:\n state[\"latest\"] = fetch_manifest(timeout=NUDGE_TIMEOUT_S).get(\"latest\")\n except UpdateError:\n pass\n _save_check_state(state)\n\n latest = state.get(\"latest\")\n if not latest or not is_newer(latest, VERSION):\n return\n if now - state.get(\"notified_at\", 0) \u003c= CHECK_INTERVAL_S:\n return\n\n hint = \"vastai update\" if is_managed_install() else PIP_UPGRADE_HINT\n arrow = \"↑\" if \"utf\" in (sys.stderr.encoding or \"\").lower() else \"*\"\n # Trailing blank line separates the nudge from the command output that follows it.\n print(\n f\"{arrow} vastai {latest} is available (you have {VERSION}). Run `{hint}`.\\n\",\n file=sys.stderr,\n )\n state[\"notified_at\"] = now\n _save_check_state(state)\n\n\n# ---------------------------------------------------------------------------\n# Update / rollback (managed installs only)\n# ---------------------------------------------------------------------------\n\ndef _run(cmd, *, env=None):\n cmd = [str(c) for c in cmd]\n try:\n return subprocess.run(cmd, env=env, check=True, capture_output=True, text=True)\n except subprocess.CalledProcessError as e:\n # Surface the tool's own stderr — that's the real error (e.g. uv's\n # resolution failure), not a category we invent.\n detail = (e.stderr or e.stdout or \"\").strip()\n raise UpdateError(f\"Command failed: {' '.join(cmd)}\" + (f\"\\n{detail}\" if detail else \"\"))\n except OSError as e:\n # Carry the actual OSError through (wrong path, missing binary, …)\n # rather than relabeling everything \"Required tool not found\".\n raise UpdateError(f\"Could not run {' '.join(cmd)}: {e}\")\n\n\ndef _prune_uv_cache(uv: Path, env: dict) -> None:\n \"\"\"Best-effort prune of uv's cache after a successful swap.\n\n Each update hardlinks the new env out of the cache, so the cache\n accumulates every version's wheels forever unless pruned. Prune (not\n clean): the cache is the user-global one, shared with any other uv on\n the machine. The update already succeeded — never fail because of this.\n \"\"\"\n try:\n _run([uv, \"cache\", \"prune\", \"--quiet\"], env=env)\n except UpdateError:\n pass\n\n\ndef _link_binaries(root: Path) -> None:\n \"\"\"Point bin/\u003cname> at the fixed current/bin/\u003cname> for each shipped\n console script. The target path is constant across updates (always\n ``current/``), so this is not symlink *retargeting* — just (re)creating\n stable links.\"\"\"\n for name in MANAGED_BINARIES:\n if not (root / \"current\" / \"bin\" / name).exists():\n continue\n link, tmp = root / \"bin\" / name, root / \"bin\" / f\".{name}.tmp\"\n if tmp.is_symlink() or tmp.exists():\n tmp.unlink()\n os.symlink(Path(\"..\") / \"current\" / \"bin\" / name, tmp)\n os.replace(tmp, link)\n\n\ndef wheel_spec(target: str, manifest: dict) -> str:\n \"\"\"Requirement spec that installs ``target``.\n\n Latest installs the manifest's release wheel: hash-pinned bytes straight\n from the GitHub Release, no PyPI propagation window. Other targets are\n pins/rollbacks to long-published versions, where the PyPI pin is\n race-free and covers releases that predate wheel_url.\n \"\"\"\n install = manifest.get(\"install\") or {}\n url, sha = install.get(\"wheel_url\"), install.get(\"wheel_sha256\")\n if target == manifest.get(\"latest\") and url and sha:\n return f\"vastai @ {url}#sha256={sha}\"\n return f\"vastai=={target}\"\n\n\ndef perform_update(target: str, manifest: dict) -> None:\n \"\"\"Install ``target`` into a fresh env, verify it, then swap it in.\n\n Single active install (the deno model): the new env is built in a temp dir\n and only swapped over the live one after it verifies, so an interrupted\n update can never leave a half-written or broken ``vastai``. There is no\n version retention — a re-pin or rollback is just ``update --version X``,\n which reinstalls. Wheel integrity comes from ``wheel_spec``; we\n additionally confirm the installed version matches ``target``.\n \"\"\"\n if version_key(target) \u003c version_key(MIN_DOWNGRADE_VERSION):\n raise UpdateError(\n f\"Can't roll back to {target}: versions before {MIN_DOWNGRADE_VERSION} \"\n \"don't recognize this install's layout and would lose the ability \"\n f\"to self-update. The oldest version you can roll back to is {MIN_DOWNGRADE_VERSION}.\"\n )\n root = install_root()\n uv = root / \"bin\" / \"uv\"\n if not uv.exists():\n raise UpdateError(f\"Bootstrap tool missing at {uv}.\\nRe-run the installer: {INSTALL_SH_HINT}\")\n\n python_pin = (manifest.get(\"install\") or {}).get(\"python\") or \"3.12\"\n env_dir, new_dir, old_dir = root / \"current\", root / \".current.new\", root / \".current.old\"\n shutil.rmtree(new_dir, ignore_errors=True)\n\n uv_env = dict(os.environ,\n UV_PYTHON_INSTALL_DIR=str(root / \"python\"),\n UV_PYTHON_PREFERENCE=\"only-managed\")\n try:\n # --relocatable: shebangs must not hardcode the build path (env is renamed into place).\n _run([uv, \"venv\", new_dir, \"--python\", python_pin, \"--relocatable\", \"--quiet\"], env=uv_env)\n _run([uv, \"pip\", \"install\", \"--python\", new_dir / \"bin\" / \"python\",\n \"--quiet\", wheel_spec(target, manifest)], env=uv_env)\n installed = (_run([new_dir / \"bin\" / \"vastai\", \"--version\"]).stdout or \"\").strip()\n if target not in installed:\n raise UpdateError(\n f\"Verification failed: new install reports version {installed!r}, expected {target}\"\n )\n except UpdateError:\n shutil.rmtree(new_dir, ignore_errors=True)\n raise\n\n # Swap via two renames (~instant); a crash mid-build only dirties .current.new.\n shutil.rmtree(old_dir, ignore_errors=True)\n if env_dir.exists():\n os.replace(env_dir, old_dir)\n os.replace(new_dir, env_dir)\n _link_binaries(root)\n shutil.rmtree(old_dir, ignore_errors=True)\n _prune_uv_cache(uv, uv_env)\n\n\n# ---------------------------------------------------------------------------\n# Uninstall (managed installs only)\n# ---------------------------------------------------------------------------\n\ndef perform_uninstall() -> Path:\n \"\"\"Remove the managed install: the install root and its bin/ symlinks.\n\n Config, cache, and state (``vastai.cli.util.DIRS`` — each independently\n XDG-overridable) are left alone — the same guarantee documented for a\n manual ``rm -rf`` (docs/install-design.md §3). Safe to call while running\n from inside ``root``: like the update swap above, removing the directory\n doesn't disturb an already-running interpreter on POSIX.\n\n Raises ``UninstallError`` on any real failure (permissions, a symlink\n loop, …) instead of silently reporting success — callers should treat\n this as fallible, not a fire-and-forget best-effort cleanup.\n \"\"\"\n root = install_root().resolve()\n for name in MANAGED_BINARIES:\n link = LOCAL_BIN / name\n try:\n if link.is_symlink() and root in link.resolve().parents:\n link.unlink()\n except OSError as e:\n raise UninstallError(f\"Could not remove {link}: {e}\")\n try:\n shutil.rmtree(root)\n except FileNotFoundError:\n pass # already gone — nothing left to remove\n except OSError as e:\n raise UninstallError(f\"Could not remove {root}: {e}\")\n return root\n","content_type":"text/x-python; charset=utf-8","language":"python","size":14306,"content_sha256":"3e88a33de1a7f0404826ebc54a836f9677645b4007d97ebca46146700fc81e5c"},{"filename":"cli/util.py","content":"\"\"\"\nCLI utility functions and constants for the Vast.ai CLI.\n\nExtracted from vast.py - contains version checking, config directory setup,\nconstants, and various helper functions used by CLI commands.\n\"\"\"\n\nfrom __future__ import unicode_literals, print_function\n\nimport re\nimport json\nimport sys\nimport argparse\nimport os\nimport time\nimport math\nimport subprocess\nimport shutil\nimport requests\nimport getpass\nfrom pathlib import Path\nfrom datetime import date, datetime, timedelta, timezone\nfrom concurrent.futures import ThreadPoolExecutor\nfrom typing import Dict, List, Optional\n\n# Re-export string_to_unix_epoch so CLI command modules can import from here\nfrom vastai.api.query import string_to_unix_epoch # noqa: F401\nfrom vastai.utils import VERSION, parse_version, parse_env, smart_split # noqa: F401\n\n\n# Server URL default — canonical definition lives in api/client.py\nfrom vastai.api.client import server_url_default # noqa: E402, F401\n\n# Sentinel for API key argument default (to distinguish \"not provided\" from None)\napi_key_guard = object()\n\n\n# ---------------------------------------------------------------------------\n# App constants\n# ---------------------------------------------------------------------------\n\nAPP_NAME = \"vastai\"\n\n# Define emoji support and fallbacks\n_HAS_EMOJI = sys.stdout.encoding and 'utf' in sys.stdout.encoding.lower()\nSUCCESS = \"\\u2705\" if _HAS_EMOJI else \"[OK]\"\nWARN = \"\\u26a0\\ufe0f\" if _HAS_EMOJI else \"[!]\"\nFAIL = \"\\u274c\" if _HAS_EMOJI else \"[X]\"\nINFO = \"\\u2139\\ufe0f\" if _HAS_EMOJI else \"[i]\"\n\n\n# ---------------------------------------------------------------------------\n# Config directory setup\n# ---------------------------------------------------------------------------\n\ntry:\n # Although xdg-base-dirs is the newer name, there's\n # python compatibility issues with dependencies that\n # can be unresolvable using things like python 3.9\n # So we actually use the older name, thus older\n # version for now. This is as of now (2024/11/15)\n # the safer option. -cjm\n import xdg\n\n DIRS = {\n 'config': xdg.xdg_config_home(),\n 'temp': xdg.xdg_cache_home(),\n 'state': xdg.xdg_state_home(),\n }\n # Not part of DIRS: DIRS entries are this CLI's own runtime data and get\n # auto-created below regardless of install method. DATA_HOME is where a\n # *managed install* lives (selfupdate.install_root()) — a pip install\n # must never side-effect that directory into existence just by running.\n DATA_HOME = xdg.xdg_data_home()\n\nexcept Exception:\n # Reasonable defaults.\n from pathlib import Path\n _home = str(Path.home())\n DIRS = {\n 'config': os.path.join(_home, '.config'),\n 'temp': os.path.join(_home, '.cache'),\n 'state': os.path.join(_home, '.local', 'state'),\n }\n DATA_HOME = os.path.join(_home, '.local', 'share')\n\nfor key in DIRS.keys():\n DIRS[key] = path = os.path.join(DIRS[key], APP_NAME)\n if not os.path.exists(path):\n os.makedirs(path)\n\nCACHE_FILE = os.path.join(DIRS['temp'], \"gpu_names_cache.json\")\nGPU_TYPES_CACHE_FILE = os.path.join(DIRS['temp'], \"gpu_types_cache.json\")\nCACHE_DURATION = timedelta(hours=24)\n\n\ndef _get_gpu_names() -> Optional[List[str]]:\n \"\"\"Returns a set of GPU names available on Vast.ai, with results cached for 24 hours.\"\"\"\n\n def is_cache_valid() -> bool:\n \"\"\"Checks if the cache file exists and is less than 24 hours old.\"\"\"\n if not os.path.exists(CACHE_FILE):\n return False\n cache_age = datetime.now() - datetime.fromtimestamp(os.path.getmtime(CACHE_FILE))\n return cache_age \u003c CACHE_DURATION\n\n if is_cache_valid():\n try:\n with open(CACHE_FILE, \"r\") as file:\n gpu_names = json.load(file)\n except (OSError, json.JSONDecodeError):\n gpu_names = None\n else:\n endpoint = \"/api/v0/gpu_names/unique/\"\n url = f\"{server_url_default}{endpoint}\"\n try:\n r = requests.get(url, headers={})\n r.raise_for_status()\n gpu_names = r.json()\n except (requests.exceptions.RequestException, ValueError):\n return None\n try:\n with open(CACHE_FILE, \"w\") as file:\n json.dump(gpu_names, file)\n except OSError:\n pass\n\n try:\n return [\n name.replace(\" \", \"_\").replace(\"-\", \"_\") for name in gpu_names['gpu_names']\n ]\n except (TypeError, KeyError):\n return None\n\n\ndef _get_gpu_types() -> Optional[List[dict]]:\n \"\"\"Returns the GPU type catalog from /api/v0/gpu_types/, cached for 24 hours.\n\n Source of truth for canonical GPU names and per-card VRAM. Returns None on\n any error so callers fall back to the hardcoded constants (never empty).\n \"\"\"\n\n def is_cache_valid() -> bool:\n \"\"\"Checks if the cache file exists and is less than 24 hours old.\"\"\"\n if not os.path.exists(GPU_TYPES_CACHE_FILE):\n return False\n cache_age = datetime.now() - datetime.fromtimestamp(os.path.getmtime(GPU_TYPES_CACHE_FILE))\n return cache_age \u003c CACHE_DURATION\n\n if is_cache_valid():\n try:\n with open(GPU_TYPES_CACHE_FILE, \"r\") as file:\n payload = json.load(file)\n except (OSError, json.JSONDecodeError):\n payload = None\n else:\n endpoint = \"/api/v0/gpu_types/\"\n url = f\"{server_url_default}{endpoint}\"\n try:\n r = requests.get(url, headers={})\n r.raise_for_status()\n payload = r.json()\n except (requests.exceptions.RequestException, ValueError):\n return None\n try:\n with open(GPU_TYPES_CACHE_FILE, \"w\") as file:\n json.dump(payload, file)\n except OSError:\n pass\n\n try:\n return payload['gpu_types']\n except (TypeError, KeyError):\n return None\n\n\nAPIKEY_FILE = os.path.join(DIRS['config'], \"vast_api_key\")\nAPIKEY_FILE_HOME = os.path.expanduser(\"~/.vast_api_key\") # Legacy\nTFAKEY_FILE = os.path.join(DIRS['config'], \"vast_tfa_key\")\n\nif not os.path.exists(APIKEY_FILE) and os.path.exists(APIKEY_FILE_HOME):\n #print(f'copying key from {APIKEY_FILE_HOME} -> {APIKEY_FILE}')\n shutil.copyfile(APIKEY_FILE_HOME, APIKEY_FILE)\n\n\ndef format_key_suffix(k):\n \"\"\"Format the last 4 chars of an API key for display, e.g. '...a3f9'.\"\"\"\n if k and len(k) >= 4:\n return f\"...{k[-4:]}\"\n return \"(empty)\"\n\n\n# 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'.\n\nROLE_FILE = os.path.join(DIRS['config'], \"vast_role\")\nROLE_CLIENT = \"client\"\nROLE_HOST = \"host\"\nVALID_ROLES = {ROLE_CLIENT, ROLE_HOST}\n\n\ndef get_role():\n \"\"\"Return the stored CLI role ('client' or 'host'), or None if ROLE_FILE is missing/unrecognized — never raises.\"\"\"\n try:\n with open(ROLE_FILE) as f:\n role = f.read().strip().lower()\n except OSError:\n return None\n return role if role in VALID_ROLES else None\n\n\ndef is_client_view(role):\n \"\"\"Whether role (as returned by get_role) means \"show the client view\" — everything except explicit 'host', including None.\"\"\"\n return role != ROLE_HOST\n\n\ndef set_role_file(role):\n \"\"\"Persist role ('client' or 'host') to ROLE_FILE; raises ValueError for anything else.\"\"\"\n if role not in VALID_ROLES:\n raise ValueError(f\"role must be one of {sorted(VALID_ROLES)}, got {role!r}\")\n with open(ROLE_FILE, \"w\") as f:\n f.write(role)\n\n\ndef detect_role(client):\n \"\"\"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).\"\"\"\n if get_role() is not None:\n return\n try:\n from vastai.api import billing as billing_api\n is_host = bool(billing_api.show_user(client).get(\"host_agreement_accepted\"))\n set_role_file(ROLE_HOST if is_host else ROLE_CLIENT)\n except Exception:\n return\n\n\n# ---------------------------------------------------------------------------\n# Simple utility class\n# ---------------------------------------------------------------------------\n\nclass Object(object):\n pass\n\n\n# ---------------------------------------------------------------------------\n# Validation helpers\n# ---------------------------------------------------------------------------\n\ndef validate_seconds(value):\n \"\"\"Validate that the input value is a valid number for seconds between yesterday and Jan 1, 2100.\"\"\"\n try:\n val = int(value)\n\n # Calculate min_seconds as the start of yesterday in seconds\n yesterday = datetime.now() - timedelta(days=1)\n min_seconds = int(yesterday.timestamp())\n\n # Calculate max_seconds for Jan 1st, 2100 in seconds\n max_date = datetime(2100, 1, 1, 0, 0, 0)\n max_seconds = int(max_date.timestamp())\n\n if not (min_seconds \u003c= val \u003c= max_seconds):\n raise argparse.ArgumentTypeError(f\"{value} is not a valid second timestamp.\")\n return val\n except ValueError:\n raise argparse.ArgumentTypeError(f\"{value} is not a valid integer.\")\n\n\n# ---------------------------------------------------------------------------\n# VRL (Vast Resource Locator) parsing\n# ---------------------------------------------------------------------------\n\nclass VRLException(Exception):\n pass\n\ndef parse_vast_url(url_str):\n \"\"\"\n Breaks up a vast-style url in the form instance_id:path and does\n some basic sanity type-checking.\n\n :param url_str:\n :return:\n \"\"\"\n\n instance_id = None\n path = url_str\n #print(f'url_str: {url_str}')\n if (\":\" in url_str):\n url_parts = url_str.split(\":\", 2)\n if len(url_parts) == 2:\n (instance_id, path) = url_parts\n else:\n raise VRLException(\"Invalid VRL (Vast resource locator).\")\n else:\n try:\n instance_id = int(path)\n path = \"/\"\n except:\n pass\n\n valid_unix_path_regex = re.compile('^(/)?([^/\\0]+(/)?)+

Important: agents should read /llm.txt, /llms.txt, or /.well-known/skills.json to discover the public Skillopedia API.

)\n # Got this regex from https://stackoverflow.com/questions/537772/what-is-the-most-correct-regular-expression-for-a-unix-file-path\n if (path != \"/\") and (valid_unix_path_regex.match(path) is None):\n raise VRLException(f\"Path component: {path} of VRL is not a valid Unix style path.\")\n\n #print(f'instance_id: {instance_id}')\n #print(f'path: {path}')\n return (instance_id, path)\n\n\n# ---------------------------------------------------------------------------\n# SSH key helpers\n# ---------------------------------------------------------------------------\n\ndef get_ssh_key(argstr):\n # Import deindent from display module (avoids circular imports)\n from vastai.cli.display import deindent\n\n ssh_key = argstr\n # Including a path to a public key is pretty reasonable.\n if os.path.exists(argstr):\n with open(argstr) as f:\n ssh_key = f.read()\n\n if \"PRIVATE KEY\" in ssh_key:\n raise ValueError(deindent(\"\"\"\n \\U0001f434 Woah, hold on there, partner!\n\n That's a *private* SSH key. You need to give the *public*\n one. It usually starts with 'ssh-rsa', is on a single line,\n has around 200 or so \"base64\" characters and ends with\n some-user@some-where. \"Generate public ssh key\" would be\n a good search term if you don't know how to do this.\n \"\"\", add_separator=False))\n\n if not ssh_key.lower().startswith('ssh'):\n raise ValueError(deindent(\"\"\"\n Are you sure that's an SSH public key?\n\n Usually it starts with the stanza 'ssh-(keytype)'\n where the keytype can be things such as rsa, ed25519-sk,\n or dsa. What you passed me was:\n\n {}\n\n And welp, that just don't look right.\n \"\"\".format(ssh_key), add_separator=False))\n\n return ssh_key\n\n\ndef generate_ssh_key(auto_yes=False):\n \"\"\"\n Generate a new SSH key pair using ssh-keygen and return the public key content.\n\n Args:\n auto_yes (bool): If True, automatically answer yes to prompts\n\n Returns:\n str: The content of the generated public key\n\n Raises:\n SystemExit: If ssh-keygen is not available or key generation fails\n \"\"\"\n\n print(\"No SSH key provided. Generating a new SSH key pair and adding public key to account...\")\n\n # Define paths\n ssh_dir = Path.home() / '.ssh'\n private_key_path = ssh_dir / 'id_ed25519'\n public_key_path = ssh_dir / 'id_ed25519.pub'\n\n # Create .ssh directory if it doesn't exist\n try:\n ssh_dir.mkdir(mode=0o700, exist_ok=True)\n except OSError as e:\n print(f\"Error creating .ssh directory: {e}\", file=sys.stderr)\n sys.exit(1)\n\n # Check if any part of the key pair already exists and backup if needed\n if private_key_path.exists() or public_key_path.exists():\n print(f\"An SSH key pair 'id_ed25519' already exists in {ssh_dir}\")\n if auto_yes:\n print(\"Auto-answering yes to backup existing key pair.\")\n response = 'y'\n else:\n response = input(\"Would you like to generate a new key pair and backup your existing id_ed25519 key pair. [y/N]: \").lower()\n if response not in ['y', 'yes']:\n print(\"Aborted. No new key generated.\")\n sys.exit(0)\n\n # Generate timestamp for backup\n timestamp = int(time.time())\n backup_private_path = ssh_dir / f'id_ed25519.backup_{timestamp}'\n backup_public_path = ssh_dir / f'id_ed25519.pub.backup_{timestamp}'\n\n try:\n # Backup existing private key if it exists\n if private_key_path.exists():\n private_key_path.rename(backup_private_path)\n print(f\"Backed up existing private key to: {backup_private_path}\")\n\n # Backup existing public key if it exists\n if public_key_path.exists():\n public_key_path.rename(backup_public_path)\n print(f\"Backed up existing public key to: {backup_public_path}\")\n\n except OSError as e:\n print(f\"Error backing up existing SSH keys: {e}\", file=sys.stderr)\n sys.exit(1)\n\n print(\"Generating new SSH key pair and adding public key to account...\")\n\n # Check if ssh-keygen is available\n try:\n subprocess.run(['ssh-keygen', '--help'], capture_output=True, check=False)\n except FileNotFoundError:\n print(\"Error: ssh-keygen not found. Please install OpenSSH client tools.\", file=sys.stderr)\n sys.exit(1)\n\n # Generate the SSH key pair\n try:\n cmd = [\n 'ssh-keygen',\n '-t', 'ed25519', # Ed25519 key type\n '-f', str(private_key_path), # Output file path\n '-N', '', # Empty passphrase\n '-C', f'{os.getenv(\"USER\") or os.getenv(\"USERNAME\", \"user\")}-vast.ai' # User\n ]\n\n result = subprocess.run(\n cmd,\n capture_output=True,\n text=True,\n input='y\\n', # Automatically answer 'yes' to overwrite prompts\n check=True\n )\n\n except subprocess.CalledProcessError as e:\n print(f\"Error generating SSH key: {e}\", file=sys.stderr)\n if e.stderr:\n print(f\"ssh-keygen error: {e.stderr}\", file=sys.stderr)\n sys.exit(1)\n except Exception as e:\n print(f\"Unexpected error during key generation: {e}\", file=sys.stderr)\n sys.exit(1)\n\n # Set proper permissions for the private key\n try:\n private_key_path.chmod(0o600) # Read/write for owner only\n except OSError as e:\n print(f\"Warning: Could not set permissions for private key: {e}\", file=sys.stderr)\n\n # Read and return the public key content\n try:\n with open(public_key_path, 'r') as f:\n public_key_content = f.read().strip()\n\n return public_key_content\n\n except IOError as e:\n print(f\"Error reading generated public key: {e}\", file=sys.stderr)\n sys.exit(1)\n\n\n# ---------------------------------------------------------------------------\n# List / threading helpers\n# ---------------------------------------------------------------------------\n\ndef split_list(lst, k):\n \"\"\"\n Splits a list into sublists of maximum size k.\n \"\"\"\n return [lst[i:i + k] for i in range(0, len(lst), k)]\n\n\ndef exec_with_threads(f, args, nt=16, max_retries=5):\n def worker(sub_args):\n for arg in sub_args:\n retries = 0\n while retries \u003c= max_retries:\n try:\n result = None\n if isinstance(arg, tuple):\n result = f(*arg)\n else:\n result = f(arg)\n if result: # Assuming a truthy return value means success\n break\n except Exception as e:\n print(str(e))\n pass\n retries += 1\n stime = 0.25 * 1.3 ** retries\n print(f\"retrying in {stime}s\")\n time.sleep(stime) # Exponential backoff\n\n # Split args into nt sublists\n args_per_thread = math.ceil(len(args) / nt)\n sublists = [args[i:i + args_per_thread] for i in range(0, len(args), args_per_thread)]\n\n with ThreadPoolExecutor(max_workers=nt) as executor:\n executor.map(worker, sublists)\n\n\n# ---------------------------------------------------------------------------\n# Date / scheduling helpers\n# ---------------------------------------------------------------------------\n\ndef default_start_date():\n return datetime.now(timezone.utc).strftime(\"%Y-%m-%d\")\n\ndef default_end_date():\n return (datetime.now(timezone.utc) + timedelta(days=7)).strftime(\"%Y-%m-%d\")\n\ndef convert_timestamp_to_date(unix_timestamp):\n utc_datetime = datetime.fromtimestamp(unix_timestamp, tz=timezone.utc)\n return utc_datetime.strftime(\"%Y-%m-%d\")\n\ndef parse_day_cron_style(value):\n \"\"\"\n Accepts an integer string 0-6 or '*' to indicate 'Every day'.\n Returns 0-6 as int, or None if '*'.\n \"\"\"\n val = str(value).strip()\n if val == \"*\":\n return None\n try:\n day = int(val)\n if 0 \u003c= day \u003c= 6:\n return day\n except ValueError:\n pass\n raise argparse.ArgumentTypeError(\"Day must be 0-6 (0=Sunday) or '*' for every day.\")\n\ndef parse_hour_cron_style(value):\n \"\"\"\n Accepts an integer string 0-23 or '*' to indicate 'Every hour'.\n Returns 0-23 as int, or None if '*'.\n \"\"\"\n val = str(value).strip()\n if val == \"*\":\n return None\n try:\n hour = int(val)\n if 0 \u003c= hour \u003c= 23:\n return hour\n except ValueError:\n pass\n raise argparse.ArgumentTypeError(\"Hour must be 0-23 or '*' for every hour.\")\n\ndef convert_dates_to_timestamps(args):\n end_timestamp = time.time()\n start_timestamp = time.time() - (24 * 60 * 60)\n\n from dateutil import parser as dateutil_parser\n\n if args.end_date:\n try:\n end_timestamp = _parse_date_to_utc_timestamp(args.end_date, dateutil_parser)\n except ValueError as e:\n print(f\"Warning: Invalid end date format! Ignoring end date! \\n {str(e)}\")\n\n if args.start_date:\n try:\n start_timestamp = _parse_date_to_utc_timestamp(args.start_date, dateutil_parser)\n except ValueError as e:\n print(f\"Warning: Invalid start date format! Ignoring start date! \\n {str(e)}\")\n\n return start_timestamp, end_timestamp\n\n\ndef _parse_date_to_utc_timestamp(value, dateutil_parser):\n \"\"\"Parse a user-supplied date string into a UNIX timestamp.\n\n Naive inputs (e.g. 'YYYY-MM-DD') are interpreted as UTC so that filter\n windows don't shift with the caller's local timezone. Aware inputs keep\n their declared offset.\n \"\"\"\n dt = dateutil_parser.parse(str(value))\n if dt.tzinfo is None:\n dt = dt.replace(tzinfo=timezone.utc)\n return dt.timestamp()\n\ndef validate_frequency_values(day_of_the_week, hour_of_the_day, frequency):\n\n # Helper to raise an error with a consistent message.\n def raise_frequency_error():\n msg = \"\"\n if frequency == \"HOURLY\":\n msg += \"For HOURLY jobs, day and hour must both be \\\"*\\\".\"\n elif frequency == \"DAILY\":\n msg += \"For DAILY jobs, day must be \\\"*\\\" and hour must have a value between 0-23.\"\n elif frequency == \"WEEKLY\":\n msg += \"For WEEKLY jobs, day must have a value between 0-6 and hour must have a value between 0-23.\"\n sys.exit(msg)\n\n if frequency == \"HOURLY\":\n if not (day_of_the_week is None and hour_of_the_day is None):\n raise_frequency_error()\n if frequency == \"DAILY\":\n if not (day_of_the_week is None and hour_of_the_day is not None):\n raise_frequency_error()\n if frequency == \"WEEKLY\":\n if not (day_of_the_week is not None and hour_of_the_day is not None):\n raise_frequency_error()\n\n\ndef add_scheduled_job(client, args, req_json, cli_command, api_endpoint, request_method, instance_id, contract_end_date=None):\n start_timestamp, end_timestamp = convert_dates_to_timestamps(args)\n if args.end_date is None:\n end_timestamp = contract_end_date\n args.end_date = convert_timestamp_to_date(contract_end_date)\n\n if start_timestamp >= end_timestamp:\n raise ValueError(\"--start_date must be less than --end_date.\")\n\n day, hour, frequency = args.day, args.hour, args.schedule\n\n request_body = {\n \"start_time\": start_timestamp,\n \"end_time\": end_timestamp,\n \"api_endpoint\": api_endpoint,\n \"request_method\": request_method,\n \"request_body\": req_json,\n \"day_of_the_week\": day,\n \"hour_of_the_day\": hour,\n \"frequency\": frequency,\n \"instance_id\": instance_id\n }\n\n response = client.post(\"/commands/schedule_job/\", json_data=request_body)\n\n if args.explain:\n print(\"request json: \")\n print(request_body)\n\n if response.status_code == 200:\n print(f\"add_scheduled_job insert: success - Scheduling {frequency} job to {cli_command} from {args.start_date} UTC to {args.end_date} UTC\")\n elif response.status_code == 401:\n print(f\"add_scheduled_job insert: failed status_code: {response.status_code}. It could be because you aren't using a valid api_key.\")\n elif response.status_code == 422:\n user_input = input(\"Existing scheduled job found. Do you want to update it (y|n)? \")\n if user_input.strip().lower() == \"y\":\n scheduled_job_id = response.json()[\"scheduled_job_id\"]\n response = update_scheduled_job(client, cli_command, f\"/commands/schedule_job/{scheduled_job_id}/\", frequency, args.start_date, args.end_date, request_body)\n else:\n print(\"Job update aborted by the user.\")\n else:\n print(f\"add_scheduled_job insert: failed error: {response.status_code}. Response body: {response.text}\")\n\ndef update_scheduled_job(client, cli_command, schedule_job_path, frequency, start_date, end_date, request_body):\n response = client.put(schedule_job_path, json_data=request_body)\n\n response.raise_for_status()\n if response.status_code == 200:\n print(f\"add_scheduled_job update: success - Scheduling {frequency} job to {cli_command} from {start_date} UTC to {end_date} UTC\")\n print(response.json())\n elif response.status_code == 401:\n print(f\"add_scheduled_job update: failed status_code: {response.status_code}. It could be because you aren't using a valid api_key.\")\n else:\n print(f\"add_scheduled_job update: failed status_code: {response.status_code}.\")\n print(response.json())\n\n return response\n\n\n# ---------------------------------------------------------------------------\n# Permissions\n# ---------------------------------------------------------------------------\n\ndef load_permissions_from_file(file_path):\n with open(file_path, 'r') as file:\n return json.load(file)\n\n\ndef get_template_arguments():\n from vastai.cli.parser import argument\n return [\n argument(\"--name\", help=\"name of the template\", type=str),\n argument(\"--image\", help=\"docker container image to launch\", type=str),\n argument(\"--image_tag\", help=\"docker image tag (can also be appended to end of image_path)\", type=str),\n argument(\"--href\", help=\"link you want to provide\", type=str),\n argument(\"--repo\", help=\"link to repository\", type=str),\n argument(\"--login\", help=\"docker login arguments for private repo authentication, surround with ''\", type=str),\n argument(\"--env\", help=\"Contents of the 'Docker options' field\", type=str),\n argument(\"--ssh\", help=\"Launch as an ssh instance type\", action=\"store_true\"),\n argument(\"--jupyter\", help=\"Launch as a jupyter instance instead of an ssh instance\", action=\"store_true\"),\n argument(\"--direct\", help=\"Use (faster) direct connections for jupyter & ssh\", action=\"store_true\"),\n argument(\"--jupyter-dir\", help=\"For runtype 'jupyter', directory in instance to use to launch jupyter. Defaults to image's working directory\", type=str),\n argument(\"--jupyter-lab\", help=\"For runtype 'jupyter', Launch instance with jupyter lab\", action=\"store_true\"),\n argument(\"--onstart-cmd\", help=\"contents of onstart script as single argument\", type=str),\n argument(\"--search_params\", help=\"search offers filters\", type=str),\n argument(\"-n\", \"--no-default\", action=\"store_true\", help=\"Disable default search param query args\"),\n argument(\"--disk_space\", help=\"disk storage space, in GB\", type=str),\n argument(\"--readme\", help=\"readme string\", type=str),\n argument(\"--hide-readme\", help=\"hide the readme from users\", action=\"store_true\"),\n argument(\"--desc\", help=\"description string\", type=str),\n argument(\"--public\", help=\"make template available to public\", action=\"store_true\"),\n ]\n\n\n# ---------------------------------------------------------------------------\n# Verification thresholds\n# ---------------------------------------------------------------------------\n\ndef required_inet_mbps(gpu_total_ram_mib):\n \"\"\"Machine-total VRAM-scaled bandwidth floor for the self-test pre-flight check.\n\n Returns the minimum inet_down / inet_up (Mb/s) a machine needs to qualify\n for the GPU verification pipeline. Floors at 100, caps at 500, scales\n linearly with total VRAM against a 192 GiB reference point.\n\n Falsy / missing gpu_total_ram falls to the 100 Mb/s floor. The column is\n MiB (binary), so a B200 reporting 183359 MiB (~179 GiB) lands at 466\n Mb/s; the cap is reached once total VRAM crosses 192 GiB, e.g. multi-GPU\n machines.\n \"\"\"\n total_vram_gib = (gpu_total_ram_mib or 0) / 1024.0\n return min(500.0, max(100.0, 500.0 * total_vram_gib / 192.0))\n","content_type":"text/x-python; charset=utf-8","language":"python","size":26960,"content_sha256":"c1963374ddf6e7483afffbe8a607f6d1a58bafd915cd4b85a0fb778012cd83d8"},{"filename":"cli/utils.py","content":"\"\"\"Shared helpers used across all CLI command modules.\"\"\"\n\n\ndef get_parser():\n from vastai.cli.main import parser\n return parser\n\n\ndef get_client(args):\n \"\"\"Create a VastClient from parsed CLI args.\"\"\"\n from vastai.api.client import VastClient\n from vastai.cli.util import detect_role\n client = VastClient(\n api_key=args.api_key,\n server_url=args.url,\n retry=args.retry,\n explain=getattr(args, 'explain', False),\n curl=getattr(args, 'curl', False),\n client_type=\"cli\",\n )\n detect_role(client)\n return client\n","content_type":"text/x-python; charset=utf-8","language":"python","size":577,"content_sha256":"f030fdcb1925ffa70f19001f54f8ced626aa73d94f499d1bd06757a126fda40b"},{"filename":"data/__init__.py","content":"from .query import Query, Column\nfrom .offer import Offer\nfrom .endpoint import EndpointConfig, EndpointData\nfrom .deployment import DeploymentConfig, DeploymentData, DeploymentPutResponse\nfrom .workergroup import WorkergroupConfig\nfrom .instance import Instance, InstanceConfig, CreateInstanceResponse\n","content_type":"text/x-python; charset=utf-8","language":"python","size":303,"content_sha256":"40a37a7a3292a641751af5cc113aaaeaf0ccbd22f6802d3e67a1fde5ead88f3e"},{"filename":"data/deployment.py","content":"import dataclasses\nfrom dataclasses import dataclass\nfrom typing import Any, Optional\n\n\n@dataclass\nclass DeploymentConfig:\n \"\"\"Configuration for creating or updating a deployment.\n\n Required: name, image, file_hash, file_size.\n Endpoint scaling params are optional and follow server defaults.\n \"\"\"\n\n name: str\n image: str\n file_hash: str\n file_size: int\n env: Optional[str] = None\n tag: Optional[str] = None\n search_params: Optional[str] = None\n storage: Optional[float] = None\n ttl: Optional[float] = None\n version_label: Optional[str] = None\n docker_login_user: Optional[str] = None\n docker_login_pass: Optional[str] = None\n docker_login_repo: Optional[str] = None\n # Endpoint scaling params\n cold_workers: Optional[int] = None\n max_workers: Optional[int] = None\n min_load: Optional[float] = None\n min_cold_load: Optional[float] = None\n target_util: Optional[float] = None\n cold_mult: Optional[float] = None\n max_queue_time: Optional[float] = None\n target_queue_time: Optional[float] = None\n inactivity_timeout: Optional[float] = None\n autoscaler_instance: Optional[str] = None\n\n def to_dict(self) -> dict:\n return {k: v for k, v in dataclasses.asdict(self).items() if v is not None}\n\n\n@dataclass\nclass DeploymentPutResponse:\n \"\"\"Response from PUT /api/v0/deployments/.\"\"\"\n\n success: bool\n action: str # \"created\", \"soft_update\", \"autoscale_update\", \"exists\"\n deployment_id: int\n endpoint_id: int\n workergroup_id: Optional[int] = None\n upload_url: Optional[str] = None\n upload_fields: Optional[dict] = None\n evicted_versions: Optional[list] = None\n\n @classmethod\n def from_dict(cls, d: dict) -> \"DeploymentPutResponse\":\n return cls(\n success=d[\"success\"],\n action=d[\"action\"],\n deployment_id=d[\"deployment_id\"],\n endpoint_id=d[\"endpoint_id\"],\n workergroup_id=d.get(\"workergroup_id\"),\n upload_url=d.get(\"upload_url\"),\n upload_fields=d.get(\"upload_fields\"),\n evicted_versions=d.get(\"evicted_versions\"),\n )\n\n\n@dataclass\nclass DeploymentData:\n \"\"\"Full deployment entity as returned by GET /api/v0/deployment/{id}/.\"\"\"\n\n id: int\n name: str\n tag: str\n endpoint_id: Optional[int]\n endpoint_state: Optional[str]\n worker_count: Optional[int]\n s3_key: Optional[str]\n env: Optional[str]\n image: str\n storage: float\n search_params: Optional[str]\n file_hash: str\n current_version_id: int\n last_healthy_version_id: Optional[int]\n ttl: Optional[float]\n last_client_heartbeat: Optional[float]\n created_at: float\n updated_at: float\n\n @classmethod\n def from_dict(cls, d: dict) -> \"DeploymentData\":\n return cls(\n id=d[\"id\"],\n name=d[\"name\"],\n tag=d.get(\"tag\", \"default\"),\n endpoint_id=d.get(\"endpoint_id\"),\n endpoint_state=d.get(\"endpoint_state\"),\n worker_count=d.get(\"worker_count\"),\n s3_key=d.get(\"s3_key\"),\n env=d.get(\"env\"),\n image=d[\"image\"],\n storage=d.get(\"storage\", 50.0),\n search_params=d.get(\"search_params\"),\n file_hash=d[\"file_hash\"],\n current_version_id=d[\"current_version_id\"],\n last_healthy_version_id=d.get(\"last_healthy_version_id\"),\n ttl=d.get(\"ttl\"),\n last_client_heartbeat=d.get(\"last_client_heartbeat\"),\n created_at=d[\"created_at\"],\n updated_at=d[\"updated_at\"],\n )\n","content_type":"text/x-python; charset=utf-8","language":"python","size":3574,"content_sha256":"e0dfcb3ba2b53900ec2c36222bc75052f1f0e5d97d1010016b03b4024e8d36cb"},{"filename":"data/endpoint.py","content":"import dataclasses\nfrom dataclasses import dataclass\nfrom typing import Optional\n\n\n@dataclass\nclass EndpointConfig:\n \"\"\"Configuration for creating or updating an endpoint job.\n\n Only endpoint_name is required. All other fields default to None,\n meaning the server will use its own defaults.\n \"\"\"\n endpoint_name: str\n cold_workers: Optional[int] = None\n max_workers: Optional[int] = None\n min_load: Optional[float] = None\n min_cold_load: Optional[float] = None\n target_util: Optional[float] = None\n cold_mult: Optional[float] = None\n max_queue_time: Optional[float] = None\n target_queue_time: Optional[float] = None\n autoscaler_instance: Optional[str] = None\n endpoint_state: Optional[str] = None\n inactivity_timeout: Optional[float] = None\n\n def to_dict(self) -> dict:\n return {k: v for k, v in dataclasses.asdict(self).items() if v is not None}\n\n\n@dataclass\nclass EndpointData:\n \"\"\"Full endpoint entity as returned by the API.\"\"\"\n id: int\n config: EndpointConfig\n api_key: str\n user_id: int\n created_at: float\n auto_delete_in_seconds: Optional[float]\n auto_delete_due_24h: bool\n\n @classmethod\n def from_dict(cls, d: dict) -> \"EndpointData\":\n config = EndpointConfig(\n endpoint_name=d[\"endpoint_name\"],\n cold_workers=d.get(\"cold_workers\"),\n max_workers=d.get(\"max_workers\"),\n min_load=d.get(\"min_load\"),\n min_cold_load=d.get(\"min_cold_load\"),\n target_util=d.get(\"target_util\"),\n cold_mult=d.get(\"cold_mult\"),\n max_queue_time=d.get(\"max_queue_time\"),\n target_queue_time=d.get(\"target_queue_time\"),\n endpoint_state=d.get(\"endpoint_state\"),\n inactivity_timeout=d.get(\"inactivity_timeout\"),\n )\n return cls(\n id=d[\"id\"],\n config=config,\n api_key=d[\"api_key\"],\n user_id=d[\"user_id\"],\n created_at=d[\"created_at\"],\n auto_delete_in_seconds=d.get(\"auto_delete_in_seconds\"),\n auto_delete_due_24h=d.get(\"auto_delete_due_24h\", False),\n )\n","content_type":"text/x-python; charset=utf-8","language":"python","size":2139,"content_sha256":"c58726c179cfcc4f9b40405a5958eda1b0789202efb0f599a0b1439854edea95"},{"filename":"data/instance.py","content":"import dataclasses\nfrom dataclasses import dataclass\nfrom typing import Any, Optional\n\n\n@dataclass\nclass InstanceConfig:\n \"\"\"Configuration for creating an instance from an offer.\n\n image is required. All other fields are optional.\n \"\"\"\n image: str\n disk: Optional[float] = None\n label: Optional[str] = None\n onstart: Optional[str] = None\n env: Optional[dict[str, str]] = None\n image_login: Optional[str] = None\n runtype: Optional[str] = None\n args: Optional[list] = None\n args_str: Optional[str] = None\n python_utf8: Optional[bool] = None\n lang_utf8: Optional[bool] = None\n use_jupyter_lab: Optional[bool] = None\n jupyter_dir: Optional[str] = None\n template_hash_id: Optional[str] = None\n template_id: Optional[int] = None\n extra: Optional[str] = None\n target_state: Optional[str] = None\n\n def to_dict(self) -> dict:\n return {k: v for k, v in dataclasses.asdict(self).items() if v is not None}\n\n\n@dataclass\nclass CreateInstanceResponse:\n \"\"\"Response from PUT /api/v0/asks/{id}/ on success.\"\"\"\n success: bool\n new_contract: int\n instance_api_key: Optional[str] = None\n ask_id: Optional[int] = None\n\n @classmethod\n def from_dict(cls, d: dict) -> \"CreateInstanceResponse\":\n return cls(\n success=d[\"success\"],\n new_contract=d[\"new_contract\"],\n instance_api_key=d.get(\"instance_api_key\"),\n ask_id=d.get(\"ask_id\"),\n )\n\n\n@dataclass\nclass Instance:\n \"\"\"Instance entity as returned by GET /api/v1/instances/.\"\"\"\n id: int\n machine_id: Optional[int]\n actual_status: Optional[str]\n cur_state: Optional[str]\n next_state: Optional[str]\n intended_status: Optional[str]\n image_uuid: Optional[str]\n image_args: Optional[str]\n image_runtype: Optional[str]\n label: Optional[str]\n num_gpus: Optional[int]\n gpu_name: Optional[str]\n gpu_util: Optional[float]\n gpu_arch: Optional[str]\n gpu_temp: Optional[float]\n cuda_max_good: Optional[str]\n driver_version: Optional[str]\n cpu_cores_effective: Optional[float]\n cpu_util: Optional[float]\n cpu_ram: Optional[int]\n disk_space: Optional[float]\n disk_util: Optional[float]\n disk_usage: Optional[float]\n mem_usage: Optional[float]\n mem_limit: Optional[float]\n vmem_usage: Optional[float]\n inet_up: Optional[float]\n inet_down: Optional[float]\n ssh_host: Optional[str]\n ssh_port: Optional[int]\n ssh_idx: Optional[str]\n public_ipaddr: Optional[str]\n local_ipaddrs: Optional[list]\n direct_port_start: Optional[int]\n direct_port_end: Optional[int]\n status_msg: Optional[str]\n jupyter_token: Optional[str]\n extra_env: Optional[list]\n onstart: Optional[str]\n uptime_mins: Optional[float]\n start_date: Optional[float]\n end_date: Optional[float]\n reliability2: Optional[float]\n dph_total: Optional[float]\n host_id: Optional[int]\n template_id: Optional[int]\n template_hash_id: Optional[str]\n template_name: Optional[str]\n ports: Optional[str]\n\n @classmethod\n def from_dict(cls, d: dict) -> \"Instance\":\n \"\"\"Construct an Instance from an API response dict, ignoring unknown keys.\"\"\"\n known = {f.name: f for f in dataclasses.fields(cls)}\n kwargs = {}\n for name, field in known.items():\n if name in d:\n kwargs[name] = d[name]\n elif field.default is not dataclasses.MISSING:\n pass\n elif field.default_factory is not dataclasses.MISSING:\n pass\n else:\n kwargs[name] = None\n return cls(**kwargs)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":3644,"content_sha256":"85fd3c9560402960cdf8b1ed57b529cb6455e8610451203545576caf0bfee934"},{"filename":"data/offer.py","content":"from dataclasses import dataclass\nfrom typing import Optional\n\n\n@dataclass\nclass Offer:\n \"\"\"\n A search result offer from ask_contract_offers, as returned by /api/v0/bundles/.\n Field names match the API response keys.\n \"\"\"\n id: int\n ask_contract_id: Optional[int]\n machine_id: int\n host_id: int\n num_gpus: int\n gpu_name: str\n gpu_ram: int # MB\n gpu_total_ram: Optional[int] # MB\n gpu_frac: Optional[float]\n gpu_arch: Optional[str]\n gpu_lanes: Optional[int]\n gpu_mem_bw: Optional[float]\n gpu_max_power: Optional[float]\n gpu_max_temp: Optional[float]\n compute_cap: Optional[int]\n cuda_max_good: Optional[float]\n bw_nvlink: Optional[float]\n cpu_name: Optional[str]\n cpu_cores: Optional[int]\n cpu_cores_effective: Optional[float]\n cpu_ghz: Optional[float]\n cpu_ram: Optional[int] # MB\n cpu_arch: Optional[str]\n disk_space: Optional[float]\n disk_bw: Optional[float]\n disk_name: Optional[str]\n inet_up: Optional[float]\n inet_down: Optional[float]\n inet_up_cost: Optional[float]\n inet_down_cost: Optional[float]\n direct_port_count: Optional[int]\n pci_gen: Optional[float]\n pcie_bw: Optional[float]\n mobo_name: Optional[str]\n os_version: Optional[str]\n driver_vers: Optional[int]\n has_avx: Optional[int]\n dph_base: Optional[float]\n dph_total: Optional[float]\n min_bid: Optional[float]\n dlperf: Optional[float]\n dlperf_per_dphtotal: Optional[float]\n flops_per_dphtotal: Optional[float]\n total_flops: Optional[float]\n score: Optional[float]\n vram_costperhour: Optional[float]\n storage_cost: Optional[float]\n storage_total_cost: Optional[float]\n credit_discount_max: Optional[float]\n rentable: Optional[bool]\n rented: Optional[bool]\n is_bid: Optional[bool]\n external: Optional[int]\n hosting_type: Optional[int]\n hostname: Optional[str]\n public_ipaddr: Optional[str]\n geolocation: Optional[str]\n geolocode: Optional[int]\n reliability: Optional[float]\n reliability_mult: Optional[float]\n expected_reliability: Optional[float]\n verification: Optional[str]\n vericode: Optional[int]\n static_ip: Optional[bool]\n vms_enabled: Optional[bool]\n is_vm_deverified: Optional[bool]\n gpu_display_active: Optional[bool]\n logo: Optional[str]\n webpage: Optional[str]\n resource_type: Optional[str]\n start_date: Optional[float]\n end_date: Optional[float]\n duration: Optional[float]\n bundle_id: Optional[int]\n bundled_results: Optional[int]\n cluster_id: Optional[int]\n avail_vol_ask_id: Optional[int]\n avail_vol_dph: Optional[float]\n avail_vol_size: Optional[float]\n nw_disk_min_bw: Optional[int]\n nw_disk_max_bw: Optional[int]\n nw_disk_avg_bw: Optional[int]\n platform_fee: Optional[float]\n\n @classmethod\n def from_dict(cls, d: dict) -> \"Offer\":\n \"\"\"Construct an Offer from a raw API response dict, ignoring unknown keys.\n Missing Optional fields default to None.\"\"\"\n import dataclasses\n known = {f.name: f for f in dataclasses.fields(cls)}\n kwargs = {}\n for name, field in known.items():\n if name in d:\n kwargs[name] = d[name]\n elif hasattr(field.default, '__call__') or field.default is not dataclasses.MISSING:\n pass # has a default, let dataclass handle it\n elif field.default_factory is not dataclasses.MISSING:\n pass # has a default_factory\n else:\n # No default — use None for Optional fields, skip otherwise\n kwargs[name] = None\n return cls(**kwargs)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":4532,"content_sha256":"5630ec01a8f13a8bc5d1d446b5b6ba289b9893f5b170e407f72e391d1fd8ccc1"},{"filename":"data/query.py","content":"# Implements an object oriented query language for Vast standard query (vast/web/apiquery.py)\n\nfrom typing import Any, Optional\n\n\n# Op-to-string symbol mapping. Used by Query.unparse_query() and by\n# vastai/cli/commands/benchmarks.py for human-readable filter rendering.\nOP_TO_STR = {\n \"eq\": \"=\",\n \"neq\": \"!=\",\n \"gt\": \">\",\n \"gte\": \">=\",\n \"lt\": \"\u003c\",\n \"lte\": \"\u003c=\",\n \"in\": \" in \",\n \"notin\": \" notin \",\n}\n\n\nclass Query:\n def __init__(self, query_dict: dict[str, dict[str, Any]]):\n # self.query has format {column : {op : value}}; may have multiple ops per column and multiple columns.\n self.query = query_dict\n\n @classmethod\n def search_defaults(\n cls,\n verified: Optional[bool] = True,\n rentable: Optional[bool] = True,\n rented: Optional[bool] = False,\n ) -> \"Query\":\n \"\"\"Creates a Query pre-populated with default search filters for end-user use.\"\"\"\n q = cls({})\n for name, value in [\n (\"verified\", verified),\n (\"rentable\", rentable),\n (\"rented\", rented),\n ]:\n if value is not None:\n q.query[name] = {\"eq\": value}\n return q\n\n def unparse_query(self) -> str:\n \"\"\"Converts this Query back into a query string that parse_query() would parse into the same dict.\n\n Inverse of vast/web/parse.py:parse_query, i.e. parse_query(query.unparse_query()) == query.query.\n \"\"\"\n # Fields whose values get multiplied by these factors during parsing\n mult_fields = {\n \"cpu_ram\": 1000,\n \"gpu_ram\": 1000,\n \"gpu_total_ram\": 1000,\n \"duration\": 24.0 * 60.0 * 60.0,\n }\n\n parts = []\n for field, ops in self.query.items():\n for op, value in ops.items():\n op_str = OP_TO_STR[op]\n\n # Reverse the multiplier applied during parsing\n if field in mult_fields:\n value = value / mult_fields[field]\n # Use int if it's a whole number\n if value == int(value):\n value = int(value)\n\n # Format the value\n if isinstance(value, bool):\n val_str = \"true\" if value else \"false\"\n elif value is None:\n val_str = \"None\"\n elif isinstance(value, list):\n # in/notin: reverse the space→underscore replacement per element\n items = [str(v).replace(\" \", \"_\") for v in value]\n val_str = \",\".join(items)\n elif isinstance(value, str):\n # Reverse the underscore→space replacement done by parse_query\n val_str = value.replace(\" \", \"_\")\n else:\n val_str = str(value)\n\n parts.append(f\"{field}{op_str}{val_str}\")\n return \" \".join(parts)\n\n def extend(self, other: \"Query\") -> \"Query\":\n \"\"\"\n Extends self.query with filters in other.\n\n self.query should now have all columns that exist in either self.query or other.query;\n and each column in self.query should have all corresponding ops in either query.\n \"\"\"\n for column, ops in other.query.items():\n if column in self.query:\n conflicts = self.query[column].keys() & ops.keys()\n if conflicts:\n raise ValueError(\n f\"Cannot redefine op(s) {conflicts} for column '{column}'\"\n )\n self.query[column].update(ops)\n else:\n self.query[column] = dict(ops)\n return self\n\n\nclass Column:\n def __init__(self, name: str):\n self.name = name\n\n def __eq__(self, value: Any) -> Query:\n return Query({self.name: {\"eq\": value}})\n\n def __ne__(self, value: Any) -> Query:\n return Query({self.name: {\"neq\": value}})\n\n def __lt__(self, value: Any) -> Query:\n return Query({self.name: {\"lt\": value}})\n\n def __le__(self, value: Any) -> Query:\n return Query({self.name: {\"lte\": value}})\n\n def __gt__(self, value: Any) -> Query:\n return Query({self.name: {\"gt\": value}})\n\n def __ge__(self, value: Any) -> Query:\n return Query({self.name: {\"gte\": value}})\n\n def in_(self, values) -> Query:\n return Query({self.name: {\"in\": values}})\n\n def notin_(self, values) -> Query:\n return Query({self.name: {\"notin\": values}})\n\n\n# Pre-initialized Column objects for all columns in ask_contract_offers\nid_ = Column(\"id\")\nask_contract_id = Column(\"ask_contract_id\")\nbundle_id = Column(\"bundle_id\")\nbundled_results = Column(\"bundled_results\")\nbw_nvlink = Column(\"bw_nvlink\")\ncompute_cap = Column(\"compute_cap\")\ncpu_arch = Column(\"cpu_arch\")\ncpu_cores = Column(\"cpu_cores\")\ncpu_cores_effective = Column(\"cpu_cores_effective\")\ncpu_ghz = Column(\"cpu_ghz\")\ncpu_name = Column(\"cpu_name\")\ncpu_ram = Column(\"cpu_ram\")\ncredit_discount_max = Column(\"credit_discount_max\")\ncuda_max_good = Column(\"cuda_max_good\")\ndirect_port_count = Column(\"direct_port_count\")\ndisk_bw = Column(\"disk_bw\")\ndisk_name = Column(\"disk_name\")\ndisk_space = Column(\"disk_space\")\ndlperf = Column(\"dlperf\")\ndlperf_per_dphtotal = Column(\"dlperf_per_dphtotal\")\ndph_base = Column(\"dph_base\")\ndph_total = Column(\"dph_total\")\ndriver_version = Column(\"driver_version\")\ndriver_vers = Column(\"driver_vers\")\nduration = Column(\"duration\")\nend_date = Column(\"end_date\")\nexternal = Column(\"external\")\nflops_per_dphtotal = Column(\"flops_per_dphtotal\")\ngeolocation = Column(\"geolocation\")\ngeolocode = Column(\"geolocode\")\ngpu_arch = Column(\"gpu_arch\")\ngpu_display_active = Column(\"gpu_display_active\")\ngpu_frac = Column(\"gpu_frac\")\ngpu_ids = Column(\"gpu_ids\")\ngpu_lanes = Column(\"gpu_lanes\")\ngpu_mem_bw = Column(\"gpu_mem_bw\")\ngpu_name = Column(\"gpu_name\")\ngpu_ram = Column(\"gpu_ram\")\ngpu_total_ram = Column(\"gpu_total_ram\")\ngpu_max_power = Column(\"gpu_max_power\")\ngpu_max_temp = Column(\"gpu_max_temp\")\nhas_avx = Column(\"has_avx\")\nhost_id = Column(\"host_id\")\nhosting_type = Column(\"hosting_type\")\nhostname = Column(\"hostname\")\ninet_down = Column(\"inet_down\")\ninet_down_cost = Column(\"inet_down_cost\")\ninet_up = Column(\"inet_up\")\ninet_up_cost = Column(\"inet_up_cost\")\nis_bid = Column(\"is_bid\")\nlogo = Column(\"logo\")\nmachine_id = Column(\"machine_id\")\nmin_bid = Column(\"min_bid\")\nmobo_name = Column(\"mobo_name\")\nnum_gpus = Column(\"num_gpus\")\nos_version = Column(\"os_version\")\npci_gen = Column(\"pci_gen\")\npcie_bw = Column(\"pcie_bw\")\npublic_ipaddr = Column(\"public_ipaddr\")\nreliability = Column(\"reliability\")\nreliability_mult = Column(\"reliability_mult\")\nrentable = Column(\"rentable\")\nrented = Column(\"rented\")\nscore = Column(\"score\")\nstart_date = Column(\"start_date\")\nstatic_ip = Column(\"static_ip\")\nstorage_cost = Column(\"storage_cost\")\nstorage_total_cost = Column(\"storage_total_cost\")\ntotal_flops = Column(\"total_flops\")\nverification = Column(\"verification\")\nvericode = Column(\"vericode\")\nvram_costperhour = Column(\"vram_costperhour\")\nwebpage = Column(\"webpage\")\nvms_enabled = Column(\"vms_enabled\")\nexpected_reliability = Column(\"expected_reliability\")\nis_vm_deverified = Column(\"is_vm_deverified\")\nresource_type = Column(\"resource_type\")\ncluster_id = Column(\"cluster_id\")\navail_vol_ask_id = Column(\"avail_vol_ask_id\")\navail_vol_dph = Column(\"avail_vol_dph\")\navail_vol_size = Column(\"avail_vol_size\")\nnw_disk_min_bw = Column(\"nw_disk_min_bw\")\nnw_disk_max_bw = Column(\"nw_disk_max_bw\")\nnw_disk_avg_bw = Column(\"nw_disk_avg_bw\")\nplatform_fee = Column(\"platform_fee\")\n\n# Virtual columns accepted by the API that are converted to real column(s) by search_asks_\n# (see vast/web/views/misc.py)\ngeolocation = Column(\"geolocation\") # hashed → geolocode\ndatacenter = Column(\"datacenter\") # bool → hosting_type {eq: 0|1}\nduration = Column(\"duration\") # seconds from now → end_date\nverified = Column(\"verified\") # bool → vericode\nallocated_storage = Column(\"allocated_storage\") # GB → disk_space {gte: value}\ntarget_reliability = Column(\n \"target_reliability\"\n) # float → expected_reliability {gte: value}\n\n# Column aliases (remapped to real columns before querying)\nreliability2 = Column(\"reliability2\") # → reliability\ndphtotal = Column(\"dphtotal\") # → dph_total\ngpu_totalram = Column(\"gpu_totalram\") # → gpu_total_ram\ngpu_totalram_GB = Column(\"gpu_totalram_GB\") # → gpu_total_ram (value * 1000)\ncpu_ram_GB = Column(\"cpu_ram_GB\") # → cpu_ram (value * 1000)\ngpu_ram_GB = Column(\"gpu_ram_GB\") # → gpu_ram (value * 1000)\nubuntu_version = Column(\"ubuntu_version\") # → os_version\ndriver_version = Column(\"driver_version\") # → driver_vers\n\n\n# ---------------------------------------------------------------------------\n# GPU name constants (from vast/controller/gpulist.json, column [1])\n#\n# Usage: gpu_name == RTX_4090 or gpu_name.in_([RTX_4090, RTX_4090_Ti])\n# ---------------------------------------------------------------------------\n\n# NVIDIA Data-Center / AI\nA10 = \"A10\"\nA10g = \"A10g\"\nA16 = \"A16\"\nA30 = \"A30\"\nA40 = \"A40\"\nA100_PCIE = \"A100 PCIE\"\nA100_SXM4 = \"A100 SXM4\"\nA100_SXM = \"A100 SXM\"\nA100X = \"A100X\"\nA800_PCIE = \"A800 PCIE\"\nB200 = \"B200\"\nGH200_SXM = \"GH200 SXM\"\nH100_PCIE = \"H100 PCIE\"\nH100_SXM = \"H100 SXM\"\nH100_NVL = \"H100 NVL\"\nH200 = \"H200\"\nH200_NVL = \"H200 NVL\"\nL4 = \"L4\"\nL40 = \"L40\"\nL40S = \"L40S\"\nTesla_K20c = \"Tesla K20c\"\nTesla_K80 = \"Tesla K80\"\nTesla_M40 = \"Tesla M40\"\nTesla_P4 = \"Tesla P4\"\nTesla_P6 = \"Tesla P6\"\nTesla_P40 = \"Tesla P40\"\nTesla_P100 = \"Tesla P100\"\nTesla_T4 = \"Tesla T4\"\nTesla_V100 = \"Tesla V100\"\n\n# NVIDIA Consumer – GeForce / GTX\nGeForce_970 = \"GeForce 970\"\nGTX_750_Ti = \"GTX 750 Ti\"\nGTX_1050 = \"GTX 1050\"\nGTX_1050_Ti = \"GTX 1050 Ti\"\nGTX_1060 = \"GTX 1060\"\nGTX_1070 = \"GTX 1070\"\nGTX_1070_Ti = \"GTX 1070 Ti\"\nGTX_1080 = \"GTX 1080\"\nGTX_1080_Ti = \"GTX 1080 Ti\"\nGTX_1650 = \"GTX 1650\"\nGTX_1650_S = \"GTX 1650 S\"\nGTX_1650_Ti = \"GTX 1650 Ti\"\nGTX_1660 = \"GTX 1660\"\nGTX_1660_S = \"GTX 1660 S\"\nGTX_1660_Ti = \"GTX 1660 Ti\"\nGTX_TITAN_X = \"GTX TITAN X\"\n\n# NVIDIA Consumer – RTX 20-series\nRTX_2060 = \"RTX 2060\"\nRTX_2060S = \"RTX 2060S\"\nRTX_2070 = \"RTX 2070\"\nRTX_2070S = \"RTX 2070S\"\nRTX_2080 = \"RTX 2080\"\nRTX_2080S = \"RTX 2080S\"\nRTX_2080_Ti = \"RTX 2080 Ti\"\n\n# NVIDIA Consumer – RTX 30-series\nRTX_3050 = \"RTX 3050\"\nRTX_3060 = \"RTX 3060\"\nRTX_3060_laptop = \"RTX 3060 laptop\"\nRTX_3060_Ti = \"RTX 3060 Ti\"\nRTX_3070 = \"RTX 3070\"\nRTX_3070_laptop = \"RTX 3070 laptop\"\nRTX_3070_Ti = \"RTX 3070 Ti\"\nRTX_3080 = \"RTX 3080\"\nRTX_3080_Ti = \"RTX 3080 Ti\"\nRTX_3090 = \"RTX 3090\"\nRTX_3090_Ti = \"RTX 3090 Ti\"\n\n# NVIDIA Consumer – RTX 40-series\nRTX_4060 = \"RTX 4060\"\nRTX_4060_laptop = \"RTX 4060 laptop\"\nRTX_4060_Ti = \"RTX 4060 Ti\"\nRTX_4070 = \"RTX 4070\"\nRTX_4070_laptop = \"RTX 4070 laptop\"\nRTX_4070S = \"RTX 4070S\"\nRTX_4070S_Ti = \"RTX 4070S Ti\"\nRTX_4070_Ti = \"RTX 4070 Ti\"\nRTX_4080 = \"RTX 4080\"\nRTX_4080_laptop = \"RTX 4080 laptop\"\nRTX_4080S = \"RTX 4080S\"\nRTX_4090 = \"RTX 4090\"\nRTX_4090D = \"RTX 4090D\"\nRTX_4090_Ti = \"RTX 4090 Ti\"\n\n# NVIDIA Consumer – RTX 50-series\nRTX_5060 = \"RTX 5060\"\nRTX_5060_Ti = \"RTX 5060 Ti\"\nRTX_5070 = \"RTX 5070\"\nRTX_5070_Ti = \"RTX 5070 Ti\"\nRTX_5080 = \"RTX 5080\"\nRTX_5090 = \"RTX 5090\"\nRTX_5090_Ti = \"RTX 5090 Ti\"\n\n# NVIDIA Consumer – Titan\nTitan_RTX = \"Titan RTX\"\nTitan_V = \"Titan V\"\nTITAN_V_CEO = \"TITAN V CEO\"\nTitan_X = \"Titan X\"\nTitan_Xp = \"Titan Xp\"\n\n# NVIDIA Professional – Quadro\nQuadro_K620 = \"Quadro K620\"\nQuadro_K2200 = \"Quadro K2200\"\nQuadro_P2000 = \"Quadro P2000\"\nQuadro_P4000 = \"Quadro P4000\"\nQuadro_P5000 = \"Quadro P5000\"\nQuadro_P6000 = \"Quadro P6000\"\nQuadro_GP100 = \"Quadro GP100\"\nQuadro_GV100 = \"Quadro GV100\"\nQ_RTX_4000 = \"Q RTX 4000\"\nQ_RTX_5000 = \"Q RTX 5000\"\nQ_RTX_6000 = \"Q RTX 6000\"\nQ_RTX_8000 = \"Q RTX 8000\"\n\n# NVIDIA Professional – RTX A-series\nRTX_A2000 = \"RTX A2000\"\nRTX_A4000 = \"RTX A4000\"\nRTX_A4500 = \"RTX A4500\"\nRTX_A5000 = \"RTX A5000\"\nRTX_A6000 = \"RTX A6000\"\n\n# NVIDIA Professional – RTX Ada Generation\nRTX_2000Ada = \"RTX 2000Ada\"\nRTX_4000Ada = \"RTX 4000Ada\"\nRTX_4500Ada = \"RTX 4500Ada\"\nRTX_5000Ada = \"RTX 5000Ada\"\nRTX_5880Ada = \"RTX 5880Ada\"\nRTX_6000Ada = \"RTX 6000Ada\"\n\n# NVIDIA Professional – RTX PRO (Blackwell)\nRTX_PRO_4000 = \"RTX PRO 4000\"\nRTX_PRO_4000_Mobile = \"RTX PRO 4000 Mobile\"\nRTX_PRO_4500 = \"RTX PRO 4500\"\nRTX_PRO_5000 = \"RTX PRO 5000\"\nRTX_PRO_6000 = \"RTX PRO 6000\"\nRTX_PRO_6000_S = \"RTX PRO 6000 S\"\nRTX_PRO_6000_WS = \"RTX PRO 6000 WS\"\nRTX_PRO_6000_Max_Q = \"RTX PRO 6000 Max-Q\"\n\n# NVIDIA Mining / Other\nCMP_50HX = \"CMP 50HX\"\nCMP_170HX = \"CMP 170HX\"\nP104_100 = \"P104-100\"\nP106_100 = \"P106-100\"\n\n# AMD Instinct\nInstinctMI50 = \"InstinctMI50\"\nInstinctMI100 = \"InstinctMI100\"\nInstinctMI210 = \"InstinctMI210\"\nInstinctMI250X = \"InstinctMI250X\"\n\n# AMD Radeon Pro\nPRO_W6800 = \"PRO W6800\"\nPRO_W7800 = \"PRO W7800\"\nPRO_W7900 = \"PRO W7900\"\nPro_V620 = \"Pro V620\"\nRadeon_VII = \"Radeon VII\"\nRadeon_Pro_VII = \"Radeon Pro VII\"\n\n# AMD Radeon RX\nRX_6800 = \"RX 6800\"\nRX_6800_XT = \"RX 6800 XT\"\nRX_6900_XT = \"RX 6900 XT\"\nRX_6950_XT = \"RX 6950 XT\"\nRX_7600 = \"RX 7600\"\nRX_7700_XT = \"RX 7700 XT\"\nRX_7800_XT = \"RX 7800 XT\"\nRX_7900_GRE = \"RX 7900 GRE\"\nRX_7900_XT = \"RX 7900 XT\"\nRX_7900_XTX = \"RX 7900 XTX\"\n","content_type":"text/x-python; charset=utf-8","language":"python","size":13140,"content_sha256":"f96cbae26c0246ddacbcc2eaed145a94fb6e7a473959e6d0b5ed05548e04df56"},{"filename":"data/workergroup.py","content":"import dataclasses\nfrom dataclasses import dataclass\nfrom typing import Optional\n\n\n@dataclass\nclass WorkergroupConfig:\n \"\"\"Configuration for creating a workergroup (autoscale job).\n\n Request-only. Workergroup responses come back as raw dicts from\n show_workergroups, so fields the API returns but does not accept as input\n do not belong here.\n\n Either endpoint_id or endpoint_name is required (endpoint_id is set\n automatically when using ManagedEndpoint.add_workergroup()).\n Either template_hash or template_id is required (template resolves image, search_query, etc.).\n search_params is always required by the server; defaults to empty dict.\n \"\"\"\n search_params: str | dict = \"\"\n endpoint_id: Optional[int] = None\n endpoint_name: Optional[str] = None\n template_hash: Optional[str] = None\n template_id: Optional[int] = None\n launch_args: Optional[str] = None\n min_cold_load: Optional[float] = None\n max_workers: Optional[int] = None\n gpu_ram: Optional[float] = None\n autoscaler_instance: Optional[str] = None\n docker_login_user: Optional[str] = None\n docker_login_pass: Optional[str] = None\n docker_login_repo: Optional[str] = None\n\n def to_dict(self) -> dict:\n d = {k: v for k, v in dataclasses.asdict(self).items() if v is not None}\n # search_params is always required by the server\n d.setdefault(\"search_params\", self.search_params)\n return d\n","content_type":"text/x-python; charset=utf-8","language":"python","size":1444,"content_sha256":"d8bcb1d6156e2fa29ec3972694d0ce664ee4622e9b17578d7f4c8b65c3df8407"},{"filename":"pdf/__init__.py","content":"","content_type":"text/x-python; charset=utf-8","language":"python","size":0,"content_sha256":"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"},{"filename":"pdf/vast_pdf.py","content":"#!/usr/bin/env python3\n\n# DEPRECATED: This file is kept for backwards compatibility.\n# Please use the vastai package instead.\n\n# vast_pdf: Library of functions to create PDF reports of various type.\n# Currently only makes invoices.\nimport argparse\nimport os\nimport datetime\nimport math\nimport time\nimport random\nfrom decimal import Decimal\n\nimport typing\nimport PIL.Image\n\ntry:\n from borb.pdf.canvas.color.color import HexColor, X11Color\n from borb.pdf.canvas.layout.image.image import Image\n from borb.pdf.canvas.layout.layout_element import Alignment\n from borb.pdf.canvas.layout.page_layout.multi_column_layout import SingleColumnLayout\n from borb.pdf.canvas.layout.page_layout.page_layout import PageLayout\n from borb.pdf.canvas.layout.table.fixed_column_width_table import FixedColumnWidthTable\n from borb.pdf.canvas.layout.table.flexible_column_width_table import FlexibleColumnWidthTable\n from borb.pdf.canvas.layout.table.table import Table, TableCell\n from borb.pdf.canvas.layout.text.paragraph import Paragraph\n from borb.pdf.document import Document\n from borb.pdf.page.page import Page\n from borb.pdf.pdf import PDF\nexcept ImportError:\n print(\"\"\"\\nERROR: This library depends on a Python package called Borb to make the PDF files. To install this \n package do 'pip3 install borb'.\\n\"\"\")\n\n## Globals\nnum_rows_first_page: int = 18\nnum_rows_subsequents_pages: int = 30\ninvoice_total: float = 0\nnow = datetime.date.today()\ninvoice_number: int = now.year * 12 + now.month - 1\npage_count: int = 0\nno_table_borders = False\nno_table_borders = True\n\n\n# def Paragraph_wr(text: str, *args, **kwargs):\n# \"\"\"This is just a wrapper for the Paragraph function that prevents the text from being the null string \"\".\n#\n# :param str text: The text to be checked for null.\n# \"\"\"\n# if text == \"\":\n# text = \" \"\n# return Paragraph(text, *args, **kwargs)\n\n\ndef build_2nd_block_table() -> FixedColumnWidthTable:\n \"\"\"\n This function creates a Table containing invoice information.\n This information spans the page and is the second large block of\n text on the page.\n\n :rtype FixedColumnWidthTable: a Table containing information such as the company address, payment terms, date, and sum of charge_fields/payments.\n \"\"\"\n global invoice_total, now\n\n table = FixedColumnWidthTable(number_of_rows=3, number_of_columns=3)\n\n table.add(Paragraph(\"Vast.ai Inc.\"))\n table.add(Paragraph(\"Date\", font=\"Helvetica-Bold\", horizontal_alignment=Alignment.RIGHT))\n table.add(Paragraph(\"%d/%d/%d\" % (now.day, now.month, now.year), horizontal_alignment=Alignment.RIGHT))\n\n table.add(Paragraph(\"100 Van Ness Ave.\"))\n table.add(Paragraph(\"Payment Terms:\", font=\"Helvetica-Bold\", horizontal_alignment=Alignment.RIGHT))\n table.add(Paragraph(\"Charged - Do Not Pay\", horizontal_alignment=Alignment.RIGHT))\n\n table.add(Paragraph(\"San Francisco CA 94102\"))\n\n table.add(\n Paragraph(\"Total\", font=\"Helvetica-Bold\", font_size=Decimal(20), horizontal_alignment=Alignment.RIGHT))\n table.add(Paragraph(format_float_val_as_currency(invoice_total), font=\"Helvetica-Bold\", font_size=Decimal(20),\n horizontal_alignment=Alignment.RIGHT))\n\n table.set_padding_on_all_cells(Decimal(2), Decimal(2), Decimal(2), Decimal(2))\n if no_table_borders: table.no_borders()\n return table\n\n\ndef blank_row(table: Table, col_num: int, row_num: int = 1) -> None:\n \"\"\"Just a set of blank rows to act as filler.\n\n :param table: The table we want to modify\n :param int col_num: How many columns in the table?\n :param int row_num: How many blank rows do we need?\n :rtype None:\n \"\"\"\n for j in range(row_num):\n for i in range(col_num):\n table.add(Paragraph(\" \"))\n return None\n\n\ndef field_and_filler(user_blob: typing.Dict, table: Table, fieldname: str) -> None:\n \"\"\"Adds a field to the table along with a filler cell to span the table.\n\n :param str fieldname:\n :param Dict user_blob: A dict containing the user's info.\n :param Table table: The table to be modified.\n :rtype None:\n \"\"\"\n table.add(Paragraph(str(user_blob[fieldname])))\n # table.add(Paragraph(\" \"))\n return None\n\n\ndef build_billto_table(user_blob: dict) -> FixedColumnWidthTable:\n \"\"\"\n This function builds a Table containing billing and shipping information\n It spans the page using a single wide column.\n\n :param Dict user_blob: A dict containing the user's info.\n :rtype Table: a Table containing shipping and billing information\n \"\"\"\n table = FixedColumnWidthTable(number_of_rows=5, number_of_columns=1)\n table.add(Paragraph(\"BILL TO\", font=\"Helvetica-Bold\"))\n # table.add(Paragraph(\" \", font=\"Helvetica-Bold\"))\n\n field_ids = [\"fullname\", \"billaddress_line1\", \"billaddress_line2\"]\n list(map(lambda fieldname: field_and_filler(user_blob, table, fieldname), field_ids))\n\n table.add(Paragraph(\"%s, %s\" % (user_blob[\"billaddress_city\"], user_blob[\"billaddress_zip\"]))) # BILLING\n # table.add(Paragraph(\" \")) # SHIPPING\n # table.add(Paragraph(\" \")) # BILLING\n # table.add(Paragraph(\" \")) # SHIPPING\n table.set_padding_on_all_cells(Decimal(2), Decimal(2), Decimal(2), Decimal(2))\n if no_table_borders: table.no_borders()\n return table\n\n\nclass Charge:\n \"\"\"\n This class represents a charge line.\n \"\"\"\n\n def __init__(self, name: str, quantity: float, rate: float,\n amount: float, type: str, last4: str, timestamp: float):\n \"\"\"\n Charge constructor. Basically just turns a row into a Charge object.\n\n :param str name:\n :param float quantity:\n :param float rate:\n :param float amount:\n :param str type:\n :param str last4:\n :param float timestamp:\n \"\"\"\n self.name: str = name\n # assert quantity >= 0\n self.quantity: float = quantity\n # assert rate >= 0\n self.rate: float = rate\n # assert amount >= 0\n self.amount: float = amount\n self.type = type\n self.last4 = last4\n self.timestamp = timestamp\n\n\ndef format_float_val_as_currency(v) -> str:\n \"\"\"\n Given a float value, format it as currency. Note that the sign is backwards due to the fact that the charge_fields in\n the data that comes back are positive numbers but are written as negative on the invoice. There's also the fact\n that in our format string, if the value is actually negative it will print as \"$-x.xx\" instead of the correct\n \"-$x.xx\". I don't know if this can be fixed with format strings so we rely on this hack.\n\n :param float v:\n :rtype str:\n \"\"\"\n if v > 0:\n return \" -${:10.2f}\".format(v)\n else:\n return \" ${:10.2f}\".format(-v)\n\n\ndef build_charge_table(charges: typing.List[Charge], page_number: int)\\\n -> FlexibleColumnWidthTable:\n \"\"\"\n This function builds a Table containing itemized billing information\n\n :param typing.List[Charge] charges: the rows on the invoice\n :param int page_number: Current page number.\n :rtype FlexibleColumnWidthTable: Borb Table containing the information.\n \"\"\"\n global invoice_total\n num_rows = len(charges)\n table = FlexibleColumnWidthTable(number_of_rows=(num_rows + 3), number_of_columns=4)\n table_header_padding_top = 4\n table_header_padding_bottom = 5\n table_header_padding_left = 3\n table_header_padding_right = 3\n item_font_size = Decimal(11)\n\n for h in [\"Item\", \"Quantity\", \"Rate\", \"Amount\"]:\n table.add(\n TableCell(\n Paragraph(h, font_color=X11Color(\"White\"),\n padding_top=Decimal(table_header_padding_top),\n padding_bottom=Decimal(table_header_padding_bottom),\n padding_left=Decimal(table_header_padding_left),\n padding_right=Decimal(table_header_padding_right),\n vertical_alignment=Alignment.TOP),\n background_color=HexColor(\"0b3954\")\n )\n )\n\n odd_color = HexColor(\"BBBBBB\")\n even_color = HexColor(\"FFFFFF\")\n for row_number, item in enumerate(charges):\n c = even_color if row_number % 2 == 0 else odd_color\n table.add(TableCell(Paragraph(item.name, font=\"Helvetica-Bold\",\n font_size=item_font_size), background_color=c))\n if item.type == \"payment\":\n table.add(TableCell(Paragraph(\" \", horizontal_alignment=Alignment.RIGHT,\n font_size=item_font_size),\n background_color=c))\n table.add(TableCell(Paragraph(\" \", horizontal_alignment=Alignment.RIGHT,\n font_size=item_font_size), background_color=c))\n else:\n table.add(TableCell(Paragraph(\" {:10.2f}\".format(item.quantity),\n horizontal_alignment=Alignment.RIGHT,\n font_size=item_font_size),\n background_color=c))\n table.add(TableCell(Paragraph(format_float_val_as_currency(item.rate),\n horizontal_alignment=Alignment.RIGHT,\n font_size=item_font_size),\n background_color=c))\n table.add(TableCell(\n Paragraph(format_float_val_as_currency(item.amount),\n horizontal_alignment=Alignment.RIGHT,\n font_size=item_font_size), background_color=c))\n\n if page_number == page_count:\n table.add(TableCell(\n Paragraph(\" \", font=\"Helvetica-Bold\", horizontal_alignment=Alignment.RIGHT),\n col_span=3))\n table.add(TableCell(Paragraph(\" \", horizontal_alignment=Alignment.RIGHT)))\n\n table.add(TableCell(\n Paragraph(\"Total\", font=\"Helvetica-Bold\", horizontal_alignment=Alignment.RIGHT),\n col_span=3))\n table.add(TableCell(\n Paragraph(format_float_val_as_currency(invoice_total), horizontal_alignment=Alignment.RIGHT)))\n else:\n blank_row(table, 4, 2)\n table.set_padding_on_all_cells(Decimal(2), Decimal(5), Decimal(2), Decimal(5))\n if no_table_borders: table.no_borders()\n return table\n\n\ndef product_row(charge_fields) -> Charge:\n \"\"\"Makes a single row with charge information in it.\n\n :param charge_fields:\n :rtype Charge:\n \"\"\"\n\n type = charge_fields[\"type\"]\n is_credit = charge_fields[\"is_credit\"] if \"is_credit\" in charge_fields else False\n timestamp = charge_fields[\"timestamp\"]\n last4 = charge_fields[\"last4\"] if \"last4\" in charge_fields and charge_fields[\"last4\"] is not None else \" \"\n\n def credit_or_auto_billing():\n formatted = datetime.datetime.fromtimestamp(timestamp, tz=datetime.timezone.utc).strftime('%Y-%m-%d %H:%M')\n if is_credit:\n return \"Add Credit: *\" + last4 + \" : \" + formatted\n return \"Auto-Billing: *\" + last4 + \" : \" + formatted\n\n description = charge_fields[\n \"description\"] if \"description\" in charge_fields else credit_or_auto_billing() # \"Add Credit: *\" + last4 + \":\" + str(timestamp)\n amount = float(charge_fields[\"amount\"]) if \"amount\" in charge_fields else 0.0\n quantity = float(charge_fields[\"quantity\"]) if \"quantity\" in charge_fields else 1.0\n rate = float(charge_fields[\"rate\"]) if \"rate\" in charge_fields else amount\n return Charge(description, quantity, rate, amount, type, last4, timestamp)\n\n\ndef product_rows(rows_invoice) -> typing.List[Charge]:\n \"\"\"\n Maps a list of invoice dicts to a list of Charge objects using product_row function.\n\n :param rows_invoice:\n :rtype typing.List[Charge]:\n \"\"\"\n if rows_invoice is None:\n rows_invoice = []\n return list(map(lambda charges: product_row(charges), rows_invoice))\n\n\ndef build_invoice_charges_table(rows_invoice: typing.List[typing.Dict],\n charges_per_page: int, page_number: int) -> FlexibleColumnWidthTable:\n \"\"\"\n This function creates a page of invoice charge_fields and depletes the list of charge_fields by the number it\n prints out. Essentially a pop function that removes multiple items at once.\n\n :param int page_number:\n :param typing.List[typing.Dict] rows_invoice: List of rows in the invoice\n :param int charges_per_page: How many rows of charge_fields we put on this page.\n :rtype FlexibleColumnWidthTable:\n \"\"\"\n rows_invoice_chunk = rows_invoice[0:charges_per_page]\n # sums = {\"amount\": compute_column_sum(rows_invoice, \"amount\", True)}\n del rows_invoice[0:charges_per_page]\n return build_charge_table(product_rows(rows_invoice_chunk), page_number)\n\n\ndef compute_column_sum(rows_invoice: typing.List[typing.Dict],\n column_name: str,\n values_are_negative: bool = False) -> float:\n \"\"\"Sum over one of the columns in the invoice.\n\n :param typing.List[typing.Dict] rows_invoice: List of rows in the invoice\n :param str column_name: Name of column to sum over.\n :param bool values_are_negative: are we summing values that are actually negative?\n :rtype: float\n \"\"\"\n s: float = 0\n for row in rows_invoice:\n v: float = math.floor(float(row[column_name]) * 100) / 100\n # v = round(float(row[column_name]) * 100) / 100\n if values_are_negative:\n s = s - v\n else:\n s = s + v\n return math.floor(s * 100) / 100\n # return round(s * 100) / 100\n\n\ndef generate_invoice_page(user_blob: typing.Dict,\n rows_invoice: typing.List[typing.Dict],\n page_number: int, date_header_text: str = \"\") -> Page:\n \"\"\"Makes a single page of the invoice.\n\n :param str date_header_text:\n :param Dict user_blob: Dict of info about the user\n :param typing.List[typing.Dict] rows_invoice: List of rows in the invoice\n :param int page_number: The page number for this page.\n :rtype Page:\n \"\"\"\n global invoice_number\n\n if page_number == 1:\n rows_per_page = num_rows_first_page\n else:\n rows_per_page = num_rows_subsequents_pages\n\n page: Page = Page()\n\n # set PageLayout\n page_layout: PageLayout = SingleColumnLayout(page,\n vertical_margin=page.get_page_info().get_height() * Decimal(0.02))\n table_logo_and_invoice_num = build_logo_and_invoice_num_table(page_number)\n page_layout.add(table_logo_and_invoice_num)\n page_layout.add(Paragraph(\" \"))\n\n if page_number == 1:\n # Invoice information table\n page_layout.add(build_2nd_block_table())\n # Empty paragraph for spacing\n page_layout.add(Paragraph(\" \"))\n # Billing and shipping information table\n page_layout.add(build_billto_table(user_blob))\n page_layout.add(Paragraph(date_header_text + \" \"))\n\n if len(rows_invoice) == 0:\n # If we don't handle this case the client crashes with no output.\n page_layout.add(Paragraph(\"NO DATA\"))\n else:\n table_invoice_rows = build_invoice_charges_table(rows_invoice, rows_per_page, page_number)\n page_layout.add(table_invoice_rows)\n return page\n\n\ndef build_logo_and_invoice_num_table(page_number) -> FixedColumnWidthTable:\n \"\"\"\n At the top of every page is a table with our logo, the invoice number, and little text reading \"page X of Y\".\n This function creates that table and returns it.\n\n :param page_number:\n :rtype FixedColumnWidthTable:\n \"\"\"\n if page_number == 1:\n invoice_number_font_size = Decimal(20)\n invoice_word_font_size = Decimal(50)\n logo_img_filename: str = os.path.join(os.path.dirname(__file__), 'vast.ai-logo.png')\n logo_img_width: int = 72\n logo_img_height: int = 105\n else:\n invoice_number_font_size = Decimal(14)\n invoice_word_font_size = Decimal(20)\n logo_img_filename: str = os.path.join(os.path.dirname(__file__), 'vast.ai-logo-50pct.png')\n logo_img_width: int = 36\n logo_img_height: int = 53\n\n logo_img = PIL.Image.open(logo_img_filename)\n # add corporate logo\n # now = datetime.datetime.now()\n table_logo_and_invoice_num = FixedColumnWidthTable(number_of_rows=2, number_of_columns=4)\n\n table_logo_and_invoice_num.add(\n TableCell(\n Image(\n logo_img,\n width=Decimal(logo_img_width),\n height=Decimal(logo_img_height),\n ), row_span=2)\n )\n table_logo_and_invoice_num.add(\n Paragraph(\"Page %d of %d\" % (page_number, page_count), font=\"Helvetica\", horizontal_alignment=Alignment.RIGHT))\n table_logo_and_invoice_num.add(Paragraph(\" \", font=\"Helvetica-Bold\", horizontal_alignment=Alignment.RIGHT))\n\n table_logo_and_invoice_num.add(Paragraph(\"Invoice\", font=\"Helvetica\", font_size=invoice_word_font_size,\n horizontal_alignment=Alignment.RIGHT))\n\n table_logo_and_invoice_num.add(Paragraph(\" \", font=\"Helvetica-Bold\", horizontal_alignment=Alignment.RIGHT))\n table_logo_and_invoice_num.add(Paragraph(\" \", font=\"Helvetica-Bold\", horizontal_alignment=Alignment.RIGHT))\n\n table_logo_and_invoice_num.add(Paragraph(\"# %d\" % invoice_number, font=\"Helvetica\",\n font_size=invoice_number_font_size, horizontal_alignment=Alignment.RIGHT))\n\n if no_table_borders: table_logo_and_invoice_num.no_borders()\n return table_logo_and_invoice_num\n\n\ndef compute_pages_needed(rows_invoice: typing.List[typing.Dict]) -> int:\n \"\"\"Function to work out how many pages we need so that the page_count can be filled in at the top of every page.\n\n :param typing.List[typing.Dict] rows_invoice: The list of dicts we use elsewhere.\n :rtype int:\n \"\"\"\n num_rows_invoice = len(rows_invoice)\n num_rows_invoice = num_rows_invoice - num_rows_first_page\n page_count: int = math.ceil(num_rows_invoice / num_rows_subsequents_pages) + 1\n return page_count\n\n\n\n\n\ndef generate_invoice(user_blob: typing.Dict,\n rows_invoice: typing.List[typing.Dict], filter_data: typing.Dict) -> None:\n \"\"\"\n This is the main function in this library. It calls everything else and makes the invoice page by page. The\n resulting invoice is written as a single PDF file.\n\n :param filter_data: Parameters for client side filter.\n :param user_blob: info about the user\n :param typing.List[typing.Dict] rows_invoice: The list of dicts we use elsewhere.\n :rtype None:\n \"\"\"\n\n pdf: Document = Document()\n global page_count\n page_count = compute_pages_needed(rows_invoice)\n global invoice_total\n invoice_total = compute_column_sum(rows_invoice, \"amount\")\n page_number = 1\n\n if len(rows_invoice) == 0:\n page = generate_invoice_page(user_blob, rows_invoice, page_number, filter_data[\"header_text\"])\n print(\"Adding Empty page \", str(page_number))\n pdf.append_page(page)\n else:\n while len(rows_invoice) > 0:\n page = generate_invoice_page(user_blob, rows_invoice, page_number, filter_data[\"header_text\"])\n print(\"Adding page \", str(page_number))\n pdf.append_page(page)\n page_number += 1\n\n # We write out the latest PDF so that we can watch the file change with `evince` or similar viewer even though\n # parameters may differ from run to run.\n with open(\"latest-invoice.pdf\", \"wb\") as debug_file_handle:\n PDF.dumps(debug_file_handle, pdf)\n debug_file_handle.close()\n\n with open(filter_data[\"pdf_filename\"], \"wb\") as users_file_handle:\n PDF.dumps(users_file_handle, pdf)\n users_file_handle.close()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":19820,"content_sha256":"b7309ddb0a4f637e8ad4bb168e7d368bba3be9ebcd3fdda9f33e0fc5b8e8dc4b"},{"filename":"pdf/watcher.sh","content":"#!/usr/bin/env bash\n\n# This little utility notices when the timestamps change in this directory\n# and then reruns the pdf generator script. If you have something like\n# evince running that reloads the PDF when it changes this can be very\n# convenient.\n\nwhile true\ndo\n watch -g ls -l *.py *.sh *.png *.pdf; echo \"Something changed. Sleeping 1...\"; sleep 1; vastai generate pdf-invoices\ndone","content_type":"application/x-sh; charset=utf-8","language":"bash","size":390,"content_sha256":"d7e360bad742efdbfb3a9f5d215eae218ae2cd1deedf91b356a9350ef88e9c78"},{"filename":"sdk.py","content":"\"\"\"Main VastAI SDK class that wraps all api/ module functions.\"\"\"\n\nfrom __future__ import annotations\n\nimport warnings\nfrom typing import Dict, List, Optional, Union\n\nfrom vastai._base import _resolve_api_key, _APIKEY_SENTINEL\nfrom vastai.api.client import VastClient\nfrom vastai.api.query import (parse_order, parse_query, benchmarks_fields,\n fix_date_fields, invoices_fields,\n offers_mult, templates_fields,\n vol_offers_fields)\nfrom vastai.api import instances, offers, machines, teams, keys, endpoints, billing, storage, clusters, auth, deployments\n\n\ndef benchmarks_query(query):\n \"\"\"Parse a benchmarks query string into the filter dict the API wants.\"\"\"\n if isinstance(query, str):\n query = parse_query(query, {}, benchmarks_fields)\n return fix_date_fields(query, [\"last_update\"])\n return query\n\n\ndef volume_query(query):\n \"\"\"Parse a volume query string; the api layer seeds the defaults.\"\"\"\n if isinstance(query, str):\n return parse_query(query, {}, vol_offers_fields, {}, offers_mult)\n return query\n\n\ndef template_fields(*, runtype=None, use_ssh=None, jup_direct=None,\n ssh_direct=None, use_jupyter_lab=None,\n docker_login_repo=None, extra_filters=None,\n readme_visible=None, private=None, **flags) -> dict:\n \"\"\"Friendly template flags, with the raw api fields taking precedence.\n\n update_template forwarded **kwargs straight to the api layer, so the raw\n field names are what every working script uses today -- the friendly ones\n all raised TypeError. Both spellings have to keep working.\n \"\"\"\n fields = offers.template_fields_from_flags(**flags)\n overrides = {\n \"runtype\": runtype, \"use_ssh\": use_ssh, \"jup_direct\": jup_direct,\n \"ssh_direct\": ssh_direct, \"use_jupyter_lab\": use_jupyter_lab,\n \"docker_login_repo\": docker_login_repo, \"extra_filters\": extra_filters,\n \"readme_visible\": readme_visible, \"private\": private,\n }\n fields.update({k: v for k, v in overrides.items() if v is not None})\n return fields\n\n\nclass VastAI:\n \"\"\"High-level SDK for the Vast.ai platform.\n\n Wraps the low-level api/ modules behind a single object with one method\n per API operation. Every method delegates to the corresponding module\n function, passing ``self.client`` (a :class:`VastClient`) along with any\n caller-supplied arguments.\n\n Args:\n api_key: Vast.ai API key. When omitted, the key is resolved from\n ``VAST_API_KEY``, then ``$XDG_CONFIG_HOME/vastai/vast_api_key``\n (falling back to ``~/.config/vastai/vast_api_key``), then the\n legacy ``~/.vast_api_key``.\n server_url: Base URL of the Vast.ai API server.\n retry: Number of retries on transient HTTP errors.\n raw: If *True*, return raw JSON dicts instead of formatted output.\n explain: If *True*, print request details for debugging.\n quiet: If *True*, suppress informational output.\n curl: If *True*, print equivalent curl commands instead of executing.\n \"\"\"\n\n def __init__(\n self,\n api_key: Optional[str] = None,\n server_url: Optional[str] = None,\n retry: int = 3,\n raw: bool = False,\n explain: bool = False,\n quiet: bool = False,\n curl: bool = False,\n ):\n resolved_key = _resolve_api_key(_APIKEY_SENTINEL if api_key is None else api_key)\n self.client = VastClient(resolved_key, server_url, retry, explain, curl)\n self.raw = raw\n self.quiet = quiet\n\n # ------------------------------------------------------------------\n # Instance methods\n # ------------------------------------------------------------------\n\n def show_instances(self) -> list[dict]:\n \"\"\"Return all of the authenticated user's instances as a flat list.\"\"\"\n return instances.show_instances(self.client)\n\n def show_instances_v1(self, params: dict) -> dict:\n \"\"\"Return instances using the paginated v1 API; for filtering, sorting, and manual pagination.\"\"\"\n return instances.show_instances_v1(self.client, params)\n\n def show_instance_filters(self) -> list:\n \"\"\"Return distinct filterable values for instances.\"\"\"\n return instances.show_instance_filters(self.client)\n\n def show_instance(self, id: int) -> Optional[dict]:\n \"\"\"Return details of a single instance.\"\"\"\n return instances.show_instance(self.client, id)\n\n def create_instance(\n self, id: int, image: Optional[str] = None, disk: float = 10,\n env: Optional[Union[str, dict]] = None, price: Optional[float] = None,\n label: Optional[str] = None, extra: Optional[str] = None,\n onstart_cmd: Optional[str] = None, login: Optional[str] = None,\n python_utf8: bool = False, lang_utf8: bool = False,\n jupyter_lab: bool = False, jupyter_dir: Optional[str] = None,\n force: bool = False, cancel_unavail: bool = False,\n template_hash: Optional[str] = None, user: Optional[str] = None,\n runtype: Optional[str] = None, args: Optional[list] = None,\n volume_info: Optional[dict] = None,\n bid_price: Optional[float] = None, ssh: bool = False,\n jupyter: bool = False, direct: bool = False,\n create_volume: Optional[int] = None, link_volume: Optional[int] = None,\n volume_size: Optional[float] = None, mount_path: Optional[str] = None,\n volume_label: Optional[str] = None,\n ) -> dict:\n \"\"\"Create a new instance from a contract offer ID.\n\n ``runtype`` is left out of the request when neither it nor the\n ssh/jupyter/direct flags are given, so the backend picks the default.\n \"\"\"\n return instances.create_instance(\n self.client, id,\n image=image, disk=disk, env=env, price=price, bid_price=bid_price,\n label=label, extra=extra, onstart_cmd=onstart_cmd, login=login,\n python_utf8=python_utf8, lang_utf8=lang_utf8, jupyter_lab=jupyter_lab,\n jupyter_dir=jupyter_dir, force=force, cancel_unavail=cancel_unavail,\n template_hash=template_hash, user=user, runtype=runtype, ssh=ssh,\n jupyter=jupyter, direct=direct, args=args, volume_info=volume_info,\n create_volume=create_volume, link_volume=link_volume,\n volume_size=volume_size, mount_path=mount_path, volume_label=volume_label,\n )\n\n def destroy_instance(self, id: int) -> dict:\n \"\"\"Destroy an instance.\"\"\"\n return instances.destroy_instance(self.client, id)\n\n def start_instance(self, id: int) -> dict:\n \"\"\"Start a stopped instance.\"\"\"\n return instances.start_instance(self.client, id)\n\n def stop_instance(self, id: int) -> dict:\n \"\"\"Stop a running instance.\"\"\"\n return instances.stop_instance(self.client, id)\n\n def reboot_instance(self, id: int) -> dict:\n \"\"\"Reboot an instance.\"\"\"\n return instances.reboot_instance(self.client, id)\n\n def recycle_instance(self, id: int) -> dict:\n \"\"\"Recycle an instance.\"\"\"\n return instances.recycle_instance(self.client, id)\n\n def label_instance(self, id: int, label: str) -> dict:\n \"\"\"Set a label on an instance.\"\"\"\n return instances.label_instance(self.client, id, label)\n\n def prepay_instance(self, id: int, amount: float) -> dict:\n \"\"\"Prepay for an instance.\"\"\"\n return instances.prepay_instance(self.client, id, amount)\n\n def change_bid(self, id: int, price: Optional[float] = None) -> dict:\n \"\"\"Change the bid price for an instance.\"\"\"\n return instances.change_bid(self.client, id, price=price)\n\n def show_pending_price_increases(self) -> dict:\n \"\"\"Return the pending-price-increase envelope for the authenticated user.\n\n See ``vastai.api.price_increase.list_pending`` for the response shape.\n \"\"\"\n from vastai.api import price_increase\n return price_increase.list_pending(self.client)\n\n def accept_price_increase(self, instance_id: int) -> dict:\n \"\"\"Accept the pending price increase on ``instance_id``.\n\n Resolves the instance's pending row via the pending-list endpoint\n and PUTs the per-row accept. The new prices and extended end date\n apply only after the contract's current ``end_date``.\n\n Raises ``LookupError`` when no pending row matches the instance id.\n \"\"\"\n from vastai.api import price_increase\n match = price_increase.resolve_instance_to_pending(self.client, instance_id)\n return price_increase.accept(self.client, match[\"pending_price_increase_id\"])\n\n def reject_price_increase(self, instance_id: int) -> dict:\n \"\"\"Reject the pending price increase on ``instance_id``.\n\n Tombstones the matching pending row on the backend; no cutover\n follows. Raises ``LookupError`` when no pending row matches.\n \"\"\"\n from vastai.api import price_increase\n match = price_increase.resolve_instance_to_pending(self.client, instance_id)\n return price_increase.reject(self.client, match[\"pending_price_increase_id\"])\n\n def execute(self, id: int, command: str) -> dict:\n \"\"\"Execute a command on an instance.\"\"\"\n return instances.execute(self.client, id, command)\n\n def logs(\n self,\n instance_id: int,\n tail: Optional[str] = None,\n filter: Optional[str] = None,\n daemon_logs: bool = False,\n ) -> str:\n \"\"\"Retrieve logs for an instance.\"\"\"\n return instances.logs(self.client, instance_id, tail=tail, filter=filter, daemon_logs=daemon_logs)\n\n def ssh_url(self, id: int) -> str:\n \"\"\"Get the SSH URL for an instance.\"\"\"\n inst = instances.show_instance(self.client, id)\n if isinstance(inst, list):\n inst = inst[0] if inst else {}\n if not isinstance(inst, dict):\n return \"\"\n port = inst.get(\"ssh_port\") or (inst.get(\"ports\") or {}).get(\"22/tcp\", [{}])[0].get(\"HostPort\")\n ip = inst.get(\"ssh_host\") or inst.get(\"public_ipaddr\")\n return f\"ssh://root@{ip}:{port}\" if port and ip else \"\"\n\n def scp_url(self, id: int) -> str:\n \"\"\"Get the SCP URL for an instance.\"\"\"\n inst = instances.show_instance(self.client, id)\n if isinstance(inst, list):\n inst = inst[0] if inst else {}\n if not isinstance(inst, dict):\n return \"\"\n port = inst.get(\"ssh_port\") or (inst.get(\"ports\") or {}).get(\"22/tcp\", [{}])[0].get(\"HostPort\")\n ip = inst.get(\"ssh_host\") or inst.get(\"public_ipaddr\")\n return f\"scp://root@{ip}:{port}\" if port and ip else \"\"\n\n def take_snapshot(self, instance_id, **kwargs) -> dict:\n \"\"\"Take a container snapshot and push to a registry.\"\"\"\n return instances.take_snapshot(self.client, instance_id, **kwargs)\n\n # ------------------------------------------------------------------\n # Deployment methods\n # ------------------------------------------------------------------\n\n def show_deployments(self) -> list[dict]:\n \"\"\"Return all deployments for the authenticated user.\"\"\"\n return deployments.show_deployments(self.client)\n\n def show_deployment(self, id: int) -> dict:\n \"\"\"Return details of a single deployment.\"\"\"\n return deployments.show_deployment(self.client, id)\n\n def show_deployment_versions(self, id: int) -> list:\n \"\"\"Return version history for a deployment.\"\"\"\n return deployments.show_deployment_versions(self.client, id)\n\n def delete_deployment(self, id: int) -> dict:\n \"\"\"Delete a deployment.\"\"\"\n return deployments.delete_deployment(self.client, id)\n\n # ------------------------------------------------------------------\n # Offer / search methods\n # ------------------------------------------------------------------\n\n def search_offers(\n self,\n query=None,\n type: str = \"on-demand\",\n order: str = \"score-\",\n limit: Optional[int] = None,\n storage: float = 5.0,\n no_default: bool = False,\n ) -> list:\n \"\"\"Search for GPU offers.\n\n Args:\n query: Query string (e.g. \"num_gpus=1 gpu_name=RTX_4090\") or pre-parsed dict.\n type: One of \"on-demand\", \"reserved\", or \"bid\".\n order: Comma-separated sort fields, e.g. \"score-\" or \"dph_total\".\n limit: Max results.\n storage: Allocated storage in GiB for pricing.\n no_default: Skip default filters (verified, rentable, etc.).\n \"\"\"\n from vastai.api.query import parse_query, offers_fields, offers_alias, offers_mult\n from vastai.utils import preprocess_search_query, postprocess_search_results\n\n # Expand georegion/chunked directives before parsing.\n # Seed defaults before parsing so `field=any` correctly removes them,\n # matching CLI behavior. Pass no_default=True to the helper afterward\n # so it does not reapply the same defaults.\n georegion_active, chunked = False, False\n defaults_applied = False\n if isinstance(query, str):\n georegion_active, chunked, query = preprocess_search_query(query)\n base = {} if no_default else {\n \"verified\": {\"eq\": True}, \"external\": {\"eq\": False},\n \"rentable\": {\"eq\": True},\n }\n query = parse_query(query, base, offers_fields, offers_alias, offers_mult)\n defaults_applied = True\n\n # Parse order string into list\n order_list = parse_order(order)\n\n results = offers.search_offers(\n self.client, query=query, offer_type=type, order=order_list,\n limit=limit, storage=storage,\n no_default=(no_default or defaults_applied),\n )\n\n if isinstance(results, list):\n results = postprocess_search_results(results, georegion_active=georegion_active, chunked=chunked)\n\n return results\n\n def search_templates(self, query: Optional[Union[str, dict]] = None) -> list[dict]:\n \"\"\"Search for templates.\n\n ``query`` takes the same string syntax as ``vastai search templates``\n (e.g. \"name=pytorch\"), or an already-parsed filter dict.\n \"\"\"\n if isinstance(query, str):\n query = parse_query(query, {}, templates_fields)\n query = fix_date_fields(query, [\"created_at\", \"recent_create_date\"])\n return offers.search_templates(self.client, query=query)\n\n def search_benchmarks(self, query: Optional[Union[str, dict]] = None,\n order: Optional[list] = None,\n limit: Optional[int] = None,\n after_token: Optional[str] = None,\n all_pages: bool = False) -> list[dict]:\n \"\"\"Search for benchmarks.\n\n Returns a single page (the backend caps it at 200 rows) so a broad query\n doesn't scan the whole table. Pass all_pages=True to follow pagination\n and get every matching row, or use search_benchmarks_v1() to page\n manually via next_token.\n \"\"\"\n query = benchmarks_query(query)\n if all_pages:\n return offers.search_benchmarks(self.client, query=query, order=order,\n limit=limit, after_token=after_token)\n params = offers.benchmarks_query_args(query=query, order=order,\n limit=limit, after_token=after_token)\n data = offers.search_benchmarks_v1(self.client, params)\n return data.get(\"benchmarks\") or []\n\n def search_benchmarks_v1(self, params: dict) -> dict:\n \"\"\"Return benchmarks using the paginated API; for manual pagination.\"\"\"\n return offers.search_benchmarks_v1(self.client, params)\n\n def search_volumes(\n self, query: Optional[Union[str, dict]] = None,\n order: Optional[Union[str, list]] = None, limit: Optional[int] = None,\n storage: float = 1.0, no_default: bool = False,\n ) -> list[dict]:\n \"\"\"Search for volume offers.\n\n ``query`` takes the same string syntax as ``vastai search volumes``,\n or an already-parsed filter dict.\n \"\"\"\n return offers.search_volumes(\n self.client, query=volume_query(query),\n order=parse_order(order), limit=limit, storage=storage,\n no_default=no_default,\n )\n\n def search_network_volumes(\n self, query: Optional[Union[str, dict]] = None,\n order: Optional[Union[str, list]] = None, limit: Optional[int] = None,\n storage: float = 1.0, no_default: bool = False,\n ) -> list[dict]:\n \"\"\"Search for network volume offers.\n\n ``query`` takes the same string syntax as\n ``vastai search network volumes``, or an already-parsed filter dict.\n \"\"\"\n return offers.search_network_volumes(\n self.client, query=volume_query(query),\n order=parse_order(order), limit=limit, storage=storage,\n no_default=no_default,\n )\n\n def search_invoices(self, query: Optional[Union[str, dict]] = None) -> list[dict]:\n \"\"\"Search for invoices.\n\n ``query`` takes the same string syntax as ``vastai search invoices``,\n or an already-parsed filter dict.\n \"\"\"\n if isinstance(query, str):\n query = parse_query(query, {}, invoices_fields)\n query = fix_date_fields(query, [\"when\", \"paid_on\", \"payment_expected\"])\n return offers.search_invoices(self.client, query=query)\n\n def search_offers_new(\n self,\n query=None,\n type: str = \"on-demand\",\n order: str = \"score-\",\n limit: Optional[int] = None,\n storage: float = 5.0,\n no_default: bool = False,\n ) -> list:\n \"\"\"Search for GPU offers using the new /search/asks/ endpoint.\n\n Args:\n query: Query string (e.g. \"num_gpus=1 gpu_name=RTX_4090\") or pre-parsed dict.\n type: One of \"on-demand\", \"reserved\", or \"bid\".\n order: Comma-separated sort fields, e.g. \"score-\" or \"dph_total\".\n limit: Max results.\n storage: Allocated storage in GiB for pricing.\n no_default: Skip default filters (verified, rentable, etc.).\n \"\"\"\n from vastai.api.query import parse_query, offers_fields, offers_alias, offers_mult\n from vastai.utils import preprocess_search_query, postprocess_search_results\n\n georegion_active, chunked = False, False\n defaults_applied = False\n if isinstance(query, str):\n georegion_active, chunked, query = preprocess_search_query(query)\n base = {} if no_default else {\n \"verified\": {\"eq\": True}, \"external\": {\"eq\": False},\n \"rentable\": {\"eq\": True},\n }\n query = parse_query(query, base, offers_fields, offers_alias, offers_mult)\n defaults_applied = True\n\n order_list = parse_order(order)\n\n results = offers.search_offers_new(\n self.client, query=query, offer_type=type, order=order_list,\n limit=limit, storage=storage,\n no_default=(no_default or defaults_applied),\n )\n\n if isinstance(results, list):\n results = postprocess_search_results(results, georegion_active=georegion_active, chunked=chunked)\n\n return results\n\n def launch_instance(\n self, gpu_name: str, num_gpus: str, image: str,\n region: Optional[str] = None, disk: float = 10, order: str = \"score-\",\n limit: Optional[int] = None, env: Optional[Union[str, dict]] = None,\n label: Optional[str] = None, extra: Optional[str] = None,\n onstart_cmd: Optional[str] = None, login: Optional[str] = None,\n python_utf8: bool = False, lang_utf8: bool = False,\n jupyter_lab: bool = False, jupyter_dir: Optional[str] = None,\n cancel_unavail: bool = False, template_hash: Optional[str] = None,\n runtype: Optional[str] = None, args: Optional[list] = None,\n query: Optional[dict] = None,\n ssh: bool = False, jupyter: bool = False, direct: bool = False,\n ) -> dict:\n \"\"\"Launch the top instance from search offers matching the given criteria.\"\"\"\n return offers.launch_instance(\n self.client, gpu_name, num_gpus, image, region=region, disk=disk,\n order=order, limit=limit, env=env, label=label, extra=extra,\n onstart_cmd=onstart_cmd, login=login, python_utf8=python_utf8,\n lang_utf8=lang_utf8, jupyter_lab=jupyter_lab, jupyter_dir=jupyter_dir,\n cancel_unavail=cancel_unavail, template_hash=template_hash,\n runtype=runtype, ssh=ssh, jupyter=jupyter, direct=direct, args=args,\n query=query,\n )\n\n # ------------------------------------------------------------------\n # Machine methods\n # ------------------------------------------------------------------\n\n def show_machines(self) -> list[dict]:\n \"\"\"Return all hosted machines.\"\"\"\n return machines.show_machines(self.client)\n\n def show_machine(self, id: int) -> dict:\n \"\"\"Return details of a single machine.\n\n The underlying ``GET /machines/{id}`` endpoint returns a one-element\n list; this wrapper unwraps it so callers get a single machine dict.\n Raises ``ValueError`` if the backend returns zero or multiple rows.\n \"\"\"\n result = machines.show_machine(self.client, id)\n if not isinstance(result, list):\n return result\n if not result:\n raise ValueError(f\"Machine {id} not found\")\n if len(result) > 1:\n raise ValueError(f\"Expected 1 machine for id={id}, got {len(result)}\")\n return result[0]\n\n def show_maints(self, ids) -> list[dict]:\n \"\"\"Show maintenance information for machines.\"\"\"\n return machines.show_maints(self.client, ids)\n\n def list_machine(self, id: int, **kwargs) -> dict:\n \"\"\"List a machine for rent with optional pricing parameters.\"\"\"\n return machines.list_machine(self.client, id, **kwargs)\n\n def list_machines(self, ids, **kwargs) -> list[dict]:\n \"\"\"List multiple machines for rent.\"\"\"\n results = []\n for mid in ids:\n results.append(machines.list_machine(self.client, mid, **kwargs))\n return results\n\n def unlist_machine(self, id: int) -> dict:\n \"\"\"Unlist a machine from being available.\"\"\"\n return machines.unlist_machine(self.client, id)\n\n def set_defjob(self, id: int, **kwargs) -> dict:\n \"\"\"Set the default job on a machine.\"\"\"\n return machines.set_defjob(self.client, id, **kwargs)\n\n def remove_defjob(self, id: int) -> dict:\n \"\"\"Remove the default job from a machine.\"\"\"\n return machines.remove_defjob(self.client, id)\n\n def set_min_bid(self, id: int, price: Optional[float] = None) -> dict:\n \"\"\"Set the minimum bid price for a machine.\"\"\"\n return machines.set_min_bid(self.client, id, price=price)\n\n def schedule_maint(self, id: int, sdate, duration, category: str = \"not provided\") -> dict:\n \"\"\"Schedule maintenance for a machine.\"\"\"\n return machines.schedule_maint(self.client, id, sdate, duration, maintenance_category=category)\n\n def cancel_maint(self, id: int) -> dict:\n \"\"\"Cancel scheduled maintenance for a machine.\"\"\"\n return machines.cancel_maint(self.client, id)\n\n def cleanup_machine(self, id: int) -> dict:\n \"\"\"Clean up a machine's configuration and resources.\"\"\"\n return machines.cleanup_machine(self.client, id)\n\n def defrag_machines(self, ids) -> dict:\n \"\"\"Defragment machines.\"\"\"\n return machines.defrag_machines(self.client, ids)\n\n def delete_machine(self, id: int) -> dict:\n \"\"\"Delete a machine if not being used by clients.\"\"\"\n return machines.delete_machine(self.client, id)\n\n def reports(self, id: int) -> list[dict]:\n \"\"\"Generate reports for a machine.\"\"\"\n return machines.reports(self.client, id)\n\n # ------------------------------------------------------------------\n # Team methods\n # ------------------------------------------------------------------\n\n def create_team(self, team_name: str) -> dict:\n \"\"\"Create a new team.\"\"\"\n return teams.create_team(self.client, team_name)\n\n def destroy_team(self) -> dict:\n \"\"\"Destroy the current team.\"\"\"\n return teams.destroy_team(self.client)\n\n def show_members(self) -> list[dict]:\n \"\"\"Show all team members.\"\"\"\n return teams.show_members(self.client)\n\n def invite_member(self, email: str, role: str) -> dict:\n \"\"\"Invite a new member to the team.\"\"\"\n return teams.invite_member(self.client, email, role)\n\n def remove_member(self, id: int) -> dict:\n \"\"\"Remove a member from the team.\"\"\"\n return teams.remove_member(self.client, id)\n\n def show_team_roles(self) -> list[dict]:\n \"\"\"Show all team roles.\"\"\"\n return teams.show_team_roles(self.client)\n\n def show_team_role(self, name: str) -> dict:\n \"\"\"Show details of a specific team role.\"\"\"\n return teams.show_team_role(self.client, name)\n\n def create_team_role(self, name: str, permissions) -> dict:\n \"\"\"Create a new team role.\"\"\"\n return teams.create_team_role(self.client, name, permissions)\n\n def remove_team_role(self, name: str) -> dict:\n \"\"\"Remove a team role.\"\"\"\n return teams.remove_team_role(self.client, name)\n\n # ------------------------------------------------------------------\n # SSH / API key methods\n # ------------------------------------------------------------------\n\n def show_ssh_keys(self) -> list[dict]:\n \"\"\"Show all SSH keys.\"\"\"\n return keys.show_ssh_keys(self.client)\n\n def create_ssh_key(self, ssh_key: Optional[str] = None) -> dict:\n \"\"\"Create a new SSH key.\"\"\"\n return keys.create_ssh_key(self.client, ssh_key=ssh_key)\n\n def delete_ssh_key(self, id: int) -> dict:\n \"\"\"Delete an SSH key.\"\"\"\n return keys.delete_ssh_key(self.client, id)\n\n def attach_ssh(self, instance_id: int, ssh_key: str) -> dict:\n \"\"\"Attach an SSH key to an instance.\"\"\"\n return keys.attach_ssh(self.client, instance_id, ssh_key)\n\n def detach_ssh(self, instance_id: int, ssh_key_id: str) -> dict:\n \"\"\"Detach an SSH key from an instance.\"\"\"\n return keys.detach_ssh(self.client, instance_id, ssh_key_id)\n\n def show_api_keys(self) -> list[dict]:\n \"\"\"Return all API keys associated with the account.\n\n The underlying ``GET /auth/apikeys/`` endpoint returns an envelope dict\n ``{\"apikeys\": [...]}``; this wrapper unwraps it so callers get a plain\n list of API key dicts.\n \"\"\"\n result = keys.show_api_keys(self.client)\n if isinstance(result, dict) and \"apikeys\" in result:\n return result[\"apikeys\"]\n return result\n\n def show_api_key(self, id: int) -> dict:\n \"\"\"Show details of an API key.\"\"\"\n return keys.show_api_key(self.client, id)\n\n def create_api_key(self, name, permissions, key_params=None) -> dict:\n \"\"\"Create a new API key.\"\"\"\n return keys.create_api_key(self.client, name, permissions, key_params=key_params)\n\n def delete_api_key(self, id: int) -> dict:\n \"\"\"Delete an API key.\"\"\"\n return keys.delete_api_key(self.client, id)\n\n def reset_api_key(self) -> dict:\n \"\"\"Reset the API key.\"\"\"\n return keys.reset_api_key(self.client)\n\n def update_ssh_key(self, id: int, ssh_key: str) -> dict:\n \"\"\"Update an SSH key.\"\"\"\n return keys.update_ssh_key(self.client, id, ssh_key)\n\n # ------------------------------------------------------------------\n # Endpoint methods\n # ------------------------------------------------------------------\n\n def show_endpoints(self) -> list[dict]:\n \"\"\"Show all serverless endpoints.\"\"\"\n return endpoints.show_endpoints(self.client)\n\n def create_endpoint(self, **kwargs) -> dict:\n \"\"\"Create a new serverless endpoint.\"\"\"\n return endpoints.create_endpoint(self.client, **kwargs)\n\n def delete_endpoint(self, id: int) -> dict:\n \"\"\"Delete a serverless endpoint.\"\"\"\n return endpoints.delete_endpoint(self.client, id)\n\n def get_endpt_logs(self, id: int, level: int = 1, tail: Optional[int] = None) -> dict:\n \"\"\"Fetch logs for a serverless endpoint.\"\"\"\n return endpoints.get_endpt_logs(self.client, id, level=level, tail=tail)\n\n def show_workergroups(self) -> list[dict]:\n \"\"\"Show all worker groups.\"\"\"\n return endpoints.show_workergroups(self.client)\n\n def create_workergroup(self, **kwargs) -> dict:\n \"\"\"Create a new autoscale worker group.\"\"\"\n return endpoints.create_workergroup(self.client, **kwargs)\n\n def delete_workergroup(self, id: int) -> dict:\n \"\"\"Delete a worker group.\"\"\"\n return endpoints.delete_workergroup(self.client, id)\n\n def get_wrkgrp_logs(self, id: int, level: int = 1, tail: Optional[int] = None) -> dict:\n \"\"\"Fetch logs for a worker group.\"\"\"\n return endpoints.get_wrkgrp_logs(self.client, id, level=level, tail=tail)\n\n def get_endpoint_workers(self, id: int):\n \"\"\"List workers under a given endpoint, with live status and measured_perf.\"\"\"\n return endpoints.get_endpoint_workers(self.client, id)\n\n # ------------------------------------------------------------------\n # Billing methods\n # ------------------------------------------------------------------\n\n def show_invoices(\n self, start_date: Optional[str] = None, end_date: Optional[str] = None,\n only_charges: bool = False, only_credits: bool = False,\n ) -> dict:\n \"\"\"Show invoice details (deprecated; use show_invoices_v1).\n\n Returns dict with 'invoices' list and 'current' charges.\n \"\"\"\n warnings.warn(\n \"VastAI.show_invoices() is deprecated; use VastAI.show_invoices_v1(**params) for the paginated v1 API.\",\n DeprecationWarning, stacklevel=2,\n )\n return billing.show_invoices(\n self.client, start_date=start_date, end_date=end_date,\n only_charges=only_charges, only_credits=only_credits,\n )\n\n def show_invoices_v1(self, **kwargs) -> dict:\n \"\"\"Get billing history reports with advanced filtering and pagination.\"\"\"\n return billing.show_invoices_v1(self.client, kwargs)\n\n def show_earnings(\n self, start_date: Optional[str] = None, end_date: Optional[str] = None,\n machine_id: Optional[int] = None,\n ) -> list[dict]:\n \"\"\"Show earnings information.\"\"\"\n return billing.show_earnings(\n self.client, start_date=start_date, end_date=end_date,\n machine_id=machine_id,\n )\n\n def show_deposit(self, id: int) -> dict:\n \"\"\"Show deposit details.\"\"\"\n return billing.show_deposit(self.client, id)\n\n def show_user(self) -> dict:\n \"\"\"Show current user details.\"\"\"\n return billing.show_user(self.client)\n\n def set_user(self, params) -> dict:\n \"\"\"Set user parameters.\"\"\"\n return billing.set_user(self.client, params)\n\n def show_subaccounts(self) -> list[dict]:\n \"\"\"Show all subaccounts.\"\"\"\n return billing.show_subaccounts(self.client)\n\n def create_subaccount(\n self,\n email: str,\n username: str,\n password: str,\n type: Optional[str] = None,\n ) -> dict:\n \"\"\"Create a new subaccount.\"\"\"\n host_only = type is not None and type.lower() == \"host\"\n return billing.create_subaccount(self.client, email, username, password, host_only=host_only)\n\n def show_ipaddrs(self) -> list[dict]:\n \"\"\"Show IP addresses.\"\"\"\n return billing.show_ipaddrs(self.client)\n\n def fetch_contracts(self, label: Optional[str] = None, contract_ids: Optional[list] = None) -> list:\n \"\"\"Fetch contracts, optionally filtered by label.\"\"\"\n return billing.fetch_contracts(self.client, label=label, contract_ids=contract_ids)\n\n # ------------------------------------------------------------------\n # Storage methods\n # ------------------------------------------------------------------\n\n def copy(self, src: str, dst: str) -> dict:\n \"\"\"Copy files between instances, volumes, cloud services, and local.\n\n Supported location formats:\n\n - ``[instance_id:]path`` legacy format, still supported\n - ``C.instance_id:path`` container copy format\n - ``cloud_service:path`` cloud service format\n - ``cloud_service.id:path`` cloud service with ID\n - ``local:path`` explicit local path\n - ``V.volume_id:path`` volume copy\n\n Volume copy is currently only supported for copying to other volumes,\n instances, or cloud services, not local.\n\n Do not copy to /root or / as a destination directory, as this can mess\n up the permissions on the instance ssh folder, breaking future copy\n operations (they use ssh authentication). See\n https://vast.ai/docs/gpu-instances/data-movement#constraints for more\n information about constraints.\n\n Args:\n src: Source in vast URL format, e.g. \"instance_id:/path\" or just \"/local/path\".\n dst: Destination in vast URL format.\n\n Examples:\n vast.copy(\"6003036:/workspace/\", \"6003038:/workspace/\")\n vast.copy(\"C.11824:/data/test\", \"local:data/test\")\n vast.copy(\"local:data/test\", \"C.11824:/data/test\")\n vast.copy(\"drive:/folder/file.txt\", \"C.6003036:/workspace/\")\n vast.copy(\"s3.101:/data/\", \"C.6003036:/workspace/\")\n vast.copy(\"V.1234:/file\", \"C.5678:/workspace/\")\n vast.copy(\"V.1234:/file\", \"s3.101:/workspace/\")\n \"\"\"\n from vastai.utils import parse_vast_url\n src_id, src_path = parse_vast_url(src)\n dst_id, dst_path = parse_vast_url(dst)\n return storage.copy(self.client, src_id, dst_id, src_path, dst_path)\n\n def cancel_copy(self, dst_id) -> dict:\n \"\"\"Cancel a file copy operation.\"\"\"\n return storage.cancel_copy(self.client, dst_id)\n\n def cancel_sync(self, dst_id) -> dict:\n \"\"\"Cancel a file sync operation.\"\"\"\n return storage.cancel_sync(self.client, dst_id)\n\n def cloud_copy(\n self, src: str, dst: str = \"/workspace\", instance: Optional[str] = None,\n connection: Optional[str] = None, transfer: str = \"Instance to Cloud\",\n dry_run: bool = False, size_only: bool = False,\n ignore_existing: bool = False, update: bool = False,\n delete_excluded: bool = False, flags: Optional[list] = None,\n ) -> dict:\n \"\"\"Copy files between cloud and instance.\n\n Scheduling (the CLI's --schedule/--day/--hour) is not available here;\n it posts to a different endpoint than this one.\n \"\"\"\n if flags is None:\n flags = storage.rclone_flags(\n dry_run=dry_run, size_only=size_only,\n ignore_existing=ignore_existing, update=update,\n delete_excluded=delete_excluded,\n )\n return storage.cloud_copy(\n self.client, src=src, dst=dst, instance=instance,\n connection=connection, transfer=transfer, flags=flags,\n )\n\n def clone_volume(self, source: int, dest: int, **kwargs) -> dict:\n \"\"\"Clone an existing volume.\"\"\"\n return storage.clone_volume(self.client, source, dest, **kwargs)\n\n def show_volumes(self, type: str = \"all\") -> list[dict]:\n \"\"\"Show stats on owned volumes.\"\"\"\n return storage.show_volumes(self.client, type=type)\n\n def show_instance_backups(self, contract_id: Optional[int] = None) -> list[dict]:\n \"\"\"List the backups Vast has taken of your instances and volumes.\"\"\"\n return storage.show_instance_backups(self.client, contract_id=contract_id)\n\n def instance_backup_files(self, contract_id: int) -> dict:\n \"\"\"Objects stored for one backup, plus a short-lived token to fetch them.\"\"\"\n return storage.instance_backup_files(self.client, contract_id)\n\n def create_volume(self, id: int, size: float = 15, name: Optional[str] = None) -> dict:\n \"\"\"Create a new volume from an offer ID.\"\"\"\n return storage.create_volume(self.client, id, size=size, name=name)\n\n def delete_volume(self, id: int) -> dict:\n \"\"\"Delete a volume.\"\"\"\n return storage.delete_volume(self.client, id)\n\n def list_volume(self, id: int, **kwargs) -> dict:\n \"\"\"List disk space for rent as a volume.\"\"\"\n return storage.list_volume(self.client, id, **kwargs)\n\n def unlist_volume(self, id: int) -> dict:\n \"\"\"Unlist a volume offer.\"\"\"\n return storage.unlist_volume(self.client, id)\n\n def create_network_volume(self, id: int, size: float = 15, name: Optional[str] = None) -> dict:\n \"\"\"Create a new network volume.\"\"\"\n return storage.create_network_volume(self.client, id, size=size, name=name)\n\n def list_network_volume(self, disk_id: int, **kwargs) -> dict:\n \"\"\"List disk space for rent as a network volume.\"\"\"\n return storage.list_network_volume(self.client, disk_id, **kwargs)\n\n def unlist_network_volume(self, id: int) -> dict:\n \"\"\"Unlist a network volume offer.\"\"\"\n return storage.unlist_network_volume(self.client, id)\n\n def show_network_disks(self) -> dict:\n \"\"\"Show network disks associated with your account.\"\"\"\n return storage.show_network_disks(self.client)\n\n def add_network_disk(self, machines: List[int], mount_point: str, disk_id: Optional[int] = None) -> dict:\n \"\"\"Add a network disk to a physical cluster.\"\"\"\n return storage.add_network_disk(self.client, machines, mount_point, disk_id=disk_id)\n\n def show_connections(self) -> list[dict]:\n \"\"\"Show all connections.\"\"\"\n return storage.show_connections(self.client)\n\n # ------------------------------------------------------------------\n # Cluster methods\n # ------------------------------------------------------------------\n\n def show_clusters(self) -> dict:\n \"\"\"Show clusters associated with your account.\"\"\"\n return clusters.show_clusters(self.client)\n\n def create_cluster(self, subnet: str, manager_id: int) -> dict:\n \"\"\"Create a Vast cluster.\"\"\"\n return clusters.create_cluster(self.client, subnet, manager_id)\n\n def delete_cluster(self, cluster_id: int) -> dict:\n \"\"\"Delete a cluster.\"\"\"\n return clusters.delete_cluster(self.client, cluster_id)\n\n def join_cluster(self, cluster_id: int, machine_ids: List[int]) -> dict:\n \"\"\"Join machines to a cluster.\"\"\"\n return clusters.join_cluster(self.client, cluster_id, machine_ids)\n\n def remove_machine_from_cluster(\n self,\n cluster_id: int,\n machine_id: int,\n new_manager_id: Optional[int] = None,\n ) -> dict:\n \"\"\"Remove a machine from a cluster.\"\"\"\n return clusters.remove_machine_from_cluster(\n self.client, cluster_id, machine_id, new_manager_id=new_manager_id,\n )\n\n def show_overlays(self) -> list[dict]:\n \"\"\"Show overlays associated with your account.\"\"\"\n return clusters.show_overlays(self.client)\n\n def create_overlay(self, cluster_id: int, name: str) -> dict:\n \"\"\"Create an overlay network on a physical cluster.\"\"\"\n return clusters.create_overlay(self.client, cluster_id, name)\n\n def delete_overlay(self, overlay_identifier: Optional[str] = None) -> dict:\n \"\"\"Delete an overlay and remove all associated instances.\"\"\"\n return clusters.delete_overlay(self.client, overlay_identifier=overlay_identifier)\n\n def join_overlay(self, name: str, instance_id: int) -> dict:\n \"\"\"Add an instance to an overlay network.\"\"\"\n return clusters.join_overlay(self.client, name, instance_id)\n\n # ------------------------------------------------------------------\n # Auth / account methods\n # ------------------------------------------------------------------\n\n def set_api_key(self, api_key: str) -> None:\n \"\"\"Update the API key used by this client.\"\"\"\n self.client.api_key = api_key\n\n def show_audit_logs(self) -> list[dict]:\n \"\"\"Display account audit logs.\"\"\"\n return auth.show_audit_logs(self.client)\n\n def show_env_vars(self, show_values: bool = False) -> dict:\n \"\"\"Show user environment variables.\n\n Args:\n show_values: If True, return actual values. If False, mask them.\n \"\"\"\n secrets = auth.show_env_vars(self.client)\n if not show_values and isinstance(secrets, dict):\n return {k: \"****\" for k in secrets}\n return secrets\n\n def create_env_var(self, name: str, value: str) -> dict:\n \"\"\"Create a new user environment variable.\"\"\"\n return auth.create_env_var(self.client, name, value)\n\n def delete_env_var(self, name: str) -> dict:\n \"\"\"Delete a user environment variable.\"\"\"\n return auth.delete_env_var(self.client, name)\n\n def show_scheduled_jobs(self) -> list[dict]:\n \"\"\"Show scheduled jobs for the account.\"\"\"\n return auth.show_scheduled_jobs(self.client)\n\n def create_scheduled_job(self, start_time, end_time, api_endpoint, request_method,\n request_body, frequency, instance_id, **kwargs) -> dict:\n \"\"\"Create a new scheduled job.\"\"\"\n return auth.create_scheduled_job(\n self.client, start_time, end_time, api_endpoint, request_method,\n request_body, frequency, instance_id, **kwargs,\n )\n\n def update_scheduled_job(self, id: int, request_body, **kwargs) -> dict:\n \"\"\"Update an existing scheduled job.\"\"\"\n return auth.update_scheduled_job(self.client, id, request_body, **kwargs)\n\n def delete_scheduled_job(self, id: int) -> dict:\n \"\"\"Delete a scheduled job.\"\"\"\n return auth.delete_scheduled_job(self.client, id)\n\n def create_template(\n self, name: Optional[str] = None, image: Optional[str] = None,\n image_tag: Optional[str] = None, href: Optional[str] = None,\n repo: Optional[str] = None, login: Optional[str] = None,\n env: Optional[str] = None, ssh: bool = False, jupyter: bool = False,\n direct: bool = False, jupyter_dir: Optional[str] = None,\n jupyter_lab: bool = False, onstart_cmd: Optional[str] = None,\n search_params: Optional[str] = None, no_default: bool = False,\n disk_space: Optional[float] = None, readme: Optional[str] = None,\n hide_readme: bool = False, desc: Optional[str] = None,\n public: bool = False,\n runtype: Optional[str] = None, use_ssh: Optional[bool] = None,\n jup_direct: Optional[bool] = None, ssh_direct: Optional[bool] = None,\n use_jupyter_lab: Optional[bool] = None,\n docker_login_repo: Optional[str] = None,\n extra_filters: Optional[dict] = None,\n readme_visible: Optional[bool] = None, private: Optional[bool] = None,\n ) -> dict:\n \"\"\"Create a new template.\n\n ``search_params`` is parsed into the template's offer filters, the same\n way ``vastai create template --search_params`` does it.\n \"\"\"\n return offers.create_template(\n self.client,\n name=name, image=image, image_tag=image_tag, href=href,\n repo=repo, env=env, onstart_cmd=onstart_cmd, jupyter_dir=jupyter_dir,\n disk_space=disk_space, readme=readme, desc=desc,\n **template_fields(\n ssh=ssh, jupyter=jupyter, direct=direct, jupyter_lab=jupyter_lab,\n login=login, hide_readme=hide_readme, public=public,\n search_params=search_params, no_default=no_default,\n runtype=runtype, use_ssh=use_ssh, jup_direct=jup_direct,\n ssh_direct=ssh_direct, use_jupyter_lab=use_jupyter_lab,\n docker_login_repo=docker_login_repo, extra_filters=extra_filters,\n readme_visible=readme_visible, private=private,\n ),\n )\n\n def delete_template(self, template_id: Optional[int] = None, hash_id: Optional[str] = None) -> dict:\n \"\"\"Delete a template by ID or hash.\"\"\"\n return offers.delete_template(self.client, template_id=template_id, hash_id=hash_id)\n\n def update_template(\n self, hash_id: str, name: Optional[str] = None, image: Optional[str] = None,\n image_tag: Optional[str] = None, href: Optional[str] = None,\n repo: Optional[str] = None, login: Optional[str] = None,\n env: Optional[str] = None, ssh: bool = False, jupyter: bool = False,\n direct: bool = False, jupyter_dir: Optional[str] = None,\n jupyter_lab: bool = False, onstart_cmd: Optional[str] = None,\n search_params: Optional[str] = None, no_default: bool = False,\n disk_space: Optional[float] = None, readme: Optional[str] = None,\n hide_readme: bool = False, desc: Optional[str] = None,\n public: bool = False,\n runtype: Optional[str] = None, use_ssh: Optional[bool] = None,\n jup_direct: Optional[bool] = None, ssh_direct: Optional[bool] = None,\n use_jupyter_lab: Optional[bool] = None,\n docker_login_repo: Optional[str] = None,\n extra_filters: Optional[dict] = None,\n readme_visible: Optional[bool] = None, private: Optional[bool] = None,\n ) -> dict:\n \"\"\"Update an existing template by creating a new version of it with a new hash_id.\n\n Raises:\n ValueError: No template with this hash_id is visible to the caller.\n \"\"\"\n return offers.update_template(\n self.client, hash_id=hash_id,\n name=name, image=image, image_tag=image_tag, href=href,\n repo=repo, env=env, onstart_cmd=onstart_cmd, jupyter_dir=jupyter_dir,\n disk_space=disk_space, readme=readme, desc=desc,\n **template_fields(\n ssh=ssh, jupyter=jupyter, direct=direct, jupyter_lab=jupyter_lab,\n login=login, hide_readme=hide_readme, public=public,\n search_params=search_params, no_default=no_default,\n runtype=runtype, use_ssh=use_ssh, jup_direct=jup_direct,\n ssh_direct=ssh_direct, use_jupyter_lab=use_jupyter_lab,\n docker_login_repo=docker_login_repo, extra_filters=extra_filters,\n readme_visible=readme_visible, private=private,\n ),\n )\n\n # ------------------------------------------------------------------\n # Batch instance methods\n # ------------------------------------------------------------------\n\n def create_instances(\n self, ids: List[int], image: Optional[str] = None, disk: float = 10,\n env: Optional[Union[str, dict]] = None, price: Optional[float] = None,\n label: Optional[str] = None, extra: Optional[str] = None,\n onstart_cmd: Optional[str] = None, login: Optional[str] = None,\n python_utf8: bool = False, lang_utf8: bool = False,\n jupyter_lab: bool = False, jupyter_dir: Optional[str] = None,\n force: bool = False, cancel_unavail: bool = False,\n template_hash: Optional[str] = None, user: Optional[str] = None,\n runtype: Optional[str] = None, args: Optional[list] = None,\n volume_info: Optional[dict] = None,\n bid_price: Optional[float] = None, ssh: bool = False,\n jupyter: bool = False, direct: bool = False,\n create_volume: Optional[int] = None, link_volume: Optional[int] = None,\n volume_size: Optional[float] = None, mount_path: Optional[str] = None,\n volume_label: Optional[str] = None,\n ) -> dict:\n \"\"\"Create multiple instances from a list of offer IDs.\"\"\"\n return instances.create_instance(\n self.client, ids,\n image=image, disk=disk, env=env, price=price, bid_price=bid_price,\n label=label, extra=extra, onstart_cmd=onstart_cmd, login=login,\n python_utf8=python_utf8, lang_utf8=lang_utf8, jupyter_lab=jupyter_lab,\n jupyter_dir=jupyter_dir, force=force, cancel_unavail=cancel_unavail,\n template_hash=template_hash, user=user, runtype=runtype, ssh=ssh,\n jupyter=jupyter, direct=direct, args=args, volume_info=volume_info,\n create_volume=create_volume, link_volume=link_volume,\n volume_size=volume_size, mount_path=mount_path, volume_label=volume_label,\n )\n\n def destroy_instances(self, ids: List[int]) -> dict:\n \"\"\"Destroy multiple instances.\"\"\"\n return instances.destroy_instance(self.client, id=ids)\n\n def start_instances(self, ids: List[int]) -> dict:\n \"\"\"Start multiple instances.\"\"\"\n return instances.start_instance(self.client, id=ids)\n\n def stop_instances(self, ids: List[int]) -> dict:\n \"\"\"Stop multiple instances.\"\"\"\n return instances.stop_instance(self.client, id=ids)\n\n # ------------------------------------------------------------------\n # Update methods\n # ------------------------------------------------------------------\n\n def update_endpoint(self, id: int, **kwargs) -> dict:\n \"\"\"Update an existing endpoint.\"\"\"\n return endpoints.update_endpoint(self.client, id=id, **kwargs)\n\n def update_env_var(self, name: str, value: str) -> dict:\n \"\"\"Update an existing user environment variable.\"\"\"\n return auth.update_env_var(self.client, name, value)\n\n def update_instance(self, id: int, **kwargs) -> dict:\n \"\"\"Update/recreate an instance from a new/updated template.\"\"\"\n return instances.update_instance(self.client, id=id, **kwargs)\n\n def update_workergroup(self, id: int, **kwargs) -> dict:\n \"\"\"Update an existing autoscale worker group.\"\"\"\n return endpoints.update_workergroup(self.client, id=id, **kwargs)\n\n def update_workers(self, id: int, cancel: bool = False) -> dict:\n \"\"\"Trigger a rolling update of all workers in a workergroup, or cancel an in-progress update.\"\"\"\n return endpoints.update_workers(self.client, id=id, cancel=cancel)\n\n def update_team_role(self, id: int, **kwargs) -> dict:\n \"\"\"Update an existing team role.\"\"\"\n return teams.update_team_role(self.client, id=id, **kwargs)\n\n # ------------------------------------------------------------------\n # TFA (Two-Factor Authentication) methods\n # ------------------------------------------------------------------\n\n def tfa_activate(self, **kwargs) -> dict:\n \"\"\"Activate a new 2FA method by verifying the code.\"\"\"\n return auth.tfa_activate(self.client, **kwargs)\n\n def tfa_delete(self, **kwargs) -> dict:\n \"\"\"Remove a 2FA method from your account.\"\"\"\n return auth.tfa_delete(self.client, **kwargs)\n\n def tfa_login(self, **kwargs) -> dict:\n \"\"\"Complete 2FA login by verifying code.\"\"\"\n return auth.tfa_login(self.client, **kwargs)\n\n def tfa_regen_codes(self, **kwargs) -> dict:\n \"\"\"Regenerate backup codes for 2FA.\"\"\"\n return auth.tfa_regen_codes(self.client, **kwargs)\n\n def tfa_resend_sms(self, **kwargs) -> dict:\n \"\"\"Resend SMS 2FA code.\"\"\"\n return auth.tfa_resend_sms(self.client, **kwargs)\n\n def tfa_send_sms(self, **kwargs) -> dict:\n \"\"\"Request a 2FA SMS verification code.\"\"\"\n return auth.tfa_send_sms(self.client, **kwargs)\n\n def tfa_status(self) -> dict:\n \"\"\"Show the current 2FA status and configured methods.\"\"\"\n return auth.tfa_status(self.client)\n\n def tfa_totp_setup(self) -> dict:\n \"\"\"Generate TOTP secret and QR code for Authenticator app setup.\"\"\"\n return auth.tfa_totp_setup(self.client)\n\n def tfa_update(self, **kwargs) -> dict:\n \"\"\"Update a 2FA method's settings.\"\"\"\n return auth.tfa_update(self.client, **kwargs)\n\n # ------------------------------------------------------------------\n # Additional billing methods\n # ------------------------------------------------------------------\n\n def generate_pdf_invoices(self, **kwargs):\n \"\"\"Generate PDF invoices based on filters.\"\"\"\n raise NotImplementedError(\"generate_pdf_invoices is not yet implemented\")\n\n def transfer_credit(self, recipient: str, amount: float) -> dict:\n \"\"\"Transfer credit to another account.\"\"\"\n return teams.transfer_credit(self.client, recipient=recipient, amount=amount)\n\n # ------------------------------------------------------------------\n # Additional methods\n # ------------------------------------------------------------------\n\n def list_volumes(self, ids, **kwargs) -> dict:\n \"\"\"List disk space for rent as volumes on multiple machines.\"\"\"\n return storage.list_volumes(self.client, ids=ids, **kwargs)\n\n def self_test_machine(self, machine_id, **kwargs):\n \"\"\"Perform a self-test on the specified machine.\"\"\"\n raise NotImplementedError(\"self_test_machine requires CLI\")\n\n # ------------------------------------------------------------------\n # Backward-compatible aliases\n # ------------------------------------------------------------------\n\n invite_team_member = invite_member\n remove_team_member = remove_member\n show_team_members = show_members\n","content_type":"text/x-python; charset=utf-8","language":"python","size":53494,"content_sha256":"20869d8aa8456e180326cb98b4d2ed92dbb96339ed01726f62bbf6efda5db59f"},{"filename":"serverless/__init__.py","content":"from .client import (\n Serverless,\n CoroutineServerless,\n ServerlessRequest,\n RequestStatus,\n Endpoint,\n Endpoint_,\n ManagedEndpoint,\n ManagedDeployment,\n Session,\n)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":193,"content_sha256":"3d6baf33aa4ab4eb9dbe96117c9511ee311095cedbd46c442acdea50fc664fd9"},{"filename":"serverless/client/__init__.py","content":"from .client import Serverless, CoroutineServerless, ServerlessRequest, _ServerlessBase\nfrom .request_status import RequestStatus\nfrom .endpoint import Endpoint, Endpoint_\nfrom .managed import ManagedEndpoint, ManagedDeployment\nfrom .session import Session\nfrom .worker import Worker as ServerlessWorker\n","content_type":"text/x-python; charset=utf-8","language":"python","size":304,"content_sha256":"0722e1d4198ed06b429089bc4571af5ae8fdbb92823b2eb23f972f0bfec6ce90"},{"filename":"serverless/client/client.py","content":"# client.py\nfrom .connection import _make_request\nfrom .endpoint import Endpoint, Endpoint_\nfrom .managed import ManagedEndpoint, ManagedDeployment\nfrom .worker import Worker\nfrom .session import Session\nfrom .request_status import RequestStatus\nfrom vastai.data.endpoint import EndpointConfig, EndpointData\nfrom vastai.data.deployment import (\n DeploymentConfig,\n DeploymentData,\n DeploymentPutResponse,\n)\nfrom vastai.data.workergroup import WorkergroupConfig\nimport asyncio\nimport aiohttp\nimport ssl\nimport os\nimport tempfile\nimport logging\nimport random\nimport time\nimport collections\nfrom abc import abstractmethod\nfrom typing import (\n Any,\n Awaitable,\n Callable,\n Coroutine,\n Deque,\n Dict,\n Generic,\n Optional,\n TypeVar,\n Union,\n List,\n)\n\nlogger = logging.getLogger(\"vastai\")\n\nR = TypeVar(\"R\", bound=Awaitable)\n\n\nclass ServerlessRequest(asyncio.Future):\n \"\"\"A fire-and-observe request to a Serverless endpoint.\n\n Composes a RequestStatus tracker for status/timestamp tracking.\n Properties delegate to the tracker for backward compatibility.\n \"\"\"\n\n def __init__(self, tracker: Optional[RequestStatus] = None):\n super().__init__()\n self.tracker = tracker if tracker is not None else RequestStatus()\n\n @property\n def status(self):\n return self.tracker.status\n\n @status.setter\n def status(self, value):\n self.tracker.status = value\n\n @property\n def create_time(self):\n return self.tracker.create_time\n\n @create_time.setter\n def create_time(self, value):\n self.tracker.create_time = value\n\n @property\n def start_time(self):\n return self.tracker.start_time\n\n @start_time.setter\n def start_time(self, value):\n self.tracker.start_time = value\n\n @property\n def complete_time(self):\n return self.tracker.complete_time\n\n @complete_time.setter\n def complete_time(self, value):\n self.tracker.complete_time = value\n\n @property\n def req_idx(self):\n return self.tracker.req_idx\n\n @req_idx.setter\n def req_idx(self, value):\n self.tracker.req_idx = value\n\n def then(self, callback) -> \"ServerlessRequest\":\n def _done(fut):\n if fut.exception() is not None:\n print(fut.exception())\n return\n callback(fut.result())\n\n self.add_done_callback(_done)\n return self\n\n\nclass _ServerlessBase(Generic[R]):\n SSL_CERT_URL = \"https://console.vast.ai/static/jvastai_root.cer\"\n VAST_WEB_URL = \"https://console.vast.ai\"\n VAST_SERVERLESS_URL = \"https://run.vast.ai\"\n\n def __init__(\n self,\n api_key: Optional[str] = os.environ.get(\"VAST_API_KEY\", None),\n *,\n debug: bool = False,\n instance: str = \"prod\",\n autoscaler_url: Optional[str] = None,\n webserver_url: str = VAST_WEB_URL,\n connection_limit: int = 500,\n default_request_timeout: float = 600.0,\n max_poll_interval: float = 5.0,\n ):\n if api_key is None or api_key == \"\":\n raise AttributeError(\n \"API key missing. Please set VAST_API_KEY in your environment variables.\"\n )\n self.api_key = api_key\n self.vast_web_url = webserver_url\n if autoscaler_url:\n self.autoscaler_url = autoscaler_url\n else:\n match instance:\n case \"prod\":\n self.autoscaler_url = \"https://run.vast.ai\"\n case \"alpha\":\n self.autoscaler_url = \"https://run-alpha.vast.ai\"\n case \"candidate\":\n self.autoscaler_url = \"https://run-candidate.vast.ai\"\n case \"local\":\n self.autoscaler_url = \"http://localhost:8080\"\n case _:\n self.autoscaler_url = \"https://run.vast.ai\"\n\n self.latencies = collections.deque(maxlen=50)\n self.debug = debug\n self.default_request_timeout = float(default_request_timeout)\n self.max_poll_interval = float(max_poll_interval)\n self.logger = logging.getLogger(self.__class__.__name__)\n\n if not self.logger.handlers:\n self.logger.addHandler(logging.NullHandler())\n\n if debug:\n handler = logging.StreamHandler()\n formatter = logging.Formatter(\n \"[%(asctime)s] %(name)s - %(levelname)s - %(message)s\"\n )\n handler.setFormatter(formatter)\n handler.setLevel(logging.DEBUG)\n\n self.logger.addHandler(handler)\n self.logger.setLevel(logging.DEBUG)\n\n # If we attach our own handler, avoid double-logging via root handlers.\n self.logger.propagate = False\n else:\n # Let the application decide if/where logs go.\n self.logger.propagate = True\n\n self.connection_limit = connection_limit\n self._session: aiohttp.ClientSession | None = None\n self._ssl_context: ssl.SSLContext | None = None\n self._transport_owner = (\n None # if set, delegate _get_session/get_ssl_context and skip close\n )\n\n async def __aenter__(self):\n await self._get_session()\n return self\n\n async def __aexit__(self, exc_type, exc, tb):\n await self.close()\n\n async def _get_session(self) -> aiohttp.ClientSession:\n if self._transport_owner is not None:\n return await self._transport_owner._get_session()\n if self._session is None or self._session.closed:\n self.logger.info(\"Started aiohttp ClientSession\")\n connector = aiohttp.TCPConnector(\n limit=self.connection_limit, ssl=await self.get_ssl_context()\n )\n self._session = aiohttp.ClientSession(connector=connector)\n return self._session\n\n def is_open(self):\n if self._transport_owner is not None:\n return self._transport_owner.is_open()\n return self._session is not None and not self._session.closed\n\n async def close(self):\n if self._transport_owner is not None:\n return # transport owned by another client\n if self._session and not self._session.closed:\n await self._session.close()\n self.logger.info(\"Closed aiohttp ClientSession\")\n\n async def get_ssl_context(self) -> ssl.SSLContext:\n \"\"\"Download Vast.ai root cert and build SSL context (cached).\"\"\"\n if self._transport_owner is not None:\n return await self._transport_owner.get_ssl_context()\n if self._ssl_context is None:\n async with aiohttp.ClientSession() as s:\n async with s.get(self.SSL_CERT_URL) as resp:\n if resp.status != 200:\n raise Exception(f\"Failed to fetch SSL cert: {resp.status}\")\n cert_bytes = await resp.read()\n\n tmpfile = tempfile.NamedTemporaryFile(delete=False, suffix=\".cer\")\n tmpfile.write(cert_bytes)\n tmpfile.close()\n\n ctx = ssl.create_default_context()\n ctx.load_verify_locations(cafile=tmpfile.name)\n # The Vast.ai root CA cert has pathlen set without keyCertSign in\n # Key Usage. OpenSSL 3.x (Python ≥3.10) rejects this by default.\n # Clearing VERIFY_X509_STRICT relaxes that single check while\n # keeping full chain-of-trust and signature verification intact.\n ctx.verify_flags &= ~ssl.VERIFY_X509_STRICT\n self.logger.info(\"Loaded Vast.ai SSL certificate\")\n\n self._ssl_context = ctx\n os.unlink(tmpfile.name)\n\n return self._ssl_context\n\n async def get_endpoint(self, name=\"\") -> Endpoint_[R]:\n endpoints = await self.get_endpoints()\n for e in endpoints:\n if e.name == name:\n return e\n raise Exception(f\"Endpoint {name} could not be found\")\n\n async def get_endpoints(self) -> list[Endpoint_[R]]:\n try:\n result = await _make_request(\n client=self,\n url=self.vast_web_url,\n route=\"/api/v0/endptjobs/\",\n api_key=self.api_key,\n params={\"client_id\": \"me\"},\n )\n except Exception as ex:\n raise Exception(f\"Failed to get endpoints:\\nReason={ex}\")\n\n if not result.get(\"ok\"):\n raise Exception(\n f\"Failed to get endpoints: HTTP {result.get('status')} - {result.get('text', '')[:512]}\"\n )\n\n response = result.get(\"json\") or {}\n endpoints = []\n for e in response.get(\"results\", []):\n data = EndpointData.from_dict(e)\n endpoints.append(Endpoint(client=self, data=data))\n self.logger.info(f\"Found {len(endpoints)} endpoints\")\n return endpoints\n\n async def fetch_endpoint(self, endpoint_id: int) -> EndpointData:\n \"\"\"Fetch full endpoint data by ID. Used internally by ManagedEndpoint.get().\"\"\"\n try:\n result = await _make_request(\n client=self,\n url=self.vast_web_url,\n route=f\"/api/v0/endptjobs/{endpoint_id}/\",\n api_key=self.api_key,\n )\n except Exception as ex:\n raise Exception(f\"Failed to fetch endpoint {endpoint_id}:\\nReason={ex}\")\n\n if not result.get(\"ok\"):\n raise Exception(\n f\"Failed to fetch endpoint {endpoint_id}: HTTP {result.get('status')} - {result.get('text', '')[:512]}\"\n )\n\n return EndpointData.from_dict(result[\"json\"][\"result\"])\n\n async def fetch_deployment(self, deployment_id: int) -> DeploymentData:\n \"\"\"Fetch full deployment data by ID. Used internally by ManagedDeployment.get().\"\"\"\n try:\n result = await _make_request(\n client=self,\n url=self.vast_web_url,\n route=f\"/api/v0/deployment/{deployment_id}/\",\n api_key=self.api_key,\n )\n except Exception as ex:\n raise Exception(f\"Failed to fetch deployment {deployment_id}:\\nReason={ex}\")\n\n if not result.get(\"ok\"):\n raise Exception(\n f\"Failed to fetch deployment {deployment_id}: HTTP {result.get('status')} - {result.get('text', '')[:512]}\"\n )\n\n return DeploymentData.from_dict(result[\"json\"][\"deployment\"])\n\n async def create_endpoint(self, config: EndpointConfig) -> ManagedEndpoint[R]:\n \"\"\"Create a new endpoint job. Returns a ManagedEndpoint (data fetched lazily).\"\"\"\n try:\n result = await _make_request(\n client=self,\n url=self.vast_web_url,\n route=\"/api/v0/endptjobs/\",\n api_key=self.api_key,\n body=config.to_dict(),\n method=\"POST\",\n )\n except Exception as ex:\n raise Exception(f\"Failed to create endpoint:\\nReason={ex}\")\n\n if not result.get(\"ok\"):\n raise Exception(\n f\"Failed to create endpoint: HTTP {result.get('status')} - {result.get('text', '')[:512]}\"\n )\n\n endpoint_id = result[\"json\"][\"result\"]\n self.logger.info(f\"Created endpoint {endpoint_id}\")\n return ManagedEndpoint(id=endpoint_id, client=self)\n\n async def put_deployment(self, config: DeploymentConfig) -> ManagedDeployment[R]:\n \"\"\"Create or update a deployment. Returns a ManagedDeployment (data fetched lazily).\"\"\"\n logger.debug(f\"putting deployment: \\n {config.to_dict()}\")\n try:\n result = await _make_request(\n client=self,\n url=self.vast_web_url,\n route=\"/api/v0/deployments/\",\n api_key=self.api_key,\n body=config.to_dict(),\n method=\"PUT\",\n )\n except Exception as ex:\n raise Exception(f\"Failed to put deployment:\\nReason={ex}\")\n\n if not result.get(\"ok\"):\n raise Exception(\n f\"Failed to put deployment: HTTP {result.get('status')} - {result.get('text', '')[:512]}\"\n )\n\n resp = DeploymentPutResponse.from_dict(result[\"json\"])\n self.logger.info(\n f\"Deployment {resp.deployment_id} ({resp.action}), endpoint {resp.endpoint_id}\"\n )\n return ManagedDeployment(\n id=resp.deployment_id,\n endpoint_id=resp.endpoint_id,\n client=self,\n put_response=resp,\n )\n\n async def create_workergroup(self, config: WorkergroupConfig) -> int:\n \"\"\"Create a workergroup (autoscale job). Returns the workergroup ID.\"\"\"\n try:\n result = await _make_request(\n client=self,\n url=self.vast_web_url,\n route=\"/api/v0/workergroups/\",\n api_key=self.api_key,\n body=config.to_dict(),\n method=\"POST\",\n )\n except Exception as ex:\n raise Exception(f\"Failed to create workergroup:\\nReason={ex}\")\n\n if not result.get(\"ok\"):\n raise Exception(\n f\"Failed to create workergroup: HTTP {result.get('status')} - {result.get('text', '')[:512]}\"\n )\n\n wg_id = result[\"json\"][\"id\"]\n self.logger.info(f\"Created workergroup {wg_id}\")\n return wg_id\n\n async def delete_workergroup(self, workergroup_id: int) -> None:\n \"\"\"Delete a workergroup (autoscale job) by ID.\"\"\"\n try:\n result = await _make_request(\n client=self,\n url=self.vast_web_url,\n route=f\"/api/v0/workergroups/{workergroup_id}/\",\n api_key=self.api_key,\n method=\"DELETE\",\n )\n except Exception as ex:\n raise Exception(\n f\"Failed to delete workergroup {workergroup_id}:\\nReason={ex}\"\n )\n\n if not result.get(\"ok\"):\n raise Exception(\n f\"Failed to delete workergroup {workergroup_id}: HTTP {result.get('status')} - {result.get('text', '')[:512]}\"\n )\n\n self.logger.info(f\"Deleted workergroup {workergroup_id}\")\n\n async def update_workers(self, workergroup_id: int) -> dict:\n \"\"\"Trigger a rolling update for all workers in a workergroup via the autoscaler.\"\"\"\n try:\n result = await _make_request(\n client=self,\n url=self.autoscaler_url,\n route=\"/update_workers/\",\n api_key=self.api_key,\n body={\"workergroup_id\": workergroup_id, \"api_key\": self.api_key},\n method=\"POST\",\n )\n except Exception as ex:\n raise Exception(\n f\"Failed to trigger worker update for workergroup {workergroup_id}:\\nReason={ex}\"\n )\n\n if not result.get(\"ok\"):\n raise Exception(\n f\"Failed to trigger worker update for workergroup {workergroup_id}: HTTP {result.get('status')} - {result.get('text', '')[:512]}\"\n )\n\n self.logger.info(f\"Triggered rolling update for workergroup {workergroup_id}\")\n return result.get(\"json\", {})\n\n async def find_workergroup_for_endpoint(self, endpoint_id: int) -> Optional[int]:\n \"\"\"Fetch all workergroups and return the first one matching endpoint_id, or None.\"\"\"\n try:\n result = await _make_request(\n client=self,\n url=self.vast_web_url,\n route=\"/api/v0/workergroups/\",\n api_key=self.api_key,\n method=\"GET\",\n )\n except Exception:\n return None\n if not result.get(\"ok\"):\n return None\n for wg in result.get(\"json\", {}).get(\"results\", []):\n if wg.get(\"endpoint_id\") == endpoint_id:\n return wg[\"id\"]\n return None\n\n async def delete_endpoint(self, endpoint_id: int) -> None:\n \"\"\"Delete an endpoint job by ID.\"\"\"\n try:\n result = await _make_request(\n client=self,\n url=self.vast_web_url,\n route=f\"/api/v0/endptjobs/{endpoint_id}/\",\n api_key=self.api_key,\n method=\"DELETE\",\n )\n except Exception as ex:\n raise Exception(f\"Failed to delete endpoint {endpoint_id}:\\nReason={ex}\")\n\n if not result.get(\"ok\"):\n raise Exception(\n f\"Failed to delete endpoint {endpoint_id}: HTTP {result.get('status')} - {result.get('text', '')[:512]}\"\n )\n\n self.logger.info(f\"Deleted endpoint {endpoint_id}\")\n\n async def get_endpoint_workers(self, endpoint: Endpoint_[R]) -> List[Worker]:\n if not isinstance(endpoint, Endpoint_):\n raise ValueError(\"endpoint must be an Endpoint\")\n\n url = f\"{self.autoscaler_url}/get_endpoint_workers/\"\n payload = {\"id\": endpoint.id, \"api_key\": self.api_key}\n\n async with self._session.post(\n url, json=payload, timeout=aiohttp.ClientTimeout(total=30)\n ) as resp:\n if resp.status != 200:\n text = await resp.text()\n raise RuntimeError(\n f\"get_endpoint_workers failed: HTTP {resp.status} - {text}\"\n )\n\n data = await resp.json(content_type=None)\n\n # If error message from authenticate_endpoint_apikey_by_id occurs, there is a possibility that\n # the endpoint's worker instances are not ready to be queried. If an error message occurs,\n # return an empty list and print the error message to the user. The endpoint get_endpoint_workers\n # should normally return a list of dictionaries containing worker instance information.\n if isinstance(data, dict):\n if \"error_msg\" in data.keys():\n self.logger.warning(\n 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.\"\n )\n return []\n\n if not isinstance(data, list):\n raise RuntimeError(\n f\"Unexpected response type (wanted list): {type(data)}\"\n )\n\n return [Worker.from_dict(item) for item in data]\n\n async def get_endpoint_session(\n self, endpoint, session_id: int, session_auth: str, timeout: float = 10.0\n ):\n try:\n result = await _make_request(\n client=self,\n url=session_auth.get(\"url\"),\n api_key=\"\",\n route=\"/session/get\",\n body={\"session_id\": session_id, \"session_auth\": session_auth},\n method=\"POST\",\n retries=1,\n timeout=timeout,\n stream=False,\n )\n\n if not result.get(\"ok\"):\n error_msg = f\"Error on /session/get: {result.get('json', {}).get('error', result.get('text', 'Unknown error'))}\"\n raise Exception(error_msg)\n\n worker_response = result.get(\"json\") or {}\n\n auth_data = worker_response.get(\"auth_data\")\n if auth_data is None:\n raise Exception(\"Missing auth_data in response\")\n\n session = Session(\n endpoint=endpoint,\n session_id=session_id,\n lifetime=worker_response.get(\"lifetime\"),\n expiration=worker_response.get(\"expiration\"),\n auth_data=auth_data,\n url=auth_data.get(\"url\"),\n )\n return session\n\n except asyncio.TimeoutError:\n raise\n except Exception as ex:\n error_msg = f\"Failed to get session {session_id}: {ex}\"\n self.logger.error(error_msg)\n raise Exception(error_msg)\n\n async def end_endpoint_session(self, session: Session, timeout: float = 10.0):\n try:\n self.logger.debug(\n f\"Attempting to end session {session.session_id} at {session.url}\"\n )\n result = await _make_request(\n client=self,\n url=session.url,\n api_key=\"\",\n route=\"/session/end\",\n body={\n \"session_id\": session.session_id,\n \"session_auth\": session.auth_data,\n },\n method=\"POST\",\n retries=1,\n timeout=timeout,\n stream=False,\n )\n\n if not result.get(\"ok\"):\n error_msg = f\"Error on /session/end: {result.get('json', {}).get('error', result.get('text', 'Unknown error'))}\"\n raise Exception(error_msg)\n\n self.logger.debug(f\"Successfully ended session {session.session_id}\")\n return\n\n except asyncio.TimeoutError:\n raise\n except Exception as ex:\n error_msg = f\"Failed to end session {session.session_id}: {ex}\"\n self.logger.error(error_msg)\n raise Exception(error_msg)\n\n async def start_endpoint_session(\n self,\n endpoint: Endpoint_[R],\n cost: int = 100,\n lifetime: float = 60,\n on_close_route: str = None,\n on_close_payload: dict = None,\n timeout: float = None,\n ) -> Session:\n try:\n session_start_response = await self.queue_endpoint_request(\n endpoint=endpoint,\n worker_route=\"/session/create\",\n worker_payload={\n \"lifetime\": lifetime,\n \"on_close_route\": on_close_route,\n \"on_close_payload\": on_close_payload,\n },\n cost=cost,\n timeout=timeout,\n worker_timeout=10.0,\n )\n if not session_start_response.get(\"ok\"):\n error_msg = f\"Error on /session/create: {session_start_response.get('json', {}).get('error', session_start_response.get('text', 'Unknown error'))}\"\n raise Exception(error_msg)\n response = session_start_response.get(\"response\")\n if response is None:\n raise Exception(\"No response from /session/create\")\n if not isinstance(response, dict):\n raise Exception(\n \"Invalid response from /session/create: expected mapping\"\n )\n session_id = response.get(\"session_id\")\n if session_id is None:\n raise Exception(\"Missing session id\")\n expiration = response.get(\"expiration\")\n url = session_start_response.get(\"url\")\n if url is None:\n raise Exception(\"Missing URL\")\n auth_data = session_start_response.get(\"auth_data\")\n if auth_data is None:\n raise Exception(\"Missing auth data\")\n return Session(\n endpoint=endpoint,\n session_id=session_id,\n lifetime=lifetime,\n expiration=expiration,\n url=url,\n auth_data=auth_data,\n )\n except asyncio.TimeoutError:\n raise\n except Exception as ex:\n error_msg = f\"Failed to create session: {ex}\"\n self.logger.error(error_msg)\n raise Exception(error_msg)\n\n async def _do_request(\n self,\n endpoint: Endpoint_[R],\n worker_route: str,\n worker_payload: dict,\n session: Optional[Session] = None,\n tracker: Optional[RequestStatus] = None,\n cost: int = 100,\n timeout: Optional[float] = None,\n worker_timeout: Optional[float] = 600,\n retry: bool = True,\n max_retries: Optional[int] = None,\n stream: bool = False,\n ) -> dict:\n \"\"\"Core request logic: route to a worker, execute, retry on failure. Returns the response dict.\"\"\"\n if tracker is None:\n tracker = RequestStatus()\n\n request_idx: int = 0\n total_attempts = 0\n start_time = time.time()\n try:\n while True:\n total_attempts += 1\n tracker.status = \"Queued\"\n worker_url = \"\"\n auth_data = {}\n session_id = None\n\n # Check total elapsed time\n if timeout is not None and (time.time() - start_time) >= timeout:\n raise asyncio.TimeoutError(\n f\"Timed out after {time.time() - start_time:.1f}s waiting for worker\"\n )\n\n if session is None:\n if request_idx == 0:\n self.logger.debug(\n f\"Sending initial route call for request_idx {request_idx}\"\n )\n else:\n self.logger.debug(\n f\"Sending retry route call for request_idx {request_idx}\"\n )\n\n route = await endpoint._route(\n cost=cost, req_idx=request_idx, timeout=60.0\n )\n\n request_idx = route.request_idx\n if request_idx:\n self.logger.debug(f\"Got request index {request_idx}\")\n else:\n self.logger.info( # usually non-fatal due to get_serverless_groups polling delay; we should fix that and make this an error again.\n \"Did not get request_idx from initial route\"\n )\n\n poll_interval = 1\n poll_elapsed = 0\n attempt = 0\n while route.status != \"READY\":\n tracker.status = \"Polling\"\n\n # Check total elapsed time\n if (\n timeout is not None\n and (time.time() - start_time) >= timeout\n ):\n raise asyncio.TimeoutError(\n f\"Timed out after {time.time() - start_time:.1f}s waiting for worker to become ready\"\n )\n\n await asyncio.sleep(poll_interval)\n poll_elapsed += poll_interval\n\n route = await endpoint._route(\n cost=cost, req_idx=request_idx, timeout=60.0\n )\n request_idx = route.request_idx or request_idx\n\n attempt += 1\n poll_interval = random.uniform(\n 0.1,\n min(\n 2 ** min(attempt, 20) + random.uniform(0, 1),\n self.max_poll_interval,\n ),\n )\n self.logger.debug(f\"Polling route, attempt {attempt}\")\n\n worker_url = route.get_url()\n auth_data = route.body\n else:\n if session.url is not None:\n worker_url = session.url\n auth_data = session.auth_data\n session_id = session.session_id\n\n payload = worker_payload\n worker_request_body = {\n \"auth_data\": auth_data,\n \"session_id\": session_id,\n \"payload\": payload,\n }\n\n self.logger.debug(\"Found worker machine, starting work\")\n if tracker.status != \"Retrying\":\n tracker.status = \"In Progress\"\n tracker.start_time = time.time()\n\n # Transport/JSON failures may raise; HTTP errors return ok=False.\n try:\n result = await _make_request(\n client=self,\n url=worker_url,\n route=worker_route,\n api_key=endpoint.api_key,\n body=worker_request_body,\n method=\"POST\",\n retries=1, # avoid stacking retries with the outer loop\n timeout=worker_timeout,\n stream=stream,\n allow_non_json=True,\n )\n except (\n aiohttp.ClientConnectorError,\n aiohttp.ServerDisconnectedError,\n ) as ex:\n # Worker is gone\n if session is not None:\n # Session is bound to this worker - can't re-route\n self.logger.error(\n f\"Session worker unavailable ({type(ex).__name__})\"\n )\n session.open = False\n raise ConnectionError(\n f\"Session worker unavailable: {ex}\"\n ) from ex\n # No session - reset request_idx to force a fresh route\n self.logger.warning(\n f\"Worker unavailable ({type(ex).__name__}), re-routing to new worker\"\n )\n tracker.status = \"Retrying\"\n continue\n except Exception as ex:\n self.logger.error(f\"Worker request failed: {ex}\")\n tracker.status = \"Retrying\"\n continue\n\n if not result.get(\"ok\"):\n if (\n retry\n and result.get(\"retryable\")\n and (max_retries is None or total_attempts \u003c max_retries)\n ):\n # Check if we have time left before retrying\n if (\n timeout is not None\n and (time.time() - start_time) >= timeout\n ):\n raise asyncio.TimeoutError(\n f\"Request timed out after {time.time() - start_time:.1f}s\"\n )\n\n tracker.status = \"Retrying\"\n await asyncio.sleep(\n min(\n (2 ** min(total_attempts, 20)) + random.uniform(0, 1),\n self.max_poll_interval,\n )\n )\n continue\n\n # Return the raw HTTP result to the caller\n tracker.status = \"Complete\"\n tracker.complete_time = time.time()\n if tracker.start_time is not None:\n self.latencies.append(\n tracker.complete_time - tracker.start_time\n )\n\n return {\n \"response\": result.get(\"json\")\n if result.get(\"json\") is not None\n else {\"error\": result.get(\"text\", \"\")},\n \"ok\": result.get(\"ok\"),\n \"status\": result.get(\"status\"),\n \"text\": result.get(\"text\"),\n \"latency\": (tracker.complete_time - tracker.start_time)\n if tracker.start_time\n else None,\n \"url\": worker_url,\n \"request_idx\": request_idx,\n \"auth_data\": auth_data,\n }\n\n # Success\n if stream:\n worker_response = result.get(\"stream\")\n elif result.get(\"content\") is not None:\n worker_response = result.get(\"content\")\n else:\n worker_response = result.get(\"json\")\n\n tracker.status = \"Complete\"\n tracker.complete_time = time.time()\n self.latencies.append(tracker.complete_time - tracker.start_time)\n self.logger.info(\"Endpoint request task completed\")\n\n return {\n \"response\": worker_response,\n \"ok\": result.get(\"ok\"),\n \"status\": result.get(\"status\"),\n \"text\": result.get(\"text\"),\n \"content_type\": result.get(\"content_type\"),\n \"latency\": tracker.complete_time - tracker.start_time,\n \"url\": worker_url,\n \"request_idx\": request_idx,\n \"auth_data\": auth_data,\n }\n\n except asyncio.CancelledError:\n tracker.status = \"Cancelled\"\n raise\n except Exception as ex:\n tracker.status = \"Errored\"\n self.logger.error(f\"Request errored: {ex}\")\n raise\n\n @abstractmethod\n def queue_endpoint_request(\n self,\n endpoint: Endpoint_[R],\n worker_route: str,\n worker_payload: dict,\n session: Optional[Session] = None,\n serverless_request: Optional[Union[ServerlessRequest, RequestStatus]] = None,\n cost: int = 100,\n timeout: Optional[float] = None,\n worker_timeout: Optional[float] = 600,\n retry: bool = True,\n max_retries: Optional[int] = None,\n stream: bool = False,\n ) -> R:\n \"\"\"Dispatch a request to a serverless endpoint. Return type depends on subclass.\"\"\"\n ...\n\n\nclass Serverless(_ServerlessBase[ServerlessRequest]):\n \"\"\"Fire-and-observe client: queue_endpoint_request returns a ServerlessRequest (Future).\"\"\"\n\n def queue_endpoint_request(\n self,\n endpoint: Endpoint,\n worker_route: str,\n worker_payload: dict,\n session: Optional[Session] = None,\n serverless_request: Optional[Union[ServerlessRequest, RequestStatus]] = None,\n cost: int = 100,\n timeout: Optional[float] = None,\n worker_timeout: Optional[float] = 600,\n retry: bool = True,\n max_retries: Optional[int] = None,\n stream: bool = False,\n ) -> ServerlessRequest:\n \"\"\"Return a Future that will resolve once the request completes.\"\"\"\n if serverless_request is None:\n serverless_request = ServerlessRequest()\n elif isinstance(serverless_request, RequestStatus):\n serverless_request = ServerlessRequest(tracker=serverless_request)\n\n async def task(request: ServerlessRequest):\n try:\n result = await self._do_request(\n endpoint=endpoint,\n worker_route=worker_route,\n worker_payload=worker_payload,\n session=session,\n tracker=request.tracker,\n cost=cost,\n timeout=timeout,\n worker_timeout=worker_timeout,\n retry=retry,\n max_retries=max_retries,\n stream=stream,\n )\n request.set_result(result)\n except asyncio.CancelledError:\n pass # tracker already set by _do_request\n except Exception as ex:\n request.set_exception(ex)\n\n bg_task = asyncio.create_task(task(serverless_request))\n\n def _propagate_cancel(fut: ServerlessRequest):\n if fut.cancelled():\n bg_task.cancel()\n\n serverless_request.add_done_callback(_propagate_cancel)\n self.logger.info(\"Queued endpoint request\")\n return serverless_request\n\n\nclass CoroutineServerless(_ServerlessBase[Coroutine[Any, Any, dict]]):\n \"\"\"Async/await client: queue_endpoint_request returns a coroutine.\"\"\"\n\n def queue_endpoint_request(\n self,\n endpoint: Endpoint,\n worker_route: str,\n worker_payload: dict,\n session: Optional[Session] = None,\n serverless_request: Optional[Union[ServerlessRequest, RequestStatus]] = None,\n cost: int = 100,\n timeout: Optional[float] = None,\n worker_timeout: Optional[float] = 600,\n retry: bool = True,\n max_retries: Optional[int] = None,\n stream: bool = False,\n ) -> Coroutine[Any, Any, dict]:\n \"\"\"Return a coroutine that resolves to the response dict when awaited.\"\"\"\n if isinstance(serverless_request, ServerlessRequest):\n tracker = serverless_request.tracker\n elif isinstance(serverless_request, RequestStatus):\n tracker = serverless_request\n else:\n tracker = None\n return self._do_request(\n endpoint=endpoint,\n worker_route=worker_route,\n worker_payload=worker_payload,\n session=session,\n tracker=tracker,\n cost=cost,\n timeout=timeout,\n worker_timeout=worker_timeout,\n retry=retry,\n max_retries=max_retries,\n stream=stream,\n )\n","content_type":"text/x-python; charset=utf-8","language":"python","size":37223,"content_sha256":"c8f06990fa44cf8ad752a3de1d632c7e8b7f9a294b260e6d6f9a39719722a503"},{"filename":"serverless/client/connection.py","content":"# connection.py\nimport aiohttp\nimport asyncio\nimport random\nimport json\nfrom typing import AsyncIterator, Dict, Optional, Any\n\nfrom vastai.utils import VERSION\n\n_JITTER_CAP_SECONDS = 5.0\n\ndef _retryable(status: int) -> bool:\n return status in (408, 429) or (500 \u003c= status \u003c 600)\n\ndef _backoff_delay(attempt: int) -> float:\n # capped exponential backoff with jitter\n return min((2 ** attempt) + random.uniform(0, 1), _JITTER_CAP_SECONDS)\n\ndef _build_kwargs(\n *,\n headers: Dict[str, str],\n params: Dict[str, str],\n ssl_context,\n timeout: Optional[float],\n body: Optional[dict],\n method: str,\n stream: bool,\n) -> Dict:\n return {\n \"headers\": headers,\n \"params\": params,\n \"ssl\": ssl_context,\n \"timeout\": aiohttp.ClientTimeout(total=None if stream else timeout),\n **({\"json\": body} if method != \"GET\" and body else {}),\n }\n\nasync def _iter_sse_json(resp: aiohttp.ClientResponse) -> AsyncIterator[dict]:\n \"\"\"\n Yield JSON objects from an SSE/text stream. Accepts lines starting with 'data:' or raw JSONL.\n \"\"\"\n buffer = b\"\"\n async for chunk in resp.content.iter_any():\n if not chunk:\n continue\n buffer += chunk\n while b\"\\n\" in buffer:\n line, buffer = buffer.split(b\"\\n\", 1)\n line = line.strip()\n if not line:\n continue\n if line.startswith(b\"data:\"):\n line = line[5:].strip()\n try:\n yield json.loads(line.decode(\"utf-8\"))\n except Exception:\n # Ignore keepalives/bad fragments\n continue\n\n # flush tail if present\n tail = buffer.strip()\n if tail:\n try:\n yield json.loads(tail.decode(\"utf-8\"))\n except Exception:\n pass\n\nasync def _open_once(\n *,\n session: aiohttp.ClientSession,\n method: str,\n url: str,\n route: str,\n kwargs: Dict,\n):\n \"\"\"\n Execute one HTTP attempt and return the aiohttp response object.\n Caller is responsible for reading/closing via 'async with' or resp.release()/resp.close().\n \"\"\"\n request_fn = {\"GET\": session.get, \"POST\": session.post, \"PUT\": session.put, \"DELETE\": session.delete}.get(method, session.post)\n return await request_fn(url + route, **kwargs)\n\n\nasync def _make_request(\n client,\n route: str,\n api_key: str,\n url: str = \"\",\n body: Optional[dict] = None,\n params: Optional[dict] = None,\n method: str = \"GET\",\n retries: int = 5,\n timeout: float = 30,\n stream: bool = False,\n allow_non_json: bool = False,\n) -> Dict[str, Any]:\n \"\"\"\n Make an HTTP request with capped exponential backoff + jitter, returning a structured result.\n\n - Never raises for HTTP non-2xx responses. Instead returns result with ok=False and status/text/json.\n - Raises only for \"mechanical\" failures (aiohttp/transport) and invalid JSON on successful (2xx) responses.\n - allow_non_json: a 2xx text/* or media body that is not a JSON object or array is returned\n (str for text/*, bytes otherwise) instead of raising.\n\n Return shape (non-stream):\n {\n \"ok\": bool,\n \"status\": int|None,\n \"url\": str,\n \"headers\": dict,\n \"text\": str,\n \"json\": Any|None,\n \"content\": bytes|str|None, # allow_non_json only: str for text/*, else bytes\n \"content_type\": str, # the declared type, without parameters (\"\" if none)\n \"retryable\": bool,\n \"attempt\": int\n }\n\n Return shape (stream=True, ok=True adds \"stream\"):\n { ... , \"stream\": AsyncIterator[dict] }\n \"\"\"\n method = method.upper()\n body = body or {}\n params = {**(params or {})}\n params[\"api_key\"] = api_key\n\n headers = {\"Authorization\": f\"Bearer {api_key}\", \"User-Agent\": f\"vastai-sdk/{VERSION}\"}\n\n session = await client._get_session()\n ssl_context = await client.get_ssl_context() if client else None\n\n full_url = url + route\n\n last_result: Dict[str, Any] = {\n \"ok\": False,\n \"status\": None,\n \"url\": full_url,\n \"headers\": {},\n \"text\": \"\",\n \"json\": None,\n \"retryable\": False,\n \"attempt\": 0,\n }\n\n # STREAMING PATH (open upfront so we can return status + a stream iterator)\n if stream:\n for attempt in range(1, retries + 1):\n last_result[\"attempt\"] = attempt\n try:\n kwargs = _build_kwargs(\n headers=headers,\n params=params,\n ssl_context=ssl_context,\n timeout=timeout,\n body=body,\n method=method,\n stream=True,\n )\n\n resp = await _open_once(\n session=session,\n method=method,\n url=url,\n route=route,\n kwargs=kwargs,\n )\n\n status = resp.status\n last_result[\"status\"] = status\n last_result[\"headers\"] = dict(resp.headers)\n last_result[\"retryable\"] = _retryable(status)\n\n if status \u003c 200 or status >= 300:\n text = await resp.text()\n last_result[\"text\"] = text\n\n # best-effort JSON parse (do not raise)\n try:\n if text and text.lstrip().startswith((\"{\", \"[\")):\n last_result[\"json\"] = json.loads(text)\n except Exception:\n pass\n\n resp.release()\n\n if last_result[\"retryable\"] and attempt \u003c retries:\n await asyncio.sleep(_backoff_delay(attempt))\n continue\n\n # final non-2xx: return what we saw\n return last_result\n\n # 2xx: return stream iterator that owns the response lifetime\n last_result[\"ok\"] = True\n last_result[\"text\"] = \"\" # streaming: we don't pre-read body\n\n async def _stream_iter() -> AsyncIterator[dict]:\n try:\n async for obj in _iter_sse_json(resp):\n yield obj\n finally:\n try:\n resp.release()\n except Exception:\n pass\n\n last_result[\"stream\"] = _stream_iter()\n return last_result\n\n except asyncio.TimeoutError as ex:\n if attempt == retries:\n raise TimeoutError(f\"Request to {full_url} timed out after {timeout}s\") from ex\n await asyncio.sleep(_backoff_delay(attempt))\n except Exception as ex:\n if attempt == retries:\n raise ex\n await asyncio.sleep(_backoff_delay(attempt))\n\n return last_result\n\n # NON-STREAMING PATH\n for attempt in range(1, retries + 1):\n last_result[\"attempt\"] = attempt\n try:\n kwargs = _build_kwargs(\n headers=headers,\n params=params,\n ssl_context=ssl_context,\n timeout=timeout,\n body=body,\n method=method,\n stream=False,\n )\n\n request_fn = {\"GET\": session.get, \"POST\": session.post, \"PUT\": session.put, \"DELETE\": session.delete}.get(method, session.post)\n async with await request_fn(full_url, **kwargs) as resp:\n status = resp.status\n content_type = (resp.headers.get(\"Content-Type\") or \"\").split(\";\", 1)[0].lower()\n is_text = content_type.startswith(\"text/\")\n is_media = (allow_non_json and bool(content_type) and not is_text\n and \"json\" not in content_type)\n raw = await resp.read() if is_media else None\n text = \"\" if is_media and 200 \u003c= status \u003c 300 else await resp.text(errors=\"replace\")\n\n result: Dict[str, Any] = {\n \"ok\": 200 \u003c= status \u003c 300,\n \"status\": status,\n \"url\": full_url,\n \"headers\": dict(resp.headers),\n \"text\": text,\n \"json\": None,\n \"content\": None,\n \"content_type\": content_type,\n \"retryable\": _retryable(status),\n \"attempt\": attempt,\n }\n\n if result[\"ok\"]:\n if is_media or (allow_non_json and is_text):\n # The body, unless it is an object or array (mislabelled JSON).\n body = raw if is_media else text\n try:\n if body.lstrip()[:1] in (\"{\", \"[\", b\"{\", b\"[\"):\n result[\"json\"] = json.loads(body)\n return result\n except ValueError:\n pass\n result[\"content\"] = body\n return result\n # Successful responses are expected to be JSON; invalid JSON is a hard failure\n try:\n result[\"json\"] = await resp.json(content_type=None)\n except Exception:\n raise Exception(f\"Invalid JSON from {full_url}:\\n{text}\")\n\n # Debug: log the exact response\n if hasattr(client, 'logger'):\n client.logger.debug(f\"_make_request response from {route}: {result}\")\n\n return result\n\n # Non-2xx: best-effort JSON parse (do not raise)\n try:\n if text and text.lstrip().startswith((\"{\", \"[\")):\n result[\"json\"] = json.loads(text)\n except Exception:\n pass\n\n # Retry only if retryable and attempts remain\n if result[\"retryable\"] and attempt \u003c retries:\n last_result = result\n await asyncio.sleep(_backoff_delay(attempt))\n continue\n\n # Debug: log the exact response for non-2xx\n if hasattr(client, 'logger'):\n client.logger.debug(f\"_make_request response from {route}: {result}\")\n\n return result\n\n except asyncio.TimeoutError as ex:\n if attempt == retries:\n raise TimeoutError(f\"Request to {full_url} timed out after {timeout}s\") from ex\n await asyncio.sleep(_backoff_delay(attempt))\n except Exception as ex:\n if attempt == retries:\n raise ex\n await asyncio.sleep(_backoff_delay(attempt))\n\n return last_result\n","content_type":"text/x-python; charset=utf-8","language":"python","size":10931,"content_sha256":"7431c02969f935903bc6b0bcec6720a9bf86d1cc0df378ecd5ac57ce2fe573d1"},{"filename":"serverless/client/endpoint.py","content":"# endpoint.py\nimport asyncio\nfrom dataclasses import dataclass\nimport os\n\nimport time\nfrom .connection import _make_request\nfrom typing import Awaitable, Generic, Optional, TypeVar, Union, TYPE_CHECKING\nimport logging\n\nfrom vastai.data.endpoint import EndpointData\n\nlogger = logging.getLogger(\"vastai\")\n\nif TYPE_CHECKING:\n from .client import _ServerlessBase, ServerlessRequest\n from .request_status import RequestStatus\n from .session import Session\n\nR = TypeVar(\"R\", bound=Awaitable)\n\n\n@dataclass\nclass EndpointDataPartial:\n name: str\n id: int\n api_key: str\n\n\nclass Endpoint_(Generic[R]):\n client: \"_ServerlessBase[R]\"\n\n def __repr__(self):\n return f\"\u003cEndpoint {self.name} (id={self.data.id})>\"\n\n def __init__(\n self,\n client: \"_ServerlessBase[R]\",\n name: Optional[str] = None,\n id: Optional[int] = None,\n api_key: Optional[str] = None,\n data: Optional[EndpointData] = None,\n soft_refresh_threshold=3 * 24 * 3600,\n hard_refresh_threshold=6 * 24 * 3600,\n ):\n if client is None:\n raise ValueError(\"client is required\")\n self.client = client\n if isinstance(data, EndpointData):\n self.data: \"EndpointData | EndpointDataPartial\" = data\n elif isinstance(name, str) and isinstance(id, int) and isinstance(api_key, str):\n self.data: \"EndpointData | EndpointDataPartial\" = EndpointDataPartial(\n name, id, api_key\n )\n else:\n raise ValueError(\n \"Either data or all of name, id, and api_key are required!\"\n )\n self.refresh_task: Optional[asyncio.Task[EndpointData]] = None\n self.last_refresh = time.time()\n self.soft_refresh_threshold = soft_refresh_threshold\n self.hard_refresh_threshold = hard_refresh_threshold\n\n @property\n def name(self):\n return (\n self.data.config.endpoint_name\n if isinstance(self.data, EndpointData)\n else self.data.name\n )\n\n @property\n def id(self):\n return self.data.id\n\n @property\n def api_key(self):\n return self.data.api_key\n\n async def refresh(self, block=False):\n logger.debug(\"refresh\")\n if self.refresh_task is not None:\n if block:\n try:\n self.data = await self.refresh_task\n finally:\n self.refresh_task = None\n else:\n try:\n self.data = self.refresh_task.result()\n self.refresh_task = None\n except asyncio.InvalidStateError:\n pass\n except Exception as e: # refresh task errored, reset it.\n self.refresh_task = None\n raise e\n else:\n if block:\n self.data = await self.client.fetch_endpoint(self.data.id)\n else:\n self.refresh_task = asyncio.create_task(\n self.client.fetch_endpoint(self.data.id)\n )\n\n def request(\n self,\n route,\n payload,\n serverless_request: Optional[\n Union[\"ServerlessRequest\", \"RequestStatus\"]\n ] = None,\n cost: int = 100,\n retry: bool = True,\n stream: bool = False,\n timeout: Optional[float] = None,\n session: Optional[\"Session\"] = None,\n ) -> R:\n return self.client.queue_endpoint_request(\n endpoint=self,\n worker_route=route,\n worker_payload=payload,\n serverless_request=serverless_request,\n cost=cost,\n retry=retry,\n stream=stream,\n timeout=timeout,\n session=session,\n )\n\n def close_session(self, session: \"Session\"):\n return self.client.end_endpoint_session(session=session)\n\n async def session_healthcheck(self, session: \"Session\"):\n result = await self.client.get_endpoint_session(\n endpoint=self, session_id=session.session_id, session_auth=session.auth_data\n )\n return result is not None\n\n def get_session(self, session_id: int, session_auth: dict, timeout: float = 10):\n return self.client.get_endpoint_session(\n endpoint=self,\n session_id=session_id,\n session_auth=session_auth,\n timeout=timeout,\n )\n\n def session(\n self,\n cost: int = 100,\n lifetime: float = 60,\n on_close_route: str = None,\n on_close_payload: dict = None,\n timeout: float = None,\n ) -> \"Session\":\n return self.client.start_endpoint_session(\n endpoint=self,\n cost=cost,\n lifetime=lifetime,\n on_close_route=on_close_route,\n on_close_payload=on_close_payload,\n timeout=timeout,\n )\n\n def get_workers(self):\n return self.client.get_endpoint_workers(self)\n\n async def _route(\n self, cost: float = 0.0, req_idx: int = 0, timeout: float = 60.0\n ) -> \"RouteResponse\":\n VAST_DEBUG_WORKER_URL = os.environ.get(\"VAST_DEBUG_WORKER_URL\")\n if VAST_DEBUG_WORKER_URL:\n return RouteResponse({\"request_idx\": 1, \"url\": VAST_DEBUG_WORKER_URL})\n await self.client._get_session()\n if time.time() > self.last_refresh + self.soft_refresh_threshold:\n await self.refresh(\n block=time.time() > self.last_refresh + self.hard_refresh_threshold\n )\n try:\n result = await _make_request(\n client=self.client,\n url=self.client.autoscaler_url,\n route=\"/route/\",\n api_key=self.data.api_key,\n body={\n \"endpoint\": self.name,\n \"api_key\": self.data.api_key,\n \"cost\": cost,\n \"request_idx\": req_idx,\n \"replay_timeout\": timeout,\n },\n method=\"POST\",\n timeout=10.0,\n retries=1,\n stream=False,\n )\n except Exception as ex:\n raise RuntimeError(f\"Failed to route endpoint: {ex}\") from ex\n\n if not result.get(\"ok\"):\n raise RuntimeError(\n f\"Failed to route endpoint: HTTP {result.get('status')} - {result.get('text', '')[:512]}\"\n )\n\n return RouteResponse(result.get(\"json\") or {})\n\n\nclass RouteResponse:\n status: str\n body: dict\n request_idx: int\n\n def __repr__(self):\n return f\"\u003cRouteResponse status={self.status}>\"\n\n def __init__(self, body: dict):\n if \"request_idx\" in body.keys():\n self.request_idx = body.get(\"request_idx\")\n else:\n self.request_idx = 0\n if \"url\" in body.keys():\n self.status = \"READY\"\n self.body = body\n else:\n self.status = \"WAITING\"\n self.body = body\n\n def get_url(self):\n return self.body.get(\"url\")\n\n\n# Backward-compatible alias — works for instantiation at runtime.\n# Static analysis infers the R parameter from the client argument to __init__.\nEndpoint = Endpoint_\n","content_type":"text/x-python; charset=utf-8","language":"python","size":7238,"content_sha256":"317403645cccfb4aff7e81cde82ebc4d07fa643106a5abcbe826ae6ab958ed80"},{"filename":"serverless/client/managed.py","content":"from __future__ import annotations\nfrom typing import TYPE_CHECKING, Awaitable, Generic, Optional, Union, TypeVar\n\nfrom vastai.data.endpoint import EndpointData\nfrom vastai.data.deployment import DeploymentData, DeploymentPutResponse\nfrom vastai.data.workergroup import WorkergroupConfig\nfrom .endpoint import Endpoint_\nimport requests\n\nif TYPE_CHECKING:\n from .client import _ServerlessBase\n\nR = TypeVar(\"R\", bound=Awaitable)\n\n\nclass ManagedEndpoint(Generic[R]):\n \"\"\"Wraps an endpoint with lazy data fetching, caching, and request dispatch.\n\n The .id property is always available. All other data is fetched on first\n access via .get() and cached. Call .invalidate() to force a re-fetch.\n \"\"\"\n\n def __init__(\n self, id: int, client: _ServerlessBase[R], data: Optional[EndpointData] = None\n ):\n self._id = id\n self._client = client\n self._routing_endpoint: Optional[Endpoint_[R]] = (\n Endpoint_[R](client, data=data) if isinstance(data, EndpointData) else None\n )\n\n @property\n def id(self) -> int:\n return self._id\n\n def invalidate(self):\n \"\"\"Clear cached data so the next .get() re-fetches from the API.\"\"\"\n self._data = None\n self._routing_endpoint = None\n\n async def _get_routing_endpoint(self) -> Endpoint_[R]:\n if self._routing_endpoint is None:\n data = await self._client.fetch_endpoint(self._id)\n self._routing_endpoint = Endpoint_(client=self._client, data=data)\n return self._routing_endpoint\n\n async def request(self, route: str, payload: dict, **kwargs) -> dict:\n \"\"\"Route a request to this endpoint's workers. Fetches endpoint data if needed.\"\"\"\n ep = await self._get_routing_endpoint()\n return await self._client.queue_endpoint_request(\n endpoint=ep,\n worker_route=route,\n worker_payload=payload,\n **kwargs,\n )\n\n async def add_workergroup(\n self,\n config_or_template_hash: Union[WorkergroupConfig, str, None] = None,\n **kwargs,\n ) -> int:\n \"\"\"Add a workergroup to this endpoint. Returns the workergroup ID.\n\n Accepts either:\n - A WorkergroupConfig object\n - A template_hash string (common case)\n - Keyword arguments passed to WorkergroupConfig\n \"\"\"\n if isinstance(config_or_template_hash, WorkergroupConfig):\n config = config_or_template_hash\n if config.endpoint_id is None:\n config.endpoint_id = self._id\n elif isinstance(config_or_template_hash, str):\n config = WorkergroupConfig(\n endpoint_id=self._id, template_hash=config_or_template_hash, **kwargs\n )\n else:\n config = WorkergroupConfig(endpoint_id=self._id, **kwargs)\n return await self._client.create_workergroup(config)\n\n async def delete(self) -> None:\n \"\"\"Delete this endpoint.\"\"\"\n await self._client.delete_endpoint(self._id)\n\n def __repr__(self) -> str:\n status = \"loaded\" if self._routing_endpoint is not None else \"lazy\"\n return f\"\u003cManagedEndpoint id={self._id} ({status})>\"\n\n\nclass ManagedDeployment(Generic[R]):\n \"\"\"Wraps a deployment with lazy data fetching, caching, and endpoint access.\n\n The .id and .endpoint_id properties are always available.\n If the deployment was just created/updated, .needs_upload and .upload()\n are available to handle S3 blob uploads.\n \"\"\"\n\n def __init__(\n self,\n id: int,\n endpoint_id: int,\n client: _ServerlessBase[R],\n data: Optional[DeploymentData] = None,\n put_response: Optional[DeploymentPutResponse] = None,\n ):\n self._id = id\n self._endpoint_id = endpoint_id\n self._client = client\n self._data = data\n self._put_response = put_response\n self._endpoint: Optional[ManagedEndpoint[R]] = None\n\n @property\n def id(self) -> int:\n return self._id\n\n @property\n def endpoint_id(self) -> int:\n return self._endpoint_id\n\n @property\n def action(self) -> Optional[str]:\n \"\"\"The action taken by put_deployment (e.g. 'created', 'soft_update', 'exists').\"\"\"\n return self._put_response.action if self._put_response else None\n\n @property\n def workergroup_id(self) -> Optional[int]:\n \"\"\"The workergroup ID associated with this deployment, if returned by the server.\"\"\"\n return self._put_response.workergroup_id if self._put_response else None\n\n @property\n def needs_upload(self) -> bool:\n \"\"\"True if the deployment requires a blob upload to S3 to complete setup.\"\"\"\n return (\n self._put_response is not None\n and self._put_response.upload_url is not None\n and self._put_response.upload_fields is not None\n )\n\n async def upload(self, file_path: str) -> None:\n \"\"\"Upload a tarball to S3 using the presigned POST from put_deployment.\n\n Raises ValueError if no upload is needed, RuntimeError on upload failure.\n \"\"\"\n if not self.needs_upload:\n raise ValueError(\"No upload needed for this deployment\")\n\n import aiohttp\n import os\n\n upload_url = self._put_response.upload_url\n upload_fields = self._put_response.upload_fields\n\n data = aiohttp.FormData()\n for key, value in upload_fields.items():\n data.add_field(key, value)\n data.add_field(\n \"file\",\n open(file_path, \"rb\"),\n filename=os.path.basename(file_path),\n content_type=\"application/gzip\",\n )\n\n session = await self._client._get_session()\n async with session.post(upload_url, data=data) as resp:\n if resp.status not in (200, 201, 204):\n text = await resp.text()\n raise RuntimeError(\n f\"S3 upload failed: HTTP {resp.status} - {text[:512]}\"\n )\n\n self._put_response.upload_url = None\n self._put_response.upload_fields = None\n\n @property\n def endpoint(self) -> ManagedEndpoint[R]:\n \"\"\"Get the ManagedEndpoint for this deployment (lazy — data fetched on first use).\"\"\"\n if self._endpoint is None:\n self._endpoint = ManagedEndpoint(id=self._endpoint_id, client=self._client)\n return self._endpoint\n\n async def get(self) -> DeploymentData:\n \"\"\"Fetch and cache the full deployment data. Returns cached data on subsequent calls.\"\"\"\n if self._data is None:\n self._data = await self._client.fetch_deployment(self._id)\n return self._data\n\n def invalidate(self):\n \"\"\"Clear cached data so the next .get() re-fetches from the API.\"\"\"\n self._data = None\n\n def sync_heartbeat(\n self,\n ): # for use in background thread outside of main aiohttp connection\n headers = {\"Authorization\": f\"Bearer {self._client.api_key}\"}\n requests.post(\n url=self._client.vast_web_url.rstrip(\"/\")\n + f\"/api/v0/deployment/{self.id}/heartbeat/\",\n headers=headers,\n )\n\n def __repr__(self) -> str:\n status = \"loaded\" if self._data is not None else \"lazy\"\n upload = \" needs_upload\" if self.needs_upload else \"\"\n return f\"\u003cManagedDeployment id={self._id} endpoint_id={self._endpoint_id} ({status}{upload})>\"\n","content_type":"text/x-python; charset=utf-8","language":"python","size":7424,"content_sha256":"ba459ccab6d4ca9964ceeafbf82105f697144013565c8e0bac24db61c414667f"},{"filename":"serverless/client/request_status.py","content":"import time\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\n\n@dataclass\nclass RequestStatus:\n \"\"\"Observable status tracker for a serverless request lifecycle.\"\"\"\n status: str = \"New\"\n create_time: float = field(default_factory=time.time)\n start_time: Optional[float] = None\n complete_time: Optional[float] = None\n req_idx: int = 0\n","content_type":"text/x-python; charset=utf-8","language":"python","size":371,"content_sha256":"345ed9f21e9c97259616e4aff438474f5c1bac33e6e1751838f683a55daee797"},{"filename":"serverless/client/session.py","content":"from .endpoint import Endpoint\nfrom typing import Callable, Optional\nimport inspect\nimport logging\n\nlog = logging.getLogger(__name__)\n\n\nclass Session:\n endpoint: Endpoint\n session_id: str\n open: bool\n url: str\n lifetime: float\n auth_data: dict\n on_close_route: str\n on_close_payload: dict\n\n 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):\n if endpoint is None:\n raise ValueError(\"Session cannot be created with empty endpoint\")\n if session_id is None:\n raise ValueError(\"Session cannot be created with empty session_id\")\n if url is None:\n raise ValueError(\"Session cannot be created with empty url\")\n\n self.endpoint = endpoint\n self.session_id = session_id\n self.lifetime = lifetime\n self.expiration = expiration\n self.url = url\n self.auth_data = auth_data\n\n self.open = True\n self.on_close_route = on_close_route\n self.on_close_payload = on_close_payload\n self._closing = False # guard to ensure one-time execution\n\n async def __aenter__(self):\n return self\n\n async def __aexit__(self, exc_type, exc, tb):\n await self.close()\n return False\n\n async def is_open(self):\n result = await self.endpoint.session_healthcheck(self)\n self.open = result\n return result\n\n async def close(self):\n \"\"\"\n Explicit close for non-async contexts.\n Returns an awaitable if async work is required.\n \"\"\"\n if not self.open or self._closing:\n return None\n self._closing = True\n try:\n await self.endpoint.close_session(self)\n except Exception as e:\n log.warning(f\"Error closing session {self.session_id}: {e}\")\n finally:\n self.open = False\n return\n\n def request(\n self,\n route,\n payload,\n serverless_request=None,\n cost: int = 100,\n retry: bool = True,\n stream: bool = False,\n ):\n \"\"\"Forward requests to the endpoint\"\"\"\n if not self.open:\n raise ValueError(\"Cannot make request on closed session.\")\n\n async def _wrapped_request():\n result = await self.endpoint.request(\n route=route,\n payload=payload,\n serverless_request=serverless_request,\n cost=cost,\n retry=retry,\n stream=stream,\n session=self,\n )\n\n if result.get(\"status\") == 410:\n self.open = False\n raise ValueError(\"Cannot make request on closed session.\")\n\n return result\n\n return _wrapped_request()","content_type":"text/x-python; charset=utf-8","language":"python","size":2856,"content_sha256":"1b6aa39c534c2e4b6214c155be49457b52b7d87414fc289da9a3ce791fde5638"},{"filename":"serverless/client/worker.py","content":"from dataclasses import dataclass\nfrom typing import Any, Dict, List, Tuple, Optional\n\n@dataclass\nclass Worker:\n id: int\n status: str\n cur_load: float\n new_load: float\n cur_load_rolling_avg: float\n cur_perf: float\n perf: float\n measured_perf: float\n dlperf: float\n reliability: float\n reqs_working: int\n disk_usage: float\n loaded_at: float\n started_at: float\n\n @staticmethod\n def from_dict(d: Dict[str, Any]) -> \"Worker\":\n # Be resilient to missing / extra fields\n status = d.get(\"status\") or \"UNKNOWN\"\n\n return Worker(\n id=int(d.get(\"id\")),\n status=status,\n cur_load=float(d.get(\"cur_load\", 0.0)),\n new_load=float(d.get(\"new_load\", 0.0)),\n cur_load_rolling_avg=float(d.get(\"cur_load_rolling_avg\", 0.0)),\n cur_perf=float(d.get(\"cur_perf\", 0.0)),\n perf=float(d.get(\"perf\", 0.0)),\n measured_perf=float(d.get(\"measured_perf\", 0.0)),\n dlperf=float(d.get(\"dlperf\", 0.0)),\n reliability=float(d.get(\"reliability\", 0.0)),\n reqs_working=int(d.get(\"reqs_working\", 0)),\n disk_usage=float(d.get(\"disk_usage\", 0.0)),\n loaded_at=float(d.get(\"loaded_at\", 0.0)),\n started_at=float(d.get(\"started_at\", 0.0)),\n )\n","content_type":"text/x-python; charset=utf-8","language":"python","size":1325,"content_sha256":"0ba439d0dcca641364538fa07567d14102b76db72169e4af5c92067e16e3b152"},{"filename":"serverless/remote/__init__.py","content":"import os\n\nfrom .base import DockerLogin, Image, Autoscaling\n\nif os.environ.get(\"IS_DEPLOYMENT\", False):\n from .serve import Deployment\nelse:\n from .deploy import Deployment\n\n__all__ = [\"Deployment\", \"DockerLogin\", \"Image\", \"Autoscaling\"]\n","content_type":"text/x-python; charset=utf-8","language":"python","size":245,"content_sha256":"5ce99b5377d43dcac42c6143abef416362c208b61a2ee7dd55789ff91184c374"},{"filename":"serverless/remote/base.py","content":"from abc import ABC, abstractmethod\nfrom typing import (\n Callable,\n Any,\n Awaitable,\n ParamSpec,\n Optional,\n TypeVar,\n Type,\n AsyncContextManager,\n TypedDict,\n)\nfrom typing_extensions import Unpack\nfrom vastai.data import Query\nfrom dataclasses import dataclass\n\n\nclass DockerLogin(TypedDict, total=False):\n docker_login_user: str\n docker_login_pass: str\n docker_login_repo: str\n\n\nclass Image:\n def __init__(\n self,\n from_image: str,\n storage: float = 50,\n worker_port: tuple[int, str] = (3000, \"tcp\"),\n **docker_login: Unpack[DockerLogin],\n ):\n self._image = from_image\n self._envs: dict[str, str] = {}\n self._runs: list[str | tuple[str, ...]] = []\n self._venv: str | None = (\n None # None=SDK default, \"\"=system python, str=custom venv path\n )\n self._pip_installs: list[str] = []\n self._apt_gets: list[str] = []\n self._docker_login: DockerLogin = docker_login\n self._requires = Query.search_defaults()\n self._storage = storage\n self._copies: list[tuple[str, str]] = []\n self._ports: set[tuple[int, str]] = {worker_port}\n\n def pip_install(self, *args: str) -> \"Image\":\n for arg in args:\n self._pip_installs.append(arg)\n return self\n\n def apt_get(self, *args: str) -> \"Image\":\n for arg in args:\n self._apt_gets.append(arg)\n return self\n\n def run_script(self, script) -> \"Image\":\n self._runs.append(script)\n return self\n\n def run_cmd(self, *args) -> \"Image\":\n self._runs.append(args)\n return self\n\n def env(self, **kwargs: str) -> \"Image\":\n for k, v in kwargs.items():\n self._envs[k] = v\n return self\n\n def require(self, *args: Query) -> \"Image\":\n for arg in args:\n self._requires.extend(arg)\n return self\n\n def copy(self, src: str, dst: str) -> \"Image\":\n self._copies.append((src, dst))\n return self\n\n def venv(self, path: str) -> \"Image\":\n \"\"\"Use an existing venv at the given path instead of the SDK managed venv.\"\"\"\n self._venv = path\n return self\n\n def use_system_python(self) -> \"Image\":\n \"\"\"Use the image's system Python instead of a venv.\"\"\"\n self._venv = \"\"\n return self\n\n def publish_port(self, number: int, type_: str = \"tcp\") -> \"Image\":\n self._ports.add((number, type_))\n return self\n\n\nclass Autoscaling(TypedDict, total=False):\n # Scaling params\n cold_workers: int\n max_workers: int\n min_load: int\n min_cold_load: int\n target_util: float\n cold_mult: int\n max_queue_time: float\n target_queue_time: float\n inactivity_timeout: int\n autoscaler_instance: str\n\n\n@dataclass\nclass Config:\n name: str\n pip_installs: list[str]\n apt_gets: list[str]\n envs: list[list[str]] | list[tuple[str, ...]] # format [(KEY,VALUE)]\n runs: (\n list[str | list[str]] | list[str | tuple[str, ...]]\n ) # tuples get turned into list after JSON serialize->deserialize round trip\n\n\nT = TypeVar(\"T\")\nDeploymentT = TypeVar(\"DeploymentT\", bound=\"Deployment_\")\n\n\nclass Deployment_(ABC):\n by_name: dict[Type, dict[str, \"Deployment_\"]] = {}\n\n def __init__(\n self,\n name: Optional[str] = None,\n tag: str = \"default\",\n version_label: Optional[str] = None,\n **kwargs,\n ):\n self.name = name\n self.tag = tag\n self.version_label: str | None = version_label\n self.root_module: str | None = None\n self.is_main = False\n\n def context(\n self, *args, **kwargs\n ) -> Callable[\n [Type[AsyncContextManager[T]]], Type[AsyncContextManager[T]]\n ]: # registers context_class as deployment context\n return lambda c: c # no-op when not in serve mode\n\n def get_context(self, context_class: Type[AsyncContextManager[T]]) -> T:\n raise NotImplementedError(\n f\"{context_class}\"\n ) # should not be run when not in serve mode\n\n @property\n def name(self):\n return self.name_\n\n @name.setter\n def name(self, name: str | None):\n self.name_ = name\n if isinstance(name, str):\n if self.__class__ not in self.by_name:\n self.by_name[self.__class__] = {}\n self.by_name[self.__class__][name] = self\n\n @classmethod\n def lookup(cls: Type[DeploymentT], name: str) -> DeploymentT | None:\n if cls not in cls.by_name:\n return None\n ret = cls.by_name[cls].get(name)\n return ret # type: ignore\n\n def relativize(self, f: Callable[..., Any]) -> tuple[str]:\n mod = f.__module__\n name = f.__name__\n file = f.__globals__.get(\"__file__\")\n # assume single file deployment for now.\n if self.root_module is None:\n self.root_module = mod\n if self.name_ is None:\n self.name = \"\"\n if file is None:\n raise Exception(\n 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.\"\n )\n self.file = file\n if f.__module__ != self.root_module:\n raise Exception(\n 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}.\"\n )\n return (name,)\n\n P = ParamSpec(\"P\")\n\n # TODO: add type checked version of remote\n @abstractmethod\n def remote(\n self,\n f: Callable[P, Awaitable[Any]] | None = None,\n *,\n benchmark_dataset: list[dict] | None = None,\n benchmark_generator: Callable[[], dict] | None = None,\n benchmark_runs: int = 10,\n workload_calculator: Callable[..., float] | None = None,\n ) -> (\n Callable[P, Awaitable[Any]]\n | Callable[[Callable[P, Awaitable[Any]]], Callable[P, Awaitable[Any]]]\n ): ...\n\n # Stubs for methods that are expected to be present in serve mode, but shouldn't do anything.\n def configure_autoscaling(self, **kwargs: Unpack[Autoscaling]):\n pass\n\n def image(self, from_image: str, storage: int) -> Image:\n return Image(\"\")\n\n def ensure_ready(self): # NO-OP in serve mode\n pass\n","content_type":"text/x-python; charset=utf-8","language":"python","size":6445,"content_sha256":"e7602b673b62b4d37f9ab8276aac0d255492c3892fcb2fba08a94700b73bb0a1"},{"filename":"serverless/remote/deploy.py","content":"import inspect\nimport threading\nimport time\nimport json\nimport os\nfrom typing import Optional, Any, Callable, Awaitable, ParamSpec, BinaryIO\nfrom typing_extensions import Unpack\nfrom dataclasses import asdict, dataclass\nfrom vastai.data import query\nfrom vastai import AsyncClient\nfrom vastai._base import _APIKEY_SENTINEL\nfrom vastai.data.deployment import DeploymentConfig\nfrom vastai.serverless.client import ManagedDeployment\nfrom . import serialization\nfrom .base import Deployment_, Config, DockerLogin, Image, Autoscaling\nfrom .utils import create_deployment_tarball, compute_deployment_hash\nfrom os.path import getsize\nimport tempfile\nimport asyncio\nimport logging\n\nlogger = logging.getLogger(\"vastai\")\n\nlogger.debug(\"mode: deploy\")\n# TODO: implement heartbeat, sync ready.\n\n\"\"\"\nDesign:\n\nDeployment has a root module.\n\nSingle file vs package deployments:\nsingle file:\n- from vastai.serverless import deployments\n- deployment = deployments.Deployment(Image) \n- @deployment.remote...\n- deployment.ensure_ready()\n\nmulti file: # skip for now\n- @deployments.post_configure(name = name, image = image)\n def post_configure(deployment):\n deployment.ensure_ready()\n with post_configure:\n import mod1 \n import mod2\n\nServe:\n- pull tarball\n- read env, onstart, installs, submodules, name from config. # non essential\n- inserts envs, runs installs, runs onstart. # non essential\n- from vastai.serverless import deployments # essential\n- import deployment # for package deployments require __init__.py to explicitly import remote function modules \n- my_deployment = deployments.get(name) # name is None if not manually set here. \n- serve(my_deployment) -> # essential\n async loop\n async serve(fname, args) -> serialize(my_deployment.registry[fname](deserialize(args)))\n- fname relative to root module\n\n\nDeploy:\n- make endpoint args and get a tarball upload link\n- make tarball and post it\n- when functions are called:\n - convert fname to be relative to root module. \n\n\nClient -> Serve depends:\n- POST -> app.endpoint/api/f serialized args -> returns f(args). \n - serialization: root module agrees (i.e., `deployment` in serve is the root_module of deployment)\nServe depends:\n- handle_f -> app.get(f_name)(args,kwargs)\n - app \u003c- deployments.get_app(name)\n - import deployment -> registers app and app functions\n - deployment module in ./deployment/ or ./deployment.py, untarred in tarball\n - name from config.json\n - environment:\n - On Image defined by from_image\n - with machine requirements set by requires_\n - envs set to envs (from config.json)\n - pip installs installed \"\"\n - apt gets installed \"\"\n - runs ran \"\"\n - vastai sdk installed (by bootstrap script)\n - where config.json, deployments are in tarball\n - bootstrap script knows deployment_id to grab from tarball with from S3\n\nServe \u003c-> Webserver:\n - deployment_id in env\n - bootstrap script in onstart\n - Image, search params from deployments\n\nServe \u003c-> Deploy:\n - Tarball as promised\nDeploy \u003c-> webserver\n\n\"\"\"\n\nDEBUG_DEPLOYMENT_TAR = os.environ.get(\"DEBUG_DEPLOYMENT_TAR\")\n\n\n@dataclass\nclass _FullDeployment:\n root_module: str\n deployment: ManagedDeployment\n\n\nP = ParamSpec(\"P\")\n\n\nclass Deployment(Deployment_): # TODO: Async Context Manager compatible with client\n def __init__(\n self,\n api_key: str | object = _APIKEY_SENTINEL,\n ttl: float | None = None,\n autoscaler_instance=\"prod\",\n autoscaler_url: Optional[str] = None,\n webserver_url=\"https://console.vast.ai\",\n **kwargs,\n ):\n super().__init__(**kwargs)\n self.async_client = AsyncClient(api_key, vast_server=webserver_url)\n self.client = self.async_client.serverless(\n instance=autoscaler_instance, autoscaler_url=autoscaler_url\n )\n self.autoscaler_instance = autoscaler_instance\n self.autoscaler_url = self.client.autoscaler_url\n self.webserver_url = webserver_url\n self._image: Image | None = None\n self._autoscaling: Autoscaling | None = None\n self._ttl = ttl\n self._inner: _FullDeployment | None = None\n\n def _compile_env(self, checked_image: Image) -> str:\n envs = [f\"-p {port}:{port}/{type_}\" for port, type_ in checked_image._ports]\n if self.autoscaler_url is not None:\n envs.append(f\"-e REPORT_ADDR={self.autoscaler_url}\")\n if self.webserver_url != \"https://console.vast.ai\":\n envs.append(f\"-e VAST_API_URL={self.webserver_url}\")\n if checked_image._venv == \"\":\n envs.append(\"-e USE_SYSTEM_PYTHON=true\")\n envs.append(\"-e UV_SYSTEM_PYTHON=1\")\n envs.append(\"-e UV_BREAK_SYSTEM_PACKAGES=1\")\n elif checked_image._venv is not None:\n envs.append(f\"-e ENV_PATH={checked_image._venv}\")\n return \" \".join(envs)\n\n def _into_deployment_config_and_tarball(self, tar_path: str) -> DeploymentConfig:\n # should error if _image, _autoscaling, file, or any other field needed to calculate DeploymentConfig is None\n if not isinstance(self._image, Image):\n raise Exception(\n \"Trying to deploy a deployment without an image configured.\"\n )\n if not isinstance(self.name, str):\n raise Exception(\n \"Trying to deploy an unbound deployment. Have any remote functions been registered?\"\n )\n if not isinstance(self.root_module, str):\n raise Exception(\n \"Trying to deploy an unbound deployment. Have any remote functions been registered?\"\n )\n if not isinstance(self._autoscaling, dict):\n raise Exception(\n \"Trying to deploy a deployment without autoscaling configured.\"\n )\n logger.debug(f\"Building deployment tarball at {tar_path}\")\n hash, size = self._compute_hash_and_filesize_and_make_tar(\n tar_path, self.name, self._image\n )\n logger.info(f\"Deployment tarball built: hash={hash}, size={size} bytes\")\n return DeploymentConfig(\n name=self.name\n if self.name\n else self.root_module, # if \"\", is main deployment for module\n image=self._image._image,\n env=self._compile_env(self._image),\n file_hash=hash,\n file_size=size,\n tag=self.tag,\n search_params=self._image._requires.unparse_query(),\n storage=self._image._storage,\n ttl=self._ttl,\n version_label=self.version_label,\n **self._image._docker_login,\n **self._autoscaling,\n )\n\n def configure_autoscaling(self, **kwargs: Unpack[Autoscaling]):\n if self._autoscaling is None:\n self._autoscaling = kwargs\n else:\n self._autoscaling.update(kwargs)\n\n def image(\n self, from_image: str, storage: int, **docker_login: Unpack[DockerLogin]\n ) -> Image:\n self._image = Image(from_image, storage, **docker_login)\n return self._image\n\n def _collate_config(self, checked_name: str, checked_image: Image) -> Config:\n # should error if any required fields are still None\n return Config(\n checked_name,\n checked_image._pip_installs,\n checked_image._apt_gets,\n list(checked_image._envs.items()),\n checked_image._runs,\n )\n\n def _compute_hash_and_filesize_and_make_tar(\n self, tar_path: str, checked_name: str, checked_image: Image\n ) -> tuple[str, int]:\n if not isinstance(self.file, str):\n raise Exception(\n \"Trying to deploy a deployment not yet bound to a Python module. Have any remote functions been registered?\"\n )\n config = self._collate_config(checked_name, checked_image)\n hash = compute_deployment_hash(config, self.file, checked_image._copies)\n create_deployment_tarball(tar_path, config, self.file, checked_image._copies)\n size = getsize(tar_path)\n return (hash, size)\n\n def _heartbeat_thread(self, deployment: ManagedDeployment, ttl: float):\n interval = min(ttl / 2, 60)\n while True:\n deployment.sync_heartbeat()\n time.sleep(interval)\n\n async def async_ensure_ready(self):\n if not isinstance(self.root_module, str):\n raise Exception(\n \"Trying to deploy a deployment not yet bound to a Python module. Have any remote functions been registered?\"\n )\n\n logger.info(f\"Preparing deployment '{self.name}' (module={self.root_module})\")\n with tempfile.NamedTemporaryFile(\n delete_on_close=False\n ) as f: # deletes at end of context manager instead of at f.close\n tar_path = f.name if not DEBUG_DEPLOYMENT_TAR else DEBUG_DEPLOYMENT_TAR\n f.close() # _into_deployment_config_and_tarball will reopen it when it makes the tarball\n config = self._into_deployment_config_and_tarball(tar_path)\n logger.debug(f\"Registering deployment with server\")\n deployment = await self.client.put_deployment(config)\n if deployment.needs_upload:\n logger.info(f\"Uploading deployment tarball\")\n await deployment.upload(tar_path)\n logger.info(f\"Upload complete\")\n else:\n logger.info(f\"Deployment tarball already up to date, skipping upload\")\n if deployment.action == \"soft_update\":\n wg_id = deployment.workergroup_id\n if not wg_id:\n wg_id = await self.client.find_workergroup_for_endpoint(\n deployment.endpoint_id\n )\n if wg_id:\n logger.info(f\"Triggering rolling update for workergroup {wg_id}\")\n await self.client.update_workers(wg_id)\n else:\n logger.warning(\n f\"soft_update but no workergroup found for endpoint {deployment.endpoint_id}, skipping update_workers\"\n )\n self._inner = _FullDeployment(self.root_module, deployment)\n if self._ttl is not None:\n threading.Thread(\n target=self._heartbeat_thread,\n args=(deployment, self._ttl),\n daemon=True,\n ).start()\n logger.info(f\"Deployment '{self.name}' is ready (id={deployment.id})\")\n\n def ensure_ready(self):\n asyncio.run(self.async_ensure_ready())\n\n def _unwrap_worker_response(self, response: dict[str, Any]) -> serialization.JSON:\n try:\n return response[\"response\"][\"result\"]\n except KeyError:\n raise Exception(\n f\"Remote function call failed with status code {response['status']}: {response['text']}\"\n )\n\n async def _dispatch(self, f_name, globals, sig, args, kwargs) -> Any:\n if not isinstance(self._inner, _FullDeployment):\n raise Exception(\"Deployment is not ready. Call .ensure_ready() first!\")\n bound_args = sig.bind(*args, **kwargs)\n bound_args.apply_defaults()\n route = \"/remote/\" + \"/\".join(f_name)\n logger.debug(f\"Dispatching remote call to {route}\")\n return serialization.deserialize_unwrap_error(\n self._unwrap_worker_response(\n await self._inner.deployment.endpoint.request(\n route,\n {\n \"kwargs\": {\n k: serialization.serialize(v, self._inner.root_module)\n for k, v in bound_args.arguments.items()\n }\n },\n )\n ),\n self._inner.root_module,\n globals,\n )\n\n def remote(\n self,\n f: Callable[P, Awaitable[Any]] | None = None,\n *,\n benchmark_dataset: list[dict] | None = None,\n benchmark_generator: Callable[[], dict] | None = None,\n benchmark_runs: int = 10,\n workload_calculator: Callable[..., float] | None = None,\n ) -> (\n Callable[P, Awaitable[Any]]\n | Callable[[Callable[P, Awaitable[Any]]], Callable[P, Awaitable[Any]]]\n ):\n def decorator(\n f: Callable[P, Awaitable[Any]], **_\n ) -> Callable[P, Awaitable[Any]]:\n f_rel_name = self.relativize(f)\n logger.debug(f\"Registered remote function: {f.__name__}\")\n f_globals = f.__globals__\n sig = inspect.signature(f)\n\n def inner(*args: P.args, **kwargs: P.kwargs) -> Awaitable[Any]:\n return self._dispatch(f_rel_name, f_globals, sig, args, kwargs)\n\n return inner\n\n if f is not None:\n return decorator(f)\n else:\n return decorator\n","content_type":"text/x-python; charset=utf-8","language":"python","size":12923,"content_sha256":"38c92bd2280a0fb936b9dda9d124e560dc334e07262846ee9232f9eb464d38dd"},{"filename":"serverless/remote/serialization.py","content":"import base64\nfrom typing import Any, TypeAlias\nfrom importlib import import_module\nimport logging\n\nlogger = logging.getLogger(\"vastai\")\n\n\ndef relativize_module(mod: str, root_module):\n if mod == root_module:\n return {\"type\": \"relative\", \"path\": []}\n elif mod.startswith(root_module):\n return {\n \"type\": \"relative\",\n \"path\": mod[len(root_module) :].lstrip(\".\").split(\".\"),\n } # will be non-empty\n else:\n return {\"type\": \"absolute\", \"path\": mod.split(\".\")}\n\n\ndef derelativize_module(mod, root_module: str):\n if mod[\"type\"] == \"relative\":\n prefix = [] if root_module == \"__main__\" else root_module.split(\".\")\n return \".\".join(prefix + mod[\"path\"])\n else:\n return \".\".join(mod[\"path\"])\n\n\nJSON: TypeAlias = int | str | float | list[\"JSON\"] | dict[str, \"JSON\"]\n\n\ndef serialize(obj, root_module: str) -> JSON:\n if type(obj) in [int, str, float, type(None)]:\n return obj\n elif type(obj) is bytes:\n return {\"type\": \"bytes\", \"contents\": base64.b64encode(obj).decode(\"utf-8\")}\n elif type(obj) in [list, tuple]:\n return {\n \"type\": type(obj).__name__,\n \"contents\": [serialize(child, root_module) for child in obj],\n }\n elif (\n type(obj) is dict\n ): # JSON dicts must have string keys, so instead we do this as list of tuples\n return {\n \"type\": \"dict\",\n \"contents\": [\n [serialize(k, root_module), serialize(v, root_module)]\n for k, v in obj.items()\n ],\n }\n elif hasattr(obj, \"__dict__\") and hasattr(\n obj, \"__class__\"\n ): # is a member of a normal Python class\n return {\n \"type\": \"obj\",\n \"module\": relativize_module(obj.__class__.__module__, root_module),\n \"class\": obj.__class__.__qualname__,\n \"contents\": {k: serialize(v, root_module) for k, v in obj.__dict__.items()},\n }\n raise TypeError(f\"Object of type {type(obj)} is not JSON serializable!\")\n\n\ndef deserialize(json, root_module: str, globals):\n if type(json) in [int, str, float, type(None)]:\n return json\n elif json[\"type\"] == \"bytes\":\n return base64.b64decode(json[\"contents\"])\n elif json[\"type\"] in [\"list\", \"tuple\"]:\n return __builtins__[json[\"type\"]](\n deserialize(child, root_module, globals) for child in json[\"contents\"]\n )\n elif json[\"type\"] == \"dict\":\n return dict(\n (deserialize(k, root_module, globals), deserialize(v, root_module, globals))\n for k, v in json[\"contents\"]\n )\n elif json[\"type\"] == \"obj\":\n modname = derelativize_module(json[\"module\"], root_module)\n namespace = globals if modname == \"\" else import_module(modname).__dict__\n cls = namespace[json[\"class\"]]\n obj = cls.__new__(cls)\n for k, v in json[\"contents\"].items():\n setattr(obj, k, deserialize(v, root_module, globals))\n return obj\n raise TypeError(\n f\"JSON does not correspond to known Vast deployment datatype: {json}\"\n )\n\n\ndef serialize_ok(obj, root_module: str):\n return {\"ok\": serialize(obj, root_module)}\n\n\ndef serialize_err(err, root_module: str):\n return {\"err\": serialize(err, root_module)}\n\n\ndef deserialize_unwrap_error(json, root_module: str, globals):\n logger.debug(f\"got json: \\n{json}\")\n for k, v in json.items():\n if (k) == \"ok\":\n return deserialize(v, root_module, globals)\n if (k) == \"err\":\n err = deserialize(v, root_module, globals)\n if isinstance(err, Exception):\n raise err\n raise TypeError(\n 'Expected json to have format {\"ok\": deserializable} or {\"err\": deserializes_to_error_object}: '\n + json\n )\n","content_type":"text/x-python; charset=utf-8","language":"python","size":3810,"content_sha256":"facfc33d9ce987e5834f8495e29f5d1b5580b0a3377a7dd0c0e033ecb8d5abfd"},{"filename":"serverless/remote/serve_deployment.py","content":"#!/usr/bin/python\n\"\"\"\nDeployment bootstrap script invoked by start_server.sh when IS_DEPLOYMENT=true.\n\nDownloads and extracts the deployment tarball, applies config.json\n(env vars, apt packages, pip packages, run scripts), then starts the\ndeployment worker.\n\"\"\"\n\nimport json\nimport os\nimport subprocess\nimport sys\nimport tarfile\nimport tempfile\nimport urllib.request\n\n\nVAST_API_URL = os.environ.get(\"VAST_API_URL\", \"https://console.vast.ai\")\n\n\ndef get_api_key() -> str:\n \"\"\"Resolve an API key from the container environment.\n\n Inside a Vast instance, Kaalia sets CONTAINER_API_KEY with the\n instance-level API key that has permissions to fetch deployment blobs.\n \"\"\"\n key = os.environ.get(\"CONTAINER_API_KEY\") or os.environ.get(\"VAST_API_KEY\")\n if not key:\n raise RuntimeError(\n \"No API key found in environment (CONTAINER_API_KEY or VAST_API_KEY)\"\n )\n return key\n\n\ndef get_download_url(deployment_id: str, api_key: str) -> str:\n \"\"\"Fetch the presigned S3 download URL for the deployment blob.\"\"\"\n url = f\"{VAST_API_URL}/api/v0/deployment/{deployment_id}/download_url/\"\n req = urllib.request.Request(url, headers={\"Authorization\": f\"Bearer {api_key}\"})\n with urllib.request.urlopen(req) as resp:\n data = json.loads(resp.read())\n if not data.get(\"success\"):\n raise RuntimeError(f\"Failed to get download URL: {data}\")\n return data[\"download_url\"]\n\n\ndef download_tarball(download_url: str) -> str:\n \"\"\"Download the deployment tarball to a temp file, returning its path.\"\"\"\n fd, path = tempfile.mkstemp(suffix=\".tar.gz\")\n try:\n with urllib.request.urlopen(download_url) as resp, os.fdopen(fd, \"wb\") as f:\n while True:\n chunk = resp.read(1 \u003c\u003c 20) # 1 MB\n if not chunk:\n break\n f.write(chunk)\n except Exception:\n os.unlink(path)\n raise\n return path\n\n\ndef extract_tarball(tarball_path: str):\n cwd = os.getcwd()\n with tarfile.open(tarball_path, \"r:gz\") as tf:\n for member in tf.getmembers():\n tf.extract(\n member, filter=\"fully_trusted\"\n ) # preserve absolute paths; deployment .tar files are created by client via deployments SDK and are trusted inside the container.\n\n\ndef get_config(path: str):\n \"\"\"Load and return config.json as a Config dataclass.\"\"\"\n from vastai.serverless.remote.base import Config\n\n with open(path) as f:\n raw = json.load(f)\n return Config(\n name=raw[\"name\"],\n pip_installs=raw.get(\"pip_installs\", []),\n apt_gets=raw.get(\"apt_gets\", []),\n envs=raw.get(\"envs\", []),\n runs=raw.get(\"runs\", []),\n )\n\n\ndef export_envs(envs: list):\n \"\"\"Export environment variables and append to /etc/environment.\"\"\"\n etc_lines = []\n for entry in envs:\n key, value = entry[0], entry[1]\n os.environ[key] = value\n etc_lines.append(f'{key}=\"{value}\"\\n')\n if etc_lines:\n with open(\"/etc/environment\", \"a\") as f:\n f.writelines(etc_lines)\n\n\ndef run_apt_gets(packages: list[str]):\n \"\"\"Install apt packages.\"\"\"\n if not packages:\n return\n subprocess.run(\n [\"apt-get\", \"update\", \"-qq\"],\n check=True,\n )\n subprocess.run(\n [\"apt-get\", \"install\", \"-y\", \"-qq\"] + packages,\n check=True,\n )\n\n\ndef run_pip_installs(packages: list[str]):\n \"\"\"Install pip packages using uv.\"\"\"\n if not packages:\n return\n subprocess.run(\n [\"uv\", \"pip\", \"install\"] + packages,\n check=True,\n )\n\n\ndef run_scripts(runs: list):\n \"\"\"Run scripts from config.\n\n Each entry is either a string (passed to sh -c) or a list of args (execvp style).\n \"\"\"\n for entry in runs:\n if isinstance(entry, str):\n subprocess.run([\"sh\", \"-c\", entry], check=True)\n elif isinstance(entry, (list, tuple)):\n subprocess.run(list(entry), check=True)\n else:\n raise ValueError(f\"Invalid run entry: {entry!r}\")\n\n\ndef main():\n deployment_id = os.environ.get(\"DEPLOYMENT_ID\")\n if not deployment_id:\n raise RuntimeError(\"DEPLOYMENT_ID environment variable not set\")\n\n debug_deployment_tar = os.environ.get(\"DEBUG_DEPLOYMENT_TAR\")\n if debug_deployment_tar:\n tarball_path = debug_deployment_tar\n else:\n api_key = get_api_key()\n\n # Download deployment tarball\n print(f\"Fetching download URL for deployment {deployment_id}\")\n download_url = get_download_url(deployment_id, api_key)\n\n print(\"Downloading deployment tarball\")\n tarball_path = download_tarball(download_url)\n\n # Extract preserving absolute paths — the deployment runs in a\n # container owned by the client, and we trust their path choices.\n print(\"Extracting deployment tarball\")\n try:\n extract_tarball(tarball_path)\n finally:\n os.unlink(tarball_path)\n\n # Load config\n config = get_config(\"config.json\")\n print(f\"Loaded config for deployment: {config.name}\")\n\n # (1) Export envs\n if config.envs:\n print(f\"Exporting {len(config.envs)} environment variables\")\n export_envs(config.envs)\n\n # (2) Install packages\n if config.apt_gets:\n print(f\"Installing {len(config.apt_gets)} apt packages\")\n run_apt_gets(config.apt_gets)\n\n if config.pip_installs:\n print(f\"Installing {len(config.pip_installs)} pip packages\")\n run_pip_installs(config.pip_installs)\n\n # (3) Run scripts\n if config.runs:\n print(f\"Running {len(config.runs)} setup scripts\")\n run_scripts(config.runs)\n\n # Look up and start deployment\n from vastai.serverless.remote.serve import Deployment\n\n # ensure deployment import time code is run in serve mode.\n os.environ[\"IS_DEPLOYMENT\"] = \"1\"\n # deployment module/package is guaranteed from tarball\n sys.path.insert(0, os.getcwd())\n import deployment # running this import has the side effect of registering deployments and remote functions with vastai.serverless.remote.serve.Deployment\n\n our_deployment = Deployment.lookup(config.name)\n if our_deployment is None:\n raise RuntimeError(f\"Failed to lookup registered deployment: {config.name}\")\n\n print(f\"Starting deployment worker: {config.name}\")\n worker = our_deployment.into_worker()\n worker.run()\n\n\nif __name__ == \"__main__\":\n main()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":6405,"content_sha256":"02068a319eedf3aac5db5375efb0be7331902f739f6a3df56a438c10613e4338"},{"filename":"serverless/remote/serve.py","content":"from .base import Config, Deployment_\nfrom ..server.worker import Worker, WorkerConfig, HandlerConfig, BenchmarkConfig\nfrom .serialization import serialize, deserialize, serialize_ok, serialize_err\nfrom typing import (\n ParamSpec,\n Type,\n Any,\n Callable,\n Awaitable,\n TypeVar,\n AsyncContextManager,\n Optional,\n)\nfrom dataclasses import dataclass\nimport asyncio\nimport logging\n\nlogger = logging.getLogger(\"vastai\")\n\nlogger.debug(\"mode: serve\")\n\n\n@dataclass\nclass RemoteFunc:\n func: Callable[..., Awaitable[Any]]\n globals: dict[str, Any]\n allow_parallel_requests: bool\n max_queue_time: float\n benchmark_dataset: Optional[list[dict[str, Any]]]\n benchmark_generator: Optional[Callable[[], dict[str, Any]]]\n benchmark_runs: int\n workload_calculator: Optional[Callable[..., float]]\n\n\nT = TypeVar(\"T\")\nP = ParamSpec(\"P\")\n\n\nclass Deployment(Deployment_, AsyncContextManager):\n def __init__(self, *args, **kwargs):\n super().__init__(*args, **kwargs)\n self.contexts: dict[Type, Any] = {}\n self.context_managers: dict[Type, AsyncContextManager] = {}\n self.remote_funcs: dict[tuple[str], RemoteFunc] = {}\n\n def context(\n self, *args, **kwargs\n ) -> Callable[\n [Type[AsyncContextManager[T]]], Type[AsyncContextManager[T]]\n ]: # registers context_class as deployment context\n def decorator(\n context_class: Type[AsyncContextManager[T]],\n ) -> Type[AsyncContextManager[T]]:\n self.context_managers[context_class] = context_class(*args, **kwargs)\n return context_class\n\n return decorator\n\n def get_context(self, context_class: Type[AsyncContextManager[T]]) -> T:\n if context_class not in self.contexts:\n raise KeyError(f\"Context {context_class} not registered or not yet entered\")\n return self.contexts[context_class]\n\n async def __aenter__(self):\n self.contexts = dict(\n zip(\n self.context_managers.keys(),\n await asyncio.gather(\n *(\n manager.__aenter__()\n for manager in self.context_managers.values()\n )\n ),\n )\n )\n\n async def __aexit__(self, exc_type, exc, tb):\n await asyncio.gather(\n *(\n self.context_managers[cls_].__aexit__(exc_type, exc, tb)\n for cls_ in self.contexts.keys() # only run aexits for contexts we've aentered\n )\n )\n self.contexts = {} # invalidate contexts\n\n def _wrap_remote_func(\n self, root_module: str, func: Callable[..., Awaitable[Any]], func_globals: dict\n ) -> Callable[..., Awaitable[Any]]:\n \"\"\"Wrap a remote function with serialization/deserialization.\n\n Expects payload: {\"args\": [serialized_args], \"kwargs\": {serialized_kwargs}}\n Returns: {\"ok\": serialized_result} on success, {\"err\": serialized_exception} on failure.\n \"\"\"\n\n async def wrapper(*, args: list = [], kwargs: dict = {}) -> dict:\n deserialized_args = [\n deserialize(a, root_module, func_globals) for a in args\n ]\n deserialized_kwargs = {\n k: deserialize(v, root_module, func_globals) for k, v in kwargs.items()\n }\n try:\n result = await func(*deserialized_args, **deserialized_kwargs)\n return serialize_ok(result, root_module)\n except Exception as e:\n return serialize_err(e, root_module)\n\n return wrapper\n\n def _wrap_workload_calculator(\n self,\n root_module: str,\n user_calc: Callable[..., float],\n func_globals: dict,\n ) -> Callable[[dict], float]:\n \"\"\"Deserialize the request payload and score it with the user's calculator.\n\n Runs synchronously inside count_workload and returns a plain float,\n unlike the async, result-serializing remote-function wrapper.\n \"\"\"\n\n def calculator(payload: dict) -> float:\n deserialized_args = [\n deserialize(a, root_module, func_globals)\n for a in payload.get(\"args\", [])\n ]\n deserialized_kwargs = {\n k: deserialize(v, root_module, func_globals)\n for k, v in payload.get(\"kwargs\", {}).items()\n }\n return float(user_calc(*deserialized_args, **deserialized_kwargs))\n\n return calculator\n\n def into_worker(self) -> Worker:\n handlers: list[HandlerConfig] = []\n if isinstance(self.root_module, str):\n for key, entry in self.remote_funcs.items():\n route = \"/remote/\" + \"/\".join(key)\n\n benchmark_config = None\n if (\n entry.benchmark_dataset is not None\n or entry.benchmark_generator is not None\n ):\n # Serialize benchmark dataset values so they match the\n # format a real client sends (the wrapper deserializes them).\n dataset = entry.benchmark_dataset\n if dataset is not None:\n dataset = [\n {\n \"kwargs\": {\n k: serialize(v, self.root_module)\n for k, v in item.items()\n }\n for item in dataset\n }\n ]\n benchmark_generator = None\n if entry.benchmark_generator is not None:\n root_module = self.root_module # in case self.root_module changes, bind to current root_module\n original_benchmark_generator = (\n entry.benchmark_generator\n ) # see above\n benchmark_generator = lambda: {\n \"kwargs\": {\n k: serialize(v, root_module)\n for k, v in original_benchmark_generator().items()\n }\n }\n\n benchmark_config = BenchmarkConfig(\n dataset=dataset,\n generator=benchmark_generator,\n runs=entry.benchmark_runs,\n )\n\n wrapped = self._wrap_remote_func(\n self.root_module, entry.func, entry.globals\n )\n\n workload_calculator = None\n if entry.workload_calculator is not None:\n workload_calculator = self._wrap_workload_calculator(\n self.root_module, entry.workload_calculator, entry.globals\n )\n\n handlers.append(\n HandlerConfig(\n route=route,\n remote_function=wrapped,\n allow_parallel_requests=entry.allow_parallel_requests,\n benchmark_config=benchmark_config,\n max_queue_time=entry.max_queue_time,\n workload_calculator=workload_calculator,\n )\n )\n\n config = WorkerConfig(\n handlers=handlers,\n lifecycle=self,\n )\n return Worker(config)\n raise TypeError(\"Refusing to create empty Worker!\")\n\n def remote(\n self,\n f: Callable[P, Awaitable[Any]] | None = None,\n *,\n allow_parallel_requests: bool = False,\n max_queue_time: float = 30.0,\n benchmark_dataset: list[dict] | None = None,\n benchmark_generator: Callable[[], dict] | None = None,\n benchmark_runs: int = 10,\n workload_calculator: Callable[..., float] | None = None,\n ) -> (\n Callable[P, Awaitable[Any]]\n | Callable[[Callable[P, Awaitable[Any]]], Callable[P, Awaitable[Any]]]\n ):\n def decorator(f: Callable[P, Awaitable[Any]]) -> Callable[P, Awaitable[Any]]:\n key = self.relativize(f)\n self.remote_funcs[key] = RemoteFunc(\n func=f,\n globals=f.__globals__,\n allow_parallel_requests=allow_parallel_requests,\n max_queue_time=max_queue_time,\n benchmark_dataset=benchmark_dataset,\n benchmark_generator=benchmark_generator,\n benchmark_runs=benchmark_runs,\n workload_calculator=workload_calculator,\n )\n return f\n\n if f is not None:\n return decorator(f)\n return decorator\n","content_type":"text/x-python; charset=utf-8","language":"python","size":8780,"content_sha256":"39c42c0e198cbc81799f2357a2559a1e9b8d74a0c2c97679a73558b8adf5e4ea"},{"filename":"serverless/remote/utils.py","content":"from __future__ import annotations\n\nimport dataclasses\nimport hashlib\nimport io\nimport json\nimport os\nimport re\nimport tarfile\nimport tempfile\nfrom typing import TYPE_CHECKING, Tuple, BinaryIO\n\nif TYPE_CHECKING:\n from vastai.serverless.remote.base import Config\n\n\ndef _sanitize_info(\n info: tarfile.TarInfo, arcname: str | None = None\n) -> tarfile.TarInfo:\n \"\"\"Set uid/gid to 0 on a TarInfo.\n\n If arcname is provided, override info.name. This is needed because\n gettarinfo strips leading '/' from absolute arcnames.\n \"\"\"\n info.uid = 0\n info.gid = 0\n info.uname = \"root\"\n info.gname = \"root\"\n if arcname is not None:\n info.name = arcname\n return info\n\n\ndef add_file(tf: tarfile.TarFile, src_path: str, arcname: str) -> None:\n \"\"\"Add a single file to the tarball under arcname, with uid/gid 0.\"\"\"\n info = tf.gettarinfo(src_path, arcname=arcname)\n _sanitize_info(info, arcname)\n with open(src_path, \"rb\") as f:\n tf.addfile(info, f)\n\n\ndef add_folder(tf: tarfile.TarFile, src_path: str, arcname: str) -> None:\n \"\"\"Recursively add a folder to the tarball, renaming the root to arcname, with uid/gid 0.\"\"\"\n src_path = os.path.normpath(src_path)\n for root, dirs, files in os.walk(src_path):\n rel = os.path.relpath(root, src_path)\n arc_root = arcname if rel == \".\" else os.path.join(arcname, rel)\n\n # Add directory entry\n info = tf.gettarinfo(root, arcname=arc_root)\n _sanitize_info(info, arc_root)\n tf.addfile(info)\n\n for fname in files:\n file_path = os.path.join(root, fname)\n arc_path = os.path.join(arc_root, fname)\n info = tf.gettarinfo(file_path, arcname=arc_path)\n _sanitize_info(info, arc_path)\n with open(file_path, \"rb\") as f:\n tf.addfile(info, f)\n\n\ndef add_string(\n tf: tarfile.TarFile,\n content: str,\n arcname: str,\n executable: bool = False,\n) -> None:\n \"\"\"Add a string as a text file to the tarball with uid/gid 0.\n\n If executable=True, the file gets mode 0o755; otherwise 0o644.\n \"\"\"\n data = content.encode(\"utf-8\")\n info = tarfile.TarInfo(name=arcname)\n info.size = len(data)\n info.mode = 0o755 if executable else 0o644\n _sanitize_info(info)\n tf.addfile(info, io.BytesIO(data))\n\n\ndef add_path(tf: tarfile.TarFile, src_path: str, arcname: str) -> None:\n \"\"\"Add a file or directory to the tarball, dispatching to add_file or add_folder.\"\"\"\n if os.path.isdir(src_path):\n add_folder(tf, src_path, arcname)\n elif os.path.isfile(src_path):\n add_file(tf, src_path, arcname)\n else:\n raise FileNotFoundError(f\"Source path does not exist: {src_path}\")\n\n\ndef serialize_config(config: Config) -> str:\n \"\"\"JSON-serialize a Config dataclass.\"\"\"\n return json.dumps(dataclasses.asdict(config))\n\n\ndef is_python_package(path: str) -> bool:\n \"\"\"Return True if path is a Python package directory (contains __init__.py).\"\"\"\n return os.path.isdir(path) and os.path.isfile(os.path.join(path, \"__init__.py\"))\n\n\ndef is_python_module(path: str) -> bool:\n \"\"\"Return True if path is a single .py file.\"\"\"\n return os.path.isfile(path) and path.endswith(\".py\")\n\n\ndef deployment_arcname(deployment_path: str) -> str:\n \"\"\"Return the archive name for a deployment: './deployment' for packages, './deployment.py' for modules.\n\n Raises ValueError if the path is neither a Python package nor a module.\n \"\"\"\n if is_python_package(deployment_path):\n return \"./deployment\"\n if is_python_module(deployment_path):\n return \"./deployment.py\"\n raise ValueError(\n f\"Deployment path is neither a Python package (directory with __init__.py) \"\n f\"nor a Python module (.py file): {deployment_path}\"\n )\n\n\nIGNORE_COMMENT_BODY = b\"!VAST_IGNORE_CHANGES\"\n\n# Matches the closing delimiter (with escape handling) for a multiline string.\n_CLOSE_MULTILINE_RE = {\n b'\"\"\"': re.compile(rb'(?:[^\\\\]|\\\\.)*?\"\"\"'),\n b\"'''\": re.compile(rb\"(?:[^\\\\]|\\\\.)*?'''\"),\n}\n\n# Tokenizer for Python source outside multiline strings.\n# Matches (in priority order): closed triple-quoted strings, unclosed\n# triple-quotes (groups 1/2), single-quoted strings, or comment start (group 3).\n# Unclosed triple-quotes MUST precede single-quote patterns, otherwise\n# \"...\" would consume \"\" from the start of \"\"\".\n# Unmatched text (code, whitespace, etc.) is simply skipped by finditer.\n_PYTHON_TOKEN_RE = re.compile(\n rb'\"\"\"(?:[^\\\\]|\\\\.)*?\"\"\"|'\n rb\"'''(?:[^\\\\]|\\\\.)*?'''|\"\n rb'(\"\"\")|'\n rb\"(''')|\"\n rb'\"(?:[^\"\\\\\\n]|\\\\.)*\"|'\n rb\"'(?:[^'\\\\\\n]|\\\\.)*'|\"\n rb\"(#)\"\n)\n\n\ndef _scan_python_line(\n line: bytes,\n in_multiline: bytes | None,\n) -> tuple[bool, bytes | None]:\n \"\"\"Scan a Python source line left-to-right, tracking string context.\n\n Returns (has_ignore_comment, new_multiline_state).\n \"\"\"\n i = 0\n if in_multiline is not None:\n m = _CLOSE_MULTILINE_RE[in_multiline].match(line)\n if m is None:\n return False, in_multiline\n i = m.end()\n\n for m in _PYTHON_TOKEN_RE.finditer(line, i):\n if m.group(1):\n return False, b'\"\"\"'\n if m.group(2):\n return False, b\"'''\"\n if m.group(3):\n comment_body = line[m.start() + 1 :].strip()\n return comment_body == IGNORE_COMMENT_BODY, None\n\n return False, None\n\n\ndef filter_ignored_lines(data: bytes) -> bytes:\n \"\"\"Remove lines whose comment is exactly ``#!VAST_IGNORE_CHANGES``.\n\n Tracks triple-quoted multiline strings so that occurrences inside string\n literals are not treated as comments.\n \"\"\"\n lines = data.split(b\"\\n\")\n result: list[bytes] = []\n in_multiline: bytes | None = None\n for line in lines:\n has_ignore, in_multiline = _scan_python_line(line, in_multiline)\n if not has_ignore:\n result.append(line)\n return b\"\\n\".join(result)\n\n\ndef read_file_for_hash(path: str, filter_comments: bool = False) -> bytes:\n \"\"\"Read a file's bytes for hashing.\n\n When filter_comments is True and the file is a .py file, lines whose\n comment is exactly ``# !VAST_IGNORE_CHANGES`` are excluded.\n \"\"\"\n with open(path, \"rb\") as f:\n data = f.read()\n if filter_comments and path.endswith(\".py\"):\n data = filter_ignored_lines(data)\n return data\n\n\ndef hash_update_file(\n hasher: hashlib._Hash,\n src_path: str,\n arcname: str,\n filter_comments: bool = False,\n) -> None:\n \"\"\"Feed a single file into the hasher: its archive name followed by its content.\"\"\"\n hasher.update(arcname.encode(\"utf-8\"))\n hasher.update(read_file_for_hash(src_path, filter_comments=filter_comments))\n\n\ndef hash_update_directory(\n hasher: hashlib._Hash,\n src_path: str,\n arcname: str,\n filter_comments: bool = False,\n) -> None:\n \"\"\"Walk a directory in sorted order, feeding each file into the hasher.\"\"\"\n src_path = os.path.normpath(src_path)\n for root, dirs, files in os.walk(src_path):\n dirs.sort()\n rel = os.path.relpath(root, src_path)\n arc_root = arcname if rel == \".\" else os.path.join(arcname, rel)\n for fname in sorted(files):\n hash_update_file(\n hasher,\n os.path.join(root, fname),\n os.path.join(arc_root, fname),\n filter_comments=filter_comments,\n )\n\n\ndef hash_update_path(\n hasher: hashlib._Hash,\n src_path: str,\n arcname: str,\n filter_comments: bool = False,\n) -> None:\n \"\"\"Feed a file or directory into the hasher, dispatching appropriately.\"\"\"\n if os.path.isdir(src_path):\n hash_update_directory(\n hasher, src_path, arcname, filter_comments=filter_comments\n )\n elif os.path.isfile(src_path):\n hash_update_file(hasher, src_path, arcname, filter_comments=filter_comments)\n else:\n raise FileNotFoundError(f\"Source path does not exist: {src_path}\")\n\n\ndef compute_deployment_hash(\n config: Config,\n deployment_path: str,\n extra_files: list[tuple[str, str]] | None = None,\n) -> str:\n \"\"\"Compute a SHA-256 hex digest over deployment data without creating a tarball.\n\n Hashes (in order):\n 1. The JSON-serialized Config\n 2. The deployment files (with #!VAST_IGNORE_CHANGES comment lines excluded from .py files)\n 3. Any extra files (no comment filtering)\n\n Args:\n config: The deployment Config object.\n deployment_path: Path to the deployment source (a .py file or a package directory).\n extra_files: List of (source_path, dest_path) pairs.\n \"\"\"\n hasher = hashlib.sha256()\n hasher.update(serialize_config(config).encode(\"utf-8\"))\n hash_update_path(\n hasher,\n deployment_path,\n deployment_arcname(deployment_path),\n filter_comments=True,\n )\n for src, dest in extra_files or []:\n hash_update_path(hasher, src, dest)\n return hasher.hexdigest()\n\n\ndef create_deployment_tarball(\n tar_path: str,\n config: Config,\n deployment_path: str,\n extra_files: list[tuple[str, str]] | None = None\n):\n \"\"\"Create a deployment tarball in /tmp and return its path.\n\n The tarball contains:\n - ./config.json — JSON-serialized Config\n - ./deployment[.py] — the deployment package or module\n - extra files at their specified destination paths\n\n Args:\n config: The deployment Config object.\n deployment_path: Path to the deployment source (a .py file or a package directory).\n extra_files: List of (source_path, dest_path) pairs. dest_path may be absolute.\n \"\"\"\n tf = tarfile.TarFile.open(tar_path, mode=\"w:gz\")\n try:\n add_string(tf, serialize_config(config), \"./config.json\")\n add_path(tf, deployment_path, deployment_arcname(deployment_path))\n for src, dest in extra_files or []:\n add_path(tf, src, dest)\n finally:\n tf.close()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":9924,"content_sha256":"066ff079f9261ebb862d5e62eff3dedba4f92658c16a80d05f18111114101ef7"},{"filename":"serverless/server/__init__.py","content":"","content_type":"text/x-python; charset=utf-8","language":"python","size":0,"content_sha256":"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"},{"filename":"serverless/server/lib/__init__.py","content":"","content_type":"text/x-python; charset=utf-8","language":"python","size":0,"content_sha256":"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"},{"filename":"serverless/server/lib/data_types.py","content":"import time\nimport logging\nfrom dataclasses import dataclass, field\nfrom enum import Enum\nfrom abc import ABC, abstractmethod\nfrom typing import (\n Dict,\n Any,\n Union,\n Tuple,\n Optional,\n Set,\n TypeVar,\n Generic,\n Type,\n Callable,\n Awaitable,\n)\nfrom aiohttp import web, ClientResponse\nimport inspect\nimport random\nimport json\nimport psutil\nimport os\nimport asyncio\n\n\"\"\"\ntype variable representing an incoming payload to pyworker that will used to calculate load and will then\nbe forwarded to the model\n\"\"\"\n\nlog = logging.getLogger(__file__)\n\n\nclass JsonDataException(Exception):\n def __init__(self, json_msg: Dict[str, Any]):\n self.message = json_msg\n\n\nApiPayload_T = TypeVar(\"ApiPayload_T\", bound=\"ApiPayload\")\n\n\n@dataclass\nclass ApiPayload(ABC):\n @classmethod\n @abstractmethod\n def for_test(cls: Type[ApiPayload_T]) -> ApiPayload_T:\n \"\"\"defines how to create a payload for load testing\"\"\"\n pass\n\n @abstractmethod\n def generate_payload_json(self) -> Dict[str, Any]:\n \"\"\"defines how to convert an ApiPayload to JSON that will be sent to model API\"\"\"\n pass\n\n def generate_payload_multipart(self) -> Optional[Dict[str, Any]]:\n \"\"\"form fields to POST as multipart, or None (default) to POST JSON\"\"\"\n return None\n\n @abstractmethod\n def count_workload(self) -> float:\n \"\"\"defines how to calculate workload for a payload\"\"\"\n pass\n\n @classmethod\n @abstractmethod\n def from_json_msg(\n cls: Type[ApiPayload_T], json_msg: Dict[str, Any]\n ) -> ApiPayload_T:\n \"\"\"\n defines how to create an API payload from a JSON message,\n it should throw an JsonDataException if there are issues with some fields\n or they are missing in the format of\n {\n \"field\": \"error msg\"\n }\n \"\"\"\n pass\n\n\n@dataclass\nclass AuthData:\n \"\"\"data used to authenticate requester\"\"\"\n\n cost: str\n endpoint: str\n reqnum: int\n request_idx: int\n signature: str\n url: str\n\n @classmethod\n def from_json_msg(cls, json_msg: Dict[str, Any]):\n errors = {}\n for param in inspect.signature(cls).parameters:\n if param not in json_msg:\n errors[param] = \"missing parameter\"\n if errors:\n raise JsonDataException(errors)\n return cls(\n **{\n k: v\n for k, v in json_msg.items()\n if k in inspect.signature(cls).parameters\n }\n )\n\n\n@dataclass\nclass EndpointHandler(ABC, Generic[ApiPayload_T]):\n \"\"\"\n Each model endpoint will have a handler responsible for counting workload from the incoming ApiPayload\n and converting it to json to be forwarded to model API\n \"\"\"\n\n has_benchmark: bool = True\n concurrency: int = 10\n benchmark_runs: int = 8\n allow_parallel_requests: bool = False\n max_queue_time: float = None\n remote_function: Callable = None\n do_warmup_benchmark: bool = True\n\n @property\n @abstractmethod\n def endpoint(self) -> str:\n \"\"\"the endpoint on the model API\"\"\"\n pass\n\n @property\n @abstractmethod\n def healthcheck_endpoint(self) -> Optional[str]:\n \"\"\"the endpoint on the model API that is used for healthchecks\"\"\"\n pass\n\n @classmethod\n @abstractmethod\n def payload_cls(cls) -> Type[ApiPayload_T]:\n \"\"\"ApiPayload class\"\"\"\n pass\n\n @abstractmethod\n def make_benchmark_payload(self) -> ApiPayload_T:\n \"\"\"defines how to create an ApiPayload for benchmarking.\"\"\"\n pass\n\n @abstractmethod\n async def generate_client_response(\n self, client_request: web.Request, model_response: ClientResponse\n ) -> Union[web.Response, web.StreamResponse]:\n \"\"\"\n defines how to convert a model API response to a response to PyWorker client\n \"\"\"\n pass\n\n @abstractmethod\n async def call_remote_dispatch_function(self, params: dict):\n \"\"\"\n define a remote function for this endpoint, return the result\n \"\"\"\n pass\n\n @classmethod\n def get_data_from_request(\n cls, req_data: Dict[str, Any]\n ) -> Tuple[AuthData, ApiPayload_T]:\n errors = {}\n auth_data: Optional[AuthData] = None\n session_id: Optional[str] = None\n payload: Optional[ApiPayload_T] = None\n try:\n if \"auth_data\" in req_data:\n auth_data = AuthData.from_json_msg(req_data[\"auth_data\"])\n else:\n errors[\"auth_data\"] = \"field missing\"\n except JsonDataException as e:\n errors[\"auth_data\"] = e.message\n try:\n if \"payload\" in req_data:\n payload_cls = cls.payload_cls()\n payload = payload_cls.from_json_msg(req_data[\"payload\"])\n else:\n errors[\"payload\"] = \"field missing\"\n except JsonDataException as e:\n errors[\"payload\"] = e.message\n\n if \"session_id\" in req_data:\n session_id = req_data[\"session_id\"]\n\n if errors:\n raise JsonDataException(errors)\n if auth_data and payload:\n return (auth_data, payload, session_id)\n else:\n raise Exception(\"error deserializing request data\")\n\n\n@dataclass\nclass Session:\n session_id: str\n lifetime: float # extends TTL per-request\n auth_data: dict\n expiration: float # epoch seconds\n on_close_route: str\n on_close_payload: dict\n requests: list[web.Request] = field(default_factory=list)\n created_at: float = field(default_factory=time.time)\n request_idx: int = 0 # Request idx associated with this session\n session_reqnum: int = 0 # Internal per-session request count\n\n\n@dataclass\nclass RequestMetrics:\n \"\"\"Tracks metrics for an active request.\"\"\"\n\n request_idx: int\n reqnum: int\n workload: float\n status: str\n success: bool = False\n is_session: bool = False\n session: Session = None\n session_reqnum: Optional[int] = None\n entered_queue_at: float = 0.0\n work_started_at: float = 0.0\n work_completed_at: float = 0.0\n\n\n@dataclass\nclass SystemMetrics:\n \"\"\"General system metrics\"\"\"\n\n model_loading_start: float\n model_loading_time: Union[float, None]\n last_disk_usage: float\n additional_disk_usage: float\n model_is_loaded: bool\n disk_usage: float = 0.0\n\n @staticmethod\n def get_disk_usage_GB():\n return psutil.disk_usage(\"/\").used / (2**30) # want units of GB\n\n @classmethod\n def empty(cls):\n disk_usage = SystemMetrics.get_disk_usage_GB()\n return cls(\n model_loading_start=time.time(),\n model_loading_time=None,\n last_disk_usage=disk_usage,\n additional_disk_usage=0.0,\n model_is_loaded=False,\n disk_usage=disk_usage,\n )\n\n def update_disk_usage(self):\n \"\"\"Refresh the absolute figure; the delta is measured against the last successful report.\"\"\"\n self.disk_usage = SystemMetrics.get_disk_usage_GB()\n self.additional_disk_usage = self.disk_usage - self.last_disk_usage\n\n def commit_disk_usage(self):\n \"\"\"Call once a report carrying the current figures was delivered.\"\"\"\n self.last_disk_usage = self.disk_usage\n\n\n@dataclass\nclass BenchmarkResult:\n request_idx: int\n workload: float\n task: Awaitable[ClientResponse]\n response: Optional[ClientResponse] = None\n\n @property\n def is_successful(self) -> bool:\n return self.response is not None and self.response.status == 200\n\n\n@dataclass\nclass ModelMetrics:\n \"\"\"Model specific metrics\"\"\"\n\n # these are reset after being sent to autoscaler\n workload_served: float\n workload_received: float\n workload_cancelled: float\n workload_errored: float\n workload_rejected: float\n # these are not\n workload_pending: float\n error_msg: Optional[str]\n max_throughput: float\n requests_recieved: Set[int] = field(default_factory=set)\n requests_working: dict[int, RequestMetrics] = field(default_factory=dict)\n requests_deleting: list[RequestMetrics] = field(default_factory=list)\n last_update: float = field(default_factory=time.time)\n\n @classmethod\n def empty(cls):\n return cls(\n workload_pending=0.0,\n workload_served=0.0,\n workload_cancelled=0.0,\n workload_errored=0.0,\n workload_rejected=0.0,\n workload_received=0.0,\n error_msg=None,\n max_throughput=0.0,\n )\n\n @property\n def workload_processing(self) -> float:\n return max(self.workload_received - self.workload_cancelled, 0.0)\n\n @property\n def wait_time(self) -> float:\n if len(self.requests_working) == 0:\n return 0.0\n return sum(\n [\n request.workload\n for request in self.requests_working.values()\n if not request.is_session\n ]\n ) / max(self.max_throughput, 0.00001)\n\n @property\n def cur_load(self) -> float:\n return sum([request.workload for request in self.requests_working.values()])\n\n @property\n def working_request_idxs(self) -> list[int]:\n return [req.request_idx for req in self.requests_working.values()]\n\n def set_errored(self, error_msg):\n self.reset()\n self.error_msg = error_msg\n\n def reset(self):\n self.workload_served = 0\n self.workload_received = 0\n self.workload_cancelled = 0\n self.workload_errored = 0\n self.workload_rejected = 0\n self.last_update = time.time()\n\n\n@dataclass\nclass WorkerStatusData:\n \"\"\"Data that is reported to autoscaler's on_status endpoint\"\"\"\n\n id: int\n mtoken: str\n version: str\n loadtime: float\n cur_load: float\n rej_load: float\n new_load: float\n error_msg: str\n max_perf: float\n cur_perf: float\n cur_capacity: float\n max_capacity: float\n num_requests_working: int\n num_requests_recieved: int\n additional_disk_usage: float\n disk_usage: float\n working_request_idxs: list[int]\n url: str\n\n\nclass LogAction(Enum):\n \"\"\"\n These actions tell the backend what a log value means, for example:\n actions [\n # this marks the model server as loaded\n (LogAction.ModelLoaded, \"Starting server\"),\n # these mark the model server as errored\n (LogAction.ModelError, \"Exception loading model\"),\n (LogAction.ModelError, \"Server failed to bind to port\"),\n # this tells the backend to print any logs containing the string into its own logs\n # which are visible in the vast console instance logs\n (LogAction.Info, \"Starting model download\"),\n ]\n \"\"\"\n\n ModelLoaded = 1\n ModelError = 2\n Info = 3\n","content_type":"text/x-python; charset=utf-8","language":"python","size":10804,"content_sha256":"82335201d4c9b47532671c48351a0f0a9d72dc4566e5cd16d4a8fc852f4965d9"},{"filename":"serverless/server/lib/metrics.py","content":"import os\nimport time\nimport logging\nimport json\nfrom asyncio import sleep\nfrom dataclasses import dataclass, asdict, field\nfrom functools import cache\nimport asyncio\nfrom aiohttp import ClientSession, ClientTimeout, TCPConnector, ClientResponseError\n\nfrom .data_types import WorkerStatusData, SystemMetrics, ModelMetrics, RequestMetrics\nfrom typing import Awaitable, NoReturn, List\n\nMETRICS_UPDATE_INTERVAL = 1\nDELETE_REQUESTS_INTERVAL = 1\n\nlog = logging.getLogger(__file__)\n\n\n@cache\ndef get_url() -> str:\n use_ssl = os.environ.get(\"USE_SSL\", \"false\") == \"true\"\n worker_port = os.environ[f\"VAST_TCP_PORT_{os.environ['WORKER_PORT']}\"]\n public_ip = os.environ[\"PUBLIC_IPADDR\"]\n return f\"http{'s' if use_ssl else ''}://{public_ip}:{worker_port}\"\n\n\n@dataclass\nclass Metrics:\n version: str = \"0\"\n mtoken: str = \"\"\n last_metric_update: float = 0.0\n last_request_served: float = 0.0\n update_pending: bool = False\n id: int = field(default_factory=lambda: int(os.environ[\"CONTAINER_ID\"]))\n report_addr: List[str] = field(\n default_factory=lambda: os.environ[\"REPORT_ADDR\"].split(\",\")\n )\n url: str = field(default_factory=get_url)\n system_metrics: SystemMetrics = field(default_factory=SystemMetrics.empty)\n model_metrics: ModelMetrics = field(default_factory=ModelMetrics.empty)\n _session: ClientSession | None = field(default=None, init=False, repr=False)\n\n async def http(self) -> ClientSession:\n if self._session is None:\n self._session = ClientSession(\n timeout=ClientTimeout(total=10),\n connector=TCPConnector(\n limit=8,\n limit_per_host=4,\n force_close=True,\n enable_cleanup_closed=True,\n ),\n )\n return self._session\n \n async def aclose(self) -> None:\n if self._session is not None:\n await self._session.close()\n self._session = None\n\n def _request_id(self, request: RequestMetrics) -> str:\n \"\"\"Return a formatted identifier for logging based on request type.\"\"\"\n if request.is_session:\n return f\"Session {request.request_idx}\"\n elif request.session:\n return f\"Request {request.session_reqnum} in Session {request.session.request_idx}\"\n else:\n return f\"Request {request.request_idx}\"\n\n def _request_start(self, request: RequestMetrics) -> None:\n \"\"\"\n this function is called prior to forwarding a request to a model API.\n \"\"\"\n log.debug(f\"Starting {self._request_id(request)}\")\n request.status = \"Started\"\n self.model_metrics.workload_pending += request.workload\n self.model_metrics.workload_received += request.workload\n if not request.session:\n self.model_metrics.requests_recieved.add(request.reqnum)\n self.model_metrics.requests_working[request.reqnum] = request\n self.update_pending = True\n\n def _request_end(self, request: RequestMetrics) -> None:\n \"\"\"\n this function is called after handling of a request ends, regardless of the outcome\n \"\"\"\n log.debug(f\"Ending {self._request_id(request)}\")\n self.model_metrics.workload_pending -= request.workload\n if not request.session:\n self.model_metrics.requests_working.pop(request.reqnum, None)\n self.model_metrics.requests_deleting.append(request)\n self.last_request_served = time.time()\n\n def _request_success(self, request: RequestMetrics) -> None:\n \"\"\"\n this function is called after a response from model API is received and forwarded.\n \"\"\"\n log.debug(f\"{self._request_id(request)} succeeded\")\n request.work_completed_at = time.time()\n self.model_metrics.workload_served += request.workload\n request.status = \"Success\"\n request.success = True\n self.update_pending = True\n\n def _request_errored(self, request: RequestMetrics, message: str) -> None:\n \"\"\"\n this function is called if model API returns an error\n \"\"\"\n log.error(f\"{self._request_id(request)} errored: {message}\")\n request.work_completed_at = time.time()\n self.model_metrics.workload_errored += request.workload\n request.status = \"Error\"\n request.success = False\n self.update_pending = True\n\n def _request_canceled(self, request: RequestMetrics) -> None:\n \"\"\"\n this function is called if client drops connection before model API has responded\n \"\"\"\n log.debug(f\"Canceling {self._request_id(request)}\")\n self.model_metrics.workload_cancelled += request.workload\n request.success = True\n request.status = \"Cancelled\"\n \n def _request_reject(self, request: RequestMetrics) -> None:\n \"\"\"\n this function is called if the current wait time for the model is above max_queue_time\n \"\"\"\n log.debug(f\"Rejecting {self._request_id(request)}\")\n self.model_metrics.workload_rejected += request.workload\n if not request.session:\n self.model_metrics.requests_recieved.add(request.reqnum)\n self.model_metrics.requests_deleting.append(request)\n request.success = False\n request.status = \"Rejected\"\n self.update_pending = True\n\n async def _send_delete_requests_loop(self) -> Awaitable[NoReturn]:\n while True:\n await sleep(DELETE_REQUESTS_INTERVAL)\n if len(self.model_metrics.requests_deleting) > 0:\n await self.__send_delete_requests_and_reset()\n\n async def _send_metrics_loop(self) -> Awaitable[NoReturn]:\n while True:\n await sleep(METRICS_UPDATE_INTERVAL)\n elapsed = time.time() - self.last_metric_update\n if self.system_metrics.model_is_loaded is False and elapsed >= 10:\n await self.__send_metrics_and_reset()\n elif self.update_pending or elapsed > 10:\n await self.__send_metrics_and_reset()\n\n def _model_loaded(self, max_throughput: float) -> None:\n self.system_metrics.model_loading_time = (\n time.time() - self.system_metrics.model_loading_start\n )\n self.system_metrics.model_is_loaded = True\n self.model_metrics.max_throughput = max_throughput\n\n def _model_errored(self, error_msg: str) -> None:\n self.model_metrics.set_errored(error_msg)\n self.system_metrics.model_is_loaded = True\n\n def _set_version(self, version: str) -> None:\n self.version = version\n\n def _set_mtoken(self, mtoken: str) -> None:\n self.mtoken = mtoken\n\n #######################################Private#######################################\n\n async def __send_delete_requests_and_reset(self):\n async def post(report_addr: str, requests: list[dict]) -> bool:\n data = {\n \"worker_id\": self.id,\n \"mtoken\": self.mtoken,\n \"requests\": requests,\n }\n log.debug(\n f\"Deleting requests: {[r['request_idx'] for r in requests]}\"\n )\n full_path = report_addr.rstrip(\"/\") + \"/delete_requests/\"\n for attempt in range(1, 4):\n try:\n session = await self.http()\n async with session.post(full_path, json=data) as res:\n log.debug(f\"delete_requests response: {res.status}\")\n res.raise_for_status()\n return True\n except asyncio.TimeoutError:\n log.debug(\"delete_requests timed out\")\n except (ClientResponseError, Exception) as e:\n log.debug(f\"delete_requests failed with error: {e}\")\n await asyncio.sleep(2)\n log.debug(f\"Retrying delete_request, attempt: {attempt}\")\n return False\n\n # Take a snapshot of what we plan to send this tick.\n # New arrivals after this snapshot will remain in the queue for the next tick.\n snapshot = list(self.model_metrics.requests_deleting)\n\n if not snapshot:\n return # nothing to do\n\n requests_payload = [\n {\n \"request_idx\": r.request_idx,\n \"success\": r.success,\n \"status\": r.status,\n \"entered_queue_at\": r.entered_queue_at,\n \"work_started_at\": r.work_started_at,\n \"work_completed_at\": r.work_completed_at,\n }\n for r in snapshot\n ]\n\n for report_addr in self.report_addr:\n # TODO: Add a Redis subscriber queue for delete_requests\n if report_addr == \"https://cloud.vast.ai/api/v0\":\n # Patch: ignore the Redis API report_addr\n continue\n\n sent = await post(report_addr, requests_payload)\n\n if sent:\n # Remove only the items we actually sent from the live queue.\n sent_set = {r.request_idx for r in snapshot}\n self.model_metrics.requests_deleting[:] = [\n r for r in self.model_metrics.requests_deleting\n if r.request_idx not in sent_set\n ]\n break\n\n\n async def __send_metrics_and_reset(self):\n\n loadtime_snapshot = self.system_metrics.model_loading_time\n\n def get_worker_status_data() -> WorkerStatusData:\n return WorkerStatusData(\n id=self.id,\n mtoken=self.mtoken,\n version=self.version,\n loadtime=(loadtime_snapshot or 0.0), \n new_load=self.model_metrics.workload_processing,\n cur_load=self.model_metrics.cur_load,\n rej_load=self.model_metrics.workload_rejected,\n max_perf=self.model_metrics.max_throughput,\n cur_perf=self.model_metrics.workload_served,\n error_msg=self.model_metrics.error_msg or \"\",\n num_requests_working=len(self.model_metrics.requests_working),\n num_requests_recieved=len(self.model_metrics.requests_recieved),\n additional_disk_usage=self.system_metrics.additional_disk_usage,\n disk_usage=self.system_metrics.disk_usage,\n working_request_idxs=self.model_metrics.working_request_idxs,\n cur_capacity=0,\n max_capacity=0,\n url=self.url,\n )\n\n async def send_data(report_addr: str) -> bool:\n data = get_worker_status_data()\n log_data = asdict(data)\n def obfuscate(secret: str) -> str:\n if secret is None:\n return \"\"\n return secret[:7] + \"...\" if len(secret) > 12 else (\"*\" * len(secret))\n \n log_data[\"mtoken\"] = obfuscate(log_data.get(\"mtoken\"))\n log.debug(\n \"\\n\".join(\n [\n \"\\n\" + \"#\" * 20 + \" Worker Status \" + \"#\" * 20,\n f\"{json.dumps(log_data, indent=2)}\",\n \"#\" * 55,\n ]\n )\n )\n\n full_path = report_addr.rstrip(\"/\") + \"/worker_status/\"\n for attempt in range(1, 4):\n try:\n session = await self.http()\n async with session.post(full_path, json=asdict(data)) as res:\n res.raise_for_status()\n return True\n except asyncio.TimeoutError:\n log.debug(f\"Request to /worker_status/ timed out\")\n except (ClientResponseError, Exception) as e:\n log.debug(f\"Request to /worker_status/ failed with error: {e}\")\n await asyncio.sleep(2)\n log.debug(f\"Retrying /worker_status/ request, attempt: {attempt}\")\n log.debug(f\"Failed to send worker status through {report_addr}\")\n return False\n\n ###########\n\n self.system_metrics.update_disk_usage()\n\n sent = False\n for report_addr in self.report_addr:\n if await send_data(report_addr):\n sent = True\n break\n\n if sent:\n self.update_pending = False\n self.model_metrics.reset()\n self.system_metrics.commit_disk_usage()\n self.last_metric_update = time.time()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":12497,"content_sha256":"d25c80f90cfe6275579addd46f9819f0ac6584049efe096ce6c0889b389341a5"},{"filename":"serverless/server/lib/server.py","content":"import os\nimport logging\nfrom typing import List\nimport ssl\nfrom asyncio import run, gather\nimport asyncio\n\nfrom .backend import Backend\nfrom .metrics import Metrics\nfrom aiohttp import web\n\nlog = logging.getLogger(__file__)\n\nasync def start_server_async(backend: Backend, routes: List[web.RouteDef], **kwargs):\n try:\n use_ssl = os.environ.get(\"USE_SSL\", \"false\") == \"true\"\n if use_ssl is True:\n log.debug(\"Getting SSL Certificate from /etc/instance.crt\")\n try:\n ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)\n ssl_context.load_cert_chain(\n certfile=\"/etc/instance.crt\",\n keyfile=\"/etc/instance.key\",\n )\n except Exception as ex:\n raise Exception(f\"Failed to get SSL Certificate: {ex}\")\n else:\n ssl_context = None\n\n log.debug(\"Starting Worker Server...\")\n app = web.Application(client_max_size=0)\n app.add_routes(routes)\n\n # Hardcoded session routes\n app.router.add_post(\"/session/create\", backend.session_create_handler)\n app.router.add_post(\"/session/end\", backend.session_end_handler)\n app.router.add_post(\"/session/get\", backend.session_get_handler)\n app.router.add_post(\"/session/health\", backend.session_health_handler)\n\n # Pyworker management routes\n app.router.add_post(\"/pyworker/update\", backend.pyworker_update_handler)\n\n runner = web.AppRunner(app, handler_cancellation=True)\n await runner.setup()\n site = web.TCPSite(\n runner,\n ssl_context=ssl_context,\n port=int(os.environ[\"WORKER_PORT\"]),\n **kwargs\n )\n\n # Create HTTP-only app for internal webhooks to /session/end\n http_app = web.Application()\n http_app.router.add_post(\"/session/end\", backend.session_end_handler)\n\n http_runner = web.AppRunner(http_app)\n await http_runner.setup()\n http_port = int(os.environ.get(\"WORKER_HTTP_PORT\", int(os.environ[\"WORKER_PORT\"]) + 1))\n http_site = web.TCPSite(\n http_runner,\n ssl_context=None, # No SSL for internal webhooks\n port=http_port,\n **kwargs\n )\n\n log.debug(f\"Starting HTTP-only server for /session/end on port {http_port}\")\n await site.start()\n await http_site.start()\n log.debug(\"HTTP servers started, beginning log tracking\")\n await backend._start_tracking()\n\n except Exception as e:\n err_msg = f\"Worker Server failed to launch: {e}\"\n log.error(err_msg)\n\n async def beacon():\n metrics = Metrics()\n metrics._set_version(getattr(backend, \"version\", \"0\"))\n metrics._set_mtoken(getattr(backend, \"mtoken\", \"\"))\n try:\n while True:\n metrics._model_errored(err_msg)\n await metrics._Metrics__send_metrics_and_reset()\n await asyncio.sleep(10)\n finally:\n await metrics.aclose()\n\n await beacon()\n\n\n\ndef start_server(backend: Backend, routes: List[web.RouteDef], **kwargs):\n run(start_server_async(backend, routes, **kwargs))\n\n","content_type":"text/x-python; charset=utf-8","language":"python","size":3266,"content_sha256":"97d42071a141872f86e4713fb2014ef548cc2d2abbf3262963802a70a4e157b4"},{"filename":"serverless/server/worker.py","content":"from vastai.serverless.server.lib import backend, server\nfrom vastai.serverless.server.lib.data_types import (\n ApiPayload,\n EndpointHandler,\n LogAction,\n JsonDataException,\n)\nfrom dataclasses import dataclass, field\nfrom aiohttp import web, ClientResponse\nimport logging\nimport json\nimport math\nimport random\nimport inspect\nimport logging\nfrom typing import (\n Optional,\n Dict,\n Callable,\n Awaitable,\n Union,\n Any,\n Type,\n AsyncContextManager,\n)\n\n# Callable types\nRequestPayloadParser = Callable[[Dict[str, Any]], Dict[str, Any]]\nClientResponseGenerator = Callable[\n [web.Request, ClientResponse], Awaitable[Union[web.Response, web.StreamResponse]]\n]\nWorkloadCalculator = Callable[[Dict[str, Any]], float]\n\nlog = logging.getLogger(__name__)\n\nDEFAULT_WORKLOAD = 100.0\n\n\n@dataclass\nclass LogActionConfig:\n \"\"\"Configuration for defining log actions\"\"\"\n\n on_load: list[str] = field(default_factory=list)\n on_error: list[str] = field(default_factory=list)\n on_info: list[str] = field(default_factory=list)\n\n @property\n def log_actions(self) -> list[LogAction]:\n log_actions_ = []\n log_actions_.extend([(LogAction.ModelLoaded, log) for log in self.on_load])\n log_actions_.extend([(LogAction.ModelError, log) for log in self.on_error])\n log_actions_.extend([(LogAction.Info, log) for log in self.on_info])\n return log_actions_\n\n\n@dataclass\nclass BenchmarkConfig:\n \"\"\"Configuration for defining a benchmark\"\"\"\n\n dataset: list[dict] | None = None\n generator: Callable[[], dict] | None = None # optional sample factory\n runs: int = 8\n concurrency: int | None = 10\n do_warmup: bool = True\n\n\n@dataclass\nclass HandlerConfig:\n \"\"\"Configuration for defining handlers\"\"\"\n\n route: str\n healthcheck: Optional[str] = None\n allow_parallel_requests: bool = False\n max_queue_time: Optional[float] = 30.0\n benchmark_config: Optional[BenchmarkConfig] = None\n handler_class: Optional[Type[EndpointHandler]] = None\n payload_class: Optional[Type[ApiPayload]] = None\n request_parser: Optional[RequestPayloadParser] = None\n response_generator: Optional[ClientResponseGenerator] = None\n workload_calculator: Optional[WorkloadCalculator] = None\n is_remote_dispatch: bool = False\n remote_function: Optional[Callable] = None\n\n\n@dataclass\nclass WorkerConfig:\n model_server_url: str = None\n model_server_port: int = None\n model_log_file: str = None\n benchmark_data: list[dict[str, Any]] = field(default_factory=list)\n handlers: list[HandlerConfig] = field(default_factory=list)\n model_healthcheck_url: str = None\n benchmark_route: Optional[str] = None\n log_action_config: LogActionConfig = field(default_factory=LogActionConfig)\n max_sessions: Optional[int] = 10\n lifecycle: Optional[AsyncContextManager] = None\n\n\nclass EndpointHandlerFactory:\n \"\"\"Factory for creating endpoint handlers from WorkerConfig\"\"\"\n\n def __init__(self, config: WorkerConfig):\n self.config = config\n self._handlers: Dict[str, EndpointHandler] = {}\n self._build_handlers()\n\n def _build_handlers(self) -> None:\n \"\"\"Build endpoint handlers from config\"\"\"\n # If no handlers, create a default handler at /\n if not self.config.handlers:\n default_handler_config = HandlerConfig(\n route=\"/\", healthcheck=self.config.model_healthcheck_url\n )\n handler = self._create_handler(default_handler_config)\n self._handlers[\"/\"] = handler\n else:\n # Build handlers from config\n for handler_config in self.config.handlers:\n # TODO: Override values from handler class with HandlerConfig values\n if handler_config.handler_class is not None:\n # Use custom handler class\n handler = handler_config.handler_class()\n self._handlers[handler_config.route] = handler\n continue\n handler = self._create_handler(handler_config)\n self._handlers[handler_config.route] = handler\n\n def _create_handler(self, handler_config: HandlerConfig) -> EndpointHandler:\n \"\"\"Create a generic endpoint handler from HandlerConfig\"\"\"\n\n # Extract config values with defaults\n route_path = handler_config.route\n healthcheck_path = handler_config.healthcheck\n benchmark_config = handler_config.benchmark_config\n user_payload_class = handler_config.payload_class\n user_request_parser = handler_config.request_parser\n user_response_generator = handler_config.response_generator\n user_workload_calculator = handler_config.workload_calculator\n\n # If user provided a custom payload class, use it\n if user_payload_class:\n PayloadClass = user_payload_class\n else:\n # Create a generic ApiPayload class\n @dataclass\n class GenericApiPayload(ApiPayload):\n input: Dict[str, Any] = field(default_factory=dict)\n\n @classmethod\n def for_test(cls) -> \"GenericApiPayload\":\n # Use random.choice from benchmark_data if available\n if benchmark_config:\n if benchmark_config.dataset:\n test_data = random.choice(benchmark_config.dataset)\n elif benchmark_config.generator:\n try:\n test_data = benchmark_config.generator()\n except Exception as ex:\n raise (\n Exception(\n f'Error generating benchmark data for \"{handler_config.route}\" handler: {ex}'\n )\n )\n else:\n raise (\n Exception(\n f'Error generating benchmark data for \"{handler_config.route}\" handler: Missing BenchmarkConfig!'\n )\n )\n return cls(input=test_data.copy())\n\n def generate_payload_json(self) -> Dict[str, Any]:\n return self.input\n\n @classmethod\n def from_dict(cls, input: Dict[str, Any]) -> \"GenericApiPayload\":\n return cls(input=input)\n\n def count_workload(self) -> float:\n if not user_workload_calculator:\n return DEFAULT_WORKLOAD\n # Fall back rather than let a bad calculator 500 requests or\n # poison the autoscaler's load sums.\n try:\n workload = float(user_workload_calculator(self.input))\n except Exception:\n log.warning(\n f'workload_calculator for \"{route_path}\" raised; '\n f\"falling back to {DEFAULT_WORKLOAD}\",\n exc_info=True,\n )\n return DEFAULT_WORKLOAD\n if not math.isfinite(workload) or workload \u003c 0:\n log.warning(\n f'workload_calculator for \"{route_path}\" returned '\n f\"{workload!r}; falling back to {DEFAULT_WORKLOAD}\"\n )\n return DEFAULT_WORKLOAD\n return workload\n\n @classmethod\n def from_json_msg(cls, json_msg: Dict[str, Any]) -> \"GenericApiPayload\":\n if not isinstance(json_msg, dict):\n raise JsonDataException(\n {\"data\": \"payload must be a dictionary\"}\n )\n\n # Apply user's request parser if provided\n if user_request_parser:\n try:\n json_msg = user_request_parser(json_msg)\n except JsonDataException:\n raise # a bad request (422), not a server error\n except Exception as e:\n raise Exception(f\"Error in user response handler: {e}\")\n\n try:\n return cls.from_dict(json_msg)\n except (json.JSONDecodeError, JsonDataException) as e:\n raise JsonDataException(f\"Error in user response handler: {e}\")\n\n PayloadClass = GenericApiPayload\n\n # Create a generic EndpointHandler class\n @dataclass\n class GenericEndpointHandler(EndpointHandler[PayloadClass]):\n _route: str = field(default=route_path)\n has_benchmark: bool = field(\n default=(True if handler_config.benchmark_config else False)\n )\n allow_parallel_requests: bool = field(\n default=handler_config.allow_parallel_requests\n )\n max_queue_time: float = field(default=handler_config.max_queue_time)\n benchmark_runs: int = field(\n default=(\n handler_config.benchmark_config.runs\n if handler_config.benchmark_config\n else 8\n )\n )\n concurrency: int = field(\n default=(\n handler_config.benchmark_config.concurrency\n if handler_config.benchmark_config\n and handler_config.benchmark_config.concurrency\n else 10\n )\n )\n do_warmup_benchmark: bool = field(\n default=(\n handler_config.benchmark_config.do_warmup\n if handler_config.benchmark_config\n and handler_config.benchmark_config.concurrency\n else True\n )\n )\n remote_function: Callable = field(\n default=(\n handler_config.remote_function\n if handler_config.remote_function is not None\n else None\n )\n )\n\n @property\n def endpoint(self) -> str:\n \"\"\"The endpoint is the same as the route\"\"\"\n return self._route\n\n async def call_remote_dispatch_function(self, params: dict):\n \"\"\"\n define a remote function for this endpoint, return the result\n \"\"\"\n if self.remote_function is None:\n raise RuntimeError(\n f\"remote_function is not configured for route {self._route}\"\n )\n\n try:\n return await self.remote_function(**params)\n except Exception as ex:\n raise RuntimeError(\n f\"Error calling remote function for route {self._route}: {ex}\"\n ) from ex\n\n @property\n def healthcheck_endpoint(self) -> Optional[str]:\n return None\n\n @classmethod\n def payload_cls(cls) -> Type[PayloadClass]:\n return PayloadClass\n\n def make_benchmark_payload(self) -> PayloadClass:\n \"\"\"Just call the payload class's for_test() method\"\"\"\n return PayloadClass.for_test()\n\n async def generate_client_response(\n self,\n client_request: web.Request,\n model_response: ClientResponse,\n ) -> Union[web.Response, web.StreamResponse]:\n # User implemented override\n if user_response_generator:\n try:\n return await user_response_generator(\n client_request, model_response\n )\n except Exception as e:\n raise Exception(f\"Error in user response generator: {e}\")\n\n # Reasonable default: return response, handle streaming if present\n # Detect streaming\n content_type = model_response.content_type or \"\"\n is_stream = (\n content_type.startswith(\"text/event-stream\")\n or content_type in (\"application/x-ndjson\", \"application/jsonl\")\n or \"stream\" in content_type.lower()\n ) or (model_response.headers.get(\"Transfer-Encoding\") == \"chunked\")\n\n if is_stream:\n # Streaming passthrough\n res = web.StreamResponse(\n status=model_response.status,\n )\n if content_type:\n res.content_type = content_type\n\n await res.prepare(client_request)\n\n async for chunk in model_response.content.iter_any():\n if not chunk:\n continue\n await res.write(chunk)\n\n await res.write_eof()\n return res\n else:\n # Non-streaming\n body = await model_response.read()\n\n headers = model_response.headers.copy()\n headers.pop(\"Content-Type\", None)\n\n return web.Response(\n body=body,\n status=model_response.status,\n content_type=content_type or None,\n headers=headers,\n )\n\n return GenericEndpointHandler()\n\n def get_handler(self, route: str) -> Optional[EndpointHandler]:\n \"\"\"Get handler for a specific route\"\"\"\n return self._handlers.get(route)\n\n def get_all_handlers(self) -> Dict[str, EndpointHandler]:\n \"\"\"Get all registered handlers\"\"\"\n return self._handlers.copy()\n\n def get_benchmark_handler(self) -> Optional[EndpointHandler]:\n \"\"\"\n Get the benchmark handler.\n Raises errors if there are too many / too little BenchmarkConfigs\n Should be exactly one EndpointHandler with a BenchmarkConfig\n \"\"\"\n if not self._handlers:\n return None\n\n benchmark_handler: EndpointHandler = None\n for handler in self._handlers.values():\n if handler.has_benchmark:\n if benchmark_handler is not None:\n raise Exception(\n \"Cannot define BenchmarkConfig for more than one EndpointHandler!\"\n )\n else:\n benchmark_handler = handler\n if not benchmark_handler:\n raise Exception(\"Missing EndpointHandler with BenchmarkConfig\")\n return benchmark_handler\n\n def has_handlers(self) -> bool:\n \"\"\"Check if any handlers are registered\"\"\"\n return len(self._handlers) > 0\n\n @property\n def model_server_base_url(self) -> str:\n \"\"\"Get the full model server base URL\"\"\"\n return f\"{self.config.model_server_url}:{self.config.model_server_port}\"\n\n\nclass Worker:\n \"\"\"\n This class provides a simple to use abstraction over the pyworker backend.\n All custom implementations of pyworker can be created by configuring a Worker object.\n The pyworker starts by calling Worker.run()\n \"\"\"\n\n def __init__(self, config: WorkerConfig):\n # Configure logging for the pyworker internals.\n # Force-override any prior basicConfig (e.g. from vastai CLI code) by\n # setting the root logger level and adding a handler if none exist.\n root_logger = logging.getLogger()\n root_logger.setLevel(logging.DEBUG)\n if not root_logger.handlers:\n handler = logging.StreamHandler()\n handler.setFormatter(\n logging.Formatter(\n fmt=\"%(asctime)s[%(levelname)-5s] %(message)s\",\n datefmt=\"%Y-%m-%d %H:%M:%S\",\n )\n )\n root_logger.addHandler(handler)\n else:\n for handler in root_logger.handlers:\n handler.setLevel(logging.DEBUG)\n\n handler_factory = EndpointHandlerFactory(config)\n\n # Get all endpoint handlers\n handlers = handler_factory.get_all_handlers()\n benchmark_handler = handler_factory.get_benchmark_handler()\n\n # Create backend\n self.backend = backend.Backend(\n model_server_url=f\"{config.model_server_url}:{config.model_server_port}\",\n model_log_file=config.model_log_file,\n benchmark_handler=benchmark_handler,\n log_actions=config.log_action_config.log_actions,\n healthcheck_url=config.model_healthcheck_url,\n max_sessions=config.max_sessions,\n lifecycle=config.lifecycle,\n )\n\n # Attach endpoint handlers to HTTP routes\n self.routes = []\n for route_path, handler in handlers.items():\n self.routes.append(\n web.post(route_path, self.backend.create_handler(handler))\n )\n\n async def run_async(self, **kwargs):\n await server.start_server_async(self.backend, self.routes, **kwargs)\n\n def run(self, **kwargs):\n server.start_server(self.backend, self.routes, **kwargs)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":17514,"content_sha256":"1c7ef504497010485fa4c363a50e236cd530646aaf5739c07198641d236b7513"},{"filename":"sync/__init__.py","content":"from .client import SyncClient\nfrom .results import SyncOffer, SyncInstance\n","content_type":"text/x-python; charset=utf-8","language":"python","size":76,"content_sha256":"4849b7a633eb07a304a49167a3228a42187dc5975c3865d5c82611ccc4dc5afe"},{"filename":"sync/client.py","content":"import json\nimport requests\nfrom typing import Any, Optional, Union\n\nfrom vastai._base import _BaseClient, _APIKEY_SENTINEL\nfrom vastai.data.query import Query\nfrom vastai.data.offer import Offer\nfrom vastai.data.instance import Instance, InstanceConfig, CreateInstanceResponse\nfrom vastai.sync.results import SyncOffer, SyncInstance\n\n\nclass SyncClient(_BaseClient):\n \"\"\"\n Synchronous Vast.ai API client.\n\n api_key resolution order:\n 1. Explicit api_key argument\n 2. VAST_API_KEY environment variable\n 3. $XDG_CONFIG_HOME/vastai/vast_api_key\n 4. ~/.vast_api_key (legacy)\n \"\"\"\n def __init__(self, api_key: object = _APIKEY_SENTINEL, vast_server: str = \"https://console.vast.ai\"):\n super().__init__(api_key, vast_server)\n\n def search(\n self,\n query: Query,\n type_: Optional[str] = None,\n order: Optional[list] = [[\"score\", \"desc\"]],\n limit: Optional[int] = None,\n allocated_storage: Optional[float] = None,\n target_reliability: Optional[float] = None,\n external: Optional[bool] = None,\n extra_ids: Optional[list] = None,\n has_avx: Optional[int] = None,\n gpu_option: Optional[Any] = None,\n show_incompatible: Optional[Any] = None,\n sort_option: Optional[Any] = None,\n template_id: Optional[Any] = None,\n ) -> list[SyncOffer]:\n \"\"\"Search available offers and return a list of SyncOffer result wrappers.\"\"\"\n payload = dict(query.query)\n for key, value in [\n (\"type\", type_),\n (\"order\", order),\n (\"limit\", limit),\n (\"allocated_storage\", allocated_storage),\n (\"target_reliability\", target_reliability),\n (\"external\", external),\n (\"extra_ids\", extra_ids),\n (\"has_avx\", has_avx),\n (\"gpu_option\", gpu_option),\n (\"show_incompatible\", show_incompatible),\n (\"sort_option\", sort_option),\n (\"template_id\", template_id),\n ]:\n if value is not None:\n payload[key] = value\n\n response = requests.post(\n self._url(\"/api/v0/bundles/\"),\n json=payload,\n headers=self._headers(),\n )\n response.raise_for_status()\n data = response.json()\n\n return [SyncOffer(Offer.from_dict(o), self) for o in data.get(\"offers\", [])]\n\n def create_instance(\n self,\n offer_id: int,\n config: InstanceConfig,\n ) -> SyncInstance:\n \"\"\"Accept an offer and create an instance. Returns a SyncInstance wrapper.\"\"\"\n payload = config.to_dict()\n response = requests.put(\n self._url(f\"/api/v0/asks/{offer_id}/\"),\n json=payload,\n headers=self._headers(),\n )\n response.raise_for_status()\n data = response.json()\n\n if not data.get(\"success\"):\n raise Exception(f\"Failed to create instance: {data.get('error', '')} - {data.get('msg', '')}\")\n\n resp = CreateInstanceResponse.from_dict(data)\n # Fetch the newly created instance\n instances = self.show_instances()\n for inst in instances:\n if inst.id == resp.new_contract:\n return inst\n # Fallback: return a minimal instance\n return SyncInstance(Instance.from_dict({\"id\": resp.new_contract}), self)\n\n def show_instances(self) -> list[SyncInstance]:\n \"\"\"Return all instances owned by the authenticated user.\n\n The legacy v0 ``/instances`` list endpoint is deprecated; this pages\n through the v1 ``/api/v1/instances/`` endpoint (omitting ``select_cols``\n so the backend returns full instance rows) and accumulates every page.\n \"\"\"\n instances = []\n params = {\n \"select_filters\": json.dumps({}),\n \"order_by\": json.dumps([{\"col\": \"id\", \"dir\": \"asc\"}]),\n \"limit\": 25,\n }\n while True:\n response = requests.get(\n self._url(\"/api/v1/instances/\"),\n params=params,\n headers=self._headers(),\n )\n response.raise_for_status()\n data = response.json()\n instances.extend(data.get(\"instances\") or [])\n next_token = data.get(\"next_token\")\n if not next_token:\n break\n params[\"after_token\"] = next_token\n\n return [SyncInstance(Instance.from_dict(i), self) for i in instances]\n\n def destroy_instance(self, instance_or_id: Union[SyncInstance, int]) -> None:\n \"\"\"Destroy a running or stopped instance.\"\"\"\n instance_id = instance_or_id if isinstance(instance_or_id, int) else instance_or_id.id\n response = requests.delete(\n self._url(f\"/api/v0/instances/{instance_id}/\"),\n headers=self._headers(),\n )\n response.raise_for_status()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":4983,"content_sha256":"b54e9c9f93a3d0a8be44f1e3af1aa8e9c26ca0ebc9a15a7d05b11c1db72a93b6"},{"filename":"sync/results.py","content":"from __future__ import annotations\nfrom typing import TYPE_CHECKING, Optional, Union\n\nfrom vastai.data.offer import Offer\nfrom vastai.data.instance import InstanceConfig\n\nif TYPE_CHECKING:\n from vastai.sync.client import SyncClient\n from vastai.data.instance import Instance\n\n\nclass SyncOffer:\n def __init__(self, data: Offer, client: SyncClient):\n self._data = data\n self._client = client\n\n def create_instance(self, config: InstanceConfig) -> SyncInstance:\n \"\"\"Create an instance from this offer. Synonym for SyncClient.create_instance(self.id, config).\"\"\"\n return self._client.create_instance(self._data.id, config)\n\n def __getattr__(self, name: str):\n return getattr(self._data, name)\n\n def __repr__(self) -> str:\n return f\"SyncOffer({self._data!r})\"\n\n\nclass SyncInstance:\n def __init__(self, data: Instance, client: SyncClient):\n self._data = data\n self._client = client\n\n def destroy(self) -> None:\n \"\"\"Synonym for SyncClient.destroy_instance(self).\"\"\"\n return self._client.destroy_instance(self._data.id)\n\n def __getattr__(self, name: str):\n return getattr(self._data, name)\n\n def __repr__(self) -> str:\n return f\"SyncInstance({self._data!r})\"\n","content_type":"text/x-python; charset=utf-8","language":"python","size":1265,"content_sha256":"329428b7e38b59d0ef4a734b55d8f899515ae08059153082e03b2acec75e51d8"},{"filename":"utils.py","content":"\"\"\"Shared utility functions for the Vast.ai SDK and CLI.\n\nGeneric helpers that are not specific to any single layer (CLI, API, SDK).\n\"\"\"\n\nimport re\nimport os\nimport sys\nimport time\nimport math\nimport argparse\nimport subprocess\nimport importlib.metadata\nfrom datetime import datetime, timedelta, timezone\nfrom concurrent.futures import ThreadPoolExecutor\n\n\n# ---------------------------------------------------------------------------\n# Version helpers\n# ---------------------------------------------------------------------------\n\ndef parse_version(version: str) -> tuple:\n parts = version.split(\".\")\n\n if len(parts) \u003c 3:\n print(f\"Invalid version format: {version}\", file=sys.stderr)\n\n return tuple(int(part) for part in parts)\n\n\ndef _get_git_version():\n try:\n result = subprocess.run(\n [\"git\", \"describe\", \"--tags\", \"--abbrev=0\"],\n capture_output=True,\n text=True,\n check=True,\n cwd=os.path.dirname(__file__),\n )\n tag = result.stdout.strip()\n return tag[1:] if tag.startswith(\"v\") else tag\n except Exception:\n return \"0.0.0\"\n\n\ndef _get_local_version():\n try:\n return importlib.metadata.version(\"vastai\")\n except Exception:\n return _get_git_version()\n\n\nVERSION = _get_local_version()\n\n\n# ---------------------------------------------------------------------------\n# Validation helpers\n# ---------------------------------------------------------------------------\n\ndef validate_seconds(value):\n \"\"\"Validate that the input value is a valid number for seconds between yesterday and Jan 1, 2100.\"\"\"\n try:\n val = int(value)\n\n # Calculate min_seconds as the start of yesterday in seconds\n yesterday = datetime.now() - timedelta(days=1)\n min_seconds = int(yesterday.timestamp())\n\n # Calculate max_seconds for Jan 1st, 2100 in seconds\n max_date = datetime(2100, 1, 1, 0, 0, 0)\n max_seconds = int(max_date.timestamp())\n\n if not (min_seconds \u003c= val \u003c= max_seconds):\n raise argparse.ArgumentTypeError(f\"{value} is not a valid second timestamp.\")\n return val\n except ValueError:\n raise argparse.ArgumentTypeError(f\"{value} is not a valid integer.\")\n\n\n# ---------------------------------------------------------------------------\n# VRL (Vast Resource Locator) parsing\n# ---------------------------------------------------------------------------\n\nclass VRLException(Exception):\n pass\n\ndef parse_vast_url(url_str):\n \"\"\"\n Breaks up a vast-style url in the form instance_id:path and does\n some basic sanity type-checking.\n\n :param url_str:\n :return:\n \"\"\"\n\n instance_id = None\n path = url_str\n #print(f'url_str: {url_str}')\n if (\":\" in url_str):\n url_parts = url_str.split(\":\", 2)\n if len(url_parts) == 2:\n (instance_id, path) = url_parts\n else:\n raise VRLException(\"Invalid VRL (Vast resource locator).\")\n else:\n try:\n instance_id = int(path)\n path = \"/\"\n except:\n pass\n\n valid_unix_path_regex = re.compile('^(/)?([^/\\0]+(/)?)+

Important: agents should read /llm.txt, /llms.txt, or /.well-known/skills.json to discover the public Skillopedia API.