Add vpc cluster support in Vultr (#2196)

Co-authored-by: Bihan  Rana <[email protected]>

src/dstack/_internal/core/backends/__init__.py

@@ -6,6 +6,7 @@
     BackendType.GCP,
     BackendType.REMOTE,
     BackendType.OCI,
+    BackendType.VULTR,
 ]
 BACKENDS_WITH_CREATE_INSTANCE_SUPPORT = [
     BackendType.AWS,

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

@@ -1,5 +1,5 @@
 import base64
-from typing import Any
+from typing import Any, Optional
 
 import requests
 from requests import Response
@@ -28,7 +28,21 @@ 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, user_data: str):
+    def get_vpc_for_region(self, region: str) -> Optional[str]:
+        response = self._make_request("GET", "/vpcs?per_page=500")
+        vpcs = response.json().get("vpcs", [])
+        if vpcs:
+            for vpc in vpcs:
+                if vpc["description"] == f"dstack-vpc-{region}":
+                    return vpc
+        return None
+
+    def create_vpc(self, region: str):
+        data = {"region": region, "description": f"dstack-vpc-{region}"}
+        response = self._make_request("POST", "/vpcs", data=data)
+        return response.json()["vpc"]
+
+    def launch_instance(self, region: str, plan: str, label: str, user_data: str, vpc_id: str):
         # For Bare-metals
         if "vbm" in plan:
             # "Docker on Ubuntu 22.04" is required for bare-metals.
@@ -38,6 +52,7 @@ def launch_instance(self, region: str, plan: str, label: str, user_data: str):
                 "label": label,
                 "image_id": "docker",
                 "user_data": base64.b64encode(user_data.encode()).decode(),
+                "attach_vpc": [vpc_id],
             }
             resp = self._make_request("POST", "/bare-metals", data)
             return resp.json()["bare_metal"]["id"]
@@ -50,6 +65,7 @@ def launch_instance(self, region: str, plan: str, label: str, user_data: str):
                 "label": label,
                 "os_id": 1743,
                 "user_data": base64.b64encode(user_data.encode()).decode(),
+                "attach_vpc": [vpc_id],
             }
             resp = self._make_request("POST", "/instances", data)
             return resp.json()["instance"]["id"]
@@ -60,6 +76,7 @@ def launch_instance(self, region: str, plan: str, label: str, user_data: str):
                 "label": label,
                 "image_id": "docker",
                 "user_data": base64.b64encode(user_data.encode()).decode(),
+                "attach_vpc": [vpc_id],
             }
             resp = self._make_request("POST", "/instances", data)
             return resp.json()["instance"]["id"]

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

@@ -71,11 +71,27 @@ def run_job(
     def create_instance(
         self, instance_offer: InstanceOfferWithAvailability, instance_config: InstanceConfiguration
     ) -> JobProvisioningData:
+        # create vpc
+        vpc = self.api_client.get_vpc_for_region(instance_offer.region)
+        if not vpc:
+            vpc = self.api_client.create_vpc(instance_offer.region)
+
+        subnet = vpc["v4_subnet"]
+        subnet_mask = vpc["v4_subnet_mask"]
+
+        setup_commands = [
+            f"sudo ufw allow from {subnet}/{subnet_mask}",
+            "sudo ufw reload",
+        ]
         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()),
+            user_data=get_user_data(
+                authorized_keys=instance_config.get_public_keys(),
+                backend_specific_commands=setup_commands,
+            ),
+            vpc_id=vpc["id"],
         )
 
         launched_instance = JobProvisioningData(
@@ -120,13 +136,18 @@ def update_provisioning_data(
         # Access specific fields
         instance_status = instance_data["status"]
         instance_main_ip = instance_data["main_ip"]
+        instance_internal_ip = instance_data["internal_ip"]
         if instance_status == "active":
             provisioning_data.hostname = instance_main_ip
+            provisioning_data.internal_ip = instance_internal_ip
         if instance_status == "failed":
             raise ProvisioningError("VM entered FAILED state")
 
 
 def _supported_instances(offer: InstanceOffer) -> bool:
+    # The vbm-4c-32gb plan does not support VPC, so it is excluded.
+    if offer.instance.name == "vbm-4c-32gb":
+        return False
     if offer.instance.resources.spot:
         return False
     for family in [