code: def test_iter_batches_e2e_async(ray_start_regular_shared):
      """We add time.sleep in 3 places:
      1. In the base generator to simulate streaming executor blocking on next results.
      2. In the collate_fn to simulate expensive slicing/formatting/collation
      3. In the user thread to simulate training.
      """
  
      def collate_fn(batch):
          time.sleep(2)
          return batch
  
      block_refs_iter = itertools.starmap(
          lambda block, metadata: (ray.put(block), metadata),
          block_generator(num_blocks=20, num_rows=2),
      )
      start_time = time.time()
      output_batches = iter_batches(
          block_refs_iter,
          dataset_tag="dataset",
          batch_size=None,
          collate_fn=collate_fn,
          prefetch_batches=4,
      )
      batches = []
      for batch in output_batches:
          time.sleep(1.5)
          batches.append(batch)
      end_time = time.time()
  
      # 20 batches, 1.5 second sleep. Should be less than 45 seconds, even with some
      # overhead.
      # If there was no overlap, then we would expect this to take at least 20*2.5 = 50
      assert end_time - start_time < 45, end_time - start_time
  
      assert len(batches) == 20
      assert all(len(batch) == 2 for batch in batches)
  
file_path: ray/python/ray/data/tests/block_batching/test_iter_batches.py
imports: import itertools
  import queue
  import threading
  import time
  from typing import Iterator, List, Tuple
  import pandas
  import pyarrow
  import pytest
  import ray
  from ray.data._internal.block_batching.interfaces import Batch, BlockPrefetcher
  from ray.data._internal.block_batching.iter_batches import iter_batches, prefetch_batches_locally, restore_original_order
  from ray.data.block import Block, BlockMetadata
label_desc: Use image node_id 6 for calling a global function with example usage: test_iter_batches_e2e_async(ray_start_regular_shared) and returns: batch