Function test_map_operator_ray_args

python/ray/data/tests/test_map_operator.py:380–420

(shutdown_only, use_actors)

@pytest.mark.parametrize("use_actors", [False, True])
def test_map_operator_ray_args(shutdown_only, use_actors):
    ray.shutdown()
    ray.init(num_cpus=0, num_gpus=1)
    # Create with inputs.
    input_op = InputDataBuffer(
        DataContext.get_current(), make_ref_bundles([[i] for i in range(10)])
    )
    compute_strategy = ActorPoolStrategy(size=1) if use_actors else TaskPoolStrategy()
    op = MapOperator.create(
        _mul2_map_data_prcessor,
        input_op=input_op,
        data_context=DataContext.get_current(),
        name="TestMapper",
        compute_strategy=compute_strategy,
        ray_remote_args={"num_cpus": 0, "num_gpus": 1},
    )

    # Feed data and block on exec.
    op.start(ExecutionOptions())
    if use_actors:
        # Wait for the actor to start.
        run_op_tasks_sync(op)

    while input_op.has_next():
        if use_actors:
            # For actors, we need to check capacity before adding input
            # and process tasks when the actor pool is at capacity.
            while not op.can_add_input():
                run_one_op_task(op)

        assert op.can_add_input()
        op.add_input(input_op.get_next(), 0)

    op.all_inputs_done()
    run_op_tasks_sync(op)

    # Check we don't hang and complete with num_gpus=1.
    outputs = _take_outputs(op)
    expected = [[i * 2] for i in range(10)]
    assert sorted(outputs) == expected, f"Expected {expected}, got {outputs}"
    assert op.has_completed()


@pytest.mark.parametrize("use_actors", [False, True])