import operator import pytest import ray from ray.data._internal.actor_autoscaler import ActorPoolScalingRequest from ray.data._internal.compute import ActorPoolStrategy, TaskPoolStrategy from ray.data._internal.execution.interfaces import ExecutionOptions, ExecutionResources from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer 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.output_splitter import OutputSplitter from ray.data._internal.execution.util import make_ref_bundles from ray.data.context import DataContext from ray.data.tests.conftest import * # noqa from ray.data.tests.conftest import noop_counter from ray.data.tests.test_operators import _mul2_map_data_prcessor from ray.data.tests.util import run_op_tasks_sync SMALL_STR = "hello" * 120 def test_execution_resources(ray_start_10_cpus_shared): """Unit test for ExecutionResources.""" r1 = ExecutionResources() r2 = ExecutionResources(cpu=1) r3 = ExecutionResources(gpu=1) r4 = ExecutionResources(cpu=1, gpu=1, object_store_memory=100 * 1024 * 1024) r5 = ExecutionResources( cpu=1, gpu=1, object_store_memory=1024 * 1024 * 1024, memory=64 * 1024 * 1024 ) unlimited = ExecutionResources.for_limits() # Test __eq__. assert r1 == ExecutionResources(0, 0, 0, 0) assert r2 == ExecutionResources(1, 0, 0, 0) assert r3 == ExecutionResources(0, 1, 0, 0) assert r4 == ExecutionResources(1, 1, 100 * 1024 * 1024, 0) assert r5 == ExecutionResources(1, 1, 1024 * 1024 * 1024, 64 * 1024 * 1024) assert unlimited == ExecutionResources( float("inf"), float("inf"), float("inf"), float("inf") ) # Test __repr__. assert ( repr(r1) == "ExecutionResources(cpu=0.0, gpu=0.0, object_store_memory=0.0B, memory=0.0B)" ) assert ( repr(r2) == "ExecutionResources(cpu=1, gpu=0.0, object_store_memory=0.0B, memory=0.0B)" ) assert ( repr(r3) == "ExecutionResources(cpu=0.0, gpu=1, object_store_memory=0.0B, memory=0.0B)" ) assert ( repr(r4) == "ExecutionResources(cpu=1, gpu=1, object_store_memory=100.0MiB, memory=0.0B)" ) assert ( repr(r5) == "ExecutionResources(cpu=1, gpu=1, object_store_memory=1.0GiB, memory=64.0MiB)" ) assert ( repr(unlimited) == "ExecutionResources(cpu=inf, gpu=inf, object_store_memory=inf, memory=inf)" ) # Test object_store_memory_str. assert r3.object_store_memory_str() == "0.0B" assert r4.object_store_memory_str() == "100.0MiB" assert r5.object_store_memory_str() == "1.0GiB" assert unlimited.object_store_memory_str() == "inf" # Test add. assert r1.add(r1) == r1 assert r1.add(r2) == r2 assert r2.add(r2) == ExecutionResources(cpu=2) assert r2.add(r3) == ExecutionResources(cpu=1, gpu=1) assert r4.add(r4) == ExecutionResources( cpu=2, gpu=2, object_store_memory=200 * 1024 * 1024 ) assert r5.add(r5) == ExecutionResources( cpu=2, gpu=2, object_store_memory=2 * 1024 * 1024 * 1024, memory=128 * 1024 * 1024, ) # Test subtract. assert r2.subtract(r1) == r2 assert r2.subtract(r2) == r1 assert r4.subtract(r2) == ExecutionResources( gpu=1, object_store_memory=100 * 1024 * 1024 ) assert r5.subtract(r4) == ExecutionResources( object_store_memory=924 * 1024 * 1024, memory=64 * 1024 * 1024 ) assert r4.subtract(r5) == ExecutionResources( object_store_memory=-924 * 1024 * 1024, memory=-64 * 1024 * 1024 ) assert r5.subtract(r5) == r1 # Test scale. assert r1.scale(2) == r1 assert r2.scale(2) == ExecutionResources(cpu=2) assert r3.scale(0.5) == ExecutionResources(gpu=0.5) assert r4.scale(0.5) == ExecutionResources( cpu=0.5, gpu=0.5, object_store_memory=50 * 1024 * 1024 ) assert r5.scale(0.5) == ExecutionResources( cpu=0.5, gpu=0.5, object_store_memory=512 * 1024 * 1024, memory=32 * 1024 * 1024, ) assert r5.scale(0) == r1 assert unlimited.scale(0) == r1 # Test limit. for r in [r1, r2, r3, r4, r5]: assert r.satisfies_limit(r) assert r.satisfies_limit(unlimited) assert r2.satisfies_limit(ExecutionResources.for_limits(gpu=1)) assert r3.satisfies_limit(ExecutionResources.for_limits(cpu=1)) assert r4.satisfies_limit(r5) assert not r5.satisfies_limit( ExecutionResources.for_limits(memory=63 * 1024 * 1024) ) assert r5.satisfies_limit(ExecutionResources.for_limits(memory=64 * 1024 * 1024)) assert not r5.satisfies_limit(r4) def test_resource_canonicalization_with_no_ray_remote_args(): input_op = InputDataBuffer( DataContext.get_current(), make_ref_bundles([[i] for i in range(1)]) ) op = MapOperator.create( _mul2_map_data_prcessor, input_op=input_op, data_context=DataContext.get_current(), ray_remote_args=None, ) assert op.incremental_resource_usage().cpu == 1 def test_execution_options_resource_limit(): """Test ExecutionOptions.resource_limit.""" # Test that the default resource_limits should be inf. options = ExecutionOptions() assert options.resource_limits.cpu == float("inf") assert options.resource_limits.gpu == float("inf") assert options.resource_limits.object_store_memory == float("inf") # Test when passing in the resource_limits parameter, missing # fields should be set to inf. options = ExecutionOptions(resource_limits=ExecutionResources(cpu=1)) assert options.resource_limits.cpu == 1 assert options.resource_limits.gpu == float("inf") assert options.resource_limits.object_store_memory == float("inf") # Test when modifying the resource_limits attribute, # missing fields should be set to inf. options.resource_limits = ExecutionResources(object_store_memory=100) assert options.resource_limits.cpu == float("inf") assert options.resource_limits.gpu == float("inf") assert options.resource_limits.object_store_memory == 100 def test_scheduling_strategy_overrides(ray_start_10_cpus_shared, restore_data_context): input_op = InputDataBuffer( DataContext.get_current(), make_ref_bundles([[i] for i in range(100)]) ) op = MapOperator.create( _mul2_map_data_prcessor, input_op=input_op, data_context=DataContext.get_current(), name="TestMapper", compute_strategy=TaskPoolStrategy(), ray_remote_args={"num_gpus": 2, "scheduling_strategy": "DEFAULT"}, ) assert op._ray_remote_args == {"num_gpus": 2, "scheduling_strategy": "DEFAULT"} ray.data.DataContext.get_current().scheduling_strategy = "DEFAULT" op = MapOperator.create( _mul2_map_data_prcessor, input_op=input_op, data_context=DataContext.get_current(), name="TestMapper", compute_strategy=TaskPoolStrategy(), ray_remote_args={"num_gpus": 2}, ) assert op._ray_remote_args == {"num_gpus": 2} def test_task_pool_resource_reporting(ray_start_10_cpus_shared): ctx = ray.data.DataContext.get_current() ctx._max_num_blocks_in_streaming_gen_buffer = 1 input_op = InputDataBuffer( DataContext.get_current(), make_ref_bundles([[SMALL_STR] for i in range(100)]) ) op = MapOperator.create( _mul2_map_data_prcessor, data_context=DataContext.get_current(), input_op=input_op, name="TestMapper", compute_strategy=TaskPoolStrategy(), ) op.start(ExecutionOptions(), noop_counter()) assert op.current_logical_usage() == ExecutionResources(cpu=0, gpu=0, memory=0) assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 assert op.metrics.obj_store_mem_pending_task_inputs == 0 # No tasks running yet, so pending task outputs is None. assert op.metrics.obj_store_mem_pending_task_outputs is None op.add_input(input_op.get_next(), 0) op.add_input(input_op.get_next(), 0) assert op.current_logical_usage() == ExecutionResources(cpu=2, gpu=0, memory=0) assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 assert op.metrics.obj_store_mem_pending_task_inputs == pytest.approx(1600, rel=0.5) # No sample available yet, so pending task outputs is None. assert op.metrics.obj_store_mem_pending_task_outputs is None run_op_tasks_sync(op) assert op.current_logical_usage() == ExecutionResources(cpu=0, gpu=0, memory=0) assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == pytest.approx(3200, rel=0.5) assert op.metrics.obj_store_mem_pending_task_inputs == 0 assert op.metrics.obj_store_mem_pending_task_outputs == 0 def test_task_pool_resource_reporting_with_dynamic_remote_args( ray_start_10_cpus_shared, ): """Test that current_logical_usage reflects dynamic resources from ray_remote_args_fn, not just the statically defined ray_remote_args.""" input_op = InputDataBuffer( DataContext.get_current(), make_ref_bundles([[SMALL_STR] for i in range(100)]) ) # ray_remote_args set 1 CPU, but ray_remote_args_fn overrides memory to 500 op = MapOperator.create( _mul2_map_data_prcessor, data_context=DataContext.get_current(), input_op=input_op, name="TestMapper", compute_strategy=TaskPoolStrategy(), ray_remote_args={"num_cpus": 1}, ray_remote_args_fn=lambda: {"memory": 500}, ) op.start(ExecutionOptions(), noop_counter()) assert op.current_logical_usage() == ExecutionResources(cpu=0, gpu=0, memory=0) op.add_input(input_op.get_next(), 0) op.add_input(input_op.get_next(), 0) # Should reflect actual dynamic resources: 2 tasks * (1 cpu, 500 memory) assert op.current_logical_usage() == ExecutionResources(cpu=2, gpu=0, memory=1000) run_op_tasks_sync(op) assert op.current_logical_usage() == ExecutionResources(cpu=0, gpu=0, memory=0) def test_task_pool_resource_reporting_with_bundling(ray_start_10_cpus_shared): ctx = ray.data.DataContext.get_current() ctx._max_num_blocks_in_streaming_gen_buffer = 1 input_op = InputDataBuffer( DataContext.get_current(), make_ref_bundles([[SMALL_STR] for i in range(100)]) ) op = MapOperator.create( _mul2_map_data_prcessor, input_op=input_op, data_context=DataContext.get_current(), name="TestMapper", compute_strategy=TaskPoolStrategy(), min_rows_per_bundle=3, ) op.start(ExecutionOptions(), noop_counter()) assert op.current_logical_usage() == ExecutionResources(cpu=0, gpu=0, memory=0) assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 assert op.metrics.obj_store_mem_pending_task_inputs == 0 # No tasks running yet, so pending task outputs is None. assert op.metrics.obj_store_mem_pending_task_outputs is None op.add_input(input_op.get_next(), 0) # No tasks submitted yet due to bundling. assert op.current_logical_usage() == ExecutionResources(cpu=0, gpu=0, memory=0) assert op.metrics.obj_store_mem_internal_inqueue == pytest.approx(800, rel=0.5) assert op.metrics.obj_store_mem_internal_outqueue == 0 assert op.metrics.obj_store_mem_pending_task_inputs == 0 # No tasks running yet, so pending task outputs is None. assert op.metrics.obj_store_mem_pending_task_outputs is None op.add_input(input_op.get_next(), 0) # No tasks submitted yet due to bundling. assert op.current_logical_usage() == ExecutionResources(cpu=0, gpu=0, memory=0) assert op.metrics.obj_store_mem_internal_inqueue == pytest.approx(1600, rel=0.5) assert op.metrics.obj_store_mem_internal_outqueue == 0 assert op.metrics.obj_store_mem_pending_task_inputs == 0 # No tasks running yet, so pending task outputs is None. assert op.metrics.obj_store_mem_pending_task_outputs is None op.add_input(input_op.get_next(), 0) # Task has now been submitted since we've met the minimum bundle size. assert op.current_logical_usage() == ExecutionResources(cpu=1, gpu=0, memory=0) assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 assert op.metrics.obj_store_mem_pending_task_inputs == pytest.approx(2400, rel=0.5) # No sample available yet, so pending task outputs is None. assert op.metrics.obj_store_mem_pending_task_outputs is None def test_actor_pool_scheduling(ray_start_10_cpus_shared, restore_data_context): # TODO move to test_actor_pool_map_operator.py ctx = ray.data.DataContext.get_current() ctx._max_num_blocks_in_streaming_gen_buffer = 1 # Block AP until all actors have fully started up ctx.wait_for_min_actors_s = 60 input_op = InputDataBuffer( DataContext.get_current(), make_ref_bundles([[SMALL_STR] for i in range(100)]) ) op = MapOperator.create( _mul2_map_data_prcessor, min_rows_per_bundle=None, input_op=input_op, data_context=DataContext.get_current(), name="TestMapper", compute_strategy=ActorPoolStrategy( min_size=2, max_size=10, max_tasks_in_flight_per_actor=2 ), ) # NOTE: This is blocking, until actors are fully started up op.start(ExecutionOptions(), noop_counter()) min_resource_usage, _ = op.min_max_resource_requirements() assert min_resource_usage == ExecutionResources(cpu=2, gpu=0, object_store_memory=0) # `incremental_resource_usage` should always report 0 CPU and GPU, as # it doesn't consider scaling-up. assert op.incremental_resource_usage() == ExecutionResources( cpu=0, gpu=0, object_store_memory=0 ) assert op.current_logical_usage() == ExecutionResources(cpu=2, gpu=0, memory=0) assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 assert op.metrics.obj_store_mem_pending_task_inputs == 0 # assert op.metrics.obj_store_mem_pending_task_outputs == 0 # NOTE: Until actors start up, we should not be adding the inputs to the # operator to avoid queuing up inside of it assert not op.can_add_input() # Finalize operator initialization sequence and make it schedulable run_op_tasks_sync(op, only_existing=True) # Add inputs. for i in range(4): assert op.incremental_resource_usage() == ExecutionResources( cpu=0, gpu=0, object_store_memory=0 ) # Should be able to add inputs now assert op.can_add_input() op.add_input(input_op.get_next(), 0) assert op.current_logical_usage() == ExecutionResources(cpu=2, gpu=0, memory=0) # NOTE: No queueing is happening, tasks are dispatched right away assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 assert op.metrics.obj_store_mem_pending_task_inputs > 0 # assert op.metrics.obj_store_mem_pending_task_outputs > 0 # Assert there are 4 running tasks now assert op.num_active_tasks() == 4 assert op._actor_pool.num_pending_actors() == 0 assert op._actor_pool.num_running_actors() == 2 assert op.current_logical_usage() == ExecutionResources(cpu=2, gpu=0, memory=0) assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 assert op.metrics.obj_store_mem_pending_task_inputs == pytest.approx(3200, rel=0.5) # assert op.metrics.obj_store_mem_pending_task_outputs > 0 # Indicate that no more inputs will arrive. op.all_inputs_done() # Wait until tasks are done. run_op_tasks_sync(op) min_usage = ExecutionResources() # Work is done, scale down the actor pool. for pool in op.get_autoscaling_actor_pools(): num_scaled_down = pool.scale( ActorPoolScalingRequest(delta=-pool.current_size()) ) # NOTE: Actor Pool will retain the min-size assert num_scaled_down == pool.current_size() - pool.min_size() min_usage = min_usage.add( pool.per_actor_resource_usage().scale(pool.min_size()) ) assert op.current_logical_usage() == min_usage assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == pytest.approx( 6400, rel=0.5, ) assert op.metrics.obj_store_mem_pending_task_inputs == 0 assert op.metrics.obj_store_mem_pending_task_outputs == 0 # Consume task outputs. while op.has_next(): op.get_next() # Work is done, scale down the actor pool, and outputs have been consumed. for pool in op.get_autoscaling_actor_pools(): num_scaled_down = pool.scale( ActorPoolScalingRequest(delta=-pool.current_size()) ) # NOTE: Actor Pool will retain the min-size assert num_scaled_down == pool.current_size() - pool.min_size() assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 assert op.metrics.obj_store_mem_pending_task_inputs == 0 assert op.metrics.obj_store_mem_pending_task_outputs == 0 def test_actor_pool_resource_reporting_with_dynamic_remote_args( ray_start_10_cpus_shared, ): """Test that current_logical_usage reflects dynamic resources from ray_remote_args_fn, not just the statically defined ray_remote_args.""" input_op = InputDataBuffer( DataContext.get_current(), make_ref_bundles([[SMALL_STR] for i in range(100)]) ) # ray_remote_args set 1 CPU, but ray_remote_args_fn overrides memory to 500 op = MapOperator.create( _mul2_map_data_prcessor, min_rows_per_bundle=None, input_op=input_op, data_context=DataContext.get_current(), name="TestMapper", compute_strategy=ActorPoolStrategy(min_size=2, max_size=2), # Create two actors ray_remote_args={"num_cpus": 1}, ray_remote_args_fn=lambda: {"memory": 500}, ) # Blocking until actors are fully started op.start(ExecutionOptions(), noop_counter()) run_op_tasks_sync(op, only_existing=True) # Should reflect dynamic resources: 2 actors * (1 cpu, 500 memory) assert op.current_logical_usage() == ExecutionResources(cpu=2, gpu=0, memory=1000) def test_actor_pool_scheduling_with_bundling( ray_start_10_cpus_shared, restore_data_context ): # TODO move to test_actor_pool_map_operator.py ctx = ray.data.DataContext.get_current() ctx._max_num_blocks_in_streaming_gen_buffer = 1 MIN_ROWS_PER_BUNDLE = 5 input_op = InputDataBuffer( DataContext.get_current(), make_ref_bundles([[SMALL_STR] for _ in range(100)]) ) op = MapOperator.create( _mul2_map_data_prcessor, input_op=input_op, data_context=DataContext.get_current(), name="TestMapper", compute_strategy=ActorPoolStrategy(min_size=2, max_size=10), min_rows_per_bundle=MIN_ROWS_PER_BUNDLE, ) # NOTE: This is blocking, until actor pool is fully started up op.start(ExecutionOptions(), noop_counter()) min_resource_usage, _ = op.min_max_resource_requirements() assert min_resource_usage == ExecutionResources(cpu=2, gpu=0, object_store_memory=0) # `incremental_resource_usage` should always report 0 CPU and GPU, as # it doesn't consider scaling-up. assert op.incremental_resource_usage() == ExecutionResources( cpu=0, gpu=0, object_store_memory=0 ) # Pool is idle while waiting for actors to start. assert op.current_logical_usage() == ExecutionResources(cpu=2, gpu=0, memory=0) assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 assert op.metrics.obj_store_mem_pending_task_inputs == 0 # assert op.metrics.obj_store_mem_pending_task_outputs == 0 # NOTE: Until actors start up, we should not be adding the inputs to the # operator to avoid queuing up inside of it assert not op.can_add_input() # Finalize operator initialization sequence and make it schedulable run_op_tasks_sync(op, only_existing=True) # Assert all actors are running assert op._actor_pool.num_pending_actors() == 0 assert op._actor_pool.num_running_actors() == 2 # Add inputs for i in range(MIN_ROWS_PER_BUNDLE - 1): assert op.incremental_resource_usage() == ExecutionResources( cpu=0, gpu=0, object_store_memory=0 ) # Should be able to add inputs now assert op.can_add_input() op.add_input(input_op.get_next(), 0) assert op.current_logical_usage() == ExecutionResources(cpu=2, gpu=0, memory=0) # While bundling, no tasks are scheduled assert op.num_active_tasks() == 0 assert op.metrics.obj_store_mem_pending_task_inputs == 0 # assert op.metrics.obj_store_mem_pending_task_outputs == 0 assert op.metrics.obj_store_mem_internal_inqueue == pytest.approx( (i + 1) * 800, rel=0.5 ) assert op.metrics.obj_store_mem_internal_outqueue == 0 # Adding 1 more input triggers task scheduling op.add_input(input_op.get_next(), 0) assert op.num_active_tasks() == 1 # Queue is now empty assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 # Running task has pending inputs/outputs single_task_pending_inputs = op.metrics.obj_store_mem_pending_task_inputs single_task_pending_outputs = op.metrics.obj_store_mem_pending_task_outputs assert single_task_pending_inputs > 0 # assert single_task_pending_outputs > 0 # Add more inputs, but less than necessary to launch another task for i in range(MIN_ROWS_PER_BUNDLE - 1): assert op.incremental_resource_usage() == ExecutionResources( cpu=0, gpu=0, object_store_memory=0, memory=0 ) # Should be able to add inputs now assert op.can_add_input() op.add_input(input_op.get_next(), 0) assert op.current_logical_usage() == ExecutionResources(cpu=2, gpu=0, memory=0) # While bundling, no *new* tasks are scheduled assert op.num_active_tasks() == 1 assert op.metrics.obj_store_mem_internal_inqueue == pytest.approx( (i + 1) * 800, rel=0.5 ) assert op.metrics.obj_store_mem_internal_outqueue == 0 assert ( op.metrics.obj_store_mem_pending_task_inputs == single_task_pending_inputs ) assert ( op.metrics.obj_store_mem_pending_task_outputs == single_task_pending_outputs ) # Mark inputs as completed op.all_inputs_done() # Bundler should be drained and 1 more task launched assert op.num_active_tasks() == 2 assert op._block_ref_bundler.num_blocks() == 0 assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 assert op.metrics.obj_store_mem_pending_task_inputs > 0 # assert op.metrics.obj_store_mem_pending_task_outputs > 0 # Wait until tasks are done. run_op_tasks_sync(op) # Work is done, scale down the actor pool. for pool in op.get_autoscaling_actor_pools(): num_scaled_down = pool.scale( ActorPoolScalingRequest(delta=-pool.current_size()) ) # NOTE: Actor Pool will retain the min-size assert num_scaled_down == pool.current_size() - pool.min_size() assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == pytest.approx(12000, rel=0.5) assert op.metrics.obj_store_mem_pending_task_inputs == 0 assert op.metrics.obj_store_mem_pending_task_outputs == 0 # Consume task outputs. while op.has_next(): op.get_next() min_usage = ExecutionResources() # Work is done, scale down the actor pool, and outputs have been consumed. for pool in op.get_autoscaling_actor_pools(): num_scaled_down = pool.scale( ActorPoolScalingRequest(delta=-pool.current_size()) ) # NOTE: Actor Pool will retain the min-size assert num_scaled_down == pool.current_size() - pool.min_size() min_usage = min_usage.add( pool.per_actor_resource_usage().scale(pool.min_size()) ) assert op.current_logical_usage() == min_usage assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 assert op.metrics.obj_store_mem_pending_task_inputs == 0 assert op.metrics.obj_store_mem_pending_task_outputs == 0 def test_limit_resource_reporting(ray_start_10_cpus_shared): input_op = InputDataBuffer( DataContext.get_current(), make_ref_bundles([[SMALL_STR, SMALL_STR] for i in range(2)]), ) # Two two-row bundles op = LimitOperator(3, input_op, DataContext.get_current()) op.start(ExecutionOptions(), noop_counter()) assert op.current_logical_usage() == ExecutionResources( cpu=0, gpu=0, object_store_memory=0, memory=0 ) assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 op.add_input(input_op.get_next(), 0) assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == pytest.approx(1600, rel=0.5) op.add_input(input_op.get_next(), 0) assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == pytest.approx(2400, rel=0.5) while op.has_next(): op.get_next() assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 def test_output_splitter_resource_reporting(ray_start_10_cpus_shared): input_op = InputDataBuffer( DataContext.get_current(), make_ref_bundles([[SMALL_STR] for i in range(4)]) ) op = OutputSplitter( input_op, 2, equal=False, data_context=DataContext.get_current(), locality_hints=["0", "1"], ) op.start(ExecutionOptions(actor_locality_enabled=True), noop_counter()) assert op.current_logical_usage() == ExecutionResources( cpu=0, gpu=0, object_store_memory=0, memory=0 ) assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 # 2 * n inputs are bufferred to optimize locality. for i in range(3): op.add_input(input_op.get_next(), 0) assert op.metrics.obj_store_mem_internal_inqueue == pytest.approx( 800 * (i + 1), rel=0.5 ) assert op.metrics.obj_store_mem_internal_outqueue == 0 op.add_input(input_op.get_next(), 0) assert op.metrics.obj_store_mem_internal_inqueue == pytest.approx(2400, rel=0.5) assert op.metrics.obj_store_mem_internal_outqueue == pytest.approx(800, rel=0.5) op.all_inputs_done() while op.has_next(): op.get_next() assert op.metrics.obj_store_mem_internal_inqueue == 0 assert op.metrics.obj_store_mem_internal_outqueue == 0 def test_execution_resources_to_resource_dict(): resources = ExecutionResources(cpu=1, gpu=2, object_store_memory=3, memory=4) assert resources.to_resource_dict() == { "CPU": 1, "GPU": 2, "object_store_memory": 3, "memory": 4, } def test_execution_resources_combine_sum_empty_reuses_zero(): # An empty fold returns the shared zero singleton instead of allocating. assert ExecutionResources.combine_sum([]) is ExecutionResources.zero() # Works for a one-shot generator (can't be len()'d or re-iterated). assert ExecutionResources.combine_sum(iter([])) is ExecutionResources.zero() def test_execution_resources_combine_sum(): rs = [ ExecutionResources(cpu=1, gpu=2, object_store_memory=3, memory=4), ExecutionResources(cpu=10, gpu=20, object_store_memory=30, memory=40), ] expected = ExecutionResources(cpu=11, gpu=22, object_store_memory=33, memory=44) assert ExecutionResources.combine_sum(rs) == expected # Same result from a one-shot generator. assert ExecutionResources.combine_sum(r for r in rs) == expected def test_execution_resources_combine(): rs = [ ExecutionResources(cpu=1, gpu=5, object_store_memory=3, memory=40), ExecutionResources(cpu=10, gpu=2, object_store_memory=30, memory=4), ] # Per-dimension fold with an arbitrary float op. assert ExecutionResources.combine(rs, operator.add) == ExecutionResources( 11, 7, 33, 44 ) assert ExecutionResources.combine(rs, max) == ExecutionResources(10, 5, 30, 40) assert ExecutionResources.combine(rs, min) == ExecutionResources(1, 2, 3, 4) # Single-pass over a one-shot generator. assert ExecutionResources.combine((r for r in rs), max) == ExecutionResources( 10, 5, 30, 40 ) # Empty -> None (no identity to seed a general fn with). assert ExecutionResources.combine([], operator.add) is None if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", __file__]))