Support internal IP and cluster placement for jarvislabs (#4247)

Author: r4victor

mkdocs/docs/concepts/backends.md

@@ -952,6 +952,9 @@ projects:
 
 </div>
 
+??? info "Cluster placement"
+    Fleets with `placement: cluster` are supported in the `india-chennai-01` and `india-noida-01` regions, where instances join the account's default [VPC](https://docs.jarvislabs.ai/vpc/) and communicate over private IPs. No additional configuration is required.
+
 ### CloudRift
 
 Log into your [CloudRift](https://console.cloudrift.ai/) console, click `API Keys` in the sidebar and click the button to create a new API key.

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

@@ -2,7 +2,7 @@
 import subprocess
 import tempfile
 from collections.abc import Iterable
-from typing import List, Optional, cast
+from typing import Callable, List, Optional, cast
 
 import gpuhunt
 from gpuhunt.providers.jarvislabs import JarvisLabsProvider
@@ -13,6 +13,7 @@
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
     ComputeWithInstanceVolumesSupport,
+    ComputeWithMultinodeSupport,
     ComputeWithPrivilegedSupport,
     generate_unique_instance_name,
     get_shim_commands,
@@ -47,6 +48,11 @@
 SSH_CONNECT_TIMEOUT_SECONDS = 10
 SSH_SETUP_TIMEOUT_SECONDS = 240
 SSH_LAUNCH_TIMEOUT_SECONDS = 60
+# VMs in these regions join the account's default VPC and get a private IP that is
+# routable between VMs of the same account. Other regions have no VPC, so their
+# private IPs cannot be used for inter-node communication.
+# See https://docs.jarvislabs.ai/vpc/
+VPC_REGIONS = frozenset({"india-chennai-01", "india-noida-01"})
 
 
 class JarvisLabsOfferBackendData(TypedDict):
@@ -70,6 +76,7 @@ class JarvisLabsCompute(
     ComputeWithCreateInstanceSupport,
     ComputeWithPrivilegedSupport,
     ComputeWithInstanceVolumesSupport,
+    ComputeWithMultinodeSupport,
     Compute,
 ):
     def __init__(self, config: JarvisLabsConfig):
@@ -98,6 +105,13 @@ def get_offers_modifiers(
     ) -> Iterable[OfferModifier]:
         return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)]
 
+    def get_offers_post_filter(
+        self, requirements: Requirements
+    ) -> Optional[Callable[[InstanceOfferWithAvailability], bool]]:
+        if not requirements.multinode:
+            return None
+        return lambda offer: offer.region in VPC_REGIONS
+
     def create_instance(
         self,
         instance_offer: InstanceOfferWithAvailability,
@@ -207,6 +221,7 @@ def update_provisioning_data(
             return
         provisioning_data.hostname = hostname
         provisioning_data.username = username
+        provisioning_data.internal_ip = _get_internal_ip(instance)
 
     def terminate_instance(
         self, instance_id: str, region: str, backend_data: Optional[str] = None
@@ -261,6 +276,16 @@ def _raise_failed_status(status: dict) -> None:
     raise ProvisioningError(_format_failed_status(status), status)
 
 
+def _get_internal_ip(instance: dict) -> Optional[str]:
+    # Private IPs of VMs outside a VPC are not routable between VMs.
+    if not instance.get("vpc_id"):
+        return None
+    private_ip = instance.get("private_ip")
+    if not isinstance(private_ip, str) or not private_ip:
+        return None
+    return private_ip
+
+
 def _get_ssh_username(instance: dict) -> str:
     ssh_command = instance.get("ssh_str") or instance.get("ssh_command")
     if not isinstance(ssh_command, str):