From d0ebafdd77d17f4ccc5d0effbdcd869d4437c197 Mon Sep 17 00:00:00 2001 From: peterschmidt85 Date: Tue, 10 Feb 2026 15:33:42 +0100 Subject: [PATCH] Show live progress in `dstack attach` and handle PortUsedError - Show a spinner with status table in `dstack attach` while the run is provisioning, matching the UX of `dstack apply`. - Handle `PortUsedError` gracefully in both `dstack attach` and `dstack apply` instead of showing a raw traceback. The error message suggests using `-p` to override the local port mapping. - Store the port number on `PortUsedError` for structured access. Co-authored-by: Cursor --- src/dstack/_internal/cli/commands/attach.py | 50 ++++++++++++++++--- .../cli/services/configurators/run.py | 19 ++++++- .../_internal/core/services/ssh/ports.py | 8 +-- 3 files changed, 66 insertions(+), 11 deletions(-) diff --git a/src/dstack/_internal/cli/commands/attach.py b/src/dstack/_internal/cli/commands/attach.py index e005723e14..a22d63d37c 100644 --- a/src/dstack/_internal/cli/commands/attach.py +++ b/src/dstack/_internal/cli/commands/attach.py @@ -12,8 +12,12 @@ print_finished_message, ) from dstack._internal.cli.utils.common import console, get_start_time +from dstack._internal.cli.utils.rich import MultiItemStatus +from dstack._internal.cli.utils.run import get_runs_table from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT from dstack._internal.core.errors import CLIError +from dstack._internal.core.models.runs import RunStatus +from dstack._internal.core.services.ssh.ports import PortUsedError from dstack._internal.utils.common import get_or_error from dstack.api._public.runs import Run @@ -76,15 +80,39 @@ def _command(self, args: argparse.Namespace): run = self.api.runs.get(args.run_name) if run is None: raise CLIError(f"Run {args.run_name} not found") + + # Show live progress while waiting for the run to be ready + if _is_provisioning(run): + with MultiItemStatus(f"Attaching to [code]{run.name}[/]...", console=console) as live: + while _is_provisioning(run): + live.update(get_runs_table([run])) + time.sleep(5) + run.refresh() + console.print(get_runs_table([run], verbose=run.status == RunStatus.FAILED)) + console.print( + f"\nProvisioning [code]{run.name}[/] completed [secondary]({run.status.value})[/]" + ) + + if run.status.is_finished() and run.status != RunStatus.DONE: + raise CLIError(f"Run {args.run_name} is {run.status.value}") + exit_code = 0 try: - attached = run.attach( - ssh_identity_file=args.ssh_identity_file, - bind_address=args.host, - ports_overrides=args.ports, - replica_num=args.replica, - job_num=args.job, - ) + try: + attached = run.attach( + ssh_identity_file=args.ssh_identity_file, + bind_address=args.host, + ports_overrides=args.ports, + replica_num=args.replica, + job_num=args.job, + ) + except PortUsedError as e: + console.print( + f"[error]Failed to attach: port [code]{e.port}[/code] is already in use." + f" Use [code]-p[/code] in [code]dstack attach[/code] to override the local" + f" port mapping, e.g. [code]-p {e.port + 1}:{e.port}[/code].[/]" + ) + exit(1) if not attached: raise CLIError(f"Failed to attach to run {args.run_name}") _print_attached_message( @@ -159,3 +187,11 @@ def _print_attached_message( output += f"To connect to the run via SSH, use `ssh {name}`.\n" output += "Press Ctrl+C to detach..." console.print(output) + + +def _is_provisioning(run: Run) -> bool: + return run.status in ( + RunStatus.SUBMITTED, + RunStatus.PENDING, + RunStatus.PROVISIONING, + ) diff --git a/src/dstack/_internal/cli/services/configurators/run.py b/src/dstack/_internal/cli/services/configurators/run.py index 1077eff8a9..33ba4e10d2 100644 --- a/src/dstack/_internal/cli/services/configurators/run.py +++ b/src/dstack/_internal/cli/services/configurators/run.py @@ -56,6 +56,7 @@ InvalidRepoCredentialsError, get_repo_creds_and_default_branch, ) +from dstack._internal.core.services.ssh.ports import PortUsedError from dstack._internal.utils.common import local_time from dstack._internal.utils.interpolator import InterpolatorError, VariablesInterpolator from dstack._internal.utils.logging import get_logger @@ -168,6 +169,13 @@ def apply_configuration( ) except ServerClientError as e: raise CLIError(e.msg) + except PortUsedError as e: + console.print( + f"[error]Failed to submit: port [code]{e.port}[/code] is already in use." + f" Use [code]-p[/code] in [code]dstack apply[/code] to override the local" + f" port mapping, e.g. [code]-p {e.port + 1}:{e.port}[/code].[/]" + ) + exit(1) if command_args.detach: detach_message = f"Run [code]{run.name}[/] submitted, detaching..." @@ -206,7 +214,16 @@ def apply_configuration( configurator_args, _BIND_ADDRESS_ARG, None ) try: - if run.attach(bind_address=bind_address): + try: + attached = run.attach(bind_address=bind_address) + except PortUsedError as e: + console.print( + f"[error]Failed to attach: port [code]{e.port}[/code] is already in use." + f" Use [code]-p[/code] in [code]dstack attach[/code] to override the local" + f" port mapping, e.g. [code]-p {e.port + 1}:{e.port}[/code].[/]" + ) + exit(1) + if attached: for entry in run.logs(): sys.stdout.buffer.write(entry) sys.stdout.buffer.flush() diff --git a/src/dstack/_internal/core/services/ssh/ports.py b/src/dstack/_internal/core/services/ssh/ports.py index f0716e6158..1d41bcd2c6 100644 --- a/src/dstack/_internal/core/services/ssh/ports.py +++ b/src/dstack/_internal/core/services/ssh/ports.py @@ -11,7 +11,9 @@ class PortUsedError(DstackError): - pass + def __init__(self, port: int): + self.port = port + super().__init__(f"Port {port} is already in use") class PortsLock: @@ -28,10 +30,10 @@ def acquire(self) -> "PortsLock": if not local_port: # None or 0 continue if local_port in assigned_ports: - raise PortUsedError(f"Port {local_port} is already in use") + raise PortUsedError(local_port) sock = self._listen(local_port) if sock is None: - raise PortUsedError(f"Port {local_port} is already in use") + raise PortUsedError(local_port) self.sockets[remote_port] = sock assigned_ports.add(local_port)