Update Vast.ai spot instances support

pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
     "python-multipart>=0.0.16",
     "filelock",
     "psutil",
-    "gpuhunt==0.1.26",
+    "gpuhunt==0.1.27",
     "argcomplete>=3.5.0",
     "ignore-python>=0.2.0",
     "orjson",

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

@@ -4,14 +4,31 @@
 
 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.errors import ComputeError
 from dstack._internal.core.models.common import RegistryAuth
 
 
 class VastAIRateLimitError(Exception):
     pass
 
 
+class VastAICreateInstanceError(Exception):
+    """
+    Attributes:
+        resp: Full API response.
+        created_instance_id: ID of the instance that was created despite the error. For example,
+            Vast.ai creates spot instances in the stopped state and returns an error if the bid is
+            insufficient.
+    """
+
+    def __init__(
+        self, *args, resp: str, created_instance_id: Optional[int] = None, **kwargs
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.resp = resp
+        self.created_instance_id = created_instance_id
+
+
 class VastAIAPIClient:
     def __init__(self, api_key: str):
         self.api_url = "https://console.vast.ai/api"
@@ -27,7 +44,7 @@ def create_instance(
         disk_size: int,
         registry_auth: Optional[RegistryAuth] = None,
         bid: Optional[float] = None,
-    ) -> dict:
+    ) -> int:
         """
         Args:
             instance_name: instance label
@@ -39,10 +56,10 @@ def create_instance(
                 (spot) instance is created; if None, an on-demand instance.
 
         Raises:
-            NoCapacityError: if instance cannot be created
+            VastAICreateInstanceError: if instance cannot be created
 
         Returns:
-            create instance response
+            Instance ID
         """
         image_login = None
         if registry_auth:
@@ -70,12 +87,20 @@ def create_instance(
             "force": False,
         }
         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)
-        return data
+        try:
+            data = resp.json()
+        except requests.exceptions.JSONDecodeError:
+            raise VastAICreateInstanceError(resp=resp.text)
+        if resp.status_code != 200 or not data["success"]:
+            raise VastAICreateInstanceError(
+                resp=resp.text, created_instance_id=data.get("new_contract")
+            )
+        return data["new_contract"]
 
     def destroy_instance(self, instance_id: Union[str, int]) -> None:
         resp = self.s.delete(self._url(f"/v0/instances/{instance_id}/"))
+        if resp.status_code == 429:
+            raise VastAIRateLimitError()
         try:
             data = resp.json()
         except requests.exceptions.JSONDecodeError:

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

@@ -1,3 +1,4 @@
+import time
 from typing import List, Optional
 
 import gpuhunt
@@ -12,7 +13,11 @@
 )
 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, VastAIRateLimitError
+from dstack._internal.core.backends.vastai.api_client import (
+    VastAIAPIClient,
+    VastAICreateInstanceError,
+    VastAIRateLimitError,
+)
 from dstack._internal.core.backends.vastai.models import VastAIConfig
 from dstack._internal.core.backends.vastai.profile_options import (
     VASTAI_DEFAULT_MIN_RELIABILITY,
@@ -21,8 +26,9 @@
     VastAIProfileOptions,
 )
 from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT
-from dstack._internal.core.errors import ProvisioningError
+from dstack._internal.core.errors import ComputeError, NoCapacityError, ProvisioningError
 from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.common import CoreModel
 from dstack._internal.core.models.instances import (
     InstanceAvailability,
     InstanceOfferWithAvailability,
@@ -125,23 +131,39 @@ def run_job(
         commands = get_docker_commands(
             [run.run_spec.ssh_key_pub.strip(), project_ssh_public_key.strip()]
         )
+        offer_backend_data: VastAIOfferBackendData = VastAIOfferBackendData.__response__.parse_obj(
+            instance_offer.backend_data
+        )
         bid = None
         if instance_offer.instance.resources.spot:
-            bid = instance_offer.price
-        resp = self.api_client.create_instance(
-            instance_name=instance_name,
-            bundle_id=instance_offer.instance.name,
-            image_name=job.job_spec.image_name,
-            onstart=" && ".join(commands),
-            disk_size=round(instance_offer.instance.resources.disk.size_mib / 1024),
-            registry_auth=job.job_spec.registry_auth,
-            bid=bid,
-        )
-        instance_id = resp["new_contract"]
+            if offer_backend_data.min_bid is None:
+                raise ComputeError(
+                    "VastAIOfferBackendData.min_bid is unexpectedly missing for a spot offer"
+                )
+            bid = offer_backend_data.min_bid
+        try:
+            instance_id = self.api_client.create_instance(
+                instance_name=instance_name,
+                bundle_id=instance_offer.instance.name,
+                image_name=job.job_spec.image_name,
+                onstart=" && ".join(commands),
+                disk_size=round(instance_offer.instance.resources.disk.size_mib / 1024),
+                registry_auth=job.job_spec.registry_auth,
+                bid=bid,
+            )
+        except VastAICreateInstanceError as e:
+            if e.created_instance_id is not None:
+                _terminate_instance_with_rate_limit_retry(
+                    self.api_client, e.created_instance_id, retries=4
+                )
+                logger.debug(
+                    "Terminated Vast.ai instance %s that failed to start", e.created_instance_id
+                )
+            raise NoCapacityError(e.resp) from e
         return JobProvisioningData(
             backend=instance_offer.backend,
             instance_type=instance_offer.instance,
-            instance_id=instance_id,
+            instance_id=str(instance_id),
             hostname=None,
             internal_ip=None,
             region=instance_offer.region,
@@ -183,3 +205,38 @@ def update_provisioning_data(
                 and ": OCI runtime create failed:" in resp["status_msg"]
             ):
                 raise ProvisioningError(resp["status_msg"])
+
+
+class VastAIOfferBackendData(CoreModel):
+    min_bid: float | None = None
+
+
+def _terminate_instance_with_rate_limit_retry(
+    client: VastAIAPIClient, instance_id: int, retries: int
+) -> bool:
+    for attempt in range(retries):
+        try:
+            client.destroy_instance(instance_id)
+            return True
+        except VastAIRateLimitError:
+            if attempt + 1 < retries:
+                logger.warning(
+                    "Hit rate limit when terminating Vast.ai instance %s. Attempt %s/%s",
+                    instance_id,
+                    attempt + 1,
+                    retries,
+                )
+                time.sleep(attempt + 1)
+        except BaseException as e:
+            logger.error(
+                "Failed to terminate Vast.ai instance %s. Terminate it manually to avoid accumulating charges. Error: %r",
+                instance_id,
+                e,
+            )
+            raise
+    logger.error(
+        "Failed to terminate Vast.ai instance %s after %s attempts hit rate limits. Terminate it manually to avoid accumulating charges",
+        instance_id,
+        retries,
+    )
+    return False

src/tests/_internal/core/backends/vastai/test_compute.py

@@ -23,7 +23,9 @@ def _requirements() -> Requirements:
     return Requirements(resources=ResourcesSpec())
 
 
-def _offer(*, spot: bool, price: float = 0.5) -> InstanceOfferWithAvailability:
+def _offer(
+    *, spot: bool, price: float = 0.5, min_bid: float | None = None
+) -> InstanceOfferWithAvailability:
     return InstanceOfferWithAvailability(
         backend=BackendType.VASTAI,
         instance=InstanceType(
@@ -39,6 +41,9 @@ def _offer(*, spot: bool, price: float = 0.5) -> InstanceOfferWithAvailability:
         region="Hong Kong, HK",
         price=price,
         availability=InstanceAvailability.AVAILABLE,
+        backend_data={
+            **({"min_bid": min_bid} if min_bid is not None else {}),
+        },
     )
 
 
@@ -115,17 +120,17 @@ def test_vastai_compute_can_disable_community_cloud():
 def test_vastai_run_job_bids_on_spot_offer():
     compute = VastAICompute(_config())
     compute.api_client = MagicMock()
-    compute.api_client.create_instance.return_value = {"new_contract": 123}
+    compute.api_client.create_instance.return_value = 123
 
-    _run_job(compute, _offer(spot=True, price=0.1244444))
+    _run_job(compute, _offer(spot=True, price=0.14, min_bid=0.1244444))
 
     assert compute.api_client.create_instance.call_args.kwargs["bid"] == 0.1244444
 
 
 def test_vastai_run_job_does_not_bid_on_ondemand_offer():
     compute = VastAICompute(_config())
     compute.api_client = MagicMock()
-    compute.api_client.create_instance.return_value = {"new_contract": 123}
+    compute.api_client.create_instance.return_value = 123
 
     _run_job(compute, _offer(spot=False, price=0.24))