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,7 +1005,7 @@ 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
 
@@ -1019,7 +1019,7 @@ def push(
                     )
                     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,16 @@
 
 DEPLOYING_STATUSES = ["BUILDING", "DEPLOYING", "LOADING_MODEL", "UPDATING"]
 ACTIVE_STATUS = "ACTIVE"
+SCALED_TO_ZERO_STATUS = "SCALED_TO_ZERO"
+READY_STATUSES = [ACTIVE_STATUS, SCALED_TO_ZERO_STATUS]
+TERMINAL_FAILURE_STATUSES = [
+    "BUILD_FAILED",
+    "BUILD_STOPPED",
+    "DEPLOY_FAILED",
+    "FAILED",
+    "INACTIVE",
+    "UNHEALTHY",
+]
 NO_ENVIRONMENTS_EXIST_ERROR_MESSAGING = (
     "Model hasn't been deployed yet. No environments exist."
 )

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", 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

@@ -1440,6 +1440,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+",