import json from typing import List, Optional, Tuple from unittest import mock from unittest.mock import patch import pytest from ray._common.test_utils import wait_for_condition from ray._common.utils import Timer from ray.serve._private.cluster_node_info_cache import ClusterNodeInfoCache from ray.serve._private.common import RequestProtocol from ray.serve._private.constants import ( PROXY_HEALTH_CHECK_UNHEALTHY_THRESHOLD, RAY_SERVE_FALLBACK_PROXY_GRPC_PORT, RAY_SERVE_FALLBACK_PROXY_HTTP_PORT, ) from ray.serve._private.proxy_state import ProxyState, ProxyStateManager, ProxyWrapper from ray.serve._private.test_utils import MockTimer from ray.serve.config import HTTPOptions, ProxyLocation from ray.serve.schema import LoggingConfig, ProxyStatus HEAD_NODE_ID = "node_id-index-head" class MockClusterNodeInfoCache: def __init__(self): self.alive_nodes = [] def get_alive_nodes(self): return self.alive_nodes def get_alive_node_ids(self): return {node_id for node_id, _, _ in self.alive_nodes} class FakeProxyActor: def __init__(self, *args, **kwargs): pass def ready(self): return json.dumps(["mock_worker_id", "mock_log_file_path"]) def check_health(self): pass class FakeProxyWrapper(ProxyWrapper): def __init__(self, *args, **kwargs): self.actor_handle = FakeProxyActor(*args, **kwargs) self.is_ready_response = None self.is_healthy_response = None self.is_drained_response = False self.worker_id = "mock_worker_id" self.log_file_path = "mock_log_file_path" self.shutdown = False self.num_health_checks = 0 self.num_drain_checks = 0 @property def actor_id(self) -> str: pass def is_ready(self, timeout_s: float) -> Optional[bool]: return self.is_ready_response def is_healthy(self, timeout_s: float) -> Optional[bool]: self.num_health_checks += 1 return self.is_healthy_response def is_drained(self, timeout_s: float) -> Optional[bool]: self.num_drain_checks += 1 return self.is_drained_response def is_shutdown(self): return self.shutdown def update_draining(self, draining: bool): pass def kill(self): self.shutdown = True def get_num_health_checks(self): return self.num_health_checks def get_num_drain_checks(self): return self.num_health_checks def _create_proxy_state_manager( http_options: HTTPOptions = HTTPOptions(), head_node_id: str = HEAD_NODE_ID, cluster_node_info_cache=MockClusterNodeInfoCache(), actor_proxy_wrapper_class=FakeProxyWrapper, timer=Timer(), running_native_proxies: bool = False, proxy_location=None, ) -> (ProxyStateManager, ClusterNodeInfoCache): return ( ProxyStateManager( http_options=http_options, head_node_id=head_node_id, cluster_node_info_cache=cluster_node_info_cache, logging_config=LoggingConfig(), actor_proxy_wrapper_class=actor_proxy_wrapper_class, timer=timer, running_native_proxies=running_native_proxies, proxy_location=proxy_location, ), cluster_node_info_cache, ) def _create_proxy_state( actor_proxy_wrapper_class=FakeProxyWrapper, status: ProxyStatus = ProxyStatus.STARTING, node_id: str = "mock_node_id", timer=Timer(), **kwargs, ) -> ProxyState: state = ProxyState( actor_proxy_wrapper=actor_proxy_wrapper_class(), actor_name="alice", node_id=node_id, node_ip="mock_node_ip", node_instance_id="mock_instance_id", timer=timer, ) state._set_status(status=status) return state @pytest.fixture def number_of_worker_nodes() -> int: return 100 @pytest.fixture def all_nodes(number_of_worker_nodes) -> List[Tuple[str, str, str]]: return [(HEAD_NODE_ID, "fake-head-ip", "fake-head-instance-id")] + [ (f"worker-node-id-{i}", f"fake-worker-ip-{i}", f"fake-instance-id-{i}") for i in range(number_of_worker_nodes) ] def _reconcile_and_check_proxy_status(state: ProxyState, status: ProxyStatus): state.reconcile() assert state.status == status return True def _update_and_check_proxy_state_manager( proxy_state_manager: ProxyStateManager, node_ids: List[str], statuses: List[ProxyStatus], **kwargs, ): proxy_state_manager.update(**kwargs) proxy_states = proxy_state_manager._proxy_states assert all( [ proxy_states[node_ids[idx]].status == statuses[idx] for idx in range(len(node_ids)) ] ), [proxy_state.status for proxy_state in proxy_states.values()] return True def test_node_selection_via_proxy_location(all_nodes): # `proxy_location` is the placement authority (threaded separately from the # deprecated HTTPOptions.location); covers HTTP and gRPC (single ProxyActor). all_node_ids = {node_id for node_id, _, _ in all_nodes} def mgr(**kwargs): psm, cache = _create_proxy_state_manager(**kwargs) cache.alive_nodes = all_nodes return psm assert ( mgr(proxy_location=ProxyLocation.Disabled)._get_target_nodes(all_node_ids) == [] ) assert ( mgr(proxy_location=ProxyLocation.HeadOnly)._get_target_nodes(all_node_ids) == all_nodes[:1] ) assert ( mgr(proxy_location=ProxyLocation.EveryNode)._get_target_nodes(all_node_ids) == all_nodes ) # Default (neither location nor proxy_location) resolves to EveryNode. assert mgr()._get_target_nodes(all_node_ids) == all_nodes assert mgr().get_proxy_location() == ProxyLocation.EveryNode # Explicit (deprecated) HTTPOptions.location overrides proxy_location. with pytest.warns(DeprecationWarning, match="`location` in HTTPOptions"): override = HTTPOptions(location=ProxyLocation.HeadOnly) psm = mgr(http_options=override, proxy_location=ProxyLocation.EveryNode) assert psm._get_target_nodes(all_node_ids) == all_nodes[:1] assert psm.get_proxy_location() == ProxyLocation.HeadOnly def test_node_selection(all_nodes): all_node_ids = {node_id for node_id, _, _ in all_nodes} # Test NoServer proxy_state_manager, cluster_node_info_cache = _create_proxy_state_manager( HTTPOptions(location=ProxyLocation.Disabled) ) cluster_node_info_cache.alive_nodes = all_nodes assert proxy_state_manager._get_target_nodes(all_node_ids) == [] # Test HeadOnly proxy_state_manager, cluster_node_info_cache = _create_proxy_state_manager( HTTPOptions(location=ProxyLocation.HeadOnly) ) cluster_node_info_cache.alive_nodes = all_nodes assert proxy_state_manager._get_target_nodes(all_node_ids) == all_nodes[:1] # Test EveryNode proxy_state_manager, cluster_node_info_cache = _create_proxy_state_manager( HTTPOptions(location=ProxyLocation.EveryNode) ) cluster_node_info_cache.alive_nodes = all_nodes assert proxy_state_manager._get_target_nodes(all_node_ids) == all_nodes # Test specific nodes proxy_state_manager, cluster_node_info_cache = _create_proxy_state_manager( HTTPOptions(location=ProxyLocation.EveryNode) ) cluster_node_info_cache.alive_nodes = all_nodes assert proxy_state_manager._get_target_nodes({HEAD_NODE_ID}) == [ (HEAD_NODE_ID, "fake-head-ip", "fake-head-instance-id") ] @patch("ray.serve._private.proxy_state.PROXY_HEALTH_CHECK_PERIOD_S", 5) def test_proxy_state_manager_restarts_unhealthy_proxies(all_nodes): """Test the update method in ProxyStateManager would kill and restart unhealthy proxies. """ timer = MockTimer() proxy_state_manager, cluster_node_info_cache = _create_proxy_state_manager( timer=timer ) cluster_node_info_cache.alive_nodes = all_nodes # First iteration, refresh state proxy_state_manager.update(proxy_nodes={HEAD_NODE_ID}) prev_proxy_state = proxy_state_manager._proxy_states[HEAD_NODE_ID] # Mark existing head-node proxy UNHEALTHY prev_proxy_state._set_status(ProxyStatus.UNHEALTHY) old_proxy = prev_proxy_state.actor_handle # Continuously trigger update and wait for status to be changed to HEALTHY. for _ in range(1): proxy_state_manager.update(proxy_nodes={HEAD_NODE_ID}) # Advance timer by 5 (to perform a health-check) timer.advance(5) new_proxy_state = proxy_state_manager._proxy_states[HEAD_NODE_ID] # Previous proxy's state stays UNHEALTHY assert prev_proxy_state.status == ProxyStatus.UNHEALTHY # Ensure the old proxy is getting shutdown. assert prev_proxy_state._shutting_down # New proxy's state should be STARTING assert new_proxy_state.status == ProxyStatus.STARTING assert new_proxy_state.proxy_restart_count == 1 new_proxy = new_proxy_state.actor_handle # Ensure the new proxy is completely different object than old proxy. assert new_proxy != old_proxy def test_proxy_state_reconcile_shutting_down(): proxy_state = _create_proxy_state() previous_status = proxy_state.status proxy_state.shutdown() # This should be no-op. The status of the http proxy state will not be changed. proxy_state.reconcile() current_status = proxy_state.status # Ensure the proxy state is in the shutting down state. assert proxy_state._shutting_down # Ensure the status didn't change. assert previous_status == current_status def test_proxy_state_reconcile_readiness_check_succeed(): proxy_state = _create_proxy_state() # Configure is_ready to be true proxy_state._actor_proxy_wrapper.is_ready_response = True # Ensure the proxy status before update is STARTING. assert proxy_state.status == ProxyStatus.STARTING # Ensure actor_details are set to the initial state when the proxy_state is created. assert proxy_state.actor_details.worker_id is None assert proxy_state.actor_details.log_file_path is None assert proxy_state.actor_details.status == ProxyStatus.STARTING.value # Continuously trigger update and wait for status to be changed. proxy_state.reconcile() assert proxy_state.status == ProxyStatus.HEALTHY # Ensure actor_details are updated. assert proxy_state.actor_details.worker_id == "mock_worker_id" assert proxy_state.actor_details.log_file_path == "mock_log_file_path" assert proxy_state.actor_details.status == ProxyStatus.HEALTHY.value def test_proxy_state_reconcile_readiness_check_pending(): proxy_state = _create_proxy_state() # Ensure the proxy status before update is STARTING. assert proxy_state.status == ProxyStatus.STARTING # When the proxy readiness check is pending, the proxy wrapper is_ready # will return None proxy_state._actor_proxy_wrapper.is_ready_response = None # Trigger update. The status do not change, while readiness check is pending for _ in range(10): proxy_state.reconcile() assert proxy_state.status == ProxyStatus.STARTING # Unblock is_ready call, trigger update, and wait for status change to HEALTHY. proxy_state._actor_proxy_wrapper.is_ready_response = True proxy_state.reconcile() assert proxy_state.status == ProxyStatus.HEALTHY def test_proxy_state_reconcile_readiness_check_fails(): proxy_state = _create_proxy_state() # Emulate readiness check failure proxy_state._actor_proxy_wrapper.is_ready_response = False # Ensure the proxy status before update is STARTING. assert proxy_state.status == ProxyStatus.STARTING # First failure shouldn't trigger state-transition to UNHEALTHY proxy_state.reconcile() assert proxy_state.status == ProxyStatus.STARTING # After PROXY_HEALTH_CHECK_UNHEALTHY_THRESHOLD failures, state should # transition to UNHEALTHY for _ in range(PROXY_HEALTH_CHECK_UNHEALTHY_THRESHOLD - 1): proxy_state.reconcile() assert proxy_state.status == ProxyStatus.UNHEALTHY @patch("ray.serve._private.proxy_state.PROXY_READY_CHECK_TIMEOUT_S", 5) def test_proxy_state_reconcile_health_check_succeed(): timer = MockTimer() proxy_state = _create_proxy_state(time=timer) # Emulate readiness check succeeding proxy_state._actor_proxy_wrapper.is_ready_response = True # Should transition to HEALTHY proxy_state.reconcile() assert proxy_state.status == ProxyStatus.HEALTHY # Trigger update few more times and the status continue to be HEALTHY. for _ in range(10): _reconcile_and_check_proxy_status(proxy_state, ProxyStatus.HEALTHY) # Advance timer by 5s timer.advance(5) # Health-checks should have been performed on every iteration assert proxy_state._actor_proxy_wrapper.num_health_checks == 10 @patch("ray.serve._private.proxy_state.PROXY_HEALTH_CHECK_PERIOD_S", 5) def test_proxy_state_reconcile_health_check_transient_failures(): timer = MockTimer() # Start with HEALTHY state proxy_state = _create_proxy_state(status=ProxyStatus.HEALTHY, timer=timer) # Simulate health-checks failing proxy_state._actor_proxy_wrapper.is_healthy_response = False # Reconcile PROXY_HEALTH_CHECK_UNHEALTHY_THRESHOLD - 1 times, state should # continue to stay HEALTHY for _ in range(PROXY_HEALTH_CHECK_UNHEALTHY_THRESHOLD - 1): _reconcile_and_check_proxy_status(proxy_state, ProxyStatus.HEALTHY) # Advance timer by 5 (to trigger new health-check) timer.advance(5) assert ( proxy_state._actor_proxy_wrapper.get_num_health_checks() == PROXY_HEALTH_CHECK_UNHEALTHY_THRESHOLD - 1 ) # Simulate health-checks passing proxy_state._actor_proxy_wrapper.is_healthy_response = True wait_for_condition( condition_predictor=_reconcile_and_check_proxy_status, state=proxy_state, status=ProxyStatus.HEALTHY, ) # Ensure _consecutive_health_check_failures is reset assert proxy_state._consecutive_health_check_failures == 0 @patch("ray.serve._private.proxy_state.PROXY_HEALTH_CHECK_PERIOD_S", 5) def test_proxy_state_reconcile_health_check_persistent_failures(): timer = MockTimer() # Start with HEALTHY state proxy_state = _create_proxy_state(status=ProxyStatus.HEALTHY, timer=timer) # Simulate health-checks failing proxy_state._actor_proxy_wrapper.is_healthy_response = False # Reconcile PROXY_HEALTH_CHECK_UNHEALTHY_THRESHOLD - 1 times, state should # continue to stay HEALTHY for _ in range(PROXY_HEALTH_CHECK_UNHEALTHY_THRESHOLD - 1): _reconcile_and_check_proxy_status(proxy_state, ProxyStatus.HEALTHY) # Advance timer by 5 (to trigger new health-check) timer.advance(5) assert ( proxy_state._actor_proxy_wrapper.get_num_health_checks() == PROXY_HEALTH_CHECK_UNHEALTHY_THRESHOLD - 1 ) # On PROXY_HEALTH_CHECK_UNHEALTHY_THRESHOLD iteration, state transitions to # UNHEALTHY _reconcile_and_check_proxy_status(proxy_state, ProxyStatus.UNHEALTHY) # Ensure _consecutive_health_check_failures is correct assert proxy_state._consecutive_health_check_failures == 3 @patch("ray.serve._private.proxy_state.PROXY_HEALTH_CHECK_PERIOD_S", 0) @pytest.mark.parametrize("number_of_worker_nodes", [0, 1, 2, 3]) def test_proxy_manager_update_proxies_states(all_nodes, number_of_worker_nodes): """Test update draining logics. When update nodes to inactive, head node http proxy should never be draining while worker node http proxy should change to draining. When update nodes to active, head node http proxy should continue to be healthy while worker node http proxy should be healthy. """ manager, cluster_node_info_cache = _create_proxy_state_manager( HTTPOptions(location=ProxyLocation.EveryNode) ) cluster_node_info_cache.alive_nodes = all_nodes for node_id, _, _ in all_nodes: manager._proxy_states[node_id] = _create_proxy_state( status=ProxyStatus.HEALTHY, node_id=node_id, ) node_ids = [node_id for node_id, _, _ in all_nodes] # No target proxy nodes proxy_nodes = {HEAD_NODE_ID} # Head node proxy should continue to be HEALTHY. # Worker node proxy should turn DRAINING. wait_for_condition( condition_predictor=_update_and_check_proxy_state_manager, proxy_state_manager=manager, node_ids=node_ids, statuses=[ProxyStatus.HEALTHY] + [ProxyStatus.DRAINING] * number_of_worker_nodes, proxy_nodes=proxy_nodes, ) # All nodes are target proxy nodes proxy_nodes = set(node_ids) # Head node proxy should continue to be HEALTHY. # Worker node proxy should turn HEALTHY. wait_for_condition( condition_predictor=_update_and_check_proxy_state_manager, proxy_state_manager=manager, node_ids=node_ids, statuses=[ProxyStatus.HEALTHY] * (number_of_worker_nodes + 1), proxy_nodes=proxy_nodes, ) @patch("ray.serve._private.proxy_state.PROXY_HEALTH_CHECK_PERIOD_S", 5) def test_proxy_state_reconcile_draining_success(): """Test that the proxy will remain DRAINING even if health check succeeds.""" timer = MockTimer(start_time=0) # Start with HEALTHY state proxy_state = _create_proxy_state(status=ProxyStatus.HEALTHY, timer=timer) # Simulate health-checks passing proxy_state._actor_proxy_wrapper.is_healthy_response = True # Simulate is_drained returning false proxy_state._actor_proxy_wrapper.is_drained_response = False for _ in range(10): proxy_state.reconcile(draining=True) assert proxy_state.status == ProxyStatus.DRAINING # Advance timer by 5 (to trigger new health-check, drain-check) timer.advance(5) assert proxy_state._actor_proxy_wrapper.get_num_health_checks() == 10 assert proxy_state._actor_proxy_wrapper.get_num_drain_checks() == 10 # Make sure the status is still DRAINING assert proxy_state.status == ProxyStatus.DRAINING # Simulate is_drained request to ProxyActor pending (for 5 iterations) proxy_state._actor_proxy_wrapper.is_drained_response = None for _ in range(5): proxy_state.reconcile(draining=True) assert proxy_state.status == ProxyStatus.DRAINING # Advance timer by 5 (to trigger new health-check, drain-check) timer.advance(5) assert proxy_state._actor_proxy_wrapper.get_num_health_checks() == 15 # No new drain checks will occur, since there's a pending one (not completed yet) assert proxy_state._actor_proxy_wrapper.get_num_drain_checks() == 15 # Simulate draining completed proxy_state._actor_proxy_wrapper.is_drained_response = True # Advance timer by 5 (to trigger new health-check, drain-check on next iteration) timer.advance(5) proxy_state.reconcile(draining=True) # State should transition to DRAINED assert proxy_state.status == ProxyStatus.DRAINED @patch("ray.serve._private.proxy_state.PROXY_DRAIN_CHECK_PERIOD_S", 5) @pytest.mark.parametrize("number_of_worker_nodes", [1]) def test_proxy_actor_manager_removing_proxies(all_nodes, number_of_worker_nodes): """Test the state transition from DRAINING to UNHEALTHY for the proxy actor.""" assert len(all_nodes) == 2, "There should be 2 nodes in this test" timer = MockTimer(start_time=0) manager, cluster_node_info_cache = _create_proxy_state_manager( HTTPOptions(location=ProxyLocation.EveryNode), timer=timer, ) cluster_node_info_cache.alive_nodes = all_nodes for node_id, _, _ in all_nodes: manager._proxy_states[node_id] = _create_proxy_state( status=ProxyStatus.STARTING, node_id=node_id, timer=timer, ) manager._proxy_states[node_id]._actor_proxy_wrapper.is_ready_response = True # All nodes are target proxy nodes node_ids = [node_id for node_id, _, _ in all_nodes] worker_node_id = node_ids[1] worker_proxy_state = manager._proxy_states[worker_node_id] # Reconcile all proxies states manager.update( proxy_nodes=(set(node_ids)), ) # Assert all proxies are HEALTHY proxy_statuses = [manager._proxy_states[node_id].status for node_id in node_ids] assert [ProxyStatus.HEALTHY, ProxyStatus.HEALTHY] == proxy_statuses # Reconcile proxies with empty set of target nodes (ie only proxy on the head-node # should be preserved all the other should be drained) worker_proxy_state._actor_proxy_wrapper.is_drained_response = False N = 10 for _ in range(N): manager.update( proxy_nodes={HEAD_NODE_ID}, ) timer.advance(5) # Assert that # - Head-node proxy is HEALTHY4 # - Worker node proxy is DRAINING proxy_statuses = [manager._proxy_states[node_id].status for node_id in node_ids] assert [ProxyStatus.HEALTHY, ProxyStatus.DRAINING] == proxy_statuses assert worker_proxy_state._actor_proxy_wrapper.get_num_drain_checks() == N # Mark target proxy as fully drained worker_proxy_state._actor_proxy_wrapper.is_drained_response = True # Reconcile proxies with empty set of target nodes (worker node proxy # will be shutdown by now) manager.update( proxy_nodes={HEAD_NODE_ID}, ) assert len(manager._proxy_states) == 1 assert manager._proxy_states[HEAD_NODE_ID].status == ProxyStatus.HEALTHY def test_is_ready_for_shutdown(all_nodes): """Test `is_ready_for_shutdown()` returns True the correct state. Before `shutdown()` is called, `is_ready_for_shutdown()` should return false. After `shutdown()` is called and all proxy actor are killed, `is_ready_for_shutdown()` should return true. """ manager, cluster_node_info_cache = _create_proxy_state_manager( HTTPOptions(location=ProxyLocation.EveryNode) ) cluster_node_info_cache.alive_nodes = all_nodes for node_id, _, _ in all_nodes: manager._proxy_states[node_id] = _create_proxy_state( status=ProxyStatus.HEALTHY, node_id=node_id, ) # Ensure before shutdown, manager is not shutdown assert not manager.is_ready_for_shutdown() manager.shutdown() # Ensure after shutdown, manager is shutdown and all proxy states are shutdown def check_is_ready_for_shutdown(): return manager.is_ready_for_shutdown() wait_for_condition(check_is_ready_for_shutdown) @patch("ray.serve._private.proxy_state.PROXY_HEALTH_CHECK_UNHEALTHY_THRESHOLD", 1) @pytest.mark.parametrize("number_of_worker_nodes", [1]) def test_proxy_state_manager_timing_out_on_start(number_of_worker_nodes, all_nodes): """Test update method on ProxyStateManager when the proxy state is STARTING and when the ready call takes longer than PROXY_READY_CHECK_TIMEOUT_S. The proxy state started with STARTING. After update is called, ready calls takes some time to finish. The proxy state manager will restart the proxy state after PROXY_READY_CHECK_TIMEOUT_S. After the next period of check_health call, the proxy state manager will check on backoff timeout, not immediately restarting the proxy states, and eventually set the proxy state to HEALTHY. """ fake_time = MockTimer() proxy_state_manager, cluster_node_info_cache = _create_proxy_state_manager( http_options=HTTPOptions(location=ProxyLocation.EveryNode), timer=fake_time, ) cluster_node_info_cache.alive_nodes = all_nodes node_ids = {node[0] for node in all_nodes} # Run update to create proxy states. proxy_state_manager.update(proxy_nodes=node_ids) # Ensure 2 proxies are created, one for the head node and another for the worker. assert len(proxy_state_manager._proxy_states) == len(node_ids) # Ensure the proxy state statuses before update are STARTING, set the # readiness check to failing for node_id in node_ids: proxy_state = proxy_state_manager._proxy_states[node_id] assert proxy_state.status == ProxyStatus.STARTING proxy_state._actor_proxy_wrapper.is_ready_response = False # Capture current proxy states (prior to updating) prev_proxy_states = dict(proxy_state_manager._proxy_states) # Trigger PSM to reconcile proxy_state_manager.update(proxy_nodes=node_ids) # Ensure the proxy state statuses before update are STARTING, set the # readiness check to failing for node_id in node_ids: proxy_state = proxy_state_manager._proxy_states[node_id] prev_proxy_state = prev_proxy_states[node_id] # Assert # - All proxies are restarted # - Previous proxy states are UNHEALTHY # - New proxy states are STARTING assert proxy_state != prev_proxy_state assert prev_proxy_state.status == ProxyStatus.UNHEALTHY assert proxy_state.status == ProxyStatus.STARTING # Mark new proxy readiness checks as passing proxy_state._actor_proxy_wrapper.is_ready_response = True # Capture current proxy states again (prior to updating) prev_proxy_states = dict(proxy_state_manager._proxy_states) # Trigger PSM to reconcile proxy_state_manager.update(proxy_nodes=node_ids) # Ensure the proxy state statuses before update are STARTING, set the # readiness check to failing for node_id in node_ids: proxy_state = proxy_state_manager._proxy_states[node_id] prev_proxy_state = prev_proxy_states[node_id] # Assert # - All proxies are restarted # - Previous proxy states are UNHEALTHY # - New proxy states are STARTING assert proxy_state == prev_proxy_state assert prev_proxy_state.status == ProxyStatus.HEALTHY assert proxy_state.status == ProxyStatus.HEALTHY def test_proxy_state_manager_get_targets(all_nodes): """Test the get_targets method on ProxyStateManager.""" manager, cluster_node_info_cache = _create_proxy_state_manager( HTTPOptions(location=ProxyLocation.EveryNode) ) cluster_node_info_cache.alive_nodes = all_nodes for node_id, _, _ in all_nodes: manager._proxy_states[node_id] = _create_proxy_state( status=ProxyStatus.HEALTHY, node_id=node_id, ) manager._proxy_states[all_nodes[-1][0]].try_update_status(ProxyStatus.DRAINED) targets = manager.get_targets(RequestProtocol.HTTP) assert len(targets) == len(all_nodes) - 1 assert targets[0].ip == "mock_node_ip" assert targets[0].port == 8000 assert targets[0].instance_id == "mock_instance_id" assert targets[0].name == "alice" targets = manager.get_targets(RequestProtocol.GRPC) assert len(targets) == 0 with pytest.raises(ValueError): manager.get_targets("invalid_protocol") class TestFallbackProxy: """Tests for fallback proxy lifecycle management in ProxyStateManager.""" HEAD_NODE = (HEAD_NODE_ID, "fake-head-ip", "fake-head-instance-id") WORKER_NODE = ("worker-node-id-0", "fake-worker-ip-0", "fake-worker-instance-id-0") def _create_manager_with_fallback(self, timer=None): """Helper to create a ProxyStateManager with running_native_proxies=True.""" timer = timer or Timer() cache = MockClusterNodeInfoCache() cache.alive_nodes = [self.HEAD_NODE, self.WORKER_NODE] manager, cluster_node_info_cache = _create_proxy_state_manager( http_options=HTTPOptions(location=ProxyLocation.HeadOnly), cluster_node_info_cache=cache, running_native_proxies=True, timer=timer, ) return manager, cluster_node_info_cache def test_fallback_proxy_starts_on_head_node(self): """When running_native_proxies=True, a fallback proxy should be started on the head node.""" manager, _ = self._create_manager_with_fallback() assert manager._fallback_proxy_state is None manager.update(proxy_nodes={HEAD_NODE_ID}) assert manager._fallback_proxy_state is not None assert manager._fallback_proxy_state._node_id == HEAD_NODE_ID assert manager._fallback_proxy_state._status == ProxyStatus.STARTING def test_no_fallback_proxy_when_not_native(self): """When running_native_proxies=False, no fallback proxy should be started.""" cache = MockClusterNodeInfoCache() cache.alive_nodes = [self.HEAD_NODE] manager, _ = _create_proxy_state_manager( http_options=HTTPOptions(location=ProxyLocation.HeadOnly), cluster_node_info_cache=cache, running_native_proxies=False, ) manager.update(proxy_nodes={HEAD_NODE_ID}) assert manager._fallback_proxy_state is None def test_fallback_proxy_uses_correct_ports(self): """The fallback proxy should use the fallback ports (8500/9500), not the default proxy ports (8000/9000).""" cache = MockClusterNodeInfoCache() cache.alive_nodes = [self.HEAD_NODE, self.WORKER_NODE] manager, _ = _create_proxy_state_manager( http_options=HTTPOptions(location=ProxyLocation.HeadOnly), cluster_node_info_cache=cache, running_native_proxies=True, ) with mock.patch.object( manager, "_start_proxy", wraps=manager._start_proxy ) as mock_start: manager.update(proxy_nodes={HEAD_NODE_ID}) fallback_calls = [ call for call in mock_start.call_args_list if "fallback" in call.kwargs.get("name", "") ] assert len(fallback_calls) == 1 kwargs = fallback_calls[0].kwargs assert kwargs["http_options"].port == RAY_SERVE_FALLBACK_PROXY_HTTP_PORT assert kwargs["grpc_options"].port == RAY_SERVE_FALLBACK_PROXY_GRPC_PORT def test_fallback_proxy_not_started_twice(self): """Calling update multiple times should not create a second fallback proxy.""" manager, _ = self._create_manager_with_fallback() manager.update(proxy_nodes={HEAD_NODE_ID}) first_fallback = manager._fallback_proxy_state manager.update(proxy_nodes={HEAD_NODE_ID}) assert manager._fallback_proxy_state is first_fallback @patch("ray.serve._private.proxy_state.PROXY_HEALTH_CHECK_UNHEALTHY_THRESHOLD", 1) def test_stop_unhealthy_fallback_proxy(self): """When the fallback proxy becomes unhealthy, it should be stopped and the restart count incremented.""" manager, _ = self._create_manager_with_fallback() manager.update(proxy_nodes={HEAD_NODE_ID}) assert manager._fallback_proxy_state is not None assert manager._fallback_proxy_restart_count == 0 manager._fallback_proxy_state._set_status(ProxyStatus.UNHEALTHY) manager.update(proxy_nodes={HEAD_NODE_ID}) # Unhealthy fallback proxy should be stopped and cleared # A new one should be started in the same update cycle assert manager._fallback_proxy_restart_count == 1 assert manager._fallback_proxy_state is not None assert manager._fallback_proxy_state.status == ProxyStatus.STARTING def test_stop_drained_fallback_proxy(self): """When the fallback proxy enters DRAINED status, it should be removed.""" manager, _ = self._create_manager_with_fallback() manager.update(proxy_nodes={HEAD_NODE_ID}) assert manager._fallback_proxy_state is not None manager._fallback_proxy_state._set_status(ProxyStatus.DRAINED) manager.update(proxy_nodes={HEAD_NODE_ID}) # Drained fallback proxy is stopped, then a new one is started assert manager._fallback_proxy_restart_count == 1 assert manager._fallback_proxy_state is not None assert manager._fallback_proxy_state.status == ProxyStatus.STARTING def test_fallback_proxy_included_in_handles(self): """get_proxy_handles() should include the fallback proxy.""" manager, _ = self._create_manager_with_fallback() manager.update(proxy_nodes={HEAD_NODE_ID}) handles = manager.get_proxy_handles() assert f"fallback-{HEAD_NODE_ID}" in handles assert ( handles[f"fallback-{HEAD_NODE_ID}"] is manager._fallback_proxy_state.actor_handle ) def test_fallback_proxy_included_in_names(self): """get_proxy_names() should include the fallback proxy.""" manager, _ = self._create_manager_with_fallback() manager.update(proxy_nodes={HEAD_NODE_ID}) names = manager.get_proxy_names() assert f"fallback-{HEAD_NODE_ID}" in names assert ( names[f"fallback-{HEAD_NODE_ID}"] == manager._fallback_proxy_state.actor_name ) def test_fallback_proxy_included_in_alive_actor_ids(self): """get_alive_proxy_actor_ids() should include the fallback proxy's actor ID.""" manager, _ = self._create_manager_with_fallback() manager.update(proxy_nodes={HEAD_NODE_ID}) actor_ids = manager.get_alive_proxy_actor_ids() assert manager._fallback_proxy_state.actor_id in actor_ids def test_fallback_proxy_blocks_shutdown_readiness(self): """is_ready_for_shutdown() should return False until the fallback proxy is also shut down.""" manager, _ = self._create_manager_with_fallback() manager.update(proxy_nodes={HEAD_NODE_ID}) # Shut down regular proxies but not the fallback for proxy_state in manager._proxy_states.values(): proxy_state.shutdown() proxy_state._actor_proxy_wrapper.shutdown = True # Fallback proxy is still alive, so not ready for shutdown assert not manager.is_ready_for_shutdown() # Now shut down the fallback proxy manager._fallback_proxy_state.shutdown() manager._fallback_proxy_state._actor_proxy_wrapper.shutdown = True assert manager.is_ready_for_shutdown() def test_shutdown_includes_fallback_proxy(self): """shutdown() should also shut down the fallback proxy.""" manager, _ = self._create_manager_with_fallback() manager.update(proxy_nodes={HEAD_NODE_ID}) assert manager._fallback_proxy_state is not None fallback_wrapper = manager._fallback_proxy_state._actor_proxy_wrapper manager.shutdown() assert manager._fallback_proxy_state._shutting_down assert fallback_wrapper.shutdown is True def test_fallback_proxy_reconciled_during_update(self): """The fallback proxy should be reconciled (health checked) during each update cycle.""" timer = MockTimer() manager, _ = self._create_manager_with_fallback(timer=timer) manager.update(proxy_nodes={HEAD_NODE_ID}) # Make fallback proxy ready manager._fallback_proxy_state._actor_proxy_wrapper.is_ready_response = True manager.update(proxy_nodes={HEAD_NODE_ID}) assert manager._fallback_proxy_state.status == ProxyStatus.HEALTHY # Set health check to pass and verify it runs manager._fallback_proxy_state._actor_proxy_wrapper.is_healthy_response = True timer.advance(5) manager.update(proxy_nodes={HEAD_NODE_ID}) assert manager._fallback_proxy_state.status == ProxyStatus.HEALTHY assert manager._fallback_proxy_state._actor_proxy_wrapper.num_health_checks > 0 def test_get_fallback_proxy_targets(self): """get_fallback_proxy_targets() should return the fallback proxy targets when the fallback proxy is set and healthy.""" manager, _ = self._create_manager_with_fallback() # The fallback proxy is not set. targets = manager.get_fallback_proxy_targets() assert len(targets) == 0 # The fallback proxy is not healthy. state = _create_proxy_state() manager._fallback_proxy_state = state targets = manager.get_fallback_proxy_targets() assert len(targets) == 0 # The fallback proxy is healthy. state.try_update_status(ProxyStatus.HEALTHY) targets = manager.get_fallback_proxy_targets() assert len(targets) == 1 target = targets[RequestProtocol.HTTP] assert target.ip == state.actor_details.node_ip assert target.port == 8500 assert target.instance_id == state.actor_details.node_instance_id assert target.name == state.actor_name if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", "-s", __file__]))