wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
13afa51474b122dcddeaf71832238bb5ef6ecaf1
ray-project--ray/python/ray/serve/tests/unit/test_deployment_state.py
T
wehub-resource-sync f1825c8ceb chore: import upstream snapshot with attribution
2026-07-13 13:17:40 +08:00

9646 lines
358 KiB
Python
Raw Blame History

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__]))
Reference in New Issue View Git Blame Copy Permalink