# coding: utf-8 import os import sys from queue import Queue from unittest import mock import pytest from ray._common.test_utils import wait_for_condition from ray._common.utils import binary_to_hex, hex_to_binary from ray.autoscaler._private.prom_metrics import AutoscalerPrometheusMetrics from ray.autoscaler.v2.instance_manager.cloud_providers.read_only.cloud_provider import ( ReadOnlyProvider, ) from ray.autoscaler.v2.instance_manager.subscribers.cloud_instance_updater import ( CloudInstanceUpdater, ) from ray.autoscaler.v2.instance_manager.subscribers.ray_stopper import ( # noqa RayStopper, ) from ray.autoscaler.v2.metrics_reporter import AutoscalerMetricsReporter from ray.core.generated.autoscaler_pb2 import DrainNodeReason from ray.core.generated.instance_manager_pb2 import ( Instance, InstanceUpdateEvent, TerminationRequest, ) def _get_stopped_nodes_total(metrics_reporter: AutoscalerMetricsReporter) -> float: total_samples = [ sample.value for metric in metrics_reporter._prom_metrics.stopped_nodes.collect() for sample in metric.samples if sample.name == "autoscaler_stopped_nodes_total" ] assert len(total_samples) == 1 return total_samples[0] class TestRayStopper: def test_no_op(self): mock_gcs_client = mock.MagicMock() ray_stopper = RayStopper(gcs_client=mock_gcs_client, error_queue=Queue()) ray_stopper.notify( [ InstanceUpdateEvent( instance_id="test_id", new_instance_status=Instance.REQUESTED, ) ] ) assert mock_gcs_client.drain_node.call_count == 0 @pytest.mark.parametrize( "drain_accepted", [True, False, None], ids=["drain_accepted", "drain_rejected", "drain_error"], ) def test_idle_termination(self, drain_accepted): mock_gcs_client = mock.MagicMock() if drain_accepted is None: mock_gcs_client.drain_node.side_effect = Exception("error") else: mock_gcs_client.drain_node.return_value = ( drain_accepted, f"accepted={str(drain_accepted)}", ) error_queue = Queue() ray_stopper = RayStopper(gcs_client=mock_gcs_client, error_queue=error_queue) ray_stopper.notify( [ InstanceUpdateEvent( instance_id="test_id", new_instance_status=Instance.RAY_STOP_REQUESTED, termination_request=TerminationRequest( cause=TerminationRequest.Cause.IDLE, idle_duration_ms=1000, ray_node_id="0000", ), ) ] ) def verify(): mock_gcs_client.drain_node.assert_has_calls( [ mock.call( node_id="0000", reason=DrainNodeReason.DRAIN_NODE_REASON_IDLE_TERMINATION, reason_message=( "Termination of node that's idle for 1.0 seconds." ), deadline_timestamp_ms=0, ) ] ) if drain_accepted: assert error_queue.empty() else: error = error_queue.get_nowait() assert error.im_instance_id == "test_id" return True wait_for_condition(verify) @pytest.mark.parametrize( "stop_accepted", [True, False], ids=["stop_accepted", "stop_rejected"], ) def test_preemption(self, stop_accepted): mock_gcs_client = mock.MagicMock() mock_gcs_client.drain_nodes.return_value = [0] if stop_accepted else [] error_queue = Queue() ray_stopper = RayStopper(gcs_client=mock_gcs_client, error_queue=error_queue) ray_stopper.notify( [ InstanceUpdateEvent( instance_id="i-1", new_instance_status=Instance.RAY_STOP_REQUESTED, termination_request=TerminationRequest( cause=TerminationRequest.Cause.MAX_NUM_NODE_PER_TYPE, max_num_nodes_per_type=10, ray_node_id=binary_to_hex(hex_to_binary(b"1111")), ), ), InstanceUpdateEvent( instance_id="i-2", new_instance_status=Instance.RAY_STOP_REQUESTED, termination_request=TerminationRequest( cause=TerminationRequest.Cause.MAX_NUM_NODES, max_num_nodes=100, ray_node_id=binary_to_hex(hex_to_binary(b"2222")), ), ), ] ) def verify(): mock_gcs_client.drain_nodes.assert_has_calls( [ mock.call( node_ids=[hex_to_binary(b"1111")], ), mock.call( node_ids=[hex_to_binary(b"2222")], ), ] ) if stop_accepted: assert error_queue.empty() else: error_in_ids = set() while not error_queue.empty(): error = error_queue.get_nowait() error_in_ids.add(error.im_instance_id) assert error_in_ids == {"i-1", "i-2"} return True wait_for_condition(verify) class TestCloudInstanceUpdater: def test_launch_no_op(self): mock_provider = mock.MagicMock() launcher = CloudInstanceUpdater(mock_provider) launcher.notify( [ InstanceUpdateEvent( new_instance_status=Instance.RAY_RUNNING, launch_request_id="1", instance_type="type-1", ), ] ) mock_provider.launch.assert_not_called() def test_launch_new_instances(self): mock_provider = mock.MagicMock() launcher = CloudInstanceUpdater(mock_provider) launcher.notify( [ InstanceUpdateEvent( new_instance_status=Instance.REQUESTED, launch_request_id="1", instance_type="type-1", ), InstanceUpdateEvent( new_instance_status=Instance.REQUESTED, launch_request_id="1", instance_type="type-1", ), InstanceUpdateEvent( new_instance_status=Instance.REQUESTED, launch_request_id="2", instance_type="type-1", ), InstanceUpdateEvent( new_instance_status=Instance.REQUESTED, launch_request_id="2", instance_type="type-2", ), ] ) def verify(): mock_provider.launch.assert_has_calls( [ mock.call(shape={"type-1": 2}, request_id="1"), mock.call(shape={"type-1": 1, "type-2": 1}, request_id="2"), ] ) return True wait_for_condition(verify) def test_multi_notify(self): mock_provider = mock.MagicMock() launcher = CloudInstanceUpdater(mock_provider) launcher.notify( [ InstanceUpdateEvent( new_instance_status=Instance.REQUESTED, launch_request_id="1", instance_type="type-1", ), ] ) launcher.notify( [ InstanceUpdateEvent( new_instance_status=Instance.REQUESTED, launch_request_id="2", instance_type="type-1", ), ] ) def verify(): assert mock_provider.launch.call_count == 2 mock_provider.launch.assert_has_calls( [ mock.call(shape={"type-1": 1}, request_id="1"), mock.call(shape={"type-1": 1}, request_id="2"), ] ) return True wait_for_condition(verify) def test_terminate_no_op(self): mock_provider = mock.MagicMock() launcher = CloudInstanceUpdater(mock_provider) launcher.notify( [ InstanceUpdateEvent( new_instance_status=Instance.RAY_RUNNING, instance_id="1", cloud_instance_id="c1", ), ] ) def verify(): mock_provider.terminate.assert_not_called() return True wait_for_condition(verify) def test_terminate_instances(self): mock_provider = mock.MagicMock() metrics_reporter = AutoscalerMetricsReporter( AutoscalerPrometheusMetrics(session_name="test") ) launcher = CloudInstanceUpdater(mock_provider, metrics_reporter) launcher.notify( [ InstanceUpdateEvent( new_instance_status=Instance.TERMINATING, instance_id="1", cloud_instance_id="c1", ), InstanceUpdateEvent( new_instance_status=Instance.TERMINATING, instance_id="2", cloud_instance_id="c2", ), InstanceUpdateEvent( new_instance_status=Instance.TERMINATING, instance_id="3", cloud_instance_id="c3", ), ] ) def verify(): mock_provider.terminate.assert_called_once_with( ids=["c1", "c2", "c3"], request_id=mock.ANY ) assert _get_stopped_nodes_total(metrics_reporter) == 0 return True wait_for_condition(verify) def test_count_stopped_instances_on_terminated(self): mock_provider = mock.MagicMock() metrics_reporter = AutoscalerMetricsReporter( AutoscalerPrometheusMetrics(session_name="test") ) launcher = CloudInstanceUpdater(mock_provider, metrics_reporter) launcher.notify( [ InstanceUpdateEvent( new_instance_status=Instance.TERMINATED, instance_id="1", cloud_instance_id="c1", ), InstanceUpdateEvent( new_instance_status=Instance.TERMINATED, instance_id="2", cloud_instance_id="c2", ), InstanceUpdateEvent( new_instance_status=Instance.TERMINATED, instance_id="3", ), ] ) def verify(): mock_provider.terminate.assert_not_called() assert _get_stopped_nodes_total(metrics_reporter) == 2 return True wait_for_condition(verify) class TestReadOnlyProvider: def test_terminate_raises_not_implemented_with_correct_interface(self): """ReadOnlyProvider.terminate() must accept ids and request_id kwargs. Regression test for: TypeError: ReadOnlyProvider.terminate() got an unexpected keyword argument 'ids' CloudInstanceUpdater calls provider.terminate(ids=..., request_id=...) which matches the ICloudInstanceProvider interface. ReadOnlyProvider must accept the same signature even though it raises NotImplementedError. """ provider = object.__new__(ReadOnlyProvider) # skip __init__ (needs GCS) with pytest.raises(NotImplementedError): provider.terminate(ids=["node-1", "node-2"], request_id="req-1") def test_terminate_via_cloud_instance_updater_raises_not_implemented(self): """Verify the full call path: CloudInstanceUpdater -> ReadOnlyProvider. When a TERMINATING event arrives, CloudInstanceUpdater calls provider.terminate(ids=..., request_id=...). With ReadOnlyProvider this should surface as NotImplementedError, not TypeError. """ provider = object.__new__(ReadOnlyProvider) updater = CloudInstanceUpdater(cloud_provider=provider) with pytest.raises(NotImplementedError): updater.notify( [ InstanceUpdateEvent( new_instance_status=Instance.TERMINATING, instance_id="i-1", cloud_instance_id="cloud-node-1", ), ] ) if __name__ == "__main__": if os.environ.get("PARALLEL_CI"): sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__])) else: sys.exit(pytest.main(["-sv", __file__]))