dstack 0.19.6rc1__py3-none-any.whl → 0.19.8__py3-none-any.whl

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

@@ -6,7 +6,7 @@ import threading
 from abc import ABC, abstractmethod
 from functools import lru_cache
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, List, Literal, Optional
 
 import git
 import requests
@@ -25,6 +25,7 @@ from dstack._internal.core.models.gateways import (
 )
 from dstack._internal.core.models.instances import (
     InstanceConfiguration,
+    InstanceOffer,
     InstanceOfferWithAvailability,
     SSHKey,
 )
@@ -44,6 +45,8 @@ logger = get_logger(__name__)
 DSTACK_SHIM_BINARY_NAME = "dstack-shim"
 DSTACK_RUNNER_BINARY_NAME = "dstack-runner"
 
+GoArchType = Literal["amd64", "arm64"]
+
 
 class Compute(ABC):
     """
@@ -144,6 +147,7 @@ class ComputeWithCreateInstanceSupport(ABC):
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         """
         Launches a new instance. It should return `JobProvisioningData` ASAP.
@@ -176,7 +180,7 @@ class ComputeWithCreateInstanceSupport(ABC):
         )
         instance_offer = instance_offer.copy()
         self._restrict_instance_offer_az_to_volumes_az(instance_offer, volumes)
-        return self.create_instance(instance_offer, instance_config)
+        return self.create_instance(instance_offer, instance_config, placement_group=None)
 
     def _restrict_instance_offer_az_to_volumes_az(
         self,
@@ -225,9 +229,15 @@ class ComputeWithPlacementGroupSupport(ABC):
     def create_placement_group(
         self,
         placement_group: PlacementGroup,
+        master_instance_offer: InstanceOffer,
     ) -> PlacementGroupProvisioningData:
         """
         Creates a placement group.
+
+        Args:
+            placement_group: details about the placement group to be created
+            master_instance_offer: the first instance dstack will attempt to add
+                                   to the placement group
         """
         pass
 
@@ -242,10 +252,27 @@ class ComputeWithPlacementGroupSupport(ABC):
         """
         pass
 
+    @abstractmethod
+    def is_suitable_placement_group(
+        self,
+        placement_group: PlacementGroup,
+        instance_offer: InstanceOffer,
+    ) -> bool:
+        """
+        Checks if the instance offer can be provisioned in the placement group.
+
+        Should return immediately, without performing API calls.
+
+        Can be called with an offer originating from a different backend, because some backends
+        (BackendType.DSTACK) produce offers on behalf of other backends. Should return `False`
+        in that case.
+        """
+        pass
+
 
 class ComputeWithGatewaySupport(ABC):
     """
-    Must be subclassed and imlemented to support gateways.
+    Must be subclassed and implemented to support gateways.
     """
 
     @abstractmethod
@@ -418,6 +445,21 @@ def generate_unique_volume_name(
     )
 
 
+def generate_unique_placement_group_name(
+    project_name: str,
+    fleet_name: str,
+    max_length: int = _DEFAULT_MAX_RESOURCE_NAME_LEN,
+) -> str:
+    """
+    Generates a unique placement group name valid across all backends.
+    """
+    return generate_unique_backend_name(
+        resource_name=fleet_name,
+        project_name=project_name,
+        max_length=max_length,
+    )
+
+
 def generate_unique_backend_name(
     resource_name: str,
     project_name: Optional[str],
@@ -483,13 +525,14 @@ def get_shim_env(
     base_path: Optional[PathLike] = None,
     bin_path: Optional[PathLike] = None,
     backend_shim_env: Optional[Dict[str, str]] = None,
+    arch: Optional[str] = None,
 ) -> Dict[str, str]:
     log_level = "6"  # Trace
     envs = {
         "DSTACK_SHIM_HOME": get_dstack_working_dir(base_path),
         "DSTACK_SHIM_HTTP_PORT": str(DSTACK_SHIM_HTTP_PORT),
         "DSTACK_SHIM_LOG_LEVEL": log_level,
-        "DSTACK_RUNNER_DOWNLOAD_URL": get_dstack_runner_download_url(),
+        "DSTACK_RUNNER_DOWNLOAD_URL": get_dstack_runner_download_url(arch),
         "DSTACK_RUNNER_BINARY_PATH": get_dstack_runner_binary_path(bin_path),
         "DSTACK_RUNNER_HTTP_PORT": str(DSTACK_RUNNER_HTTP_PORT),
         "DSTACK_RUNNER_SSH_PORT": str(DSTACK_RUNNER_SSH_PORT),
@@ -509,16 +552,19 @@ def get_shim_commands(
     base_path: Optional[PathLike] = None,
     bin_path: Optional[PathLike] = None,
     backend_shim_env: Optional[Dict[str, str]] = None,
+    arch: Optional[str] = None,
 ) -> List[str]:
     commands = get_shim_pre_start_commands(
         base_path=base_path,
         bin_path=bin_path,
+        arch=arch,
     )
     shim_env = get_shim_env(
         authorized_keys=authorized_keys,
         base_path=base_path,
         bin_path=bin_path,
         backend_shim_env=backend_shim_env,
+        arch=arch,
     )
     for k, v in shim_env.items():
         commands += [f'export "{k}={v}"']
@@ -539,35 +585,63 @@ def get_dstack_runner_version() -> str:
     return version or "latest"
 
 
-def get_dstack_runner_download_url() -> str:
-    if url := os.environ.get("DSTACK_RUNNER_DOWNLOAD_URL"):
-        return url
-    build = get_dstack_runner_version()
-    if settings.DSTACK_VERSION is not None:
-        bucket = "dstack-runner-downloads"
-    else:
-        bucket = "dstack-runner-downloads-stgn"
-    return (
-        f"https://{bucket}.s3.eu-west-1.amazonaws.com/{build}/binaries/dstack-runner-linux-amd64"
-    )
-
-
-def get_dstack_shim_download_url() -> str:
-    if url := os.environ.get("DSTACK_SHIM_DOWNLOAD_URL"):
-        return url
-    build = get_dstack_runner_version()
-    if settings.DSTACK_VERSION is not None:
-        bucket = "dstack-runner-downloads"
-    else:
-        bucket = "dstack-runner-downloads-stgn"
-    return f"https://{bucket}.s3.eu-west-1.amazonaws.com/{build}/binaries/dstack-shim-linux-amd64"
+def normalize_arch(arch: Optional[str] = None) -> GoArchType:
+    """
+    Converts the given free-form architecture string to the Go GOARCH format.
+    Only 64-bit x86 and ARM are supported. If the word size is not specified (e.g., `x86`, `arm`),
+    64-bit is implied.
+    If the arch is not specified, falls back to `amd64`.
+    """
+    if not arch:
+        return "amd64"
+    arch_lower = arch.lower()
+    if "32" in arch_lower or arch_lower in ["i386", "i686"]:
+        raise ValueError(f"32-bit architectures are not supported: {arch}")
+    if arch_lower.startswith("x86") or arch_lower.startswith("amd"):
+        return "amd64"
+    if arch_lower.startswith("arm") or arch_lower.startswith("aarch"):
+        return "arm64"
+    raise ValueError(f"Unsupported architecture: {arch}")
+
+
+def get_dstack_runner_download_url(arch: Optional[str] = None) -> str:
+    url_template = os.environ.get("DSTACK_RUNNER_DOWNLOAD_URL")
+    if not url_template:
+        if settings.DSTACK_VERSION is not None:
+            bucket = "dstack-runner-downloads"
+        else:
+            bucket = "dstack-runner-downloads-stgn"
+        url_template = (
+            f"https://{bucket}.s3.eu-west-1.amazonaws.com"
+            "/{version}/binaries/dstack-runner-linux-{arch}"
+        )
+    version = get_dstack_runner_version()
+    arch = normalize_arch(arch)
+    return url_template.format(version=version, arch=arch)
+
+
+def get_dstack_shim_download_url(arch: Optional[str] = None) -> str:
+    url_template = os.environ.get("DSTACK_SHIM_DOWNLOAD_URL")
+    if not url_template:
+        if settings.DSTACK_VERSION is not None:
+            bucket = "dstack-runner-downloads"
+        else:
+            bucket = "dstack-runner-downloads-stgn"
+        url_template = (
+            f"https://{bucket}.s3.eu-west-1.amazonaws.com"
+            "/{version}/binaries/dstack-shim-linux-{arch}"
+        )
+    version = get_dstack_runner_version()
+    arch = normalize_arch(arch)
+    return url_template.format(version=version, arch=arch)
 
 
 def get_shim_pre_start_commands(
     base_path: Optional[PathLike] = None,
     bin_path: Optional[PathLike] = None,
+    arch: Optional[str] = None,
 ) -> List[str]:
-    url = get_dstack_shim_download_url()
+    url = get_dstack_shim_download_url(arch)
     dstack_shim_binary_path = get_dstack_shim_binary_path(bin_path)
     dstack_working_dir = get_dstack_working_dir(base_path)
     return [

dstack/_internal/core/backends/base/offers.py

@@ -2,6 +2,7 @@ from dataclasses import asdict
 from typing import Callable, List, Optional
 
 import gpuhunt
+from pydantic import parse_obj_as
 
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.instances import (
@@ -11,13 +12,14 @@ from dstack._internal.core.models.instances import (
     InstanceType,
     Resources,
 )
-from dstack._internal.core.models.resources import DEFAULT_DISK, Memory, Range
+from dstack._internal.core.models.resources import DEFAULT_DISK, CPUSpec, Memory, Range
 from dstack._internal.core.models.runs import Requirements
 
 # Offers not supported by all dstack versions are hidden behind one or more flags.
 # This list enables the flags that are currently supported.
 SUPPORTED_GPUHUNT_FLAGS = [
     "oci-spot",
+    "lambda-arm",
 ]
 
 
@@ -71,6 +73,7 @@ def catalog_item_to_offer(
     if disk_size_mib is None:
         return None
     resources = Resources(
+        cpu_arch=item.cpu_arch,
         cpus=item.cpu,
         memory_mib=round(item.memory * 1024),
         gpus=gpus,
@@ -90,6 +93,9 @@ def catalog_item_to_offer(
 
 
 def offer_to_catalog_item(offer: InstanceOffer) -> gpuhunt.CatalogItem:
+    cpu_arch = offer.instance.resources.cpu_arch
+    if cpu_arch is None:
+        cpu_arch = gpuhunt.CPUArchitecture.X86
     gpu_count = len(offer.instance.resources.gpus)
     gpu_vendor = None
     gpu_name = None
@@ -104,6 +110,7 @@ def offer_to_catalog_item(offer: InstanceOffer) -> gpuhunt.CatalogItem:
         instance_name=offer.instance.name,
         location=offer.region,
         price=offer.price,
+        cpu_arch=cpu_arch,
         cpu=offer.instance.resources.cpus,
         memory=offer.instance.resources.memory_mib / 1024,
         gpu_count=gpu_count,
@@ -125,8 +132,11 @@ def requirements_to_query_filter(req: Optional[Requirements]) -> gpuhunt.QueryFi
 
     res = req.resources
     if res.cpu:
-        q.min_cpu = res.cpu.min
-        q.max_cpu = res.cpu.max
+        # TODO: Remove in 0.20. Use res.cpu directly
+        cpu = parse_obj_as(CPUSpec, res.cpu)
+        q.cpu_arch = cpu.arch
+        q.min_cpu = cpu.count.min
+        q.max_cpu = cpu.count.max
     if res.memory:
         q.min_memory = res.memory.min
         q.max_memory = res.memory.max

dstack/_internal/core/backends/cudo/compute.py

@@ -18,6 +18,7 @@ from dstack._internal.core.models.instances import (
     InstanceConfiguration,
     InstanceOfferWithAvailability,
 )
+from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.utils.logging import get_logger
 
@@ -58,6 +59,7 @@ class CudoCompute(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         vm_id = generate_unique_instance_name(instance_config, max_length=MAX_RESOURCE_NAME_LEN)
         public_keys = instance_config.get_public_keys()

dstack/_internal/core/backends/datacrunch/compute.py

@@ -20,6 +20,7 @@ from dstack._internal.core.models.instances import (
     InstanceOffer,
     InstanceOfferWithAvailability,
 )
+from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.resources import Memory, Range
 from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.utils.logging import get_logger
@@ -85,6 +86,7 @@ class DataCrunchCompute(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         instance_name = generate_unique_instance_name(
             instance_config, max_length=MAX_INSTANCE_NAME_LEN

dstack/_internal/core/backends/gcp/auth.py

@@ -19,7 +19,7 @@ def authenticate(creds: AnyGCPCreds, project_id: Optional[str] = None) -> Tuple[
     credentials, credentials_project_id = get_credentials(creds)
     if project_id is None:
         # If project_id is not specified explicitly, try using credentials' project_id.
-        # Explicit project_id takes precedence bacause credentials' project_id may be irrelevant.
+        # Explicit project_id takes precedence because credentials' project_id may be irrelevant.
         # For example, with Workload Identity Federation for GKE, it's cluster project_id.
         project_id = credentials_project_id
     if project_id is None:

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

@@ -1,10 +1,12 @@
 import concurrent.futures
 import json
+import threading
 from collections import defaultdict
 from typing import Callable, Dict, List, Literal, Optional, Tuple
 
 import google.api_core.exceptions
 import google.cloud.compute_v1 as compute_v1
+from cachetools import TTLCache, cachedmethod
 from google.cloud import tpu_v2
 from gpuhunt import KNOWN_TPUS
 
@@ -98,6 +100,8 @@ class GCPCompute(
         self.resource_policies_client = compute_v1.ResourcePoliciesClient(
             credentials=self.credentials
         )
+        self._extra_subnets_cache_lock = threading.Lock()
+        self._extra_subnets_cache = TTLCache(maxsize=30, ttl=60)
 
     def get_offers(
         self, requirements: Optional[Requirements] = None
@@ -166,6 +170,7 @@ class GCPCompute(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         instance_name = generate_unique_instance_name(
             instance_config, max_length=gcp_resources.MAX_RESOURCE_NAME_LEN
@@ -192,18 +197,16 @@ class GCPCompute(
             config=self.config,
             region=instance_offer.region,
         )
-        extra_subnets = _get_extra_subnets(
-            subnetworks_client=self.subnetworks_client,
-            config=self.config,
+        extra_subnets = self._get_extra_subnets(
             region=instance_offer.region,
             instance_type_name=instance_offer.instance.name,
         )
         placement_policy = None
-        if instance_config.placement_group_name is not None:
+        if placement_group is not None:
             placement_policy = gcp_resources.get_placement_policy_resource_name(
                 project_id=self.config.project_id,
                 region=instance_offer.region,
-                placement_policy=instance_config.placement_group_name,
+                placement_policy=placement_group.name,
             )
         labels = {
             "owner": "dstack",
@@ -406,6 +409,7 @@ class GCPCompute(
     def create_placement_group(
         self,
         placement_group: PlacementGroup,
+        master_instance_offer: InstanceOffer,
     ) -> PlacementGroupProvisioningData:
         policy = compute_v1.ResourcePolicy(
             name=placement_group.name,
@@ -440,6 +444,16 @@ class GCPCompute(
                 raise PlacementGroupInUseError()
             raise
 
+    def is_suitable_placement_group(
+        self,
+        placement_group: PlacementGroup,
+        instance_offer: InstanceOffer,
+    ) -> bool:
+        return (
+            placement_group.configuration.backend == BackendType.GCP
+            and placement_group.configuration.region == instance_offer.region
+        )
+
     def create_gateway(
         self,
         configuration: GatewayComputeConfiguration,
@@ -757,6 +771,38 @@ class GCPCompute(
             instance_id,
         )
 
+    @cachedmethod(
+        cache=lambda self: self._extra_subnets_cache,
+        lock=lambda self: self._extra_subnets_cache_lock,
+    )
+    def _get_extra_subnets(
+        self,
+        region: str,
+        instance_type_name: str,
+    ) -> List[Tuple[str, str]]:
+        if self.config.extra_vpcs is None:
+            return []
+        if instance_type_name == "a3-megagpu-8g":
+            subnets_num = 8
+        elif instance_type_name in ["a3-edgegpu-8g", "a3-highgpu-8g"]:
+            subnets_num = 4
+        else:
+            return []
+        extra_subnets = []
+        for vpc_name in self.config.extra_vpcs[:subnets_num]:
+            subnet = gcp_resources.get_vpc_subnet_or_error(
+                subnetworks_client=self.subnetworks_client,
+                vpc_project_id=self.config.vpc_project_id or self.config.project_id,
+                vpc_name=vpc_name,
+                region=region,
+            )
+            vpc_resource_name = gcp_resources.vpc_name_to_vpc_resource_name(
+                project_id=self.config.vpc_project_id or self.config.project_id,
+                vpc_name=vpc_name,
+            )
+            extra_subnets.append((vpc_resource_name, subnet))
+        return extra_subnets
+
 
 def _supported_instances_and_zones(
     regions: List[str],
@@ -831,36 +877,6 @@ def _get_vpc_subnet(
     )
 
 
-def _get_extra_subnets(
-    subnetworks_client: compute_v1.SubnetworksClient,
-    config: GCPConfig,
-    region: str,
-    instance_type_name: str,
-) -> List[Tuple[str, str]]:
-    if config.extra_vpcs is None:
-        return []
-    if instance_type_name == "a3-megagpu-8g":
-        subnets_num = 8
-    elif instance_type_name in ["a3-edgegpu-8g", "a3-highgpu-8g"]:
-        subnets_num = 4
-    else:
-        return []
-    extra_subnets = []
-    for vpc_name in config.extra_vpcs[:subnets_num]:
-        subnet = gcp_resources.get_vpc_subnet_or_error(
-            subnetworks_client=subnetworks_client,
-            vpc_project_id=config.vpc_project_id or config.project_id,
-            vpc_name=vpc_name,
-            region=region,
-        )
-        vpc_resource_name = gcp_resources.vpc_name_to_vpc_resource_name(
-            project_id=config.vpc_project_id or config.project_id,
-            vpc_name=vpc_name,
-        )
-        extra_subnets.append((vpc_resource_name, subnet))
-    return extra_subnets
-
-
 def _get_image_id(instance_type_name: str, cuda: bool) -> str:
     if instance_type_name == "a3-megagpu-8g":
         image_name = "dstack-a3mega-5"

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

@@ -205,12 +205,17 @@ def _get_network_interfaces(
     else:
         network_interface.access_configs = []
 
+    if extra_subnetworks:
+        # Multiple interfaces are set only for GPU VM that require gVNIC for best performance
+        network_interface.nic_type = compute_v1.NetworkInterface.NicType.GVNIC.name
+
     network_interfaces = [network_interface]
     for network, subnetwork in extra_subnetworks or []:
         network_interfaces.append(
             compute_v1.NetworkInterface(
                 network=network,
                 subnetwork=subnetwork,
+                nic_type=compute_v1.NetworkInterface.NicType.GVNIC.name,
             )
         )
     return network_interfaces
@@ -437,7 +442,7 @@ def wait_for_operation(operation: Operation, verbose_name: str = "operation", ti
         raise
     except Exception as e:
         # Write only debug logs here.
-        # The unexpected errors will be propagated and logged appropriatly by the caller.
+        # The unexpected errors will be propagated and logged appropriately by the caller.
         logger.debug("Error during %s: %s", verbose_name, e)
         raise operation.exception() or e
     return result

dstack/_internal/core/backends/lambdalabs/compute.py

@@ -20,6 +20,7 @@ from dstack._internal.core.models.instances import (
     InstanceOffer,
     InstanceOfferWithAvailability,
 )
+from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 
 MAX_INSTANCE_NAME_LEN = 60
@@ -46,7 +47,10 @@ class LambdaCompute(
         return offers_with_availability
 
     def create_instance(
-        self, instance_offer: InstanceOfferWithAvailability, instance_config: InstanceConfiguration
+        self,
+        instance_offer: InstanceOfferWithAvailability,
+        instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         instance_name = generate_unique_instance_name(
             instance_config, max_length=MAX_INSTANCE_NAME_LEN
@@ -89,7 +93,10 @@ class LambdaCompute(
         instance_info = _get_instance_info(self.api_client, provisioning_data.instance_id)
         if instance_info is not None and instance_info["status"] != "booting":
             provisioning_data.hostname = instance_info["ip"]
-            commands = get_shim_commands(authorized_keys=[project_ssh_public_key])
+            commands = get_shim_commands(
+                authorized_keys=[project_ssh_public_key],
+                arch=provisioning_data.instance_type.resources.cpu_arch,
+            )
             # shim is assumed to be run under root
             launch_command = "sudo sh -c '" + "&& ".join(commands) + "'"
             thread = Thread(
@@ -179,13 +186,18 @@ def _setup_instance(
     ssh_private_key: str,
 ):
     setup_commands = (
-        "mkdir /home/ubuntu/.dstack && "
-        "sudo apt-get update && "
-        "sudo apt-get install -y --no-install-recommends nvidia-container-toolkit && "
-        "sudo nvidia-ctk runtime configure --runtime=docker && "
-        "sudo pkill -SIGHUP dockerd"
+        "mkdir /home/ubuntu/.dstack",
+        "sudo apt-get update",
+        "sudo apt-get install -y --no-install-recommends nvidia-container-toolkit",
+        "sudo install -d -m 0755 /etc/docker",
+        # Workaround for https://github.com/NVIDIA/nvidia-container-toolkit/issues/48
+        """echo '{"exec-opts":["native.cgroupdriver=cgroupfs"]}' | sudo tee /etc/docker/daemon.json""",
+        "sudo nvidia-ctk runtime configure --runtime=docker",
+        "sudo systemctl restart docker.service",  # `systemctl reload` (`kill -HUP`) won't work
+    )
+    _run_ssh_command(
+        hostname=hostname, ssh_private_key=ssh_private_key, command=" && ".join(setup_commands)
     )
-    _run_ssh_command(hostname=hostname, ssh_private_key=ssh_private_key, command=setup_commands)
 
 
 def _launch_runner(

dstack/_internal/core/backends/local/compute.py

@@ -15,6 +15,7 @@ from dstack._internal.core.models.instances import (
     InstanceType,
     Resources,
 )
+from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
 from dstack._internal.core.models.volumes import Volume, VolumeProvisioningData
 from dstack._internal.utils.logging import get_logger
@@ -53,6 +54,7 @@ class LocalCompute(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         return JobProvisioningData(
             backend=instance_offer.backend,