import sys from copy import deepcopy from typing import Any, List, Optional, Tuple from unittest.mock import MagicMock, Mock, patch import pytest from ray._common.ray_constants import DEFAULT_MAX_CONCURRENCY_ASYNC from ray._raylet import NodeID from ray.serve._private.autoscaling_state import AutoscalingStateManager from ray.serve._private.common import ( RUNNING_REQUESTS_KEY, DeploymentHandleSource, DeploymentID, DeploymentStatus, DeploymentStatusTrigger, GangReservationResult, HandleMetricReport, ReplicaID, ReplicaMetricReport, ReplicaState, TargetCapacityDirection, TimeStampedValue, ) from ray.serve._private.config import DeploymentConfig, ReplicaConfig from ray.serve._private.constants import ( DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_S, DEFAULT_GRACEFUL_SHUTDOWN_WAIT_LOOP_S, DEFAULT_HEALTH_CHECK_PERIOD_S, DEFAULT_HEALTH_CHECK_TIMEOUT_S, DEFAULT_MAX_ONGOING_REQUESTS, RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE, RAY_SERVE_INTERNAL_DEPLOYMENT_ACTOR_NAME_ENV_VAR, RAY_SERVE_INTERNAL_DEPLOYMENT_APP_NAME_ENV_VAR, RAY_SERVE_INTERNAL_DEPLOYMENT_CODE_VERSION_ENV_VAR, RAY_SERVE_INTERNAL_DEPLOYMENT_NAME_ENV_VAR, RAY_SERVE_STATUS_GAUGE_REPORT_INTERVAL_S, ) from ray.serve._private.deployment_info import DeploymentInfo from ray.serve._private.deployment_state import ( ALL_REPLICA_STATES, CHECKPOINT_KEY, SLOW_STARTUP_WARNING_S, ActorReplicaWrapper, DeploymentActorContainer, DeploymentActorState, DeploymentActorWrapper, DeploymentReplica, DeploymentState, DeploymentStateManager, DeploymentVersion, ReplicaStartupStatus, ReplicaStateContainer, ) from ray.serve._private.exceptions import DeploymentIsBeingDeletedError from ray.serve._private.long_poll import LongPollNamespace from ray.serve._private.test_utils import ( MockDeploymentActorWrapper, MockPlacementGroup, dead_replicas_context, replica_rank_context, uninitialized_replicas_context, ) from ray.serve._private.utils import ( get_capacity_adjusted_num_replicas, get_random_string, ) from ray.serve.config import DeploymentActorConfig, GangSchedulingConfig from ray.serve.schema import ReplicaRank from ray.util.placement_group import validate_placement_group TEST_DEPLOYMENT_ID = DeploymentID(name="test_deployment", app_name="test_app") TEST_DEPLOYMENT_ID_2 = DeploymentID(name="test_deployment_2", app_name="test_app") def deployment_info( version: Optional[str] = None, num_replicas: Optional[int] = 1, user_config: Optional[Any] = None, replica_config: Optional[ReplicaConfig] = None, **config_opts, ) -> Tuple[DeploymentInfo, DeploymentVersion]: info = DeploymentInfo( version=version, start_time_ms=0, actor_name="abc", deployment_config=DeploymentConfig( num_replicas=num_replicas, user_config=user_config, **config_opts ), replica_config=replica_config or ReplicaConfig.create(lambda x: x), deployer_job_id="", ) if version is not None: code_version = version else: code_version = get_random_string() version = DeploymentVersion( code_version, info.deployment_config, info.replica_config.ray_actor_options ) return info, version def deployment_version(code_version) -> DeploymentVersion: return DeploymentVersion(code_version, DeploymentConfig(), {}) @pytest.fixture def mock_max_per_replica_retry_count(): with patch("ray.serve._private.deployment_state.MAX_PER_REPLICA_RETRY_COUNT", 2): yield 2 class FakeDeploymentReplica: """Fakes the DeploymentReplica class.""" def __init__(self, version: DeploymentVersion): self._version = version self._replica_id = ReplicaID( get_random_string(), deployment_id=DeploymentID(name="fake") ) @property def replica_id(self): return self._replica_id @property def version(self): return self._version def update_state(self, state): pass def replica(version: Optional[DeploymentVersion] = None) -> FakeDeploymentReplica: version = version or DeploymentVersion(get_random_string(), DeploymentConfig(), {}) return FakeDeploymentReplica(version) class TestReplicaStateContainer: def test_count(self): c = ReplicaStateContainer() r1, r2, r3 = ( replica(deployment_version("1")), replica(deployment_version("2")), replica(deployment_version("2")), ) c.add(ReplicaState.STARTING, r1) c.add(ReplicaState.STARTING, r2) c.add(ReplicaState.STOPPING, r3) assert c.count() == 3 # Test filtering by state. assert c.count() == c.count( states=[ReplicaState.STARTING, ReplicaState.STOPPING] ) assert c.count(states=[ReplicaState.STARTING]) == 2 assert c.count(states=[ReplicaState.STOPPING]) == 1 # Test filtering by version. assert c.count(version=deployment_version("1")) == 1 assert c.count(version=deployment_version("2")) == 2 assert c.count(version=deployment_version("3")) == 0 assert c.count(exclude_version=deployment_version("1")) == 2 assert c.count(exclude_version=deployment_version("2")) == 1 assert c.count(exclude_version=deployment_version("3")) == 3 # Test filtering by state and version. assert ( c.count(version=deployment_version("1"), states=[ReplicaState.STARTING]) == 1 ) assert ( c.count(version=deployment_version("3"), states=[ReplicaState.STARTING]) == 0 ) assert ( c.count( version=deployment_version("2"), states=[ReplicaState.STARTING, ReplicaState.STOPPING], ) == 2 ) assert ( c.count( exclude_version=deployment_version("1"), states=[ReplicaState.STARTING] ) == 1 ) assert ( c.count( exclude_version=deployment_version("3"), states=[ReplicaState.STARTING] ) == 2 ) assert ( c.count( exclude_version=deployment_version("2"), states=[ReplicaState.STARTING, ReplicaState.STOPPING], ) == 1 ) def test_get(self): c = ReplicaStateContainer() r1, r2, r3 = replica(), replica(), replica() c.add(ReplicaState.STARTING, r1) c.add(ReplicaState.STARTING, r2) c.add(ReplicaState.STOPPING, r3) assert c.get() == [r1, r2, r3] assert c.get() == c.get([ReplicaState.STARTING, ReplicaState.STOPPING]) assert c.get([ReplicaState.STARTING]) == [r1, r2] assert c.get([ReplicaState.STOPPING]) == [r3] def test_get_by_id(self): c = ReplicaStateContainer() r1, r2, r3 = replica(), replica(), replica() c.add(ReplicaState.STARTING, r1) c.add(ReplicaState.RUNNING, r2) c.add(ReplicaState.STOPPING, r3) # Found: each replica is retrievable by its ID regardless of state. assert c.get_by_id(r1.replica_id) is r1 assert c.get_by_id(r2.replica_id) is r2 assert c.get_by_id(r3.replica_id) is r3 # Not found: a replica ID that was never added returns None. unknown = replica() assert c.get_by_id(unknown.replica_id) is None # After pop: popped replicas are no longer in the index. popped = c.pop(states=[ReplicaState.RUNNING]) assert popped == [r2] assert c.get_by_id(r2.replica_id) is None # Remaining replicas are still found. assert c.get_by_id(r1.replica_id) is r1 assert c.get_by_id(r3.replica_id) is r3 # Pop everything and verify the index is fully cleared. c.pop() assert c.get_by_id(r1.replica_id) is None assert c.get_by_id(r3.replica_id) is None def test_pop_basic(self): c = ReplicaStateContainer() r1, r2, r3 = replica(), replica(), replica() c.add(ReplicaState.STARTING, r1) c.add(ReplicaState.STARTING, r2) c.add(ReplicaState.STOPPING, r3) assert c.pop() == [r1, r2, r3] assert not c.pop() def test_pop_exclude_version(self): c = ReplicaStateContainer() r1, r2, r3 = ( replica(deployment_version("1")), replica(deployment_version("1")), replica(deployment_version("2")), ) c.add(ReplicaState.STARTING, r1) c.add(ReplicaState.STARTING, r2) c.add(ReplicaState.STARTING, r3) assert c.pop(exclude_version=deployment_version("1")) == [r3] assert not c.pop(exclude_version=deployment_version("1")) assert c.pop(exclude_version=deployment_version("2")) == [r1, r2] assert not c.pop(exclude_version=deployment_version("2")) assert not c.pop() def test_pop_max_replicas(self): c = ReplicaStateContainer() r1, r2, r3 = replica(), replica(), replica() c.add(ReplicaState.STARTING, r1) c.add(ReplicaState.STARTING, r2) c.add(ReplicaState.STOPPING, r3) assert not c.pop(max_replicas=0) assert len(c.pop(max_replicas=1)) == 1 assert len(c.pop(max_replicas=2)) == 2 c.add(ReplicaState.STARTING, r1) c.add(ReplicaState.STARTING, r2) c.add(ReplicaState.STOPPING, r3) assert len(c.pop(max_replicas=10)) == 3 def test_pop_states(self): c = ReplicaStateContainer() r1, r2, r3, r4 = replica(), replica(), replica(), replica() # Check popping single state. c.add(ReplicaState.STOPPING, r1) c.add(ReplicaState.STARTING, r2) c.add(ReplicaState.STOPPING, r3) assert c.pop(states=[ReplicaState.STARTING]) == [r2] assert not c.pop(states=[ReplicaState.STARTING]) assert c.pop(states=[ReplicaState.STOPPING]) == [r1, r3] assert not c.pop(states=[ReplicaState.STOPPING]) # Check popping multiple states. Ordering of states should be # preserved. c.add(ReplicaState.STOPPING, r1) c.add(ReplicaState.STARTING, r2) c.add(ReplicaState.STOPPING, r3) c.add(ReplicaState.STARTING, r4) assert c.pop(states=[ReplicaState.STOPPING, ReplicaState.STARTING]) == [ r1, r3, r2, r4, ] assert not c.pop(states=[ReplicaState.STOPPING, ReplicaState.STARTING]) assert not c.pop(states=[ReplicaState.STOPPING]) assert not c.pop(states=[ReplicaState.STARTING]) assert not c.pop() def test_pop_integration(self): c = ReplicaStateContainer() r1, r2, r3, r4 = ( replica(deployment_version("1")), replica(deployment_version("2")), replica(deployment_version("2")), replica(deployment_version("3")), ) c.add(ReplicaState.STOPPING, r1) c.add(ReplicaState.STARTING, r2) c.add(ReplicaState.RUNNING, r3) c.add(ReplicaState.RUNNING, r4) assert not c.pop( exclude_version=deployment_version("1"), states=[ReplicaState.STOPPING] ) assert c.pop( exclude_version=deployment_version("1"), states=[ReplicaState.RUNNING], max_replicas=1, ) == [r3] assert c.pop( exclude_version=deployment_version("1"), states=[ReplicaState.RUNNING], max_replicas=1, ) == [r4] c.add(ReplicaState.RUNNING, r3) c.add(ReplicaState.RUNNING, r4) assert c.pop( exclude_version=deployment_version("1"), states=[ReplicaState.RUNNING] ) == [r3, r4] assert c.pop( exclude_version=deployment_version("1"), states=[ReplicaState.STARTING] ) == [r2] c.add(ReplicaState.STARTING, r2) c.add(ReplicaState.RUNNING, r3) c.add(ReplicaState.RUNNING, r4) assert c.pop( exclude_version=deployment_version("1"), states=[ReplicaState.RUNNING, ReplicaState.STARTING], ) == [r3, r4, r2] assert c.pop( exclude_version=deployment_version("nonsense"), states=[ReplicaState.STOPPING], ) == [r1] def _mock_deployment_actor_wrapper(deployment_id, code_version: str, name: str): """Create a MockDeploymentActorWrapper for container tests.""" config = DeploymentActorConfig(name=name, actor_class="builtins:object") return MockDeploymentActorWrapper(deployment_id, config, code_version) class TestDeploymentActorContainer: def test_add_and_count(self): dep_id = DeploymentID(name="test", app_name="app") c = DeploymentActorContainer(dep_id) w1 = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_a") w2 = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_b") w3 = _mock_deployment_actor_wrapper(dep_id, "v2", "actor_a") c.add(DeploymentActorState.STARTING, w1) c.add(DeploymentActorState.STARTING, w2) c.add(DeploymentActorState.RECOVERING, w3) assert c.count() == 3 assert c.count(states=[DeploymentActorState.STARTING]) == 2 assert c.count(states=[DeploymentActorState.RECOVERING]) == 1 assert c.count(code_version="v1") == 2 assert c.count(code_version="v2") == 1 assert c.count(code_version="v3") == 0 def test_is_empty(self): dep_id = DeploymentID(name="test", app_name="app") c = DeploymentActorContainer(dep_id) assert c.is_empty() w = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_a") c.add(DeploymentActorState.STARTING, w) assert not c.is_empty() c.pop() # remove all tracked actors assert c.is_empty() def test_add_moves_existing(self): """Adding same (code_version, name) moves from old state to new.""" dep_id = DeploymentID(name="test", app_name="app") c = DeploymentActorContainer(dep_id) w = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_a") c.add(DeploymentActorState.STARTING, w) assert c.count() == 1 assert c.count(states=[DeploymentActorState.STARTING]) == 1 c.add(DeploymentActorState.RUNNING, w) assert c.count() == 1 assert c.count(states=[DeploymentActorState.STARTING]) == 0 assert c.count(states=[DeploymentActorState.RUNNING]) == 1 def test_get(self): dep_id = DeploymentID(name="test", app_name="app") c = DeploymentActorContainer(dep_id) w1 = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_a") w2 = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_b") w3 = _mock_deployment_actor_wrapper(dep_id, "v2", "actor_a") c.add(DeploymentActorState.STARTING, w1) c.add(DeploymentActorState.STARTING, w2) c.add(DeploymentActorState.RECOVERING, w3) assert set(c.get()) == {w1, w2, w3} assert set(c.get(states=[DeploymentActorState.STARTING])) == {w1, w2} assert c.get(states=[DeploymentActorState.RECOVERING]) == [w3] assert set(c.get(code_version="v1")) == {w1, w2} assert c.get(code_version="v2") == [w3] assert c.get(code_version="v3") == [] def test_pop(self): dep_id = DeploymentID(name="test", app_name="app") c = DeploymentActorContainer(dep_id) w1 = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_a") w2 = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_b") w3 = _mock_deployment_actor_wrapper(dep_id, "v2", "actor_a") c.add(DeploymentActorState.STARTING, w1) c.add(DeploymentActorState.RUNNING, w2) c.add(DeploymentActorState.RUNNING, w3) # Pop by code_version removed = c.pop(code_version="v1") assert len(removed) == 2 states_removed = {s for s, _ in removed} assert states_removed == { DeploymentActorState.STARTING, DeploymentActorState.RUNNING, } assert c.count() == 1 assert c.get() == [w3] # Pop remaining removed = c.pop(code_version="v2") assert len(removed) == 1 assert removed[0][1].wrapper is w3 assert c.count() == 0 def test_pop_by_states(self): dep_id = DeploymentID(name="test", app_name="app") c = DeploymentActorContainer(dep_id) w1 = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_a") w2 = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_b") c.add(DeploymentActorState.STARTING, w1) c.add(DeploymentActorState.RUNNING, w2) removed = c.pop(code_version="v1", states=[DeploymentActorState.STARTING]) assert len(removed) == 1 assert removed[0][1].wrapper is w1 assert c.count() == 1 assert c.get() == [w2] def test_pop_all_versions(self): dep_id = DeploymentID(name="test", app_name="app") c = DeploymentActorContainer(dep_id) w1 = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_a") w2 = _mock_deployment_actor_wrapper(dep_id, "v2", "actor_b") c.add(DeploymentActorState.RUNNING, w1) c.add(DeploymentActorState.RUNNING, w2) removed = c.pop() assert len(removed) == 2 assert c.count() == 0 def test_get_wrapper(self): dep_id = DeploymentID(name="test", app_name="app") c = DeploymentActorContainer(dep_id) w1 = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_a") w2 = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_b") w3 = _mock_deployment_actor_wrapper(dep_id, "v2", "actor_a") c.add(DeploymentActorState.RUNNING, w1) c.add(DeploymentActorState.RUNNING, w2) c.add(DeploymentActorState.RUNNING, w3) assert c.get_wrapper("v1", "actor_a") is w1 assert c.get_wrapper("v1", "actor_b") is w2 assert c.get_wrapper("v2", "actor_a") is w3 assert c.get_wrapper("v1", "actor_c") is None assert c.get_wrapper("v3", "actor_a") is None # After pop, get_wrapper returns None c.pop(code_version="v1") assert c.get_wrapper("v1", "actor_a") is None assert c.get_wrapper("v1", "actor_b") is None assert c.get_wrapper("v2", "actor_a") is w3 def test_get_code_versions(self): dep_id = DeploymentID(name="test", app_name="app") c = DeploymentActorContainer(dep_id) assert c.get_code_versions() == set() w1 = _mock_deployment_actor_wrapper(dep_id, "v1", "actor_a") w2 = _mock_deployment_actor_wrapper(dep_id, "v2", "actor_b") c.add(DeploymentActorState.RUNNING, w1) c.add(DeploymentActorState.RUNNING, w2) assert c.get_code_versions() == {"v1", "v2"} c.pop(code_version="v1") assert c.get_code_versions() == {"v2"} class TestDeploymentActorWrapper: @pytest.mark.parametrize( "deployment_runtime_env,actor_runtime_env,conflicting_keys", [ ( { "env_vars": { RAY_SERVE_INTERNAL_DEPLOYMENT_APP_NAME_ENV_VAR: "user-app", } }, None, [RAY_SERVE_INTERNAL_DEPLOYMENT_APP_NAME_ENV_VAR], ), ( None, { "env_vars": { RAY_SERVE_INTERNAL_DEPLOYMENT_NAME_ENV_VAR: "user-deployment", RAY_SERVE_INTERNAL_DEPLOYMENT_ACTOR_NAME_ENV_VAR: "user-actor", } }, [ RAY_SERVE_INTERNAL_DEPLOYMENT_ACTOR_NAME_ENV_VAR, RAY_SERVE_INTERNAL_DEPLOYMENT_NAME_ENV_VAR, ], ), ( { "env_vars": { RAY_SERVE_INTERNAL_DEPLOYMENT_CODE_VERSION_ENV_VAR: "user-v1", } }, { "env_vars": { RAY_SERVE_INTERNAL_DEPLOYMENT_NAME_ENV_VAR: "user-deployment", } }, [ RAY_SERVE_INTERNAL_DEPLOYMENT_CODE_VERSION_ENV_VAR, RAY_SERVE_INTERNAL_DEPLOYMENT_NAME_ENV_VAR, ], ), ], ) def test_start_rejects_reserved_internal_context_env_vars( self, deployment_runtime_env, actor_runtime_env, conflicting_keys, ): mock_ready_ref = MagicMock() mock_handle = MagicMock() mock_handle.__ray_ready__ = MagicMock( remote=MagicMock(return_value=mock_ready_ref) ) mock_actor_cls = MagicMock() mock_actor_cls.options.return_value.remote.return_value = mock_handle actor_options = {} if actor_runtime_env is not None: actor_options["runtime_env"] = actor_runtime_env config = DeploymentActorConfig( name="counter", actor_class="builtins:object", actor_options=actor_options, ) wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v1", ) with patch( "ray.serve._private.deployment_state.DeploymentActorConfig.get_actor_class", return_value=mock_actor_cls, ): success, err = wrapper.start(deployment_runtime_env=deployment_runtime_env) assert success is False assert err is not None for key in conflicting_keys: assert key in err mock_actor_cls.options.assert_not_called() def test_properties(self): """Test actor_logical_name and code_version properties.""" config = DeploymentActorConfig(name="counter", actor_class="builtins:object") wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v1", ) assert wrapper.actor_logical_name == "counter" assert wrapper.code_version == "v1" def test_start_success(self): """Test start() returns (True, None) when actor creation succeeds.""" mock_ready_ref = MagicMock() mock_handle = MagicMock() mock_handle.__ray_ready__ = MagicMock( remote=MagicMock(return_value=mock_ready_ref) ) mock_actor_cls = MagicMock() mock_actor_cls.options.return_value.remote.return_value = mock_handle config = DeploymentActorConfig(name="counter", actor_class="builtins:object") wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v1", ) with patch( "ray.serve._private.deployment_state.DeploymentActorConfig.get_actor_class", return_value=mock_actor_cls, ): success, err = wrapper.start() assert success is True assert err is None assert wrapper._handle is mock_handle assert wrapper._ready_ref is mock_ready_ref def test_start_injects_internal_deployment_context_env_vars(self): """Test start() injects deployment metadata into actor runtime_env.""" mock_ready_ref = MagicMock() mock_handle = MagicMock() mock_handle.__ray_ready__ = MagicMock( remote=MagicMock(return_value=mock_ready_ref) ) mock_actor_cls = MagicMock() mock_actor_cls.options.return_value.remote.return_value = mock_handle config = DeploymentActorConfig(name="counter", actor_class="builtins:object") wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v1", ) with patch( "ray.serve._private.deployment_state.DeploymentActorConfig.get_actor_class", return_value=mock_actor_cls, ): success, err = wrapper.start( deployment_runtime_env={"env_vars": {"PARENT_ENV": "parent"}} ) assert success is True assert err is None actor_options = mock_actor_cls.options.call_args.kwargs assert "runtime_env" in actor_options assert actor_options["runtime_env"]["env_vars"] == { "PARENT_ENV": "parent", "RAY_SERVE_INTERNAL_DEPLOYMENT_APP_NAME": TEST_DEPLOYMENT_ID.app_name, "RAY_SERVE_INTERNAL_DEPLOYMENT_NAME": TEST_DEPLOYMENT_ID.name, "RAY_SERVE_INTERNAL_DEPLOYMENT_ACTOR_NAME": "counter", "RAY_SERVE_INTERNAL_DEPLOYMENT_CODE_VERSION": "v1", } def test_start_preserves_user_env_vars_when_injecting_internal_context(self): """Test internal context env injection preserves unrelated user env vars.""" mock_ready_ref = MagicMock() mock_handle = MagicMock() mock_handle.__ray_ready__ = MagicMock( remote=MagicMock(return_value=mock_ready_ref) ) mock_actor_cls = MagicMock() mock_actor_cls.options.return_value.remote.return_value = mock_handle config = DeploymentActorConfig( name="counter", actor_class="builtins:object", actor_options={"runtime_env": {"env_vars": {"CHILD_ENV": "child"}}}, ) wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v2", ) with patch( "ray.serve._private.deployment_state.DeploymentActorConfig.get_actor_class", return_value=mock_actor_cls, ): wrapper.start(deployment_runtime_env={"env_vars": {"PARENT_ENV": "parent"}}) actor_options = mock_actor_cls.options.call_args.kwargs assert actor_options["runtime_env"]["env_vars"]["PARENT_ENV"] == "parent" assert actor_options["runtime_env"]["env_vars"]["CHILD_ENV"] == "child" assert ( actor_options["runtime_env"]["env_vars"][ "RAY_SERVE_INTERNAL_DEPLOYMENT_ACTOR_NAME" ] == "counter" ) def test_start_failure(self): """Test start() returns (False, error_msg) when actor creation fails.""" mock_actor_cls = MagicMock() mock_actor_cls.options.return_value.remote.side_effect = RuntimeError( "out of resources" ) config = DeploymentActorConfig(name="counter", actor_class="builtins:object") wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v1", ) with patch( "ray.serve._private.deployment_state.DeploymentActorConfig.get_actor_class", return_value=mock_actor_cls, ): success, err = wrapper.start() assert success is False assert "out of resources" in err def test_check_ready_already_ready(self): """Test check_ready() when _ready_ref is None and _handle is set.""" config = DeploymentActorConfig(name="counter", actor_class="builtins:object") wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v1", ) wrapper._handle = object() # Any truthy value - actor is ready wrapper._ready_ref = None ready, err = wrapper.check_ready() assert ready is True assert err is None def test_check_ready_not_started(self): """Test check_ready() when _ready_ref and _handle are None.""" config = DeploymentActorConfig(name="counter", actor_class="builtins:object") wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v1", ) ready, err = wrapper.check_ready() assert ready is False assert err is None def test_check_ready_ref_not_ready(self): """Test check_ready() when _ready_ref exists but is not ready.""" with patch( "ray.serve._private.deployment_state.check_obj_ref_ready_nowait", return_value=False, ): config = DeploymentActorConfig( name="counter", actor_class="builtins:object" ) wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v1", ) wrapper._ready_ref = object() # Pending ref, not ready yet ready, err = wrapper.check_ready() assert ready is False assert err is None def test_check_ready_ref_ready_then_success(self): """Test check_ready() when _ready_ref is ready and ray.get succeeds.""" with ( patch( "ray.serve._private.deployment_state.check_obj_ref_ready_nowait", return_value=True, ), patch("ray.serve._private.deployment_state.ray.get"), ): config = DeploymentActorConfig( name="counter", actor_class="builtins:object" ) wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v1", ) wrapper._ready_ref = object() # Ref that will pass check ready, err = wrapper.check_ready() assert ready is True assert err is None assert wrapper._ready_ref is None def test_check_ready_ref_ready_then_fails(self): """Test check_ready() when ray.get on _ready_ref raises.""" with ( patch( "ray.serve._private.deployment_state.check_obj_ref_ready_nowait", return_value=True, ), patch( "ray.serve._private.deployment_state.ray.get", side_effect=RuntimeError("actor crashed"), ), ): config = DeploymentActorConfig( name="counter", actor_class="builtins:object" ) wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v1", ) wrapper._ready_ref = object() # Ref that passes check, ray.get will fail ready, err = wrapper.check_ready() assert ready is False assert "actor crashed" in err def test_kill_with_handle(self): """Test kill() when _handle is already set.""" with patch("ray.serve._private.deployment_state.ray.kill") as mock_kill: config = DeploymentActorConfig( name="counter", actor_class="builtins:object" ) fake_handle = object() # Simulates existing actor handle wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v1", recovered_handle=fake_handle, ) wrapper.kill() mock_kill.assert_called_once_with(fake_handle, no_restart=True) def test_kill_without_handle(self): """Test kill() when _handle is None - fetches via get_actor then kills.""" fake_handle = object() with ( patch( "ray.serve._private.deployment_state.ray.get_actor", return_value=fake_handle, ), patch("ray.serve._private.deployment_state.ray.kill") as mock_kill, ): config = DeploymentActorConfig( name="counter", actor_class="builtins:object" ) wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v1", ) wrapper.kill() mock_kill.assert_called_once_with(fake_handle, no_restart=True) def test_kill_actor_already_stopped(self): """Test kill() when actor is already stopped - should not raise.""" with ( patch( "ray.serve._private.deployment_state.ray.get_actor", side_effect=ValueError("actor not found"), ), patch("ray.serve._private.deployment_state.ray.kill"), ): config = DeploymentActorConfig( name="counter", actor_class="builtins:object" ) wrapper = DeploymentActorWrapper( deployment_id=TEST_DEPLOYMENT_ID, config=config, code_version="v1", ) wrapper.kill() # should not raise def check_counts( deployment_state: DeploymentState, total: Optional[int] = None, by_state: Optional[List[Tuple[ReplicaState, int]]] = None, ): replicas = { state: deployment_state._replicas.count(states=[state]) for state in ALL_REPLICA_STATES } if total is not None: assert deployment_state._replicas.count() == total, f"Replicas: {replicas}" if by_state is not None: for state, count, version in by_state: assert isinstance(state, ReplicaState) assert isinstance(count, int) and count >= 0 curr_count = deployment_state._replicas.count( version=version, states=[state] ) msg = ( f"Expected {count} for state {state} but got {curr_count}. Current " f"replicas: {replicas}" ) assert curr_count == count, msg def test_create_delete_single_replica(mock_deployment_state_manager): create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info_1, v1 = deployment_info() dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Single replica should be created. dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) # update() should not transition the state if the replica isn't ready. dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) ds._replicas.get()[0]._actor.set_ready() assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Now the replica should be marked running. dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Removing the replica should transition it to stopping. ds.delete() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STOPPING, 1, None)]) assert ds._replicas.get()[0]._actor.stopped assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ds.curr_status_info.status_trigger == DeploymentStatusTrigger.DELETING # Once it's done stopping, replica should be removed. replica = ds._replicas.get()[0] replica._actor.set_done_stopping() dsm.update() check_counts(ds, total=0) @pytest.mark.parametrize( "allocate_logs, expected_dead", [(True, 1), (False, 0)], ids=["with_logs", "no_logs"], ) def test_recent_dead_replicas_retention( mock_deployment_state_manager, allocate_logs, expected_dead ): """A stopped replica is retained for the dashboard iff it allocated a log file.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() dsm.deploy(TEST_DEPLOYMENT_ID, deployment_info()[0]) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() if allocate_logs: # set_ready() allocates the replica's log file path. ds._replicas.get()[0]._actor.set_ready() dsm.update() replica_id = ds._replicas.get()[0].replica_id.unique_id ds.delete() dsm.update() ds._replicas.get()[0]._actor.set_done_stopping() dsm.update() check_counts(ds, total=0) # Dead replicas are tracked separately, so the live list is unaffected. assert ds.list_replica_details() == [] dead = ds.list_recent_dead_replicas() assert len(dead) == expected_dead if expected_dead: assert dead[0].replica_id == replica_id assert dead[0].state == ReplicaState.STOPPED assert dead[0].log_file_path is not None @patch("ray.serve._private.deployment_state.RAY_SERVE_RETAINED_DEAD_REPLICAS", 2) def test_recent_dead_replicas_bounded(mock_deployment_state_manager): """Only the most recent RAY_SERVE_RETAINED_DEAD_REPLICAS are retained.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() dsm.deploy(TEST_DEPLOYMENT_ID, deployment_info(num_replicas=3)[0]) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] assert ds._recent_dead_replicas.maxlen == 2 # Bring up 3 replicas, then stop all of them at once. dsm.update() for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, None)]) ds.delete() dsm.update() for replica in ds._replicas.get(): replica._actor.set_done_stopping() dsm.update() check_counts(ds, total=0) # Buffer is capped at 2 even though three replicas have stopped. assert len(ds._recent_dead_replicas) == 2 def test_force_kill(mock_deployment_state_manager): create_dsm, timer, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() grace_period_s = 10 info_1, _ = deployment_info(graceful_shutdown_timeout_s=grace_period_s) dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() # Create deployment. ds._replicas.get()[0]._actor.set_ready() dsm.update() # Delete deployment. ds.delete() # Replica should remain in STOPPING until it finishes. dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STOPPING, 1, None)]) assert ds._replicas.get()[0]._actor.stopped for _ in range(10): dsm.update() # force_stop shouldn't be called until after the timer. assert not ds._replicas.get()[0]._actor.force_stopped_counter print(ds._replicas) check_counts(ds, total=1, by_state=[(ReplicaState.STOPPING, 1, None)]) # Advance the timer, now the replica should be force stopped. timer.advance(grace_period_s + 0.1) dsm.update() assert ds._replicas.get()[0]._actor.force_stopped_counter == 1 check_counts(ds, total=1, by_state=[(ReplicaState.STOPPING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ds.curr_status_info.status_trigger == DeploymentStatusTrigger.DELETING # Force stop should be called repeatedly until the replica stops. dsm.update() assert ds._replicas.get()[0]._actor.force_stopped_counter == 2 check_counts(ds, total=1, by_state=[(ReplicaState.STOPPING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ds.curr_status_info.status_trigger == DeploymentStatusTrigger.DELETING # Once the replica is done stopping, it should be removed. replica = ds._replicas.get()[0] replica._actor.set_done_stopping() dsm.update() check_counts(ds, total=0) def test_redeploy_same_version(mock_deployment_state_manager): # Redeploying with the same version and code should do nothing. create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info_1, v1 = deployment_info(version="1") assert dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v1)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Test redeploying while the initial deployment is still pending. updating = dsm.deploy(TEST_DEPLOYMENT_ID, info_1) assert not updating # Redeploying the exact same info shouldn't cause any change in status assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v1)]) # Mark the replica ready. After this, the initial goal should be complete. ds._replicas.get()[0]._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Test redeploying after the initial deployment has finished. updating = dsm.deploy(TEST_DEPLOYMENT_ID, info_1) assert not updating assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) def test_redeploy_no_version(mock_deployment_state_manager): """Redeploying with no version specified (`None`) should always redeploy the replicas. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, v1 = deployment_info(version=None) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Test redeploying while the initial deployment is still pending. assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) dsm.update() # The initial replica should be stopping. The new replica should # start without waiting for the old one to stop completely. check_counts( ds, total=2, by_state=[ (ReplicaState.STOPPING, 1, None), (ReplicaState.STARTING, 1, None), ], ) # Mark old replica as completely stopped. ds._replicas.get(states=[ReplicaState.STOPPING])[0]._actor.set_done_stopping() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Check that the new replica has started. ds._replicas.get(states=[ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Now deploy a third version after the transition has finished. b_info_3, v3 = deployment_info(version="3") assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_3) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) dsm.update() # The initial replica should be stopping. The new replica should # start without waiting for the old one to stop completely. check_counts( ds, total=2, by_state=[ (ReplicaState.STOPPING, 1, None), (ReplicaState.STARTING, 1, v3), ], ) ds._replicas.get(states=[ReplicaState.STOPPING])[0]._actor.set_done_stopping() dsm.update() ds._replicas.get(states=[ReplicaState.STARTING])[0]._actor.set_ready() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) def test_redeploy_new_version(mock_deployment_state_manager): """Redeploying with a new version should start a new replica.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, v1 = deployment_info(version="1") dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v1)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Test redeploying while the initial deployment is still pending. b_info_2, v2 = deployment_info(version="2") assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) dsm.update() # The new replica should start without waiting for the old one # to stop. check_counts( ds, total=2, by_state=[(ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v2)], ) # Mark old replica as stopped. ds._replicas.get(states=[ReplicaState.STOPPING])[0]._actor.set_done_stopping() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v2)]) # Mark new replica as ready ds._replicas.get(states=[ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v2)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Now deploy a third version after the transition has finished. b_info_3, v3 = deployment_info(version="3") dsm.deploy(TEST_DEPLOYMENT_ID, b_info_3) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) dsm.update() # New replica should start without waiting for old one to stop check_counts( ds, total=2, by_state=[(ReplicaState.STOPPING, 1, v2), (ReplicaState.STARTING, 1, v3)], ) # Mark old replica as stopped and mark new replica as ready ds._replicas.get(states=[ReplicaState.STOPPING])[0]._actor.set_done_stopping() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v3)]) ds._replicas.get(states=[ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v3)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) def test_redeploy_different_num_replicas(mock_deployment_state_manager): """Tests status changes when redeploying with different num_replicas. 1. Deploys a deployment -> checks if it's UPDATING. 2. Redeploys deployment -> checks that it's still UPDATING. 3. Makes deployment HEALTHY, and then redeploys with more replicas -> check that is becomes UPSCALING. 4. Makes deployment HEALTHY, and then redeploys with more replicas -> check that is becomes DOWNSCALING. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() version = "1" b_info_1, v1 = deployment_info(version=version, num_replicas=5) dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, by_state=[(ReplicaState.STARTING, 5, v1)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Test redeploying with a higher num_replicas while the deployment is UPDATING. b_info_2, v1 = deployment_info(version=version, num_replicas=10) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) # Redeploying while the deployment is UPDATING shouldn't change status. assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) dsm.update() check_counts(ds, by_state=[(ReplicaState.STARTING, 10, v1)]) for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, by_state=[(ReplicaState.RUNNING, 10, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Redeploy with a higher number of replicas. The status should be UPSCALING. b_info_3, v1 = deployment_info(version=version, num_replicas=20) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_3) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) dsm.update() check_counts(ds, by_state=[(ReplicaState.STARTING, 10, v1)]) for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, by_state=[(ReplicaState.RUNNING, 20, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.UPSCALE_COMPLETED ) # Redeploy with lower number of replicas. The status should be DOWNSCALING. b_info_4, v1 = deployment_info(version=version, num_replicas=5) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_4) assert ds.curr_status_info.status == DeploymentStatus.DOWNSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) dsm.update() check_counts( ds, by_state=[(ReplicaState.STOPPING, 15, v1), (ReplicaState.RUNNING, 5, v1)] ) for replica in ds._replicas.get(states=[ReplicaState.STOPPING]): replica._actor.set_done_stopping() dsm.update() check_counts(ds, total=5, by_state=[(ReplicaState.RUNNING, 5, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.DOWNSCALE_COMPLETED ) @pytest.mark.parametrize( "option,value", [ ("user_config", {"hello": "world"}), ("max_ongoing_requests", 10), ("graceful_shutdown_timeout_s", DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_S + 1), ("graceful_shutdown_wait_loop_s", DEFAULT_GRACEFUL_SHUTDOWN_WAIT_LOOP_S + 1), ("health_check_period_s", DEFAULT_HEALTH_CHECK_PERIOD_S + 1), ("health_check_timeout_s", DEFAULT_HEALTH_CHECK_TIMEOUT_S + 1), ], ) def test_deploy_new_config_same_code_version( mock_deployment_state_manager, option, value ): """Deploying a new config with the same version should not deploy a new replica.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, v1 = deployment_info(version="1") assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Create the replica initially. dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Update to a new config without changing the code version. b_info_2, v2 = deployment_info(version="1", **{option: value}) updated = dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) assert updated assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v1)]) if option in [ "user_config", "graceful_shutdown_wait_loop_s", "max_ongoing_requests", ]: dsm.update() check_counts(ds, total=1) check_counts( ds, total=1, by_state=[(ReplicaState.UPDATING, 1, v2)], ) # Mark the replica as ready. ds._replicas.get()[0]._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v2)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) def test_deploy_new_config_same_code_version_2(mock_deployment_state_manager): """Make sure we don't transition from STARTING to UPDATING directly.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, v1 = deployment_info(version="1") updated = dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) assert updated ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Create the replica initially. dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v1)]) # Update to a new config without changing the code version. b_info_2, v2 = deployment_info(version="1", user_config={"hello": "world"}) updated = dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) assert updated assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) dsm.update() # Since it's STARTING, we cannot transition to UPDATING check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v1)]) ds._replicas.get()[0]._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.UPDATING, 1, v2)]) # Mark the replica as ready. ds._replicas.get()[0]._actor.set_ready() dsm.update() check_counts(ds, total=1) check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v2)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) def test_deploy_new_config_new_version(mock_deployment_state_manager): # Deploying a new config with a new version should deploy a new replica. create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, v1 = deployment_info(version="1") assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Create the replica initially. dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Update to a new config and a new version. b_info_2, v2 = deployment_info(version="2", user_config={"hello": "world"}) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) dsm.update() # New version should start immediately without waiting for # replicas of old version to completely stop check_counts( ds, total=2, by_state=[ (ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v2), ], ) # Mark replica of old version as stopped ds._replicas.get(states=[ReplicaState.STOPPING])[0]._actor.set_done_stopping() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v2)]) # Mark new replica as ready ds._replicas.get(states=[ReplicaState.STARTING])[0]._actor.set_ready() assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Check that the new version is now running. dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v2)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) def test_initial_deploy_no_throttling(mock_deployment_state_manager): # All replicas should be started at once for a new deployment. create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, v1 = deployment_info(num_replicas=10, version="1") updated = dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) assert updated ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.STARTING, 10, v1)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() # Check that the new replicas have started. dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.RUNNING, 10, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) def test_new_version_deploy_throttling_new(mock_deployment_state_manager): """All replicas should be started at once for a new deployment. When the version is updated, it should be throttled. The throttling should apply to both code version and user config updates. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, v1 = deployment_info(num_replicas=10, version="1", user_config="1") updated = dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) assert updated ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.STARTING, 10, v1)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() # Check that the new replicas have started. dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.RUNNING, 10, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Now deploy a new version. Two old replicas should be stopped. b_info_2, v2 = deployment_info(num_replicas=10, version="2", user_config="2") assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) dsm.update() check_counts( ds, total=12, by_state=[ (ReplicaState.RUNNING, 8, v1), (ReplicaState.STOPPING, 2, v1), (ReplicaState.STARTING, 2, v2), ], ) # Mark only one of the replicas as done stopping. ds._replicas.get(states=[ReplicaState.STOPPING])[0]._actor.set_done_stopping() dsm.update() check_counts( ds, total=11, by_state=[ # Old version running (ReplicaState.RUNNING, 8, v1), # Replicas being "rolled out" (ReplicaState.STARTING, 2, v2), # Out of the picture (ReplicaState.STOPPING, 1, v1), ], ) # Mark one new replica as ready. Then the rollout should continue, # stopping another old-version-replica and starting another # new-version-replica ds._replicas.get(states=[ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() check_counts( ds, total=12, by_state=[ # Old version running (ReplicaState.RUNNING, 7, v1), # New version running (ReplicaState.RUNNING, 1, v2), # Replicas being "rolled out" (ReplicaState.STARTING, 2, v2), # Out of the picture (ReplicaState.STOPPING, 2, v1), ], ) # Mark the old replicas as done stopping. ds._replicas.get(states=[ReplicaState.STOPPING])[0]._actor.set_done_stopping() ds._replicas.get(states=[ReplicaState.STARTING])[0]._actor.set_ready() # Old replicas should be stopped and new versions started in batches of 2. new_replicas = 2 old_replicas = 8 while old_replicas: assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # 2 replicas should be stopping, and simultaneously 2 replicas # should start to fill the gap. old_replicas -= 2 dsm.update() check_counts( ds, total=12, by_state=[ # Old version running (ReplicaState.RUNNING, old_replicas, v1), # New version running (ReplicaState.RUNNING, new_replicas, v2), # New replicas being "rolled out" (ReplicaState.STARTING, 2, v2), # Out of the picture (ReplicaState.STOPPING, 2, v1), ], ) ds._replicas.get(states=[ReplicaState.STOPPING])[0]._actor.set_done_stopping() ds._replicas.get(states=[ReplicaState.STOPPING])[1]._actor.set_done_stopping() dsm.update() check_counts( ds, total=10, by_state=[ # Old version running (ReplicaState.RUNNING, old_replicas, v1), # New version running (ReplicaState.RUNNING, new_replicas, v2), # Replicas being "rolled out" (ReplicaState.STARTING, 2, v2), ], ) # Set both ready. ds._replicas.get(states=[ReplicaState.STARTING])[0]._actor.set_ready() ds._replicas.get(states=[ReplicaState.STARTING])[1]._actor.set_ready() new_replicas += 2 # All new replicas should be up and running. dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.RUNNING, 10, v2)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) def test_reconfigure_throttling(mock_deployment_state_manager): """All replicas should be started at once for a new deployment. When the version is updated, it should be throttled. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, v1 = deployment_info(num_replicas=2, version="1", user_config="1") assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() # Check that the new replicas have started. dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Now deploy a new user_config. One replica should be updated. b_info_2, v2 = deployment_info(num_replicas=2, version="1", user_config="2") assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.RUNNING, 1, v1), (ReplicaState.UPDATING, 1, v2)], ) # Mark the updating replica as ready. ds._replicas.get(states=[ReplicaState.UPDATING])[0]._actor.set_ready() # The updated replica should now be RUNNING. # The second replica should now be updated. dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.RUNNING, 1, v2), (ReplicaState.UPDATING, 1, v2)], ) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Mark the updating replica as ready. ds._replicas.get(states=[ReplicaState.UPDATING])[0]._actor.set_ready() # Both replicas should now be RUNNING. dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v2)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) @pytest.mark.parametrize( ("num_replicas", "percentage", "expected_stopping"), [ # The existing case: 50% of 10 replicas is 5. (10, 0.5, 5), # Test default percentage (20%) of 10 replicas is 2. (10, None, 2), # Test rounding down: 50% of 3 replicas is 1.5 -> 1. (3, 0.5, 1), # Test minimum of 1: 20% of 4 replicas is 0.8 -> 0, but minimum is 1. (4, 0.2, 1), # Test percentage that isn't a clean divisor. (10, 0.21, 2), # Test 100% update. All old replicas should be stopping. (5, 1.0, 5), ], ) def test_rolling_update_percentage_configurable( mock_deployment_state_manager, num_replicas, percentage, expected_stopping ): """Test that rolling_update_percentage controls how many replicas update per wave.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() deploy_kwargs = {"num_replicas": num_replicas, "version": "1"} if percentage is not None: deploy_kwargs["rolling_update_percentage"] = percentage b_info_1, v1 = deployment_info(**deploy_kwargs) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v1)] ) # Deploy new version and check that the correct number of replicas are # transitioning. deploy_kwargs["version"] = "2" b_info_2, v2 = deployment_info(**deploy_kwargs) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) dsm.update() expected_running_v1 = num_replicas - expected_stopping # When there are 0 running v1 replicas, the check_counts `by_state` # entry should be omitted. expected_by_state = [ (ReplicaState.STOPPING, expected_stopping, v1), (ReplicaState.STARTING, expected_stopping, v2), ] if expected_running_v1 > 0: expected_by_state.insert(0, (ReplicaState.RUNNING, expected_running_v1, v1)) check_counts( ds, total=num_replicas + expected_stopping, by_state=expected_by_state, ) def test_new_version_and_scale_down(mock_deployment_state_manager): # Test the case when we reduce the number of replicas and change the # version at the same time. First the number of replicas should be # turned down, then the rolling update should happen. create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, v1 = deployment_info(num_replicas=10, version="1") updated = dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) assert updated ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.STARTING, 10, v1)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() # Check that the new replicas have started. dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.RUNNING, 10, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Now deploy a new version and scale down the number of replicas to 2. # First, 8 old replicas should be stopped to bring it down to the target. b_info_2, v2 = deployment_info(num_replicas=2, version="2") assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) dsm.update() check_counts( ds, total=10, by_state=[(ReplicaState.RUNNING, 2, v1), (ReplicaState.STOPPING, 8, v1)], ) # Mark only one of the replicas as done stopping. # This should not yet trigger the rolling update because there are still # stopping replicas. ds._replicas.get(states=[ReplicaState.STOPPING])[0]._actor.set_done_stopping() dsm.update() check_counts( ds, total=9, by_state=[(ReplicaState.RUNNING, 2, v1), (ReplicaState.STOPPING, 7, v1)], ) # Stop the remaining replicas. for replica in ds._replicas.get(states=[ReplicaState.STOPPING]): replica._actor.set_done_stopping() # Now the rolling update should trigger, stopping one of the old # replicas, simultaneously starting replica of new version. dsm.update() check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, v1), (ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v2), ], ) # Mark old replica as stopped ds._replicas.get(states=[ReplicaState.STOPPING])[0]._actor.set_done_stopping() dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.RUNNING, 1, v1), (ReplicaState.STARTING, 1, v2)], ) # New version is started, final old version replica should be stopped. ds._replicas.get(states=[ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, v2), (ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v2), ], ) # Old replica finishes stopping and new replica is ready ds._replicas.get(states=[ReplicaState.STOPPING])[0]._actor.set_done_stopping() dsm.update() ds._replicas.get(states=[ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v2)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) def test_new_version_and_scale_up(mock_deployment_state_manager): # Test the case when we increase the number of replicas and change the # version at the same time. The new replicas should all immediately be # turned up. When they're up, rolling update should trigger. create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, v1 = deployment_info(num_replicas=2, version="1") updated = dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) assert updated ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() # Check that the new replicas have started. dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Now deploy a new version and scale up the number of replicas to 10. # 8 new replicas should be started. b_info_2, v2 = deployment_info(num_replicas=10, version="2") assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) dsm.update() check_counts( ds, total=10, by_state=[(ReplicaState.RUNNING, 2, v1), (ReplicaState.STARTING, 8, v2)], ) # Mark the new replicas as ready. for replica in ds._replicas.get(states=[ReplicaState.STARTING]): replica._actor.set_ready() # Now that the new version replicas are up, rolling update should start. dsm.update() check_counts( ds, total=12, by_state=[ (ReplicaState.RUNNING, 0, v1), (ReplicaState.STOPPING, 2, v1), (ReplicaState.STARTING, 2, v2), (ReplicaState.RUNNING, 8, v2), ], ) # Mark the replicas as done stopping. for replica in ds._replicas.get(states=[ReplicaState.STOPPING]): replica._actor.set_done_stopping() dsm.update() check_counts( ds, total=10, by_state=[(ReplicaState.RUNNING, 8, v2), (ReplicaState.STARTING, 2, v2)], ) # Mark the remaining replicas as ready. for replica in ds._replicas.get(states=[ReplicaState.STARTING]): replica._actor.set_ready() # All new replicas should be up and running. dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.RUNNING, 10, v2)]) dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) @pytest.mark.parametrize("target_capacity_direction", ["up", "down"]) def test_scale_num_replicas(mock_deployment_state_manager, target_capacity_direction): """Test upscaling and downscaling the number of replicas manually. Upscaling version: 1. Deploy deployment with num_replicas=3. 2. 3 replicas starting, status=UPDATING, trigger=DEPLOY. 3. It becomes healthy with 3 running replicas. 4. Update deployment to num_replicas=5. 5. 2 replicas starting, status=UPSCALING, trigger=CONFIG_UPDATE. 6. It becomes healthy with 5 running replicas, status=HEALTHY, trigger=CONFIG_UPDATE """ # State version = get_random_string() # Create deployment state manager create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Deploy deployment with 3 replicas info_1, v1 = deployment_info(num_replicas=3, version=version) dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] # status=UPDATING, status_trigger=DEPLOY dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.STARTING, 3, v1)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Set replicas ready and check statuses for replica in ds._replicas.get(): replica._actor.set_ready() # status=HEALTHY, status_trigger=DEPLOY dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # upscale or downscale the number of replicas manually new_num_replicas = 5 if target_capacity_direction == "up" else 1 info_2, _ = deployment_info(num_replicas=new_num_replicas, version=version) dsm.deploy(TEST_DEPLOYMENT_ID, info_2) dsm.update() # status=UPSCALING/DOWNSCALING, status_trigger=CONFIG_UPDATE assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) if target_capacity_direction == "up": check_counts( ds, total=5, by_state=[(ReplicaState.RUNNING, 3, v1), (ReplicaState.STARTING, 2, v1)], ) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING for replica in ds._replicas.get(): replica._actor.set_ready() else: check_counts( ds, total=3, by_state=[(ReplicaState.RUNNING, 1, v1), (ReplicaState.STOPPING, 2, v1)], ) assert ds.curr_status_info.status == DeploymentStatus.DOWNSCALING for replica in ds._replicas.get(): replica._actor.set_done_stopping() # After the upscaling/downscaling finishes # status=HEALTHY, status_trigger=UPSCALING_COMPLETED/DOWNSCALE_COMPLETED dsm.update() check_counts( ds, total=new_num_replicas, by_state=[(ReplicaState.RUNNING, new_num_replicas, v1)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ds.curr_status_info.status_trigger == ( DeploymentStatusTrigger.UPSCALE_COMPLETED if target_capacity_direction == "up" else DeploymentStatusTrigger.DOWNSCALE_COMPLETED ) @pytest.mark.parametrize("force_stop_unhealthy_replicas", [False, True]) def test_health_check( mock_deployment_state_manager, force_stop_unhealthy_replicas: bool ): create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, v1 = deployment_info(num_replicas=2, version="1") assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] ds.FORCE_STOP_UNHEALTHY_REPLICAS = force_stop_unhealthy_replicas dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() # Health check shouldn't be called until it's ready. assert not replica._actor.health_check_called # Check that the new replicas have started. dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) dsm.update() for replica in ds._replicas.get(): # Health check shouldn't be called until it's ready. assert replica._actor.health_check_called # Mark one replica unhealthy; it should be stopped. ds._replicas.get()[0]._actor.set_unhealthy() dsm.update() # SIMULTANEOUSLY a new replica should be started to try to reach # the target number of healthy replicas. check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, v1), (ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v1), ], ) stopping_replicas = ds._replicas.get(states=[ReplicaState.STOPPING]) assert len(stopping_replicas) == 1 stopping_replica = stopping_replicas[0] if force_stop_unhealthy_replicas: assert stopping_replica._actor.force_stopped_counter == 1 else: assert stopping_replica._actor.force_stopped_counter == 0 assert ds.curr_status_info.status == DeploymentStatus.UNHEALTHY # If state transitioned from healthy -> unhealthy, status driver should be none assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.HEALTH_CHECK_FAILED ) stopping_replica._actor.set_done_stopping() dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.RUNNING, 1, v1), (ReplicaState.STARTING, 1, v1)], ) assert ds.curr_status_info.status == DeploymentStatus.UNHEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.HEALTH_CHECK_FAILED ) replica = ds._replicas.get(states=[ReplicaState.STARTING])[0] replica._actor.set_ready() assert ds.curr_status_info.status == DeploymentStatus.UNHEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.HEALTH_CHECK_FAILED ) dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ds.curr_status_info.status_trigger == DeploymentStatusTrigger.UNSPECIFIED def test_health_gauge_caching(mock_deployment_state_manager): """Test that the health gauge is only set when the value changes. The _health_gauge_cache avoids redundant Gauge.set() calls on every control-loop iteration, which are expensive at scale. """ create_dsm, timer, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, v1 = deployment_info(num_replicas=2, version="1") assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) for replica in ds._replicas.get(): replica._actor.set_ready() # First update: _check_startup_replicas transitions STARTING -> RUNNING. # check_and_update_replicas hasn't seen them as RUNNING yet. dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v1)]) # Second update: check_and_update_replicas processes the RUNNING replicas # for the first time, calling check_health() and setting the gauge. dsm.update() replica_ids = [r.replica_id.unique_id for r in ds._replicas.get()] # After the second update the cache should have (value=1, timestamp) for both. for rid in replica_ids: cached_value, cached_time = ds._health_gauge_cache[rid] assert cached_value == 1 # Track how many times Gauge.set is called using a wrapper. original_set = ds.health_check_gauge.set call_count = 0 def counting_set(*args, **kwargs): nonlocal call_count call_count += 1 return original_set(*args, **kwargs) ds.health_check_gauge.set = counting_set # Subsequent updates with all-healthy replicas should NOT call Gauge.set # because the cache already has value 1 for each replica (within TTL). dsm.update() dsm.update() dsm.update() assert call_count == 0, ( f"Gauge.set was called {call_count} times for already-healthy replicas; " "expected 0 (should be cached)" ) # After the TTL expires, the gauge should be re-reported even though # the value hasn't changed. timer.advance(RAY_SERVE_STATUS_GAUGE_REPORT_INTERVAL_S + 1) dsm.update() assert call_count == len(replica_ids), ( f"Gauge.set was called {call_count} times after TTL expired; " f"expected {len(replica_ids)} (one per replica)" ) # Mark one replica unhealthy — gauge should transition to 0. call_count = 0 ds._replicas.get()[0]._actor.set_unhealthy() dsm.update() # Gauge.set should have been called at least once (for the now-unhealthy replica). assert call_count >= 1 # The stopping replica should have cache value 0. stopping = ds._replicas.get(states=[ReplicaState.STOPPING]) assert len(stopping) == 1 cached_value, _ = ds._health_gauge_cache[stopping[0].replica_id.unique_id] assert cached_value == 0 # After the stopped replica is fully removed, its cache entry should be cleaned up. stopped_id = stopping[0].replica_id.unique_id stopping[0]._actor.set_done_stopping() call_count = 0 dsm.update() assert stopped_id not in ds._health_gauge_cache def test_update_while_unhealthy(mock_deployment_state_manager): create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, v1 = deployment_info(num_replicas=2, version="1") assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() # Health check shouldn't be called until it's ready. assert not replica._actor.health_check_called # Check that the new replicas have started. dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) dsm.update() for replica in ds._replicas.get(): # Health check shouldn't be called until it's ready. assert replica._actor.health_check_called # Mark one replica unhealthy. It should be stopped. ds._replicas.get()[0]._actor.set_unhealthy() dsm.update() check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, v1), (ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v1), ], ) assert ds.curr_status_info.status == DeploymentStatus.UNHEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.HEALTH_CHECK_FAILED ) replica = ds._replicas.get(states=[ReplicaState.STOPPING])[0] replica._actor.set_done_stopping() dsm.update() check_counts( ds, total=2, by_state=[ (ReplicaState.RUNNING, 1, v1), (ReplicaState.STARTING, 1, v1), ], ) # Now deploy a new version (e.g., a rollback). This should update the status # to UPDATING and then it should eventually become healthy. b_info_2, v2 = deployment_info(num_replicas=2, version="2") assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) # The replica that was still starting should be stopped (over the # running replica). dsm.update() # Simultaneously, a replica with the new version should be started check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, v1), (ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v2), ], ) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Mark the remaining running replica of the old version as unhealthy ds._replicas.get(states=[ReplicaState.RUNNING])[0]._actor.set_unhealthy() dsm.update() # A replica of the new version should get started to try to reach # the target number of healthy replicas check_counts( ds, total=4, by_state=[(ReplicaState.STOPPING, 2, v1), (ReplicaState.STARTING, 2, v2)], ) # Check that a failure in the old version replica does not mark the # deployment as UNHEALTHY. assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) ds._replicas.get(states=[ReplicaState.STOPPING])[0]._actor.set_done_stopping() ds._replicas.get(states=[ReplicaState.STOPPING])[1]._actor.set_done_stopping() # Another replica of the new version should get started. dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v2)]) # Mark new version replicas as ready. for replica in ds._replicas.get(states=[ReplicaState.STARTING]): replica._actor.set_ready() # Both replicas should be RUNNING, deployment should be HEALTHY. dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v2)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) def _constructor_failure_loop_two_replica( dsm, ds, num_loops, replica_retry_multiplier=3 ): """Helper function to exact constructor failure loops.""" for i in range(num_loops): # Two replicas should be created. dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, None)]) assert ds._replica_constructor_retry_counter == i * 2 replica_1 = ds._replicas.get()[0] replica_2 = ds._replicas.get()[1] replica_1._actor.set_failed_to_start() replica_2._actor.set_failed_to_start() # Now the replica should be marked STOPPING after failure. dsm.update() if ds._replica_constructor_retry_counter >= replica_retry_multiplier * 2: check_counts( ds, total=2, by_state=[(ReplicaState.STOPPING, 2, None)], ) else: check_counts( ds, total=4, by_state=[ (ReplicaState.STOPPING, 2, None), (ReplicaState.STARTING, 2, None), ], ) # Once it's done stopping, replica should be removed. replica_1._actor.set_done_stopping() replica_2._actor.set_done_stopping() def test_deploy_with_consistent_constructor_failure( mock_deployment_state_manager, mock_max_per_replica_retry_count ): """ Test deploy() multiple replicas with consistent constructor failure. The deployment should get marked FAILED. """ create_dsm, timer, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, _ = deployment_info(num_replicas=2) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) loop_count = mock_max_per_replica_retry_count _constructor_failure_loop_two_replica( dsm, ds, loop_count, mock_max_per_replica_retry_count ) assert ds._replica_constructor_retry_counter == 2 * mock_max_per_replica_retry_count assert ds.curr_status_info.status == DeploymentStatus.DEPLOY_FAILED assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.REPLICA_STARTUP_FAILED ) check_counts(ds, total=2) assert ds.curr_status_info.message != "" # No more replicas should be retried. for _ in range(20): dsm.update() assert ds._replica_constructor_retry_counter == 4 check_counts(ds, total=0) timer.advance(10) # simulate time passing between each call to update def test_deploy_with_partial_constructor_failure( mock_deployment_state_manager, mock_max_per_replica_retry_count ): """ Test deploy() multiple replicas with constructor failure exceedining pre-set limit but achieved partial success with at least 1 running replica. Ensures: 1) Deployment status doesn't get marked FAILED. 2) There should be expected # of RUNNING replicas eventually that matches user intent 3) Replica counter set as -1 to stop tracking current goal as it's already completed Same testing for same test case in test_deploy.py. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, _ = deployment_info(num_replicas=2) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) _constructor_failure_loop_two_replica(dsm, ds, 1, mock_max_per_replica_retry_count) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, None)]) assert ds._replica_constructor_retry_counter == 2 assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Let one replica reach RUNNING state while the other still fails replica_1 = ds._replicas.get()[0] replica_2 = ds._replicas.get()[1] replica_1._actor.set_ready() replica_2._actor.set_failed_to_start() # Failed to start replica should be removed dsm.update() # A new replica should be brought up to take its place check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, None), (ReplicaState.STOPPING, 1, None), (ReplicaState.STARTING, 1, None), ], ) # Mark old replica as done stopping replica_2._actor.set_done_stopping() dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.RUNNING, 1, None), (ReplicaState.STARTING, 1, None)], ) # Set the starting one to fail again and trigger retry limit starting_replica = ds._replicas.get(states=[ReplicaState.STARTING])[0] starting_replica._actor.set_failed_to_start() dsm.update() # Ensure our goal returned with replica_has_started flag set assert ds._replica_has_started # Deployment should NOT be considered complete yet assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # A new replica should be brought up to take its place check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, None), (ReplicaState.STOPPING, 1, None), (ReplicaState.STARTING, 1, None), ], ) starting_replica = ds._replicas.get(states=[ReplicaState.STOPPING])[0] starting_replica._actor.set_done_stopping() dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.RUNNING, 1, None), (ReplicaState.STARTING, 1, None)], ) starting_replica = ds._replicas.get(states=[ReplicaState.STARTING])[0] starting_replica._actor.set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, None)]) # Deployment should be considered complete assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) def test_deploy_with_placement_group_failure(mock_deployment_state_manager): """ Test deploy with a placement group failure. """ def fake_create_placement_group_fn(placement_group_bundles, *args, **kwargs): """Fakes the placement_group_fn used by the scheduler. Lets the test to run without starting Ray. Raises an exception if the bundles are invalid. """ validate_placement_group(bundles=placement_group_bundles) create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm( create_placement_group_fn_override=fake_create_placement_group_fn, ) def create_deployment_state( deployment_id: DeploymentID, pg_bundles=None ) -> List[DeploymentState]: b_info, _ = deployment_info(num_replicas=3) b_info.replica_config.placement_group_bundles = pg_bundles assert dsm.deploy(deployment_id, b_info) ds = dsm._deployment_states[deployment_id] assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) return ds # Make all of ds1's replica's placement groups invalid. invalid_bundle = [{"GPU": 0}] with pytest.raises(ValueError): validate_placement_group(invalid_bundle) ds1 = create_deployment_state(TEST_DEPLOYMENT_ID, pg_bundles=invalid_bundle) ds2 = create_deployment_state(TEST_DEPLOYMENT_ID_2) # Now ds1's replicas should all fail, while ds2's replicas should run. dsm.update() check_counts(ds1, total=3, by_state=[(ReplicaState.STOPPING, 3, None)]) assert ds1._replica_constructor_retry_counter == 3 assert "Retrying 6 more time(s)" in ds1.curr_status_info.message # Set all of ds1's replicas to stopped. for replica in ds1._replicas.get(): replica._actor.set_done_stopping() check_counts(ds2, total=3, by_state=[(ReplicaState.STARTING, 3, None)]) assert ds2._replica_constructor_retry_counter == 0 # Set all of ds2's replicas to ready. for replica in ds2._replicas.get(): replica._actor.set_ready() dsm.update() assert ds1.curr_status_info.status == DeploymentStatus.UPDATING check_counts(ds1, total=3, by_state=[(ReplicaState.STOPPING, 3, None)]) assert ds1._replica_constructor_retry_counter == 6 assert "Retrying 3 more time(s)" in ds1.curr_status_info.message # Set all of ds1's replicas to stopped. for replica in ds1._replicas.get(): replica._actor.set_done_stopping() assert ds2.curr_status_info.status == DeploymentStatus.HEALTHY check_counts(ds2, total=3, by_state=[(ReplicaState.RUNNING, 3, None)]) assert ds2._replica_constructor_retry_counter == 0 dsm.update() assert ds1.curr_status_info.status == DeploymentStatus.UPDATING check_counts(ds1, total=3, by_state=[(ReplicaState.STOPPING, 3, None)]) assert ds1._replica_constructor_retry_counter == 9 assert "Retrying 0 more time(s)" in ds1.curr_status_info.message # Set all of ds1's replicas to stopped. for replica in ds1._replicas.get(): replica._actor.set_done_stopping() dsm.update() # All replicas have failed to initialize 3 times. The deployment should # stop trying to initialize replicas. assert ds1.curr_status_info.status == DeploymentStatus.DEPLOY_FAILED check_counts(ds1, total=0) assert ds1._replica_constructor_retry_counter == 9 assert "The deployment failed to start" in ds1.curr_status_info.message def test_deploy_with_gang_placement_group_failure(mock_deployment_state_manager): """ Test deploy with a gang placement group creation failure. """ def failing_create_placement_group_fn(request, *args, **kwargs): raise RuntimeError("Simulated gang PG creation failure") create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm( create_placement_group_fn_override=failing_create_placement_group_fn, ) b_info, _ = deployment_info( num_replicas=4, gang_scheduling_config=GangSchedulingConfig(gang_size=2), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] assert ds.curr_status_info.status == DeploymentStatus.UPDATING # Each dsm.update() call attempts to create gang PGs, fails, and # increments the retry counter by 1. The threshold is # min(max_constructor_retry_count, target_num_replicas * MAX_PER_REPLICA_RETRY_COUNT) # = min(inf, 4 * 2) = 8. threshold = ds._failed_to_start_threshold for i in range(1, threshold + 1): dsm.update() assert "Gang scheduling failed" in ds.curr_status_info.message if i < threshold: assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert f"Retrying {threshold - i} more time(s)" in ( ds.curr_status_info.message ) # After reaching the threshold, the next update should fail the deployment. dsm.update() assert ds.curr_status_info.status == DeploymentStatus.DEPLOY_FAILED assert "The deployment failed to start" in ds.curr_status_info.message def _deployment_actors_config(): """DeploymentActorConfig for tests. Uses import path to avoid Ray init.""" return [ DeploymentActorConfig( name="counter", actor_class="ray.serve.tests.test_deployment_actors:SharedCounter", init_kwargs={"start": 0}, ), ] def _deployment_actors_config_two(): """Two deployment actors for partial-failure tests.""" return [ DeploymentActorConfig( name="counter", actor_class="ray.serve.tests.test_deployment_actors:SharedCounter", init_kwargs={"start": 0}, ), DeploymentActorConfig( name="cache", actor_class="ray.serve.tests.test_deployment_actors:SharedCounter", init_kwargs={"start": 0}, ), ] def _get_deployment_actor_wrapper( ds: DeploymentState, code_version: str, actor_name: str = "counter", ): wrapper = ds._deployment_actors.get_wrapper(code_version, actor_name) if wrapper is None: raise KeyError( f"No deployment actor wrapper for version={code_version!r} " f"name={actor_name!r}" ) return wrapper class TestDeploymentActors: """Deployment actor tests using setter methods on wrapper instances.""" def test_deploy_with_deployment_actors_deferred_replica_creation( self, mock_deployment_state_manager ): """Replicas are not created until deployment actors are ready.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=2, deployment_actors=_deployment_actors_config(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] for _ in range(5): dsm.update() check_counts(ds, total=0) _get_deployment_actor_wrapper(ds, "1").set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, None)]) for r in ds._replicas.get(): r._actor.set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_deploy_with_deployment_actor_failure(self, mock_deployment_state_manager): """Deployment actor constructor failure transitions to DEPLOY_FAILED.""" create_dsm, _, _, _ = mock_deployment_state_manager error_msg = "Deployment actor 'counter' failed: constructor error" dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=2, deployment_actors=_deployment_actors_config(), max_constructor_retry_count=2, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] for i in range(2): dsm.update() _get_deployment_actor_wrapper(ds, "1").set_failed_to_start(error_msg) dsm.update() if i < 1: assert ds.curr_status_info.status == DeploymentStatus.UPDATING else: assert ds.curr_status_info.status == DeploymentStatus.DEPLOY_FAILED assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.DEPLOYMENT_ACTOR_FAILED ) assert error_msg in ds.curr_status_info.message def test_delete_deployment_calls_stop_deployment_actors( self, mock_deployment_state_manager ): """Deleting a deployment calls force_stop on deployment actors.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() wrapper_v1 = _get_deployment_actor_wrapper(ds, "1") wrapper_v1.set_ready() dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY ds.delete() dsm.update() ds._replicas.get()[0]._actor.set_done_stopping() dsm.update() assert wrapper_v1.killed def test_cleanup_orphaned_deployment_actors_on_version_change( self, mock_deployment_state_manager ): """Orphaned version's deployment actor is killed after rollout.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info1, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() wrapper_v1 = _get_deployment_actor_wrapper(ds, "1") wrapper_v1.set_ready() dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY info2, _ = deployment_info( version="2", num_replicas=1, deployment_actors=_deployment_actors_config(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info2) dsm.update() wrapper_v2 = _get_deployment_actor_wrapper(ds, "2") wrapper_v2.set_ready() dsm.update() ds._replicas.get()[0]._actor.set_done_stopping() dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert wrapper_v1.killed assert not wrapper_v2.killed def test_deploy_without_deployment_actors_creates_replicas_immediately( self, mock_deployment_state_manager ): """Deployments without deployment_actors create replicas immediately.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info(version="1", num_replicas=2) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, None)]) for r in ds._replicas.get(): r._actor.set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_deployment_actor_start_retry_then_succeed( self, mock_deployment_state_manager ): """Deployment actor fails N times then succeeds; replicas created.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=2, deployment_actors=_deployment_actors_config(), max_constructor_retry_count=5, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() _get_deployment_actor_wrapper(ds, "1").set_failed_to_start("fail") dsm.update() assert ds._deployment_actor_retry_counter == 1 dsm.update() _get_deployment_actor_wrapper(ds, "1").set_failed_to_start("fail") dsm.update() assert ds._deployment_actor_retry_counter == 2 dsm.update() _get_deployment_actor_wrapper(ds, "1").set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, None)]) for r in ds._replicas.get(): r._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_deployment_actor_retry_counter_reset_on_redeploy( self, mock_deployment_state_manager ): """Redeploy resets _deployment_actor_retry_counter.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info1, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), max_constructor_retry_count=5, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() _get_deployment_actor_wrapper(ds, "1").set_failed_to_start("fail") dsm.update() assert ds._deployment_actor_retry_counter == 1 info2, _ = deployment_info( version="2", num_replicas=1, deployment_actors=_deployment_actors_config(), max_constructor_retry_count=5, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info2) assert ds._deployment_actor_retry_counter == 0 def test_deployment_actor_terminal_failure(self, mock_deployment_state_manager): """After threshold failures, deployment stops retrying.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), max_constructor_retry_count=3, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] threshold = ds._deployment_actor_failed_to_start_threshold for _ in range(threshold): dsm.update() _get_deployment_actor_wrapper(ds, "1").set_failed_to_start( "persistent error" ) dsm.update() assert ds.deployment_actor_terminally_failed() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.DEPLOY_FAILED assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.DEPLOYMENT_ACTOR_FAILED ) def test_deployment_actor_failed_handled_when_already_deploy_failed( self, mock_deployment_state_manager ): """Repeated ticks with DEPLOYMENT_ACTOR_FAILED when already DEPLOY_FAILED. When deployment is already in DEPLOY_FAILED (due to deployment actor failure), check_curr_status hits deployment_actor_terminally_failed() again on each tick. handle_transition must handle DEPLOYMENT_ACTOR_FAILED in the DEPLOY_FAILED block to avoid returning None and crashing. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), max_constructor_retry_count=2, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Reach DEPLOY_FAILED via deployment actor terminal failure for _ in range(2): dsm.update() _get_deployment_actor_wrapper(ds, "1").set_failed_to_start( "persistent error" ) dsm.update() assert ds.curr_status_info.status == DeploymentStatus.DEPLOY_FAILED assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.DEPLOYMENT_ACTOR_FAILED ) # Repeated ticks: deployment_actor_terminally_failed() stays True, so # check_curr_status calls handle_transition(DEPLOYMENT_ACTOR_FAILED) # while status is already DEPLOY_FAILED. Must not crash. for _ in range(5): dsm.update() assert ds.curr_status_info is not None assert ds.curr_status_info.status == DeploymentStatus.DEPLOY_FAILED assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.DEPLOYMENT_ACTOR_FAILED ) def test_no_deployment_actors_not_terminally_failed( self, mock_deployment_state_manager ): """Deployment without deployment_actors must not be terminally failed.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info(version="1", num_replicas=1) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] assert not ds.deployment_actor_terminally_failed() assert not ds._terminally_failed() ds._deployment_actor_retry_counter = ( ds._deployment_actor_failed_to_start_threshold + 1 ) assert not ds.deployment_actor_terminally_failed() dsm.update() assert ds.curr_status_info.status != DeploymentStatus.DEPLOY_FAILED def test_deployment_actor_partial_failure(self, mock_deployment_state_manager): """One of two deployment actors fails; DEPLOY_FAILED after threshold.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=2, deployment_actors=_deployment_actors_config_two(), max_constructor_retry_count=2, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] for _ in range(2): dsm.update() _get_deployment_actor_wrapper(ds, "1", "counter").set_ready() _get_deployment_actor_wrapper(ds, "1", "cache").set_failed_to_start( "cache actor failed" ) dsm.update() assert ds.curr_status_info.status == DeploymentStatus.DEPLOY_FAILED assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.DEPLOYMENT_ACTOR_FAILED ) assert "cache actor failed" in ds.curr_status_info.message def test_deployment_actor_partial_failure_preserves_running( self, mock_deployment_state_manager ): """When one actor fails, already-RUNNING actors must stay tracked.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=2, deployment_actors=_deployment_actors_config_two(), max_constructor_retry_count=3, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() assert ( ds._deployment_actors.count("1", states=[DeploymentActorState.STARTING]) == 2 ) _get_deployment_actor_wrapper(ds, "1", "counter").set_ready() _get_deployment_actor_wrapper(ds, "1", "cache").set_failed_to_start( "cache actor failed" ) dsm.update() assert ( ds._deployment_actors.count("1", states=[DeploymentActorState.RUNNING]) == 1 ) counter_wrapper = ds._deployment_actors.get_wrapper("1", "counter") assert counter_wrapper is not None def test_deployment_actor_recovery_from_checkpoint( self, mock_deployment_state_manager ): """recover_target_state_from_checkpoint restores deployment actors.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, v1 = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() _get_deployment_actor_wrapper(ds, "1").set_ready() dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY dsm.save_checkpoint() checkpoint = dsm._kv_store.get(CHECKPOINT_KEY) assert checkpoint is not None mock_handle = MagicMock() with patch("ray.get_actor") as mock_get_actor: mock_get_actor.return_value = mock_handle new_dsm = create_dsm([ds._replicas.get()[0].replica_id.to_full_id_str()]) new_ds = new_dsm._deployment_states[TEST_DEPLOYMENT_ID] assert ( new_ds._deployment_actors.count( "1", states=[DeploymentActorState.RECOVERING] ) == 1 ) recovered_wrapper = _get_deployment_actor_wrapper(new_ds, "1") assert recovered_wrapper._handle is mock_handle recovered_wrapper.set_ready() new_dsm.update() assert ( new_ds._deployment_actors.count("1", states=[DeploymentActorState.RUNNING]) == 1 ) check_counts(new_ds, total=1, by_state=[(ReplicaState.RECOVERING, 1, v1)]) def test_deployment_actor_recovery_get_actor_raises_value_error( self, mock_deployment_state_manager ): """When ray.get_actor raises ValueError during recovery, actor is recreated.""" from ray.serve._private.deployment_state import CHECKPOINT_KEY create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, v1 = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() _get_deployment_actor_wrapper(ds, "1").set_ready() dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY dsm.save_checkpoint() assert dsm._kv_store.get(CHECKPOINT_KEY) is not None # ray.get_actor raises ValueError when actor not found during recovery. with patch("ray.get_actor") as mock_get_actor: mock_get_actor.side_effect = ValueError("Actor not found") new_dsm = create_dsm([ds._replicas.get()[0].replica_id.to_full_id_str()]) new_ds = new_dsm._deployment_states[TEST_DEPLOYMENT_ID] # No deployment actors recovered (ValueError path skips add). assert ( new_ds._deployment_actors.count( "1", states=[ DeploymentActorState.RECOVERING, DeploymentActorState.RUNNING, ], ) == 0 ) # Next update: start_deployment_actors recreates the missing actor. new_dsm.update() assert ( new_ds._deployment_actors.count("1", states=[DeploymentActorState.STARTING]) == 1 ) _get_deployment_actor_wrapper(new_ds, "1").set_ready() new_dsm.update() new_ds._replicas.get()[0]._actor.set_ready() new_dsm.update() assert new_ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_deployment_actor_multiple_not_ready_until_all_ready( self, mock_deployment_state_manager ): """Replicas not created until ready_count == len(configs).""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=2, deployment_actors=_deployment_actors_config_two(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=0) _get_deployment_actor_wrapper(ds, "1", "counter").set_ready() dsm.update() check_counts(ds, total=0) _get_deployment_actor_wrapper(ds, "1", "cache").set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, None)]) def test_deployment_actor_deletion_gate_with_starting( self, mock_deployment_state_manager ): """Deletion does not complete until deployment actor is stopped.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() wrapper = _get_deployment_actor_wrapper(ds, "1") wrapper.set_ready() dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY ds.delete() dsm.update() ds._replicas.get()[0]._actor.set_done_stopping() dsm.update() assert wrapper.killed dsm.update() assert TEST_DEPLOYMENT_ID not in dsm._deployment_states def test_deletion_after_version_update_before_new_actors_ready( self, mock_deployment_state_manager ): """Deletion completes when delete is called after version update but before new version's deployment actors became ready. Bug: With 2+ replicas, version update stops old replicas one-by-one (rollout). After one update we have 1 RUNNING v1 replica. Delete before v2 actors become ready. stop_deployment_actors_if_needed removes v2 actors (versions_to_keep only has v1 from replicas). Without the fix, check_deployment_actors_ready would return False and block _get_target_replica_delta, so downscale never runs for the remaining RUNNING replica and the deployment gets stuck in deletion. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info1, _ = deployment_info( version="1", num_replicas=2, deployment_actors=_deployment_actors_config(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() _get_deployment_actor_wrapper(ds, "1").set_ready() dsm.update() for r in ds._replicas.get(): r._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Version update: deploy v2. v2 deployment actors start but don't become ready. info2, _ = deployment_info( version="2", num_replicas=2, deployment_actors=_deployment_actors_config(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info2) dsm.update() # rollout stops 1 v1 replica; we have 1 STOPPING, 1 RUNNING. v2 actor STARTING. assert _get_deployment_actor_wrapper(ds, "2") is not None running = ds._replicas.get(states=[ReplicaState.RUNNING]) assert ( len(running) >= 1 ), "Need at least 1 RUNNING v1 replica to trigger the bug" # Delete before v2 actors become ready. stop_deployment_actors_if_needed # removes v2 actors. Without the fix, deployment actors block blocks # downscaling of the remaining RUNNING replica(s). ds.delete() for _ in range(30): dsm.update() if TEST_DEPLOYMENT_ID not in dsm._deployment_states: break stopping = ds._replicas.get(states=[ReplicaState.STOPPING]) for s in stopping: s._actor.set_done_stopping() assert TEST_DEPLOYMENT_ID not in dsm._deployment_states, ( "Deployment should complete deletion; without the fix it gets stuck " "because check_deployment_actors_ready blocks downscaling." ) def test_deployment_actor_retry_counter_reset_on_success( self, mock_deployment_state_manager ): """On successful readiness, retry counter is reset.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), max_constructor_retry_count=5, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() _get_deployment_actor_wrapper(ds, "1").set_ready() dsm.update() assert ds._deployment_actor_retry_counter == 0 ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_deployment_actor_start_failure_increments_counter( self, mock_deployment_state_manager ): """Failure via set_failed_to_start increments retry counter.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), max_constructor_retry_count=3, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] for i in range(3): dsm.update() _get_deployment_actor_wrapper(ds, "1").set_failed_to_start( "constructor failed" ) dsm.update() if i < 2: assert ds._deployment_actor_retry_counter == i + 1 assert ds.curr_status_info.status == DeploymentStatus.DEPLOY_FAILED def test_deployment_actor_check_ready_failure_increments_counter( self, mock_deployment_state_manager ): """check_ready failure increments counter; DEPLOY_FAILED after threshold.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), max_constructor_retry_count=3, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] for i in range(3): dsm.update() _get_deployment_actor_wrapper(ds, "1").set_failed_to_start("check failed") dsm.update() assert ds._deployment_actor_retry_counter == i + 1 assert ds.curr_status_info.status == DeploymentStatus.DEPLOY_FAILED def test_deployment_actor_check_health_healthy_stays_running( self, mock_deployment_state_manager ): """``check_and_update_deployment_actors`` re-adds healthy RUNNING actors.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() w = _get_deployment_actor_wrapper(ds, "1") w.set_ready() dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert not w.killed dsm.update() assert not w.killed assert ds._deployment_actors.get_wrapper("1", "counter") is w def test_deployment_actor_check_health_unhealthy_kills_and_recreates( self, mock_deployment_state_manager ): """Failed ``check_health`` kills the actor without startup retry counter; recovers.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), max_constructor_retry_count=5, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() w = _get_deployment_actor_wrapper(ds, "1") w.set_ready() dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY w.set_health_ok(False) # Isolate health reconciliation; a full ``dsm.update()`` would also scale # and start a replacement actor in the same tick. ds.check_and_update_deployment_actors() assert w.killed assert ds._deployment_actors.get_wrapper("1", "counter") is None assert ds._deployment_actor_retry_counter == 0 assert ds._in_transition is True assert ds.curr_status_info.status == DeploymentStatus.UNHEALTHY dsm.update() w2 = _get_deployment_actor_wrapper(ds, "1") assert w2 is not w assert not w2.killed w2.set_ready() dsm.update() assert ds._deployment_actor_retry_counter == 0 assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_deployment_actor_reset_health_state_after_running_on_ready( self, mock_deployment_state_manager ): """``check_deployment_actors_ready`` resets health bookkeeping when RUNNING.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() w = _get_deployment_actor_wrapper(ds, "1") assert w.reset_health_state_after_running_count == 0 w.set_ready() dsm.update() assert w.reset_health_state_after_running_count == 1 def test_deployment_actor_check_health_mixed_two_actors( self, mock_deployment_state_manager ): """Only unhealthy deployment actors are killed; others stay RUNNING.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config_two(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() w_counter = _get_deployment_actor_wrapper(ds, "1", "counter") w_cache = _get_deployment_actor_wrapper(ds, "1", "cache") w_counter.set_ready() w_cache.set_ready() dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY w_cache.set_health_ok(False) ds.check_and_update_deployment_actors() assert w_cache.killed assert not w_counter.killed assert ds._deployment_actors.get_wrapper("1", "cache") is None assert ds._deployment_actors.get_wrapper("1", "counter") is w_counter def test_deployment_actors_satisfied_for_target_requires_all_slots_filled( self, mock_deployment_state_manager ): """Missing deployment actor slots are unsatisfied until all are tracked again.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config_two(), ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() w_counter = _get_deployment_actor_wrapper(ds, "1", "counter") w_cache = _get_deployment_actor_wrapper(ds, "1", "cache") w_counter.set_ready() w_cache.set_ready() dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ds._deployment_actors_satisfied_for_target() w_cache.set_health_ok(False) ds.check_and_update_deployment_actors() assert not ds._deployment_actors_satisfied_for_target() assert ds._replicas.count(states=[ReplicaState.RUNNING]) == 1 def test_deployment_actor_health_check_non_target_kills_without_startup_counter( self, mock_deployment_state_manager ): """Like old replicas: non-target unhealthy actors are stopped, no startup counter.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info1, _ = deployment_info( version="1", num_replicas=1, deployment_actors=_deployment_actors_config(), max_constructor_retry_count=3, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() _get_deployment_actor_wrapper(ds, "1").set_ready() dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY info2, _ = deployment_info( version="2", num_replicas=1, deployment_actors=_deployment_actors_config(), max_constructor_retry_count=3, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, info2) dsm.update() w_v2 = _get_deployment_actor_wrapper(ds, "2") w_v2.set_ready() dsm.update() ds._replicas.get()[0]._actor.set_done_stopping() dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ds._deployment_actor_retry_counter == 0 w_orphan = _mock_deployment_actor_wrapper( TEST_DEPLOYMENT_ID, "1", "orphan_leftover" ) w_orphan.set_health_ok(False) ds._deployment_actors.add(DeploymentActorState.RUNNING, w_orphan) ds.check_and_update_deployment_actors() assert w_orphan.killed assert ds._deployment_actor_retry_counter == 0 assert ds._deployment_actors.get_wrapper("1", "orphan_leftover") is None assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_deploy_with_transient_constructor_failure(mock_deployment_state_manager): """ Test deploy() multiple replicas with transient constructor failure. Ensures: 1) Deployment status gets marked as RUNNING. 2) There should be expected # of RUNNING replicas eventually that matches user intent. 3) Replica counter set as -1 to stop tracking current goal as it's already completed. Same testing for same test case in test_deploy.py. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, _ = deployment_info(num_replicas=2) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Burn 4 retries from both replicas. _constructor_failure_loop_two_replica(dsm, ds, 2) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) assert ds._replica_constructor_retry_counter == 4 # Let both replicas succeed in last try. dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) ds._replicas.get()[0]._actor.set_ready() ds._replicas.get()[1]._actor.set_ready() # Everything should be running and healthy now dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, None)]) assert ds._replica_constructor_retry_counter == 0 assert ds._replica_has_started assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) def test_recover_state_from_replica_names(mock_deployment_state_manager): """Test recover deployment state.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Deploy deployment with version "1" and one replica info1, v1 = deployment_info(version="1") target_state_changed = dsm.deploy(TEST_DEPLOYMENT_ID, info1) assert target_state_changed dsm.save_checkpoint() ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Single replica of version `version1` should be created and in STARTING state dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v1)]) mocked_replica = ds._replicas.get()[0] # The same replica should transition to RUNNING mocked_replica._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v1)]) # (simulate controller crashed!) Create a new deployment state # manager, and it should call _recover_from_checkpoint new_dsm: DeploymentStateManager = create_dsm( [mocked_replica.replica_id.to_full_id_str()] ) # New deployment state should be created and one replica should # be RECOVERING with last-checkpointed target version `version1` new_ds = new_dsm._deployment_states[TEST_DEPLOYMENT_ID] check_counts(new_ds, total=1, by_state=[(ReplicaState.RECOVERING, 1, v1)]) # Get the new mocked replica. Note that this represents a newly # instantiated class keeping track of the state of the replica, # but pointing to the same replica actor new_mocked_replica = new_ds._replicas.get()[0] new_mocked_replica._actor.set_ready(v1) any_recovering = new_dsm.update() check_counts(new_ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v1)]) assert not any_recovering # Make sure replica ID is the same, meaning the actor is the same assert mocked_replica.replica_id == new_mocked_replica.replica_id def test_recover_during_rolling_update(mock_deployment_state_manager): """Test controller crashes before a replica is updated to new version. During recovery, the controller should wait for the version to be fetched from the replica actor. Once it is fetched and the controller realizes the replica has an outdated version, it should be stopped and a new replica should be started with the target version. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm = create_dsm() # Step 1: Create some deployment info with actors in running state info1, v1 = deployment_info(version="1") target_state_changed = dsm.deploy(TEST_DEPLOYMENT_ID, info1) assert target_state_changed dsm.save_checkpoint() ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Single replica of version `version1` should be created and in STARTING state dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v1)]) mocked_replica = ds._replicas.get()[0] # The same replica should transition to RUNNING mocked_replica._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v1)]) # Now execute a rollout: upgrade the version to "2". info2, v2 = deployment_info(version="2") target_state_changed = dsm.deploy(TEST_DEPLOYMENT_ID, info2) assert target_state_changed # In real code this checkpoint would be done by the caller of .deploy() dsm.save_checkpoint() # Before the replica could be stopped and restarted, simulate # controller crashed! A new deployment state manager should be # created, and it should call _recover_from_checkpoint new_dsm = create_dsm([mocked_replica.replica_id.to_full_id_str()]) # New deployment state should be created and one replica should # be RECOVERING with last-checkpointed target version "2" new_ds = new_dsm._deployment_states[TEST_DEPLOYMENT_ID] check_counts(new_ds, total=1, by_state=[(ReplicaState.RECOVERING, 1, v2)]) for _ in range(3): new_dsm.update() check_counts(new_ds, total=1, by_state=[(ReplicaState.RECOVERING, 1, v2)]) # Get the new mocked replica. Note that this represents a newly # instantiated class keeping track of the state of the replica, # but pointing to the same replica actor new_mocked_replica = new_ds._replicas.get()[0] # Recover real version "1" (simulate previous actor not yet stopped) new_mocked_replica._actor.set_ready(v1) # At this point the replica is running new_dsm.update() # Then deployment state manager notices the replica has outdated version -> stops it new_dsm.update() # Also, a replica of version "2" should be started check_counts( new_ds, total=2, by_state=[(ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v2)], ) new_mocked_replica._actor.set_done_stopping() # Mark old replica as stopped. new_dsm.update() check_counts(new_ds, total=1, by_state=[(ReplicaState.STARTING, 1, v2)]) new_mocked_replica_version2 = new_ds._replicas.get()[0] new_mocked_replica_version2._actor.set_ready() new_dsm.update() check_counts(new_ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v2)]) # Make sure replica name is different, meaning a different "actor" was started assert mocked_replica.replica_id != new_mocked_replica_version2.replica_id def test_actor_uninitialized_before_recover(mock_deployment_state_manager): """Test replica actor was found alive but never finished initialization. Mirrors the production scenario where the previous controller crashed between actor creation and the first `initialize_and_get_metadata(rank=...)` call. `recover()` is non- blocking: the new controller fires `was_initialized` asynchronously, `check_ready()` observes the False response in the reconcile loop, kills the actor, and the reconciler replaces it with a fresh replica. The controller-side deploy-failure counter must NOT be bumped, since the underlying cause is a previous controller crash, not user code. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm = create_dsm() info1, v1 = deployment_info(version="1") target_state_changed = dsm.deploy(TEST_DEPLOYMENT_ID, info1) assert target_state_changed dsm.save_checkpoint() ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v1)]) mocked_replica = ds._replicas.get()[0] replica_id = mocked_replica.replica_id mocked_replica._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v1)]) # Mark this replica as not-yet-initialized so that the new controller's # async `was_initialized` probe will return False. uninitialized_replicas_context.add(replica_id) new_dsm = create_dsm([replica_id.to_full_id_str()]) new_ds = new_dsm._deployment_states[TEST_DEPLOYMENT_ID] # `recover()` is non-blocking: the replica enters RECOVERING and the # probe is observed in the reconcile loop. check_counts(new_ds, total=1, by_state=[(ReplicaState.RECOVERING, 1, v1)]) starting_failures_before = new_ds._replica_constructor_retry_counter # Next update cycle: probe says False -> the replica is force-stopped # (STOPPING) and a fresh replica is started in its place (STARTING) in # the same reconcile pass. new_dsm.update() check_counts( new_ds, total=2, by_state=[(ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v1)], ) # The drop must not bump the deploy-failure counter -- the underlying # cause is a previous controller crash, not user code. assert new_ds._replica_constructor_retry_counter == starting_failures_before # The fresh replica should have a new replica id. starting_replicas = new_ds._replicas.get(states=[ReplicaState.STARTING]) assert len(starting_replicas) == 1 assert starting_replicas[0].replica_id != replica_id # Drain the STOPPING replica and confirm we're left with just the # replacement. stopping_replica = new_ds._replicas.get(states=[ReplicaState.STOPPING])[0] stopping_replica._actor.set_done_stopping() new_dsm.update() check_counts(new_ds, total=1, by_state=[(ReplicaState.STARTING, 1, v1)]) uninitialized_replicas_context.remove(replica_id) def test_actor_died_before_recover(mock_deployment_state_manager): """Test replica actor died before controller could recover it. * Deploy app / 1 deployment / 1 replica * (Simulated) Controller crashes. * Controller recovers, and tries to recover replicas from actor names. * (Simulated) The single replica from before has died before controller could recover it. * There should be 0 replicas in the deployment. * In the following control loop update cycle, the controller adds a new replica to match target state. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm = create_dsm() # Create some deployment info with actors in running state info1, v1 = deployment_info(version="1") target_state_changed = dsm.deploy(TEST_DEPLOYMENT_ID, info1) assert target_state_changed dsm.save_checkpoint() ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Single replica of version `version1` should be created and in STARTING state dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v1)]) mocked_replica = ds._replicas.get()[0] replica_id = mocked_replica.replica_id # The same replica should transition to RUNNING mocked_replica._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v1)]) # Set dead replicas context. When the controller recovers and tries # to recover replicas from actor names, the replica actor wrapper # will fail to recover. dead_replicas_context.add(replica_id) # Simulate controller crashed! A new deployment state manager should # be created, and it should call _recover_from_checkpoint new_dsm = create_dsm([replica_id.to_full_id_str()]) # Replica should fail to recover (simulate failed to get handle to # actor), meaning replica has died. new_ds = new_dsm._deployment_states[TEST_DEPLOYMENT_ID] check_counts(new_ds, total=0) # Since the previous replica is now marked dead (because controller # failed to recover it), a new replica should be added to meet # target state. new_dsm.update() check_counts(new_ds, total=1, by_state=[(ReplicaState.STARTING, 1, v1)]) dead_replicas_context.remove(replica_id) def test_shutdown(mock_deployment_state_manager): """ Test that shutdown waits for all deployments to be deleted and they are force-killed without a grace period. """ create_dsm, timer, _, _ = mock_deployment_state_manager dsm = create_dsm() grace_period_s = 10 b_info_1, _ = deployment_info( graceful_shutdown_timeout_s=grace_period_s, ) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Single replica should be created. dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) ds._replicas.get()[0]._actor.set_ready() # Now the replica should be marked running. dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) # Test shutdown flow assert not ds._replicas.get()[0]._actor.stopped # Before shutdown, `is_ready_for_shutdown()` should return False assert not dsm.is_ready_for_shutdown() dsm.shutdown() timer.advance(grace_period_s + 0.1) dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STOPPING, 1, None)]) assert ds._replicas.get()[0]._actor.stopped assert len(dsm.get_deployment_statuses()) > 0 # Once it's done stopping, replica should be removed. replica = ds._replicas.get()[0] replica._actor.set_done_stopping() dsm.update() check_counts(ds, total=0) assert len(dsm.get_deployment_statuses()) == 0 # After all deployments shutdown, `is_ready_for_shutdown()` should return True assert dsm.is_ready_for_shutdown() def test_shutdown_blocks_deploy(mock_deployment_state_manager): """After shutdown, deploy() should be a no-op and not create new deployments.""" create_dsm, timer, _, _ = mock_deployment_state_manager dsm = create_dsm() dsm.shutdown() b_info_1, _ = deployment_info(num_replicas=3) assert not dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) assert TEST_DEPLOYMENT_ID not in dsm._deployment_states dsm.update() assert TEST_DEPLOYMENT_ID not in dsm._deployment_states assert dsm.is_ready_for_shutdown() def test_shutdown_blocks_autoscale(mock_deployment_state_manager): """After shutdown, autoscale() should be a no-op.""" create_dsm, timer, _, _ = mock_deployment_state_manager dsm = create_dsm() b_info_1, _ = deployment_info() dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) dsm.update() ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] ds._replicas.get()[0]._actor.set_ready() dsm.update() dsm.shutdown() assert not dsm.autoscale(TEST_DEPLOYMENT_ID, 5) assert ds._target_state.target_num_replicas == 0 def test_shutdown_blocks_set_target_num_replicas(mock_deployment_state_manager): """After shutdown, set_target_num_replicas() should be a no-op.""" create_dsm, timer, _, _ = mock_deployment_state_manager dsm = create_dsm() b_info_1, _ = deployment_info() dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) dsm.update() ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] ds._replicas.get()[0]._actor.set_ready() dsm.update() dsm.shutdown() dsm.set_target_num_replicas(TEST_DEPLOYMENT_ID, 10) assert ds._target_state.target_num_replicas == 0 def test_shutdown_does_not_delete_checkpoint(mock_deployment_state_manager): """Tests checkpoint must survive `shutdown() and `is_ready_for_shutdown(). Only an explicit `delete_checkpoint() call should remove it. """ create_dsm, timer, _, _ = mock_deployment_state_manager dsm = create_dsm() grace_period_s = 10 b_info_1, _ = deployment_info(graceful_shutdown_timeout_s=grace_period_s) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) dsm.save_checkpoint() ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Single replica should be created and become running. dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) ds._replicas.get()[0]._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) # Checkpoint should exist after save. assert dsm._kv_store.get(CHECKPOINT_KEY) is not None # shutdown() must NOT delete the checkpoint. dsm.shutdown() assert dsm._kv_store.get(CHECKPOINT_KEY) is not None # save_checkpoint() after shutdown should be a no-op. pre_shutdown_checkpoint = dsm._kv_store.get(CHECKPOINT_KEY) dsm.save_checkpoint() assert dsm._kv_store.get(CHECKPOINT_KEY) is pre_shutdown_checkpoint timer.advance(grace_period_s + 0.1) dsm.update() replica = ds._replicas.get()[0] replica._actor.set_done_stopping() dsm.update() assert dsm.is_ready_for_shutdown() # is_ready_for_shutdown() must NOT delete the checkpoint. assert dsm._kv_store.get(CHECKPOINT_KEY) is not None # Only delete_checkpoint() should remove it. dsm.delete_checkpoint() assert dsm._kv_store.get(CHECKPOINT_KEY) is None def test_resource_requirements_none(): """Ensure resource_requirements doesn't break if a requirement is None""" class FakeActor: actor_resources = {"num_cpus": 2, "fake": None} placement_group_bundles = None available_resources = {} # Make a DeploymentReplica just to accesss its resource_requirement function replica_id = ReplicaID("asdf123", DeploymentID(name="test")) replica = DeploymentReplica(replica_id, None) replica._actor = FakeActor() # resource_requirements() should not error replica.resource_requirements() class TestActorReplicaWrapper: def test_default_value(self): actor_replica = ActorReplicaWrapper( version=deployment_version("1"), replica_id=ReplicaID( "abc123", deployment_id=DeploymentID(name="test_deployment", app_name="test_app"), ), ) assert ( actor_replica.graceful_shutdown_timeout_s == DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_S ) assert actor_replica.max_ongoing_requests == DEFAULT_MAX_ONGOING_REQUESTS assert actor_replica.health_check_period_s == DEFAULT_HEALTH_CHECK_PERIOD_S assert actor_replica.health_check_timeout_s == DEFAULT_HEALTH_CHECK_TIMEOUT_S def test_max_concurrency_override(self): actor_replica = ActorReplicaWrapper( version=deployment_version("1"), replica_id=ReplicaID( "abc123", deployment_id=DeploymentID(name="test_deployment", app_name="test_app"), ), ) max_ongoing_requests = DEFAULT_MAX_CONCURRENCY_ASYNC + 1 d_info, _ = deployment_info(max_ongoing_requests=max_ongoing_requests) replica_scheduling_request = actor_replica.start( d_info, assign_rank_callback=lambda x: 0 ) assert ( "max_concurrency" in replica_scheduling_request.actor_options and replica_scheduling_request.actor_options["max_concurrency"] == max_ongoing_requests ) def test_get_active_node_ids(mock_deployment_state_manager): """Test get_active_node_ids() are collecting the correct node ids When there are no running replicas, both methods should return empty results. When the replicas are in the RUNNING state, get_running_replica_node_ids() should return a list of all node ids. `get_active_node_ids()` should return a set of all node ids. """ node1 = NodeID.from_random().hex() node2 = NodeID.from_random().hex() node_ids = (node1, node2, node2) create_dsm, _, cluster_node_info_cache, _ = mock_deployment_state_manager dsm = create_dsm() cluster_node_info_cache.add_node(node1) cluster_node_info_cache.add_node(node2) # Deploy deployment with version "1" and 3 replicas info1, v1 = deployment_info(version="1", num_replicas=3) assert dsm.deploy(TEST_DEPLOYMENT_ID, info1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # When the replicas are in the STARTING state, `get_active_node_ids()` should # return a set of node ids. dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.STARTING, 3, v1)]) mocked_replicas = ds._replicas.get() for idx, mocked_replica in enumerate(mocked_replicas): mocked_replica._actor.set_node_id(node_ids[idx]) assert ds.get_active_node_ids() == set(node_ids) assert dsm.get_active_node_ids() == set(node_ids) # When the replicas are in RUNNING state, `get_active_node_ids()` should # return a set of `node_ids`. for mocked_replica in mocked_replicas: mocked_replica._actor.set_ready() dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v1)]) assert ds.get_active_node_ids() == set(node_ids) assert dsm.get_active_node_ids() == set(node_ids) for _ in mocked_replicas: ds._stop_one_running_replica_for_testing() dsm.update() check_counts( ds, total=6, by_state=[(ReplicaState.STOPPING, 3, v1), (ReplicaState.STARTING, 3, v1)], ) def test_get_active_node_ids_none(mock_deployment_state_manager): """Test get_active_node_ids() are not collecting none node ids. When the running replicas has None as the node id, `get_active_node_ids()` should not include it in the set. """ node1 = NodeID.from_random().hex() node2 = NodeID.from_random().hex() node_ids = (node1, node2, node2) create_dsm, _, cluster_node_info_cache, _ = mock_deployment_state_manager dsm = create_dsm() cluster_node_info_cache.add_node(node1) cluster_node_info_cache.add_node(node2) # Deploy deployment with version "1" and 3 replicas info1, v1 = deployment_info(version="1", num_replicas=3) assert dsm.deploy(TEST_DEPLOYMENT_ID, info1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # When the replicas are in the STARTING state, `get_active_node_ids()` should # return a set of node ids. dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.STARTING, 3, v1)]) mocked_replicas = ds._replicas.get() for idx, mocked_replica in enumerate(mocked_replicas): mocked_replica._actor.set_node_id(node_ids[idx]) assert ds.get_active_node_ids() == set(node_ids) assert dsm.get_active_node_ids() == set(node_ids) # When the replicas are in the RUNNING state and are having None node id, # `get_active_node_ids()` should return empty set. for mocked_replica in mocked_replicas: mocked_replica._actor.set_node_id(None) mocked_replica._actor.set_ready() dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v1)]) assert None not in ds.get_active_node_ids() assert None not in dsm.get_active_node_ids() def test_get_deployment_ids(mock_deployment_state_manager): create_dsm, _, _, _ = mock_deployment_state_manager dsm = create_dsm() assert dsm.get_deployment_ids() == [] info1, _ = deployment_info(version="1", num_replicas=1) info2, _ = deployment_info(version="2", num_replicas=1) assert dsm.deploy(TEST_DEPLOYMENT_ID, info1) assert dsm.deploy(TEST_DEPLOYMENT_ID_2, info2) assert dsm.get_deployment_ids() == [TEST_DEPLOYMENT_ID, TEST_DEPLOYMENT_ID_2] def test_get_node_id_to_alive_replica_ids(mock_deployment_state_manager): node1 = NodeID.from_random().hex() node2 = NodeID.from_random().hex() create_dsm, _, cluster_node_info_cache, _ = mock_deployment_state_manager dsm = create_dsm() cluster_node_info_cache.add_node(node1) cluster_node_info_cache.add_node(node2) info1, v1 = deployment_info(version="1", num_replicas=2) info2, v2 = deployment_info(version="2", num_replicas=1) assert dsm.deploy(TEST_DEPLOYMENT_ID, info1) assert dsm.deploy(TEST_DEPLOYMENT_ID_2, info2) ds1 = dsm._deployment_states[TEST_DEPLOYMENT_ID] ds2 = dsm._deployment_states[TEST_DEPLOYMENT_ID_2] dsm.update() check_counts(ds1, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) check_counts(ds2, total=1, by_state=[(ReplicaState.STARTING, 1, v2)]) replicas1 = ds1._replicas.get() replicas2 = ds2._replicas.get() replicas1[0]._actor.set_node_id(node1) replicas1[1]._actor.set_node_id(node2) replicas2[0]._actor.set_node_id(node1) node_id_to_alive_replica_ids = dsm.get_node_id_to_alive_replica_ids() assert type(node_id_to_alive_replica_ids) is dict assert node_id_to_alive_replica_ids == { node1: { replicas1[0].replica_id.unique_id, replicas2[0].replica_id.unique_id, }, node2: {replicas1[1].replica_id.unique_id}, } replicas1[0]._actor.set_ready() replicas1[1]._actor.set_ready() replicas2[0]._actor.set_node_id(None) replicas2[0]._actor.set_ready() dsm.update() node_id_to_alive_replica_ids = dsm.get_node_id_to_alive_replica_ids() assert type(node_id_to_alive_replica_ids) is dict assert node_id_to_alive_replica_ids == { node1: {replicas1[0].replica_id.unique_id}, node2: {replicas1[1].replica_id.unique_id}, } def test_dump_replica_states_for_testing(mock_deployment_state_manager): create_dsm, _, _, _ = mock_deployment_state_manager dsm = create_dsm() info1, _ = deployment_info(version="1", num_replicas=1) assert dsm.deploy(TEST_DEPLOYMENT_ID, info1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] assert dsm._dump_replica_states_for_testing(TEST_DEPLOYMENT_ID) is ds._replicas with pytest.raises(KeyError): dsm._dump_replica_states_for_testing(TEST_DEPLOYMENT_ID_2) def test_stop_one_running_replica_for_testing(mock_deployment_state_manager): create_dsm, _, _, _ = mock_deployment_state_manager dsm = create_dsm() info1, _ = deployment_info(version="1", num_replicas=1) assert dsm.deploy(TEST_DEPLOYMENT_ID, info1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() replica = ds._replicas.get()[0] replica._actor.set_ready() dsm.update() assert len(ds._replicas.get([ReplicaState.RUNNING])) == 1 dsm._stop_one_running_replica_for_testing(TEST_DEPLOYMENT_ID) assert len(ds._replicas.get([ReplicaState.RUNNING])) == 0 assert len(ds._replicas.get([ReplicaState.STOPPING])) == 1 class TestAutoscaling: def scale( self, dsm: DeploymentStateManager, asm: AutoscalingStateManager, deployment_ids: List[DeploymentID], ): if not deployment_ids: return app_name = deployment_ids[0].app_name assert all(dep_id.app_name == app_name for dep_id in deployment_ids) deployment_to_target_num_replicas = { dep_id: dsm.get_deployment_details(dep_id).target_num_replicas for dep_id in deployment_ids } decisions = asm.get_decision_num_replicas( app_name, deployment_to_target_num_replicas ) for deployment_id, decision_num_replicas in decisions.items(): dsm.autoscale(deployment_id, decision_num_replicas) @pytest.mark.parametrize("target_capacity_direction", ["up", "down"]) def test_basic_autoscaling( self, mock_deployment_state_manager, target_capacity_direction ): """Test autoscaling up and down. Upscaling version: 1. Deploy deployment with autoscaling limits [0,6], initial_replicas=3, target=1. 2. It becomes healthy with 3 running replicas. 3. Set average request metrics to 2 (compare to target=1). 4. Deployment autoscales, 3 replicas starting, status=UPSCALING, trigger=AUTOSCALE. 5. It becomes healthy with 6 running replicas, status=HEALTHY, trigger=UPSCALE. """ # Create deployment state manager create_dsm, timer, _, asm = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() asm: AutoscalingStateManager = asm # Deploy deployment with 3 replicas info, _ = deployment_info( autoscaling_config={ "target_ongoing_requests": 1, "min_replicas": 0, "max_replicas": 6, "initial_replicas": 3, "upscale_delay_s": 0, "downscale_delay_s": 0, "metrics_interval_s": 100, "look_back_period_s": 200, } ) dsm.deploy(TEST_DEPLOYMENT_ID, info) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] # status=UPDATING, status_trigger=DEPLOY dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.STARTING, 3, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Set replicas ready and check statuses for replica in ds._replicas.get(): replica._actor.set_ready() # status=HEALTHY, status_trigger=DEPLOY dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) req_per_replica = 2 if target_capacity_direction == "up" else 0 replicas = ds._replicas.get() if RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE: handle_metric_report = HandleMetricReport( deployment_id=TEST_DEPLOYMENT_ID, handle_id="random", actor_id="actor_id", handle_source=DeploymentHandleSource.UNKNOWN, queued_requests=[TimeStampedValue(timer.time() - 0.1, 0)], aggregated_queued_requests=0, aggregated_metrics={ RUNNING_REQUESTS_KEY: { replica._actor.replica_id.to_full_id_str(): req_per_replica for replica in replicas } }, metrics={ RUNNING_REQUESTS_KEY: { replica._actor.replica_id.to_full_id_str(): [ TimeStampedValue(timer.time() - 0.1, req_per_replica) ] for replica in replicas } }, timestamp=timer.time(), ) asm.record_request_metrics_for_handle(handle_metric_report) else: for replica in replicas: replica_metric_report = ReplicaMetricReport( replica_id=replica._actor.replica_id, aggregated_metrics={RUNNING_REQUESTS_KEY: req_per_replica}, metrics={ RUNNING_REQUESTS_KEY: [ TimeStampedValue(timer.time() - 0.1, req_per_replica) ] }, timestamp=timer.time(), ) asm.record_request_metrics_for_replica(replica_metric_report) self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) # status=UPSCALING/DOWNSCALING, status_trigger=AUTOSCALE dsm.update() if target_capacity_direction == "up": check_counts( ds, total=6, by_state=[ (ReplicaState.RUNNING, 3, None), (ReplicaState.STARTING, 3, None), ], ) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.AUTOSCALING ) # Advance timer by 60 seconds; this should exceed the slow startup # warning threshold. The message should be updated, but the status # should remain upscaling/autoscaling timer.advance(60) dsm.update() check_counts( ds, total=6, by_state=[ (ReplicaState.RUNNING, 3, None), (ReplicaState.STARTING, 3, None), ], ) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.AUTOSCALING ) assert "have taken more than" in ds.curr_status_info.message # Set replicas ready for replica in ds._replicas.get(): replica._actor.set_ready() else: # Due to two-stage downscaling one of the replicas will still be running check_counts( ds, total=3, by_state=[ (ReplicaState.STOPPING, 2, None), (ReplicaState.RUNNING, 1, None), ], ) # Trigger the second stage of downscaling self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.STOPPING, 3, None)]) assert ds.curr_status_info.status == DeploymentStatus.DOWNSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.AUTOSCALING ) for replica in ds._replicas.get(): replica._actor.set_done_stopping() dsm.update() astate = asm._app_autoscaling_states[ TEST_DEPLOYMENT_ID.app_name ]._deployment_autoscaling_states[TEST_DEPLOYMENT_ID] assert len(astate._replica_metrics) == 0 # status=HEALTHY, status_trigger=UPSCALE/DOWNSCALE dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ds.curr_status_info.status_trigger == ( DeploymentStatusTrigger.UPSCALE_COMPLETED if target_capacity_direction == "up" else DeploymentStatusTrigger.DOWNSCALE_COMPLETED ) # Make sure autoscaling state is removed when deployment is deleted dsm.delete_deployment(TEST_DEPLOYMENT_ID) dsm.update() for replica in ds._replicas.get(): replica._actor.set_done_stopping() dsm.update() assert TEST_DEPLOYMENT_ID not in dsm._deployment_states @pytest.mark.parametrize( "target_startup_status", [ ReplicaStartupStatus.PENDING_ALLOCATION, ReplicaStartupStatus.PENDING_INITIALIZATION, ], ) def test_downscaling_reclaiming_starting_replicas_first( self, target_startup_status, mock_deployment_state_manager, ): """This test asserts that when downscaling first any non-running replicas are scavenged, before stopping fully running replicas More context on the issue could be found in: https://github.com/ray-project/ray/issues/43034 """ # Create deployment state manager create_dsm, timer, _, asm = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() asm: AutoscalingStateManager = asm # Deploy deployment with 3 replicas info, _ = deployment_info( autoscaling_config={ "target_ongoing_requests": 1, "min_replicas": 0, "max_replicas": 6, "initial_replicas": 3, "upscale_delay_s": 0, "downscale_delay_s": 0, "metrics_interval_s": 100, "look_back_period_s": 200, } ) dsm.deploy(TEST_DEPLOYMENT_ID, info) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] # status=UPDATING, status_trigger=DEPLOY dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.STARTING, 3, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Set replicas as SUCCESSFUL and check statuses for replica in ds._replicas.get(): replica._actor.set_ready() # status=HEALTHY, status_trigger=DEPLOY dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Fetch all currently running replicas running_replicas = ds._replicas.get(states=[ReplicaState.RUNNING]) replicas = ds._replicas.get() if RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE: handle_metric_report = HandleMetricReport( deployment_id=TEST_DEPLOYMENT_ID, handle_id="random", actor_id="actor_id", handle_source=DeploymentHandleSource.UNKNOWN, queued_requests=[TimeStampedValue(timer.time() - 0.1, 0)], aggregated_queued_requests=0, aggregated_metrics={ RUNNING_REQUESTS_KEY: { replica._actor.replica_id.to_full_id_str(): 2 for replica in replicas } }, metrics={ RUNNING_REQUESTS_KEY: { replica._actor.replica_id.to_full_id_str(): [ TimeStampedValue(timer.time() - 0.1, 2) ] for replica in replicas } }, timestamp=timer.time(), ) asm.record_request_metrics_for_handle(handle_metric_report) else: for replica in replicas: replica_metric_report = ReplicaMetricReport( replica_id=replica._actor.replica_id, aggregated_metrics={RUNNING_REQUESTS_KEY: 2}, metrics={ RUNNING_REQUESTS_KEY: [TimeStampedValue(timer.time() - 0.1, 2)] }, timestamp=timer.time(), ) asm.record_request_metrics_for_replica(replica_metric_report) # status=UPSCALING, status_trigger=AUTOSCALE self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) dsm.update() check_counts( ds, total=6, by_state=[ (ReplicaState.RUNNING, 3, None), (ReplicaState.STARTING, 3, None), ], ) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING assert ds.curr_status_info.status_trigger == DeploymentStatusTrigger.AUTOSCALING # Set replicas as PENDING_INITIALIZATION: actors have been # successfully allocated, but replicas are still pending # successful initialization for replica in ds._replicas.get(): replica._actor.set_status(target_startup_status) # Advance timer by 60 seconds; this should exceed the slow startup # warning threshold. The message should be updated, but the status # should remain upscaling/autoscaling timer.advance(60) dsm.update() check_counts( ds, total=6, by_state=[ (ReplicaState.RUNNING, 3, None), (ReplicaState.STARTING, 3, None), ], ) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING assert ds.curr_status_info.status_trigger == DeploymentStatusTrigger.AUTOSCALING if target_startup_status == ReplicaStartupStatus.PENDING_INITIALIZATION: expected_message = ( "Deployment 'test_deployment' in application 'test_app' has 3 replicas " f"that have taken more than {SLOW_STARTUP_WARNING_S}s to initialize.\n" "This may be caused by a slow __init__ or reconfigure method." ) elif target_startup_status == ReplicaStartupStatus.PENDING_ALLOCATION: expected_message = ( "Deployment 'test_deployment' in application 'test_app' " "has 3 replicas that have taken more than 30s to be scheduled. " "This may be due to waiting for the cluster to auto-scale or for " "a runtime environment to be installed. " "Resources required for each replica: " '{"CPU": 0.1}, ' "total resources available: " "{}. " "Use `ray status` for more details." ) else: raise RuntimeError(f"Got unexpected status: {target_startup_status}") assert expected_message == ds.curr_status_info.message # Now, trigger downscaling attempting to reclaim half (3) of the replicas replicas = ds._replicas.get(states=[ReplicaState.RUNNING]) if RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE: handle_metric_report = HandleMetricReport( deployment_id=TEST_DEPLOYMENT_ID, handle_id="random", actor_id="actor_id", handle_source=DeploymentHandleSource.UNKNOWN, queued_requests=[TimeStampedValue(timer.time() - 0.1, 0)], aggregated_queued_requests=0, aggregated_metrics={ RUNNING_REQUESTS_KEY: { replica._actor.replica_id.to_full_id_str(): 1 for replica in replicas } }, metrics={ RUNNING_REQUESTS_KEY: { replica._actor.replica_id.to_full_id_str(): [ TimeStampedValue(timer.time() - 0.1, 1) ] for replica in replicas } }, timestamp=timer.time(), ) asm.record_request_metrics_for_handle(handle_metric_report) else: for replica in replicas: replica_metric_report = ReplicaMetricReport( replica_id=replica._actor.replica_id, aggregated_metrics={RUNNING_REQUESTS_KEY: 1}, metrics={ RUNNING_REQUESTS_KEY: [TimeStampedValue(timer.time() - 0.1, 1)] }, timestamp=timer.time(), ) asm.record_request_metrics_for_replica(replica_metric_report) # status=DOWNSCALING, status_trigger=AUTOSCALE self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) dsm.update() check_counts( ds, total=6, by_state=[ (ReplicaState.RUNNING, 3, None), (ReplicaState.STOPPING, 3, None), ], ) # Assert that no RUNNING replicas are being stopped assert running_replicas == ds._replicas.get(states=[ReplicaState.RUNNING]) assert ds.curr_status_info.status == DeploymentStatus.DOWNSCALING assert ds.curr_status_info.status_trigger == DeploymentStatusTrigger.AUTOSCALING for replica in ds._replicas.get(): replica._actor.set_done_stopping() # status=HEALTHY, status_trigger=UPSCALE/DOWNSCALE dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.DOWNSCALE_COMPLETED ) def test_update_autoscaling_config(self, mock_deployment_state_manager): """Test updating the autoscaling config. 1. Deploy deployment with autoscaling limits [0,6] and initial replicas = 3. 2. It becomes healthy with 3 running replicas. 3. Update autoscaling config to limits [6,10]. 4. 3 new replicas should be STARTING, and deployment status should be UPDATING. 5. It becomes healthy with 6 running replicas. """ # Create deployment state manager create_dsm, timer, _, asm = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() asm: AutoscalingStateManager = asm # Deploy deployment with 3 replicas info1, _ = deployment_info( autoscaling_config={ "target_ongoing_requests": 1, "min_replicas": 0, "max_replicas": 6, "initial_replicas": 3, "upscale_delay_s": 0, "downscale_delay_s": 0, }, version="1", ) dsm.deploy(TEST_DEPLOYMENT_ID, info1) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Set replicas ready dsm.update() for replica in ds._replicas.get(): replica._actor.set_ready() # status=HEALTHY, status_trigger=DEPLOY dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Num ongoing requests = 1, status should remain HEALTHY replicas = ds._replicas.get() if RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE: handle_metric_report = HandleMetricReport( deployment_id=TEST_DEPLOYMENT_ID, handle_id="random", actor_id="actor_id", handle_source=DeploymentHandleSource.UNKNOWN, queued_requests=[TimeStampedValue(timer.time() - 0.1, 0)], aggregated_queued_requests=0, aggregated_metrics={ RUNNING_REQUESTS_KEY: { replica._actor.replica_id.to_full_id_str(): 1 for replica in replicas } }, metrics={ RUNNING_REQUESTS_KEY: { replica._actor.replica_id.to_full_id_str(): [ TimeStampedValue(timer.time() - 0.1, 1) ] for replica in replicas } }, timestamp=timer.time(), ) asm.record_request_metrics_for_handle(handle_metric_report) else: for replica in replicas: replica_metric_report = ReplicaMetricReport( replica_id=replica._actor.replica_id, aggregated_metrics={RUNNING_REQUESTS_KEY: 1}, metrics={ RUNNING_REQUESTS_KEY: [TimeStampedValue(timer.time() - 0.1, 1)] }, timestamp=timer.time(), ) asm.record_request_metrics_for_replica(replica_metric_report) check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # Update autoscaling config info2, _ = deployment_info( autoscaling_config={ "target_ongoing_requests": 1, "min_replicas": 6, "max_replicas": 10, "upscale_delay_s": 0, "downscale_delay_s": 0, }, version="1", ) dsm.deploy(TEST_DEPLOYMENT_ID, info2) # 3 new replicas should be starting, status should be UPDATING (not upscaling) self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) dsm.update() check_counts( ds, total=6, by_state=[ (ReplicaState.RUNNING, 3, None), (ReplicaState.STARTING, 3, None), ], ) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Set replicas ready dsm.update() for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=6, by_state=[(ReplicaState.RUNNING, 6, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) def test_replicas_fail_during_initial_scale_from_zero( self, mock_deployment_state_manager ): """Test the following case: - An "erroneous" deployment (w/ autoscaling enabled) is deployed with initial replicas set to 0. Since no replicas are started, no errors have occurred from trying to start the replicas yet. - A request is sent, triggering an upscale. - The controller tries to start new replicas, but fails because of a constructor error. In this case, the deployment should transition to UNHEALTHY and stop retrying after a threshold. """ create_dsm, timer, _, asm = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() asm: AutoscalingStateManager = asm # Deploy deployment with 1 initial replica info, _ = deployment_info( autoscaling_config={ "target_ongoing_requests": 1, "min_replicas": 0, "max_replicas": 2, "initial_replicas": 0, "upscale_delay_s": 0, "downscale_delay_s": 0, "metrics_interval_s": 100, "look_back_period_s": 200, } ) dsm.deploy(TEST_DEPLOYMENT_ID, info) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Send request metrics to controller to make the deployment upscale handle_metric_report = HandleMetricReport( deployment_id=TEST_DEPLOYMENT_ID, handle_id="random", actor_id="actor_id", handle_source=DeploymentHandleSource.UNKNOWN, queued_requests=[TimeStampedValue(timer.time() - 0.1, 1)], aggregated_queued_requests=1, aggregated_metrics={}, metrics={}, timestamp=timer.time(), ) asm.record_request_metrics_for_handle(handle_metric_report) self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) # The controller should try to start a new replica. If that replica repeatedly # fails to start, the deployment should transition to UNHEALTHY and NOT retry # replicas anymore for i in range(10): dsm.update() if i < 3: check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.AUTOSCALING ) # Set replica failed to start replica = ds._replicas.get()[0] replica._actor.set_failed_to_start() dsm.update() if i < 2: check_counts( ds, total=2, by_state=[ (ReplicaState.STOPPING, 1, None), (ReplicaState.STARTING, 1, None), ], ) else: check_counts( ds, total=1, by_state=[(ReplicaState.STOPPING, 1, None)] ) # Set replica finished stopping replica._actor.set_done_stopping() else: check_counts(ds, total=0) assert ds.curr_status_info.status == DeploymentStatus.UNHEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.REPLICA_STARTUP_FAILED ) def test_replicas_fail_during_subsequent_scale_from_zero( self, mock_deployment_state_manager ): """Test the following case: - An autoscaling deployment is deployed and it reaches HEALTHY with a non-zero number of replicas. - After a period of no traffic, the deployment scales down to 0. - New traffic is sent, triggering an upscale. - The controller tries to start new replicas, but for some reason some replicas fail to start because of transient errors In this case, the deployment should transition to UNHEALTHY and keep retrying, since at least one replica of this version has successfully started in the past, meaning we don't know if it is an unrecoverable user code error. """ create_dsm, timer, _, asm = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() asm: AutoscalingStateManager = asm # Deploy deployment with 1 initial replica info, _ = deployment_info( autoscaling_config={ "target_ongoing_requests": 1, "min_replicas": 0, "max_replicas": 2, "initial_replicas": 1, "upscale_delay_s": 0, "downscale_delay_s": 0, "metrics_interval_s": 100, "look_back_period_s": 200, } ) dsm.deploy(TEST_DEPLOYMENT_ID, info) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Expected: status=UPDATING, status_trigger=CONFIG_UPDATED_STARTED dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) # Set replicas ready and check statuses # Expected: status=HEALTHY, status_trigger=CONFIG_UPDATED_COMPLETED ds._replicas.get()[0]._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_COMPLETED ) # There are no requests, so the deployment should be downscaled to zero. self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STOPPING, 1, None)]) ds._replicas.get()[0]._actor.set_done_stopping() dsm.update() check_counts(ds, total=0) # Send request metrics to controller to make the deployment upscale handle_metric_report = HandleMetricReport( deployment_id=TEST_DEPLOYMENT_ID, handle_id="random", actor_id="actor_id", handle_source=DeploymentHandleSource.UNKNOWN, queued_requests=[TimeStampedValue(timer.time() - 0.1, 1)], aggregated_queued_requests=1, aggregated_metrics={}, metrics={}, timestamp=timer.time(), ) asm.record_request_metrics_for_handle(handle_metric_report) self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) # The controller should try to start a new replica. If that replica repeatedly # fails to start, the deployment should transition to UNHEALTHY. Meanwhile # the controller should continue retrying after 3 times. for i in range(10): dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) if i < 3: assert ds.curr_status_info.status == DeploymentStatus.UPSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.AUTOSCALING ) else: assert ds.curr_status_info.status == DeploymentStatus.UNHEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.REPLICA_STARTUP_FAILED ) # Set replica failed to start replica = ds._replicas.get()[0] replica._actor.set_failed_to_start() dsm.update() check_counts( ds, total=2, by_state=[ (ReplicaState.STOPPING, 1, None), (ReplicaState.STARTING, 1, None), ], ) # Set replica finished stopping replica._actor.set_done_stopping() @pytest.mark.skipif( not RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE, reason="Testing handle metrics behavior.", ) def test_handle_metrics_timeout(self, mock_deployment_state_manager): create_dsm, timer, _, asm = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() asm: AutoscalingStateManager = asm # Deploy, start with 1 replica info, _ = deployment_info( autoscaling_config={ "target_ongoing_requests": 1, "min_replicas": 0, "max_replicas": 6, "initial_replicas": 1, "upscale_delay_s": 0, "downscale_delay_s": 0, } ) dsm.deploy(TEST_DEPLOYMENT_ID, info) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() ds._replicas.get()[0]._actor.set_ready() asm.drop_stale_handle_metrics(dsm.get_alive_replica_actor_ids()) dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) # Record 2 requests/replica -> trigger upscale handle_metric_report = HandleMetricReport( deployment_id=TEST_DEPLOYMENT_ID, handle_id="random", actor_id="actor_id", handle_source=DeploymentHandleSource.UNKNOWN, queued_requests=[TimeStampedValue(timer.time() - 0.1, 0)], aggregated_queued_requests=0, aggregated_metrics={ RUNNING_REQUESTS_KEY: { ds._replicas.get()[0]._actor.replica_id.to_full_id_str(): 2 } }, metrics={ RUNNING_REQUESTS_KEY: { ds._replicas.get()[0]._actor.replica_id.to_full_id_str(): [ TimeStampedValue(timer.time() - 0.1, 2) ] } }, timestamp=timer.time(), ) asm.record_request_metrics_for_handle(handle_metric_report) asm.drop_stale_handle_metrics(dsm.get_alive_replica_actor_ids()) self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) dsm.update() check_counts( ds, total=2, by_state=[ (ReplicaState.RUNNING, 1, None), (ReplicaState.STARTING, 1, None), ], ) assert asm.get_total_num_requests_for_deployment(TEST_DEPLOYMENT_ID) == 2 ds._replicas.get([ReplicaState.STARTING])[0]._actor.set_ready() asm.drop_stale_handle_metrics(dsm.get_alive_replica_actor_ids()) self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, None)]) assert asm.get_total_num_requests_for_deployment(TEST_DEPLOYMENT_ID) == 2 # Simulate handle was on an actor that died. 10 seconds later # the handle fails to push metrics timer.advance(10) asm.drop_stale_handle_metrics(dsm.get_alive_replica_actor_ids()) self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, None)]) assert asm.get_total_num_requests_for_deployment(TEST_DEPLOYMENT_ID) == 2 # Another 10 seconds later handle still fails to push metrics. At # this point the data from the handle should be invalidated. As a # result, the replicas should scale back down to 0. timer.advance(10) asm.drop_stale_handle_metrics(dsm.get_alive_replica_actor_ids()) # The first update will trigger the first stage of downscaling to 1 self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) dsm.update() check_counts( ds, total=2, by_state=[ (ReplicaState.STOPPING, 1, None), (ReplicaState.RUNNING, 1, None), ], ) # The second update will trigger the second stage of downscaling from 1 to 0 self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STOPPING, 2, None)]) assert asm.get_total_num_requests_for_deployment(TEST_DEPLOYMENT_ID) == 0 @pytest.mark.skipif( not RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE, reason="Testing handle metrics behavior.", ) def test_handle_metrics_on_dead_serve_actor(self, mock_deployment_state_manager): """Metrics for handles on dead serve actors should be dropped.""" create_dsm, timer, _, asm = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() asm: AutoscalingStateManager = asm d_id1 = DeploymentID("d1", "app") d_id2 = DeploymentID("d2", "app") # Deploy, start with 1 replica info1, _ = deployment_info( autoscaling_config={ "target_ongoing_requests": 1, "min_replicas": 0, "max_replicas": 6, "initial_replicas": 1, "upscale_delay_s": 0, "downscale_delay_s": 0, }, ) info2, _ = deployment_info(health_check_period_s=0.1) dsm.deploy(d_id1, info1) dsm.deploy(d_id2, info2) ds1: DeploymentState = dsm._deployment_states[d_id1] ds2: DeploymentState = dsm._deployment_states[d_id2] # One replica each asm.drop_stale_handle_metrics(dsm.get_alive_replica_actor_ids()) dsm.update() ds1._replicas.get()[0]._actor.set_ready() ds2._replicas.get()[0]._actor.set_ready() ds2._replicas.get()[0]._actor.set_actor_id("d2_replica_actor_id") asm.drop_stale_handle_metrics(dsm.get_alive_replica_actor_ids()) dsm.update() check_counts(ds1, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) check_counts(ds2, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) # Record 2 requests/replica (sent from d2 replica) -> trigger upscale handle_metric_report = HandleMetricReport( deployment_id=d_id1, handle_id="random", actor_id="d2_replica_actor_id", handle_source=DeploymentHandleSource.REPLICA, queued_requests=[TimeStampedValue(timer.time() - 0.1, 0)], aggregated_queued_requests=0, aggregated_metrics={ RUNNING_REQUESTS_KEY: { ds1._replicas.get()[0]._actor.replica_id.to_full_id_str(): 2 } }, metrics={ RUNNING_REQUESTS_KEY: { ds1._replicas.get()[0]._actor.replica_id.to_full_id_str(): [ TimeStampedValue(timer.time() - 0.1, 2) ] } }, timestamp=timer.time(), ) asm.record_request_metrics_for_handle(handle_metric_report) asm.drop_stale_handle_metrics(dsm.get_alive_replica_actor_ids()) self.scale(dsm, asm, [d_id1, d_id2]) dsm.update() check_counts( ds1, total=2, by_state=[ (ReplicaState.RUNNING, 1, None), (ReplicaState.STARTING, 1, None), ], ) assert asm.get_total_num_requests_for_deployment(d_id1) == 2 ds1._replicas.get([ReplicaState.STARTING])[0]._actor.set_ready() asm.drop_stale_handle_metrics(dsm.get_alive_replica_actor_ids()) self.scale(dsm, asm, [d_id1, d_id2]) dsm.update() check_counts(ds1, total=2, by_state=[(ReplicaState.RUNNING, 2, None)]) assert asm.get_total_num_requests_for_deployment(d_id1) == 2 # d2 replica died ds2._replicas.get()[0]._actor.set_unhealthy() asm.drop_stale_handle_metrics(dsm.get_alive_replica_actor_ids()) self.scale(dsm, asm, [d_id1, d_id2]) dsm.update() check_counts( ds2, total=2, by_state=[ (ReplicaState.STARTING, 1, None), (ReplicaState.STOPPING, 1, None), ], ) ds2._replicas.get(states=[ReplicaState.STOPPING])[0]._actor.set_done_stopping() asm.drop_stale_handle_metrics(dsm.get_alive_replica_actor_ids()) self.scale(dsm, asm, [d_id1, d_id2]) dsm.update() check_counts(ds2, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) # Now that the d2 replica is dead, its metrics should be dropped. # Consequently d1 should scale down to 0 replicas asm.drop_stale_handle_metrics(dsm.get_alive_replica_actor_ids()) self.scale(dsm, asm, [d_id1, d_id2]) dsm.update() # Due to two-stage downscaling one of the replicas will still be running check_counts( ds1, total=2, by_state=[ (ReplicaState.STOPPING, 1, None), (ReplicaState.RUNNING, 1, None), ], ) # Trigger the second stage of downscaling self.scale(dsm, asm, [d_id1, d_id2]) dsm.update() check_counts(ds1, total=2, by_state=[(ReplicaState.STOPPING, 2, None)]) def test_autoscaling_timestamps(self, mock_deployment_state_manager): """Test that last_scale_up_time and last_scale_down_time are properly tracked. This test verifies that: 1. Timestamps are None initially 2. last_scale_up_time is set after a scale-up event 3. last_scale_down_time is set after a scale-down event 4. Timestamps are available in AutoscalingContext """ # Create deployment state manager create_dsm, timer, _, asm = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() asm: AutoscalingStateManager = asm # Deploy deployment with autoscaling info, _ = deployment_info( autoscaling_config={ "target_ongoing_requests": 1, "min_replicas": 1, "max_replicas": 5, "initial_replicas": 2, "upscale_delay_s": 0, "downscale_delay_s": 0, "metrics_interval_s": 100, "look_back_period_s": 200, } ) dsm.deploy(TEST_DEPLOYMENT_ID, info) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Make replicas ready dsm.update() for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() # Get autoscaling state app_state = asm._app_autoscaling_states[TEST_DEPLOYMENT_ID.app_name] dep_autoscaling_state = app_state._deployment_autoscaling_states[ TEST_DEPLOYMENT_ID ] # Initially, timestamps should be None ctx = dep_autoscaling_state.get_autoscaling_context(2) assert ctx.last_scale_up_time is None assert ctx.last_scale_down_time is None # Trigger scale-up by setting high request metrics replicas = ds._replicas.get() req_per_replica = 5 # High load to trigger scale-up if RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE: handle_metric_report = HandleMetricReport( deployment_id=TEST_DEPLOYMENT_ID, handle_id="test_handle", actor_id="test_actor", handle_source=DeploymentHandleSource.UNKNOWN, queued_requests=[TimeStampedValue(timer.time() - 0.1, 0)], aggregated_queued_requests=0, aggregated_metrics={ RUNNING_REQUESTS_KEY: { replica._actor.replica_id.to_full_id_str(): req_per_replica for replica in replicas } }, metrics={ RUNNING_REQUESTS_KEY: { replica._actor.replica_id.to_full_id_str(): [ TimeStampedValue(timer.time() - 0.1, req_per_replica) ] for replica in replicas } }, timestamp=timer.time(), ) asm.record_request_metrics_for_handle(handle_metric_report) else: for replica in replicas: replica_metric_report = ReplicaMetricReport( replica_id=replica._actor.replica_id, aggregated_metrics={RUNNING_REQUESTS_KEY: req_per_replica}, metrics={ RUNNING_REQUESTS_KEY: [ TimeStampedValue(timer.time() - 0.1, req_per_replica) ] }, timestamp=timer.time(), ) asm.record_request_metrics_for_replica(replica_metric_report) # Record time before scale-up time_before_scale_up = timer.time() # Trigger autoscaling decision self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) # After scale-up, last_scale_up_time should be set and greater than the time before ctx_after_scale_up = dep_autoscaling_state.get_autoscaling_context(5) assert ctx_after_scale_up.last_scale_up_time is not None assert ctx_after_scale_up.last_scale_up_time >= time_before_scale_up assert ctx_after_scale_up.last_scale_down_time is None scale_up_time = ctx_after_scale_up.last_scale_up_time # Advance timer to simulate time passing timer.advance(10) # Set replicas ready dsm.update() for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() # Now trigger scale-down by setting low request metrics replicas = ds._replicas.get() req_per_replica = 0 # No load to trigger scale-down if RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE: handle_metric_report = HandleMetricReport( deployment_id=TEST_DEPLOYMENT_ID, handle_id="test_handle", actor_id="test_actor", handle_source=DeploymentHandleSource.UNKNOWN, queued_requests=[TimeStampedValue(timer.time() - 0.1, 0)], aggregated_queued_requests=0, aggregated_metrics={ RUNNING_REQUESTS_KEY: { replica._actor.replica_id.to_full_id_str(): req_per_replica for replica in replicas } }, metrics={ RUNNING_REQUESTS_KEY: { replica._actor.replica_id.to_full_id_str(): [ TimeStampedValue(timer.time() - 0.1, req_per_replica) ] for replica in replicas } }, timestamp=timer.time(), ) asm.record_request_metrics_for_handle(handle_metric_report) else: for replica in replicas: replica_metric_report = ReplicaMetricReport( replica_id=replica._actor.replica_id, aggregated_metrics={RUNNING_REQUESTS_KEY: req_per_replica}, metrics={ RUNNING_REQUESTS_KEY: [ TimeStampedValue(timer.time() - 0.1, req_per_replica) ] }, timestamp=timer.time(), ) asm.record_request_metrics_for_replica(replica_metric_report) # Record time before scale-down time_before_scale_down = timer.time() # Trigger autoscaling decision for scale-down self.scale(dsm, asm, [TEST_DEPLOYMENT_ID]) # After scale-down, last_scale_down_time should be set and greater than the time before ctx_after_scale_down = dep_autoscaling_state.get_autoscaling_context(1) assert ( ctx_after_scale_down.last_scale_up_time == scale_up_time ) # Should remain unchanged assert ctx_after_scale_down.last_scale_down_time is not None assert ctx_after_scale_down.last_scale_down_time >= time_before_scale_down assert ctx_after_scale_down.last_scale_down_time > scale_up_time class TestTargetCapacity: """ Tests related to the `target_capacity` field that adjusts the target num_replicas. """ def update_target_capacity( self, deployment_state: DeploymentState, curr_deployment_info: DeploymentInfo, target_capacity: Optional[float], target_capacity_direction: Optional[TargetCapacityDirection], ): new_deployment_info = deepcopy(curr_deployment_info) new_deployment_info.set_target_capacity( new_target_capacity=target_capacity, new_target_capacity_direction=target_capacity_direction, ) updating = deployment_state.deploy(new_deployment_info) assert updating @pytest.mark.parametrize( "num_replicas,target_capacity,expected_output", [ (10, None, 10), (10, 100, 10), (10, 99, 10), (10, 50, 5), (10, 1, 1), (10, 0, 0), (10, 25, 3), (1, None, 1), (1, 100, 1), (1, 1, 1), (1, 0, 0), (1, 23, 1), (3, 20, 1), (3, 40, 1), (3, 70, 2), (3, 90, 3), (0, None, 0), (0, 1, 0), (0, 99, 0), (0, 100, 0), ], ) def test_get_capacity_adjusted_num_replicas( self, num_replicas: int, target_capacity: Optional[float], expected_output: int ): result = get_capacity_adjusted_num_replicas(num_replicas, target_capacity) assert isinstance(result, int) assert result == expected_output def test_initial_deploy(self, mock_deployment_state_manager): """Deploy with target_capacity set, should apply immediately.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, _ = deployment_info(num_replicas=2) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] self.update_target_capacity( ds, b_info_1, target_capacity=50, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_target_capacity_100_no_effect(self, mock_deployment_state_manager): """ Deploy with no target_capacity set, then set to 100. Should take no effect. Then go back to no target_capacity, should still have no effect. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() code_version = "arbitrary_version" b_info_1, _ = deployment_info(num_replicas=2, version=code_version) # Initially deploy with no target_capacity set. assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Now update target_capacity to 100, should have no effect. self.update_target_capacity( ds, b_info_1, target_capacity=100, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Now update target_capacity back to None, should have no effect. self.update_target_capacity( ds, b_info_1, target_capacity=None, target_capacity_direction=None, ) dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_target_capacity_0(self, mock_deployment_state_manager): """Deploy with target_capacity set to 0. Should have no replicas.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() b_info_1, _ = deployment_info(num_replicas=100) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] self.update_target_capacity( ds, b_info_1, target_capacity=0, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts(ds, total=0) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_reduce_target_capacity(self, mock_deployment_state_manager): """ Deploy with target capacity set to 100, then reduce to 50, then reduce to 0. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() code_version = "arbitrary_version" b_info_1, _ = deployment_info(num_replicas=10, version=code_version) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Start with target_capacity 100. self.update_target_capacity( ds, b_info_1, target_capacity=100, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.STARTING, 10, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.RUNNING, 10, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Reduce target_capacity to 50, half the replicas should be stopped. self.update_target_capacity( ds, b_info_1, target_capacity=50, target_capacity_direction=TargetCapacityDirection.DOWN, ) dsm.update() check_counts( ds, total=10, by_state=[ (ReplicaState.RUNNING, 5, None), (ReplicaState.STOPPING, 5, None), ], ) assert ds.curr_status_info.status == DeploymentStatus.DOWNSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get([ReplicaState.STOPPING]): replica._actor.set_done_stopping() dsm.update() check_counts(ds, total=5, by_state=[(ReplicaState.RUNNING, 5, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Reduce target_capacity to 1, all but 1 of the replicas should be stopped. self.update_target_capacity( ds, b_info_1, target_capacity=1, target_capacity_direction=TargetCapacityDirection.DOWN, ) dsm.update() check_counts( ds, total=5, by_state=[ (ReplicaState.RUNNING, 1, None), (ReplicaState.STOPPING, 4, None), ], ) assert ds.curr_status_info.status == DeploymentStatus.DOWNSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get([ReplicaState.STOPPING]): replica._actor.set_done_stopping() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) # Reduce target_capacity to 0, all replicas should be stopped. self.update_target_capacity( ds, b_info_1, target_capacity=0, target_capacity_direction=TargetCapacityDirection.DOWN, ) dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STOPPING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.DOWNSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get([ReplicaState.STOPPING]): replica._actor.set_done_stopping() dsm.update() check_counts(ds, total=0) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_increase_target_capacity(self, mock_deployment_state_manager): """ Deploy with target_capacity set to 0, then increase to 1, then increase to 50, then increase to 100. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() code_version = "arbitrary_version" b_info_1, _ = deployment_info(num_replicas=10, version=code_version) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Start with target_capacity set to 0, should have no replicas start up. self.update_target_capacity( ds, b_info_1, target_capacity=0, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts(ds, total=0) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Increase target_capacity to 1, should have 1 replica start up. self.update_target_capacity( ds, b_info_1, target_capacity=1, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Set target_capacity to 50, should have 4 more replicas start up. self.update_target_capacity( ds, b_info_1, target_capacity=50, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts( ds, total=5, by_state=[ (ReplicaState.RUNNING, 1, None), (ReplicaState.STARTING, 4, None), ], ) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=5, by_state=[(ReplicaState.RUNNING, 5, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Set target_capacity to 100, should have 5 more replicas start up. self.update_target_capacity( ds, b_info_1, target_capacity=100, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts( ds, total=10, by_state=[ (ReplicaState.RUNNING, 5, None), (ReplicaState.STARTING, 5, None), ], ) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.RUNNING, 10, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_clear_target_capacity(self, mock_deployment_state_manager): """Deploy with target_capacity set, should apply immediately.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() code_version = "arbitrary_version" b_info_1, _ = deployment_info(num_replicas=10, version=code_version) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Start with target_capacity set to 50, should have 5 replicas start up. self.update_target_capacity( ds, b_info_1, target_capacity=50, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts(ds, total=5, by_state=[(ReplicaState.STARTING, 5, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPDATING for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=5, by_state=[(ReplicaState.RUNNING, 5, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Clear target_capacity, should have 5 more replicas start up. self.update_target_capacity( ds, b_info_1, target_capacity=None, target_capacity_direction=None, ) dsm.update() check_counts( ds, total=10, by_state=[ (ReplicaState.RUNNING, 5, None), (ReplicaState.STARTING, 5, None), ], ) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.RUNNING, 10, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_target_num_replicas_is_zero(self, mock_deployment_state_manager): """ If the target `num_replicas` is zero (i.e., scale-to-zero is enabled and it's autoscaled down), then replicas should remain at zero regardless of target_capacity. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Set num_replicas to 0. code_version = "arbitrary_version" b_info_1, _ = deployment_info(num_replicas=0, version=code_version) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Start with target_capacity of 50. self.update_target_capacity( ds, b_info_1, target_capacity=50, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts(ds, total=0) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=0) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY dsm.update() check_counts(ds, total=0) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Regardless of target_capacity, should stay at 0 replicas. self.update_target_capacity( ds, b_info_1, target_capacity=None, target_capacity_direction=None, ) dsm.update() check_counts(ds, total=0) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY self.update_target_capacity( ds, b_info_1, target_capacity=0, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts(ds, total=0) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY self.update_target_capacity( ds, b_info_1, target_capacity=50, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts(ds, total=0) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY self.update_target_capacity( ds, b_info_1, target_capacity=100, target_capacity_direction=TargetCapacityDirection.UP, ) check_counts(ds, total=0) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Now scale back up to 1 replica. b_info_2, _ = deployment_info(num_replicas=1, version=code_version) self.update_target_capacity( ds, b_info_2, target_capacity=100, target_capacity_direction=TargetCapacityDirection.UP, ) ds._target_state.num_replicas = 1 dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # TODO(edoakes): this test should be updated to go through the autoscaling policy. def test_target_capacity_with_changing_num_replicas( self, mock_deployment_state_manager ): """ Test that target_capacity works with changing num_replicas (emulating autoscaling). """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Set num_replicas to 0. code_version = "arbitrary_version" b_info_1, _ = deployment_info(num_replicas=2, version=code_version) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Start with target_capacity set to 0, should have 0 replica start up # regardless of the autoscaling decision. self.update_target_capacity( ds, b_info_1, target_capacity=0, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts(ds, total=0) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY self.update_target_capacity( ds, b_info_1, target_capacity=1, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING # TODO (shrekris): once this test uses the autoscaling logic, this # status trigger should be DeploymentStatusTrigger.AUTOSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Increase the target number of replicas. Should still only have 1. b_info_2, _ = deployment_info(num_replicas=10, version=code_version) self.update_target_capacity( ds, b_info_2, target_capacity=1, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Increase target_capacity to 50, should have 4 more replicas start up. self.update_target_capacity( ds, b_info_2, target_capacity=50, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts( ds, total=5, by_state=[ (ReplicaState.RUNNING, 1, None), (ReplicaState.STARTING, 4, None), ], ) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING # TODO (shrekris): once this test uses the autoscaling logic, this # status trigger should be DeploymentStatusTrigger.AUTOSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=5, by_state=[(ReplicaState.RUNNING, 5, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Reduce num_replicas and remove target_capacity, should stay the same. b_info_3, _ = deployment_info(num_replicas=5, version=code_version) self.update_target_capacity( ds, b_info_3, target_capacity=None, target_capacity_direction=None, ) dsm.update() check_counts( ds, total=5, by_state=[(ReplicaState.RUNNING, 5, None)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.UPSCALE_COMPLETED ) dsm.update() check_counts(ds, total=5, by_state=[(ReplicaState.RUNNING, 5, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Set target_capacity to 50 and increase num_replicas to 6, should have 2 stop. b_info_4, _ = deployment_info(num_replicas=6, version=code_version) self.update_target_capacity( ds, b_info_4, target_capacity=50, target_capacity_direction=TargetCapacityDirection.UP, ) dsm.update() check_counts( ds, total=5, by_state=[ (ReplicaState.RUNNING, 3, None), (ReplicaState.STOPPING, 2, None), ], ) assert ds.curr_status_info.status == DeploymentStatus.DOWNSCALING # TODO (shrekris): once this test uses the autoscaling logic, this # status trigger should be DeploymentStatusTrigger.AUTOSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get([ReplicaState.STOPPING]): replica._actor.set_done_stopping() dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Unset target capacity, should scale back up to 6. self.update_target_capacity( ds, b_info_4, target_capacity=None, target_capacity_direction=None, ) dsm.update() check_counts( ds, total=6, by_state=[ (ReplicaState.RUNNING, 3, None), (ReplicaState.STARTING, 3, None), ], ) assert ds.curr_status_info.status == DeploymentStatus.UPSCALING # TODO (shrekris): once this test uses the autoscaling logic, this # status trigger should be DeploymentStatusTrigger.AUTOSCALING assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.CONFIG_UPDATE_STARTED ) for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() check_counts(ds, total=6, by_state=[(ReplicaState.RUNNING, 6, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY class TestStopReplicasOnDrainingNodes: """Test the behavior when draining node(s).""" def test_draining_start_then_stop_replica(self, mock_deployment_state_manager): """A new replica should be started before stopping old replica. If the new replica starts quickly, the replica on the draining node should then be gracefully stopped after the new replica transitions to RUNNING. """ create_dsm, timer, cluster_node_info_cache, _ = mock_deployment_state_manager node_1 = NodeID.from_random().hex() node_2 = NodeID.from_random().hex() cluster_node_info_cache.add_node(node_1) cluster_node_info_cache.add_node(node_2) dsm: DeploymentStateManager = create_dsm() timer.reset(0) b_info_1, v1 = deployment_info( num_replicas=2, graceful_shutdown_timeout_s=20, version="1" ) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) # Drain node-2 with deadline 60. Since the replicas are still # starting and we don't know the actor node id yet nothing happens cluster_node_info_cache.draining_nodes = {node_2: 60 * 1000} dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) one_replica, another_replica = ds._replicas.get() one_replica._actor.set_node_id(node_1) one_replica._actor.set_ready() another_replica._actor.set_node_id(node_2) another_replica._actor.set_ready() # Try to start a new replica before initiating the graceful stop # process for the replica on the draining node dsm.update() check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, v1), (ReplicaState.PENDING_MIGRATION, 1, v1), (ReplicaState.STARTING, 1, v1), ], ) # 5 seconds later, the replica hasn't started yet. The replica on # the draining node should not start graceful termination yet. timer.advance(5) dsm.update() check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, v1), (ReplicaState.PENDING_MIGRATION, 1, v1), (ReplicaState.STARTING, 1, v1), ], ) # Simulate it took 5 more seconds for the new replica to be started timer.advance(5) ds._replicas.get([ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 2, v1), (ReplicaState.STOPPING, 1, v1), ], ) # After replica on draining node stops, deployment is healthy with 2 # running replicas. another_replica._actor.set_done_stopping() dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v1)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_draining_stop_replica_before_deadline(self, mock_deployment_state_manager): """If the new replacement replica takes a long time to start, the replica on the draining node should start gracefully terminating ahead of time. The graceful termination should be initiated `graceful_shutdown_timeout_s` seconds before the draining node's deadline, even if the new replica hasn't transitioned to RUNNING yet. """ create_dsm, timer, cluster_node_info_cache, _ = mock_deployment_state_manager node_1 = NodeID.from_random().hex() node_2 = NodeID.from_random().hex() cluster_node_info_cache.add_node(node_1) cluster_node_info_cache.add_node(node_2) dsm: DeploymentStateManager = create_dsm() timer.reset(0) b_info_1, v1 = deployment_info( num_replicas=2, graceful_shutdown_timeout_s=20, version="1" ) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) # Drain node-2 with deadline 60. Since the replicas are still # starting and we don't know the actor node id yet nothing happens cluster_node_info_cache.draining_nodes = {node_2: 60 * 1000} dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) one_replica, another_replica = ds._replicas.get() one_replica._actor.set_node_id(node_1) one_replica._actor.set_ready() another_replica._actor.set_node_id(node_2) another_replica._actor.set_ready() # Try to start a new replica before initiating the graceful stop # process for the replica on the draining node dsm.update() check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, v1), (ReplicaState.PENDING_MIGRATION, 1, v1), (ReplicaState.STARTING, 1, v1), ], ) # Simulate the replica is not yet started after 40 seconds. The # replica on node-2 should start graceful termination even though # a new replica hasn't come up yet. timer.advance(40) dsm.update() check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, v1), (ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v1), ], ) # Mark replica as finished stopping. another_replica._actor.set_done_stopping() dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.STARTING, 1, v1), (ReplicaState.RUNNING, 1, v1)], ) # 5 seconds later, the replica finally starts. timer.advance(5) ds._replicas.get([ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_draining_multiple_nodes(self, mock_deployment_state_manager): """Test multiple nodes draining at the same time. We should choose to stop replicas on nodes with the earliest deadlines when new replicas are started. """ create_dsm, timer, cluster_node_info_cache, _ = mock_deployment_state_manager node_1 = NodeID.from_random().hex() node_2 = NodeID.from_random().hex() node_3 = NodeID.from_random().hex() node_4 = NodeID.from_random().hex() cluster_node_info_cache.add_node(node_1) cluster_node_info_cache.add_node(node_2) cluster_node_info_cache.add_node(node_3) cluster_node_info_cache.add_node(node_4) dsm: DeploymentStateManager = create_dsm() timer.reset(0) b_info_1, v1 = deployment_info( num_replicas=4, graceful_shutdown_timeout_s=20, version="1" ) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=4, by_state=[(ReplicaState.STARTING, 4, v1)]) # Drain node-2 with deadline 60. Since the replicas are still # starting and we don't know the actor node id yet nothing happens cluster_node_info_cache.draining_nodes = { node_2: 60 * 1000, node_3: 100 * 1000, node_4: 40 * 1000, } dsm.update() check_counts(ds, total=4, by_state=[(ReplicaState.STARTING, 4, v1)]) for i, replica in enumerate(ds._replicas.get()): replica._actor.set_node_id([node_1, node_2, node_3, node_4][i]) replica._actor.set_ready() # Try to start new replicas before initiating the graceful stop # process for the replica on the draining node dsm.update() check_counts( ds, total=7, by_state=[ (ReplicaState.RUNNING, 1, v1), (ReplicaState.PENDING_MIGRATION, 3, v1), (ReplicaState.STARTING, 3, v1), ], ) # First new replica transitions to RUNNING. timer.advance(5) ds._replicas.get([ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() check_counts( ds, total=7, by_state=[ (ReplicaState.RUNNING, 2, v1), (ReplicaState.STOPPING, 1, v1), (ReplicaState.PENDING_MIGRATION, 2, v1), (ReplicaState.STARTING, 2, v1), ], ) # The replica on node_4 should be selected for graceful termination, # because node_4 has the earliest deadline. stopping_replica = ds._replicas.get([ReplicaState.STOPPING])[0] assert stopping_replica.actor_node_id == node_4 stopping_replica._actor.set_done_stopping() dsm.update() # Second new replica transitions to RUNNING. timer.advance(5) ds._replicas.get([ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() check_counts( ds, total=6, by_state=[ (ReplicaState.RUNNING, 3, v1), (ReplicaState.STOPPING, 1, v1), (ReplicaState.PENDING_MIGRATION, 1, v1), (ReplicaState.STARTING, 1, v1), ], ) # The replica on node_2 should be selected for graceful termination, # because node_2 has the second earliest deadline. stopping_replica = ds._replicas.get([ReplicaState.STOPPING])[0] assert stopping_replica.actor_node_id == node_2 stopping_replica._actor.set_done_stopping() dsm.update() # Third new replica transitions to RUNNING. timer.advance(5) ds._replicas.get([ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() check_counts( ds, total=5, by_state=[ (ReplicaState.RUNNING, 4, v1), (ReplicaState.STOPPING, 1, v1), ], ) # The replica on node_3 should be selected for graceful termination # last because node_3 has the latest deadline. stopping_replica = ds._replicas.get([ReplicaState.STOPPING])[0] assert stopping_replica.actor_node_id == node_3 stopping_replica._actor.set_done_stopping() dsm.update() # Finally all 4 replicas are running. check_counts(ds, total=4, by_state=[(ReplicaState.RUNNING, 4, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_replicas_unhealthy_on_draining_node(self, mock_deployment_state_manager): """Replicas pending migration should be stopped if unhealthy.""" create_dsm, timer, cluster_node_info_cache, _ = mock_deployment_state_manager node_1 = NodeID.from_random().hex() node_2 = NodeID.from_random().hex() cluster_node_info_cache.add_node(node_1) cluster_node_info_cache.add_node(node_2) dsm: DeploymentStateManager = create_dsm() timer.reset(0) b_info_1, v1 = deployment_info( num_replicas=2, graceful_shutdown_timeout_s=20, version="1" ) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) # Drain node-2 with deadline 60. cluster_node_info_cache.draining_nodes = {node_2: 60 * 1000} dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) one_replica, another_replica = ds._replicas.get() one_replica._actor.set_node_id(node_1) another_replica._actor.set_node_id(node_2) one_replica._actor.set_ready() another_replica._actor.set_ready() # Try to start a new replica before initiating the graceful stop # process for the replica on the draining node dsm.update() check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, v1), (ReplicaState.PENDING_MIGRATION, 1, v1), (ReplicaState.STARTING, 1, v1), ], ) # 5 seconds later, the new replica hasn't started but the # replica on the draining node has become unhealthy. It should # be stopped. timer.advance(5) ds._replicas.get([ReplicaState.PENDING_MIGRATION])[0]._actor.set_unhealthy() dsm.update() check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, v1), (ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v1), ], ) # Unhealthy replica is stopped. ds._replicas.get([ReplicaState.STOPPING])[0]._actor.set_done_stopping() check_counts( ds, total=3, by_state=[(ReplicaState.RUNNING, 1, v1), (ReplicaState.STARTING, 1, v1)], ) # New replica starts. ds._replicas.get([ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v1)]) def test_starting_replica_on_draining_node(self, mock_deployment_state_manager): """When a node gets drained, replicas in STARTING state should be stopped.""" create_dsm, timer, cluster_node_info_cache, _ = mock_deployment_state_manager node_1 = NodeID.from_random().hex() node_2 = NodeID.from_random().hex() cluster_node_info_cache.add_node(node_1) cluster_node_info_cache.add_node(node_2) dsm: DeploymentStateManager = create_dsm() timer.reset(0) b_info_1, v1 = deployment_info( num_replicas=2, graceful_shutdown_timeout_s=20, version="1" ) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) # Mark replica on node_1 as ready, but replica on node_2 is # still starting one_replica, another_replica = ds._replicas.get() one_replica._actor.set_node_id(node_1) another_replica._actor.set_node_id(node_2) one_replica._actor.set_ready() dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.RUNNING, 1, v1), (ReplicaState.STARTING, 1, v1)], ) # Drain node-2. The starting replica should be stopped immediately # without waiting for the replica to start. cluster_node_info_cache.draining_nodes = {node_2: 60 * 1000} dsm.update() check_counts( ds, total=3, by_state=[ (ReplicaState.RUNNING, 1, v1), (ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v1), ], ) stopping_replica = ds._replicas.get([ReplicaState.STOPPING])[0] assert stopping_replica.actor_node_id == node_2 # Finish stopping old replica stopping_replica._actor.set_done_stopping() dsm.update() starting_replica = ds._replicas.get([ReplicaState.STARTING])[0] assert starting_replica.actor_node_id != node_2 # Finish starting new replica starting_replica._actor.set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_in_place_update_during_draining(self, mock_deployment_state_manager): """Test that pending migration replicas of old versions are updated.""" create_dsm, timer, cluster_node_info_cache, _ = mock_deployment_state_manager node_1 = NodeID.from_random().hex() node_2 = NodeID.from_random().hex() cluster_node_info_cache.add_node(node_1) cluster_node_info_cache.add_node(node_2) dsm: DeploymentStateManager = create_dsm() timer.reset(0) b_info_1, v1 = deployment_info( num_replicas=10, graceful_shutdown_timeout_s=20, version="1" ) assert dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.STARTING, 10, v1)]) replicas = ds._replicas.get() replicas[0]._actor.set_node_id(node_2) replicas[0]._actor.set_ready() for r in replicas[1:]: r._actor.set_node_id(node_1) r._actor.set_ready() dsm.update() check_counts(ds, total=10, by_state=[(ReplicaState.RUNNING, 10, v1)]) # Drain node-2 with deadline 60. cluster_node_info_cache.draining_nodes = {node_2: 60 * 1000} dsm.update() check_counts( ds, total=11, by_state=[ (ReplicaState.RUNNING, 9, v1), (ReplicaState.PENDING_MIGRATION, 1, v1), (ReplicaState.STARTING, 1, v1), ], ) # Deploy a new version. The STARTING and PENDING_MIGRATION # replicas of the old version should be stopped. migrating_replica = ds._replicas.get([ReplicaState.PENDING_MIGRATION])[0] b_info_2, v2 = deployment_info( num_replicas=10, graceful_shutdown_timeout_s=20, version="2" ) dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) dsm.update() check_counts( ds, total=12, by_state=[ (ReplicaState.RUNNING, 9, v1), (ReplicaState.STOPPING, 2, v1), (ReplicaState.STARTING, 1, v2), ], ) assert migrating_replica.actor_details.state == ReplicaState.STOPPING # Rolling update should continue ds._replicas.get([ReplicaState.STOPPING])[0]._actor.set_done_stopping() ds._replicas.get([ReplicaState.STOPPING])[1]._actor.set_done_stopping() dsm.update() ds._replicas.get([ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() check_counts( ds, total=12, by_state=[ # Old and new running replicas (ReplicaState.RUNNING, 7, v1), (ReplicaState.RUNNING, 1, v2), # Being rolling updated (ReplicaState.STOPPING, 2, v1), (ReplicaState.STARTING, 2, v2), ], ) def test_docs_path_not_updated_for_different_version(mock_deployment_state_manager): # Create deployment state manager create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info_1, v1 = deployment_info(version="1") dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v1)]) test_docs_path = "/test/docs/path" # Set replicas ready and check statuses for replica in ds._replicas.get([ReplicaState.STARTING]): replica._actor.set_ready() replica._actor.set_docs_path(test_docs_path) assert ds.docs_path is None # status=HEALTHY, status_trigger=DEPLOY dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v1)]) assert ds.docs_path == test_docs_path # Deploy a new version info_2, v2 = deployment_info(version="2") dsm.deploy(TEST_DEPLOYMENT_ID, info_2) dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v2)], ) assert ds.docs_path == test_docs_path test_docs_path_new = "/test/docs/path/2" # Set done stopping replicas ready and check statuses for replica in ds._replicas.get([ReplicaState.STOPPING]): replica._actor.set_done_stopping() # status=HEALTHY, status_trigger=DEPLOY dsm.update() for replica in ds._replicas.get([ReplicaState.STARTING]): replica._actor.set_ready() replica._actor.set_docs_path(test_docs_path_new) assert ds.docs_path == test_docs_path dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v2)]) assert ds.docs_path == test_docs_path_new # Deploy a new version with None docs path info_3, v3 = deployment_info(version="3") dsm.deploy(TEST_DEPLOYMENT_ID, info_3) dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.STOPPING, 1, v2), (ReplicaState.STARTING, 1, v3)], ) assert ds.docs_path == test_docs_path_new for replica in ds._replicas.get([ReplicaState.STOPPING]): replica._actor.set_done_stopping() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v3)]) assert ds.docs_path == test_docs_path_new for replica in ds._replicas.get([ReplicaState.STARTING]): replica._actor.set_ready() replica._actor.set_docs_path(None) dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v3)]) assert ds.docs_path is None def test_set_target_num_replicas_api(mock_deployment_state_manager): """Test the new set_target_num_replicas API for scaling deployments.""" # Create deployment state manager create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Deploy initial deployment with 1 replica info_1, v1 = deployment_info(version="1", num_replicas=1) dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, v1)]) assert ds.target_num_replicas == 1 # Test scaling up using the new API dsm.set_target_num_replicas(TEST_DEPLOYMENT_ID, 3) assert ds.target_num_replicas == 3 dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.STARTING, 3, v1)]) def test_set_target_num_replicas_nonexistent_deployment(mock_deployment_state_manager): """Test that scaling nonexistent deployment raises KeyError.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() nonexistent_id = DeploymentID("nonexistent", "test_app") with pytest.raises(ValueError, match="Deployment.*not found"): dsm.set_target_num_replicas(nonexistent_id, 3) def test_set_target_num_replicas_during_upgrade(mock_deployment_state_manager): """Test setting target replicas while an upgrade is ongoing.""" # Create deployment state manager create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Deploy initial deployment (v1) with 2 replicas info_1, v1 = deployment_info(version="1", num_replicas=2) dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) assert ds.target_num_replicas == 2 # Get replicas to RUNNING state for replica in ds._replicas.get([ReplicaState.STARTING]): replica._actor.set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v1)]) # Start an upgrade to v2 with 2 replicas info_2, v2 = deployment_info(version="2", num_replicas=2) dsm.deploy(TEST_DEPLOYMENT_ID, info_2) dsm.update() check_counts( ds, total=3, by_state=[ (ReplicaState.STARTING, 1, v2), (ReplicaState.RUNNING, 1, v1), (ReplicaState.STOPPING, 1, v1), ], ) assert ds.target_num_replicas == 2 # Scale up to 5 replicas in the middle of the upgrade. dsm.set_target_num_replicas(TEST_DEPLOYMENT_ID, 5) assert ds.target_num_replicas == 5 def dsm_update(): for replica in ds._replicas.get([ReplicaState.STOPPING]): replica._actor.set_done_stopping() for replica in ds._replicas.get([ReplicaState.STARTING]): replica._actor.set_ready() dsm.update() dsm_update() check_counts( ds, total=5, by_state=[ (ReplicaState.STARTING, 3, v2), (ReplicaState.RUNNING, 1, v1), (ReplicaState.RUNNING, 1, v2), ], ) dsm_update() check_counts( ds, total=6, by_state=[ (ReplicaState.STARTING, 1, v2), (ReplicaState.RUNNING, 4, v2), (ReplicaState.STOPPING, 1, v1), ], ) dsm_update() check_counts(ds, total=5, by_state=[(ReplicaState.RUNNING, 5, v2)]) assert ds.target_num_replicas == 5 def test_set_target_num_replicas_deleting_deployment(mock_deployment_state_manager): """Test scaling deployment that is being deleted.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Deploy an initial deployment info, v1 = deployment_info(num_replicas=2, version="v1") dsm.deploy(TEST_DEPLOYMENT_ID, info) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, v1)]) # Delete the deployment dsm.delete_deployment(TEST_DEPLOYMENT_ID) # The deployment status should be DELETING statuses = dsm.get_deployment_statuses([TEST_DEPLOYMENT_ID]) assert statuses[0].status_trigger == DeploymentStatusTrigger.DELETING # Scaling should fail with pytest.raises(DeploymentIsBeingDeletedError): dsm.set_target_num_replicas(TEST_DEPLOYMENT_ID, 3) class TestDeploymentRankManagerIntegrationE2E: """End-to-end integration tests for rank functionality through deployment state manager.""" def _set_replicas_ready( self, ds: DeploymentState, replica_states: List[ReplicaState] ): """Helper to set replicas in given states to ready.""" for replica in ds._replicas.get(replica_states): replica._actor.set_ready() def _set_replicas_done_stopping(self, ds: DeploymentState): """Helper to set stopping replicas as done stopping.""" for replica in ds._replicas.get([ReplicaState.STOPPING]): replica._actor.set_done_stopping() def test_scaling_up_and_down_scenario(self, mock_deployment_state_manager): """Test a realistic scaling scenario through deployment state manager.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Start with 3 replicas info_1, v1 = deployment_info(num_replicas=3, version="1") dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Create initial replicas dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.STARTING, 3, v1)]) # Set replicas ready self._set_replicas_ready(ds, [ReplicaState.STARTING]) dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Check initial ranks are 0, 1, 2 ranks_mapping = ds._get_replica_ranks_mapping() ranks = sorted([r.rank for r in ranks_mapping.values()]) assert ranks == [0, 1, 2], f"Expected ranks [0, 1, 2], got {ranks}" # Scale down to 2 replicas - this should trigger rank reassignment info_2, _ = deployment_info(num_replicas=2, version="1") dsm.deploy(TEST_DEPLOYMENT_ID, info_2) dsm.update() # One replica should be stopping check_counts( ds, total=3, by_state=[(ReplicaState.RUNNING, 2, v1), (ReplicaState.STOPPING, 1, v1)], ) # Complete the scale down self._set_replicas_done_stopping(ds) dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Trigger rank consistency check with one more update dsm.update() # After scaling down and reaching healthy status, ranks should be contiguous [0, 1] ranks_mapping = ds._get_replica_ranks_mapping() ranks = sorted([r.rank for r in ranks_mapping.values()]) assert ranks == [0, 1], f"Expected ranks [0, 1] after scale down, got {ranks}" # Scale back up to 3 replicas - new replica should reuse available rank info_3, _ = deployment_info(num_replicas=3, version="1") dsm.deploy(TEST_DEPLOYMENT_ID, info_3) dsm.update() # Should have one new starting replica check_counts( ds, total=3, by_state=[(ReplicaState.RUNNING, 2, v1), (ReplicaState.STARTING, 1, v1)], ) # Set new replica ready self._set_replicas_ready(ds, [ReplicaState.STARTING]) dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Trigger rank consistency check with one more update dsm.update() # Final ranks should be contiguous [0, 1, 2] ranks_mapping = ds._get_replica_ranks_mapping() ranks = sorted([r.rank for r in ranks_mapping.values()]) assert ranks == [0, 1, 2], f"Expected final ranks [0, 1, 2], got {ranks}" def test_controller_recovery_with_scattered_ranks( self, mock_deployment_state_manager ): """Test controller recovery with existing replica ranks through deployment state manager.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Deploy with 3 replicas info_1, v1 = deployment_info(num_replicas=3, version="1") target_state_changed = dsm.deploy(TEST_DEPLOYMENT_ID, info_1) assert target_state_changed dsm.save_checkpoint() ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Create replicas and get them running dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.STARTING, 3, v1)]) self._set_replicas_ready(ds, [ReplicaState.STARTING]) dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v1)]) # Get the actual replica objects (not just IDs) replicas = ds._replicas.get([ReplicaState.RUNNING]) replica_ids = [replica.replica_id for replica in replicas] # Simulate controller crashed! Create a new deployment state manager # with the existing replica IDs to trigger recovery new_dsm: DeploymentStateManager = create_dsm( [replica_id.to_full_id_str() for replica_id in replica_ids] ) # New deployment state should be created and replicas should be RECOVERING new_ds = new_dsm._deployment_states[TEST_DEPLOYMENT_ID] check_counts(new_ds, total=3, by_state=[(ReplicaState.RECOVERING, 3, v1)]) # Complete recovery - set replicas ready self._set_replicas_ready(new_ds, [ReplicaState.RECOVERING]) new_dsm.update() check_counts(new_ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v1)]) assert new_ds.curr_status_info.status == DeploymentStatus.HEALTHY # At this point ranks should be scattered but all values [0, 1, 2] should be present ranks_mapping = new_ds._get_replica_ranks_mapping() ranks = sorted([r.rank for r in ranks_mapping.values()]) assert ranks == [0, 1, 2], "Should have recovered scattered ranks" # Trigger rank consistency check with one more update - this should reorder if needed new_dsm.update() # After rank consistency check, ranks should still be [0, 1, 2] final_ranks_mapping = new_ds._get_replica_ranks_mapping() final_ranks = sorted([r.rank for r in final_ranks_mapping.values()]) assert final_ranks == [ 0, 1, 2, ], f"Expected contiguous ranks [0, 1, 2] after consistency check, got {final_ranks}" # Clean up replica_rank_context.clear() def test_rank_recovery_after_lightweight_reconfigure( self, mock_deployment_state_manager ): """Lightweight reconfigure must pass the ReplicaRank object, not the int. Regression test for the lightweight-update branch calling `replica.reconfigure(version, rank=current_rank.rank)` (a bare int). The replica stored the int and reported it back on controller recovery, where `_recover_rank_impl` failed with `'int' object has no attribute 'rank'`, leaving every surviving replica without a rank and `_check_rank_consistency_impl` looping "Rank system is in an invalid state". """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() replica_rank_context.clear() # Deploy 3 replicas with a user_config and get them running. info_1, v1 = deployment_info(num_replicas=3, version="1", user_config="1") assert dsm.deploy(TEST_DEPLOYMENT_ID, info_1) dsm.save_checkpoint() ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.STARTING, 3, v1)]) self._set_replicas_ready(ds, [ReplicaState.STARTING]) dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v1)]) # Lightweight update: same code version, new user_config. Replicas # are reconfigured in place (rolling), not restarted. info_2, v2 = deployment_info(num_replicas=3, version="1", user_config="2") assert dsm.deploy(TEST_DEPLOYMENT_ID, info_2) dsm.save_checkpoint() for _ in range(10): dsm.update() self._set_replicas_ready(ds, [ReplicaState.UPDATING]) check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v2)]) # The reconfigure must have handed each replica the full ReplicaRank # object; this is what replicas report back on controller recovery. replicas = ds._replicas.get([ReplicaState.RUNNING]) for replica in replicas: stored_rank = replica_rank_context[replica.replica_id.unique_id] assert isinstance( stored_rank, ReplicaRank ), f"Replica stored {stored_rank!r} instead of a ReplicaRank" # Simulate controller crash and recover from the live replicas. replica_ids = [replica.replica_id for replica in replicas] new_dsm: DeploymentStateManager = create_dsm( [replica_id.to_full_id_str() for replica_id in replica_ids] ) new_ds = new_dsm._deployment_states[TEST_DEPLOYMENT_ID] check_counts(new_ds, total=3, by_state=[(ReplicaState.RECOVERING, 3, v2)]) self._set_replicas_ready(new_ds, [ReplicaState.RECOVERING]) new_dsm.update() check_counts(new_ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v2)]) # Every recovered replica must have its rank restored. for replica_id in replica_ids: assert new_ds._rank_manager.has_replica_rank( replica_id.unique_id ), f"Rank for {replica_id.unique_id} was not recovered" ranks_mapping = new_ds._get_replica_ranks_mapping() ranks = sorted([r.rank for r in ranks_mapping.values()]) assert ranks == [0, 1, 2], f"Expected recovered ranks [0, 1, 2], got {ranks}" # Clean up replica_rank_context.clear() def test_complex_reassignment_scenario(self, mock_deployment_state_manager): """Test complex reassignment with many gaps through deployment state manager.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Deploy with 4 replicas info_1, v1 = deployment_info(num_replicas=4, version="1") target_state_changed = dsm.deploy(TEST_DEPLOYMENT_ID, info_1) assert target_state_changed dsm.save_checkpoint() ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Create replicas and get them running dsm.update() check_counts(ds, total=4, by_state=[(ReplicaState.STARTING, 4, v1)]) self._set_replicas_ready(ds, [ReplicaState.STARTING]) dsm.update() check_counts(ds, total=4, by_state=[(ReplicaState.RUNNING, 4, v1)]) # Get the actual replica objects replicas = ds._replicas.get([ReplicaState.RUNNING]) replica_ids = [replica.replica_id for replica in replicas] # Simulate very scattered ranks in global context: 0, 3, 7, 10 global replica_rank_context replica_rank_context.clear() replica_rank_context[replica_ids[0].unique_id] = ReplicaRank( rank=0, node_rank=0, local_rank=0 ) replica_rank_context[replica_ids[1].unique_id] = ReplicaRank( rank=3, node_rank=0, local_rank=1 ) replica_rank_context[replica_ids[2].unique_id] = ReplicaRank( rank=7, node_rank=0, local_rank=2 ) replica_rank_context[replica_ids[3].unique_id] = ReplicaRank( rank=10, node_rank=0, local_rank=3 ) # Simulate controller crashed! Create a new deployment state manager # with the existing replica IDs to trigger recovery new_dsm: DeploymentStateManager = create_dsm( [replica_id.to_full_id_str() for replica_id in replica_ids] ) # New deployment state should be created and replicas should be RECOVERING new_ds = new_dsm._deployment_states[TEST_DEPLOYMENT_ID] check_counts(new_ds, total=4, by_state=[(ReplicaState.RECOVERING, 4, v1)]) # Complete recovery - set replicas ready self._set_replicas_ready(new_ds, [ReplicaState.RECOVERING]) new_dsm.update() check_counts(new_ds, total=4, by_state=[(ReplicaState.RUNNING, 4, v1)]) assert new_ds.curr_status_info.status == DeploymentStatus.HEALTHY # Trigger rank consistency check with one more update new_dsm.update() # After reassignment, ranks should be contiguous [0, 1, 2, 3] ranks_mapping = new_ds._get_replica_ranks_mapping() ranks = sorted([r.rank for r in ranks_mapping.values()]) assert ranks == [ 0, 1, 2, 3, ], f"Expected reassigned ranks [0, 1, 2, 3], got {ranks}" def test_rank_consistency_during_version_rollout( self, mock_deployment_state_manager ): """Test that rank consistency is maintained during version rollouts.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Start with 3 replicas of version 1 info_1, v1 = deployment_info(num_replicas=3, version="1") dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Create and ready initial replicas dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.STARTING, 3, v1)]) self._set_replicas_ready(ds, [ReplicaState.STARTING]) dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Verify initial ranks are contiguous ranks_mapping = ds._get_replica_ranks_mapping() initial_ranks = sorted([r.rank for r in ranks_mapping.values()]) assert initial_ranks == [0, 1, 2] # Deploy version 2 - this should trigger rolling update info_2, v2 = deployment_info(num_replicas=3, version="2") dsm.deploy(TEST_DEPLOYMENT_ID, info_2) dsm.update() # Complete the rolling update step by step while True: # Set any new starting replicas ready starting_replicas = ds._replicas.get([ReplicaState.STARTING]) if starting_replicas: self._set_replicas_ready(ds, [ReplicaState.STARTING]) # Complete any stopping replicas stopping_replicas = ds._replicas.get([ReplicaState.STOPPING]) if stopping_replicas: self._set_replicas_done_stopping(ds) dsm.update() # Check if rolling update is complete running_replicas = ds._replicas.get([ReplicaState.RUNNING]) if len(running_replicas) == 3 and all( r.version == v2 for r in running_replicas ): break # After rolling update is complete, deployment should be healthy assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Trigger rank consistency check with one more update dsm.update() # After rolling update, verify ranks are still contiguous final_ranks_mapping = ds._get_replica_ranks_mapping() final_ranks = sorted([r.rank for r in final_ranks_mapping.values()]) assert final_ranks == [ 0, 1, 2, ], f"Expected contiguous ranks [0, 1, 2] after rollout, got {final_ranks}" def test_rank_assignment_with_replica_failures(self, mock_deployment_state_manager): """Test rank handling when replicas fail during startup.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Deploy with 3 replicas info_1, v1 = deployment_info(num_replicas=3, version="1") dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Create initial replicas dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.STARTING, 3, v1)]) # Make first two replicas ready, but let the third fail starting_replicas = ds._replicas.get([ReplicaState.STARTING]) starting_replicas[0]._actor.set_ready() starting_replicas[1]._actor.set_ready() starting_replicas[2]._actor.set_failed_to_start() dsm.update() running_count = ds._replicas.count(states=[ReplicaState.RUNNING]) stopping_count = ds._replicas.count(states=[ReplicaState.STOPPING]) assert running_count == 2, "Should have 2 running replicas" assert stopping_count == 1, "Should have 1 stopping replica" self._set_replicas_done_stopping(ds) dsm.update() starting_count = ds._replicas.count(states=[ReplicaState.STARTING]) assert starting_count == 1, "Should have 1 starting replica" self._set_replicas_ready(ds, [ReplicaState.STARTING]) dsm.update() # second update to reassign ranks dsm.update() # Final verification - should have 3 running replicas (ignore failed/stopping replicas) running_replicas = ds._replicas.get([ReplicaState.RUNNING]) assert ( len(running_replicas) == 3 ), f"Expected 3 running replicas, got {len(running_replicas)}" # Verify that ranks are properly assigned and unique for running replicas ranks_mapping = ds._get_replica_ranks_mapping() # Filter ranks to only include those for running replicas running_replica_ids = [ replica.replica_id.unique_id for replica in running_replicas ] running_replica_ranks = [ ranks_mapping[replica_id].rank for replica_id in running_replica_ids if replica_id in ranks_mapping ] # The ranks should be assigned to all running replicas assert set(running_replica_ranks) == { 0, 1, 2, }, f"Expected ranks [0, 1, 2], got {[r.rank for r in ranks_mapping.values()]}" def test_rank_recovery_skips_when_already_assigned( self, mock_deployment_state_manager ): """Verify that recover_rank is skipped when a replica's rank is already assigned.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Deploy 3 replicas: STARTING -> RUNNING (ranks get assigned). info_1, v1 = deployment_info(num_replicas=3, version="1") dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds: DeploymentState = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.STARTING, 3, v1)]) self._set_replicas_ready(ds, [ReplicaState.STARTING]) dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v1)]) # Record the actor names and ranks for recovery. actor_names = [r.replica_id.to_full_id_str() for r in ds._replicas.get()] original_ranks = { r.replica_id.unique_id: ds._rank_manager.get_replica_rank( r.replica_id.unique_id ) for r in ds._replicas.get() } dsm.save_checkpoint() # Simulate controller crash: create a new DSM with the live actor names. new_dsm: DeploymentStateManager = create_dsm(actor_names) new_ds = new_dsm._deployment_states[TEST_DEPLOYMENT_ID] check_counts(new_ds, total=3, by_state=[(ReplicaState.RECOVERING, 3, v1)]) # Enable strict rank error mode so duplicate recover_rank raises. new_ds._rank_manager._fail_on_rank_error = True # Pre-populate 1 replica's rank in the new rank manager, simulating # the scenario where the rank was never released. target = new_ds._replicas.get(states=[ReplicaState.RECOVERING])[0] target_id = target.replica_id.unique_id target_rank = original_ranks[target_id] new_ds._rank_manager.recover_rank(target_id, target.actor_node_id, target_rank) assert new_ds._rank_manager.has_replica_rank(target_id) # Mark all recovering replicas as ready. self._set_replicas_ready(new_ds, [ReplicaState.RECOVERING]) # This update() calls _check_startup_replicas(RECOVERING). For the # pre-populated replica, has_replica_rank returns True so recover_rank # is SKIPPED. new_dsm.update() check_counts(new_ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v1)]) # Verify all ranks were recovered correctly. for rid, expected_rank in original_ranks.items(): assert new_ds._rank_manager.get_replica_rank(rid) == expected_rank class TestGetOutboundDeployments: def test_basic_outbound_deployments(self, mock_deployment_state_manager): """Test that outbound deployments are returned.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() deployment_id = DeploymentID(name="test_deployment", app_name="test_app") b_info_1, _ = deployment_info(num_replicas=1) dsm.deploy(deployment_id, b_info_1) # Create a RUNNING replica ds = dsm._deployment_states[deployment_id] dsm.update() # Transitions to STARTING for replica in ds._replicas.get([ReplicaState.STARTING]): replica._actor.set_ready() dsm.update() # Transitions to RUNNING # Set outbound deployments on the mock replica running_replicas = ds._replicas.get([ReplicaState.RUNNING]) assert len(running_replicas) == 1 d1 = DeploymentID(name="dep1", app_name="test_app") d2 = DeploymentID(name="dep2", app_name="test_app") running_replicas[0]._actor._outbound_deployments = [d1, d2] outbound_deployments = ds.get_outbound_deployments() assert outbound_deployments == [d1, d2] # Verify it's accessible through DeploymentStateManager assert dsm.get_deployment_outbound_deployments(deployment_id) == [ d1, d2, ] def test_deployment_state_manager_returns_none_for_nonexistent_deployment( self, mock_deployment_state_manager ): """Test that DeploymentStateManager returns None for nonexistent deployments.""" ( create_dsm, timer, cluster_node_info_cache, autoscaling_state_manager, ) = mock_deployment_state_manager dsm = create_dsm() deployment_id = DeploymentID(name="nonexistent", app_name="test_app") assert dsm.get_deployment_outbound_deployments(deployment_id) is None def test_returns_none_if_replicas_are_not_running( self, mock_deployment_state_manager ): """Test that DeploymentStateManager returns None if replicas are not running.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() deployment_id = DeploymentID(name="test_deployment", app_name="test_app") b_info_1, _ = deployment_info(num_replicas=2) dsm.deploy(deployment_id, b_info_1) ds = dsm._deployment_states[deployment_id] dsm.update() replicas = ds._replicas.get([ReplicaState.STARTING]) assert len(replicas) == 2 d1 = DeploymentID(name="dep1", app_name="test_app") d2 = DeploymentID(name="dep2", app_name="test_app") d3 = DeploymentID(name="dep3", app_name="test_app") d4 = DeploymentID(name="dep4", app_name="test_app") replicas[0]._actor._outbound_deployments = [d1, d2] replicas[1]._actor._outbound_deployments = [d3, d4] dsm.update() outbound_deployments = ds.get_outbound_deployments() assert outbound_deployments is None # Set replicas ready replicas[0]._actor.set_ready() dsm.update() outbound_deployments = ds.get_outbound_deployments() assert outbound_deployments == [d1, d2] def test_only_considers_replicas_matching_target_version( self, mock_deployment_state_manager ): """Test that only replicas with target version are considered. When a new version is deployed, old version replicas that are still running should not be included in the outbound deployments result. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Deploy version 1 b_info_1, v1 = deployment_info(version="1") dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() # Get v1 replica to RUNNING state ds._replicas.get()[0]._actor.set_ready() dsm.update() # Set outbound deployments for v1 replica d1 = DeploymentID(name="dep1", app_name="test_app") d2 = DeploymentID(name="dep2", app_name="test_app") ds._replicas.get()[0]._actor._outbound_deployments = [d1, d2] # Verify v1 outbound deployments are returned assert ds.get_outbound_deployments() == [d1, d2] # Deploy version 2 - this triggers rolling update b_info_2, v2 = deployment_info(version="2") dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) dsm.update() # Now we have v1 stopping and v2 starting check_counts( ds, total=2, by_state=[(ReplicaState.STOPPING, 1, v1), (ReplicaState.STARTING, 1, v2)], ) # Key test: Even though v1 replica exists (stopping), it should not be # included because target version is v2. Since v2 is not RUNNING yet, # should return None. assert ds.get_outbound_deployments() is None # Set outbound deployments for v2 replica and mark it ready d3 = DeploymentID(name="dep3", app_name="test_app") ds._replicas.get(states=[ReplicaState.STARTING])[ 0 ]._actor._outbound_deployments = [d3] ds._replicas.get(states=[ReplicaState.STARTING])[0]._actor.set_ready() dsm.update() # Now v2 is running. Should only return v2's outbound deployments (d3), # not v1's outbound deployments (d1, d2). assert ds.get_outbound_deployments() == [d3] def test_broadcast_skips_work_when_replicas_unchanged(mock_deployment_state_manager): """Test that broadcast_running_replicas_if_changed() skips all work in steady state when _broadcasted_replicas_set_changed is False and _request_routing_info_updated is False.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info_1, v1 = deployment_info() dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # _broadcasted_replicas_set_changed should be True after deploy. assert ds._broadcasted_replicas_set_changed is True # Bring deployment to healthy state. dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # After update(), the broadcast should have cleared the flag. assert ds._broadcasted_replicas_set_changed is False assert ds._request_routing_info_updated is False # Call broadcast again — should be a no-op (early return). # We verify by patching get_running_replica_infos; if the fast path # works, the patched method will NOT be called. with patch.object(ds, "get_running_replica_infos") as mock_get_infos: ds.broadcast_running_replicas_if_changed() mock_get_infos.assert_not_called() # Now stop a replica (sets _broadcasted_replicas_set_changed = True). ds._stop_one_running_replica_for_testing() assert ds._broadcasted_replicas_set_changed is True # broadcast should now do the full check. ds.broadcast_running_replicas_if_changed() # Flag should be cleared after the broadcast. assert ds._broadcasted_replicas_set_changed is False def test_broadcast_runs_when_routing_info_updated(mock_deployment_state_manager): """Test that broadcast_running_replicas_if_changed() runs when _request_routing_info_updated is True even if _broadcasted_replicas_set_changed is False.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info_1, v1 = deployment_info() dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Bring deployment to healthy state. dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds._broadcasted_replicas_set_changed is False # Simulate routing info update. ds._request_routing_info_updated = True # broadcast should NOT take the fast path. with patch.object( ds, "get_running_replica_infos", wraps=ds.get_running_replica_infos ) as mock_get_infos: ds.broadcast_running_replicas_if_changed() mock_get_infos.assert_called_once() # Both flags should be cleared after broadcast. assert ds._broadcasted_replicas_set_changed is False assert ds._request_routing_info_updated is False def test_broadcasted_replicas_set_changed_flag_set_on_state_transitions( mock_deployment_state_manager, ): """Test that _broadcasted_replicas_set_changed is set correctly during replica state transitions.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info_1, v1 = deployment_info(num_replicas=2) dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Flag should be True after deploy (_set_target_state sets it). assert ds._broadcasted_replicas_set_changed is True # After update, replicas are STARTING. dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.STARTING, 2, None)]) # broadcast clears the flag. assert ds._broadcasted_replicas_set_changed is False # Set replicas ready — this transitions them to RUNNING in _check_startup_replicas. for r in ds._replicas.get(): r._actor.set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, None)]) # Flag should be cleared again after broadcast. assert ds._broadcasted_replicas_set_changed is False # Now fail a health check (sets _broadcasted_replicas_set_changed via _stop_replica). ds._replicas.get()[0]._actor.set_unhealthy() dsm.update() # Flag set by _stop_replica then cleared by broadcast. assert ds._broadcasted_replicas_set_changed is False # Verify a replica was stopped. assert ds._replicas.count(states=[ReplicaState.STOPPING]) >= 1 def test_broadcasted_replicas_set_changed_flag_set_on_lightweight_broadcast_config_update( mock_deployment_state_manager, ): """Regression test: when a config change requires a long-poll broadcast (e.g. max_ongoing_requests changed) but does NOT require an actor restart or reconfigure, the _broadcasted_replicas_set_changed flag must still be set by the requires_long_poll_broadcast path so the broadcast is not skipped. This guards against a future scenario where a broadcast-affecting field is changed to a lighter update type that doesn't trigger actor_updating. """ create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Deploy v1 and bring to healthy steady state. b_info_1, v1 = deployment_info(version="1") dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ds._broadcasted_replicas_set_changed is False # Deploy v2 with a different max_ongoing_requests. b_info_2, v2 = deployment_info(version="1", max_ongoing_requests=42) dsm.deploy(TEST_DEPLOYMENT_ID, b_info_2) # _set_target_state also sets _broadcasted_replicas_set_changed. Clear it so we can # isolate whether _stop_or_update_outdated_version_replicas sets the # flag via the requires_long_poll_broadcast path. ds._broadcasted_replicas_set_changed = False # Patch the running replica's mock actor so reconfigure() returns False # (simulating a broadcast-needed but no-actor-update scenario). replica = ds._replicas.get()[0] original_reconfigure = replica._actor.reconfigure def patched_reconfigure(version, rank=None): # Perform the version/rank bookkeeping but report no actor update. original_reconfigure(version, rank=rank) return False replica._actor.reconfigure = patched_reconfigure # Confirm preconditions: the version change requires a broadcast. assert v1.requires_long_poll_broadcast(v2) # Directly call _stop_or_update_outdated_version_replicas (the method # that checks requires_long_poll_broadcast) so we can inspect the flag # before broadcast clears it. ds._stop_or_update_outdated_version_replicas() # The replica should stay RUNNING (no actor restart, no UPDATING state) # because our patched reconfigure() returns False. check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, v2)]) # Key assertion: _broadcasted_replicas_set_changed was set by the # requires_long_poll_broadcast path, NOT by actor_updating. assert ds._broadcasted_replicas_set_changed is True # Now broadcast and verify it fires (clearing the flag). ds.broadcast_running_replicas_if_changed() assert ds._broadcasted_replicas_set_changed is False ds._long_poll_host.notify_changed.assert_called() # Verify the fast path works: no further broadcast on next tick. ds._long_poll_host.notify_changed.reset_mock() with patch.object(ds, "get_running_replica_infos") as mock_get_infos: ds.broadcast_running_replicas_if_changed() mock_get_infos.assert_not_called() def test_broadcast_deferred_while_replicas_recovering(mock_deployment_state_manager): """Regression test: During controller recovery, broadcast_running_replicas_if_changed() must be deferred until all RECOVERING replicas have transitioned, then fire once with the complete set. More Context: https://github.com/ray-project/ray/issues/62728 """ def targets_key_was_broadcast(mock_lph): """Return True if notify_changed was called with a DEPLOYMENT_TARGETS key.""" for call in mock_lph.notify_changed.call_args_list: keys_dict = call[0][0] if any( isinstance(k, tuple) and k[0] == LongPollNamespace.DEPLOYMENT_TARGETS for k in keys_dict ): return True return False create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() # Deploy 3 replicas and bring to steady state. info_1, v1 = deployment_info(num_replicas=3, version="1") dsm.deploy(TEST_DEPLOYMENT_ID, info_1) dsm.save_checkpoint() ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() for r in ds._replicas.get(): r._actor.set_ready() dsm.update() check_counts(ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Simulate controller restart replica_ids = [r.replica_id.to_full_id_str() for r in ds._replicas.get()] new_dsm: DeploymentStateManager = create_dsm(actor_names=replica_ids) new_ds = new_dsm._deployment_states[TEST_DEPLOYMENT_ID] # All 3 replicas should be RECOVERING. check_counts(new_ds, total=3, by_state=[(ReplicaState.RECOVERING, 3, v1)]) # While replicas are RECOVERING, broadcast must not fire, even though the # _broadcasted_replicas_set_changed=True. assert new_ds._broadcasted_replicas_set_changed is True new_ds._long_poll_host.notify_changed.reset_mock() new_ds.broadcast_running_replicas_if_changed() assert not targets_key_was_broadcast(new_ds._long_poll_host) # Partially complete recovery: 2/3 replicas RUNNING. # Broadcast still suppressed because 1 replica is RECOVERING. new_ds._long_poll_host.notify_changed.reset_mock() recovering = new_ds._replicas.get(states=[ReplicaState.RECOVERING]) for r in recovering[:2]: r._actor.set_ready() new_dsm.update() check_counts( new_ds, total=3, by_state=[(ReplicaState.RUNNING, 2, v1), (ReplicaState.RECOVERING, 1, v1)], ) assert not targets_key_was_broadcast(new_ds._long_poll_host) # All replicas finish recovery, so broadcast fires new_ds._long_poll_host.notify_changed.reset_mock() for r in new_ds._replicas.get(states=[ReplicaState.RECOVERING]): r._actor.set_ready() new_dsm.update() check_counts(new_ds, total=3, by_state=[(ReplicaState.RUNNING, 3, v1)]) assert targets_key_was_broadcast(new_ds._long_poll_host) # The payload must contain all 3 running replicas. for call in new_ds._long_poll_host.notify_changed.call_args_list: keys_dict = call[0][0] key = (LongPollNamespace.DEPLOYMENT_TARGETS, TEST_DEPLOYMENT_ID) if key in keys_dict: assert len(keys_dict[key].running_replicas) == 3 break def test_in_transition_cleared_at_steady_state(mock_deployment_state_manager): """Test that _in_transition is cleared once a deployment reaches HEALTHY steady state and that subsequent ticks skip expensive work.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info_1, v1 = deployment_info() dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Flag must be True after deploy. assert ds._in_transition is True # STARTING phase: flag stays True. dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.STARTING, 1, None)]) assert ds._in_transition is True # Replica becomes ready → transitions to RUNNING. ds._replicas.get()[0]._actor.set_ready() dsm.update() check_counts(ds, total=1, by_state=[(ReplicaState.RUNNING, 1, None)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Steady state reached: flag must now be False. assert ds._in_transition is False def test_in_transition_skips_expensive_methods(mock_deployment_state_manager): """When _in_transition is False, check_curr_status and scale_deployment_replicas should be no-ops.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info_1, v1 = deployment_info() dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Bring to steady state. dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds._in_transition is False # check_curr_status should return fast (False, False). assert ds.check_curr_status() == (False, False) # scale_deployment_replicas should return fast ([], None). assert ds.scale_deployment_replicas() == ([], None) # Verify check_and_update_replicas runs health checks but skips # startup/stopping by patching _check_startup_replicas. with patch.object(ds, "_check_startup_replicas") as mock_startup: ds.check_and_update_replicas() mock_startup.assert_not_called() def test_in_transition_set_on_health_check_failure( mock_deployment_state_manager, ): """A health check failure during steady state must re-enable reconciliation so the controller can recover.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info_1, v1 = deployment_info(num_replicas=2) dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Bring to steady state. dsm.update() for r in ds._replicas.get(): r._actor.set_ready() dsm.update() assert ds._in_transition is False # Fail a health check. ds._replicas.get()[0]._actor.set_unhealthy() dsm.update() # Flag must be set (via _stop_replica) and there should be a # STOPPING replica being processed. assert ds._replicas.count(states=[ReplicaState.STOPPING]) >= 1 # The flag may already be cleared again if check_curr_status ran # in the same tick and found reconciliation needed, so just verify # the system is still functioning — the deployment should not be # stuck. assert ds._in_transition is True def test_in_transition_set_on_target_state_change( mock_deployment_state_manager, ): """Changing the target state (redeploy, autoscale) must set the flag.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info_1, v1 = deployment_info(version="1") dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Bring to steady state. dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds._in_transition is False # Redeploy with a new version. info_2, v2 = deployment_info(version="2") dsm.deploy(TEST_DEPLOYMENT_ID, info_2) assert ds._in_transition is True def test_routing_stats_change_triggers_broadcast(mock_deployment_state_manager): """Routing stats changes during health checks must set _broadcasted_replicas_set_changed so that the broadcast fast path does not skip the update.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() info_1, v1 = deployment_info() dsm.deploy(TEST_DEPLOYMENT_ID, info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Bring to steady state. dsm.update() ds._replicas.get()[0]._actor.set_ready() dsm.update() assert ds._broadcasted_replicas_set_changed is False assert ds._request_routing_info_updated is False # Simulate routing stats changing on the replica actor. ds._replicas.get()[0]._actor.get_routing_stats = lambda: {"new_key": 42} # Run health checks (STEP 1 of update loop). ds.check_and_update_replicas() # The flag must be set because routing_stats changed. assert ds._broadcasted_replicas_set_changed is True # Broadcast should now run (not take the fast path). with patch.object( ds, "get_running_replica_infos", wraps=ds.get_running_replica_infos ) as mock_get_infos: ds.broadcast_running_replicas_if_changed() mock_get_infos.assert_called_once() assert ds._broadcasted_replicas_set_changed is False def test_pending_migration_prevents_in_transition_clear( mock_deployment_state_manager, ): create_dsm, timer, cluster_node_info_cache, _ = mock_deployment_state_manager node_1 = NodeID.from_random().hex() node_2 = NodeID.from_random().hex() cluster_node_info_cache.add_node(node_1) cluster_node_info_cache.add_node(node_2) dsm: DeploymentStateManager = create_dsm() timer.reset(0) b_info_1, v1 = deployment_info( num_replicas=2, graceful_shutdown_timeout_s=20, version="1" ) dsm.deploy(TEST_DEPLOYMENT_ID, b_info_1) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Start replicas on different nodes. dsm.update() one_replica, another_replica = ds._replicas.get() one_replica._actor.set_node_id(node_1) one_replica._actor.set_ready() another_replica._actor.set_node_id(node_2) another_replica._actor.set_ready() dsm.update() check_counts(ds, total=2, by_state=[(ReplicaState.RUNNING, 2, v1)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY assert ds._in_transition is False # Drain node_2: one replica transitions to PENDING_MIGRATION and a # replacement STARTING replica is created. cluster_node_info_cache.draining_nodes = {node_2: 60 * 1000} dsm.update() assert ds._replicas.count(states=[ReplicaState.PENDING_MIGRATION]) == 1 assert ds._replicas.count(states=[ReplicaState.STARTING]) == 1 assert ds._in_transition is True # Node stops draining before the replacement is ready. The # PENDING_MIGRATION replica should move back to RUNNING. cluster_node_info_cache.draining_nodes = {} dsm.update() pending_migration_count = ds._replicas.count( states=[ReplicaState.PENDING_MIGRATION] ) assert pending_migration_count == 0, ( f"Expected 0 PENDING_MIGRATION replicas but found {pending_migration_count}. " "The replica is stuck because _in_transition was incorrectly cleared." ) class TestIsGangDeploymentProperty: """Tests for DeploymentState._is_gang_deployment property.""" @pytest.mark.parametrize( "gang_scheduling_config, expected_value", [ pytest.param(GangSchedulingConfig(gang_size=2), True, id="gang"), pytest.param(None, False, id="non-gang"), ], ) def test_is_gang_deployment( self, mock_deployment_state_manager, gang_scheduling_config, expected_value, ): """_is_gang_deployment reflects whether gang scheduling is configured.""" create_dsm, _, _, _ = mock_deployment_state_manager create_dsm_kwargs = {} deployment_info_kwargs = {"num_replicas": 2} if gang_scheduling_config is not None: create_dsm_kwargs[ "create_placement_group_fn_override" ] = lambda *args, **kwargs: Mock() deployment_info_kwargs["gang_scheduling_config"] = gang_scheduling_config dsm: DeploymentStateManager = create_dsm(**create_dsm_kwargs) b_info, _ = deployment_info(**deployment_info_kwargs) dsm.deploy(TEST_DEPLOYMENT_ID, b_info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] assert ds._is_gang_deployment is expected_value class TestScaleDeploymentGangReplicas: def test_stopping_replicas_skip_upscale(self, mock_deployment_state_manager): """Skips upscale while gang replicas are stopping after startup failures, then recovers to healthy.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm( create_placement_group_fn_override=lambda *args, **kwargs: Mock(), ) gang_size = 2 target_replicas = 2 deployment_id = DeploymentID(name="gang_stopping_skip", app_name="app") info, version = deployment_info( num_replicas=target_replicas, version="v1", gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size), ) dsm.deploy(deployment_id, info) ds = dsm._deployment_states[deployment_id] dsm.update() check_counts( ds, total=target_replicas, by_state=[(ReplicaState.STARTING, 2, version)] ) for replica in ds._replicas.get([ReplicaState.STARTING]): replica._actor.set_failed_to_start() dsm._deployment_scheduler.schedule_gang_placement_groups = Mock(return_value={}) captured_upscales = {} original_schedule = dsm._deployment_scheduler.schedule def schedule_with_capture(upscales, downscales): captured_upscales.update(upscales) return original_schedule(upscales, downscales) dsm._deployment_scheduler.schedule = Mock(side_effect=schedule_with_capture) dsm.update() assert captured_upscales == {} dsm._deployment_scheduler.schedule_gang_placement_groups.assert_not_called() check_counts( ds, total=target_replicas, by_state=[(ReplicaState.STOPPING, 2, version)] ) assert ds.curr_status_info.status == DeploymentStatus.UPDATING for replica in ds._replicas.get([ReplicaState.STOPPING]): replica._actor.set_done_stopping() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.UPDATING dsm._deployment_scheduler.schedule_gang_placement_groups = Mock( return_value={ deployment_id: GangReservationResult( success=True, gang_pgs=[Mock(name="pg-0")], gang_ids=["g0"], gang_pg_names=["SERVE_GANG::pg-0"], ) } ) dsm.update() check_counts( ds, total=target_replicas, by_state=[(ReplicaState.STARTING, 2, version)] ) for replica in ds._replicas.get([ReplicaState.STARTING]): replica._actor.set_ready() dsm.update() check_counts( ds, total=target_replicas, by_state=[(ReplicaState.RUNNING, 2, version)] ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_gang_reservation_failure_records_startup_failure( self, mock_deployment_state_manager ): """Keeps upscale empty and records reservation failure details before recovering to healthy.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() deployment_id = DeploymentID(name="gang_reservation_fail", app_name="app") error_msg = "simulated gang placement reservation failure" info, version = deployment_info( num_replicas=4, version="v1", gang_scheduling_config=GangSchedulingConfig(gang_size=2), ) dsm.deploy(deployment_id, info) ds = dsm._deployment_states[deployment_id] dsm._deployment_scheduler.schedule_gang_placement_groups = Mock( return_value={ deployment_id: GangReservationResult( success=False, error_message=error_msg ) } ) captured_upscales = {} original_schedule = dsm._deployment_scheduler.schedule def schedule_with_capture(upscales, downscales): captured_upscales.update(upscales) return original_schedule(upscales, downscales) dsm._deployment_scheduler.schedule = Mock(side_effect=schedule_with_capture) dsm.update() assert captured_upscales == {} check_counts(ds, total=0) assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert "Gang scheduling failed" in ds.curr_status_info.message assert error_msg in ds.curr_status_info.message dsm._deployment_scheduler.schedule_gang_placement_groups = Mock( return_value={ deployment_id: GangReservationResult( success=True, gang_pgs=[Mock(name="pg-0"), Mock(name="pg-1")], gang_ids=["g0", "g1"], gang_pg_names=["SERVE_GANG::pg-0", "SERVE_GANG::pg-1"], ) } ) dsm.update() check_counts(ds, total=4, by_state=[(ReplicaState.STARTING, 4, version)]) for replica in ds._replicas.get([ReplicaState.STARTING]): replica._actor.set_ready() dsm.update() check_counts(ds, total=4, by_state=[(ReplicaState.RUNNING, 4, version)]) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_successful_gang_reservation(self, mock_deployment_state_manager): """Creates expected gang scheduling requests and reaches healthy when all replicas become ready.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() gang_size = 2 num_gangs = 2 target_replicas = gang_size * num_gangs deployment_id = DeploymentID(name="gang_success_sched", app_name="app") gang_pgs = [Mock(name="pg-0"), Mock(name="pg-1")] info, version = deployment_info( num_replicas=target_replicas, version="v1", gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size), ) dsm.deploy(deployment_id, info) ds = dsm._deployment_states[deployment_id] gang_ids = ["gang_0", "gang_1"] gang_pg_names = ["SERVE_GANG::pg-0", "SERVE_GANG::pg-1"] dsm._deployment_scheduler.schedule_gang_placement_groups = Mock( return_value={ deployment_id: GangReservationResult( success=True, gang_pgs=gang_pgs, gang_ids=gang_ids, gang_pg_names=gang_pg_names, ) } ) captured_upscales = {} original_schedule = dsm._deployment_scheduler.schedule def schedule_with_capture(upscales, downscales): captured_upscales.update(upscales) return original_schedule(upscales, downscales) dsm._deployment_scheduler.schedule = Mock(side_effect=schedule_with_capture) dsm.update() assert deployment_id in captured_upscales scheduling_requests = captured_upscales[deployment_id] assert len(scheduling_requests) == target_replicas assert {r.gang_placement_group for r in scheduling_requests} == set(gang_pgs) assert sorted(r.gang_pg_index for r in scheduling_requests) == [0, 0, 1, 1] check_counts( ds, total=target_replicas, by_state=[(ReplicaState.STARTING, target_replicas, version)], ) starting_replicas = ds._replicas.get([ReplicaState.STARTING]) assert len(starting_replicas) == target_replicas gang_to_replicas = {} for replica in starting_replicas: gang_to_replicas.setdefault(replica.gang_context.gang_id, []).append( replica ) assert len(gang_to_replicas) == num_gangs for gang_id, replicas in gang_to_replicas.items(): assert len(replicas) == gang_size member_ids = {r.replica_id.unique_id for r in replicas} assert sorted(r.gang_context.rank for r in replicas) == list( range(gang_size) ) for replica in replicas: gang_context = replica.gang_context assert gang_context.gang_id == gang_id assert gang_context.world_size == gang_size assert set(gang_context.member_replica_ids) == member_ids assert gang_context.pg_name in gang_pg_names for replica in starting_replicas: replica._actor.set_ready() dsm.update() check_counts( ds, total=target_replicas, by_state=[(ReplicaState.RUNNING, target_replicas, version)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_gang_sibling_cleanup_on_startup_failure( self, mock_deployment_state_manager ): """Stops gang siblings when one member fails startup to avoid partial gangs, then recovers to healthy.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm( create_placement_group_fn_override=lambda *args, **kwargs: Mock(), ) gang_size = 2 target_replicas = 4 deployment_id = DeploymentID(name="gang_sibling_cleanup", app_name="app") info, version = deployment_info( num_replicas=target_replicas, version="v1", gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size), ) dsm.deploy(deployment_id, info) ds = dsm._deployment_states[deployment_id] dsm.update() starting_replicas = ds._replicas.get([ReplicaState.STARTING]) initial_context_by_replica = { r.replica_id.unique_id: ( r.gang_context.gang_id, r.gang_context.rank, r.gang_context.world_size, tuple(r.gang_context.member_replica_ids), ) for r in starting_replicas } gang_to_replicas = {} for replica in starting_replicas: gang_to_replicas.setdefault(replica.gang_context.gang_id, []).append( replica ) failed_gang_id, failed_gang_members = next(iter(gang_to_replicas.items())) failed_gang_members[0]._actor.set_failed_to_start() failed_gang_members[1]._actor.set_ready() dsm.update() stopping_replicas = ds._replicas.get([ReplicaState.STOPPING]) starting_replicas = ds._replicas.get([ReplicaState.STARTING]) assert len(stopping_replicas) == gang_size assert all(r.gang_context.gang_id == failed_gang_id for r in stopping_replicas) assert all(r.gang_context.gang_id != failed_gang_id for r in starting_replicas) surviving_gang_ids = {r.gang_context.gang_id for r in starting_replicas} assert len(surviving_gang_ids) == 1 for replica in starting_replicas: context_snapshot = initial_context_by_replica[replica.replica_id.unique_id] assert context_snapshot == ( replica.gang_context.gang_id, replica.gang_context.rank, replica.gang_context.world_size, tuple(replica.gang_context.member_replica_ids), ) check_counts( ds, total=target_replicas, by_state=[ (ReplicaState.STOPPING, 2, version), (ReplicaState.STARTING, 2, version), ], ) for replica in ds._replicas.get([ReplicaState.STOPPING]): replica._actor.set_done_stopping() for replica in ds._replicas.get([ReplicaState.STARTING]): replica._actor.set_ready() dsm.update() check_counts( ds, total=target_replicas, by_state=[ (ReplicaState.RUNNING, 2, version), (ReplicaState.STARTING, 2, version), ], ) for replica in ds._replicas.get([ReplicaState.STARTING]): replica._actor.set_ready() dsm.update() check_counts( ds, total=target_replicas, by_state=[(ReplicaState.RUNNING, target_replicas, version)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_gang_startup_failure_per_gang_counter(self, mock_deployment_state_manager): """When a gang of replicas fails to start, the failure counter should increment once per gang and not once per replica.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm( create_placement_group_fn_override=lambda *args, **kwargs: Mock(), ) gang_size = 2 target_replicas = 2 deployment_id = DeploymentID(name="gang_startup_threshold", app_name="app") with patch( "ray.serve._private.deployment_state.MAX_PER_REPLICA_RETRY_COUNT", 2 ): info, _ = deployment_info( num_replicas=target_replicas, version="v1", max_constructor_retry_count=10, gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size), ) dsm.deploy(deployment_id, info) ds = dsm._deployment_states[deployment_id] # Set by _failed_to_start_threshold -> min(max_constructor_retry_count, target_replicas * MAX_PER_REPLICA_RETRY_COUNT) = min(10, 2*2) = 4 expected_threshold = 4 assert ds._failed_to_start_threshold == expected_threshold def run_failure_cycle(): """Run one full cycle: start replicas → fail → stop → clean up.""" dsm.update() starting = ds._replicas.get([ReplicaState.STARTING]) for replica in starting: replica._actor.set_failed_to_start() # Transition failed replicas to STOPPING dsm.update() # Complete stopping for replica in ds._replicas.get([ReplicaState.STOPPING]): replica._actor.set_done_stopping() dsm.update() num_failure_cycles = 2 for _ in range(num_failure_cycles): run_failure_cycle() assert ds.curr_status_info.status == DeploymentStatus.UPDATING assert ds._replica_constructor_retry_counter == num_failure_cycles def test_terminally_failed_deployment_skips_gang_reservation( self, mock_deployment_state_manager ): """Does not reserve gang placement groups after terminal failure, and can recover on redeploy.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm() deployment_id = DeploymentID(name="gang_terminal_failure", app_name="app") info, _ = deployment_info( num_replicas=2, gang_scheduling_config=GangSchedulingConfig(gang_size=2), ) dsm.deploy(deployment_id, info) ds = dsm._deployment_states[deployment_id] dsm._deployment_scheduler.schedule_gang_placement_groups = Mock( return_value={ deployment_id: GangReservationResult( success=False, error_message="simulated gang reservation failure" ) } ) for _ in range(20): dsm.update() if ds.curr_status_info.status == DeploymentStatus.DEPLOY_FAILED: break assert ds.curr_status_info.status == DeploymentStatus.DEPLOY_FAILED dsm._deployment_scheduler.schedule_gang_placement_groups.reset_mock() dsm.update() dsm._deployment_scheduler.schedule_gang_placement_groups.assert_not_called() recovery_info, recovery_version = deployment_info( num_replicas=2, version="v2", gang_scheduling_config=GangSchedulingConfig(gang_size=2), ) dsm.deploy(deployment_id, recovery_info) dsm._deployment_scheduler.schedule_gang_placement_groups = Mock( return_value={ deployment_id: GangReservationResult( success=True, gang_pgs=[Mock(name="pg-recovery")], gang_ids=["g0"], gang_pg_names=["SERVE_GANG::pg-recovery"], ) } ) dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.STARTING, 2, recovery_version)] ) for replica in ds._replicas.get([ReplicaState.STARTING]): replica._actor.set_ready() dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.RUNNING, 2, recovery_version)] ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_gang_downscale_stops_complete_gangs(self, mock_deployment_state_manager): """Downscaling a gang deployment stops complete gangs and recovers to healthy.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm( create_placement_group_fn_override=lambda *args, **kwargs: Mock(), ) gang_size = 2 initial_replicas = 4 deployment_id = DeploymentID(name="gang_downscale", app_name="app") info, version = deployment_info( num_replicas=initial_replicas, version="v1", gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size), ) dsm.deploy(deployment_id, info) ds = dsm._deployment_states[deployment_id] # Start all replicas and reach HEALTHY dsm.update() for replica in ds._replicas.get([ReplicaState.STARTING]): replica._actor.set_ready() dsm.update() check_counts( ds, total=initial_replicas, by_state=[(ReplicaState.RUNNING, initial_replicas, version)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Downscale to 2 replicas (remove 1 gang) new_info, new_version = deployment_info( num_replicas=2, version="v1", gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size), ) dsm.deploy(deployment_id, new_info) dsm.update() # Verify exactly 1 complete gang (2 replicas) is stopping stopping = ds._replicas.get([ReplicaState.STOPPING]) running = ds._replicas.get([ReplicaState.RUNNING]) assert len(stopping) == 2 assert len(running) == 2 # The 2 stopping replicas must belong to the same gang stopping_gang_ids = {r.gang_context.gang_id for r in stopping} assert len(stopping_gang_ids) == 1 # The 2 running replicas must belong to the same (surviving) gang running_gang_ids = {r.gang_context.gang_id for r in running} assert len(running_gang_ids) == 1 assert stopping_gang_ids != running_gang_ids # Complete stopping and verify recovery to HEALTHY for replica in ds._replicas.get([ReplicaState.STOPPING]): replica._actor.set_done_stopping() dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.RUNNING, 2, new_version)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_gang_downscale_prefers_pending_gang(self, mock_deployment_state_manager): """Downscaling prefers the gang that still has a pending replica.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm( create_placement_group_fn_override=lambda *args, **kwargs: Mock(), ) gang_size = 2 initial_replicas = 4 deployment_id = DeploymentID(name="gang_downscale_pending", app_name="app") info, version = deployment_info( num_replicas=initial_replicas, version="v1", gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size), ) dsm.deploy(deployment_id, info) ds = dsm._deployment_states[deployment_id] # First update creates all 4 replicas in STARTING state dsm.update() starting = ds._replicas.get([ReplicaState.STARTING]) assert len(starting) == initial_replicas # Mark 3 of 4 replicas ready, leaving 1 replica from the second gang still pending. gangs: dict = {} for replica in starting: gangs.setdefault(replica.gang_context.gang_id, []).append(replica) gang_id1, gang_id2 = list(gangs.keys()) for replica in gangs[gang_id1]: replica._actor.set_ready() gangs[gang_id2][0]._actor.set_ready() dsm.update() check_counts( ds, total=initial_replicas, by_state=[ (ReplicaState.RUNNING, 3, version), (ReplicaState.STARTING, 1, version), ], ) # Downscale to 2 replicas — should prefer gang 2 (has a pending member) new_info, new_version = deployment_info( num_replicas=2, version="v1", gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size), ) dsm.deploy(deployment_id, new_info) dsm.update() stopping = ds._replicas.get([ReplicaState.STOPPING]) running = ds._replicas.get([ReplicaState.RUNNING]) assert len(stopping) == 2 assert len(running) == 2 stopping_gang_ids = {r.gang_context.gang_id for r in stopping} assert stopping_gang_ids == {gang_id2} running_gang_ids = {r.gang_context.gang_id for r in running} assert running_gang_ids == {gang_id1} # Complete stopping and verify recovery to HEALTHY for replica in ds._replicas.get([ReplicaState.STOPPING]): replica._actor.set_done_stopping() dsm.update() check_counts( ds, total=2, by_state=[(ReplicaState.RUNNING, 2, new_version)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY class TestGangHealthCheck: def _deploy_gang(self, mock_deployment_state_manager, gang_size, num_replicas): """Deploy gang-scheduled replicas and wait for them to become RUNNING.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm( create_placement_group_fn_override=lambda *args, **kwargs: MockPlacementGroup( *args, **kwargs ), ) b_info, v1 = deployment_info( version="1", num_replicas=num_replicas, gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size), ) dsm.deploy(TEST_DEPLOYMENT_ID, b_info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] # Reserves gang PGs and creates replicas dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.STARTING, num_replicas, v1)] ) # Capture replica references and wait for them to become RUNNING replicas = ds._replicas.get() for replica in replicas: replica._actor.set_ready() dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v1)] ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY # Group captured replicas by gang gangs = {} for r in replicas: assert r.gang_context is not None gangs.setdefault(r.gang_context.gang_id, []).append(r) return dsm, ds, v1, gangs @pytest.mark.parametrize("force_stop_unhealthy", [True, False]) def test_restart_gang_entire_gang_stopped( self, mock_deployment_state_manager, force_stop_unhealthy ): """Unhealthy gang is force-stopped regardless of FORCE_STOP_UNHEALTHY_REPLICAS; healthy gangs are unaffected.""" gang_size = 2 num_replicas = 4 num_gangs = num_replicas // gang_size dsm, ds, v1, gangs = self._deploy_gang( mock_deployment_state_manager, gang_size, num_replicas ) assert len(gangs) == num_gangs gang_ids = list(gangs.keys()) target_gang = gangs[gang_ids[0]] healthy_gang = gangs[gang_ids[1]] ds.FORCE_STOP_UNHEALTHY_REPLICAS = force_stop_unhealthy # Initialize health checks, then mark one replica in the target gang as unhealthy. dsm.update() target_gang[0]._actor.set_unhealthy() dsm.update() # Both replicas of the affected gang should be stopping (force-stopped). check_counts( ds, total=num_replicas, by_state=[ (ReplicaState.RUNNING, gang_size, v1), (ReplicaState.STOPPING, gang_size, v1), ], ) for r in target_gang: assert r._actor.force_stopped_counter == 1 # Healthy gang replicas should still be running. for r in healthy_gang: assert r._actor.force_stopped_counter == 0 assert ds.curr_status_info.status == DeploymentStatus.UNHEALTHY assert ( ds.curr_status_info.status_trigger == DeploymentStatusTrigger.HEALTH_CHECK_FAILED ) assert "UNHEALTHY" in ds.curr_status_info.message # After the stopped replicas finish stopping, new replicas should start. for r in target_gang: r._actor.set_done_stopping() dsm.update() check_counts( ds, total=num_replicas, by_state=[ (ReplicaState.RUNNING, gang_size, v1), (ReplicaState.STARTING, gang_size, v1), ], ) # New replicas become ready -> deployment should recover to HEALTHY. for r in ds._replicas.get([ReplicaState.STARTING]): r._actor.set_ready() dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v1)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_restart_gang_multiple_unhealthy_gang_replicas( self, mock_deployment_state_manager ): """Verify gang replicas are force-stopped once when there are multiple unhealthy replicas in the same gang.""" gang_size = 2 num_replicas = 4 dsm, ds, v1, gangs = self._deploy_gang( mock_deployment_state_manager, gang_size, num_replicas ) gang_ids = list(gangs.keys()) target_gang = gangs[gang_ids[0]] healthy_gang = gangs[gang_ids[1]] # Initialize health checks, then mark both replicas unhealthy. dsm.update() for r in target_gang: r._actor.set_unhealthy() dsm.update() check_counts( ds, total=num_replicas, by_state=[ (ReplicaState.RUNNING, gang_size, v1), (ReplicaState.STOPPING, gang_size, v1), ], ) for r in target_gang: assert r._actor.force_stopped_counter == 1 for r in healthy_gang: assert r._actor.force_stopped_counter == 0 assert ds.curr_status_info.status == DeploymentStatus.UNHEALTHY # Finish stopping -> new replicas start -> become ready -> HEALTHY. for r in target_gang: r._actor.set_done_stopping() dsm.update() for r in ds._replicas.get([ReplicaState.STARTING]): r._actor.set_ready() dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v1)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_restart_gang_multiple_gangs_failing(self, mock_deployment_state_manager): """Multiple gangs with unhealthy replicas are all stopped; surviving gang is untouched.""" gang_size = 2 num_replicas = 6 # 3 gangs dsm, ds, v1, gangs = self._deploy_gang( mock_deployment_state_manager, gang_size, num_replicas ) assert len(gangs) == 3 gang_ids = list(gangs.keys()) failed_gang_0 = gangs[gang_ids[0]] failed_gang_1 = gangs[gang_ids[1]] surviving_gang = gangs[gang_ids[2]] dsm.update() failed_gang_0[0]._actor.set_unhealthy() failed_gang_1[0]._actor.set_unhealthy() dsm.update() check_counts( ds, total=num_replicas, by_state=[ (ReplicaState.RUNNING, gang_size, v1), (ReplicaState.STOPPING, gang_size * 2, v1), ], ) for r in failed_gang_0 + failed_gang_1: assert r._actor.force_stopped_counter == 1 for r in surviving_gang: assert r._actor.force_stopped_counter == 0 assert ds.curr_status_info.status == DeploymentStatus.UNHEALTHY # Finish stopping -> new replicas start -> become ready -> HEALTHY. for r in failed_gang_0 + failed_gang_1: r._actor.set_done_stopping() dsm.update() check_counts( ds, total=num_replicas, by_state=[ (ReplicaState.RUNNING, gang_size, v1), (ReplicaState.STARTING, gang_size * 2, v1), ], ) for r in ds._replicas.get([ReplicaState.STARTING]): r._actor.set_ready() dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v1)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY class TestGangRollingUpdate: def _deploy_gang(self, mock_deployment_state_manager, gang_size, num_replicas): """Deploy a gang-scheduled deployment and advance to HEALTHY.""" create_dsm, _, _, _ = mock_deployment_state_manager dsm: DeploymentStateManager = create_dsm( create_placement_group_fn_override=lambda *args, **kwargs: MockPlacementGroup( *args, **kwargs ), ) info, _ = deployment_info( version="v1", num_replicas=num_replicas, gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size), ) dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm.update() for replica in ds._replicas.get(): replica._actor.set_ready() dsm.update() assert ds.curr_status_info.status == DeploymentStatus.HEALTHY return dsm, ds def _deploy_new_version(self, dsm, gang_size, num_replicas, version, **kwargs): """Deploy a new version and return its version tag.""" info, v = deployment_info( version=version, num_replicas=num_replicas, gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size), **kwargs, ) dsm.deploy(TEST_DEPLOYMENT_ID, info) return v def _mock_gang_pgs(self, dsm, gang_size, num_new): """Mock gang PG reservation for new replicas.""" n = num_new // gang_size dsm._deployment_scheduler.schedule_gang_placement_groups = Mock( return_value={ TEST_DEPLOYMENT_ID: GangReservationResult( success=True, gang_pgs=[Mock() for _ in range(n)], gang_ids=[f"new_gang_{i}" for i in range(n)], gang_pg_names=[f"SERVE_GANG::pg-{i}" for i in range(n)], ) } ) def _finish_stopping(self, ds): """Mark all STOPPING replicas as done.""" for r in ds._replicas.get(states=[ReplicaState.STOPPING]): r._actor.set_done_stopping() def _finish_starting(self, ds): """Mark all STARTING replicas as ready.""" for r in ds._replicas.get(states=[ReplicaState.STARTING]): r._actor.set_ready() def _advance_wave(self, dsm, ds, gang_size): """Complete one rolling-update wave: finish stops, mock PGs, start new.""" self._finish_stopping(ds) self._mock_gang_pgs(dsm, gang_size, gang_size) dsm.update() self._finish_starting(ds) def test_stop_gang_atomically(self, mock_deployment_state_manager): """Stops one complete gang per wave, never partially tearing down a gang.""" gang_size, num_replicas = 2, 4 dsm, ds = self._deploy_gang( mock_deployment_state_manager, gang_size, num_replicas ) v2 = self._deploy_new_version(dsm, gang_size, num_replicas, "v2") dsm.update() # First wave: exactly one gang (2 replicas) stops stopping = ds._replicas.get(states=[ReplicaState.STOPPING]) assert len(stopping) == gang_size assert len({r.gang_context.gang_id for r in stopping}) == 1 assert all(r.version != v2 for r in stopping) # Complete wave 1, trigger wave 2 self._advance_wave(dsm, ds, gang_size) dsm.update() # Second old gang now stopping assert ( ds._replicas.count(exclude_version=v2, states=[ReplicaState.STOPPING]) == gang_size ) # Complete wave 2 self._advance_wave(dsm, ds, gang_size) dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v2)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_rollout_size_rounding(self, mock_deployment_state_manager): """With gang_size=3, num_replicas=9, rollout_size rounds up from 1 to 3.""" gang_size, num_replicas = 3, 9 dsm, ds = self._deploy_gang( mock_deployment_state_manager, gang_size, num_replicas ) v2 = self._deploy_new_version(dsm, gang_size, num_replicas, "v2") dsm.update() # Exactly one gang (3 replicas) stops stopping = ds._replicas.get(states=[ReplicaState.STOPPING]) assert len(stopping) == gang_size assert len({r.gang_context.gang_id for r in stopping}) == 1 # Advance through all 3 waves for _ in range(3): self._advance_wave(dsm, ds, gang_size) dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v2)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_reconfigure(self, mock_deployment_state_manager): """Config-only change (same code version) reconfigures in place, no stops.""" gang_size, num_replicas = 2, 4 dsm, ds = self._deploy_gang( mock_deployment_state_manager, gang_size, num_replicas ) self._deploy_new_version( dsm, gang_size, num_replicas, "v1", user_config={"key": "new_value"} ) dsm.update() assert ds._replicas.count(states=[ReplicaState.STOPPING]) == 0 running = ds._replicas.count(states=[ReplicaState.RUNNING]) updating = ds._replicas.count(states=[ReplicaState.UPDATING]) assert running + updating == num_replicas def test_starting_replicas(self, mock_deployment_state_manager): """Rapid v1->v2->v3: v2 STARTING replicas are stopped when v3 arrives.""" gang_size, num_replicas = 2, 4 dsm, ds = self._deploy_gang( mock_deployment_state_manager, gang_size, num_replicas ) # Deploy v2, complete first wave's stops, but leave new replicas STARTING. v2 = self._deploy_new_version(dsm, gang_size, num_replicas, "v2") dsm.update() self._finish_stopping(ds) self._mock_gang_pgs(dsm, gang_size, gang_size) dsm.update() assert ( ds._replicas.count(version=v2, states=[ReplicaState.STARTING]) == gang_size ) # Deploy v3 while v2 replicas are still STARTING. v3 = self._deploy_new_version(dsm, gang_size, num_replicas, "v3") dsm.update() assert ( ds._replicas.count(version=v2, states=[ReplicaState.STOPPING]) == gang_size ) # State: 2 v1 RUNNING, 2 v2 STOPPING. # The rolling-update budget (pending_replicas=2) is already exhausted # by the missing slots, so v1 won't be stopped until new v3 replicas # fill those slots first. # Wave 1: Finish v2 stops -> update starts 2 v3 (v1 not stopped yet). self._finish_stopping(ds) self._mock_gang_pgs(dsm, gang_size, gang_size) dsm.update() assert ( ds._replicas.count(version=v3, states=[ReplicaState.STARTING]) == gang_size ) assert ds._replicas.count(states=[ReplicaState.STOPPING]) == 0 # Wave 2: v3 replicas ready -> update stops the old v1 gang. self._finish_starting(ds) dsm.update() assert ( ds._replicas.count(version=v3, states=[ReplicaState.RUNNING]) == gang_size ) assert ds._replicas.count(states=[ReplicaState.STOPPING]) == gang_size # Wave 3: Finish v1 stops -> update starts remaining v3 replicas. self._finish_stopping(ds) self._mock_gang_pgs(dsm, gang_size, gang_size) dsm.update() self._finish_starting(ds) dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v3)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_multi_gang_stop_per_wave(self, mock_deployment_state_manager): """Validate that multiple gangs are stopped per wave.""" gang_size, num_replicas = 2, 20 dsm, ds = self._deploy_gang( mock_deployment_state_manager, gang_size, num_replicas ) self._deploy_new_version(dsm, gang_size, num_replicas, "v2") dsm.update() stopping = ds._replicas.get(states=[ReplicaState.STOPPING]) # rollout_size = max(1, int(0.2 * 20)) = 4 assert len(stopping) == 4 stopping_gang_ids = {r.gang_context.gang_id for r in stopping} assert len(stopping_gang_ids) == 2 def test_recovering_member_skips_gang_update(self, mock_deployment_state_manager): """Gang with a RECOVERING member is skipped during rolling update.""" gang_size, num_replicas = 2, 4 dsm, ds = self._deploy_gang( mock_deployment_state_manager, gang_size, num_replicas ) # Move one replica to RECOVERING so its gang is incomplete. running = ds._replicas.pop(states=[ReplicaState.RUNNING]) recovering_replica = running[0] for r in running[1:]: ds._replicas.add(ReplicaState.RUNNING, r) ds._replicas.add(ReplicaState.RECOVERING, recovering_replica) v2 = self._deploy_new_version(dsm, gang_size, num_replicas, "v2") dsm.update() # An incomplete gang (with a RECOVERING member) cannot be stopped. # Additionally, the budget is insufficient because the RECOVERING # replica increases the pending count, preventing any stops. assert ds._replicas.count(states=[ReplicaState.STOPPING]) == 0 # Recover the replica so the gang is complete again recovering_replica._actor.set_ready() dsm.update() # Wave 1: first complete gang stops stopping = ds._replicas.get(states=[ReplicaState.STOPPING]) assert len(stopping) == gang_size assert len({r.gang_context.gang_id for r in stopping}) == 1 # Complete wave 1, trigger wave 2 self._advance_wave(dsm, ds, gang_size) dsm.update() # Wave 2: second gang stops assert ( ds._replicas.count(exclude_version=v2, states=[ReplicaState.STOPPING]) == gang_size ) # Complete wave 2 self._advance_wave(dsm, ds, gang_size) dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v2)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY class TestGangDraining: """Test gang-aware migration when nodes are draining.""" def _deploy_gang( self, mock_deployment_state_manager, gang_size, num_replicas, nodes ): """Deploy a gang deployment and assign replicas to nodes round-robin.""" create_dsm, timer, cluster_node_info_cache, _ = mock_deployment_state_manager for node in nodes: cluster_node_info_cache.add_node(node) dsm: DeploymentStateManager = create_dsm( create_placement_group_fn_override=lambda *args, **kwargs: Mock(), ) timer.reset(0) info, v = deployment_info( num_replicas=num_replicas, version="v1", graceful_shutdown_timeout_s=20, gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size), ) dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] num_gangs = num_replicas // gang_size dsm._deployment_scheduler.schedule_gang_placement_groups = Mock( return_value={ TEST_DEPLOYMENT_ID: GangReservationResult( success=True, gang_pgs=[Mock() for _ in range(num_gangs)], gang_ids=[f"gang_{i}" for i in range(num_gangs)], gang_pg_names=[f"SERVE_GANG::pg-{i}" for i in range(num_gangs)], ) } ) dsm.update() # Assign nodes round-robin and mark ready replicas = ds._replicas.get([ReplicaState.STARTING]) for i, replica in enumerate(replicas): replica._actor.set_node_id(nodes[i % len(nodes)]) replica._actor.set_ready() dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY return dsm, ds, timer, cluster_node_info_cache, v def _mock_gang_pgs(self, dsm, gang_size, num_new): """Mock gang PG reservation for new replicas.""" n = num_new // gang_size dsm._deployment_scheduler.schedule_gang_placement_groups = Mock( return_value={ TEST_DEPLOYMENT_ID: GangReservationResult( success=True, gang_pgs=[Mock() for _ in range(n)], gang_ids=[f"new_gang_{i}" for i in range(n)], gang_pg_names=[f"SERVE_GANG::pg-new-{i}" for i in range(n)], ) } ) def test_entire_gang_migration(self, mock_deployment_state_manager): """When gang members' nodes drain, ALL members move to PENDING_MIGRATION. When the deadline is up, the entire gang is stopped.""" gang_size, num_replicas = 2, 4 node_1 = "node-1" node_2 = "node-2" node_3 = "node-3" # non-draining target # graceful_shutdown_timeout_s=20 dsm, ds, timer, cache, v1 = self._deploy_gang( mock_deployment_state_manager, gang_size, num_replicas, [node_1, node_2] ) cache.add_node(node_3) # Drain both nodes with a 40s deadline, graceful_shutdown_timeout=20s. # Gangs are stopped when: curr_time >= deadline - graceful_shutdown_timeout. deadline_ms = 40 * 1000 cache.draining_nodes = {node_1: deadline_ms, node_2: deadline_ms} # Block replacements so we can test the deadline path in isolation. dsm._deployment_scheduler.schedule_gang_placement_groups = Mock(return_value={}) dsm.update() assert ( ds._replicas.count(states=[ReplicaState.PENDING_MIGRATION]) == num_replicas ) # Advance 15s — still before deadline - timeout (40 - 20 = 20s). # Gangs should remain PENDING_MIGRATION (no deadline-triggered stops). timer.advance(15) dsm.update() assert ( ds._replicas.count(states=[ReplicaState.PENDING_MIGRATION]) == num_replicas ) assert ds._replicas.count(states=[ReplicaState.STOPPING]) == 0 # Advance past deadline - graceful_shutdown_timeout (total 25s > 20s). # All old gangs should now be stopped due to deadline expiry. timer.advance(10) dsm.update() assert ds._replicas.count(states=[ReplicaState.STOPPING]) == num_replicas for r in ds._replicas.get([ReplicaState.STOPPING]): r._actor.set_done_stopping() # Schedule replacements and complete self._mock_gang_pgs(dsm, gang_size, num_replicas) dsm.update() for r in ds._replicas.get([ReplicaState.STARTING]): r._actor.set_node_id(node_3) r._actor.set_ready() dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v1)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_partial_node_drain_migrates_entire_gang( self, mock_deployment_state_manager ): """When only one node in a gang is draining, ALL gang members migrate.""" gang_size, num_replicas = 2, 2 node_1 = "node-1" node_2 = "node-2" node_3 = "node-3" dsm, ds, timer, cache, v1 = self._deploy_gang( mock_deployment_state_manager, gang_size, num_replicas, [node_1, node_2] ) cache.add_node(node_3) # Drain only node_2. The gang has members on both nodes, # so ALL members (including the one on healthy node_1) must migrate. dsm._deployment_scheduler.schedule_gang_placement_groups = Mock(return_value={}) cache.draining_nodes = {node_2: 60 * 1000} dsm.update() # Both replicas should be PENDING_MIGRATION, not just the one on node_2 assert ( ds._replicas.count(states=[ReplicaState.PENDING_MIGRATION]) == num_replicas ) assert ds._replicas.count(states=[ReplicaState.RUNNING]) == 0 # Schedule replacements on non-draining node_3 self._mock_gang_pgs(dsm, gang_size, num_replicas) dsm.update() for r in ds._replicas.get([ReplicaState.STARTING]): r._actor.set_node_id(node_3) r._actor.set_ready() # Advance past deadline so old replicas are stopped timer.advance(50) dsm.update() for r in ds._replicas.get([ReplicaState.STOPPING]): r._actor.set_done_stopping() dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v1)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_gang_recovery_when_node_stops_draining( self, mock_deployment_state_manager ): """Gang stays PENDING_MIGRATION while ANY member's node is draining, and returns to RUNNING when all members' nodes stop draining.""" gang_size, num_replicas = 2, 2 node_1 = "node-1" node_2 = "node-2" dsm, ds, _, cache, v1 = self._deploy_gang( mock_deployment_state_manager, gang_size, num_replicas, [node_1, node_2] ) # Prevent the scheduler from creating replacements during draining. dsm._deployment_scheduler.schedule_gang_placement_groups = Mock(return_value={}) # Drain both nodes — all members go to PENDING_MIGRATION. cache.draining_nodes = {node_1: 60 * 1000, node_2: 60 * 1000} dsm.update() assert ( ds._replicas.count(states=[ReplicaState.PENDING_MIGRATION]) == num_replicas ) # Stop draining node_1 but keep node_2 draining. # The gang still has a member on node_2, so it stays PENDING_MIGRATION. cache.draining_nodes = {node_2: 60 * 1000} dsm.update() assert ( ds._replicas.count(states=[ReplicaState.PENDING_MIGRATION]) == num_replicas ) # Stop draining completely — gang should recover to RUNNING. cache.draining_nodes = {} dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v1)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_starting_replicas_stopped(self, mock_deployment_state_manager): """STARTING gang members on draining nodes are stopped immediately.""" gang_size, num_replicas = 2, 2 node_1 = "node-1" node_2 = "node-2" create_dsm, timer, cache, _ = mock_deployment_state_manager cache.add_node(node_1) cache.add_node(node_2) dsm: DeploymentStateManager = create_dsm( create_placement_group_fn_override=lambda *args, **kwargs: Mock(), ) timer.reset(0) info, v1 = deployment_info( num_replicas=num_replicas, version="v1", gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size), ) dsm.deploy(TEST_DEPLOYMENT_ID, info) ds = dsm._deployment_states[TEST_DEPLOYMENT_ID] dsm._deployment_scheduler.schedule_gang_placement_groups = Mock( return_value={ TEST_DEPLOYMENT_ID: GangReservationResult( success=True, gang_pgs=[Mock()], gang_ids=["gang_0"], gang_pg_names=["SERVE_GANG::pg-0"], ) } ) dsm.update() # Assign nodes but DON'T mark ready — replicas stay STARTING replicas = ds._replicas.get([ReplicaState.STARTING]) replicas[0]._actor.set_node_id(node_1) replicas[1]._actor.set_node_id(node_2) # Drain node_2 while replicas are STARTING cache.draining_nodes = {node_2: 60 * 1000} dsm.update() # Both STARTING replicas in the gang should be stopped stopping = ds._replicas.count(states=[ReplicaState.STOPPING]) assert stopping == gang_size # Complete stopping and start new replicas for r in ds._replicas.get([ReplicaState.STOPPING]): r._actor.set_done_stopping() self._mock_gang_pgs(dsm, gang_size, num_replicas) dsm.update() for r in ds._replicas.get([ReplicaState.STARTING]): r._actor.set_node_id(node_1) r._actor.set_ready() dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v1)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY def test_gang_excess_migration_stops_complete_gangs( self, mock_deployment_state_manager ): """When replacement replicas are ready, excess gangs are stopped atomically.""" gang_size, num_replicas = 2, 4 node_1 = "node-1" node_2 = "node-2" dsm, ds, _, cache, v1 = self._deploy_gang( mock_deployment_state_manager, gang_size, num_replicas, [node_1] ) cache.add_node(node_2) # Prevent replacements during initial draining dsm._deployment_scheduler.schedule_gang_placement_groups = Mock(return_value={}) # Drain node_1 with far deadline so no deadline-triggered stops cache.draining_nodes = {node_1: 600 * 1000} dsm.update() assert ( ds._replicas.count(states=[ReplicaState.PENDING_MIGRATION]) == num_replicas ) # Start one gang worth of replacements and make them RUNNING self._mock_gang_pgs(dsm, gang_size, gang_size) dsm.update() for r in ds._replicas.get([ReplicaState.STARTING]): r._actor.set_node_id(node_2) r._actor.set_ready() dsm.update() # One complete gang of old replicas should be stopping (excess) stopping = ds._replicas.get([ReplicaState.STOPPING]) assert len(stopping) == gang_size # Verify they're all from the same gang gang_ids = {r.gang_context.gang_id for r in stopping} assert len(gang_ids) == 1 # Complete first gang: stop old, start+ready replacement, stop second old gang for r in ds._replicas.get([ReplicaState.STOPPING]): r._actor.set_done_stopping() self._mock_gang_pgs(dsm, gang_size, gang_size) dsm.update() for r in ds._replicas.get([ReplicaState.STARTING]): r._actor.set_node_id(node_2) r._actor.set_ready() dsm.update() # Clean up the second gang's STOPPING replicas for r in ds._replicas.get([ReplicaState.STOPPING]): r._actor.set_done_stopping() dsm.update() check_counts( ds, total=num_replicas, by_state=[(ReplicaState.RUNNING, num_replicas, v1)], ) assert ds.curr_status_info.status == DeploymentStatus.HEALTHY if __name__ == "__main__": sys.exit(pytest.main(["-v", "-s", __file__]))