diff --git a/truss/api/definitions.py b/truss/api/definitions.py index 2cbf1c105..d48994a77 100644 --- a/truss/api/definitions.py +++ b/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.") diff --git a/truss/cli/cli.py b/truss/cli/cli.py index c147342a2..24ce95701 100644 --- a/truss/cli/cli.py +++ b/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,13 +1005,13 @@ 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...", @@ -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}." ) diff --git a/truss/remote/baseten/core.py b/truss/remote/baseten/core.py index 2249f6748..817d3794a 100644 --- a/truss/remote/baseten/core.py +++ b/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." ) diff --git a/truss/remote/baseten/utils/status.py b/truss/remote/baseten/utils/status.py index a8ff75d2b..11b87e0b4 100644 --- a/truss/remote/baseten/utils/status.py +++ b/truss/remote/baseten/utils/status.py @@ -23,6 +23,7 @@ "LOADING_MODEL", "ACTIVE", "UPDATING", + "SCALED_TO_ZERO", "WAKING_UP", ] diff --git a/truss/tests/api/test_model_deployment.py b/truss/tests/api/test_model_deployment.py new file mode 100644 index 000000000..ba3541bd7 --- /dev/null +++ b/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) diff --git a/truss/tests/cli/test_cli.py b/truss/tests/cli/test_cli.py index cd122e21f..ee425bf27 100644 --- a/truss/tests/cli/test_cli.py +++ b/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+", diff --git a/truss/tests/cli/test_model_log_watcher.py b/truss/tests/cli/test_model_log_watcher.py new file mode 100644 index 000000000..c282f0d9c --- /dev/null +++ b/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