wehub-resource-sync
477 lines
18 KiB
Python
477 lines
18 KiB
Python
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__]))
|