Migrate off deprecated Vast.ai /api/v0/instances

src/dstack/_internal/core/backends/vastai/api_client.py

@@ -1,37 +1,22 @@
-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
 from dstack._internal.core.errors import ComputeError, NoCapacityError
 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 = []

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()