code: def test_streaming_generator_backpressure_multiple_objects(
      shutdown_only, backpressure_size
  ):
      """Verify backpressure works when it needs more than 1 objects
      to backpressure.
      """
      ray.init()
  
      @ray.remote
      class Reporter:
          def __init__(self):
              self.reported = set()
  
          def report(self, i):
              self.reported.add(i)
  
          def reported(self):
              return self.reported
  
      TOTAL_RETURN = 10
  
      @ray.remote(_generator_backpressure_num_objects=backpressure_size)
      def f(reporter):
          for i in range(TOTAL_RETURN):
              print("yield ", i)
              ray.get(reporter.report.remote(i))
              yield np.random.rand(1024)  # 8KB
  
      reporter = Reporter.remote()
  
      def check_reported(i):
          return i in ray.get(reporter.reported.remote())
  
      gen = f.remote(reporter)
      # It is backpressured for every 2 objects.
      for i in range(
          0, TOTAL_RETURN - backpressure_size, backpressure_size
      ):
          print("iteration ", i)
          r, _ = ray.wait([gen])
          assert len(r) == 1
          for j in range(backpressure_size):
              wait_for_condition(lambda: check_reported(i + j))
          wait_for_condition(
              lambda: not check_reported(i + backpressure_size)
          )
          # Wait a little bit to make sure it is backpressured.
          time.sleep(2)
          wait_for_condition(
              lambda: not check_reported(i + backpressure_size)
          )
          # Consume the ref -> task will progress.
          for j in range(backpressure_size):
              ray.get(next(gen))
              # We cannot have more values than TOTAL_RETURN - 1, so
              # we should use min here.
              wait_for_condition(
                  lambda: check_reported(
                      min(i + backpressure_size + j, TOTAL_RETURN - 1)
                  )
              )
  
file_path: ray/python/ray/tests/test_streaming_generator_backpressure.py
imports: import asyncio
  import pytest
  import numpy
  import sys
  import time
  import signal
  import os
  import ray
  from ray._private.test_utils import wait_for_condition
  from ray.util.state import list_tasks
label_desc: Use image node_id 2 for calling a global function with example usage: test_streaming_generator_backpressure_multiple_objects(shutdown_only, backpressure_size) and returns: unknown, self