Fix review issues in Vultr integration

Update supported_instance with regex match

Author: Bihan Rana

setup.py

@@ -48,7 +48,7 @@ def get_long_description():
     "python-multipart>=0.0.16",
     "filelock",
     "psutil",
-    "gpuhunt>=0.0.17,<0.1.0",
+    "gpuhunt>=0.0.18,<0.1.0",
 ]
 
 GATEWAY_AND_SERVER_COMMON_DEPS = [

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

@@ -1,10 +1,10 @@
 import base64
-from typing import Any, List
+from typing import Any
 
 import requests
 from requests import Response
 
-from dstack._internal.core.errors import BackendInvalidCredentialsError
+from dstack._internal.core.errors import BackendError, BackendInvalidCredentialsError
 
 API_URL = "https://api.vultr.com/v2"
 
@@ -28,13 +28,7 @@ def get_instance(self, instance_id: str, plan_type: str):
             response = self._make_request("GET", f"/instances/{instance_id}")
             return response.json()["instance"]
 
-    def launch_instance(
-        self, region: str, plan: str, label: str, startup_script: str, public_keys: List[str]
-    ):
-        # Fetch or create startup script ID
-        script_id: str = self.get_startup_script_id(startup_script)
-        # Fetch or create SSH key IDs
-        sshkey_ids: List[str] = self.get_sshkey_id(public_keys)
+    def launch_instance(self, region: str, plan: str, label: str, user_data: str):
         # For Bare-metals
         if "vbm" in plan:
             # "Docker on Ubuntu 22.04" is required for bare-metals.
@@ -43,8 +37,7 @@ def launch_instance(
                 "plan": plan,
                 "label": label,
                 "image_id": "docker",
-                "script_id": script_id,
-                "sshkey_id": sshkey_ids,
+                "user_data": base64.b64encode(user_data.encode()).decode(),
             }
             resp = self._make_request("POST", "/bare-metals", data)
             return resp.json()["bare_metal"]["id"]
@@ -56,8 +49,7 @@ def launch_instance(
                 "plan": plan,
                 "label": label,
                 "os_id": 1743,
-                "script_id": script_id,
-                "sshkey_id": sshkey_ids,
+                "user_data": base64.b64encode(user_data.encode()).decode(),
             }
             resp = self._make_request("POST", "/instances", data)
             return resp.json()["instance"]["id"]
@@ -67,64 +59,11 @@ def launch_instance(
                 "plan": plan,
                 "label": label,
                 "image_id": "docker",
-                "script_id": script_id,
-                "sshkey_id": sshkey_ids,
+                "user_data": base64.b64encode(user_data.encode()).decode(),
             }
             resp = self._make_request("POST", "/instances", data)
             return resp.json()["instance"]["id"]
 
-    def get_startup_script_id(self, startup_script: str) -> str:
-        script_name = "dstack-shim-script"
-        encoded_script = base64.b64encode(startup_script.encode()).decode()
-
-        # Get the list of startup scripts
-        response = self._make_request("GET", "/startup-scripts")
-        scripts = response.json()["startup_scripts"]
-
-        # Find the script by name
-        existing_script = next((s for s in scripts if s["name"] == script_name), None)
-
-        if existing_script:
-            # Update the existing script
-            startup_id = existing_script["id"]
-            update_payload = {
-                "name": script_name,
-                "script": encoded_script,
-            }
-            self._make_request("PATCH", f"/startup-scripts/{startup_id}", update_payload)
-        else:
-            # Create a new script
-            create_payload = {
-                "name": script_name,
-                "type": "boot",
-                "script": encoded_script,
-            }
-            create_response = self._make_request("POST", "/startup-scripts", create_payload)
-            startup_id = create_response.json()["startup_script"]["id"]
-
-        return startup_id
-
-    def get_sshkey_id(self, ssh_ids: List[str]) -> List[str]:
-        # Fetch existing SSH keys
-        response = self._make_request("GET", "/ssh-keys")
-        ssh_keys = response.json()["ssh_keys"]
-
-        ssh_key_ids = []
-        existing_keys = {key["ssh_key"]: key["id"] for key in ssh_keys}
-
-        for ssh_key in ssh_ids:
-            if ssh_key in existing_keys:
-                # SSH key already exists, add its id to the list
-                ssh_key_ids.append(existing_keys[ssh_key])
-            else:
-                # Create new SSH key
-                create_payload = {"name": "dstack-ssh-key", "ssh_key": ssh_key}
-                create_response = self._make_request("POST", "/ssh-keys", create_payload)
-                new_ssh_key_id = create_response.json()["ssh_key"]["id"]
-                ssh_key_ids.append(new_ssh_key_id)
-
-        return ssh_key_ids
-
     def terminate_instance(self, instance_id: str, plan_type: str):
         if plan_type == "bare-metal":
             # Terminate bare-metal instance
@@ -151,4 +90,10 @@ def _make_request(self, method: str, path: str, data: Any = None) -> Response:
                 requests.codes.unauthorized,
             ):
                 raise BackendInvalidCredentialsError(e.response.text)
+            if e.response is not None and e.response.status_code in (
+                requests.codes.bad_request,
+                requests.codes.internal_server_error,
+                requests.codes.not_found,
+            ):
+                raise BackendError(e.response.text)
             raise

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

@@ -1,12 +1,13 @@
 import json
+import re
 from typing import List, Optional
 
 import requests
 
 from dstack._internal.core.backends.base import Compute
 from dstack._internal.core.backends.base.compute import (
     get_instance_name,
-    get_shim_commands,
+    get_user_data,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
 from dstack._internal.core.backends.vultr.api_client import VultrApiClient
@@ -16,6 +17,7 @@
 from dstack._internal.core.models.instances import (
     InstanceAvailability,
     InstanceConfiguration,
+    InstanceOffer,
     InstanceOfferWithAvailability,
     SSHKey,
 )
@@ -37,6 +39,8 @@ def get_offers(
         offers = get_catalog_offers(
             backend=BackendType.VULTR,
             requirements=requirements,
+            locations=self.config.regions or None,
+            extra_filter=_supported_instances,
         )
         offers = [
             InstanceOfferWithAvailability(
@@ -66,19 +70,12 @@ def run_job(
     def create_instance(
         self, instance_offer: InstanceOfferWithAvailability, instance_config: InstanceConfiguration
     ) -> JobProvisioningData:
-        public_keys = instance_config.get_public_keys()
-        commands = get_shim_commands(authorized_keys=public_keys)
-        shim_commands = "#!/bin/sh\n" + " ".join([" && ".join(commands)])
-        try:
-            instance_id = self.api_client.launch_instance(
-                region=instance_offer.region,
-                label=instance_config.instance_name,
-                plan=instance_offer.instance.name,
-                startup_script=shim_commands,
-                public_keys=public_keys,
-            )
-        except KeyError as e:
-            raise BackendError(e)
+        instance_id = self.api_client.launch_instance(
+            region=instance_offer.region,
+            label=instance_config.instance_name,
+            plan=instance_offer.instance.name,
+            user_data=get_user_data(authorized_keys=instance_config.get_public_keys()),
+        )
 
         launched_instance = JobProvisioningData(
             backend=instance_offer.backend,
@@ -126,3 +123,29 @@ def update_provisioning_data(
             provisioning_data.hostname = instance_main_ip
         if instance_status == "failed":
             raise ProvisioningError("VM entered FAILED state")
+
+
+def _supported_instances(offer: InstanceOffer) -> bool:
+    if offer.instance.resources.spot:
+        return False
+    for family in [
+        # Bare Metal - GPU
+        r"vbm-\d+c-\d+gb-\d+-(a100|h100|l40|mi300x)-gpu",
+        # Bare Metal - AMD CPU
+        r"vbm-\d+c-\d+gb-amd",
+        # Bare Metal - Intel CPU
+        r"vbm-\d+c-\d+gb(-v\d+)?",
+        # Cloud GPU
+        r"vcg-(a16|a40|l40s|a100)-\d+c-\d+g-\d+vram",
+        # Cloud Compute - Regular Performance
+        r"vc2-\d+c-\d+gb(-sc1)?",
+        # Cloud Compute - High Frequency
+        r"vhf-\d+c-\d+gb(-sc1)?",
+        # Cloud Compute - High Performance
+        r"vhp-\d+c-\d+gb-(intel|amd)(-sc1)?",
+        # Optimized Cloud Compute
+        r"voc-[cgms]-\d+c-\d+gb-\d+s-amd(-sc1)?",
+    ]:
+        if re.fullmatch(family, offer.instance.name):
+            return True
+    return False

src/dstack/_internal/core/models/backends/vultr.py

@@ -3,13 +3,12 @@
 from pydantic.fields import Field
 from typing_extensions import Annotated, Literal
 
-from dstack._internal.core.models.backends.base import ConfigElement, ConfigMultiElement
+from dstack._internal.core.models.backends.base import ConfigMultiElement
 from dstack._internal.core.models.common import CoreModel
 
 
 class VultrConfigInfo(CoreModel):
     type: Literal["vultr"] = "vultr"
-    project_id: str
     regions: Optional[List[str]] = None
 
 
@@ -33,11 +32,9 @@ class VultrConfigInfoWithCreds(VultrConfigInfo):
 class VultrConfigInfoWithCredsPartial(CoreModel):
     type: Literal["vultr"] = "vultr"
     creds: Optional[AnyVultrCreds]
-    project_id: Optional[str]
     regions: Optional[List[str]]
 
 
 class VultrConfigValues(CoreModel):
     type: Literal["vultr"] = "vultr"
     regions: Optional[ConfigMultiElement]
-    project_id: Optional[ConfigElement]

src/dstack/_internal/server/background/tasks/process_instances.py

@@ -911,6 +911,6 @@ def _get_instance_timeout_interval(
         return timedelta(seconds=1200)
     if backend_type == BackendType.OCI and instance_type_name.startswith("BM."):
         return timedelta(seconds=1200)
-    if backend_type == BackendType.VULTR:
-        return timedelta(seconds=1800)
+    if backend_type == BackendType.VULTR and instance_type_name.startswith("vbm"):
+        return timedelta(seconds=3300)
     return timedelta(seconds=600)

src/dstack/_internal/server/background/tasks/process_running_jobs.py

@@ -656,6 +656,6 @@ def _get_runner_timeout_interval(backend_type: BackendType, instance_type_name:
         return timedelta(seconds=1200)
     if backend_type == BackendType.OCI and instance_type_name.startswith("BM."):
         return timedelta(seconds=1200)
-    if backend_type == BackendType.VULTR:
-        return timedelta(seconds=1800)
+    if backend_type == BackendType.VULTR and instance_type_name.startswith("vbm"):
+        return timedelta(seconds=3300)
     return timedelta(seconds=600)

src/dstack/_internal/server/services/backends/configurators/vultr.py

@@ -24,16 +24,7 @@
     raise_invalid_credentials_error,
 )
 
-REGIONS = [
-    "no-luster-1",
-    "se-smedjebacken-1",
-    "gb-london-1",
-    "se-stockholm-1",
-    "us-newyork-1",
-    "us-santaclara-1",
-]
-
-DEFAULT_REGION = "no-luster-1"
+REGIONS = []
 
 
 class VultrConfigurator(Configurator):
@@ -44,9 +35,7 @@ def get_config_values(self, config: VultrConfigInfoWithCredsPartial) -> VultrCon
         if config.creds is None:
             return config_values
         self._validate_vultr_api_key(config.creds.api_key)
-        config_values.regions = self._get_regions_element(
-            selected=config.regions or [DEFAULT_REGION]
-        )
+        config_values.regions = self._get_regions_element(selected=config.regions or [])
         return config_values
 
     def create_backend(

src/dstack/_internal/server/services/config.py

@@ -436,7 +436,6 @@ class VultrConfig(CoreModel):
         Optional[List[str]],
         Field(description="The list of Vultr regions. Omit to use all regions"),
     ] = None
-    project_id: Annotated[str, Field(description="The project ID")]
     creds: Annotated[AnyVultrCreds, Field(description="The credentials")]