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

Code before:
            return cluster_resources

        # Create operator that requests 15GB memory
        o1 = InputDataBuffer(DataContext.get_current(), [])
        o2 = mock_map_op(
            o1,
            ray_remote_args={"num_cpus": 1, "memory": 15 * 1024**3},
            name="HighMemoryTask",
        )

        topo = build_streaming_topology(o2, ExecutionOptions())
        options = ExecutionOptions()
        options.resource_limits = cluster_resources

        resource_manager = ResourceManager(
            topo, options, get_total_resources, DataContext.get_current()
        )
        resource_manager.update_usages()

        # Task cannot be submitted because it exceeds memory limit
        allocator = resource_manager._op_resource_allocator
        if allocator is not None:
            can_submit = allocator.can_submit_new_task(o2)
            assert not can_submit, (
                "Task should be blocked: requires 15GB but only 10GB available"
            )


class TestResourceAllocatorUnblockingStreamingOutputBackpressure:
    """Tests for OpResourceAllocator._should_unblock_streaming_output_backpressure."""

    def test_unblock_backpressure_terminal_operator(self, restore_data_context):
        """Terminal operator (no downstream) should always unblock."""
        o1 = InputDataBuffer(DataContext.get_current(), [])
        o2 = mock_map_op(o1)

        topo = build_streaming_topology(o2, ExecutionOptions())

        resource_manager = ResourceManager(
            topo, ExecutionOptions(), MagicMock(), DataContext.get_current()
        )
        allocator = resource_manager._op_resource_allocator

        # o2 is terminal (no output_dependencies), should always unblock
        assert allocator._should_unblock_streaming_output_backpressure(o2) is True

        # Add o3 operator - o2 is no longer terminal
        o3 = mock_map_op(o2)

        topo = build_streaming_topology(o3, ExecutionOptions())


Review comment (refactor):
Breaking abstraction barriers like accessing internal attributes can lead to brittle tests. Would it make sense to instead call `create_resource_allocator`, assert that it's not None, and test against the directly constructed allocator?

Code after:
        o2 = mock_map_op(
            o1,
            ray_remote_args={"num_cpus": 1, "memory": 2000},
            name="HighMemoryTask",
        )

        topo = build_streaming_topology(o2, ExecutionOptions())
        options = ExecutionOptions()

        resource_manager = ResourceManager(
            topology=topo,
            options=options,
            get_total_resources=lambda: cluster_resources,
            data_context=DataContext.get_current(),
        )
        resource_manager.update_usages()

        # Task cannot be submitted because it exceeds memory limit
        allocator = create_resource_allocator(
            resource_manager, DataContext.get_current()
        )
        assert allocator is not None
        allocator.update_budgets(limits=resource_manager.get_global_limits())
        can_submit = allocator.can_submit_new_task(o2)
        assert (
            not can_submit
        ), "Task should be blocked: requires 2000 bytes but only 1000 bytes memory available"


class TestResourceAllocatorUnblockingStreamingOutputBackpressure:
    """Tests for OpResourceAllocator._should_unblock_streaming_output_backpressure."""

    def test_unblock_backpressure_terminal_operator(self, restore_data_context):
        """Terminal operator (no downstream) should always unblock."""
        o1 = InputDataBuffer(DataContext.get_current(), [])
        o2 = mock_map_op(o1)

        topo = build_streaming_topology(o2, ExecutionOptions())

        resource_manager = ResourceManager(
            topo, ExecutionOptions(), MagicMock(), DataContext.get_current()
        )
        allocator = resource_manager._op_resource_allocator

        # o2 is terminal (no output_dependencies), should always unblock
        assert allocator._should_unblock_streaming_output_backpressure(o2) is True

        # Add o3 operator - o2 is no longer terminal
        o3 = mock_map_op(o2)

        topo = build_streaming_topology(o3, ExecutionOptions())