diff --git a/truss/cli/cli.py b/truss/cli/cli.py index a5a6e918f..3311c31db 100644 --- a/truss/cli/cli.py +++ b/truss/cli/cli.py @@ -26,6 +26,7 @@ ModelVersionId, ) from truss.remote.baseten.service import BasetenService +from truss.remote.baseten.utils.status import get_displayable_status from truss.remote.remote_cli import inquire_model_name, inquire_remote_name from truss.remote.remote_factory import USER_TRUSSRC_PATH, RemoteFactory from truss.truss_config import Build, ModelServer @@ -284,18 +285,19 @@ def _create_chains_table(service) -> Tuple[rich.table.Table, List[str]]: statuses = [] # After reversing, the first one is the entrypoint (per order in service). for i, (name, status, logs_url) in enumerate(reversed(service.get_info())): - if status == ACTIVE_STATUS: + displayable_status = get_displayable_status(status) + if displayable_status == ACTIVE_STATUS: spinner_name = "active" - elif status in DEPLOYING_STATUSES: - if status == "BUILDING": + elif displayable_status in DEPLOYING_STATUSES: + if displayable_status == "BUILDING": spinner_name = "building" - elif status == "LOADING": + elif displayable_status == "LOADING": spinner_name = "loading" else: spinner_name = "deploying" else: spinner_name = "failed" - spinner = rich.spinner.Spinner(spinner_name, text=status) + spinner = rich.spinner.Spinner(spinner_name, text=displayable_status) if i == 0: display_name = f"{name} (entrypoint)" else: @@ -304,7 +306,7 @@ def _create_chains_table(service) -> Tuple[rich.table.Table, List[str]]: table.add_row(spinner, display_name, logs_url) if i == 0: # Add section divider after entrypoint. table.add_section() - statuses.append(status) + statuses.append(displayable_status) return table, statuses diff --git a/truss/remote/baseten/utils/status.py b/truss/remote/baseten/utils/status.py new file mode 100644 index 000000000..d7ffbb601 --- /dev/null +++ b/truss/remote/baseten/utils/status.py @@ -0,0 +1,29 @@ +STATUS_TO_DISPLAYABLE = { + "BUILDING_MODEL": "BUILDING", + "DEPLOYING_MODEL": "DEPLOYING", + "MODEL_DEPLOY_FAILED": "DEPLOY_FAILED", + "MODEL_LOADING": "LOADING_MODEL", + "MODEL_READY": "ACTIVE", + "MODEL_UNHEALTHY": "UNHEALTHY", + "BUILDING_MODEL_FAILED": "BUILD_FAILED", + "BUILDING_MODEL_STOPPED": "BUILD_STOPPED", + "DEACTIVATING_MODEL": "DEACTIVATING", + "DEACTIVATED_MODEL": "INACTIVE", + "MODEL_DNE_ERROR": "FAILED", + "UPDATING": "UPDATING", + "MIGRATING_WORKLOAD_PLANES": "UPDATING", + "SCALED_TO_ZERO": "SCALED_TO_ZERO", + "SCALING_FROM_ZERO": "WAKING_UP", +} + + +def get_displayable_status(status: str) -> str: + """ + TODO: Remove this method once Chains is supported in the REST API + + This is used by the `truss chains deploy` command right now to + print the right status. Once Chains are supported by the REST API, the + Baseten REST API will return status strings matching the ones here, so we don't + need to do any mapping. + """ + return STATUS_TO_DISPLAYABLE[status]