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

Code before:
            topo2, ExecutionOptions(), MagicMock(), DataContext.get_current()
        )

        # o5's downstream (o6, o7) has no blocking materializing ops
        assert resource_manager2._is_blocking_materializing_op(o5) is False
        assert resource_manager2._is_blocking_materializing_op(o7) is False

    def test_memory_limit_blocks_task_submission(self, restore_data_context):
        """Test that tasks are blocked when memory limit is exceeded."""
        # Cluster has 10GB memory available
        cluster_resources = ExecutionResources(cpu=10, gpu=0, memory=10 * 1024**3)

        def get_total_resources():
            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())


Review comment (refactor):
Is specifying the `resource_limits` necessary here? IIUC this test should still verify the correct behavior even if we remove this (`resource_limits` is an optional configuration), and removing it would simplify the test

Code after:
        )

        # o5's downstream (o6, o7) has no blocking materializing ops
        assert resource_manager2._is_blocking_materializing_op(o5) is False
        assert resource_manager2._is_blocking_materializing_op(o7) is False

    def test_memory_limit_blocks_task_submission(self, restore_data_context):
        """Test that tasks are blocked when memory limit is exceeded."""
        # Cluster has 1000 bytes of memory
        cluster_resources = ExecutionResources(cpu=1, gpu=0, memory=1000)

        # Request 2000 bytes memory
        o1 = InputDataBuffer(DataContext.get_current(), [])
        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())