Function test_map_kwargs

python/ray/data/tests/test_map_operator.py:486–534

Test propagating additional kwargs to map tasks.

(ray_start_regular_shared, use_actors)

@pytest.mark.parametrize("use_actors", [False, True])
def test_map_kwargs(ray_start_regular_shared, use_actors):
    """Test propagating additional kwargs to map tasks."""
    foo = 1
    bar = np.random.random(1024 * 1024)
    kwargs = {
        "foo": foo,  # Pass by value
        "bar": ray.put(bar),  # Pass by ObjectRef
    }

    def map_fn(block_iter: Iterable[Block], ctx: TaskContext) -> Iterable[Block]:
        nonlocal foo, bar
        assert ctx.kwargs["foo"] == foo
        # bar should be automatically deref'ed.
        assert np.array_equal(ctx.kwargs["bar"], bar)

        yield from block_iter

    input_op = InputDataBuffer(
        DataContext.get_current(),
        make_ref_bundles([[i] for i in range(10)]),
    )
    compute_strategy = ActorPoolStrategy() if use_actors else TaskPoolStrategy()
    op = MapOperator.create(
        create_map_transformer_from_block_fn(map_fn),
        input_op=input_op,
        data_context=DataContext.get_current(),
        name="TestMapper",
        compute_strategy=compute_strategy,
    )
    op.add_map_task_kwargs_fn(lambda: kwargs)
    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)

    _take_outputs(op)
    assert op.has_completed()


@pytest.mark.parametrize(