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

dstack/_internal/cli/services/args.py

@@ -19,8 +19,8 @@ def port_mapping(v: str) -> PortMapping:
     return PortMapping.parse(v)
 
 
-def cpu_spec(v: str) -> resources.Range[int]:
-    return parse_obj_as(resources.Range[int], v)
+def cpu_spec(v: str) -> dict:
+    return resources.CPUSpec.parse(v)
 
 
 def memory_spec(v: str) -> resources.Range[resources.Memory]:

dstack/_internal/cli/services/configurators/fleet.py

@@ -20,6 +20,7 @@ from dstack._internal.cli.utils.rich import MultiItemStatus
 from dstack._internal.core.errors import (
     CLIError,
     ConfigurationError,
+    MethodNotAllowedError,
     ResourceNotExistsError,
     ServerClientError,
     URLNotFoundError,
@@ -321,7 +322,7 @@ def _print_plan_header(plan: FleetPlan):
             offer.instance.name,
             resources.pretty_format(),
             "yes" if resources.spot else "no",
-            f"${offer.price:g}",
+            f"${offer.price:3f}".rstrip("0").rstrip("."),
             availability,
             style=None if index == 1 else "secondary",
         )
@@ -367,7 +368,7 @@ def _apply_plan(api: Client, plan: FleetPlan) -> Fleet:
             project_name=api.project,
             plan=plan,
         )
-    except URLNotFoundError:
+    except (URLNotFoundError, MethodNotAllowedError):
         # TODO: Remove in 0.20
         return api.client.fleets.create(
             project_name=api.project,

dstack/_internal/cli/services/configurators/run.py

@@ -6,9 +6,10 @@ from pathlib import Path
 from typing import Dict, List, Optional, Set, Tuple
 
 import gpuhunt
+from pydantic import parse_obj_as
 
 import dstack._internal.core.models.resources as resources
-from dstack._internal.cli.services.args import disk_spec, gpu_spec, port_mapping
+from dstack._internal.cli.services.args import cpu_spec, disk_spec, gpu_spec, port_mapping
 from dstack._internal.cli.services.configurators.base import (
     ApplyEnvVarsConfiguratorMixin,
     BaseApplyConfigurator,
@@ -39,6 +40,7 @@ from dstack._internal.core.models.configurations import (
     TaskConfiguration,
 )
 from dstack._internal.core.models.repos.base import Repo
+from dstack._internal.core.models.resources import CPUSpec
 from dstack._internal.core.models.runs import JobSubmission, JobTerminationReason, RunStatus
 from dstack._internal.core.services.configs import ConfigManager
 from dstack._internal.core.services.diff import diff_models
@@ -52,7 +54,7 @@ from dstack.api.utils import load_profile
 _KNOWN_AMD_GPUS = {gpu.name.lower() for gpu in gpuhunt.KNOWN_AMD_GPUS}
 _KNOWN_NVIDIA_GPUS = {gpu.name.lower() for gpu in gpuhunt.KNOWN_NVIDIA_GPUS}
 _KNOWN_TPU_VERSIONS = {gpu.name.lower() for gpu in gpuhunt.KNOWN_TPUS}
-
+_KNOWN_TENSTORRENT_GPUS = {gpu.name.lower() for gpu in gpuhunt.KNOWN_TENSTORRENT_ACCELERATORS}
 _BIND_ADDRESS_ARG = "bind_address"
 
 logger = get_logger(__name__)
@@ -72,6 +74,7 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
     ):
         self.apply_args(conf, configurator_args, unknown_args)
         self.validate_gpu_vendor_and_image(conf)
+        self.validate_cpu_arch_and_image(conf)
         if repo is None:
             repo = self.api.repos.load(Path.cwd())
         config_manager = ConfigManager()
@@ -289,6 +292,14 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
             default=default_max_offers,
         )
         cls.register_env_args(configuration_group)
+        configuration_group.add_argument(
+            "--cpu",
+            type=cpu_spec,
+            help="Request CPU for the run. "
+            "The format is [code]ARCH[/]:[code]COUNT[/] (all parts are optional)",
+            dest="cpu_spec",
+            metavar="SPEC",
+        )
         configuration_group.add_argument(
             "--gpu",
             type=gpu_spec,
@@ -310,6 +321,8 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
         apply_profile_args(args, conf)
         if args.run_name:
             conf.name = args.run_name
+        if args.cpu_spec:
+            conf.resources.cpu = resources.CPUSpec.parse_obj(args.cpu_spec)
         if args.gpu_spec:
             conf.resources.gpu = resources.GPUSpec.parse_obj(args.gpu_spec)
         if args.disk_spec:
@@ -342,7 +355,7 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
 
     def validate_gpu_vendor_and_image(self, conf: BaseRunConfiguration) -> None:
         """
-        Infers `resources.gpu.vendor` if not set, requires `image` if the vendor is AMD.
+        Infers and sets `resources.gpu.vendor` if not set, requires `image` if the vendor is AMD.
         """
         gpu_spec = conf.resources.gpu
         if gpu_spec is None:
@@ -350,6 +363,7 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
         if gpu_spec.count.max == 0:
             return
         has_amd_gpu: bool
+        has_tt_gpu: bool
         vendor = gpu_spec.vendor
         if vendor is None:
             names = gpu_spec.name
@@ -362,6 +376,8 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
                         vendors.add(gpuhunt.AcceleratorVendor.NVIDIA)
                     elif name in _KNOWN_AMD_GPUS:
                         vendors.add(gpuhunt.AcceleratorVendor.AMD)
+                    elif name in _KNOWN_TENSTORRENT_GPUS:
+                        vendors.add(gpuhunt.AcceleratorVendor.TENSTORRENT)
                     else:
                         maybe_tpu_version, _, maybe_tpu_cores = name.partition("-")
                         if maybe_tpu_version in _KNOWN_TPU_VERSIONS and maybe_tpu_cores.isdigit():
@@ -380,15 +396,45 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
                 # to execute a run on an instance with an AMD accelerator with a default
                 # CUDA image, not a big deal.
                 has_amd_gpu = gpuhunt.AcceleratorVendor.AMD in vendors
+                has_tt_gpu = gpuhunt.AcceleratorVendor.TENSTORRENT in vendors
             else:
                 # If neither gpu.vendor nor gpu.name is set, assume Nvidia.
                 vendor = gpuhunt.AcceleratorVendor.NVIDIA
                 has_amd_gpu = False
+                has_tt_gpu = False
             gpu_spec.vendor = vendor
         else:
             has_amd_gpu = vendor == gpuhunt.AcceleratorVendor.AMD
+            has_tt_gpu = vendor == gpuhunt.AcceleratorVendor.TENSTORRENT
         if has_amd_gpu and conf.image is None:
-            raise ConfigurationError("`image` is required if `resources.gpu.vendor` is AMD.")
+            raise ConfigurationError("`image` is required if `resources.gpu.vendor` is `amd`")
+        if has_tt_gpu and conf.image is None:
+            raise ConfigurationError(
+                "`image` is required if `resources.gpu.vendor` is `tenstorrent`"
+            )
+
+    def validate_cpu_arch_and_image(self, conf: BaseRunConfiguration) -> None:
+        """
+        Infers `resources.cpu.arch` if not set, requires `image` if the architecture is ARM.
+        """
+        # TODO: Remove in 0.20. Use conf.resources.cpu directly
+        cpu_spec = parse_obj_as(CPUSpec, conf.resources.cpu)
+        arch = cpu_spec.arch
+        if arch is None:
+            gpu_spec = conf.resources.gpu
+            if (
+                gpu_spec is not None
+                and gpu_spec.vendor in [None, gpuhunt.AcceleratorVendor.NVIDIA]
+                and gpu_spec.name
+                and any(map(gpuhunt.is_nvidia_superchip, gpu_spec.name))
+            ):
+                arch = gpuhunt.CPUArchitecture.ARM
+            else:
+                arch = gpuhunt.CPUArchitecture.X86
+        # NOTE: We don't set the inferred resources.cpu.arch for compatibility with older servers.
+        # Servers with ARM support set the arch using the same logic.
+        if arch == gpuhunt.CPUArchitecture.ARM and conf.image is None:
+            raise ConfigurationError("`image` is required if `resources.cpu.arch` is `arm`")
 
 
 class RunWithPortsConfigurator(BaseRunConfigurator):

dstack/_internal/cli/utils/fleet.py

@@ -79,7 +79,9 @@ def get_fleets_table(
                 "BACKEND": backend,
                 "REGION": region,
                 "RESOURCES": resources,
-                "PRICE": f"${instance.price:.4}" if instance.price is not None else "",
+                "PRICE": f"${instance.price:.4f}".rstrip("0").rstrip(".")
+                if instance.price is not None
+                else "",
                 "STATUS": status,
                 "CREATED": format_date(instance.created),
                 "ERROR": error,

dstack/_internal/cli/utils/run.py

@@ -1,3 +1,4 @@
+import shutil
 from typing import Any, Dict, List, Optional, Union
 
 from rich.markup import escape
@@ -36,7 +37,7 @@ def print_run_plan(
 
     req = job_plan.job_spec.requirements
     pretty_req = req.pretty_format(resources_only=True)
-    max_price = f"${req.max_price:g}" if req.max_price else "-"
+    max_price = f"${req.max_price:3f}".rstrip("0").rstrip(".") if req.max_price else "-"
     max_duration = (
         format_pretty_duration(job_plan.job_spec.max_duration)
         if job_plan.job_spec.max_duration
@@ -94,14 +95,12 @@ def print_run_plan(
             props.add_row(th("Inactivity duration"), inactivity_duration)
     props.add_row(th("Reservation"), run_spec.configuration.reservation or "-")
 
-    offers = Table(box=None)
+    offers = Table(box=None, expand=shutil.get_terminal_size(fallback=(120, 40)).columns <= 110)
     offers.add_column("#")
-    offers.add_column("BACKEND")
-    offers.add_column("REGION")
-    offers.add_column("INSTANCE TYPE")
-    offers.add_column("RESOURCES")
-    offers.add_column("SPOT")
-    offers.add_column("PRICE")
+    offers.add_column("BACKEND", style="grey58", ratio=2)
+    offers.add_column("RESOURCES", ratio=4)
+    offers.add_column("INSTANCE TYPE", style="grey58", no_wrap=True, ratio=2)
+    offers.add_column("PRICE", style="grey58", ratio=1)
     offers.add_column()
 
     job_plan.offers = job_plan.offers[:max_offers] if max_offers else job_plan.offers
@@ -122,14 +121,12 @@ def print_run_plan(
             instance += f" ({offer.blocks}/{offer.total_blocks})"
         offers.add_row(
             f"{i}",
-            offer.backend.replace("remote", "ssh"),
-            offer.region,
+            offer.backend.replace("remote", "ssh") + " (" + offer.region + ")",
+            r.pretty_format(include_spot=True),
             instance,
-            r.pretty_format(),
-            "yes" if r.spot else "no",
-            f"${offer.price:g}",
+            f"${offer.price:.4f}".rstrip("0").rstrip("."),
             availability,
-            style=None if i == 1 else "secondary",
+            style=None if i == 1 or not include_run_properties else "secondary",
         )
     if job_plan.total_offers > len(job_plan.offers):
         offers.add_row("", "...", style="secondary")
@@ -141,7 +138,8 @@ def print_run_plan(
         if job_plan.total_offers > len(job_plan.offers):
             console.print(
                 f"[secondary] Shown {len(job_plan.offers)} of {job_plan.total_offers} offers, "
-                f"${job_plan.max_price:g} max[/]"
+                f"${job_plan.max_price:3f}".rstrip("0").rstrip(".")
+                + "max[/]"
             )
         console.print()
     else:
@@ -151,19 +149,18 @@ def print_run_plan(
 def get_runs_table(
     runs: List[Run], verbose: bool = False, format_date: DateFormatter = pretty_date
 ) -> Table:
-    table = Table(box=None)
-    table.add_column("NAME", style="bold", no_wrap=True)
-    table.add_column("BACKEND", style="grey58")
+    table = Table(box=None, expand=shutil.get_terminal_size(fallback=(120, 40)).columns <= 110)
+    table.add_column("NAME", style="bold", no_wrap=True, ratio=2)
+    table.add_column("BACKEND", style="grey58", ratio=2)
+    table.add_column("RESOURCES", ratio=3 if not verbose else 2)
     if verbose:
-        table.add_column("INSTANCE", no_wrap=True)
-    table.add_column("RESOURCES")
+        table.add_column("INSTANCE", no_wrap=True, ratio=1)
+        table.add_column("RESERVATION", no_wrap=True, ratio=1)
+    table.add_column("PRICE", style="grey58", ratio=1)
+    table.add_column("STATUS", no_wrap=True, ratio=1)
+    table.add_column("SUBMITTED", style="grey58", no_wrap=True, ratio=1)
     if verbose:
-        table.add_column("RESERVATION", no_wrap=True)
-    table.add_column("PRICE", no_wrap=True)
-    table.add_column("STATUS", no_wrap=True)
-    table.add_column("SUBMITTED", style="grey58", no_wrap=True)
-    if verbose:
-        table.add_column("ERROR", no_wrap=True)
+        table.add_column("ERROR", no_wrap=True, ratio=2)
 
     for run in runs:
         run_error = _get_run_error(run)
@@ -202,10 +199,10 @@ def get_runs_table(
                 job_row.update(
                     {
                         "BACKEND": f"{jpd.backend.value.replace('remote', 'ssh')} ({jpd.region})",
-                        "INSTANCE": instance,
                         "RESOURCES": resources.pretty_format(include_spot=True),
+                        "INSTANCE": instance,
                         "RESERVATION": jpd.reservation,
-                        "PRICE": f"${jpd.price:.4}",
+                        "PRICE": f"${jpd.price:.4f}".rstrip("0").rstrip("."),
                     }
                 )
             if len(run.jobs) == 1:

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

@@ -159,6 +159,7 @@ class AWSCompute(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         project_name = instance_config.project_name
         ec2_resource = self.session.resource("ec2", region_name=instance_offer.region)
@@ -248,7 +249,7 @@ class AWSCompute(
                         spot=instance_offer.instance.resources.spot,
                         subnet_id=subnet_id,
                         allocate_public_ip=allocate_public_ip,
-                        placement_group_name=instance_config.placement_group_name,
+                        placement_group_name=placement_group.name if placement_group else None,
                         enable_efa=enable_efa,
                         max_efa_interfaces=max_efa_interfaces,
                         reservation_id=instance_config.reservation,
@@ -291,6 +292,7 @@ class AWSCompute(
     def create_placement_group(
         self,
         placement_group: PlacementGroup,
+        master_instance_offer: InstanceOffer,
     ) -> PlacementGroupProvisioningData:
         ec2_client = self.session.client("ec2", region_name=placement_group.configuration.region)
         logger.debug("Creating placement group %s...", placement_group.name)
@@ -323,6 +325,16 @@ class AWSCompute(
                 raise e
         logger.debug("Deleted placement group %s", placement_group.name)
 
+    def is_suitable_placement_group(
+        self,
+        placement_group: PlacementGroup,
+        instance_offer: InstanceOffer,
+    ) -> bool:
+        return (
+            placement_group.configuration.backend == BackendType.AWS
+            and placement_group.configuration.region == instance_offer.region
+        )
+
     def create_gateway(
         self,
         configuration: GatewayComputeConfiguration,

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

@@ -62,6 +62,7 @@ from dstack._internal.core.models.instances import (
     InstanceOfferWithAvailability,
     InstanceType,
 )
+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
@@ -109,6 +110,7 @@ class AzureCompute(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         instance_name = generate_unique_instance_name(
             instance_config, max_length=azure_resources.MAX_RESOURCE_NAME_LEN
@@ -136,6 +138,10 @@ class AzureCompute(
             location=location,
         )
 
+        managed_identity_resource_group, managed_identity_name = parse_vm_managed_identity(
+            self.config.vm_managed_identity
+        )
+
         base_tags = {
             "owner": "dstack",
             "dstack_project": instance_config.project_name,
@@ -159,7 +165,8 @@ class AzureCompute(
             network_security_group=network_security_group,
             network=network,
             subnet=subnet,
-            managed_identity=None,
+            managed_identity_name=managed_identity_name,
+            managed_identity_resource_group=managed_identity_resource_group,
             image_reference=_get_image_ref(
                 compute_client=self._compute_client,
                 location=location,
@@ -255,7 +262,8 @@ class AzureCompute(
             network_security_group=network_security_group,
             network=network,
             subnet=subnet,
-            managed_identity=None,
+            managed_identity_name=None,
+            managed_identity_resource_group=None,
             image_reference=_get_gateway_image_ref(),
             vm_size="Standard_B1ms",
             instance_name=instance_name,
@@ -338,6 +346,21 @@ def get_resource_group_network_subnet_or_error(
     return resource_group, network_name, subnet_name
 
 
+def parse_vm_managed_identity(
+    vm_managed_identity: Optional[str],
+) -> Tuple[Optional[str], Optional[str]]:
+    if vm_managed_identity is None:
+        return None, None
+    try:
+        resource_group, managed_identity = vm_managed_identity.split("/")
+        return resource_group, managed_identity
+    except Exception:
+        raise ComputeError(
+            "`vm_managed_identity` specified in incorrect format."
+            " Supported format: 'managedIdentityResourceGroup/managedIdentityName'"
+        )
+
+
 def _parse_config_vpc_id(vpc_id: str) -> Tuple[str, str]:
     resource_group, network_name = vpc_id.split("/")
     return resource_group, network_name
@@ -466,7 +489,8 @@ def _launch_instance(
     network_security_group: str,
     network: str,
     subnet: str,
-    managed_identity: Optional[str],
+    managed_identity_name: Optional[str],
+    managed_identity_resource_group: Optional[str],
     image_reference: ImageReference,
     vm_size: str,
     instance_name: str,
@@ -488,6 +512,20 @@ def _launch_instance(
         public_ip_address_configuration = VirtualMachinePublicIPAddressConfiguration(
             name="public_ip_config",
         )
+    managed_identity = None
+    if managed_identity_name is not None:
+        if managed_identity_resource_group is None:
+            managed_identity_resource_group = resource_group
+        managed_identity = VirtualMachineIdentity(
+            type=ResourceIdentityType.USER_ASSIGNED,
+            user_assigned_identities={
+                azure_utils.get_managed_identity_id(
+                    subscription_id,
+                    managed_identity_resource_group,
+                    managed_identity_name,
+                ): UserAssignedIdentitiesValue(),
+            },
+        )
     try:
         poller = compute_client.virtual_machines.begin_create_or_update(
             resource_group,
@@ -552,16 +590,7 @@ def _launch_instance(
                 ),
                 priority="Spot" if spot else "Regular",
                 eviction_policy="Delete" if spot else None,
-                identity=None
-                if managed_identity is None
-                else VirtualMachineIdentity(
-                    type=ResourceIdentityType.USER_ASSIGNED,
-                    user_assigned_identities={
-                        azure_utils.get_managed_identity_id(
-                            subscription_id, resource_group, managed_identity
-                        ): UserAssignedIdentitiesValue()
-                    },
-                ),
+                identity=managed_identity,
                 user_data=base64.b64encode(user_data.encode()).decode(),
                 tags=tags,
             ),

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

@@ -4,6 +4,7 @@ from typing import List, Optional, Tuple
 
 import azure.core.exceptions
 from azure.core.credentials import TokenCredential
+from azure.mgmt import msi as msi_mgmt
 from azure.mgmt import network as network_mgmt
 from azure.mgmt import resource as resource_mgmt
 from azure.mgmt import subscription as subscription_mgmt
@@ -97,6 +98,7 @@ class AzureConfigurator(Configurator):
         self._check_config_locations(config)
         self._check_config_tags(config)
         self._check_config_resource_group(config=config, credential=credential)
+        self._check_config_vm_managed_identity(config=config, credential=credential)
         self._check_config_vpc(config=config, credential=credential)
 
     def create_backend(
@@ -260,6 +262,25 @@ class AzureConfigurator(Configurator):
                     except BackendError as e:
                         raise ServerClientError(e.args[0])
 
+    def _check_config_vm_managed_identity(
+        self, config: AzureBackendConfigWithCreds, credential: auth.AzureCredential
+    ):
+        try:
+            resource_group, identity_name = compute.parse_vm_managed_identity(
+                config.vm_managed_identity
+            )
+        except BackendError as e:
+            raise ServerClientError(e.args[0])
+        if resource_group is None or identity_name is None:
+            return
+        msi_client = msi_mgmt.ManagedServiceIdentityClient(credential, config.subscription_id)
+        try:
+            msi_client.user_assigned_identities.get(resource_group, identity_name)
+        except azure.core.exceptions.ResourceNotFoundError:
+            raise ServerClientError(
+                f"Managed identity {identity_name} not found in resource group {resource_group}"
+            )
+
     def _set_client_creds_tenant_id(
         self,
         creds: AzureClientCreds,

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

@@ -62,6 +62,15 @@ class AzureBackendConfig(CoreModel):
             )
         ),
     ] = None
+    vm_managed_identity: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "The managed identity to associate with provisioned VMs."
+                " Must have a format `managedIdentityResourceGroup/managedIdentityName`"
+            )
+        ),
+    ] = None
     tags: Annotated[
         Optional[Dict[str, str]],
         Field(description="The tags that will be assigned to resources created by `dstack`"),