diff --git a/src/dstack/_internal/core/backends/vastai/api_client.py b/src/dstack/_internal/core/backends/vastai/api_client.py index 712ff7211..6b8a041d7 100644 --- a/src/dstack/_internal/core/backends/vastai/api_client.py +++ b/src/dstack/_internal/core/backends/vastai/api_client.py @@ -1,9 +1,6 @@ -import threading -import time -from typing import List, Optional, Union +from typing import Optional, Union import requests -from requests.adapters import HTTPAdapter, Retry import dstack._internal.utils.docker as docker from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT @@ -11,27 +8,15 @@ from dstack._internal.core.models.common import RegistryAuth +class VastAIRateLimitError(Exception): + pass + + class VastAIAPIClient: def __init__(self, api_key: str): - self.api_url = "https://console.vast.ai/api/v0".rstrip("/") + self.api_url = "https://console.vast.ai/api" self.api_key = api_key self.s = requests.Session() # TODO: set adequate timeout everywhere the session is used - retries = Retry( - total=5, - backoff_factor=1, - status_forcelist=[429, 504], - ) - self.s.mount(prefix=self._url("/instances/"), adapter=HTTPAdapter(max_retries=retries)) - self.lock = threading.Lock() - self.instances_cache_ts: float = 0 - self.instances_cache: List[dict] = [] - - def get_bundle(self, bundle_id: Union[str, int]) -> Optional[dict]: - resp = self.s.post(self._url("/bundles/"), json={"id": {"eq": bundle_id}}) - resp.raise_for_status() - data = resp.json() - offers = data["offers"] - return offers[0] if offers else None def create_instance( self, @@ -80,14 +65,13 @@ def create_instance( "create_from": None, "force": False, } - resp = self.s.put(self._url(f"/asks/{bundle_id}/"), json=payload) + resp = self.s.put(self._url(f"/v0/asks/{bundle_id}/"), json=payload) if resp.status_code != 200 or not (data := resp.json())["success"]: raise NoCapacityError(resp.text) - self._invalidate_cache() return data def destroy_instance(self, instance_id: Union[str, int]) -> None: - resp = self.s.delete(self._url(f"/instances/{instance_id}/")) + resp = self.s.delete(self._url(f"/v0/instances/{instance_id}/")) try: data = resp.json() except requests.exceptions.JSONDecodeError: @@ -96,44 +80,20 @@ def destroy_instance(self, instance_id: Union[str, int]) -> None: if data.get("error") != "no_such_instance": raise ComputeError(resp.text) - def get_instances(self, cache_ttl: float = 3.0) -> List[dict]: - with self.lock: - if time.time() - self.instances_cache_ts > cache_ttl: - resp = self.s.get(self._url("/instances/")) - resp.raise_for_status() - data = resp.json() - self.instances_cache_ts = time.time() - self.instances_cache = data["instances"] - return self.instances_cache - def get_instance(self, instance_id: Union[str, int]) -> Optional[dict]: - instances = self.get_instances() - for instance in instances: - if instance["id"] == int(instance_id): - return instance - return None - - def request_logs(self, instance_id: Union[str, int]) -> dict: - resp = self.s.put( - self._url(f"/instances/request_logs/{instance_id}/"), json={"tail": "1000"} - ) + resp = self.s.get(self._url(f"/v0/instances/{instance_id}/")) + if resp.status_code == 429: + raise VastAIRateLimitError() resp.raise_for_status() data = resp.json() - if not data["success"]: - raise requests.HTTPError(data) - return data + return data["instances"] def auth_test(self) -> bool: try: - self.get_instances() + self.s.get(self._url("/v1/instances/")).raise_for_status() return True except requests.HTTPError: return False def _url(self, path): return f"{self.api_url}/{path.lstrip('/')}?api_key={self.api_key}" - - def _invalidate_cache(self): - with self.lock: - self.instances_cache_ts = 0 - self.instances_cache = [] diff --git a/src/dstack/_internal/core/backends/vastai/compute.py b/src/dstack/_internal/core/backends/vastai/compute.py index fc1a95707..17268b6cd 100644 --- a/src/dstack/_internal/core/backends/vastai/compute.py +++ b/src/dstack/_internal/core/backends/vastai/compute.py @@ -12,7 +12,7 @@ ) from dstack._internal.core.backends.base.offers import get_catalog_offers from dstack._internal.core.backends.base.profile_options import get_backend_profile_options -from dstack._internal.core.backends.vastai.api_client import VastAIAPIClient +from dstack._internal.core.backends.vastai.api_client import VastAIAPIClient, VastAIRateLimitError from dstack._internal.core.backends.vastai.models import VastAIConfig from dstack._internal.core.backends.vastai.profile_options import ( VASTAI_DEFAULT_MIN_RELIABILITY, @@ -162,7 +162,14 @@ def update_provisioning_data( project_ssh_public_key: str, project_ssh_private_key: str, ): - resp = self.api_client.get_instance(provisioning_data.instance_id) + try: + resp = self.api_client.get_instance(provisioning_data.instance_id) + except VastAIRateLimitError: + logger.warning( + "Reached Vast.ai rate limit when updating instance %s provisioning data", + provisioning_data.instance_id, + ) + return if resp is not None: if resp["actual_status"] == "running": provisioning_data.hostname = resp["public_ipaddr"].strip()