diff --git a/src/dstack/_internal/cli/services/configurators/fleet.py b/src/dstack/_internal/cli/services/configurators/fleet.py index 2865078f2..76aee535c 100644 --- a/src/dstack/_internal/cli/services/configurators/fleet.py +++ b/src/dstack/_internal/cli/services/configurators/fleet.py @@ -4,7 +4,6 @@ from typing import List, Optional import requests -from rich.console import Group from rich.table import Table from dstack._internal.cli.services.configurators.base import ( @@ -17,6 +16,7 @@ console, ) from dstack._internal.cli.utils.fleet import get_fleets_table +from dstack._internal.cli.utils.rich import MultiItemStatus from dstack._internal.core.errors import ConfigurationError, ResourceNotExistsError from dstack._internal.core.models.configurations import ApplyConfigurationType from dstack._internal.core.models.fleets import ( @@ -118,10 +118,12 @@ def apply_configuration( console.print("Fleet configuration submitted. Exiting...") return try: - with console.status("") as live: + with MultiItemStatus( + f"Provisioning [code]{fleet.name}[/]...", console=console + ) as live: while not _finished_provisioning(fleet): table = get_fleets_table([fleet]) - live.update(Group(f"Provisioning [code]{fleet.name}[/]...\n", table)) + live.update(table) time.sleep(LIVE_TABLE_PROVISION_INTERVAL_SECS) fleet = self.api.client.fleets.get(self.api.project, fleet.name) except KeyboardInterrupt: @@ -133,7 +135,6 @@ def apply_configuration( else: console.print("Exiting... Fleet provisioning will continue in the background.") return - console.print() console.print( get_fleets_table( [fleet], diff --git a/src/dstack/_internal/cli/services/configurators/gateway.py b/src/dstack/_internal/cli/services/configurators/gateway.py index 0e8429e7c..8b2ab7824 100644 --- a/src/dstack/_internal/cli/services/configurators/gateway.py +++ b/src/dstack/_internal/cli/services/configurators/gateway.py @@ -2,7 +2,6 @@ import time from typing import List -from rich.console import Group from rich.table import Table from dstack._internal.cli.services.configurators.base import BaseApplyConfigurator @@ -12,6 +11,7 @@ console, ) from dstack._internal.cli.utils.gateway import get_gateways_table +from dstack._internal.cli.utils.rich import MultiItemStatus from dstack._internal.core.errors import ResourceNotExistsError from dstack._internal.core.models.configurations import ApplyConfigurationType from dstack._internal.core.models.gateways import ( @@ -100,10 +100,12 @@ def apply_configuration( console.print("Gateway configuration submitted. Exiting...") return try: - with console.status("") as live: + with MultiItemStatus( + f"Provisioning [code]{gateway.name}[/]...", console=console + ) as live: while not _finished_provisioning(gateway): table = get_gateways_table([gateway], include_created=True) - live.update(Group(f"Provisioning [code]{gateway.name}[/]...\n", table)) + live.update(table) time.sleep(LIVE_TABLE_PROVISION_INTERVAL_SECS) gateway = self.api.client.gateways.get(self.api.project, gateway.name) except KeyboardInterrupt: @@ -116,7 +118,6 @@ def apply_configuration( else: console.print("Exiting... Gateway provisioning will continue in the background.") return - console.print() console.print( get_gateways_table( [gateway], diff --git a/src/dstack/_internal/cli/services/configurators/run.py b/src/dstack/_internal/cli/services/configurators/run.py index 2adbcb985..6c647a9a7 100644 --- a/src/dstack/_internal/cli/services/configurators/run.py +++ b/src/dstack/_internal/cli/services/configurators/run.py @@ -6,7 +6,6 @@ from typing import Dict, List, Optional, Set, Tuple import gpuhunt -from rich.console import Group import dstack._internal.core.models.resources as resources from dstack._internal.cli.services.args import disk_spec, gpu_spec, port_mapping @@ -19,6 +18,7 @@ confirm_ask, console, ) +from dstack._internal.cli.utils.rich import MultiItemStatus from dstack._internal.cli.utils.run import get_runs_table, print_run_plan from dstack._internal.core.errors import ( CLIError, @@ -171,7 +171,7 @@ def apply_configuration( try: # We can attach to run multiple times if it goes from running to pending (retried). while True: - with console.status("") as live: + with MultiItemStatus(f"Launching [code]{run.name}[/]...", console=console) as live: while run.status in ( RunStatus.SUBMITTED, RunStatus.PENDING, @@ -179,11 +179,10 @@ def apply_configuration( RunStatus.TERMINATING, ): table = get_runs_table([run]) - live.update(Group(f"Launching [code]{run.name}[/]...\n", table)) + live.update(table) time.sleep(5) run.refresh() - console.print() console.print( get_runs_table( [run], diff --git a/src/dstack/_internal/cli/services/configurators/volume.py b/src/dstack/_internal/cli/services/configurators/volume.py index ef20ba4f5..41d961519 100644 --- a/src/dstack/_internal/cli/services/configurators/volume.py +++ b/src/dstack/_internal/cli/services/configurators/volume.py @@ -2,7 +2,6 @@ import time from typing import List -from rich.console import Group from rich.table import Table from dstack._internal.cli.services.configurators.base import BaseApplyConfigurator @@ -11,6 +10,7 @@ confirm_ask, console, ) +from dstack._internal.cli.utils.rich import MultiItemStatus from dstack._internal.cli.utils.volume import get_volumes_table from dstack._internal.core.errors import ResourceNotExistsError from dstack._internal.core.models.configurations import ApplyConfigurationType @@ -98,10 +98,12 @@ def apply_configuration( console.print("Volume configuration submitted. Exiting...") return try: - with console.status("") as live: + with MultiItemStatus( + f"Provisioning [code]{volume.name}[/]...", console=console + ) as live: while not _finished_provisioning(volume): table = get_volumes_table([volume]) - live.update(Group(f"Provisioning [code]{volume.name}[/]...\n", table)) + live.update(table) time.sleep(LIVE_TABLE_PROVISION_INTERVAL_SECS) volume = self.api.client.volumes.get(self.api.project, volume.name) except KeyboardInterrupt: @@ -113,7 +115,6 @@ def apply_configuration( else: console.print("Exiting... Volume provisioning will continue in the background.") return - console.print() console.print( get_volumes_table( [volume], diff --git a/src/dstack/_internal/cli/utils/rich.py b/src/dstack/_internal/cli/utils/rich.py index 891eef4d7..fb2895ab3 100644 --- a/src/dstack/_internal/cli/utils/rich.py +++ b/src/dstack/_internal/cli/utils/rich.py @@ -1,8 +1,12 @@ import logging from datetime import datetime +from types import TracebackType from typing import TYPE_CHECKING, Callable, Iterable, List, Optional, Union +from rich.console import Group +from rich.live import Live from rich.logging import RichHandler +from rich.spinner import Spinner from rich.text import Text from rich.traceback import Traceback @@ -122,3 +126,31 @@ def prepend_path(self, message, record): path = path + "[/link]" message = f"[log.path]{path}[/] " + message return message + + +class MultiItemStatus: + """An alternative to rich.status.Status that allows extra renderables below the spinner""" + + def __init__(self, status: "RenderableType", *, console: Optional["Console"] = None) -> None: + self._spinner = Spinner("dots", text=status, style="status.spinner") + self._live = Live( + renderable=self._spinner, + console=console, + refresh_per_second=12.5, + transient=True, + ) + + def update(self, *renderables: "RenderableType") -> None: + self._live.update(renderable=Group(self._spinner, *renderables)) + + def __enter__(self) -> "MultiItemStatus": + self._live.start() + return self + + def __exit__( + self, + exc_type: Optional[type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> None: + self._live.stop()