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

Code before:
    def test_on_data_ready_exception(self, ray_start_regular_shared):
        streaming_gen = create_stub_streaming_gen(
            block_nbytes=[128 * MiB],
            raise_exception=AssertionError("Block generation failed"),
        )

        def verify_exception(exc, task_exec_stats, task_exec_driver_stats):
            assert isinstance(exc, AssertionError)
            assert task_exec_stats is None
            assert task_exec_driver_stats is None

        data_op_task = DataOpTask(
            0,
            streaming_gen,
            task_done_callback=verify_exception,
        )

        with pytest.raises(AssertionError, match="Block generation failed"):
            while not data_op_task.has_finished:
                ray.wait([streaming_gen], fetch_local=False)
                data_op_task.on_data_ready(None)

    @pytest.mark.parametrize(
        "preempt_on", ["block_ready_callback", "metadata_ready_callback"]
    )
    def test_on_data_ready_with_preemption_during_call(
        self,
        preempt_on: Union[
            Literal["block_ready_callback"], Literal["metadata_ready_callback"]
        ],
        ray_start_cluster_enabled,
        ensure_block_metadata_stored_in_plasma,
    ):
        """Test that ``on_data_ready`` works when a node dies during its execution."""
        # Shutdown Ray incase it's already initialized.
        ray.shutdown()

        # Create a single-worker-node cluster with 1 logical CPU.
        cluster = ray_start_cluster_enabled
        head_node = cluster.add_node(num_cpus=0)  # noqa: F841
        cluster.wait_for_nodes()
        ray.init()

        worker_node = cluster.add_node(num_cpus=1)
        cluster.wait_for_nodes()

        # Create a streaming generator that produces a single 128 MiB output block, and
        # configure it so that it preempts the worker node in the specified callback.
        streaming_gen = create_stub_streaming_gen(block_nbytes=[128 * MiB])

        def remove_and_add_back_worker_node(_):

Review comment (suggestion):
In general, I think we should avoid testing against internal attributes, but won't block since I know this has been in-flight for a while

Code after:
    def test_on_data_ready_exception(self, ray_start_regular_shared):
        streaming_gen = create_stub_streaming_gen(
            block_nbytes=[128 * MiB],
            raise_exception=AssertionError("Block generation failed"),
        )

        def verify_exception(exc, task_exec_stats, task_exec_driver_stats):
            assert isinstance(exc, AssertionError)
            assert task_exec_stats is None
            assert task_exec_driver_stats is None

        data_op_task = DataOpTask(
            0,
            streaming_gen,
            task_done_callback=verify_exception,
        )

        with pytest.raises(AssertionError, match="Block generation failed"):
            while not data_op_task.has_finished:
                ray.wait([streaming_gen], fetch_local=False)
                data_op_task.on_data_ready(None)

    def test_operator_name_parameter(self, ray_start_regular_shared):
        streaming_gen = create_stub_streaming_gen(block_nbytes=[1])
        task = DataOpTask(0, streaming_gen, operator_name="MapBatches(fn)")
        assert task._operator_name == "MapBatches(fn)"

        streaming_gen2 = create_stub_streaming_gen(block_nbytes=[1])
        task_default = DataOpTask(1, streaming_gen2)
        assert task_default._operator_name == "Unknown"

    @pytest.mark.parametrize(
        "preempt_on", ["block_ready_callback", "metadata_ready_callback"]
    )
    def test_on_data_ready_with_preemption_during_call(
        self,
        preempt_on: Union[
            Literal["block_ready_callback"], Literal["metadata_ready_callback"]
        ],
        ray_start_cluster_enabled,
        ensure_block_metadata_stored_in_plasma,
    ):
        """Test that ``on_data_ready`` works when a node dies during its execution."""
        # Shutdown Ray incase it's already initialized.
        ray.shutdown()

        # Create a single-worker-node cluster with 1 logical CPU.
        cluster = ray_start_cluster_enabled
        head_node = cluster.add_node(num_cpus=0)  # noqa: F841
        cluster.wait_for_nodes()
        ray.init()