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

Code before:
        actor = PoolWorker.options(_labels=labels).remote("node1")
        return actor, actor.get_location.remote()

    pool = _ActorPool(
        create_actor_fn,
        ExecutionResources(cpu=1),
        min_size=1,
        max_size=4,
        initial_size=1,
        max_actor_concurrency=1,
        max_tasks_in_flight_per_actor=4,
        _map_worker_cls_name="MapWorker(TestOp)",
    )

    with caplog.at_level(logging.DEBUG, logger=logger_name):
        caplog.clear()
        pool.scale(ActorPoolScalingRequest(delta=1, reason="scale up test"))
        assert re.search(
            r"Scaling up MapWorker\(TestOp\) actor pool by 1 \(reason=scale up test",
            caplog.text,
        ), f"Expected scale-up log with map worker name; got: {caplog.text}"

        with freeze_time() as frozen_time:
            frozen_time.tick(
                datetime.timedelta(
                    seconds=_ActorPool._ACTOR_POOL_SCALE_DOWN_DEBOUNCE_PERIOD_S + 1
                )
            )
            caplog.clear()
            pool.scale(
                ActorPoolScalingRequest.downscale(delta=-1, reason="scale down test")
            )
            assert re.search(
                r"Scaled down MapWorker\(TestOp\) actor pool by \d+ "
                r"\(reason=scale down test",
                caplog.text,
            ), f"Expected scale-down log with map worker name; got: {caplog.text}"


def _create_bundle_with_single_row(row):
    block = pa.Table.from_pylist([row])
    block_ref = ray.put(block)
    metadata = BlockAccessor.for_block(block).get_metadata()
    schema = BlockAccessor.for_block(block).schema()
    return RefBundle([(block_ref, metadata)], owns_blocks=False, schema=schema)


@pytest.mark.parametrize("min_rows_per_bundle", [2, None])
def test_internal_input_queue_is_empty_after_early_completion(
    ray_start_regular_shared, min_rows_per_bundle, restore_data_context
):

Review comment (refactor):
This is breaking abstraction barriers and accessing an internal constant. If I refactor `_ActorPool` a month from now and change the debouncing logic, this test will break even if I don't break any of `_ActorPool`'s public contracts.

An alternative is to expos a seam in `_ActorPool`:
* Add a parameter for the debouncing period s to the constructor that defaults to _ACTOR_POOL_SCALE_DOWN_DEBOUNCE_PERIOD_S
* In this test, set that parameter to 0
* Then, you shouldn't need to use freeze_time or anything

Code after:
        actor = PoolWorker.options(_labels=labels).remote("node1")
        return actor, actor.get_location.remote()

    pool = _ActorPool(
        create_actor_fn,
        ExecutionResources(cpu=1),
        min_size=1,
        max_size=1,
        initial_size=1,
        max_actor_concurrency=1,
        max_tasks_in_flight_per_actor=1,
        map_worker_cls_name="MapWorker(TestOp)",
        debounce_period_s=0,
    )

    with caplog.at_level(logging.DEBUG, logger=logger_name):
        caplog.clear()
        pool.scale(ActorPoolScalingRequest(delta=1, reason="scale up test"))
        assert re.search(
            r"Scaling up MapWorker\(TestOp\) actor pool by 1 \(reason=scale up test",
            caplog.text,
        ), f"Expected scale-up log with map worker name; got: {caplog.text}"

        caplog.clear()
        pool.scale(
            ActorPoolScalingRequest.downscale(delta=-1, reason="scale down test")
        )
        assert re.search(
            r"Scaled down MapWorker\(TestOp\) actor pool by \d+ "
            r"\(reason=scale down test",
            caplog.text,
        ), f"Expected scale-down log with map worker name; got: {caplog.text}"


def _create_bundle_with_single_row(row):
    block = pa.Table.from_pylist([row])
    block_ref = ray.put(block)
    metadata = BlockAccessor.for_block(block).get_metadata()
    schema = BlockAccessor.for_block(block).schema()
    return RefBundle([(block_ref, metadata)], owns_blocks=False, schema=schema)


@pytest.mark.parametrize("min_rows_per_bundle", [2, None])
def test_internal_input_queue_is_empty_after_early_completion(
    ray_start_regular_shared, min_rows_per_bundle, restore_data_context
):
    data_context = ray.data.DataContext.get_current()
    # Block until actors are ready so we can schedule tasks
    data_context.wait_for_min_actors_s = 60

    op = MapOperator.create(