Skip to content

Commit d0ebafd

Browse files
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 <cursoragent@cursor.com>
1 parent fb4a4da commit d0ebafd

File tree

3 files changed

+66
-11
lines changed

3 files changed

+66
-11
lines changed

src/dstack/_internal/cli/commands/attach.py

Lines changed: 43 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,12 @@
1212
print_finished_message,
1313
)
1414
from dstack._internal.cli.utils.common import console, get_start_time
15+
from dstack._internal.cli.utils.rich import MultiItemStatus
16+
from dstack._internal.cli.utils.run import get_runs_table
1517
from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT
1618
from dstack._internal.core.errors import CLIError
19+
from dstack._internal.core.models.runs import RunStatus
20+
from dstack._internal.core.services.ssh.ports import PortUsedError
1721
from dstack._internal.utils.common import get_or_error
1822
from dstack.api._public.runs import Run
1923

@@ -76,15 +80,39 @@ def _command(self, args: argparse.Namespace):
7680
run = self.api.runs.get(args.run_name)
7781
if run is None:
7882
raise CLIError(f"Run {args.run_name} not found")
83+
84+
# Show live progress while waiting for the run to be ready
85+
if _is_provisioning(run):
86+
with MultiItemStatus(f"Attaching to [code]{run.name}[/]...", console=console) as live:
87+
while _is_provisioning(run):
88+
live.update(get_runs_table([run]))
89+
time.sleep(5)
90+
run.refresh()
91+
console.print(get_runs_table([run], verbose=run.status == RunStatus.FAILED))
92+
console.print(
93+
f"\nProvisioning [code]{run.name}[/] completed [secondary]({run.status.value})[/]"
94+
)
95+
96+
if run.status.is_finished() and run.status != RunStatus.DONE:
97+
raise CLIError(f"Run {args.run_name} is {run.status.value}")
98+
7999
exit_code = 0
80100
try:
81-
attached = run.attach(
82-
ssh_identity_file=args.ssh_identity_file,
83-
bind_address=args.host,
84-
ports_overrides=args.ports,
85-
replica_num=args.replica,
86-
job_num=args.job,
87-
)
101+
try:
102+
attached = run.attach(
103+
ssh_identity_file=args.ssh_identity_file,
104+
bind_address=args.host,
105+
ports_overrides=args.ports,
106+
replica_num=args.replica,
107+
job_num=args.job,
108+
)
109+
except PortUsedError as e:
110+
console.print(
111+
f"[error]Failed to attach: port [code]{e.port}[/code] is already in use."
112+
f" Use [code]-p[/code] in [code]dstack attach[/code] to override the local"
113+
f" port mapping, e.g. [code]-p {e.port + 1}:{e.port}[/code].[/]"
114+
)
115+
exit(1)
88116
if not attached:
89117
raise CLIError(f"Failed to attach to run {args.run_name}")
90118
_print_attached_message(
@@ -159,3 +187,11 @@ def _print_attached_message(
159187
output += f"To connect to the run via SSH, use `ssh {name}`.\n"
160188
output += "Press Ctrl+C to detach..."
161189
console.print(output)
190+
191+
192+
def _is_provisioning(run: Run) -> bool:
193+
return run.status in (
194+
RunStatus.SUBMITTED,
195+
RunStatus.PENDING,
196+
RunStatus.PROVISIONING,
197+
)

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

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@
5656
InvalidRepoCredentialsError,
5757
get_repo_creds_and_default_branch,
5858
)
59+
from dstack._internal.core.services.ssh.ports import PortUsedError
5960
from dstack._internal.utils.common import local_time
6061
from dstack._internal.utils.interpolator import InterpolatorError, VariablesInterpolator
6162
from dstack._internal.utils.logging import get_logger
@@ -168,6 +169,13 @@ def apply_configuration(
168169
)
169170
except ServerClientError as e:
170171
raise CLIError(e.msg)
172+
except PortUsedError as e:
173+
console.print(
174+
f"[error]Failed to submit: port [code]{e.port}[/code] is already in use."
175+
f" Use [code]-p[/code] in [code]dstack apply[/code] to override the local"
176+
f" port mapping, e.g. [code]-p {e.port + 1}:{e.port}[/code].[/]"
177+
)
178+
exit(1)
171179

172180
if command_args.detach:
173181
detach_message = f"Run [code]{run.name}[/] submitted, detaching..."
@@ -206,7 +214,16 @@ def apply_configuration(
206214
configurator_args, _BIND_ADDRESS_ARG, None
207215
)
208216
try:
209-
if run.attach(bind_address=bind_address):
217+
try:
218+
attached = run.attach(bind_address=bind_address)
219+
except PortUsedError as e:
220+
console.print(
221+
f"[error]Failed to attach: port [code]{e.port}[/code] is already in use."
222+
f" Use [code]-p[/code] in [code]dstack attach[/code] to override the local"
223+
f" port mapping, e.g. [code]-p {e.port + 1}:{e.port}[/code].[/]"
224+
)
225+
exit(1)
226+
if attached:
210227
for entry in run.logs():
211228
sys.stdout.buffer.write(entry)
212229
sys.stdout.buffer.flush()

src/dstack/_internal/core/services/ssh/ports.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,9 @@
1111

1212

1313
class PortUsedError(DstackError):
14-
pass
14+
def __init__(self, port: int):
15+
self.port = port
16+
super().__init__(f"Port {port} is already in use")
1517

1618

1719
class PortsLock:
@@ -28,10 +30,10 @@ def acquire(self) -> "PortsLock":
2830
if not local_port: # None or 0
2931
continue
3032
if local_port in assigned_ports:
31-
raise PortUsedError(f"Port {local_port} is already in use")
33+
raise PortUsedError(local_port)
3234
sock = self._listen(local_port)
3335
if sock is None:
34-
raise PortUsedError(f"Port {local_port} is already in use")
36+
raise PortUsedError(local_port)
3537
self.sockets[remote_port] = sock
3638
assigned_ports.add(local_port)
3739

0 commit comments

Comments
 (0)