Support dynamic run waiting CLI status with extra renderables

Author: r4victor

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

@@ -28,7 +28,12 @@
 from dstack._internal.cli.services.resources import apply_resources_args, register_resources_args
 from dstack._internal.cli.utils.common import 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.cli.utils.run import (
+    RunWaitStatus,
+    get_run_wait_status,
+    get_runs_table,
+    print_run_plan,
+)
 from dstack._internal.core.errors import (
     CLIError,
     ConfigurationError,
@@ -192,10 +197,14 @@ def apply_configuration(
         try:
             # We can attach to run multiple times if it goes from running to pending (retried).
             while True:
-                with MultiItemStatus(f"Launching [code]{run.name}[/]...", console=console) as live:
+                with MultiItemStatus(_get_apply_status(run), console=console) as live:
                     while not _is_ready_to_attach(run):
                         table = get_runs_table([run])
-                        live.update(table)
+                        live.update(
+                            table,
+                            *_get_apply_wait_renderables(run),
+                            status=_get_apply_status(run),
+                        )
                         time.sleep(5)
                         run.refresh()
 
@@ -801,6 +810,26 @@ def _print_service_urls(run: Run) -> None:
     console.print()
 
 
+def _get_apply_status(run: Run) -> str:
+    wait_status = get_run_wait_status(run._run)
+    if wait_status is None:
+        return f"Launching [code]{run.name}[/]..."
+    return f"[code]{run.name}[/] is {wait_status.value}..."
+
+
+def _get_apply_wait_renderables(run: Run) -> list[str]:
+    wait_status = get_run_wait_status(run._run)
+    if wait_status is RunWaitStatus.WAITING_FOR_REQUESTS and run._run.service is not None:
+        return [f"Service URL: [link={run.service_url}]{run.service_url}[/]"]
+    if (
+        wait_status is RunWaitStatus.WAITING_FOR_SCHEDULE
+        and run._run.next_triggered_at is not None
+    ):
+        next_run = run._run.next_triggered_at.astimezone().strftime("%Y-%m-%d %H:%M %Z")
+        return [f"Next run: {next_run}"]
+    return []
+
+
 def _print_dev_environment_connection_info(run: Run) -> None:
     if not FeatureFlags.CLI_PRINT_JOB_CONNECTION_INFO:
         return

src/dstack/_internal/cli/utils/rich.py

@@ -140,7 +140,11 @@ def __init__(self, status: "RenderableType", *, console: Optional["Console"] = N
             transient=True,
         )
 
-    def update(self, *renderables: "RenderableType") -> None:
+    def update(
+        self, *renderables: "RenderableType", status: Optional["RenderableType"] = None
+    ) -> None:
+        if status is not None:
+            self._spinner.update(text=status)
         self._live.update(renderable=Group(self._spinner, *renderables))
 
     def __enter__(self) -> "MultiItemStatus":

src/dstack/_internal/cli/utils/run.py

@@ -1,4 +1,5 @@
 import shutil
+from enum import Enum
 from typing import Any, Dict, List, Optional
 
 from rich.markup import escape
@@ -49,6 +50,11 @@
 from dstack.api import Run
 
 
+class RunWaitStatus(str, Enum):
+    WAITING_FOR_REQUESTS = "waiting for requests"
+    WAITING_FOR_SCHEDULE = "waiting for schedule"
+
+
 def print_offers_json(run_plan: RunPlan, run_spec):
     """Print offers information in JSON format."""
     job_plan = run_plan.job_plans[0]
@@ -200,6 +206,40 @@ def th(s: str) -> str:
         console.print(NO_FLEETS_WARNING if no_fleets else NO_OFFERS_WARNING)
 
 
+def get_run_wait_status(run: CoreRun) -> Optional[RunWaitStatus]:
+    # Only synthesize a CLI-specific waiting state when the server did not provide
+    # a more specific run-level message such as "retrying".
+    if run.status_message not in ("", run.status.value):
+        return None
+
+    if run.status == RunStatus.PENDING and run.next_triggered_at is not None:
+        return RunWaitStatus.WAITING_FOR_SCHEDULE
+
+    if _is_waiting_for_requests(run):
+        return RunWaitStatus.WAITING_FOR_REQUESTS
+
+    return None
+
+
+def _is_waiting_for_requests(run: CoreRun) -> bool:
+    if run.run_spec.configuration.type != "service":
+        return False
+    if run.service is None or run.next_triggered_at is not None:
+        return False
+    if run.status not in (RunStatus.SUBMITTED, RunStatus.PENDING):
+        return False
+    return not any(_is_job_active(job.job_submissions[-1].status) for job in run.jobs)
+
+
+def _is_job_active(status: JobStatus) -> bool:
+    return status in (
+        JobStatus.SUBMITTED,
+        JobStatus.PROVISIONING,
+        JobStatus.PULLING,
+        JobStatus.RUNNING,
+    )
+
+
 def _format_run_status(run) -> str:
     status_text = (
         run.latest_job_submission.status_message

src/tests/_internal/cli/services/configurators/test_run.py

@@ -1,4 +1,5 @@
 import argparse
+from datetime import datetime, timezone
 from textwrap import dedent
 from typing import List, Optional, Tuple
 from unittest.mock import Mock
@@ -9,6 +10,8 @@
 from dstack._internal.cli.services.configurators import get_run_configurator_class
 from dstack._internal.cli.services.configurators.run import (
     BaseRunConfigurator,
+    _get_apply_status,
+    _get_apply_wait_renderables,
     render_run_spec_diff,
 )
 from dstack._internal.core.errors import ConfigurationError
@@ -18,11 +21,25 @@
     BaseRunConfiguration,
     DevEnvironmentConfiguration,
     PortMapping,
+    ScalingSpec,
+    ServiceConfiguration,
     TaskConfiguration,
 )
 from dstack._internal.core.models.envs import Env
 from dstack._internal.core.models.profiles import Profile
-from dstack._internal.server.testing.common import get_run_spec
+from dstack._internal.core.models.resources import Range
+from dstack._internal.core.models.runs import RunStatus, ServiceSpec
+from dstack._internal.server.services import encryption  # noqa: F401  # import for side-effect
+from dstack._internal.server.services.runs import run_model_to_run
+from dstack._internal.server.testing.common import (
+    create_project,
+    create_repo,
+    create_run,
+    create_user,
+    get_run_spec,
+)
+from dstack.api import Run
+from dstack.api.server import APIClient
 
 
 class TestApplyArgs:
@@ -401,3 +418,74 @@ def test_no_diff(self):
         old = get_run_spec(run_name="test", repo_id="test")
         new = get_run_spec(run_name="test", repo_id="test")
         assert render_run_spec_diff(old, new) is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+class TestApplyStatusHelpers:
+    async def test_waiting_for_requests_status_and_renderables(self, session):
+        project = await create_project(session=session)
+        user = await create_user(session=session)
+        repo = await create_repo(session=session, project_id=project.id)
+        run_spec = get_run_spec(
+            run_name="service-run",
+            repo_id=repo.name,
+            configuration=ServiceConfiguration(
+                type="service",
+                image="ubuntu:latest",
+                commands=["echo hello"],
+                port=80,
+                replicas=Range[int](min=0, max=1),
+                scaling=ScalingSpec(metric="rps", target=1),
+            ),
+        )
+        run_model = await create_run(
+            session=session,
+            project=project,
+            repo=repo,
+            user=user,
+            run_name="service-run",
+            run_spec=run_spec,
+            status=RunStatus.PENDING,
+        )
+        run_model.service_spec = ServiceSpec(url="/proxy/services/test/service-run/").json()
+        await session.commit()
+        await session.refresh(run_model)
+
+        api_run = Run(
+            api_client=Mock(spec=APIClient, base_url="http://127.0.0.1:3000"),
+            project=project.name,
+            run=run_model_to_run(run_model),
+        )
+
+        assert _get_apply_status(api_run) == "[code]service-run[/] is waiting for requests..."
+        assert _get_apply_wait_renderables(api_run) == [
+            "Service URL: [link=http://127.0.0.1:3000/proxy/services/test/service-run/]http://127.0.0.1:3000/proxy/services/test/service-run/[/]"
+        ]
+
+    async def test_waiting_for_schedule_status_and_renderables(self, session):
+        project = await create_project(session=session)
+        user = await create_user(session=session)
+        repo = await create_repo(session=session, project_id=project.id)
+        run_model = await create_run(
+            session=session,
+            project=project,
+            repo=repo,
+            user=user,
+            run_name="scheduled-run",
+            status=RunStatus.PENDING,
+            next_triggered_at=datetime(2023, 1, 2, 3, 10, tzinfo=timezone.utc),
+        )
+        await session.refresh(run_model)
+
+        api_run = Run(
+            api_client=Mock(spec=APIClient),
+            project=project.name,
+            run=run_model_to_run(run_model),
+        )
+        next_run = datetime(2023, 1, 2, 3, 10, tzinfo=timezone.utc)
+        api_run._run.next_triggered_at = next_run
+
+        assert _get_apply_status(api_run) == "[code]scheduled-run[/] is waiting for schedule..."
+        expected_next_run = next_run.astimezone().strftime("%Y-%m-%d %H:%M %Z")
+        assert _get_apply_wait_renderables(api_run) == [f"Next run: {expected_next_run}"]

src/tests/_internal/cli/utils/test_run.py

@@ -10,7 +10,9 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import selectinload
 
-from dstack._internal.cli.utils.run import get_runs_table
+from dstack._internal.cli.utils.run import (
+    get_runs_table,
+)
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.configurations import (
     AnyRunConfiguration,