code: def test_generator_wait(shutdown_only):
      """
      Make sure the generator works with ray.wait.
      """
      ray.init(num_cpus=8)
  
      @ray.remote
      def f(sleep_time):
          for i in range(2):
              time.sleep(sleep_time)
              yield i
  
      @ray.remote
      def g(sleep_time):
          time.sleep(sleep_time)
          return 10
  
      gen = f.remote(1)
  
      """
      Test basic cases.
      """
      for expected_rval in [0, 1]:
          s = time.time()
          r, ur = ray.wait([gen], num_returns=1)
          print(time.time() - s)
          assert len(r) == 1
          assert ray.get(next(r[0])) == expected_rval
          assert len(ur) == 0
  
      # Should raise a stop iteration.
      for _ in range(3):
          s = time.time()
          r, ur = ray.wait([gen], num_returns=1)
          print(time.time() - s)
          assert len(r) == 1
          with pytest.raises(StopIteration):
              assert next(r[0]) == 0
          assert len(ur) == 0
  
      gen = f.remote(0)
      # Wait until the generator task finishes
      ray.get(gen._generator_ref)
      for i in range(2):
          r, ur = ray.wait([gen], timeout=0)
          assert len(r) == 1
          assert len(ur) == 0
          assert ray.get(next(r[0])) == i
  
      """
      Test the case ref is mixed with regular object ref.
      """
      gen = f.remote(0)
      ref = g.remote(3)
      ready, unready = [], [gen, ref]
      result_set = set()
      while unready:
          ready, unready = ray.wait(unready)
          print(ready, unready)
          assert len(ready) == 1
          for r in ready:
              if isinstance(r, ObjectRefGenerator):
                  try:
                      ref = next(r)
                      print(ref)
                      print(ray.get(ref))
                      result_set.add(ray.get(ref))
                  except StopIteration:
                      pass
                  else:
                      unready.append(r)
              else:
                  result_set.add(ray.get(r))
  
      assert result_set == {0, 1, 10}
  
      """
      Test timeout.
      """
      gen = f.remote(3)
      ref = g.remote(1)
      ready, unready = ray.wait([gen, ref], timeout=2)
      assert len(ready) == 1
      assert len(unready) == 1
  
      """
      Test num_returns
      """
      gen = f.remote(1)
      ref = g.remote(1)
      ready, unready = ray.wait([ref, gen], num_returns=2)
      assert len(ready) == 2
      assert len(unready) == 0
  
imports: import asyncio
  import pytest
  import numpy
  import sys
  import time
  from collections import Counter
  from starlette.responses import StreamingResponse
  from starlette.requests import Request
  from fastapi import FastAPI
  from ray import serve
  from pydantic import BaseModel
  import ray
  from ray._raylet import ObjectRefGenerator
  from ray.exceptions import WorkerCrashedError
  from ray._private.test_utils import run_string_as_driver_nonblocking
  from ray.util.state import list_actors
label_desc: Use image node_id 3 for calling a global function with example usage: test_generator_wait(shutdown_only) and returns: int