import functools import math import time import types import unittest from collections import defaultdict from unittest.mock import MagicMock, patch import pytest import ray from ray.data._internal.execution.backpressure_policy import ( ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY, ConcurrencyCapBackpressurePolicy, ) from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer from ray.data._internal.execution.operators.task_pool_map_operator import ( TaskPoolMapOperator, ) from ray.data._internal.execution.resource_manager import ResourceManager from ray.data.context import DataContext from ray.data.tests.conftest import mock_all_to_all_op from ray.util.annotations import RayDeprecationWarning class TestConcurrencyCapBackpressurePolicy(unittest.TestCase): """Tests for ConcurrencyCapBackpressurePolicy.""" @classmethod def setUpClass(cls): cls._cluster_cpus = 10 ray.init(num_cpus=cls._cluster_cpus) data_context = ray.data.DataContext.get_current() data_context.set_config( ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY, [ConcurrencyCapBackpressurePolicy], ) @classmethod def tearDownClass(cls): ray.shutdown() data_context = ray.data.DataContext.get_current() data_context.remove_config(ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY) def _mock_resource_manager(self): """Helper to create a resource manager mock with real method bindings.""" rm = MagicMock() rm.is_op_eligible = types.MethodType(ResourceManager.is_op_eligible, rm) rm._get_downstream_ineligible_ops = types.MethodType( ResourceManager._get_downstream_ineligible_ops, rm ) rm._is_blocking_materializing_op = types.MethodType( ResourceManager._is_blocking_materializing_op, rm ) return rm def test_basic(self): concurrency = 16 input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()]) map_op_no_concurrency = TaskPoolMapOperator( map_transformer=MagicMock(), data_context=DataContext.get_current(), input_op=input_op, ) map_op = TaskPoolMapOperator( map_transformer=MagicMock(), data_context=DataContext.get_current(), input_op=map_op_no_concurrency, max_concurrency=concurrency, ) map_op.metrics.num_tasks_running = 0 map_op.metrics.num_tasks_finished = 0 topology = { map_op: MagicMock(), input_op: MagicMock(), map_op_no_concurrency: MagicMock(), } mock_resource_manager = MagicMock() # Return None to skip dynamic output queue size backpressure check mock_resource_manager.get_op_usage.return_value = None mock_resource_manager.get_budget.return_value = None mock_resource_manager.is_op_eligible.return_value = False policy = ConcurrencyCapBackpressurePolicy( DataContext.get_current(), topology, mock_resource_manager, ) self.assertEqual(policy._concurrency_caps[map_op], concurrency) self.assertTrue(math.isinf(policy._concurrency_caps[input_op])) self.assertTrue(math.isinf(policy._concurrency_caps[map_op_no_concurrency])) # Gradually increase num_tasks_running to the cap. for i in range(1, concurrency + 1): self.assertTrue(policy.can_add_input(map_op)) map_op.metrics.num_tasks_running = i # Now num_tasks_running reaches the cap, so can_add_input should return False. self.assertFalse(policy.can_add_input(map_op)) map_op.metrics.num_tasks_running = concurrency / 2 self.assertEqual(policy.can_add_input(map_op), True) def _create_record_time_actor(self): @ray.remote(num_cpus=0) class RecordTimeActor: def __init__(self): self._start_time = defaultdict(lambda: []) self._end_time = defaultdict(lambda: []) def record_start_time(self, index): self._start_time[index].append(time.time()) def record_end_time(self, index): self._end_time[index].append(time.time()) def get_start_and_end_time_for_op(self, index): return min(self._start_time[index]), max(self._end_time[index]) def get_start_and_end_time_for_all_tasks_of_op(self, index): return self._start_time[index], self._end_time[index] actor = RecordTimeActor.remote() return actor def _get_map_func(self, actor, index): def map_func(data, actor, index): actor.record_start_time.remote(index) yield data actor.record_end_time.remote(index) return functools.partial(map_func, actor=actor, index=index) def test_e2e_normal(self): """A simple E2E test with ConcurrencyCapBackpressurePolicy enabled.""" actor = self._create_record_time_actor() map_func1 = self._get_map_func(actor, 1) map_func2 = self._get_map_func(actor, 2) # Create a dataset with 2 map ops. Each map op has N tasks, where N is # the number of cluster CPUs. N = self.__class__._cluster_cpus ds = ray.data.range(N, override_num_blocks=N) # Use different `num_cpus` to make sure they don't fuse. ds = ds.map_batches(map_func1, batch_size=None, num_cpus=1, concurrency=1) ds = ds.map_batches(map_func2, batch_size=None, num_cpus=1.1, concurrency=1) res = ds.take_all() self.assertEqual(len(res), N) # We recorded the start and end time of each op, # check that these 2 ops are executed interleavingly. # This means that the executor didn't allocate all resources to the first # op in the beginning. start1, end1 = ray.get(actor.get_start_and_end_time_for_op.remote(1)) start2, end2 = ray.get(actor.get_start_and_end_time_for_op.remote(2)) assert start1 < start2 < end1 < end2, (start1, start2, end1, end2) def test_can_add_input_with_dynamic_output_queue_size_backpressure_disabled(self): """Test can_add_input when dynamic output queue size backpressure is disabled.""" input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()]) map_op = TaskPoolMapOperator( map_transformer=MagicMock(), data_context=DataContext.get_current(), input_op=input_op, max_concurrency=5, ) map_op.metrics.num_tasks_running = 3 topology = {map_op: MagicMock(), input_op: MagicMock()} # Create policy with dynamic output queue size backpressure disabled policy = ConcurrencyCapBackpressurePolicy( DataContext.get_current(), topology, MagicMock(), # resource_manager ) policy.enable_dynamic_output_queue_size_backpressure = False # Should only check against configured concurrency cap self.assertTrue(policy.can_add_input(map_op)) # 3 < 5 map_op.metrics.num_tasks_running = 5 self.assertFalse(policy.can_add_input(map_op)) # 5 >= 5 def test_can_add_input_with_non_map_operator(self): """Test can_add_input with non-MapOperator (should use basic cap check).""" input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()]) input_op.metrics.num_tasks_running = 1 topology = {input_op: MagicMock()} policy = ConcurrencyCapBackpressurePolicy( DataContext.get_current(), topology, MagicMock(), # resource_manager ) # InputDataBuffer has infinite concurrency cap, so should always allow self.assertTrue(policy.can_add_input(input_op)) def test_can_add_input_with_ineligible_op(self): """Test can_add_input when op is not eligible for backpressure.""" input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()]) map_op = TaskPoolMapOperator( map_transformer=MagicMock(), data_context=DataContext.get_current(), input_op=input_op, max_concurrency=5, ) map_op.metrics.num_tasks_running = 3 topology = {map_op: MagicMock(), input_op: MagicMock()} mock_resource_manager = self._mock_resource_manager() # Override to test policy behavior when op is not eligible mock_resource_manager.is_op_eligible = MagicMock(return_value=False) policy = ConcurrencyCapBackpressurePolicy( DataContext.get_current(), topology, mock_resource_manager, ) policy.enable_dynamic_output_queue_size_backpressure = True # Should skip dynamic backpressure and use basic cap check self.assertTrue(policy.can_add_input(map_op)) # 3 < 5 map_op.metrics.num_tasks_running = 5 self.assertFalse(policy.can_add_input(map_op)) # 5 >= 5 def test_can_add_input_with_materializing_downstream_op(self): """Test can_add_input when downstream op is a materializing operator.""" input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()]) map_op = TaskPoolMapOperator( map_transformer=MagicMock(), data_context=DataContext.get_current(), input_op=input_op, max_concurrency=5, ) map_op.metrics.num_tasks_running = 3 # Create materializing downstream op (automatically adds to map_op._output_dependencies) mock_all_to_all_op(map_op) topology = {map_op: MagicMock(), input_op: MagicMock()} mock_resource_manager = self._mock_resource_manager() policy = ConcurrencyCapBackpressurePolicy( DataContext.get_current(), topology, mock_resource_manager, ) policy.enable_dynamic_output_queue_size_backpressure = True # Should skip dynamic backpressure and use basic cap check # to avoid starving materializing operators self.assertTrue(policy.can_add_input(map_op)) # 3 < 5 map_op.metrics.num_tasks_running = 5 self.assertFalse(policy.can_add_input(map_op)) # 5 >= 5 @patch( "ray.data._internal.execution.backpressure_policy." "concurrency_cap_backpressure_policy.get_available_object_store_budget_fraction" ) def test_can_add_input_with_object_store_memory_usage_ratio_above_threshold( self, mock_get_budget_fraction ): """Test can_add_input when object store memory usage ratio is above threshold.""" input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()]) map_op = TaskPoolMapOperator( map_transformer=MagicMock(), data_context=DataContext.get_current(), input_op=input_op, max_concurrency=5, ) map_op.metrics.num_tasks_running = 3 topology = {map_op: MagicMock(), input_op: MagicMock()} mock_resource_manager = self._mock_resource_manager() # Mock available object store memory budget fraction above threshold to skip dynamic backpressure threshold = ( ConcurrencyCapBackpressurePolicy.AVAILABLE_OBJECT_STORE_BUDGET_THRESHOLD ) # Set fraction above threshold to skip dynamic backpressure mock_get_budget_fraction.return_value = threshold + 0.05 policy = ConcurrencyCapBackpressurePolicy( DataContext.get_current(), topology, mock_resource_manager, ) policy.enable_dynamic_output_queue_size_backpressure = True # Initialize EWMA state to verify it's not updated when ratio > threshold initial_level = 100.0 initial_dev = 20.0 policy._q_level_nbytes[map_op] = initial_level policy._q_level_dev[map_op] = initial_dev # Should skip dynamic backpressure and use basic cap check # EWMA state should not be updated (early return) self.assertTrue(policy.can_add_input(map_op)) # 3 < 5 self.assertEqual(policy._q_level_nbytes[map_op], initial_level) self.assertEqual(policy._q_level_dev[map_op], initial_dev) map_op.metrics.num_tasks_running = 5 self.assertFalse(policy.can_add_input(map_op)) # 5 >= 5 # EWMA state should still not be updated self.assertEqual(policy._q_level_nbytes[map_op], initial_level) self.assertEqual(policy._q_level_dev[map_op], initial_dev) @patch( "ray.data._internal.execution.backpressure_policy." "concurrency_cap_backpressure_policy.get_available_object_store_budget_fraction" ) def test_can_add_input_with_object_store_memory_usage_ratio_below_threshold( self, mock_get_budget_fraction ): """Test can_add_input when object store memory usage ratio is below threshold.""" input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()]) map_op = TaskPoolMapOperator( map_transformer=MagicMock(), data_context=DataContext.get_current(), input_op=input_op, max_concurrency=5, ) map_op.metrics.num_tasks_running = 3 topology = {map_op: MagicMock(), input_op: MagicMock()} mock_resource_manager = self._mock_resource_manager() # Mock available object store memory budget fraction below threshold to apply dynamic backpressure threshold = ( ConcurrencyCapBackpressurePolicy.AVAILABLE_OBJECT_STORE_BUDGET_THRESHOLD ) # Set fraction below threshold to apply dynamic backpressure mock_get_budget_fraction.return_value = threshold - 0.05 # Mock queue size methods mock_resource_manager.get_mem_op_internal.return_value = 100 mock_resource_manager.get_mem_op_outputs.return_value = 200 policy = ConcurrencyCapBackpressurePolicy( DataContext.get_current(), topology, mock_resource_manager, ) policy.enable_dynamic_output_queue_size_backpressure = True # Should proceed with dynamic backpressure logic # Initialize EWMA state for the operator with a different level # so we can verify the update happens (queue size is 300) initial_level = 200.0 initial_dev = 50.0 policy._q_level_nbytes[map_op] = initial_level policy._q_level_dev[map_op] = initial_dev result = policy.can_add_input(map_op) # With queue size 300, initial level=200, dev=50, bounds=[150, 250] # Queue size 300 is above the upper bound, so should backoff. # running=3, backoff by 1 -> effective_cap=2 # running=3 < effective_cap=2 should be False self.assertFalse(result) # EWMA state should be updated when ratio < threshold # Level should move toward 300 (queue size) self.assertNotEqual(policy._q_level_nbytes[map_op], initial_level) # Dev should also be updated self.assertNotEqual(policy._q_level_dev[map_op], initial_dev) @patch( "ray.data._internal.execution.backpressure_policy." "concurrency_cap_backpressure_policy.get_available_object_store_budget_fraction" ) def test_can_add_input_effective_cap_calculation(self, mock_get_budget_fraction): """Test that effective cap calculation works correctly with different queue sizes.""" input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()]) map_op = TaskPoolMapOperator( map_transformer=MagicMock(), data_context=DataContext.get_current(), input_op=input_op, max_concurrency=8, ) map_op.metrics.num_tasks_running = 4 topology = {map_op: MagicMock(), input_op: MagicMock()} mock_resource_manager = self._mock_resource_manager() threshold = ( ConcurrencyCapBackpressurePolicy.AVAILABLE_OBJECT_STORE_BUDGET_THRESHOLD ) # Set fraction below threshold to apply dynamic backpressure mock_get_budget_fraction.return_value = threshold - 0.05 policy = ConcurrencyCapBackpressurePolicy( DataContext.get_current(), topology, mock_resource_manager, ) policy.enable_dynamic_output_queue_size_backpressure = True # Test different queue sizes using policy constants test_cases = [ # (internal_usage, downstream_usage, level, dev, expected_result, description) ( 50, 50, 5000.0, 200.0, True, "low_queue_below_lower_bound", ), # 100 < 5000 - 2*200 = 4600, ramp up ( 200, 200, 400.0, 50.0, False, "medium_queue_in_hold_region", ), # 400 in [300, 500], hold ( 300, 300, 200.0, 50.0, False, "high_queue_above_upper_bound", ), # 600 > 200 + 2*50 = 300, backoff ] for ( internal_usage, downstream_usage, level, dev, expected_result, description, ) in test_cases: with self.subTest(description=description): mock_resource_manager.get_mem_op_internal.return_value = internal_usage mock_resource_manager.get_mem_op_outputs.return_value = downstream_usage mock_resource_manager.get_op_outputs_object_store_usage_with_downstream.return_value = ( downstream_usage ) # Initialize EWMA state policy._q_level_nbytes[map_op] = level policy._q_level_dev[map_op] = dev result = policy.can_add_input(map_op) assert ( result == expected_result ), f"Expected {expected_result} for {description}" def test_emits_deprecation_warning_when_dynamic_backpressure_enabled( restore_data_context, ): ctx = DataContext.get_current() ctx.enable_dynamic_output_queue_size_backpressure = True input_op = InputDataBuffer(ctx, input_data=[MagicMock()]) topology = {input_op: MagicMock()} with pytest.warns(RayDeprecationWarning, match="deprecated"): ConcurrencyCapBackpressurePolicy(ctx, topology, MagicMock()) if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", __file__]))