dstack 0.19.30__py3-none-any.whl → 0.19.31__py3-none-any.whl

dstack/_internal/cli/commands/__init__.py

@@ -7,6 +7,7 @@ from typing import ClassVar, Optional
 from rich_argparse import RichHelpFormatter
 
 from dstack._internal.cli.services.completion import ProjectNameCompleter
+from dstack._internal.cli.utils.common import configure_logging
 from dstack._internal.core.errors import CLIError
 from dstack.api import Client
 
@@ -52,9 +53,16 @@ class BaseCommand(ABC):
 
     @abstractmethod
     def _command(self, args: argparse.Namespace):
+        self._configure_logging()
         if not self.ACCEPT_EXTRA_ARGS and args.extra_args:
             raise CLIError(f"Unrecognized arguments: {shlex.join(args.extra_args)}")
 
+    def _configure_logging(self) -> None:
+        """
+        Override this method to configure command-specific logging
+        """
+        configure_logging()
+
 
 class APIBaseCommand(BaseCommand):
     api: Client

dstack/_internal/cli/commands/project.py

@@ -1,11 +1,12 @@
 import argparse
+from typing import Any, Union
 
 from requests import HTTPError
 from rich.table import Table
 
 import dstack.api.server
 from dstack._internal.cli.commands import BaseCommand
-from dstack._internal.cli.utils.common import confirm_ask, console
+from dstack._internal.cli.utils.common import add_row_from_dict, confirm_ask, console
 from dstack._internal.core.errors import ClientError, CLIError
 from dstack._internal.core.services.configs import ConfigManager
 from dstack._internal.utils.logging import get_logger
@@ -58,6 +59,10 @@ class ProjectCommand(BaseCommand):
         # List subcommand
         list_parser = subparsers.add_parser("list", help="List configured projects")
         list_parser.set_defaults(subfunc=self._list)
+        for parser in [self._parser, list_parser]:
+            parser.add_argument(
+                "-v", "--verbose", action="store_true", help="Show more information"
+            )
 
         # Set default subcommand
         set_default_parser = subparsers.add_parser("set-default", help="Set default project")
@@ -122,30 +127,32 @@ class ProjectCommand(BaseCommand):
         table = Table(box=None)
         table.add_column("PROJECT", style="bold", no_wrap=True)
         table.add_column("URL", style="grey58")
-        table.add_column("USER", style="grey58")
+        if args.verbose:
+            table.add_column("USER", style="grey58")
         table.add_column("DEFAULT", justify="center")
 
         for project_config in config_manager.list_project_configs():
             project_name = project_config.name
             is_default = project_name == default_project.name if default_project else False
-
-            # Get username from API
-            try:
-                api_client = dstack.api.server.APIClient(
-                    base_url=project_config.url, token=project_config.token
-                )
-                user_info = api_client.users.get_my_user()
-                username = user_info.username
-            except ClientError:
-                username = "(invalid token)"
-
-            table.add_row(
-                project_name,
-                project_config.url,
-                username,
-                "✓" if is_default else "",
-                style="bold" if is_default else None,
-            )
+            row: dict[Union[str, int], Any] = {
+                "PROJECT": project_name,
+                "URL": project_config.url,
+                "DEFAULT": "✓" if is_default else "",
+            }
+
+            if args.verbose:
+                # Get username from API
+                try:
+                    api_client = dstack.api.server.APIClient(
+                        base_url=project_config.url, token=project_config.token
+                    )
+                    user_info = api_client.users.get_my_user()
+                    username = user_info.username
+                except ClientError:
+                    username = "(invalid token)"
+                row["USER"] = username
+
+            add_row_from_dict(table, row, style="bold" if is_default else None)
 
         console.print(table)
 

dstack/_internal/cli/commands/server.py

@@ -82,3 +82,8 @@ class ServerCommand(BaseCommand):
             log_level=uvicorn_log_level,
             workers=1,
         )
+
+    def _configure_logging(self) -> None:
+        # Server logging is configured in the FastAPI lifespan function.
+        # No need to configure CLI logging.
+        pass

dstack/_internal/cli/main.py

@@ -22,7 +22,7 @@ from dstack._internal.cli.commands.server import ServerCommand
 from dstack._internal.cli.commands.stats import StatsCommand
 from dstack._internal.cli.commands.stop import StopCommand
 from dstack._internal.cli.commands.volume import VolumeCommand
-from dstack._internal.cli.utils.common import _colors, configure_logging, console
+from dstack._internal.cli.utils.common import _colors, console
 from dstack._internal.cli.utils.updates import check_for_updates
 from dstack._internal.core.errors import ClientError, CLIError, ConfigurationError, SSHError
 from dstack._internal.core.services.ssh.client import get_ssh_client_info
@@ -39,8 +39,6 @@ def main():
     RichHelpFormatter.styles["argparse.groups"] = "bold grey74"
     RichHelpFormatter.styles["argparse.text"] = "grey74"
 
-    configure_logging()
-
     parser = argparse.ArgumentParser(
         description=(
             "Not sure where to start?"

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

@@ -24,6 +24,7 @@ from dstack._internal.core.backends.base.compute import (
     ComputeWithMultinodeSupport,
     ComputeWithPlacementGroupSupport,
     ComputeWithPrivateGatewaySupport,
+    ComputeWithPrivilegedSupport,
     ComputeWithReservationSupport,
     ComputeWithVolumeSupport,
     generate_unique_gateway_instance_name,
@@ -90,6 +91,7 @@ def _ec2client_cache_methodkey(self, ec2_client, *args, **kwargs):
 class AWSCompute(
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     ComputeWithMultinodeSupport,
     ComputeWithReservationSupport,
     ComputeWithPlacementGroupSupport,

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

@@ -43,6 +43,7 @@ from dstack._internal.core.backends.base.compute import (
     ComputeWithCreateInstanceSupport,
     ComputeWithGatewaySupport,
     ComputeWithMultinodeSupport,
+    ComputeWithPrivilegedSupport,
     generate_unique_gateway_instance_name,
     generate_unique_instance_name,
     get_gateway_user_data,
@@ -78,6 +79,7 @@ CONFIGURABLE_DISK_SIZE = Range[Memory](min=Memory.parse("30GB"), max=Memory.pars
 class AzureCompute(
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     ComputeWithMultinodeSupport,
     ComputeWithGatewaySupport,
     Compute,

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

@@ -5,14 +5,16 @@ import string
 import threading
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
+from enum import Enum
 from functools import lru_cache
 from pathlib import Path
-from typing import Callable, Dict, List, Literal, Optional
+from typing import Callable, Dict, List, Optional
 
 import git
 import requests
 import yaml
 from cachetools import TTLCache, cachedmethod
+from gpuhunt import CPUArchitecture
 
 from dstack._internal import settings
 from dstack._internal.core.backends.base.offers import filter_offers_by_requirements
@@ -65,7 +67,21 @@ NVIDIA_GPUS_REQUIRING_PROPRIETARY_KERNEL_MODULES = frozenset(
     ]
 )
 
-GoArchType = Literal["amd64", "arm64"]
+
+class GoArchType(str, Enum):
+    """
+    A subset of GOARCH values
+    """
+
+    AMD64 = "amd64"
+    ARM64 = "arm64"
+
+    def to_cpu_architecture(self) -> CPUArchitecture:
+        if self == self.AMD64:
+            return CPUArchitecture.X86
+        if self == self.ARM64:
+            return CPUArchitecture.ARM
+        assert False, self
 
 
 class Compute(ABC):
@@ -304,6 +320,15 @@ class ComputeWithCreateInstanceSupport(ABC):
             ]
 
 
+class ComputeWithPrivilegedSupport:
+    """
+    Must be subclassed to support runs with `privileged: true`.
+    All VM-based Computes (that is, Computes that use the shim) should subclass this mixin.
+    """
+
+    pass
+
+
 class ComputeWithMultinodeSupport:
     """
     Must be subclassed to support multinode tasks and cluster fleets.
@@ -704,14 +729,14 @@ def normalize_arch(arch: Optional[str] = None) -> GoArchType:
     If the arch is not specified, falls back to `amd64`.
     """
     if not arch:
-        return "amd64"
+        return GoArchType.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"
+        return GoArchType.AMD64
     if arch_lower.startswith("arm") or arch_lower.startswith("aarch"):
-        return "arm64"
+        return GoArchType.ARM64
     raise ValueError(f"Unsupported architecture: {arch}")
 
 
@@ -727,8 +752,7 @@ def get_dstack_runner_download_url(arch: Optional[str] = None) -> str:
             "/{version}/binaries/dstack-runner-linux-{arch}"
         )
     version = get_dstack_runner_version()
-    arch = normalize_arch(arch)
-    return url_template.format(version=version, arch=arch)
+    return url_template.format(version=version, arch=normalize_arch(arch).value)
 
 
 def get_dstack_shim_download_url(arch: Optional[str] = None) -> str:
@@ -743,8 +767,7 @@ def get_dstack_shim_download_url(arch: Optional[str] = None) -> str:
             "/{version}/binaries/dstack-shim-linux-{arch}"
         )
     version = get_dstack_runner_version()
-    arch = normalize_arch(arch)
-    return url_template.format(version=version, arch=arch)
+    return url_template.format(version=version, arch=normalize_arch(arch).value)
 
 
 def get_setup_cloud_instance_commands(

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

@@ -22,6 +22,7 @@ from dstack._internal.utils.common import get_or_error
 SUPPORTED_GPUHUNT_FLAGS = [
     "oci-spot",
     "lambda-arm",
+    "gcp-a4",
 ]
 
 

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

@@ -4,6 +4,7 @@ from dstack._internal.core.backends.base.compute import (
     Compute,
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     get_shim_commands,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
@@ -27,6 +28,7 @@ logger = get_logger(__name__)
 class CloudRiftCompute(
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     Compute,
 ):
     def __init__(self, config: CloudRiftConfig):

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

@@ -6,6 +6,7 @@ from dstack._internal.core.backends.base.backend import Compute
 from dstack._internal.core.backends.base.compute import (
     ComputeWithCreateInstanceSupport,
     ComputeWithFilteredOffersCached,
+    ComputeWithPrivilegedSupport,
     generate_unique_instance_name,
     get_shim_commands,
 )
@@ -32,6 +33,7 @@ MAX_RESOURCE_NAME_LEN = 30
 class CudoCompute(
     ComputeWithFilteredOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     Compute,
 ):
     def __init__(self, config: CudoConfig):

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

@@ -8,6 +8,7 @@ from dstack._internal.core.backends.base.backend import Compute
 from dstack._internal.core.backends.base.compute import (
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     generate_unique_instance_name,
     get_shim_commands,
 )
@@ -39,6 +40,7 @@ CONFIGURABLE_DISK_SIZE = Range[Memory](min=IMAGE_SIZE, max=None)
 class DataCrunchCompute(
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     Compute,
 ):
     def __init__(self, config: DataCrunchConfig):

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

@@ -7,6 +7,7 @@ from dstack._internal.core.backends.base.backend import Compute
 from dstack._internal.core.backends.base.compute import (
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     generate_unique_instance_name,
     get_user_data,
 )
@@ -40,6 +41,7 @@ DOCKER_INSTALL_COMMANDS = [
 class BaseDigitalOceanCompute(
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     Compute,
 ):
     def __init__(self, config: BaseDigitalOceanConfig, api_url: str, type: BackendType):

dstack/_internal/core/backends/features.py

@@ -4,6 +4,7 @@ from dstack._internal.core.backends.base.compute import (
     ComputeWithMultinodeSupport,
     ComputeWithPlacementGroupSupport,
     ComputeWithPrivateGatewaySupport,
+    ComputeWithPrivilegedSupport,
     ComputeWithReservationSupport,
     ComputeWithVolumeSupport,
 )
@@ -38,6 +39,10 @@ BACKENDS_WITH_CREATE_INSTANCE_SUPPORT = _get_backends_with_compute_feature(
     configurator_classes=_configurator_classes,
     compute_feature_class=ComputeWithCreateInstanceSupport,
 )
+BACKENDS_WITH_PRIVILEGED_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithPrivilegedSupport,
+)
 BACKENDS_WITH_MULTINODE_SUPPORT = [BackendType.REMOTE] + _get_backends_with_compute_feature(
     configurator_classes=_configurator_classes,
     compute_feature_class=ComputeWithMultinodeSupport,

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

@@ -23,6 +23,7 @@ from dstack._internal.core.backends.base.compute import (
     ComputeWithMultinodeSupport,
     ComputeWithPlacementGroupSupport,
     ComputeWithPrivateGatewaySupport,
+    ComputeWithPrivilegedSupport,
     ComputeWithVolumeSupport,
     generate_unique_gateway_instance_name,
     generate_unique_instance_name,
@@ -90,6 +91,7 @@ class GCPVolumeDiskBackendData(CoreModel):
 class GCPCompute(
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     ComputeWithMultinodeSupport,
     ComputeWithPlacementGroupSupport,
     ComputeWithGatewaySupport,
@@ -111,8 +113,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)
+        self._usable_subnets_cache_lock = threading.Lock()
+        self._usable_subnets_cache = TTLCache(maxsize=1, ttl=120)
 
     def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
         regions = get_or_error(self.config.regions)
@@ -203,12 +205,12 @@ class GCPCompute(
         disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024)
         # Choose any usable subnet in a VPC.
         # Configuring a specific subnet per region is not supported yet.
-        subnetwork = _get_vpc_subnet(
-            subnetworks_client=self.subnetworks_client,
-            config=self.config,
+        subnetwork = self._get_vpc_subnet(instance_offer.region)
+        extra_subnets = self._get_extra_subnets(
             region=instance_offer.region,
+            instance_type_name=instance_offer.instance.name,
         )
-        extra_subnets = self._get_extra_subnets(
+        roce_subnets = self._get_roce_subnets(
             region=instance_offer.region,
             instance_type_name=instance_offer.instance.name,
         )
@@ -330,6 +332,7 @@ class GCPCompute(
                 network=self.config.vpc_resource_name,
                 subnetwork=subnetwork,
                 extra_subnetworks=extra_subnets,
+                roce_subnetworks=roce_subnets,
                 allocate_public_ip=allocate_public_ip,
                 placement_policy=placement_policy,
             )
@@ -339,6 +342,13 @@ class GCPCompute(
                 # If the request succeeds, we'll probably timeout and update_provisioning_data() will get hostname.
                 operation = self.instances_client.insert(request=request)
                 gcp_resources.wait_for_extended_operation(operation, timeout=30)
+            except google.api_core.exceptions.BadRequest as e:
+                if "Network profile only allows resource creation in location" in e.message:
+                    # A hack to find the correct RoCE VPC zone by trial and error.
+                    # Could be better to find it via the API.
+                    logger.debug("Got GCP error when provisioning a VM: %s", e)
+                    continue
+                raise
             except (
                 google.api_core.exceptions.ServiceUnavailable,
                 google.api_core.exceptions.NotFound,
@@ -487,11 +497,7 @@ class GCPCompute(
         )
         # Choose any usable subnet in a VPC.
         # Configuring a specific subnet per region is not supported yet.
-        subnetwork = _get_vpc_subnet(
-            subnetworks_client=self.subnetworks_client,
-            config=self.config,
-            region=configuration.region,
-        )
+        subnetwork = self._get_vpc_subnet(configuration.region)
 
         labels = {
             "owner": "dstack",
@@ -793,10 +799,6 @@ 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,
@@ -808,15 +810,16 @@ class GCPCompute(
             subnets_num = 8
         elif instance_type_name in ["a3-edgegpu-8g", "a3-highgpu-8g"]:
             subnets_num = 4
+        elif instance_type_name == "a4-highgpu-8g":
+            subnets_num = 1  # 1 main + 1 extra + 8 RoCE
         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,
+                usable_subnets=self._list_usable_subnets(),
             )
             vpc_resource_name = gcp_resources.vpc_name_to_vpc_resource_name(
                 project_id=self.config.vpc_project_id or self.config.project_id,
@@ -825,6 +828,58 @@ class GCPCompute(
             extra_subnets.append((vpc_resource_name, subnet))
         return extra_subnets
 
+    def _get_roce_subnets(
+        self,
+        region: str,
+        instance_type_name: str,
+    ) -> List[Tuple[str, str]]:
+        if not self.config.roce_vpcs:
+            return []
+        if instance_type_name == "a4-highgpu-8g":
+            nics_num = 8
+        else:
+            return []
+        roce_vpc = self.config.roce_vpcs[0]  # roce_vpcs is validated to have at most 1 item
+        subnets = gcp_resources.get_vpc_subnets(
+            vpc_name=roce_vpc,
+            region=region,
+            usable_subnets=self._list_usable_subnets(),
+        )
+        if len(subnets) < nics_num:
+            raise ComputeError(
+                f"{instance_type_name} requires {nics_num} RoCE subnets,"
+                f" but only {len(subnets)} are available in VPC {roce_vpc}"
+            )
+        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=roce_vpc,
+        )
+        nic_subnets = []
+        for subnet in subnets[:nics_num]:
+            nic_subnets.append((vpc_resource_name, subnet))
+        return nic_subnets
+
+    @cachedmethod(
+        cache=lambda self: self._usable_subnets_cache,
+        lock=lambda self: self._usable_subnets_cache_lock,
+    )
+    def _list_usable_subnets(self) -> list[compute_v1.UsableSubnetwork]:
+        # To avoid hitting the `ListUsable requests per minute` system limit, we fetch all subnets
+        # at once and cache them
+        return gcp_resources.list_project_usable_subnets(
+            subnetworks_client=self.subnetworks_client,
+            project_id=self.config.vpc_project_id or self.config.project_id,
+        )
+
+    def _get_vpc_subnet(self, region: str) -> Optional[str]:
+        if self.config.vpc_name is None:
+            return None
+        return gcp_resources.get_vpc_subnet_or_error(
+            vpc_name=self.config.vpc_name,
+            region=region,
+            usable_subnets=self._list_usable_subnets(),
+        )
+
 
 def _supported_instances_and_zones(
     regions: List[str],
@@ -867,8 +922,8 @@ def _has_gpu_quota(quotas: Dict[str, float], resources: Resources) -> bool:
     gpu = resources.gpus[0]
     if _is_tpu(gpu.name):
         return True
-    if gpu.name == "H100":
-        # H100 and H100_MEGA quotas are not returned by `regions_client.list`
+    if gpu.name in ["B200", "H100"]:
+        # B200, H100 and H100_MEGA quotas are not returned by `regions_client.list`
         return True
     quota_name = f"NVIDIA_{gpu.name}_GPUS"
     if gpu.name == "A100" and gpu.memory_mib == 80 * 1024:
@@ -889,21 +944,6 @@ def _unique_instance_name(instance: InstanceType) -> str:
     return f"{name}-{gpu.name}-{gpu.memory_mib}"
 
 
-def _get_vpc_subnet(
-    subnetworks_client: compute_v1.SubnetworksClient,
-    config: GCPConfig,
-    region: str,
-) -> Optional[str]:
-    if config.vpc_name is None:
-        return None
-    return gcp_resources.get_vpc_subnet_or_error(
-        subnetworks_client=subnetworks_client,
-        vpc_project_id=config.vpc_project_id or config.project_id,
-        vpc_name=config.vpc_name,
-        region=region,
-    )
-
-
 @dataclass
 class GCPImage:
     id: str

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

@@ -202,5 +202,5 @@ class GCPConfigurator(
             )
         except BackendError as e:
             raise ServerClientError(e.args[0])
-        # Not checking config.extra_vpc so that users are not required to configure subnets for all regions
+        # Not checking config.extra_vpcs and config.roce_vpcs so that users are not required to configure subnets for all regions
         # but only for regions they intend to use. Validation will be done on provisioning.

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

@@ -41,11 +41,24 @@ class GCPBackendConfig(CoreModel):
         Optional[List[str]],
         Field(
             description=(
-                "The names of additional VPCs used for GPUDirect. Specify eight VPCs to maximize bandwidth."
+                "The names of additional VPCs used for multi-NIC instances, such as those that support GPUDirect."
+                " Specify eight VPCs to maximize bandwidth in clusters with eight-GPU instances."
                 " Each VPC must have a subnet and a firewall rule allowing internal traffic across all subnets"
             )
         ),
     ] = None
+    roce_vpcs: Annotated[
+        Optional[List[str]],
+        Field(
+            description=(
+                "The names of additional VPCs with the RoCE network profile."
+                " Used for RDMA on GPU instances that support the MRDMA interface type."
+                " A VPC should have eight subnets to maximize the bandwidth in clusters"
+                " with eight-GPU instances."
+            ),
+            max_items=1,  # The currently supported instance types only need one VPC with eight subnets.
+        ),
+    ] = None
     vpc_project_id: Annotated[
         Optional[str],
         Field(description="The shared VPC hosted project ID. Required for shared VPC only"),