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

Code before:

    op.get_next()
    gc.collect()

    # `InputDataBuffer` should still hold a reference to the input block even after
    # `get_next` is called.
    assert len(gc.get_referrers(block_ref)) > 0


def test_union_operator_throttling_disabled(ray_start_regular_shared):
    """Test that UnionOperator has throttling disabled.

    UnionOperator only manipulates bundle metadata and doesn't launch any tasks,
    so it should not be allocated resources.
    """
    ctx = DataContext.get_current()
    union_op = UnionOperator(ctx)

    # Verify that throttling_disabled() returns True
    assert union_op.throttling_disabled() is True


def test_union_operator_not_allocated_resources(ray_start_regular_shared):
    ctx = DataContext.get_current()
    original_enabled = ctx.op_resource_reservation_enabled
    try:
        ctx.op_resource_reservation_enabled = True
        input_op1 = InputDataBuffer(ctx, make_ref_bundles([[1, 2, 3]]))
        input_op2 = InputDataBuffer(ctx, make_ref_bundles([[4, 5, 6]]))
        union_op = UnionOperator(ctx, input_op1, input_op2)

        topo = build_streaming_topology(union_op, ExecutionOptions())
        resource_manager = ResourceManager(topo, ExecutionOptions(), MagicMock(), ctx)
        allocator = create_resource_allocator(resource_manager, ctx)
        assert allocator is not None

        allocator.update_budgets(
            limits=ExecutionResources(cpu=4, gpu=0, object_store_memory=1000)
        )
        allocation = allocator.get_allocation(union_op)
        assert allocation is None
    finally:
        ctx.op_resource_reservation_enabled = original_enabled


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))


Review comment (suggestion):
Rather than adding a `try-catch`, use the `restore_data_context` fixture: https://docs.ray.io/en/master/data/contributing/how-to-write-tests.html#use-fixtures-skip-try-finally

Code after:

    op.get_next()
    gc.collect()

    # `InputDataBuffer` should still hold a reference to the input block even after
    # `get_next` is called.
    assert len(gc.get_referrers(block_ref)) > 0


def test_union_operator_throttling_disabled(ray_start_regular_shared):
    """Test that UnionOperator has throttling disabled.

    UnionOperator only manipulates bundle metadata and doesn't launch any tasks,
    so it should not be allocated resources.
    """
    ctx = DataContext.get_current()
    union_op = UnionOperator(ctx)

    # Verify that throttling_disabled() returns True
    assert union_op.throttling_disabled() is True


def test_union_operator_not_allocated_resources(
    ray_start_regular_shared, restore_data_context
):
    ctx = DataContext.get_current()
    ctx.op_resource_reservation_enabled = True
    input_op1 = InputDataBuffer(ctx, make_ref_bundles([[1, 2, 3]]))
    input_op2 = InputDataBuffer(ctx, make_ref_bundles([[4, 5, 6]]))
    union_op = UnionOperator(ctx, input_op1, input_op2)

    topo = build_streaming_topology(union_op, ExecutionOptions())
    resource_manager = ResourceManager(topo, ExecutionOptions(), MagicMock(), ctx)
    allocator = create_resource_allocator(resource_manager, ctx)
    assert allocator is not None

    allocator.update_budgets(
        limits=ExecutionResources(cpu=1, gpu=0, object_store_memory=1000)
    )
    allocation = allocator.get_allocation(union_op)
    assert allocation is None


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))