[Serve] Let batch handler yield unhashable objects (ray-project#36353)

This change lets users yield unhashable objects in their `@serve.batch`-decorated generators.

python/ray/serve/batching.py

@@ -179,8 +179,8 @@ async def _consume_func_generator(
  if future is FINISHED_TOKEN:
  # This caller has already terminated.
  next_futures.append(FINISHED_TOKEN)
- elif result in {StopIteration, StopAsyncIteration}:
- # User's code returned SENTINEL.VALUE. No values left
+ elif result in [StopIteration, StopAsyncIteration]:
+ # User's code returned sentinel. No values left
  # for caller. Terminate iteration for caller.
  future.set_exception(StopAsyncIteration)
  next_futures.append(FINISHED_TOKEN)

python/ray/serve/tests/test_batching.py

@@ -412,12 +412,18 @@ class Textgen:
  @serve.batch(max_batch_size=4, batch_wait_timeout_s=1000)
  async def batch_handler(self, prompts: List[str]):
  for _ in range(NUM_YIELDS):
- prompt_responses = prompts
+ # Check that the batch handler can yield unhashable types
+ prompt_responses = [{"value": prompt} for prompt in prompts]
  yield prompt_responses
 
+ async def value_extractor(self, prompt_responses):
+ async for prompt_response in prompt_responses:
+ yield prompt_response["value"]
+
  async def __call__(self, request):
  prompt = request.query_params["prompt"]
- return StreamingResponse(self.batch_handler(prompt))
+ response_values = self.value_extractor(self.batch_handler(prompt))
+ return StreamingResponse(response_values)
 
  serve.run(Textgen.bind())