code: def test_streaming_backpressure_e2e(restore_data_context):
      # This test case is particularly challenging since there is a large input->output
      # increase in data size: https://github.com/ray-project/ray/issues/34041
      class TestSlow:
          def __call__(self, df: np.ndarray):
              time.sleep(2)
              return {"id": np.random.randn(1, 20, 1024, 1024)}
  
      class TestFast:
          def __call__(self, df: np.ndarray):
              time.sleep(0.5)
              return {"id": np.random.randn(1, 20, 1024, 1024)}
  
      ctx = ray.init(object_store_memory=4e9)
      ds = ray.data.range_tensor(
          20, shape=(3, 1024, 1024), parallelism=20
      )
  
      pipe = ds.map_batches(
          TestFast,
          batch_size=1,
          num_cpus=0.5,
          compute=ray.data.ActorPoolStrategy(size=2),
      ).map_batches(
          TestSlow,
          batch_size=1,
          compute=ray.data.ActorPoolStrategy(size=1),
      )
  
      for batch in pipe.iter_batches(batch_size=1, prefetch_batches=2):
          ...
  
      # If backpressure is not working right, we will spill.
      meminfo = memory_summary(
          ctx.address_info["address"], stats_only=True
      )
      assert "Spilled" not in meminfo, meminfo
  
file_path: ray/python/ray/data/tests/test_streaming_backpressure_edge_case.py
imports: import time
  import numpy
  import pandas
  import pytest
  import ray
  from ray._private.internal_api import memory_summary
  from ray.data.block import BlockMetadata
  from ray.data.datasource import Datasource, ReadTask
  from ray.data.tests.conftest import *
  from ray.tests.conftest import *
label_desc: Use image node_id 2 for calling a global function with example usage: test_streaming_backpressure_e2e(restore_data_context) and returns: dict, dict