dstack 0.19.9__py3-none-any.whl → 0.19.11__py3-none-any.whl

dstack/_internal/cli/commands/config.py

@@ -14,7 +14,7 @@ logger = get_logger(__name__)
 
 class ConfigCommand(BaseCommand):
     NAME = "config"
-    DESCRIPTION = "Configure CLI"
+    DESCRIPTION = "Configure CLI (deprecated; use `dstack project`)"
 
     def _register(self):
         super()._register()

dstack/_internal/cli/commands/metrics.py

@@ -39,8 +39,6 @@ class MetricsCommand(APIBaseCommand):
         run = self.api.runs.get(run_name=args.run_name)
         if run is None:
             raise CLIError(f"Run {args.run_name} not found")
-        if run.status.is_finished():
-            raise CLIError(f"Run {args.run_name} is finished")
         metrics = _get_run_jobs_metrics(api=self.api, run=run)
 
         if not args.watch:
@@ -55,8 +53,6 @@ class MetricsCommand(APIBaseCommand):
                     run = self.api.runs.get(run_name=args.run_name)
                     if run is None:
                         raise CLIError(f"Run {args.run_name} not found")
-                    if run.status.is_finished():
-                        raise CLIError(f"Run {args.run_name} is finished")
                     metrics = _get_run_jobs_metrics(api=self.api, run=run)
         except KeyboardInterrupt:
             pass
@@ -78,11 +74,12 @@ def _get_run_jobs_metrics(api: Client, run: Run) -> List[JobMetrics]:
 def _get_metrics_table(run: Run, metrics: List[JobMetrics]) -> Table:
     table = Table(box=None)
     table.add_column("NAME", style="bold", no_wrap=True)
+    table.add_column("STATUS")
     table.add_column("CPU")
     table.add_column("MEMORY")
     table.add_column("GPU")
 
-    run_row: Dict[Union[str, int], Any] = {"NAME": run.name}
+    run_row: Dict[Union[str, int], Any] = {"NAME": run.name, "STATUS": run.status.value}
     if len(run._run.jobs) != 1:
         add_row_from_dict(table, run_row)
 
@@ -101,9 +98,9 @@ def _get_metrics_table(run: Run, metrics: List[JobMetrics]) -> Table:
             cpu_usage = f"{cpu_usage:.0f}%"
         memory_usage = _get_metric_value(job_metrics, "memory_working_set_bytes")
         if memory_usage is not None:
-            memory_usage = f"{round(memory_usage / 1024 / 1024)}MB"
+            memory_usage = _format_memory(memory_usage, 2)
             if resources is not None:
-                memory_usage += f"/{resources.memory_mib}MB"
+                memory_usage += f"/{_format_memory(resources.memory_mib * 1024 * 1024, 2)}"
         gpu_metrics = ""
         gpus_detected_num = _get_metric_value(job_metrics, "gpus_detected_num")
         if gpus_detected_num is not None:
@@ -113,13 +110,16 @@ def _get_metrics_table(run: Run, metrics: List[JobMetrics]) -> Table:
                 if gpu_memory_usage is not None:
                     if i != 0:
                         gpu_metrics += "\n"
-                    gpu_metrics += f"#{i} {round(gpu_memory_usage / 1024 / 1024)}MB"
+                    gpu_metrics += f"gpu={i} mem={_format_memory(gpu_memory_usage, 2)}"
                     if resources is not None:
-                        gpu_metrics += f"/{resources.gpus[i].memory_mib}MB"
-                    gpu_metrics += f" {gpu_util_percent}% Util"
+                        gpu_metrics += (
+                            f"/{_format_memory(resources.gpus[i].memory_mib * 1024 * 1024, 2)}"
+                        )
+                    gpu_metrics += f" util={gpu_util_percent}%"
 
         job_row: Dict[Union[str, int], Any] = {
             "NAME": f"  replica={job.job_spec.replica_num} job={job.job_spec.job_num}",
+            "STATUS": job.job_submissions[-1].status.value,
             "CPU": cpu_usage or "-",
             "MEMORY": memory_usage or "-",
             "GPU": gpu_metrics or "-",
@@ -136,3 +136,18 @@ def _get_metric_value(job_metrics: JobMetrics, name: str) -> Optional[Any]:
         if metric.name == name:
             return metric.values[-1]
     return None
+
+
+def _format_memory(memory_bytes: int, decimal_places: int) -> str:
+    """See test_format_memory in tests/_internal/cli/commands/test_metrics.py for examples."""
+    memory_mb = memory_bytes / 1024 / 1024
+    if memory_mb >= 1024:
+        value = memory_mb / 1024
+        unit = "GB"
+    else:
+        value = memory_mb
+        unit = "MB"
+
+    if decimal_places == 0:
+        return f"{round(value)}{unit}"
+    return f"{value:.{decimal_places}f}".rstrip("0").rstrip(".") + unit

dstack/_internal/cli/commands/offer.py

@@ -84,6 +84,8 @@ class OfferCommand(APIBaseCommand):
         job_plan = run_plan.job_plans[0]
 
         if args.format == "json":
+            # FIXME: Should use effective_run_spec from run_plan,
+            # since the spec can be changed by the server and plugins
             output = {
                 "project": run_plan.project_name,
                 "user": run_plan.user,

dstack/_internal/cli/commands/project.py

@@ -0,0 +1,161 @@
+import argparse
+
+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.core.errors import ClientError, CLIError
+from dstack._internal.core.services.configs import ConfigManager
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+class ProjectCommand(BaseCommand):
+    NAME = "project"
+    DESCRIPTION = "Manage projects configs"
+
+    def _register(self):
+        super()._register()
+        subparsers = self._parser.add_subparsers(dest="subcommand", help="Command to execute")
+
+        # Add subcommand
+        add_parser = subparsers.add_parser("add", help="Add or update a project config")
+        add_parser.add_argument(
+            "--name", type=str, help="The name of the project to configure", required=True
+        )
+        add_parser.add_argument("--url", type=str, help="Server url", required=True)
+        add_parser.add_argument("--token", type=str, help="User token", required=True)
+        add_parser.add_argument(
+            "-y",
+            "--yes",
+            help="Don't ask for confirmation (e.g. update the config)",
+            action="store_true",
+        )
+        add_parser.add_argument(
+            "-n",
+            "--no",
+            help="Don't ask for confirmation (e.g. do not update the config)",
+            action="store_true",
+        )
+        add_parser.set_defaults(subfunc=self._add)
+
+        # Delete subcommand
+        delete_parser = subparsers.add_parser("delete", help="Delete a project config")
+        delete_parser.add_argument(
+            "--name", type=str, help="The name of the project to delete", required=True
+        )
+        delete_parser.add_argument(
+            "-y",
+            "--yes",
+            help="Don't ask for confirmation",
+            action="store_true",
+        )
+        delete_parser.set_defaults(subfunc=self._delete)
+
+        # List subcommand
+        list_parser = subparsers.add_parser("list", help="List configured projects")
+        list_parser.set_defaults(subfunc=self._list)
+
+        # Set default subcommand
+        set_default_parser = subparsers.add_parser("set-default", help="Set default project")
+        set_default_parser.add_argument(
+            "name", type=str, help="The name of the project to set as default"
+        )
+        set_default_parser.set_defaults(subfunc=self._set_default)
+
+    def _command(self, args: argparse.Namespace):
+        if not hasattr(args, "subfunc"):
+            args.subfunc = self._list
+        args.subfunc(args)
+
+    def _add(self, args: argparse.Namespace):
+        config_manager = ConfigManager()
+        api_client = dstack.api.server.APIClient(base_url=args.url, token=args.token)
+        try:
+            api_client.projects.get(args.name)
+        except HTTPError as e:
+            if e.response.status_code == 403:
+                raise CLIError("Forbidden. Ensure the token is valid.")
+            elif e.response.status_code == 404:
+                raise CLIError(f"Project '{args.name}' not found.")
+            else:
+                raise e
+        default_project = config_manager.get_project_config()
+        if (
+            default_project is None
+            or default_project.name != args.name
+            or default_project.url != args.url
+            or default_project.token != args.token
+        ):
+            set_it_as_default = (
+                (
+                    args.yes
+                    or not default_project
+                    or confirm_ask(f"Set '{args.name}' as your default project?")
+                )
+                if not args.no
+                else False
+            )
+            config_manager.configure_project(
+                name=args.name, url=args.url, token=args.token, default=set_it_as_default
+            )
+            config_manager.save()
+        logger.info(
+            f"Configuration updated at {config_manager.config_filepath}", {"show_path": False}
+        )
+
+    def _delete(self, args: argparse.Namespace):
+        config_manager = ConfigManager()
+        if args.yes or confirm_ask(f"Are you sure you want to delete project '{args.name}'?"):
+            config_manager.delete_project(args.name)
+            config_manager.save()
+            console.print("[grey58]OK[/]")
+
+    def _list(self, args: argparse.Namespace):
+        config_manager = ConfigManager()
+        default_project = config_manager.get_project_config()
+
+        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")
+        table.add_column("DEFAULT", justify="center")
+
+        for project_name in config_manager.list_projects():
+            project_config = config_manager.get_project_config(project_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,
+            )
+
+        console.print(table)
+
+    def _set_default(self, args: argparse.Namespace):
+        config_manager = ConfigManager()
+        project_config = config_manager.get_project_config(args.name)
+        if project_config is None:
+            raise CLIError(f"Project '{args.name}' not found")
+
+        config_manager.configure_project(
+            name=args.name, url=project_config.url, token=project_config.token, default=True
+        )
+        config_manager.save()
+        console.print("[grey58]OK[/]")

dstack/_internal/cli/commands/ps.py

@@ -36,10 +36,17 @@ class PsCommand(APIBaseCommand):
             help="Watch statuses of runs in realtime",
             action="store_true",
         )
+        self._parser.add_argument(
+            "-n",
+            "--last",
+            help="Show only the last N runs. Implies --all",
+            type=int,
+            default=None,
+        )
 
     def _command(self, args: argparse.Namespace):
         super()._command(args)
-        runs = self.api.runs.list(all=args.all)
+        runs = self.api.runs.list(all=args.all, limit=args.last)
         if not args.watch:
             console.print(run_utils.get_runs_table(runs, verbose=args.verbose))
             return
@@ -49,6 +56,6 @@ class PsCommand(APIBaseCommand):
                 while True:
                     live.update(run_utils.get_runs_table(runs, verbose=args.verbose))
                     time.sleep(LIVE_TABLE_PROVISION_INTERVAL_SECS)
-                    runs = self.api.runs.list(all=args.all)
+                    runs = self.api.runs.list(all=args.all, limit=args.last)
         except KeyboardInterrupt:
             pass

dstack/_internal/cli/main.py

@@ -15,6 +15,7 @@ from dstack._internal.cli.commands.init import InitCommand
 from dstack._internal.cli.commands.logs import LogsCommand
 from dstack._internal.cli.commands.metrics import MetricsCommand
 from dstack._internal.cli.commands.offer import OfferCommand
+from dstack._internal.cli.commands.project import ProjectCommand
 from dstack._internal.cli.commands.ps import PsCommand
 from dstack._internal.cli.commands.server import ServerCommand
 from dstack._internal.cli.commands.stats import StatsCommand
@@ -69,6 +70,7 @@ def main():
     OfferCommand.register(subparsers)
     LogsCommand.register(subparsers)
     MetricsCommand.register(subparsers)
+    ProjectCommand.register(subparsers)
     PsCommand.register(subparsers)
     ServerCommand.register(subparsers)
     StatsCommand.register(subparsers)

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

@@ -105,7 +105,7 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
             changed_fields = []
             if run_plan.action == ApplyAction.UPDATE:
                 diff = diff_models(
-                    run_plan.run_spec.configuration,
+                    run_plan.get_effective_run_spec().configuration,
                     run_plan.current_resource.run_spec.configuration,
                 )
                 changed_fields = list(diff.keys())

dstack/_internal/cli/utils/updates.py

@@ -57,10 +57,22 @@ def _is_last_check_time_outdated() -> bool:
     )
 
 
+def is_update_available(current_version: str, latest_version: str) -> bool:
+    """
+    Return True if latest_version is newer than current_version.
+    Pre-releases are only considered if the current version is also a pre-release.
+    """
+    _current_version = pkg_version.parse(str(current_version))
+    _latest_version = pkg_version.parse(str(latest_version))
+    return _current_version < _latest_version and (
+        not _latest_version.is_prerelease or _current_version.is_prerelease
+    )
+
+
 def _check_version():
     latest_version = get_latest_version()
     if latest_version is not None:
-        if pkg_version.parse(str(version.__version__)) < pkg_version.parse(latest_version):
+        if is_update_available(version.__version__, latest_version):
             console.print(f"A new version of dstack is available: [code]{latest_version}[/]\n")
 
 

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

@@ -611,9 +611,12 @@ class AWSCompute(
                 raise e
         logger.debug("Deleted EBS volume %s", volume.configuration.name)
 
-    def attach_volume(self, volume: Volume, instance_id: str) -> VolumeAttachmentData:
+    def attach_volume(
+        self, volume: Volume, provisioning_data: JobProvisioningData
+    ) -> VolumeAttachmentData:
         ec2_client = self.session.client("ec2", region_name=volume.configuration.region)
 
+        instance_id = provisioning_data.instance_id
         device_names = aws_resources.list_available_device_names(
             ec2_client=ec2_client, instance_id=instance_id
         )
@@ -646,9 +649,12 @@ class AWSCompute(
         logger.debug("Attached EBS volume %s to instance %s", volume.volume_id, instance_id)
         return VolumeAttachmentData(device_name=device_name)
 
-    def detach_volume(self, volume: Volume, instance_id: str, force: bool = False):
+    def detach_volume(
+        self, volume: Volume, provisioning_data: JobProvisioningData, force: bool = False
+    ):
         ec2_client = self.session.client("ec2", region_name=volume.configuration.region)
 
+        instance_id = provisioning_data.instance_id
         logger.debug("Detaching EBS volume %s from instance %s", volume.volume_id, instance_id)
         attachment_data = get_or_error(volume.get_attachment_data_for_instance(instance_id))
         try:
@@ -667,9 +673,10 @@ class AWSCompute(
             raise e
         logger.debug("Detached EBS volume %s from instance %s", volume.volume_id, instance_id)
 
-    def is_volume_detached(self, volume: Volume, instance_id: str) -> bool:
+    def is_volume_detached(self, volume: Volume, provisioning_data: JobProvisioningData) -> bool:
         ec2_client = self.session.client("ec2", region_name=volume.configuration.region)
 
+        instance_id = provisioning_data.instance_id
         logger.debug("Getting EBS volume %s status", volume.volume_id)
         response = ec2_client.describe_volumes(VolumeIds=[volume.volume_id])
         volumes_infos = response.get("Volumes")
@@ -819,18 +826,23 @@ def _get_regions_to_zones(session: boto3.Session, regions: List[str]) -> Dict[st
 
 def _supported_instances(offer: InstanceOffer) -> bool:
     for family in [
+        "m7i.",
+        "c7i.",
+        "r7i.",
+        "t3.",
         "t2.small",
         "c5.",
         "m5.",
-        "g4dn.",
-        "g5.",
+        "p5.",
+        "p5e.",
+        "p4d.",
+        "p4de.",
+        "p3.",
         "g6.",
         "g6e.",
         "gr6.",
-        "p3.",
-        "p4d.",
-        "p4de.",
-        "p5.",
+        "g5.",
+        "g4dn.",
     ]:
         if offer.instance.name.startswith(family):
             return True

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

@@ -391,9 +391,9 @@ class VMImageVariant(enum.Enum):
 
 
 _SUPPORTED_VM_SERIES_PATTERNS = [
-    r"D(\d+)s_v3",  # Dsv3-series
-    r"E(\d+)i?s_v4",  # Esv4-series
-    r"E(\d+)-(\d+)s_v4",  # Esv4-series (constrained vCPU)
+    r"D(\d+)s_v6",  # Dsv6-series (general purpose)
+    r"E(\d+)i?s_v6",  # Esv6-series (memory optimized)
+    r"F(\d+)s_v2",  # Fsv2-series (compute optimized)
     r"NC(\d+)s_v3",  # NCv3-series [V100 16GB]
     r"NC(\d+)as_T4_v3",  # NCasT4_v3-series [T4]
     r"ND(\d+)rs_v2",  # NDv2-series [8xV100 32GB]
@@ -401,6 +401,11 @@ _SUPPORTED_VM_SERIES_PATTERNS = [
     r"NC(\d+)ads_A100_v4",  # NC A100 v4-series [A100 80GB]
     r"ND(\d+)asr_v4",  # ND A100 v4-series [8xA100 40GB]
     r"ND(\d+)amsr_A100_v4",  # NDm A100 v4-series [8xA100 80GB]
+    # Deprecated series
+    # TODO: Remove after several releases
+    r"D(\d+)s_v3",  # Dsv3-series (general purpose)
+    r"E(\d+)i?s_v4",  # Esv4-series (memory optimized)
+    r"E(\d+)-(\d+)s_v4",  # Esv4-series (constrained vCPU)
 ]
 _SUPPORTED_VM_SERIES_PATTERN = (
     "^Standard_(" + "|".join(f"({s})" for s in _SUPPORTED_VM_SERIES_PATTERNS) + ")$"

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

@@ -19,6 +19,7 @@ from dstack._internal.core.consts import (
     DSTACK_RUNNER_SSH_PORT,
     DSTACK_SHIM_HTTP_PORT,
 )
+from dstack._internal.core.models.configurations import DEFAULT_REPO_DIR
 from dstack._internal.core.models.gateways import (
     GatewayComputeConfiguration,
     GatewayProvisioningData,
@@ -335,7 +336,9 @@ class ComputeWithVolumeSupport(ABC):
         """
         raise NotImplementedError()
 
-    def attach_volume(self, volume: Volume, instance_id: str) -> VolumeAttachmentData:
+    def attach_volume(
+        self, volume: Volume, provisioning_data: JobProvisioningData
+    ) -> VolumeAttachmentData:
         """
         Attaches a volume to the instance.
         If the volume is not found, it should raise `ComputeError()`.
@@ -344,7 +347,9 @@ class ComputeWithVolumeSupport(ABC):
         """
         raise NotImplementedError()
 
-    def detach_volume(self, volume: Volume, instance_id: str, force: bool = False):
+    def detach_volume(
+        self, volume: Volume, provisioning_data: JobProvisioningData, force: bool = False
+    ):
         """
         Detaches a volume from the instance.
         Implement only if compute may return `VolumeProvisioningData.detachable`.
@@ -352,7 +357,7 @@ class ComputeWithVolumeSupport(ABC):
         """
         raise NotImplementedError()
 
-    def is_volume_detached(self, volume: Volume, instance_id: str) -> bool:
+    def is_volume_detached(self, volume: Volume, provisioning_data: JobProvisioningData) -> bool:
         """
         Checks if a volume was detached from the instance.
         If `detach_volume()` may fail to detach volume,
@@ -754,7 +759,7 @@ def get_docker_commands(
             f" --ssh-port {DSTACK_RUNNER_SSH_PORT}"
             " --temp-dir /tmp/runner"
             " --home-dir /root"
-            " --working-dir /workflow"
+            f" --working-dir {DEFAULT_REPO_DIR}"
         ),
     ]
 

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

@@ -649,13 +649,24 @@ class GCPCompute(
             pass
         logger.debug("Deleted persistent disk for volume %s", volume.name)
 
-    def attach_volume(self, volume: Volume, instance_id: str) -> VolumeAttachmentData:
+    def attach_volume(
+        self, volume: Volume, provisioning_data: JobProvisioningData
+    ) -> VolumeAttachmentData:
+        instance_id = provisioning_data.instance_id
         logger.debug(
             "Attaching persistent disk for volume %s to instance %s",
             volume.volume_id,
             instance_id,
         )
+        if not gcp_resources.instance_type_supports_persistent_disk(
+            provisioning_data.instance_type.name
+        ):
+            raise ComputeError(
+                f"Instance type {provisioning_data.instance_type.name} does not support Persistent disk volumes"
+            )
+
         zone = get_or_error(volume.provisioning_data).availability_zone
+        is_tpu = _is_tpu_provisioning_data(provisioning_data)
         try:
             disk = self.disk_client.get(
                 project=self.config.project_id,
@@ -663,18 +674,16 @@ class GCPCompute(
                 disk=volume.volume_id,
             )
             disk_url = disk.self_link
+        except google.api_core.exceptions.NotFound:
+            raise ComputeError("Persistent disk found")
 
-            # This method has no information if the instance is a TPU or a VM,
-            # so we first try to see if there is a TPU with such name
-            try:
+        try:
+            if is_tpu:
                 get_node_request = tpu_v2.GetNodeRequest(
                     name=f"projects/{self.config.project_id}/locations/{zone}/nodes/{instance_id}",
                 )
                 tpu_node = self.tpu_client.get_node(get_node_request)
-            except google.api_core.exceptions.NotFound:
-                tpu_node = None
 
-            if tpu_node is not None:
                 # Python API to attach a disk to a TPU is not documented,
                 # so we follow the code from the gcloud CLI:
                 # https://github.com/twistedpair/google-cloud-sdk/blob/26ab5a281d56b384cc25750f3279a27afe5b499f/google-cloud-sdk/lib/googlecloudsdk/command_lib/compute/tpus/tpu_vm/util.py#L113
@@ -711,7 +720,6 @@ class GCPCompute(
                 attached_disk.auto_delete = False
                 attached_disk.device_name = f"pd-{volume.volume_id}"
                 device_name = attached_disk.device_name
-
                 operation = self.instances_client.attach_disk(
                     project=self.config.project_id,
                     zone=zone,
@@ -720,13 +728,16 @@ class GCPCompute(
                 )
                 gcp_resources.wait_for_extended_operation(operation, "persistent disk attachment")
         except google.api_core.exceptions.NotFound:
-            raise ComputeError("Persistent disk or instance not found")
+            raise ComputeError("Disk or instance not found")
         logger.debug(
             "Attached persistent disk for volume %s to instance %s", volume.volume_id, instance_id
         )
         return VolumeAttachmentData(device_name=device_name)
 
-    def detach_volume(self, volume: Volume, instance_id: str, force: bool = False):
+    def detach_volume(
+        self, volume: Volume, provisioning_data: JobProvisioningData, force: bool = False
+    ):
+        instance_id = provisioning_data.instance_id
         logger.debug(
             "Detaching persistent disk for volume %s from instance %s",
             volume.volume_id,
@@ -734,17 +745,16 @@ class GCPCompute(
         )
         zone = get_or_error(volume.provisioning_data).availability_zone
         attachment_data = get_or_error(volume.get_attachment_data_for_instance(instance_id))
-        # This method has no information if the instance is a TPU or a VM,
-        # so we first try to see if there is a TPU with such name
-        try:
-            get_node_request = tpu_v2.GetNodeRequest(
-                name=f"projects/{self.config.project_id}/locations/{zone}/nodes/{instance_id}",
-            )
-            tpu_node = self.tpu_client.get_node(get_node_request)
-        except google.api_core.exceptions.NotFound:
-            tpu_node = None
+        is_tpu = _is_tpu_provisioning_data(provisioning_data)
+        if is_tpu:
+            try:
+                get_node_request = tpu_v2.GetNodeRequest(
+                    name=f"projects/{self.config.project_id}/locations/{zone}/nodes/{instance_id}",
+                )
+                tpu_node = self.tpu_client.get_node(get_node_request)
+            except google.api_core.exceptions.NotFound:
+                raise ComputeError("Instance not found")
 
-        if tpu_node is not None:
             source_disk = (
                 f"projects/{self.config.project_id}/zones/{zone}/disks/{volume.volume_id}"
             )
@@ -815,6 +825,11 @@ def _supported_instances_and_zones(
         if _is_tpu(offer.instance.name) and not _is_single_host_tpu(offer.instance.name):
             return False
         for family in [
+            "m4-",
+            "c4-",
+            "n4-",
+            "h3-",
+            "n2-",
             "e2-medium",
             "e2-standard-",
             "e2-highmem-",
@@ -1001,3 +1016,11 @@ def _get_tpu_data_disk_for_volume(project_id: str, volume: Volume) -> tpu_v2.Att
         mode=tpu_v2.AttachedDisk.DiskMode.READ_WRITE,
     )
     return attached_disk
+
+
+def _is_tpu_provisioning_data(provisioning_data: JobProvisioningData) -> bool:
+    is_tpu = False
+    if provisioning_data.backend_data:
+        backend_data_dict = json.loads(provisioning_data.backend_data)
+        is_tpu = backend_data_dict.get("is_tpu", False)
+    return is_tpu

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

@@ -140,7 +140,10 @@ def create_instance_struct(
     initialize_params = compute_v1.AttachedDiskInitializeParams()
     initialize_params.source_image = image_id
     initialize_params.disk_size_gb = disk_size
-    initialize_params.disk_type = f"zones/{zone}/diskTypes/pd-balanced"
+    if instance_type_supports_persistent_disk(machine_type):
+        initialize_params.disk_type = f"zones/{zone}/diskTypes/pd-balanced"
+    else:
+        initialize_params.disk_type = f"zones/{zone}/diskTypes/hyperdisk-balanced"
     disk.initialize_params = initialize_params
     instance.disks = [disk]
 
@@ -421,7 +424,7 @@ def wait_for_extended_operation(
 
     if operation.error_code:
         # 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: [Code: %s]: %s",
             verbose_name,
@@ -462,3 +465,16 @@ def get_placement_policy_resource_name(
     placement_policy: str,
 ) -> str:
     return f"projects/{project_id}/regions/{region}/resourcePolicies/{placement_policy}"
+
+
+def instance_type_supports_persistent_disk(instance_type_name: str) -> bool:
+    return not any(
+        instance_type_name.startswith(series)
+        for series in [
+            "m4-",
+            "c4-",
+            "n4-",
+            "h3-",
+            "v6e",
+        ]
+    )

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

@@ -110,8 +110,10 @@ class LocalCompute(
     def delete_volume(self, volume: Volume):
         pass
 
-    def attach_volume(self, volume: Volume, instance_id: str):
+    def attach_volume(self, volume: Volume, provisioning_data: JobProvisioningData):
         pass
 
-    def detach_volume(self, volume: Volume, instance_id: str, force: bool = False):
+    def detach_volume(
+        self, volume: Volume, provisioning_data: JobProvisioningData, force: bool = False
+    ):
         pass