Fix wait_for_active treating SCALED_TO_ZERO as a deployment failure

truss/api/definitions.py

@@ -3,7 +3,7 @@
 import pydantic
 
 from truss.remote.baseten import service
-from truss.remote.baseten.core import ACTIVE_STATUS, DEPLOYING_STATUSES
+from truss.remote.baseten.core import READY_STATUSES, TERMINAL_FAILURE_STATUSES
 
 
 class ModelDeployment:
@@ -36,10 +36,10 @@ def wait_for_active(self, timeout_seconds: int = 600) -> bool:
             ):
                 raise TimeoutError("Deployment timed out.")
 
-            if deployment_status == ACTIVE_STATUS:
+            if deployment_status in READY_STATUSES:
                 return True
 
-            if deployment_status not in DEPLOYING_STATUSES:
+            if deployment_status in TERMINAL_FAILURE_STATUSES:
                 raise ValueError(f"Deployment failed with status: {deployment_status}")
 
         raise RuntimeError("Error polling deployment status.")

truss/cli/cli.py

@@ -34,8 +34,8 @@
 from truss.cli.utils import common, self_upgrade
 from truss.cli.utils.output import console, error_console, json_command
 from truss.remote.baseten.core import (
-    ACTIVE_STATUS,
-    DEPLOYING_STATUSES,
+    READY_STATUSES,
+    TERMINAL_FAILURE_STATUSES,
     ModelId,
     ModelIdentifier,
     ModelName,
@@ -1005,21 +1005,21 @@ def push(
                     f"[bold green]Deploying...Current Status: {deployment_status}"
                 )
 
-                if deployment_status == ACTIVE_STATUS:
+                if deployment_status in READY_STATUSES:
                     console.print("Deployment succeeded.", style="bold green")
                     break
 
                 # For --watch (dev deployments), enter watch mode early
                 # once past BUILDING, so user can iterate on code
-                if watch_after_push and deployment_status in ("LOADING_MODEL"):
+                if watch_after_push and deployment_status in ("LOADING_MODEL",):
                     console.print(
                         f"Deployment status: {deployment_status}. "
                         "Entering watch mode early for faster iteration...",
                         style="bold blue",
                     )
                     break
 
-                if deployment_status not in DEPLOYING_STATUSES:
+                if deployment_status in TERMINAL_FAILURE_STATUSES:
                     exc = RuntimeError(
                         f"Deployment failed with status {deployment_status}."
                     )

truss/remote/baseten/core.py

@@ -28,6 +28,18 @@
 
 DEPLOYING_STATUSES = ["BUILDING", "DEPLOYING", "LOADING_MODEL", "UPDATING"]
 ACTIVE_STATUS = "ACTIVE"
+SCALED_TO_ZERO_STATUS = "SCALED_TO_ZERO"
+READY_STATUSES = frozenset([ACTIVE_STATUS, SCALED_TO_ZERO_STATUS])
+TERMINAL_FAILURE_STATUSES = frozenset(
+    [
+        "BUILD_FAILED",
+        "BUILD_STOPPED",
+        "DEPLOY_FAILED",
+        "FAILED",
+        "INACTIVE",
+        "UNHEALTHY",
+    ]
+)
 NO_ENVIRONMENTS_EXIST_ERROR_MESSAGING = (
     "Model hasn't been deployed yet. No environments exist."
 )

truss/remote/baseten/utils/status.py

@@ -23,6 +23,7 @@
     "LOADING_MODEL",
     "ACTIVE",
     "UPDATING",
+    "SCALED_TO_ZERO",
     "WAKING_UP",
 ]
 

truss/tests/api/test_model_deployment.py

@@ -0,0 +1,55 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from truss.api.definitions import ModelDeployment
+from truss.remote.baseten.core import TERMINAL_FAILURE_STATUSES
+from truss.remote.baseten.service import BasetenService
+
+
+def _make_deployment(statuses):
+    mock_service = MagicMock(spec=BasetenService)
+    mock_service.model_id = "model_id"
+    mock_service.model_version_id = "version_id"
+    mock_service.poll_deployment_status.return_value = iter(statuses)
+    return ModelDeployment(mock_service)
+
+
+def test_wait_for_active_returns_true_on_active():
+    deployment = _make_deployment(["BUILDING", "DEPLOYING", "ACTIVE"])
+
+    assert deployment.wait_for_active() is True
+
+
+def test_wait_for_active_returns_true_on_scaled_to_zero():
+    deployment = _make_deployment(["BUILDING", "DEPLOYING", "SCALED_TO_ZERO"])
+
+    assert deployment.wait_for_active() is True
+
+
+def test_wait_for_active_keeps_polling_while_waking_up():
+    deployment = _make_deployment(["WAKING_UP", "WAKING_UP", "ACTIVE"])
+
+    assert deployment.wait_for_active() is True
+
+
+def test_wait_for_active_keeps_polling_on_unknown_status():
+    deployment = _make_deployment(["SOME_NEW_BACKEND_STATUS", "ACTIVE"])
+
+    assert deployment.wait_for_active() is True
+
+
+@pytest.mark.parametrize("status", sorted(TERMINAL_FAILURE_STATUSES))
+def test_wait_for_active_raises_on_terminal_failure(status):
+    deployment = _make_deployment(["BUILDING", status])
+
+    with pytest.raises(ValueError, match=f"Deployment failed with status: {status}"):
+        deployment.wait_for_active()
+
+
+def test_wait_for_active_raises_on_timeout():
+    deployment = _make_deployment(["SOME_NEW_BACKEND_STATUS", "ACTIVE"])
+
+    with patch("truss.api.definitions.time.time", side_effect=[0, 100]):
+        with pytest.raises(TimeoutError, match="Deployment timed out."):
+            deployment.wait_for_active(timeout_seconds=10)

truss/tests/cli/test_cli.py

@@ -1036,6 +1036,81 @@ def test_push_watch_with_tail_starts_background_tail(
     )
 
 
+def _invoke_push_watch(truss_dir, remote, statuses):
+    runner = CliRunner()
+
+    mock_service = MagicMock(spec=BasetenService)
+    mock_service.is_draft = True
+    mock_service.logs_url = "https://example.com/logs"
+    mock_service.model_id = "model_id"
+    mock_service.model_version_id = "version_id"
+    mock_service.poll_deployment.return_value = iter(
+        [{"status": status} for status in statuses]
+    )
+    remote.push = Mock(return_value=mock_service)
+
+    mock_resolve = Mock(
+        return_value=(
+            {
+                "id": "model_id",
+                "name": "model_name",
+                "hostname": "https://model.api.baseten.co",
+            },
+            [{"id": "version_id", "is_draft": True}],
+        )
+    )
+
+    with patch("truss.cli.cli.RemoteFactory.create", return_value=remote):
+        remote.api.get_teams = Mock(return_value={})
+        with patch("truss.cli.cli.resolve_model_team_name", return_value=(None, None)):
+            with patch("truss.cli.cli.resolve_model_for_watch", mock_resolve):
+                with patch("truss.cli.cli._start_watch_mode"):
+                    with patch("truss.cli.utils.common.start_keepalive"):
+                        return runner.invoke(
+                            truss_cli,
+                            [
+                                "push",
+                                str(truss_dir),
+                                "--remote",
+                                "baseten",
+                                "--model-name",
+                                "model_name",
+                                "--watch",
+                            ],
+                        )
+
+
+def test_push_watch_enters_watch_mode_early_on_loading_model(
+    custom_model_truss_dir_with_pre_and_post,
+    remote,
+    mock_baseten_requests,
+    mock_upload_truss,
+    mock_create_truss_service,
+):
+    result = _invoke_push_watch(
+        custom_model_truss_dir_with_pre_and_post, remote, ["LOADING_MODEL", "ACTIVE"]
+    )
+
+    assert result.exit_code == 0, result.output
+    assert "Entering watch mode early" in result.output
+
+
+def test_push_watch_ignores_status_that_is_substring_of_loading_model(
+    custom_model_truss_dir_with_pre_and_post,
+    remote,
+    mock_baseten_requests,
+    mock_upload_truss,
+    mock_create_truss_service,
+):
+    result = _invoke_push_watch(
+        custom_model_truss_dir_with_pre_and_post, remote, ["LOADING", "ACTIVE"]
+    )
+
+    assert result.exit_code == 0, result.output
+    assert "Entering watch mode early" not in result.output
+    assert "Deployment succeeded." in result.output
+
+
 def test_push_watch_with_publish_fails():
     """Test that --watch with --publish fails."""
     mock_truss = Mock()
@@ -1440,6 +1515,52 @@ def test_push_json_output_wait_success(
     assert data["deployment"] == deployment_response
 
 
+@pytest.mark.skipif(
+    sys.version_info < (3, 10),
+    reason="needs click>=8.2 CliRunner stdout/stderr split, only available on Python 3.10+",
+)
+def test_push_json_output_wait_scaled_to_zero(
+    custom_model_truss_dir_with_pre_and_post, remote
+):
+    runner = CliRunner()
+    deployment_response = {"status": "SCALED_TO_ZERO", "id": "deploy_id", "replicas": 0}
+    mock_service = _make_mock_service()
+    mock_service.poll_deployment.return_value = iter(
+        [{"status": "BUILDING"}, deployment_response]
+    )
+    remote.push = Mock(return_value=mock_service)
+    result = _invoke_push_json(
+        runner, custom_model_truss_dir_with_pre_and_post, remote, ["--wait"]
+    )
+
+    assert result.exit_code == 0
+    data = json.loads(result.stdout)
+    assert data["deployment"] == deployment_response
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 10),
+    reason="needs click>=8.2 CliRunner stdout/stderr split, only available on Python 3.10+",
+)
+def test_push_json_output_wait_unknown_status_keeps_polling(
+    custom_model_truss_dir_with_pre_and_post, remote
+):
+    runner = CliRunner()
+    deployment_response = {"status": "ACTIVE", "id": "deploy_id", "replicas": 1}
+    mock_service = _make_mock_service()
+    mock_service.poll_deployment.return_value = iter(
+        [{"status": "SOME_NEW_BACKEND_STATUS"}, deployment_response]
+    )
+    remote.push = Mock(return_value=mock_service)
+    result = _invoke_push_json(
+        runner, custom_model_truss_dir_with_pre_and_post, remote, ["--wait"]
+    )
+
+    assert result.exit_code == 0
+    data = json.loads(result.stdout)
+    assert data["deployment"] == deployment_response
+
+
 @pytest.mark.skipif(
     sys.version_info < (3, 10),
     reason="needs click>=8.2 CliRunner stdout/stderr split, only available on Python 3.10+",

truss/tests/cli/test_model_log_watcher.py

@@ -0,0 +1,48 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+from truss.cli.logs.model_log_watcher import ModelDeploymentLogWatcher
+from truss.remote.baseten.api import BasetenApi
+
+
+def _make_watcher(status):
+    api = MagicMock(spec=BasetenApi)
+    api.get_deployment.return_value = {"status": status, "is_development": False}
+    watcher = ModelDeploymentLogWatcher(api, "model_id", "deployment_id")
+    watcher.before_polling()
+    return watcher
+
+
+@pytest.mark.parametrize(
+    "status",
+    ["BUILDING", "DEPLOYING", "LOADING_MODEL", "ACTIVE", "UPDATING", "WAKING_UP"],
+)
+def test_should_poll_again_while_running(status):
+    assert _make_watcher(status).should_poll_again() is True
+
+
+def test_should_poll_again_while_scaled_to_zero():
+    assert _make_watcher("SCALED_TO_ZERO").should_poll_again() is True
+
+
+@pytest.mark.parametrize(
+    "status", ["BUILD_FAILED", "DEPLOY_FAILED", "INACTIVE", "DEACTIVATING"]
+)
+def test_should_not_poll_again_once_stopped(status):
+    assert _make_watcher(status).should_poll_again() is False
+
+
+def test_post_poll_refreshes_status():
+    api = MagicMock(spec=BasetenApi)
+    api.get_deployment.side_effect = [
+        {"status": "SCALED_TO_ZERO", "is_development": False},
+        {"status": "INACTIVE", "is_development": False},
+    ]
+    watcher = ModelDeploymentLogWatcher(api, "model_id", "deployment_id")
+
+    watcher.before_polling()
+    assert watcher.should_poll_again() is True
+
+    watcher.post_poll()
+    assert watcher.should_poll_again() is False