Repository: ray-project/ray
File: python/ray/data/tests/test_streaming_executor.py (Python)

Code before:
        make_map_transformer(lambda block: [b * -1 for b in block]),
        o1,
        DataContext.get_current(),
        compute_strategy=ray.data.ActorPoolStrategy(size=8),
    )
    o3 = MapOperator.create(
        make_map_transformer(lambda block: [b * 2 for b in block]),
        o2,
        DataContext.get_current(),
        compute_strategy=ray.data.ActorPoolStrategy(size=4),
    )
    _validate_dag(o3, ExecutionResources.for_limits())
    _validate_dag(o3, ExecutionResources.for_limits(cpu=20))
    _validate_dag(o3, ExecutionResources.for_limits(gpu=0))
    with pytest.raises(ValueError):
        _validate_dag(o3, ExecutionResources.for_limits(cpu=10))


# Mock the `scale_up` method to avoid creating and leaking resources.
@patch(
    "ray.data._internal.execution.operators.actor_pool_map_operator._ActorPool.scale",
    return_value=1,
)
def test_configure_output_locality_is_deprecated_noop(
    mock_scale_up, ray_start_regular_shared
):
    inputs = make_ref_bundles([[x] for x in range(20)])
    o1 = InputDataBuffer(DataContext.get_current(), inputs)
    o2 = MapOperator.create(
        make_map_transformer(lambda block: [b * -1 for b in block]),
        o1,
        DataContext.get_current(),
    )
    o3 = MapOperator.create(
        make_map_transformer(lambda block: [b * 2 for b in block]),
        o2,
        DataContext.get_current(),
        compute_strategy=ray.data.ActorPoolStrategy(size=1),
    )
    # Baseline scheduling without locality_with_output.
    build_streaming_topology(o3, ExecutionOptions())
    s1_baseline = o2._get_dynamic_ray_remote_args()["scheduling_strategy"]
    s2_baseline = o3._get_dynamic_ray_remote_args()["scheduling_strategy"]

    # Deprecated option should warn and remain a no-op for scheduling.
    with pytest.warns(DeprecationWarning, match="no-op"):
        build_streaming_topology(o3, ExecutionOptions(locality_with_output=True))
    s1 = o2._get_dynamic_ray_remote_args()["scheduling_strategy"]
    s2 = o3._get_dynamic_ray_remote_args()["scheduling_strategy"]
    assert s1 == s1_baseline
    assert s2 == s2_baseline

Review comment (refactor):
I think you could substantially simplify this test. I think you can probably test this with just a few lines:
```python
def test_configure_output_locality_is_deprecated_noop(..., restore_data_context):
    with pytest.warns(...):
        ray.data.DataContext().get_current().locality_hints = True
```

Code after:
        make_map_transformer(lambda block: [b * -1 for b in block]),
        o1,
        DataContext.get_current(),
        compute_strategy=ray.data.ActorPoolStrategy(size=8),
    )
    o3 = MapOperator.create(
        make_map_transformer(lambda block: [b * 2 for b in block]),
        o2,
        DataContext.get_current(),
        compute_strategy=ray.data.ActorPoolStrategy(size=4),
    )
    _validate_dag(o3, ExecutionResources.for_limits())
    _validate_dag(o3, ExecutionResources.for_limits(cpu=20))
    _validate_dag(o3, ExecutionResources.for_limits(gpu=0))
    with pytest.raises(ValueError):
        _validate_dag(o3, ExecutionResources.for_limits(cpu=10))


def test_configure_output_locality_is_deprecated_noop(restore_data_context):
    data_context = ray.data.DataContext.get_current()

    with pytest.warns(DeprecationWarning, match="no-op"):
        data_context.execution_options.locality_with_output = True
    assert data_context.execution_options.locality_with_output is False


class OpBufferQueueTest(unittest.TestCase):
    def test_multi_threading(self):
        num_blocks = 5_000
        num_splits = 8
        num_per_split = num_blocks // num_splits
        ref_bundles = make_ref_bundles([[[i]] for i in range(num_blocks)])

        queue = OpBufferQueue()
        for i, ref_bundle in enumerate(ref_bundles):
            ref_bundle.output_split_idx = i % num_splits
            queue.append(ref_bundle)

        def consume(output_split_idx):
            nonlocal queue

            count = 0
            while queue.has_next(output_split_idx):
                ref_bundle = queue.pop(output_split_idx)
                count += 1
                assert ref_bundle is not None
                assert ref_bundle.output_split_idx == output_split_idx
            assert count == num_per_split
            return True

        with ThreadPoolExecutor(max_workers=num_splits) as executor: