wehub-resource-sync
1087 lines
44 KiB
Python
1087 lines
44 KiB
Python
import math
|
|
import time
|
|
from dataclasses import replace
|
|
from datetime import timedelta
|
|
from typing import Any, Dict, Optional
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from freezegun import freeze_time
|
|
|
|
import ray
|
|
from ray.data._internal.compute import ComputeStrategy
|
|
from ray.data._internal.execution.block_ref_counter import BlockRefCounter
|
|
from ray.data._internal.execution.interfaces import BlockEntry, PhysicalOperator
|
|
from ray.data._internal.execution.interfaces.execution_options import (
|
|
ExecutionOptions,
|
|
ExecutionResources,
|
|
)
|
|
from ray.data._internal.execution.interfaces.physical_operator import (
|
|
ObjectStoreUsage,
|
|
TaskExecDriverStats,
|
|
)
|
|
from ray.data._internal.execution.operators.base_physical_operator import (
|
|
AllToAllOperator,
|
|
)
|
|
from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
|
|
from ray.data._internal.execution.operators.join import JoinOperator
|
|
from ray.data._internal.execution.operators.limit_operator import LimitOperator
|
|
from ray.data._internal.execution.operators.map_operator import MapOperator
|
|
from ray.data._internal.execution.operators.union_operator import UnionOperator
|
|
from ray.data._internal.execution.resource_manager import (
|
|
ResourceManager,
|
|
create_resource_allocator,
|
|
)
|
|
from ray.data._internal.execution.streaming_executor_state import (
|
|
IdleDetector,
|
|
OutputBackpressureGuard,
|
|
build_streaming_topology,
|
|
)
|
|
from ray.data._internal.execution.util import make_ref_bundles
|
|
from ray.data.block import TaskExecWorkerStats
|
|
from ray.data.context import DataContext
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.data.tests.conftest import noop_counter
|
|
|
|
|
|
def mock_map_op(
|
|
input_op: PhysicalOperator,
|
|
ray_remote_args: Optional[Dict[str, Any]] = None,
|
|
compute_strategy: Optional[ComputeStrategy] = None,
|
|
name="Map",
|
|
):
|
|
op = MapOperator.create(
|
|
MagicMock(),
|
|
input_op,
|
|
DataContext.get_current(),
|
|
ray_remote_args=ray_remote_args or {},
|
|
compute_strategy=compute_strategy,
|
|
name=name,
|
|
)
|
|
return op
|
|
|
|
|
|
def mock_union_op(input_ops):
|
|
op = UnionOperator(
|
|
DataContext.get_current(),
|
|
*input_ops,
|
|
)
|
|
op.start = MagicMock(side_effect=lambda *_: None)
|
|
return op
|
|
|
|
|
|
def mock_join_op(left_input_op, right_input_op):
|
|
left_input_op._logical_operators = [MagicMock()]
|
|
right_input_op._logical_operators = [MagicMock()]
|
|
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle._get_total_cluster_resources"
|
|
) as mock:
|
|
mock.return_value = ExecutionResources(cpu=1)
|
|
|
|
op = JoinOperator(
|
|
DataContext.get_current(),
|
|
left_input_op,
|
|
right_input_op,
|
|
("id",),
|
|
("id",),
|
|
"inner",
|
|
num_partitions=1,
|
|
partition_size_hint=1,
|
|
)
|
|
|
|
op.start = MagicMock(side_effect=lambda *_: None)
|
|
return op
|
|
|
|
|
|
def mock_all_to_all_op(input_op, name="MockShuffle"):
|
|
"""Create a mock AllToAllOperator (shuffle) for testing."""
|
|
op = AllToAllOperator(
|
|
bulk_fn=MagicMock(),
|
|
input_op=input_op,
|
|
data_context=DataContext.get_current(),
|
|
name=name,
|
|
)
|
|
return op
|
|
|
|
|
|
def _resource_manager_for_limits_only_test(
|
|
options: ExecutionOptions,
|
|
get_total_resources,
|
|
):
|
|
"""``ResourceManager`` requires a valid single-sink topology; these tests only
|
|
call ``get_global_limits()`` and never iterate real operators."""
|
|
sink = MagicMock(spec=PhysicalOperator)
|
|
sink.output_dependencies = []
|
|
topology = {sink: MagicMock()}
|
|
return ResourceManager(
|
|
topology,
|
|
options,
|
|
get_total_resources,
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
|
|
|
|
class TestResourceManager:
|
|
"""Unit tests for ResourceManager."""
|
|
|
|
def test_global_limits(self):
|
|
cluster_resources = {"CPU": 10, "GPU": 5, "object_store_memory": 1000}
|
|
default_object_store_memory_limit = math.ceil(
|
|
cluster_resources["object_store_memory"]
|
|
* ResourceManager.DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION
|
|
)
|
|
|
|
def get_total_resources():
|
|
return ExecutionResources.from_resource_dict(cluster_resources)
|
|
|
|
# Test default resource limits.
|
|
# When no resource limits are set, the resource limits should default to
|
|
# the cluster resources for CPU/GPU, and
|
|
# DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION of cluster object store memory.
|
|
options = ExecutionOptions()
|
|
resource_manager = _resource_manager_for_limits_only_test(
|
|
options, get_total_resources
|
|
)
|
|
expected = ExecutionResources(
|
|
cpu=cluster_resources["CPU"],
|
|
gpu=cluster_resources["GPU"],
|
|
object_store_memory=default_object_store_memory_limit,
|
|
)
|
|
assert resource_manager.get_global_limits() == expected
|
|
|
|
# Test setting resource_limits
|
|
options = ExecutionOptions()
|
|
options.resource_limits = ExecutionResources(
|
|
cpu=1, gpu=2, object_store_memory=100
|
|
)
|
|
resource_manager = _resource_manager_for_limits_only_test(
|
|
options, get_total_resources
|
|
)
|
|
expected = ExecutionResources(
|
|
cpu=1,
|
|
gpu=2,
|
|
object_store_memory=100,
|
|
)
|
|
assert resource_manager.get_global_limits() == expected
|
|
|
|
# Test setting exclude_resources
|
|
# The actual limit should be the default limit minus the excluded resources.
|
|
options = ExecutionOptions()
|
|
options.exclude_resources = ExecutionResources(
|
|
cpu=1, gpu=2, object_store_memory=100
|
|
)
|
|
resource_manager = _resource_manager_for_limits_only_test(
|
|
options, get_total_resources
|
|
)
|
|
expected = ExecutionResources(
|
|
cpu=cluster_resources["CPU"] - 1,
|
|
gpu=cluster_resources["GPU"] - 2,
|
|
object_store_memory=default_object_store_memory_limit - 100,
|
|
)
|
|
assert resource_manager.get_global_limits() == expected
|
|
|
|
# Test that we don't support setting both resource_limits
|
|
# and exclude_resources.
|
|
with pytest.raises(ValueError):
|
|
options = ExecutionOptions()
|
|
options.resource_limits = ExecutionResources(cpu=2)
|
|
options.exclude_resources = ExecutionResources(cpu=1)
|
|
options.validate()
|
|
|
|
def test_global_limits_cache(self):
|
|
get_total_resources = MagicMock(return_value=ExecutionResources(4, 1, 0))
|
|
|
|
cache_interval_s = 0.1
|
|
with patch.object(
|
|
ResourceManager,
|
|
"GLOBAL_LIMITS_UPDATE_INTERVAL_S",
|
|
cache_interval_s,
|
|
):
|
|
resource_manager = _resource_manager_for_limits_only_test(
|
|
ExecutionOptions(),
|
|
get_total_resources,
|
|
)
|
|
expected_resource = ExecutionResources(4, 1, 0)
|
|
# The first call should call ray.cluster_resources().
|
|
assert resource_manager.get_global_limits() == expected_resource
|
|
assert get_total_resources.call_count == 1
|
|
# The second call should return the cached value.
|
|
assert resource_manager.get_global_limits() == expected_resource
|
|
assert get_total_resources.call_count == 1
|
|
time.sleep(cache_interval_s)
|
|
# After the cache interval, the third call should call
|
|
# ray.cluster_resources() again.
|
|
assert resource_manager.get_global_limits() == expected_resource
|
|
assert get_total_resources.call_count == 2
|
|
|
|
def test_update_usage(self):
|
|
"""Test calculating op_usage."""
|
|
o1 = InputDataBuffer(DataContext.get_current(), [])
|
|
o2 = mock_map_op(o1)
|
|
o3 = mock_map_op(o2)
|
|
topo = build_streaming_topology(o3, ExecutionOptions(), noop_counter())
|
|
|
|
# Mock different metrics that contribute to the resource usage.
|
|
mock_cpu = {
|
|
o1: 0,
|
|
o2: 5,
|
|
o3: 8,
|
|
}
|
|
mock_pending_task_outputs = {
|
|
o1: 0,
|
|
o2: 100,
|
|
o3: 200,
|
|
}
|
|
mock_internal_outqueue = {
|
|
o1: 0,
|
|
o2: 300,
|
|
o3: 400,
|
|
}
|
|
mock_external_outqueue_sizes = {
|
|
o1: 100,
|
|
o2: 500,
|
|
o3: 600,
|
|
}
|
|
mock_internal_inqueue = {
|
|
o1: 0,
|
|
o2: 700,
|
|
o3: 800,
|
|
}
|
|
mock_pending_task_inputs = {
|
|
o1: 0,
|
|
o2: 900,
|
|
o3: 1000,
|
|
}
|
|
|
|
for op in [o1, o2, o3]:
|
|
op.current_logical_usage = MagicMock(
|
|
return_value=ExecutionResources(cpu=mock_cpu[op], gpu=0, memory=0)
|
|
)
|
|
op.running_logical_usage = MagicMock(
|
|
return_value=ExecutionResources(cpu=mock_cpu[op], gpu=0, memory=0)
|
|
)
|
|
op.pending_logical_usage = MagicMock(return_value=ExecutionResources.zero())
|
|
op.extra_resource_usage = MagicMock(return_value=ExecutionResources.zero())
|
|
op._metrics = MagicMock(
|
|
obj_store_mem_pending_task_outputs=mock_pending_task_outputs[op],
|
|
obj_store_mem_internal_outqueue=mock_internal_outqueue[op],
|
|
obj_store_mem_internal_inqueue=mock_internal_inqueue[op],
|
|
obj_store_mem_pending_task_inputs=mock_pending_task_inputs[op],
|
|
)
|
|
op._metrics.obj_store_mem_internal_inqueue_for_input = MagicMock(
|
|
return_value=mock_internal_inqueue[op],
|
|
)
|
|
ref_bundle = MagicMock(
|
|
size_bytes=MagicMock(return_value=mock_external_outqueue_sizes[op]),
|
|
output_split_idx=None,
|
|
)
|
|
topo[op].add_output(ref_bundle)
|
|
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
MagicMock(),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
resource_manager._op_resource_allocator = None
|
|
resource_manager.update_usages()
|
|
|
|
global_cpu = 0
|
|
global_mem = 0
|
|
for op in [o1, o2, o3]:
|
|
if op == o1:
|
|
# Resource usage of InputDataBuffer doesn't count.
|
|
expected_mem = 0
|
|
else:
|
|
expected_mem = (
|
|
mock_pending_task_outputs[op]
|
|
+ mock_internal_outqueue[op]
|
|
+ mock_external_outqueue_sizes[op]
|
|
)
|
|
for next_op in op.output_dependencies:
|
|
expected_mem += (
|
|
+mock_internal_inqueue[next_op]
|
|
+ mock_pending_task_inputs[next_op]
|
|
)
|
|
op_usage = resource_manager.get_op_usage(op)
|
|
assert op_usage.cpu == mock_cpu[op]
|
|
assert op_usage.gpu == 0
|
|
assert op_usage.object_store_memory == expected_mem
|
|
if op != o1:
|
|
# _mem_op_internal only includes pending_task_outputs
|
|
assert (
|
|
resource_manager._mem_op_internal[op]
|
|
== mock_pending_task_outputs[op]
|
|
)
|
|
assert (
|
|
resource_manager._mem_op_outputs[op]
|
|
== expected_mem - resource_manager._mem_op_internal[op]
|
|
)
|
|
global_cpu += mock_cpu[op]
|
|
global_mem += expected_mem
|
|
|
|
assert resource_manager.get_global_usage() == ExecutionResources(
|
|
global_cpu, 0, global_mem
|
|
)
|
|
|
|
def test_object_store_usage(self, restore_data_context):
|
|
input = make_ref_bundles([[x] for x in range(1)])[0]
|
|
# Set block metadata size_bytes to 1 (rather than mocking the method on the
|
|
# instance, which doesn't survive dataclasses.replace in OpBufferQueue.pop).
|
|
entry = input.blocks[0]
|
|
input = replace(
|
|
input,
|
|
blocks=[BlockEntry(entry.ref, replace(entry.metadata, size_bytes=1))],
|
|
)
|
|
|
|
o1 = InputDataBuffer(DataContext.get_current(), [input])
|
|
o2 = mock_map_op(o1)
|
|
o3 = mock_map_op(o2)
|
|
|
|
topo = build_streaming_topology(o3, ExecutionOptions(), noop_counter())
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
MagicMock(return_value=ExecutionResources.zero()),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
ray.data.DataContext.get_current()._max_num_blocks_in_streaming_gen_buffer = 1
|
|
ray.data.DataContext.get_current().target_max_block_size = 2
|
|
|
|
resource_manager.update_usages()
|
|
assert resource_manager.get_op_usage(o1).object_store_memory == 0
|
|
assert resource_manager.get_op_usage(o2).object_store_memory == 0
|
|
assert resource_manager.get_op_usage(o3).object_store_memory == 0
|
|
|
|
# Objects in an operator's internal inqueue typically count toward the previous
|
|
# operator's object store memory usage. However, data from an
|
|
# `InputDataBuffer` aren't counted because they were created outside of this
|
|
# execution.
|
|
o2.metrics.on_input_queued(input, input_index=0)
|
|
resource_manager.update_usages()
|
|
assert resource_manager.get_op_usage(o1).object_store_memory == 0
|
|
assert resource_manager.get_op_usage(o2).object_store_memory == 0
|
|
assert resource_manager.get_op_usage(o3).object_store_memory == 0
|
|
|
|
# During no-sample phase, obj_store_mem_pending_task_outputs uses fallback
|
|
# estimate based on target_max_block_size.
|
|
o2.metrics.on_input_dequeued(input, input_index=0)
|
|
o2.metrics.on_task_submitted(0, input)
|
|
resource_manager.update_usages()
|
|
assert resource_manager.get_op_usage(o1).object_store_memory == 0
|
|
# No sample available yet, returns None
|
|
assert o2.metrics.obj_store_mem_pending_task_outputs is None
|
|
op2_usage = resource_manager.get_op_usage(o2).object_store_memory
|
|
# When pending task outputs is None, it's treated as 0
|
|
assert op2_usage == 0
|
|
assert resource_manager.get_op_usage(o3).object_store_memory == 0
|
|
|
|
# When the task finishes, we move the data from the streaming generator to the
|
|
# operator's internal outqueue.
|
|
o2.metrics.on_output_queued(input)
|
|
o2.metrics.on_task_finished(
|
|
0,
|
|
None,
|
|
TaskExecWorkerStats(task_wall_time_s=0.0),
|
|
TaskExecDriverStats(task_output_backpressure_s=0),
|
|
)
|
|
resource_manager.update_usages()
|
|
assert resource_manager.get_op_usage(o1).object_store_memory == 0
|
|
assert resource_manager.get_op_usage(o2).object_store_memory == 1
|
|
assert resource_manager.get_op_usage(o3).object_store_memory == 0
|
|
|
|
o2.metrics.on_output_dequeued(input)
|
|
topo[o2].output_queue.append(input)
|
|
resource_manager.update_usages()
|
|
assert resource_manager.get_op_usage(o1).object_store_memory == 0
|
|
assert resource_manager.get_op_usage(o2).object_store_memory == 1
|
|
assert resource_manager.get_op_usage(o3).object_store_memory == 0
|
|
|
|
# Objects in the current operator's internal inqueue count towards the previous
|
|
# operator's object store memory usage.
|
|
# NOTE: `pop()` returns a copy of the bundle (via `dataclasses.replace`), so we
|
|
# must use the returned reference for subsequent o3 metric calls.
|
|
o3_input = topo[o2].output_queue.pop()
|
|
o3.metrics.on_input_queued(o3_input, input_index=0)
|
|
resource_manager.update_usages()
|
|
assert resource_manager.get_op_usage(o1).object_store_memory == 0
|
|
assert resource_manager.get_op_usage(o2).object_store_memory == 1
|
|
assert resource_manager.get_op_usage(o3).object_store_memory == 0
|
|
|
|
# Task inputs count toward the previous operator's object store memory
|
|
# usage. During no-sample phase, pending task outputs uses fallback estimate.
|
|
o3.metrics.on_input_dequeued(o3_input, input_index=0)
|
|
o3.metrics.on_task_submitted(0, o3_input)
|
|
resource_manager.update_usages()
|
|
assert resource_manager.get_op_usage(o1).object_store_memory == 0
|
|
assert resource_manager.get_op_usage(o2).object_store_memory == 1
|
|
# No sample available yet, returns None
|
|
assert o3.metrics.obj_store_mem_pending_task_outputs is None
|
|
op3_usage = resource_manager.get_op_usage(o3).object_store_memory
|
|
# When pending task outputs is None, it's treated as 0
|
|
assert op3_usage == 0
|
|
|
|
# Task inputs no longer count once the task is finished.
|
|
o3.metrics.on_output_queued(o3_input)
|
|
o3.metrics.on_task_finished(
|
|
0,
|
|
None,
|
|
TaskExecWorkerStats(task_wall_time_s=0.0),
|
|
TaskExecDriverStats(task_output_backpressure_s=0),
|
|
)
|
|
resource_manager.update_usages()
|
|
assert resource_manager.get_op_usage(o1).object_store_memory == 0
|
|
assert resource_manager.get_op_usage(o2).object_store_memory == 0
|
|
assert resource_manager.get_op_usage(o3).object_store_memory == 1
|
|
|
|
def test_object_store_accounting_delegates_to_op(self, restore_data_context):
|
|
"""``ResourceManager`` must dispatch to ``op.estimate_object_store_usage`` so subclasses can override the accounting."""
|
|
# Real upstream so the override op has a valid input dependency.
|
|
input = make_ref_bundles([[x] for x in range(1)])[0]
|
|
upstream = InputDataBuffer(DataContext.get_current(), [input])
|
|
|
|
# Subclass that overrides the accounting to return hard-coded
|
|
# values — bypasses the generic metrics+state computation.
|
|
override = mock_map_op(upstream)
|
|
override.estimate_object_store_usage = lambda state: ObjectStoreUsage(
|
|
internal=42, outputs=100
|
|
)
|
|
|
|
topo = build_streaming_topology(override, ExecutionOptions(), noop_counter())
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
MagicMock(return_value=ExecutionResources.zero()),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
|
|
resource_manager.update_usages()
|
|
|
|
# The override's hard-coded values flow through unchanged into
|
|
# both the per-component dicts and the aggregated op usage.
|
|
assert resource_manager.get_mem_op_internal(override) == 42
|
|
assert resource_manager.get_mem_op_outputs(override) == 100
|
|
assert resource_manager.get_op_usage(override).object_store_memory == 42 + 100
|
|
|
|
def test_get_completed_ops_usage(self, restore_data_context):
|
|
"""Test that _get_completed_ops_usage returns total usage of completed ops."""
|
|
o1 = InputDataBuffer(DataContext.get_current(), [])
|
|
o2 = mock_map_op(o1)
|
|
o3 = LimitOperator(1, o2, DataContext.get_current())
|
|
o4 = mock_map_op(o3)
|
|
o5 = mock_map_op(o4)
|
|
|
|
topo = build_streaming_topology(o5, ExecutionOptions(), noop_counter())
|
|
|
|
o1.mark_execution_finished()
|
|
o2.mark_execution_finished()
|
|
|
|
op_usages = {
|
|
o1: ExecutionResources.zero(),
|
|
o2: ExecutionResources(cpu=2, object_store_memory=50),
|
|
o3: ExecutionResources(cpu=1, object_store_memory=25),
|
|
o4: ExecutionResources.zero(),
|
|
o5: ExecutionResources.zero(),
|
|
}
|
|
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
MagicMock(),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
resource_manager.get_op_usage = MagicMock(side_effect=lambda op: op_usages[op])
|
|
|
|
# o2 is completed and o3 is downstream ineligible (LimitOperator)
|
|
# Total usage should be o2 + o3
|
|
completed_ops_usage = resource_manager._get_completed_ops_usage()
|
|
assert completed_ops_usage == ExecutionResources(cpu=3, object_store_memory=75)
|
|
|
|
def test_get_completed_ops_usage_complex_graph(self, restore_data_context):
|
|
"""
|
|
o1 (InputDataBuffer)
|
|
|
|
|
v
|
|
o2 (MapOperator, completed)
|
|
|
|
|
v
|
|
o3 (LimitOperator)
|
|
|
|
|
v o4 (InputDataBuffer)
|
|
| |
|
|
| v
|
|
| o5 (MapOperator, completed)
|
|
| |
|
|
v v
|
|
o6 (UnionOperator) <--
|
|
|
|
|
v
|
|
o8 (JoinOperator) <-- o7 (InputDataBuffer, completed)
|
|
"""
|
|
o1 = InputDataBuffer(DataContext.get_current(), [])
|
|
o2 = mock_map_op(o1)
|
|
o3 = LimitOperator(1, o2, DataContext.get_current())
|
|
o4 = InputDataBuffer(DataContext.get_current(), [])
|
|
o5 = mock_map_op(o4)
|
|
o6 = mock_union_op([o3, o5])
|
|
o7 = InputDataBuffer(DataContext.get_current(), [])
|
|
o8 = mock_join_op(o7, o6)
|
|
|
|
topo = build_streaming_topology(o8, ExecutionOptions(), noop_counter())
|
|
|
|
o1.mark_execution_finished()
|
|
o2.mark_execution_finished()
|
|
o4.mark_execution_finished()
|
|
o5.mark_execution_finished()
|
|
o7.mark_execution_finished()
|
|
|
|
op_usages = {
|
|
o1: ExecutionResources.zero(),
|
|
o2: ExecutionResources(cpu=2, object_store_memory=150),
|
|
o3: ExecutionResources(cpu=2, object_store_memory=50),
|
|
o4: ExecutionResources.zero(),
|
|
o5: ExecutionResources(cpu=3, object_store_memory=100),
|
|
o6: ExecutionResources.zero(),
|
|
o7: ExecutionResources(cpu=1, object_store_memory=100),
|
|
o8: ExecutionResources.zero(),
|
|
}
|
|
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
MagicMock(),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
resource_manager.get_op_usage = MagicMock(side_effect=lambda op: op_usages[op])
|
|
|
|
# Completed ops: o2, o5, o7
|
|
# Downstream ineligible: o3 (LimitOperator after o2)
|
|
# Total usage should be o2 + o3 + o5 + o7
|
|
completed_ops_usage = resource_manager._get_completed_ops_usage()
|
|
|
|
assert completed_ops_usage == ExecutionResources(cpu=8, object_store_memory=400)
|
|
|
|
def test_external_consumer_bytes_attributed_to_terminal_operator(
|
|
self, restore_data_context
|
|
):
|
|
"""External consumer bytes (e.g., iterator prefetch buffers) are charged
|
|
to the terminal operator's object store usage, not as a global deduction."""
|
|
cluster_resources = ExecutionResources(cpu=10, gpu=0, object_store_memory=1000)
|
|
|
|
o1 = InputDataBuffer(DataContext.get_current(), [])
|
|
o2 = mock_map_op(o1)
|
|
o3 = mock_map_op(o2)
|
|
|
|
topo = build_streaming_topology(o3, ExecutionOptions(), noop_counter())
|
|
|
|
o1.mark_execution_finished()
|
|
o2.mark_execution_finished()
|
|
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
lambda: cluster_resources,
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
|
|
for op in [o1, o2, o3]:
|
|
op.current_logical_usage = MagicMock(return_value=ExecutionResources.zero())
|
|
op.running_logical_usage = MagicMock(return_value=ExecutionResources.zero())
|
|
op.pending_logical_usage = MagicMock(return_value=ExecutionResources.zero())
|
|
|
|
assert resource_manager._op_resource_allocator is not None
|
|
|
|
resource_manager.update_usages()
|
|
baseline_terminal = resource_manager.get_op_usage(o3).object_store_memory
|
|
baseline_upstream = resource_manager.get_op_usage(o2).object_store_memory
|
|
|
|
def _available_pool_object_store():
|
|
return (
|
|
resource_manager.get_global_limits()
|
|
.subtract(resource_manager._get_completed_ops_usage())
|
|
.max(ExecutionResources.zero())
|
|
.object_store_memory
|
|
)
|
|
|
|
pool_before = _available_pool_object_store()
|
|
|
|
resource_manager.set_external_consumer_bytes(200)
|
|
resource_manager.update_usages()
|
|
|
|
assert (
|
|
resource_manager.get_op_usage(o3).object_store_memory
|
|
== baseline_terminal + 200
|
|
)
|
|
assert (
|
|
resource_manager.get_op_usage(o2).object_store_memory == baseline_upstream
|
|
)
|
|
assert _available_pool_object_store() == pool_before
|
|
|
|
resource_manager.set_external_consumer_bytes(0)
|
|
resource_manager.update_usages()
|
|
assert (
|
|
resource_manager.get_op_usage(o3).object_store_memory == baseline_terminal
|
|
)
|
|
|
|
# Very large external bytes: terminal usage reflects them; update still succeeds.
|
|
resource_manager.set_external_consumer_bytes(999999)
|
|
resource_manager.update_usages()
|
|
assert (
|
|
resource_manager.get_op_usage(o3).object_store_memory
|
|
== baseline_terminal + 999999
|
|
)
|
|
|
|
def test_set_external_consumer_bytes_rejects_negative(self, restore_data_context):
|
|
resource_manager = _resource_manager_for_limits_only_test(
|
|
ExecutionOptions(),
|
|
MagicMock(return_value=ExecutionResources.zero()),
|
|
)
|
|
with pytest.raises(AssertionError):
|
|
resource_manager.set_external_consumer_bytes(-1)
|
|
|
|
def test_external_consumer_bytes_input_data_buffer_sink(self, restore_data_context):
|
|
"""When the execute DAG is only an InputDataBuffer, prefetch bytes still
|
|
attach to that terminal sink instead of being dropped by the
|
|
InputDataBuffer early return."""
|
|
buf = InputDataBuffer(DataContext.get_current(), [])
|
|
topo = build_streaming_topology(buf, ExecutionOptions(), noop_counter())
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
lambda: ExecutionResources(cpu=10, gpu=0, object_store_memory=1000),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
buf.current_logical_usage = MagicMock(return_value=ExecutionResources.zero())
|
|
buf.running_logical_usage = MagicMock(return_value=ExecutionResources.zero())
|
|
buf.pending_logical_usage = MagicMock(return_value=ExecutionResources.zero())
|
|
|
|
resource_manager.update_usages()
|
|
assert resource_manager.get_op_usage(buf).object_store_memory == 0
|
|
|
|
resource_manager.set_external_consumer_bytes(150)
|
|
resource_manager.update_usages()
|
|
assert resource_manager.get_op_usage(buf).object_store_memory == 150
|
|
|
|
def test_external_consumer_bytes_surfaced_in_op_usage_str(
|
|
self, restore_data_context
|
|
):
|
|
"""The terminal operator's verbose usage string should include
|
|
external_consumer=... when an external consumer is registered, so users
|
|
can see how much of the operator's object-store memory is held by a
|
|
downstream iterator vs. the operator's own queues."""
|
|
cluster_resources = ExecutionResources(cpu=10, gpu=0, object_store_memory=1000)
|
|
|
|
o1 = InputDataBuffer(DataContext.get_current(), [])
|
|
o2 = mock_map_op(o1)
|
|
o3 = mock_map_op(o2)
|
|
|
|
topo = build_streaming_topology(o3, ExecutionOptions(), noop_counter())
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
lambda: cluster_resources,
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
|
|
for op in [o1, o2, o3]:
|
|
op.current_logical_usage = MagicMock(return_value=ExecutionResources.zero())
|
|
op.running_logical_usage = MagicMock(return_value=ExecutionResources.zero())
|
|
op.pending_logical_usage = MagicMock(return_value=ExecutionResources.zero())
|
|
|
|
resource_manager.update_usages()
|
|
|
|
# No external consumer yet: nothing extra in the usage string.
|
|
terminal_str = resource_manager.get_op_usage_str(o3, verbose=True)
|
|
upstream_str = resource_manager.get_op_usage_str(o2, verbose=True)
|
|
assert "external_consumer=" not in terminal_str
|
|
assert "external_consumer=" not in upstream_str
|
|
|
|
# Register an external consumer. Only the terminal operator's string
|
|
# should pick up `external_consumer=...`.
|
|
resource_manager.set_external_consumer_bytes(200)
|
|
resource_manager.update_usages()
|
|
terminal_str = resource_manager.get_op_usage_str(o3, verbose=True)
|
|
upstream_str = resource_manager.get_op_usage_str(o2, verbose=True)
|
|
assert "external_consumer=200.0B" in terminal_str
|
|
assert "external_consumer=" not in upstream_str
|
|
|
|
# The field is inside the existing `(in=...,out=...)` parenthetical.
|
|
assert ",external_consumer=" in terminal_str
|
|
|
|
# Non-verbose output omits the field (existing format unchanged).
|
|
terminal_str_brief = resource_manager.get_op_usage_str(o3, verbose=False)
|
|
assert "external_consumer=" not in terminal_str_brief
|
|
|
|
def test_topology_rejects_multiple_terminal_operators(self, restore_data_context):
|
|
ctx = DataContext.get_current()
|
|
a = PhysicalOperator("a", [], ctx)
|
|
b = PhysicalOperator("b", [], ctx)
|
|
topology = {a: MagicMock(), b: MagicMock()}
|
|
with pytest.raises(ValueError, match="Expected exactly one terminal operator"):
|
|
ResourceManager(
|
|
topology,
|
|
ExecutionOptions(),
|
|
MagicMock(return_value=ExecutionResources.zero()),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
|
|
def test_topology_rejects_empty_topology(self, restore_data_context):
|
|
with pytest.raises(ValueError, match="topology must be non-empty"):
|
|
ResourceManager(
|
|
{},
|
|
ExecutionOptions(),
|
|
MagicMock(return_value=ExecutionResources.zero()),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
|
|
def test_topology_rejects_no_terminal_operator(self, restore_data_context):
|
|
# Every op has a downstream in this dict, so there should be no operator with empty
|
|
# output_dependencies (e.g. a 2-node cycle). Real streaming DAGs from
|
|
# build_streaming_topology always have a unique sink.
|
|
a = MagicMock(spec=PhysicalOperator)
|
|
b = MagicMock(spec=PhysicalOperator)
|
|
a.output_dependencies = [b]
|
|
b.output_dependencies = [a]
|
|
topology = {a: MagicMock(), b: MagicMock()}
|
|
with pytest.raises(ValueError, match="No terminal operator found"):
|
|
ResourceManager(
|
|
topology,
|
|
ExecutionOptions(),
|
|
MagicMock(return_value=ExecutionResources.zero()),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
|
|
def test_is_blocking_materializing_op(self, restore_data_context):
|
|
"""Test _is_blocking_materializing_op correctly identifies blocking materializing ops.
|
|
|
|
Cases tested:
|
|
1. Operator itself is a blocking materializing op (AllToAllOperator) -> True
|
|
2. Operator has downstream ineligible blocking materializing op -> True
|
|
3. Operator with no downstream blocking materializing ops -> False
|
|
|
|
Note: AllToAllOperator.throttling_disabled() returns True, making it
|
|
ineligible for resource allocation. This means shuffle operators are
|
|
always in the "downstream ineligible" chain from eligible operators.
|
|
"""
|
|
# Build pipeline: o1 -> o2 -> o3 (limit) -> o4 (shuffle) -> o5
|
|
o1 = InputDataBuffer(DataContext.get_current(), [])
|
|
o2 = mock_map_op(o1, name="Map1")
|
|
o3 = LimitOperator(1, o2, DataContext.get_current())
|
|
o4 = mock_all_to_all_op(o3, name="Sort")
|
|
o5 = mock_map_op(o4, name="Map2")
|
|
|
|
topo = build_streaming_topology(o5, ExecutionOptions(), noop_counter())
|
|
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
MagicMock(),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
|
|
# Case 1: Shuffle operator itself is blocking materializing
|
|
assert resource_manager._is_blocking_materializing_op(o4) is True
|
|
|
|
# Case 2: Map operator before shuffle (o2) should return True because
|
|
# its downstream ineligible chain includes:
|
|
# - o3 (LimitOperator - ineligible, not in eligible types)
|
|
# - o4 (AllToAllOperator - ineligible because throttling_disabled=True)
|
|
# Since o4 is a blocking materializing op, the check returns True
|
|
assert resource_manager._is_blocking_materializing_op(o2) is True
|
|
|
|
# o3 (LimitOperator) also returns True because its downstream ineligible
|
|
# chain includes o4 (shuffle)
|
|
assert resource_manager._is_blocking_materializing_op(o3) is True
|
|
|
|
# Case 3: o5 (Map after shuffle) has no downstream ops -> False
|
|
assert resource_manager._is_blocking_materializing_op(o5) is False
|
|
|
|
# Case 4: Extend pipeline with ops that have no blocking materializing downstream
|
|
# o5 -> o6 (limit) -> o7
|
|
o6 = LimitOperator(1, o5, DataContext.get_current())
|
|
o7 = mock_map_op(o6, name="Map3")
|
|
|
|
topo2 = build_streaming_topology(o7, ExecutionOptions(), noop_counter())
|
|
resource_manager2 = ResourceManager(
|
|
topo2,
|
|
ExecutionOptions(),
|
|
MagicMock(),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
|
|
# o5's downstream (o6, o7) has no blocking materializing ops
|
|
assert resource_manager2._is_blocking_materializing_op(o5) is False
|
|
assert resource_manager2._is_blocking_materializing_op(o7) is False
|
|
|
|
def test_memory_limit_blocks_task_submission(self, restore_data_context):
|
|
"""Test that tasks are blocked when memory limit is exceeded."""
|
|
# Cluster has 1000 bytes of memory
|
|
cluster_resources = ExecutionResources(cpu=1, gpu=0, memory=1000)
|
|
|
|
# Request 2000 bytes memory
|
|
o1 = InputDataBuffer(DataContext.get_current(), [])
|
|
o2 = mock_map_op(
|
|
o1,
|
|
ray_remote_args={"num_cpus": 1, "memory": 2000},
|
|
name="HighMemoryTask",
|
|
)
|
|
|
|
topo = build_streaming_topology(o2, ExecutionOptions(), noop_counter())
|
|
options = ExecutionOptions()
|
|
|
|
resource_manager = ResourceManager(
|
|
topology=topo,
|
|
options=options,
|
|
get_total_resources=lambda: cluster_resources,
|
|
data_context=DataContext.get_current(),
|
|
block_ref_counter=BlockRefCounter(
|
|
add_object_out_of_scope_callback=lambda *_: True
|
|
),
|
|
)
|
|
resource_manager.update_usages()
|
|
|
|
# Task cannot be submitted because it exceeds memory limit
|
|
allocator = create_resource_allocator(
|
|
resource_manager, DataContext.get_current()
|
|
)
|
|
assert allocator is not None
|
|
allocator.update_budgets(limits=resource_manager.get_global_limits())
|
|
can_submit = allocator.can_submit_new_task(o2)
|
|
assert (
|
|
not can_submit
|
|
), "Task should be blocked: requires 2000 bytes but only 1000 bytes memory available"
|
|
|
|
|
|
class TestOutputBackpressureGuard:
|
|
"""Tests for OutputBackpressureGuard.should_unblock."""
|
|
|
|
def test_unblock_backpressure_terminal_operator(self, restore_data_context):
|
|
"""Terminal operator (no downstream eligible ops) with no external
|
|
consumer should always unblock (e.g., write pipeline)."""
|
|
o1 = InputDataBuffer(DataContext.get_current(), [])
|
|
o2 = mock_map_op(o1)
|
|
o3 = LimitOperator(1, o2, DataContext.get_current())
|
|
|
|
topo = build_streaming_topology(o3, ExecutionOptions(), noop_counter())
|
|
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
MagicMock(),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
guard = OutputBackpressureGuard(topo, resource_manager)
|
|
|
|
# o2 is terminal (no downstream eligible ops beyond it) and no external
|
|
# consumer — should unblock (e.g., write pipeline).
|
|
assert guard.should_unblock(o2) is True
|
|
|
|
# Add o4 operator - o2 is no longer terminal
|
|
o4 = mock_map_op(o3)
|
|
|
|
topo = build_streaming_topology(o4, ExecutionOptions(), noop_counter())
|
|
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
MagicMock(),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
guard = OutputBackpressureGuard(topo, resource_manager)
|
|
|
|
# Mock downstream (o4) having active tasks and input blocks (ie unblocking
|
|
# conditions not met)
|
|
o4.num_active_tasks = MagicMock(return_value=1)
|
|
guard._idle_detector.detect_idle = MagicMock(return_value=False)
|
|
|
|
# o2 is not terminal anymore, falls back to idle detector which returns False
|
|
assert guard.should_unblock(o2) is False
|
|
|
|
def test_no_unblock_backpressure_terminal_with_external_consumer(
|
|
self, restore_data_context
|
|
):
|
|
"""Terminal operator with an external consumer should only unblock
|
|
when consumers are starving (blocked waiting for output)."""
|
|
o1 = InputDataBuffer(DataContext.get_current(), [])
|
|
o2 = mock_map_op(o1)
|
|
o3 = LimitOperator(1, o2, DataContext.get_current())
|
|
|
|
topo = build_streaming_topology(o3, ExecutionOptions(), noop_counter())
|
|
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
MagicMock(),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
guard = OutputBackpressureGuard(topo, resource_manager)
|
|
|
|
# Register an external consumer (e.g., iter_batches or streaming_split).
|
|
resource_manager.set_external_consumer_bytes(0)
|
|
|
|
dag_output_state = topo[o3]
|
|
|
|
# No consumers waiting — should NOT unblock (prevents pileup).
|
|
dag_output_state._num_waiting_consumers = 0
|
|
assert guard.should_unblock(o2) is False
|
|
|
|
# Simulate a consumer blocked in get_output_blocking (starving).
|
|
# The output node is o3 (LimitOperator), which tracks waiting consumers.
|
|
dag_output_state._num_waiting_consumers = 1
|
|
assert guard.should_unblock(o2) is True
|
|
|
|
# Consumer gets data and stops waiting — should NOT unblock again.
|
|
dag_output_state._num_waiting_consumers = 0
|
|
assert guard.should_unblock(o2) is False
|
|
|
|
def test_unblock_backpressure_downstream_idle(self, restore_data_context):
|
|
"""Unblock when downstream is idle (no active tasks) to maintain liveness."""
|
|
o1 = InputDataBuffer(DataContext.get_current(), [])
|
|
o2 = mock_map_op(o1)
|
|
o3 = mock_map_op(o2)
|
|
|
|
topo = build_streaming_topology(o3, ExecutionOptions(), noop_counter())
|
|
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
MagicMock(),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
guard = OutputBackpressureGuard(topo, resource_manager)
|
|
o3.num_active_tasks = MagicMock(return_value=0)
|
|
|
|
# Case 1: Downstream cannot submit (resource constrained) - unblock to free resources
|
|
resource_manager.op_resource_allocator.can_submit_new_task = MagicMock(
|
|
return_value=False
|
|
)
|
|
assert guard.should_unblock(o2) is True
|
|
|
|
# Case 2: Downstream can submit but has no input blocks - unblock to produce data
|
|
resource_manager.op_resource_allocator.can_submit_new_task = MagicMock(
|
|
return_value=True
|
|
)
|
|
topo[o3].total_enqueued_input_blocks = MagicMock(return_value=0)
|
|
assert guard.should_unblock(o2) is True
|
|
|
|
def test_unblock_backpressure_fallback_to_idle_detector(self, restore_data_context):
|
|
"""When unblock conditions not met, falls back to idle detector result."""
|
|
o1 = InputDataBuffer(DataContext.get_current(), [])
|
|
o2 = mock_map_op(o1)
|
|
o3 = mock_map_op(o2)
|
|
|
|
topo = build_streaming_topology(o3, ExecutionOptions(), noop_counter())
|
|
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
MagicMock(),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
guard = OutputBackpressureGuard(topo, resource_manager)
|
|
|
|
# Case: Downstream has active tasks - falls back to idle detector
|
|
o3.num_active_tasks = MagicMock(return_value=2)
|
|
guard._idle_detector.detect_idle = MagicMock(return_value=False)
|
|
assert guard.should_unblock(o2) is False
|
|
|
|
# Case: Idle detector returns True - should unblock
|
|
guard._idle_detector.detect_idle = MagicMock(return_value=True)
|
|
assert guard.should_unblock(o2) is True
|
|
|
|
# Case: Downstream has no active tasks but has input blocks - falls back to idle detector
|
|
resource_manager.op_resource_allocator.can_submit_new_task = MagicMock(
|
|
return_value=True
|
|
)
|
|
o3.num_active_tasks = MagicMock(return_value=0)
|
|
topo[o3].total_enqueued_input_blocks = MagicMock(return_value=5)
|
|
guard._idle_detector.detect_idle = MagicMock(return_value=False)
|
|
assert guard.should_unblock(o2) is False
|
|
|
|
def test_unblock_when_resource_allocator_disabled(self, restore_data_context):
|
|
"""When the op resource allocator is disabled, the guard treats
|
|
downstream as schedulable (no budget to consult), so
|
|
"downstream resource constrained" case never fires, but the other
|
|
liveness conditions still do.
|
|
"""
|
|
# Disable resource allocator
|
|
DataContext.get_current().op_resource_reservation_enabled = False
|
|
|
|
o1 = InputDataBuffer(DataContext.get_current(), [])
|
|
o2 = mock_map_op(o1)
|
|
o3 = mock_map_op(o2)
|
|
|
|
topo = build_streaming_topology(o3, ExecutionOptions(), noop_counter())
|
|
|
|
resource_manager = ResourceManager(
|
|
topo,
|
|
ExecutionOptions(),
|
|
MagicMock(),
|
|
DataContext.get_current(),
|
|
BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
|
|
)
|
|
assert not resource_manager.op_resource_allocator_enabled()
|
|
|
|
guard = OutputBackpressureGuard(topo, resource_manager)
|
|
o3.num_active_tasks = MagicMock(return_value=0)
|
|
|
|
# "Downstream idle with empty input queue" case should fire and unblock.
|
|
topo[o3].total_enqueued_input_blocks = MagicMock(return_value=0)
|
|
assert guard.should_unblock(o2) is True
|
|
|
|
|
|
class TestIdleDetector:
|
|
"""Tests for IdleDetector."""
|
|
|
|
def test_idle_detector(self, restore_data_context):
|
|
"""Test IdleDetector behavior through its public interface."""
|
|
idle_detector = IdleDetector()
|
|
op = MagicMock()
|
|
op.metrics.num_task_outputs_generated = 0
|
|
|
|
with freeze_time() as frozen:
|
|
# First call initializes state, returns False
|
|
assert idle_detector.detect_idle(op) is False
|
|
|
|
# Call within interval returns False (rate limited)
|
|
frozen.tick(timedelta(seconds=idle_detector.DETECTION_INTERVAL_S - 1))
|
|
assert idle_detector.detect_idle(op) is False
|
|
|
|
# Call after interval with no output returns True (idle)
|
|
frozen.tick(timedelta(seconds=2))
|
|
assert idle_detector.detect_idle(op) is True
|
|
|
|
# Operator produces output - next detection returns False (active)
|
|
op.metrics.num_task_outputs_generated = 5
|
|
assert idle_detector.detect_idle(op) is False
|
|
|
|
# After output, wait for interval with no new output - returns True (idle again)
|
|
frozen.tick(timedelta(seconds=idle_detector.DETECTION_INTERVAL_S + 1))
|
|
assert idle_detector.detect_idle(op) is True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|