Vast.ai: Support spot instances

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

@@ -41,6 +41,7 @@ def create_instance(
         onstart: str,
         disk_size: int,
         registry_auth: Optional[RegistryAuth] = None,
+        bid: Optional[float] = None,
     ) -> dict:
         """
         Args:
@@ -49,6 +50,8 @@ def create_instance(
             image_name: docker image name
             onstart: commands to run on start
             registry_auth: registry auth credentials for private images
+            bid: per-machine bid price in $/hour. If set, an interruptible
+                (spot) instance is created; if None, an on-demand instance.
 
         Raises:
             NoCapacityError: if instance cannot be created
@@ -63,6 +66,7 @@ def create_instance(
         payload = {
             "client_id": "me",
             "image": image_name,
+            "price": bid,
             "disk": disk_size,
             "label": instance_name,
             "env": {

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

@@ -94,8 +94,6 @@ def get_offers_by_requirements(
             backend=BackendType.VASTAI,
             locations=self.config.regions or None,
             requirements=requirements,
-            # TODO(egor-s): spots currently not supported
-            extra_filter=lambda offer: not offer.instance.resources.spot,
             catalog=self._make_catalog(vastai_options),
         )
         offers = [
@@ -127,13 +125,17 @@ def run_job(
         commands = get_docker_commands(
             [run.run_spec.ssh_key_pub.strip(), project_ssh_public_key.strip()]
         )
+        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"]
         return JobProvisioningData(

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

@@ -1,7 +1,16 @@
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 from dstack._internal.core.backends.vastai.compute import VastAICompute
 from dstack._internal.core.backends.vastai.models import VastAIConfig, VastAICreds
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.instances import (
+    Disk,
+    Gpu,
+    InstanceAvailability,
+    InstanceOfferWithAvailability,
+    InstanceType,
+    Resources,
+)
 from dstack._internal.core.models.resources import ResourcesSpec
 from dstack._internal.core.models.runs import Requirements
 
@@ -14,6 +23,53 @@ def _requirements() -> Requirements:
     return Requirements(resources=ResourcesSpec())
 
 
+def _offer(*, spot: bool, price: float = 0.5) -> InstanceOfferWithAvailability:
+    return InstanceOfferWithAvailability(
+        backend=BackendType.VASTAI,
+        instance=InstanceType(
+            name="12345",
+            resources=Resources(
+                cpus=8,
+                memory_mib=32 * 1024,
+                gpus=[Gpu(name="RTX4090", memory_mib=24 * 1024)],
+                spot=spot,
+                disk=Disk(size_mib=100 * 1024),
+            ),
+        ),
+        region="Hong Kong, HK",
+        price=price,
+        availability=InstanceAvailability.AVAILABLE,
+    )
+
+
+def _run_job(compute: VastAICompute, offer: InstanceOfferWithAvailability):
+    run = MagicMock()
+    run.run_spec.ssh_key_pub = "ssh-rsa AAAA test"
+    job = MagicMock()
+    job.job_spec.image_name = "dstackai/base:latest"
+    job.job_spec.registry_auth = None
+    with (
+        patch(
+            "dstack._internal.core.backends.vastai.compute.generate_unique_instance_name_for_job",
+            return_value="dstack-test",
+        ),
+        patch(
+            "dstack._internal.core.backends.vastai.compute.get_docker_commands",
+            return_value=["echo hi"],
+        ),
+    ):
+        compute.run_job(
+            run=run,
+            job=job,
+            instance_offer=offer,
+            project_ssh_public_key="ssh-rsa BBBB project",
+            project_ssh_private_key="private-key",
+            volumes=[],
+            placement_group=None,
+            requirements=_requirements(),
+        )
+
+
 def test_vastai_compute_enables_community_cloud_by_default():
     with (
         patch("dstack._internal.core.backends.vastai.compute.VastAIProvider") as vast_provider_cls,
@@ -54,3 +110,23 @@ def test_vastai_compute_can_disable_community_cloud():
         vast_provider_cls.assert_called_once()
         assert vast_provider_cls.call_args.kwargs["community_cloud"] is False
         catalog_instance.add_provider.assert_called_once()
+
+
+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}
+
+    _run_job(compute, _offer(spot=True, price=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}
+
+    _run_job(compute, _offer(spot=False, price=0.24))
+
+    assert compute.api_client.create_instance.call_args.kwargs["bid"] is None