dstack 0.18.40rc1__py3-none-any.whl → 0.18.42__py3-none-any.whl

dstack/_internal/cli/commands/apply.py

@@ -6,12 +6,12 @@ from dstack._internal.cli.services.configurators import (
     get_apply_configurator_class,
     load_apply_configuration,
 )
-from dstack._internal.cli.services.configurators.base import BaseApplyConfigurator
 from dstack._internal.cli.services.repos import (
     init_default_virtual_repo,
     init_repo,
     register_init_repo_args,
 )
+from dstack._internal.cli.utils.common import console
 from dstack._internal.core.errors import CLIError
 from dstack._internal.core.models.configurations import ApplyConfigurationType
 
@@ -92,10 +92,13 @@ class ApplyCommand(APIBaseCommand):
                     configurator_class = get_apply_configurator_class(
                         ApplyConfigurationType(args.help)
                     )
-                else:
-                    configurator_class = BaseApplyConfigurator
-                configurator_class.register_args(self._parser)
+                    configurator_class.register_args(self._parser)
+                    self._parser.print_help()
+                    return
                 self._parser.print_help()
+                console.print(
+                    "\nType `dstack apply -h CONFIGURATION_TYPE` to see configuration-specific options.\n"
+                )
                 return
 
             super()._command(args)
@@ -129,5 +132,5 @@ class ApplyCommand(APIBaseCommand):
                 repo=repo,
             )
         except KeyboardInterrupt:
-            print("\nOperation interrupted by user. Exiting...")
+            console.print("\nOperation interrupted by user. Exiting...")
             exit(0)

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

@@ -1,7 +1,7 @@
 import argparse
 import os
 from abc import ABC, abstractmethod
-from typing import List, Optional, cast
+from typing import List, Optional, Union, cast
 
 from dstack._internal.cli.services.args import env_var
 from dstack._internal.core.errors import ConfigurationError
@@ -14,6 +14,8 @@ from dstack._internal.core.models.envs import Env, EnvSentinel, EnvVarTuple
 from dstack._internal.core.models.repos.base import Repo
 from dstack.api._public import Client
 
+ArgsParser = Union[argparse._ArgumentGroup, argparse.ArgumentParser]
+
 
 class BaseApplyConfigurator(ABC):
     TYPE: ApplyConfigurationType
@@ -82,7 +84,7 @@ class BaseApplyConfigurator(ABC):
 
 class ApplyEnvVarsConfiguratorMixin:
     @classmethod
-    def register_env_args(cls, parser: argparse.ArgumentParser):
+    def register_env_args(cls, parser: ArgsParser):
         parser.add_argument(
             "-e",
             "--env",

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

@@ -41,10 +41,6 @@ logger = get_logger(__name__)
 class FleetConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
     TYPE: ApplyConfigurationType = ApplyConfigurationType.FLEET
 
-    @classmethod
-    def register_args(cls, parser: argparse.ArgumentParser):
-        cls.register_env_args(parser)
-
     def apply_configuration(
         self,
         conf: FleetConfiguration,
@@ -185,20 +181,36 @@ class FleetConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
 
         console.print(f"Fleet [code]{conf.name}[/] deleted")
 
+    @classmethod
+    def register_args(cls, parser: argparse.ArgumentParser):
+        configuration_group = parser.add_argument_group(f"{cls.TYPE.value} Options")
+        configuration_group.add_argument(
+            "-n",
+            "--name",
+            dest="name",
+            help="The fleet name",
+        )
+        cls.register_env_args(configuration_group)
+
     def apply_args(self, conf: FleetConfiguration, args: argparse.Namespace, unknown: List[str]):
+        if args.name:
+            conf.name = args.name
         self.apply_env_vars(conf.env, args)
         if conf.ssh_config is None and conf.env:
             raise ConfigurationError("`env` is currently supported for SSH fleets only")
 
 
 def _preprocess_spec(spec: FleetSpec):
-    if spec.configuration.ssh_config is not None:
-        spec.configuration.ssh_config.ssh_key = _resolve_ssh_key(
-            spec.configuration.ssh_config.identity_file
-        )
-        for host in spec.configuration.ssh_config.hosts:
+    ssh_config = spec.configuration.ssh_config
+    if ssh_config is not None:
+        ssh_config.ssh_key = _resolve_ssh_key(ssh_config.identity_file)
+        if ssh_config.proxy_jump is not None:
+            ssh_config.proxy_jump.ssh_key = _resolve_ssh_key(ssh_config.proxy_jump.identity_file)
+        for host in ssh_config.hosts:
             if not isinstance(host, str):
                 host.ssh_key = _resolve_ssh_key(host.identity_file)
+                if host.proxy_jump is not None:
+                    host.proxy_jump.ssh_key = _resolve_ssh_key(host.proxy_jump.identity_file)
 
 
 def _resolve_ssh_key(ssh_key_path: Optional[str]) -> Optional[SSHKey]:

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

@@ -39,6 +39,7 @@ class GatewayConfigurator(BaseApplyConfigurator):
         unknown_args: List[str],
         repo: Optional[Repo] = None,
     ):
+        self.apply_args(conf, configurator_args, unknown_args)
         spec = GatewaySpec(
             configuration=conf,
             configuration_path=configuration_path,
@@ -170,6 +171,20 @@ class GatewayConfigurator(BaseApplyConfigurator):
 
         console.print(f"Gateway [code]{conf.name}[/] deleted")
 
+    @classmethod
+    def register_args(cls, parser: argparse.ArgumentParser):
+        configuration_group = parser.add_argument_group(f"{cls.TYPE.value} Options")
+        configuration_group.add_argument(
+            "-n",
+            "--name",
+            dest="name",
+            help="The gateway name",
+        )
+
+    def apply_args(self, conf: GatewayConfiguration, args: argparse.Namespace, unknown: List[str]):
+        if args.name:
+            conf.name = args.name
+
 
 def _get_plan(api: Client, spec: GatewaySpec) -> GatewayPlan:
     # TODO: Implement server-side /get_plan with an offer included

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

@@ -298,20 +298,21 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
 
     @classmethod
     def register_args(cls, parser: argparse.ArgumentParser):
-        parser.add_argument(
+        configuration_group = parser.add_argument_group(f"{cls.TYPE.value} Options")
+        configuration_group.add_argument(
             "-n",
             "--name",
             dest="run_name",
             help="The name of the run. If not specified, a random name is assigned",
         )
-        parser.add_argument(
+        configuration_group.add_argument(
             "--max-offers",
             help="Number of offers to show in the run plan",
             type=int,
             default=3,
         )
-        cls.register_env_args(parser)
-        parser.add_argument(
+        cls.register_env_args(configuration_group)
+        configuration_group.add_argument(
             "--gpu",
             type=gpu_spec,
             help="Request GPU for the run. "
@@ -319,7 +320,7 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
             dest="gpu_spec",
             metavar="SPEC",
         )
-        parser.add_argument(
+        configuration_group.add_argument(
             "--disk",
             type=disk_spec,
             help="Request the size range of disk for the run. Example [code]--disk 100GB..[/].",

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

@@ -38,6 +38,7 @@ class VolumeConfigurator(BaseApplyConfigurator):
         unknown_args: List[str],
         repo: Optional[Repo] = None,
     ):
+        self.apply_args(conf, configurator_args, unknown_args)
         spec = VolumeSpec(
             configuration=conf,
             configuration_path=configuration_path,
@@ -158,6 +159,20 @@ class VolumeConfigurator(BaseApplyConfigurator):
 
         console.print(f"Volume [code]{conf.name}[/] deleted")
 
+    @classmethod
+    def register_args(cls, parser: argparse.ArgumentParser):
+        configuration_group = parser.add_argument_group(f"{cls.TYPE.value} Options")
+        configuration_group.add_argument(
+            "-n",
+            "--name",
+            dest="name",
+            help="The volume name",
+        )
+
+    def apply_args(self, conf: VolumeConfiguration, args: argparse.Namespace, unknown: List[str]):
+        if args.name:
+            conf.name = args.name
+
 
 def _get_plan(api: Client, spec: VolumeSpec) -> VolumePlan:
     # TODO: Implement server-side /get_plan with an offer included

dstack/_internal/cli/services/repos.py

@@ -1,7 +1,7 @@
-from argparse import ArgumentParser, _ArgumentGroup
 from pathlib import Path
-from typing import Optional, Union
+from typing import Optional
 
+from dstack._internal.cli.services.configurators.base import ArgsParser
 from dstack._internal.core.errors import CLIError
 from dstack._internal.core.models.repos.base import Repo, RepoType
 from dstack._internal.core.models.repos.remote import GitRepoURL, RemoteRepo, RepoError
@@ -12,7 +12,7 @@ from dstack._internal.utils.path import PathLike
 from dstack.api._public import Client
 
 
-def register_init_repo_args(parser: Union[ArgumentParser, _ArgumentGroup]):
+def register_init_repo_args(parser: ArgsParser):
     parser.add_argument(
         "-t",
         "--token",

dstack/_internal/cli/utils/fleet.py

@@ -2,7 +2,7 @@ from typing import List
 
 from rich.table import Table
 
-from dstack._internal.cli.utils.common import console
+from dstack._internal.cli.utils.common import add_row_from_dict, console
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.fleets import Fleet, FleetStatus
 from dstack._internal.core.models.instances import InstanceStatus
@@ -23,6 +23,8 @@ def get_fleets_table(
         table.add_column("RESERVATION")
     table.add_column("INSTANCE")
     table.add_column("BACKEND")
+    if verbose:
+        table.add_column("REGION")
     table.add_column("RESOURCES")
     table.add_column("PRICE")
     table.add_column("STATUS")
@@ -41,6 +43,10 @@ def get_fleets_table(
                 resources = instance.instance_type.resources.pretty_format(include_spot=True)
 
             status = instance.status.value
+            total_blocks = instance.total_blocks
+            busy_blocks = instance.busy_blocks
+            if total_blocks is not None and total_blocks > 1:
+                status = f"{busy_blocks}/{total_blocks} {InstanceStatus.BUSY.value}"
             if (
                 instance.status in [InstanceStatus.IDLE, InstanceStatus.BUSY]
                 and instance.unreachable
@@ -50,40 +56,45 @@ def get_fleets_table(
             backend = instance.backend or ""
             if backend == "remote":
                 backend = "ssh"
-            if instance.region:
-                backend += f" ({instance.region})"
 
-            row = [
-                fleet.name if i == 0 else "",
-            ]
-            if verbose:
-                row.append(fleet.spec.configuration.reservation or "" if i == 0 else "")
-            row += [
-                str(instance.instance_num),
-                backend,
-                resources,
-                f"${instance.price:.4}" if instance.price is not None else "",
-                status,
-                format_date(instance.created),
-            ]
-            if verbose:
-                error = ""
-                if instance.status == InstanceStatus.TERMINATED and instance.termination_reason:
-                    error = f"{instance.termination_reason}"
-                row.append(error)
-
-            table.add_row(*row)
+            region = ""
+            if instance.region:
+                region = f"{instance.region}"
+                if verbose:
+                    if instance.availability_zone:
+                        region += f" ({instance.availability_zone})"
+                else:
+                    backend += f" ({instance.region})"
+            error = ""
+            if instance.status == InstanceStatus.TERMINATED and instance.termination_reason:
+                error = f"{instance.termination_reason}"
+            row = {
+                "FLEET": fleet.name if i == 0 else "",
+                "RESERVATION": fleet.spec.configuration.reservation or "" if i == 0 else "",
+                "INSTANCE": str(instance.instance_num),
+                "BACKEND": backend,
+                "REGION": region,
+                "RESOURCES": resources,
+                "PRICE": f"${instance.price:.4}" if instance.price is not None else "",
+                "STATUS": status,
+                "CREATED": format_date(instance.created),
+                "ERROR": error,
+            }
+            add_row_from_dict(table, row)
 
         if len(fleet.instances) == 0 and fleet.status != FleetStatus.TERMINATING:
-            row = [
-                fleet.name,
-                "-",
-                "-",
-                "-",
-                "-",
-                "-",
-                format_date(fleet.created_at),
-            ]
-            table.add_row(*row)
+            row = {
+                "FLEET": fleet.name,
+                "RESERVATION": "-",
+                "INSTANCE": "-",
+                "BACKEND": "-",
+                "REGION": "-",
+                "RESOURCES": "-",
+                "PRICE": "-",
+                "STATUS": "-",
+                "CREATED": format_date(fleet.created_at),
+                "ERROR": "-",
+            }
+            add_row_from_dict(table, row)
 
     return table

dstack/_internal/cli/utils/run.py

@@ -14,7 +14,12 @@ from dstack._internal.core.models.runs import (
     RunPlan,
 )
 from dstack._internal.core.services.profiles import get_termination
-from dstack._internal.utils.common import DateFormatter, format_pretty_duration, pretty_date
+from dstack._internal.utils.common import (
+    DateFormatter,
+    format_duration_multiunit,
+    format_pretty_duration,
+    pretty_date,
+)
 from dstack.api import Run
 
 
@@ -96,11 +101,14 @@ def print_run_plan(run_plan: RunPlan, offers_limit: int = 3):
             InstanceAvailability.BUSY,
         }:
             availability = offer.availability.value.replace("_", " ").lower()
+        instance = offer.instance.name
+        if offer.total_blocks > 1:
+            instance += f" ({offer.blocks}/{offer.total_blocks})"
         offers.add_row(
             f"{i}",
             offer.backend.replace("remote", "ssh"),
             offer.region,
-            offer.instance.name,
+            instance,
             r.pretty_format(),
             "yes" if r.spot else "no",
             f"${offer.price:g}",
@@ -155,19 +163,31 @@ def get_runs_table(
             add_row_from_dict(table, run_row)
 
         for job in run.jobs:
+            latest_job_submission = job.job_submissions[-1]
+            status = latest_job_submission.status.value
+            if verbose and latest_job_submission.inactivity_secs:
+                inactive_for = format_duration_multiunit(latest_job_submission.inactivity_secs)
+                status += f" (inactive for {inactive_for})"
             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,
-                "SUBMITTED": format_date(job.job_submissions[-1].submitted_at),
+                "STATUS": status,
+                "SUBMITTED": format_date(latest_job_submission.submitted_at),
                 "ERROR": _get_job_error(job),
             }
-            jpd = job.job_submissions[-1].job_provisioning_data
+            jpd = latest_job_submission.job_provisioning_data
             if jpd is not None:
+                resources = jpd.instance_type.resources
+                instance = jpd.instance_type.name
+                jrd = latest_job_submission.job_runtime_data
+                if jrd is not None and jrd.offer is not None:
+                    resources = jrd.offer.instance.resources
+                    if jrd.offer.total_blocks > 1:
+                        instance += f" ({jrd.offer.blocks}/{jrd.offer.total_blocks})"
                 job_row.update(
                     {
                         "BACKEND": f"{jpd.backend.value.replace('remote', 'ssh')} ({jpd.region})",
-                        "INSTANCE": jpd.instance_type.name,
-                        "RESOURCES": jpd.instance_type.resources.pretty_format(include_spot=True),
+                        "INSTANCE": instance,
+                        "RESOURCES": resources.pretty_format(include_spot=True),
                         "RESERVATION": jpd.reservation,
                         "PRICE": f"${jpd.price:.4}",
                     }

dstack/_internal/cli/utils/volume.py

@@ -2,7 +2,7 @@ from typing import List
 
 from rich.table import Table
 
-from dstack._internal.cli.utils.common import console
+from dstack._internal.cli.utils.common import add_row_from_dict, console
 from dstack._internal.core.models.volumes import Volume
 from dstack._internal.utils.common import DateFormatter, pretty_date
 
@@ -19,19 +19,40 @@ def get_volumes_table(
     table = Table(box=None)
     table.add_column("NAME", no_wrap=True)
     table.add_column("BACKEND")
+    if verbose:
+        table.add_column("REGION")
     table.add_column("STATUS")
+    if verbose:
+        table.add_column("ATTACHED")
     table.add_column("CREATED")
     if verbose:
         table.add_column("ERROR")
 
     for volume in volumes:
-        renderables = [
-            volume.name,
-            f"{volume.configuration.backend.value} ({volume.configuration.region})",
-            volume.status,
-            format_date(volume.created_at),
-        ]
+        backend = f"{volume.configuration.backend.value} ({volume.configuration.region})"
+        region = volume.configuration.region
         if verbose:
-            renderables.append(volume.status_message)
-        table.add_row(*renderables)
+            backend = volume.configuration.backend.value
+            if (
+                verbose
+                and volume.provisioning_data is not None
+                and volume.provisioning_data.availability_zone is not None
+            ):
+                region += f" ({volume.provisioning_data.availability_zone})"
+        attached = "-"
+        if volume.attachments is not None:
+            attached = ", ".join(
+                {va.instance.fleet_name for va in volume.attachments if va.instance.fleet_name}
+            )
+            attached = attached or "-"
+        row = {
+            "NAME": volume.name,
+            "BACKEND": backend,
+            "REGION": region,
+            "STATUS": volume.status,
+            "ATTACHED": attached,
+            "CREATED": format_date(volume.created_at),
+            "ERROR": volume.status_message,
+        }
+        add_row_from_dict(table, row)
     return table