code: def test_input_backpressure_e2e(restore_data_context, shutdown_only):
      # Tests that backpressure applies even when reading directly from the input
      # datasource. This relies on datasource metadata size estimation.
      @ray.remote
      class Counter:
          def __init__(self):
              self.count = 0
  
          def increment(self):
              self.count += 1
  
          def get(self):
              return self.count
  
          def reset(self):
              self.count = 0
  
      class CountingRangeDatasource(Datasource):
          def __init__(self):
              self.counter = Counter.remote()
  
          def prepare_read(self, parallelism, n):
              def range_(i):
                  ray.get(self.counter.increment.remote())
                  return [
                      pd.DataFrame(
                          {
                              "data": np.ones(
                                  (n // parallelism * 1024 * 1024,)
                              )
                          }
                      )
                  ]
  
              sz = (n // parallelism) * 1024 * 1024 * 8
              print("Block size", sz)
  
              return [
                  ReadTask(
                      lambda i=i: range_(i),
                      BlockMetadata(
                          num_rows=n // parallelism,
                          size_bytes=sz,
                          schema=None,
                          input_files=None,
                          exec_stats=None,
                      ),
                  )
                  for i in range(parallelism)
              ]
  
      source = CountingRangeDatasource()
      ctx = ray.data.DataContext.get_current()
      ctx.execution_options.resource_limits.object_store_memory = 10e6
  
      # 10GiB dataset.
      ds = ray.data.read_datasource(source, n=10000, parallelism=1000)
      it = iter(ds.iter_batches(batch_size=None, prefetch_batches=0))
      next(it)
      time.sleep(3)
      launched = ray.get(source.counter.get.remote())
  
      # If backpressure is broken we'll launch 15+.
      assert launched < 5, launched
  
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 1 for calling a global function with example usage: test_input_backpressure_e2e(restore_data_context, shutdown_only) and returns: self, unknown, list