Print message on dstack attach exit (#2358)

Author: r4victor

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

@@ -7,6 +7,10 @@
 from dstack._internal.cli.commands import APIBaseCommand
 from dstack._internal.cli.services.args import port_mapping
 from dstack._internal.cli.services.completion import RunNameCompleter
+from dstack._internal.cli.services.configurators.run import (
+    get_run_exit_code,
+    print_finished_message,
+)
 from dstack._internal.cli.utils.common import console
 from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT
 from dstack._internal.core.errors import CLIError
@@ -100,6 +104,21 @@ def _command(self, args: argparse.Namespace):
                     pass
         finally:
             run.detach()
+        # TODO: Handle run resubmissions similar to dstack apply
+
+        # After reading the logs, the run may not be marked as finished immediately.
+        # Give the run some time to transition to a finished state before exiting.
+        for _ in range(30):
+            run.refresh()
+            if run.status.is_finished():
+                print_finished_message(run)
+                exit(get_run_exit_code(run))
+            time.sleep(1)
+        console.print(
+            "[error]Lost run connection. Timed out waiting for run final status."
+            " Check `dstack ps` to see if it's done or failed."
+        )
+        exit(1)
 
 
 _IGNORED_PORTS = [DSTACK_RUNNER_HTTP_PORT]

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

@@ -238,8 +238,8 @@ def apply_configuration(
                         reattach = True
                         break
                     if run.status.is_finished():
-                        _print_finished_message(run)
-                        exit(_get_run_exit_code(run))
+                        print_finished_message(run)
+                        exit(get_run_exit_code(run))
                     time.sleep(1)
                 if not reattach:
                     console.print(
@@ -439,7 +439,7 @@ def apply_args(
     ):
         super().apply_args(conf, args, unknown)
         if args.ports:
-            conf.ports = list(merge_ports(conf.ports, args.ports).values())
+            conf.ports = list(_merge_ports(conf.ports, args.ports).values())
 
 
 class TaskConfigurator(RunWithPortsConfigurator):
@@ -475,17 +475,17 @@ def apply_args(self, conf: ServiceConfiguration, args: argparse.Namespace, unkno
         self.interpolate_run_args(conf.commands, unknown)
 
 
-def merge_ports(conf: List[PortMapping], args: List[PortMapping]) -> Dict[int, PortMapping]:
-    unique_ports_constraint([pm.container_port for pm in conf])
-    unique_ports_constraint([pm.container_port for pm in args])
+def _merge_ports(conf: List[PortMapping], args: List[PortMapping]) -> Dict[int, PortMapping]:
+    _unique_ports_constraint([pm.container_port for pm in conf])
+    _unique_ports_constraint([pm.container_port for pm in args])
     ports = {pm.container_port: pm for pm in conf}
     for pm in args:  # override conf
         ports[pm.container_port] = pm
-    unique_ports_constraint([pm.local_port for pm in ports.values() if pm.local_port is not None])
+    _unique_ports_constraint([pm.local_port for pm in ports.values() if pm.local_port is not None])
     return ports
 
 
-def unique_ports_constraint(ports: List[int]):
+def _unique_ports_constraint(ports: List[int]):
     used_ports = set()
     for i in ports:
         if i in used_ports:
@@ -514,7 +514,7 @@ def _print_service_urls(run: Run) -> None:
     console.print()
 
 
-def _print_finished_message(run: Run):
+def print_finished_message(run: Run):
     if run.status == RunStatus.DONE:
         console.print("[code]Done[/]")
         return
@@ -542,7 +542,7 @@ def _print_finished_message(run: Run):
     console.print(f"[error]{message}[/]")
 
 
-def _get_run_exit_code(run: Run) -> int:
+def get_run_exit_code(run: Run) -> int:
     if run.status == RunStatus.DONE:
         return 0
     return 1