WIP: support GCP TCPXO

Author: r4victor

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

@@ -39,11 +39,11 @@
 
 logger = get_logger(__name__)
 
-DSTACK_WORKING_DIR = "/root/.dstack"
+DSTACK_WORKING_DIR = "/etc/.dstack"
 DSTACK_SHIM_BINARY_NAME = "dstack-shim"
-DSTACK_SHIM_BINARY_PATH = f"/usr/local/bin/{DSTACK_SHIM_BINARY_NAME}"
+DSTACK_SHIM_BINARY_PATH = f"/etc/{DSTACK_SHIM_BINARY_NAME}"
 DSTACK_RUNNER_BINARY_NAME = "dstack-runner"
-DSTACK_RUNNER_BINARY_PATH = f"/usr/local/bin/{DSTACK_RUNNER_BINARY_NAME}"
+DSTACK_RUNNER_BINARY_PATH = f"/etc/{DSTACK_RUNNER_BINARY_NAME}"
 
 
 class Compute(ABC):
@@ -525,7 +525,7 @@ def get_run_shim_script(is_privileged: bool, pjrt_device: Optional[str]) -> List
     pjrt_device_env = f"--pjrt-device={pjrt_device}" if pjrt_device else ""
 
     return [
-        f"nohup dstack-shim {privileged_flag} {pjrt_device_env} &",
+        f"nohup {DSTACK_SHIM_BINARY_PATH} {privileged_flag} {pjrt_device_env} &",
     ]
 
 

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

@@ -269,7 +269,10 @@ def create_instance(
                     gpus=instance_offer.instance.resources.gpus,
                 ),
                 spot=instance_offer.instance.resources.spot,
-                user_data=get_user_data(authorized_keys),
+                user_data=get_user_data(
+                    authorized_keys,
+                    backend_specific_commands=_get_backend_specific_commands_tcpxo(),
+                ),
                 authorized_keys=authorized_keys,
                 labels=labels,
                 tags=[gcp_resources.DSTACK_INSTANCE_TAG],
@@ -805,6 +808,68 @@ def _is_single_host_tpu(instance_name: str) -> bool:
         return False
 
 
+def _get_backend_specific_commands_tcpx() -> List[str]:
+    return [
+        "cos-extensions install gpu -- --version=latest",
+        "sudo mount --bind /var/lib/nvidia /var/lib/nvidia",
+        "sudo mount -o remount,exec /var/lib/nvidia",
+        (
+            "docker run "
+            "--detach "
+            "--pull=always"
+            "--name receive-datapath-manager "
+            "--privileged "
+            "--cap-add=NET_ADMIN --network=host "
+            "--volume /var/lib/nvidia/lib64:/usr/local/nvidia/lib64 "
+            "--device /dev/nvidia0:/dev/nvidia0 --device /dev/nvidia1:/dev/nvidia1 "
+            "--device /dev/nvidia2:/dev/nvidia2 --device /dev/nvidia3:/dev/nvidia3 "
+            "--device /dev/nvidia4:/dev/nvidia4 --device /dev/nvidia5:/dev/nvidia5 "
+            "--device /dev/nvidia6:/dev/nvidia6 --device /dev/nvidia7:/dev/nvidia7 "
+            "--device /dev/nvidia-uvm:/dev/nvidia-uvm --device /dev/nvidiactl:/dev/nvidiactl "
+            "--env LD_LIBRARY_PATH=/usr/local/nvidia/lib64 "
+            "--volume /run/tcpx:/run/tcpx "
+            "--entrypoint /tcpgpudmarxd/build/app/tcpgpudmarxd "
+            "us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/tcpgpudmarxd "
+            '--gpu_nic_preset a3vm --gpu_shmem_type fd --uds_path "/run/tcpx" --setup_param "--verbose 128 2 0"'
+        ),
+        "sudo iptables -I INPUT -p tcp -m tcp -j ACCEPT",
+        "docker run --rm -v /var/lib:/var/lib us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx install --install-nccl",
+        "sudo mount --bind /var/lib/tcpx /var/lib/tcpx",
+        "sudo mount -o remount,exec /var/lib/tcpx",
+    ]
+
+
+def _get_backend_specific_commands_tcpxo() -> List[str]:
+    return [
+        "modprobe import-helper",
+        "gcloud -q auth configure-docker us-docker.pkg.dev",
+        # Install the nccl, nccl-net lib into /var/lib/tcpxo/lib64/.
+        (
+            "docker run --rm --name nccl-installer "
+            "--network=host "
+            "--volume /var/lib:/var/lib "
+            "us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/nccl-plugin-gpudirecttcpx-dev:v1.0.8-1 "
+            "install --install-nccl"
+        ),
+        # Start FasTrak receive-datapath-manager
+        (
+            "docker run "
+            "--detach "
+            "--pull=always "
+            "--name receive-datapath-manager "
+            "--cap-add=NET_ADMIN "
+            "--network=host "
+            "--privileged "
+            "--gpus all "
+            "--volume /usr/lib32:/usr/local/nvidia/lib64 "
+            "--volume /dev/dmabuf_import_helper:/dev/dmabuf_import_helper "
+            "--env LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu "
+            "us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/tcpgpudmarxd-dev:v1.0.14 "
+            "--num_hops=2 --num_nics=8 --uid= --alsologtostderr"
+        ),
+    ]
+
+
 def _get_volume_price(size: int) -> float:
     # https://cloud.google.com/compute/disks-image-pricing#persistentdisk
     # The price is different in different regions. Take max across supported regions.

src/dstack/_internal/core/backends/gcp/resources.py

@@ -8,7 +8,6 @@
 from google.api_core.operation import Operation
 from google.cloud import tpu_v2
 
-import dstack.version as version
 from dstack._internal.core.errors import BackendError, ComputeError
 from dstack._internal.core.models.instances import Gpu
 from dstack._internal.utils.common import remove_prefix
@@ -119,24 +118,14 @@ def create_instance_struct(
     subnetwork: Optional[str] = None,
     allocate_public_ip: bool = True,
 ) -> compute_v1.Instance:
-    network_interface = compute_v1.NetworkInterface()
-    network_interface.network = network
-    if subnetwork is not None:
-        network_interface.subnetwork = subnetwork
-
-    if allocate_public_ip:
-        access = compute_v1.AccessConfig()
-        access.type_ = compute_v1.AccessConfig.Type.ONE_TO_ONE_NAT.name
-        access.name = "External NAT"
-        access.network_tier = access.NetworkTier.PREMIUM.name
-        network_interface.access_configs = [access]
-    else:
-        network_interface.access_configs = []
-
     instance = compute_v1.Instance()
-    instance.network_interfaces = [network_interface]
     instance.name = instance_name
     instance.machine_type = f"zones/{zone}/machineTypes/{machine_type}"
+    instance.network_interfaces = _get_network_interfaces(
+        network=network,
+        subnetwork=subnetwork,
+        allocate_public_ip=allocate_public_ip,
+    )
 
     disk = compute_v1.AttachedDisk()
     disk.auto_delete = True
@@ -187,14 +176,45 @@ def create_instance_struct(
     return instance
 
 
-def get_image_id(cuda: bool) -> str:
-    if not cuda:
-        image_name = f"dstack-{version.base_image}"
+def _get_network_interfaces(
+    network: str,
+    subnetwork: Optional[str],
+    allocate_public_ip: bool,
+) -> List[compute_v1.NetworkInterface]:
+    network_interface = compute_v1.NetworkInterface()
+    network_interface.network = network
+    if subnetwork is not None:
+        network_interface.subnetwork = subnetwork
+    if allocate_public_ip:
+        access = compute_v1.AccessConfig()
+        access.type_ = compute_v1.AccessConfig.Type.ONE_TO_ONE_NAT.name
+        access.name = "External NAT"
+        access.network_tier = access.NetworkTier.PREMIUM.name
+        network_interface.access_configs = [access]
     else:
-        image_name = f"dstack-cuda-{version.base_image}"
-    image_name = image_name.replace(".", "-")
+        network_interface.access_configs = []
+
+    network_interfaces = [network_interface]
+    for i in range(1, 9):
+        network_interfaces.append(
+            compute_v1.NetworkInterface(
+                network=f"projects/dstack/global/networks/dstack-test-data-net-{i}",
+                subnetwork=f"projects/dstack/regions/europe-west4/subnetworks/dstack-test-data-sub-{i}",
+            )
+        )
+    return network_interfaces
 
-    return f"projects/dstack/global/images/{image_name}"
+
+def get_image_id(cuda: bool) -> str:
+    # if not cuda:
+    #     image_name = f"dstack-{version.base_image}"
+    # else:
+    #     image_name = f"dstack-cuda-{version.base_image}"
+    # image_name = image_name.replace(".", "-")
+
+    # return f"projects/dstack/global/images/{image_name}"
+    # return "projects/cos-cloud/global/images/cos-105-17412-535-78" # TCPX
+    return "projects/dstack/global/images/slurm-a3mega-20250327t101736z-cloudinit"  # TCPXO
 
 
 def get_gateway_image_id() -> str: