chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,751 @@
|
||||
import asyncio
|
||||
import math
|
||||
from collections.abc import Callable
|
||||
from concurrent.futures.thread import ThreadPoolExecutor
|
||||
from functools import partial
|
||||
from threading import Thread
|
||||
from typing import List, Optional
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import FastAPI, Request
|
||||
from starlette.responses import StreamingResponse
|
||||
|
||||
from ray import serve
|
||||
from ray._common.test_utils import SignalActor, async_wait_for_condition
|
||||
from ray.serve._private.test_utils import get_application_url
|
||||
from ray.serve.batching import _RuntimeSummaryStatistics
|
||||
from ray.serve.context import (
|
||||
_get_serve_batch_request_context,
|
||||
_get_serve_request_context,
|
||||
)
|
||||
|
||||
|
||||
def test_batching(serve_instance):
|
||||
@serve.deployment
|
||||
class BatchingExample:
|
||||
def __init__(self):
|
||||
self.count = 0
|
||||
|
||||
@serve.batch(max_batch_size=5, batch_wait_timeout_s=1)
|
||||
async def handle_batch(self, requests):
|
||||
self.count += 1
|
||||
batch_size = len(requests)
|
||||
return [self.count] * batch_size
|
||||
|
||||
async def __call__(self, request):
|
||||
return await self.handle_batch(request)
|
||||
|
||||
handle = serve.run(BatchingExample.bind())
|
||||
|
||||
result_list = [handle.remote(1) for _ in range(20)]
|
||||
# since count is only updated per batch of queries
|
||||
# If there atleast one __call__ fn call with batch size greater than 1
|
||||
# counter result will always be less than 20
|
||||
assert max([r.result() for r in result_list]) < 20
|
||||
|
||||
|
||||
def test_concurrent_batching(serve_instance):
|
||||
BATCHES_IN_FLIGHT = 2
|
||||
MAX_BATCH_SIZE = 5
|
||||
BATCH_WAIT_TIMEOUT_S = 1
|
||||
MAX_REQUESTS_IN_FLIGHT = BATCHES_IN_FLIGHT * MAX_BATCH_SIZE
|
||||
|
||||
@serve.deployment(max_ongoing_requests=MAX_REQUESTS_IN_FLIGHT * 2)
|
||||
class BatchingExample:
|
||||
def __init__(self):
|
||||
self.n_batches_in_flight = 0
|
||||
self.n_requests_in_flight = 0
|
||||
|
||||
@serve.batch(
|
||||
max_batch_size=MAX_BATCH_SIZE,
|
||||
batch_wait_timeout_s=BATCH_WAIT_TIMEOUT_S,
|
||||
max_concurrent_batches=BATCHES_IN_FLIGHT,
|
||||
)
|
||||
async def handle_batch(self, requests):
|
||||
self.n_batches_in_flight += 1
|
||||
self.n_requests_in_flight += len(requests)
|
||||
await asyncio.sleep(0.5)
|
||||
out = [
|
||||
(req_idx, self.n_batches_in_flight, self.n_requests_in_flight)
|
||||
for req_idx in requests
|
||||
]
|
||||
await asyncio.sleep(0.5)
|
||||
self.n_requests_in_flight -= len(requests)
|
||||
self.n_batches_in_flight -= 1
|
||||
return out
|
||||
|
||||
async def __call__(self, request):
|
||||
return await self.handle_batch(request)
|
||||
|
||||
handle = serve.run(BatchingExample.bind())
|
||||
|
||||
idxs = set(range(20))
|
||||
result_futures = [handle.remote(i) for i in idxs]
|
||||
result_list = [future.result() for future in result_futures]
|
||||
|
||||
out_idxs = set()
|
||||
for idx, batches_in_flight, requests_in_flight in result_list:
|
||||
out_idxs.add(idx)
|
||||
assert (
|
||||
batches_in_flight == BATCHES_IN_FLIGHT
|
||||
), f"Should have been {BATCHES_IN_FLIGHT} batches in flight at all times, got {batches_in_flight}"
|
||||
assert (
|
||||
requests_in_flight == MAX_REQUESTS_IN_FLIGHT
|
||||
), f"Should have been {MAX_REQUESTS_IN_FLIGHT} requests in flight at all times, got {requests_in_flight}"
|
||||
|
||||
assert idxs == out_idxs, "All requests should be processed"
|
||||
|
||||
|
||||
def test_batching_exception(serve_instance):
|
||||
@serve.deployment
|
||||
class NoListReturned:
|
||||
def __init__(self):
|
||||
self.count = 0
|
||||
|
||||
@serve.batch(max_batch_size=5)
|
||||
async def handle_batch(self, requests):
|
||||
return len(requests)
|
||||
|
||||
async def __call__(self, request):
|
||||
return await self.handle_batch(request)
|
||||
|
||||
# Set the max batch size.
|
||||
handle = serve.run(NoListReturned.bind())
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
assert handle.remote(1).result()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_generator_streaming_response_integration_test(serve_instance):
|
||||
NUM_YIELDS = 10
|
||||
|
||||
@serve.deployment
|
||||
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):
|
||||
# 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"]
|
||||
response_values = self.value_extractor(self.batch_handler(prompt))
|
||||
return StreamingResponse(response_values)
|
||||
|
||||
serve.run(Textgen.bind())
|
||||
|
||||
prompt_prefix = "hola"
|
||||
url = f"{get_application_url()}/?prompt={prompt_prefix}"
|
||||
with ThreadPoolExecutor() as pool:
|
||||
futs = [pool.submit(partial(httpx.get, url + str(idx))) for idx in range(4)]
|
||||
responses = [fut.result() for fut in futs]
|
||||
|
||||
for idx, response in enumerate(responses):
|
||||
assert response.status_code == 200
|
||||
assert response.text == "".join([prompt_prefix + str(idx)] * NUM_YIELDS)
|
||||
|
||||
|
||||
def test_batching_client_dropped_unary(serve_instance):
|
||||
"""Test unary batching with clients that drops the connection.
|
||||
|
||||
After requests are dropped. The next request should succeed.
|
||||
"""
|
||||
|
||||
@serve.deployment
|
||||
class ModelUnary:
|
||||
@serve.batch(max_batch_size=5)
|
||||
async def handle_batch(self, requests):
|
||||
await asyncio.sleep(0.05)
|
||||
return ["fake-response" for _ in range(len(requests))]
|
||||
|
||||
async def __call__(self, request):
|
||||
return await self.handle_batch(request)
|
||||
|
||||
serve.run(ModelUnary.bind())
|
||||
|
||||
url = f"{get_application_url()}/"
|
||||
|
||||
# Sending requests with clients that drops the connection.
|
||||
for _ in range(3):
|
||||
with pytest.raises(httpx.ReadTimeout):
|
||||
httpx.get(url, timeout=0.005)
|
||||
|
||||
# The following request should succeed.
|
||||
resp = httpx.get(url, timeout=1)
|
||||
assert resp.status_code == 200
|
||||
assert resp.text == "fake-response"
|
||||
|
||||
|
||||
def test_batching_client_dropped_streaming(serve_instance):
|
||||
"""Test streaming batching with clients that drops the connection.
|
||||
|
||||
After requests are dropped. The next request should succeed.
|
||||
"""
|
||||
|
||||
@serve.deployment
|
||||
class ModelStreaming:
|
||||
@serve.batch(max_batch_size=3)
|
||||
async def handle_batch(self, requests):
|
||||
await asyncio.sleep(0.05)
|
||||
for i in range(10):
|
||||
yield [str(i) for _ in range(len(requests))]
|
||||
|
||||
async def __call__(self, request):
|
||||
return StreamingResponse(self.handle_batch(request))
|
||||
|
||||
serve.run(ModelStreaming.bind())
|
||||
|
||||
url = "http://localhost:8000/"
|
||||
|
||||
# Sending requests with clients that drops the connection.
|
||||
for _ in range(3):
|
||||
with pytest.raises((httpx.ReadTimeout, httpx.ConnectError)):
|
||||
httpx.get(url, timeout=0.005)
|
||||
|
||||
# The following request should succeed.
|
||||
resp = httpx.get(url, timeout=1)
|
||||
assert resp.status_code == 200
|
||||
assert resp.text == "0123456789"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("max_concurrent_batches", [1, 10])
|
||||
@pytest.mark.parametrize("max_batch_size", [1, 10])
|
||||
@pytest.mark.parametrize("n_requests", [1, 10])
|
||||
async def test_observability_helpers(
|
||||
serve_instance, n_requests: int, max_batch_size: int, max_concurrent_batches: int
|
||||
) -> None:
|
||||
"""Checks observability helper methods that are used for batching.
|
||||
|
||||
Tests three observability helper methods:
|
||||
* _get_curr_iteration_start_times: gets the current iteration's start
|
||||
time.
|
||||
* _is_batching_task_alive: returns whether the batch-handler task is
|
||||
alive.
|
||||
* _get_handling_task_stack: returns the stack for the batch-handler task.
|
||||
"""
|
||||
|
||||
signal_actor = SignalActor.remote()
|
||||
|
||||
@serve.deployment(
|
||||
name="batcher", max_ongoing_requests=max_concurrent_batches * max_batch_size
|
||||
)
|
||||
class Batcher:
|
||||
@serve.batch(
|
||||
max_batch_size=max_batch_size,
|
||||
max_concurrent_batches=max_concurrent_batches,
|
||||
batch_wait_timeout_s=0.1,
|
||||
)
|
||||
async def handle_batch(self, requests):
|
||||
await signal_actor.wait.remote() # wait until the outer signal actor is released
|
||||
return [0] * len(requests)
|
||||
|
||||
async def __call__(self, request):
|
||||
return await self.handle_batch(request)
|
||||
|
||||
async def _get_curr_iteration_start_times(self) -> _RuntimeSummaryStatistics:
|
||||
return self.handle_batch._get_curr_iteration_start_times()
|
||||
|
||||
async def _is_batching_task_alive(self) -> bool:
|
||||
return await self.handle_batch._is_batching_task_alive()
|
||||
|
||||
async def _get_handling_task_stack(self) -> Optional[str]:
|
||||
return await self.handle_batch._get_handling_task_stack()
|
||||
|
||||
serve.run(target=Batcher.bind(), name="app_name")
|
||||
handle = serve.get_deployment_handle(deployment_name="batcher", app_name="app_name")
|
||||
|
||||
assert await handle._is_batching_task_alive.remote()
|
||||
|
||||
min_num_batches = min(
|
||||
math.ceil(n_requests / max_batch_size), max_concurrent_batches
|
||||
)
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
tasks1 = await send_k_requests(
|
||||
signal_actor,
|
||||
n_requests,
|
||||
min_num_batches,
|
||||
app_name="app_name",
|
||||
client=client,
|
||||
)
|
||||
prev_iter_times = await handle._get_curr_iteration_start_times.remote()
|
||||
await signal_actor.send.remote() # unblock the batch handler now that we have the iter times
|
||||
|
||||
assert len(prev_iter_times.start_times) >= min_num_batches
|
||||
assert len(await handle._get_handling_task_stack.remote()) is not None
|
||||
assert await handle._is_batching_task_alive.remote()
|
||||
|
||||
tasks2 = await send_k_requests(
|
||||
signal_actor,
|
||||
n_requests,
|
||||
min_num_batches,
|
||||
app_name="app_name",
|
||||
client=client,
|
||||
)
|
||||
new_iter_times = await handle._get_curr_iteration_start_times.remote()
|
||||
await signal_actor.send.remote() # unblock the batch handler now that we have the iter times
|
||||
|
||||
assert len(new_iter_times.start_times) >= min_num_batches
|
||||
assert len(await handle._get_handling_task_stack.remote()) is not None
|
||||
assert await handle._is_batching_task_alive.remote()
|
||||
|
||||
assert new_iter_times.min_start_time > prev_iter_times.max_start_time
|
||||
|
||||
# Cancel and await all tasks to avoid "Task exception was never retrieved" warning.
|
||||
# We don't need the HTTP responses, just need to clean up the tasks properly.
|
||||
for task in tasks1 + tasks2:
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks1, *tasks2, return_exceptions=True)
|
||||
|
||||
|
||||
async def send_k_requests(
|
||||
signal_actor: SignalActor,
|
||||
k: int,
|
||||
min_num_batches: float,
|
||||
app_name: str,
|
||||
client: httpx.AsyncClient,
|
||||
) -> List[asyncio.Task]:
|
||||
"""Send k requests and wait until at least min_num_batches are waiting.
|
||||
|
||||
Returns the list of request tasks so they can be awaited by the caller
|
||||
after unblocking the batch handler.
|
||||
"""
|
||||
await signal_actor.send.remote(True) # type: ignore[attr-defined]
|
||||
tasks = []
|
||||
for _ in range(k):
|
||||
tasks.append(
|
||||
asyncio.create_task(
|
||||
client.get(f"{get_application_url(app_name=app_name)}/")
|
||||
)
|
||||
)
|
||||
await wait_for_n_waiters(
|
||||
signal_actor, lambda num_waiters: num_waiters >= min_num_batches
|
||||
)
|
||||
return tasks
|
||||
|
||||
|
||||
async def wait_for_n_waiters(
|
||||
signal_actor: SignalActor, condition: Callable[[int], bool]
|
||||
) -> None:
|
||||
async def poll() -> bool:
|
||||
num_waiters: int = await signal_actor.cur_num_waiters.remote() # type: ignore[attr-defined]
|
||||
return condition(num_waiters)
|
||||
|
||||
return await async_wait_for_condition(poll)
|
||||
|
||||
|
||||
def test_batching_request_context(serve_instance):
|
||||
"""Test that _get_serve_batch_request_context() works correctly with batching.
|
||||
|
||||
With 6 requests and max_batch_size=3, Serve should create 2 batches processed in parallel.
|
||||
Each batch should have access to the request contexts of all requests in that batch,
|
||||
and context should be properly unset after processing.
|
||||
"""
|
||||
|
||||
@serve.deployment(max_ongoing_requests=10)
|
||||
class BatchContextTester:
|
||||
def __init__(self):
|
||||
self.batch_results = []
|
||||
|
||||
@serve.batch(
|
||||
max_batch_size=3, batch_wait_timeout_s=1.0, max_concurrent_batches=2
|
||||
)
|
||||
async def handle_batch(self, batch):
|
||||
# Store results for verification
|
||||
batch_result = {
|
||||
"batch_size": len(batch),
|
||||
"batch_request_contexts": _get_serve_batch_request_context(),
|
||||
"current_request_context": _get_serve_request_context(),
|
||||
}
|
||||
self.batch_results.append(batch_result)
|
||||
|
||||
return ["ok" for _ in range(len(batch))]
|
||||
|
||||
async def __call__(self, request):
|
||||
return await self.handle_batch(1)
|
||||
|
||||
async def get_results(self):
|
||||
return self.batch_results
|
||||
|
||||
handle = serve.run(BatchContextTester.bind())
|
||||
|
||||
def do_request():
|
||||
"""Make a request with a specific request ID."""
|
||||
url = get_application_url()
|
||||
r = httpx.post(f"{url}/")
|
||||
r.raise_for_status()
|
||||
|
||||
# Launch 6 requests. Expect 2 batches of 3 requests each.
|
||||
threads = [Thread(target=do_request) for _ in range(6)]
|
||||
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
# Get results from the deployment
|
||||
batch_results = handle.get_results.remote().result()
|
||||
|
||||
# Verify each batch has correct size and context
|
||||
total_requests_processed = 0
|
||||
request_ids_in_batch_context = set()
|
||||
|
||||
for result in batch_results:
|
||||
# Batch context should contain all 3 request contexts
|
||||
assert (
|
||||
len(result["batch_request_contexts"]) == 3
|
||||
), f"Expected 3 contexts in batch, got {result['batch_request_contexts']}"
|
||||
req_ids_in_batch_context = [
|
||||
ctx.request_id for ctx in result["batch_request_contexts"]
|
||||
]
|
||||
assert (
|
||||
len(req_ids_in_batch_context) == 3
|
||||
), f"Expected 3 batch request IDs, got {len(req_ids_in_batch_context)}"
|
||||
request_ids_in_batch_context.update(req_ids_in_batch_context)
|
||||
|
||||
# Current request context read within the batcher should be a default empty context.
|
||||
current_request_context = result["current_request_context"]
|
||||
assert current_request_context.request_id == ""
|
||||
assert current_request_context.route == ""
|
||||
assert current_request_context.app_name == ""
|
||||
assert current_request_context.multiplexed_model_id == ""
|
||||
|
||||
total_requests_processed += result["batch_size"]
|
||||
|
||||
# Verify all 6 requests were processed
|
||||
assert (
|
||||
total_requests_processed == 6
|
||||
), f"Expected 6 total requests processed, got {total_requests_processed}"
|
||||
assert (
|
||||
len(request_ids_in_batch_context) == 6
|
||||
), f"Expected 6 unique request IDs, got {len(request_ids_in_batch_context)}"
|
||||
|
||||
|
||||
def test_batch_size_fn_simple(serve_instance):
|
||||
"""Test batch_size_fn with a simple custom batch size metric."""
|
||||
|
||||
@serve.deployment
|
||||
class BatchSizeFnExample:
|
||||
def __init__(self):
|
||||
self.batches_received = []
|
||||
|
||||
@serve.batch(
|
||||
max_batch_size=100, # Set based on total size, not count
|
||||
batch_wait_timeout_s=0.5,
|
||||
batch_size_fn=lambda items: sum(item["size"] for item in items),
|
||||
)
|
||||
async def handle_batch(self, requests: List):
|
||||
# Record the batch for verification
|
||||
self.batches_received.append(requests)
|
||||
# Return results
|
||||
return [req["value"] * 2 for req in requests]
|
||||
|
||||
async def __call__(self, request):
|
||||
return await self.handle_batch(request)
|
||||
|
||||
def get_batches(self):
|
||||
return self.batches_received
|
||||
|
||||
handle = serve.run(BatchSizeFnExample.bind())
|
||||
|
||||
# Send requests with different sizes
|
||||
# Request 1: size=30, value=1
|
||||
# Request 2: size=40, value=2
|
||||
# Request 3: size=20, value=3
|
||||
# Request 4: size=25, value=4
|
||||
# Total of first 3 = 90 (< 100), but adding 4th would be 115 (> 100)
|
||||
requests = [
|
||||
{"size": 30, "value": 1},
|
||||
{"size": 40, "value": 2},
|
||||
{"size": 20, "value": 3},
|
||||
{"size": 25, "value": 4},
|
||||
]
|
||||
|
||||
result_futures = [handle.remote(req) for req in requests]
|
||||
results = [future.result() for future in result_futures]
|
||||
|
||||
# Verify results are correct
|
||||
assert results == [2, 4, 6, 8]
|
||||
|
||||
# Verify batching behavior
|
||||
batches = handle.get_batches.remote().result()
|
||||
# Should have created at least one batch
|
||||
assert len(batches) > 0
|
||||
|
||||
|
||||
def test_batch_size_fn_graph_nodes(serve_instance):
|
||||
"""Test batch_size_fn with a GNN-style use case (batching by total nodes)."""
|
||||
|
||||
class Graph:
|
||||
def __init__(self, num_nodes: int, graph_id: int):
|
||||
self.num_nodes = num_nodes
|
||||
self.graph_id = graph_id
|
||||
|
||||
@serve.deployment
|
||||
class GraphBatcher:
|
||||
def __init__(self):
|
||||
self.batch_sizes = []
|
||||
|
||||
@serve.batch(
|
||||
max_batch_size=100, # Max 100 nodes per batch
|
||||
batch_wait_timeout_s=0.5,
|
||||
batch_size_fn=lambda graphs: sum(g.num_nodes for g in graphs),
|
||||
)
|
||||
async def process_graphs(self, graphs: List[Graph]):
|
||||
# Record batch size (total nodes)
|
||||
total_nodes = sum(g.num_nodes for g in graphs)
|
||||
self.batch_sizes.append(total_nodes)
|
||||
# Return graph_id * num_nodes as result
|
||||
return [g.graph_id * g.num_nodes for g in graphs]
|
||||
|
||||
async def __call__(self, graph):
|
||||
return await self.process_graphs(graph)
|
||||
|
||||
def get_batch_sizes(self):
|
||||
return self.batch_sizes
|
||||
|
||||
handle = serve.run(GraphBatcher.bind())
|
||||
|
||||
# Create graphs with different node counts
|
||||
# Graph 1: 30 nodes, Graph 2: 40 nodes, Graph 3: 35 nodes, Graph 4: 50 nodes
|
||||
# First 3 total = 105 nodes (> 100), so should be 2 batches
|
||||
graphs = [
|
||||
Graph(num_nodes=30, graph_id=1),
|
||||
Graph(num_nodes=40, graph_id=2),
|
||||
Graph(num_nodes=35, graph_id=3),
|
||||
Graph(num_nodes=50, graph_id=4),
|
||||
]
|
||||
|
||||
result_futures = [handle.remote(g) for g in graphs]
|
||||
results = [future.result() for future in result_futures]
|
||||
|
||||
# Verify results
|
||||
assert results == [30, 80, 105, 200]
|
||||
|
||||
# Verify batch sizes respect the limit
|
||||
batch_sizes = handle.get_batch_sizes.remote().result()
|
||||
for batch_size in batch_sizes:
|
||||
# Each batch should have <= 100 nodes
|
||||
assert batch_size <= 100, f"Batch size {batch_size} exceeds limit of 100"
|
||||
|
||||
|
||||
def test_batch_size_fn_token_count(serve_instance):
|
||||
"""Test batch_size_fn with an NLP-style use case (batching by total tokens)."""
|
||||
|
||||
@serve.deployment
|
||||
class TokenBatcher:
|
||||
@serve.batch(
|
||||
max_batch_size=1000, # Max 1000 tokens per batch
|
||||
batch_wait_timeout_s=0.5,
|
||||
batch_size_fn=lambda sequences: sum(len(s.split()) for s in sequences),
|
||||
)
|
||||
async def process_sequences(self, sequences: List[str]):
|
||||
# Return word count for each sequence
|
||||
return [len(s.split()) for s in sequences]
|
||||
|
||||
async def __call__(self, sequence):
|
||||
return await self.process_sequences(sequence)
|
||||
|
||||
handle = serve.run(TokenBatcher.bind())
|
||||
|
||||
# Create sequences with different token counts
|
||||
sequences = [
|
||||
"This is a short sequence", # 5 tokens
|
||||
"This is a much longer sequence with many more words in it", # 12 tokens
|
||||
"Short", # 1 token
|
||||
"A B C D E F G H I J", # 10 tokens
|
||||
]
|
||||
|
||||
result_futures = [handle.remote(s) for s in sequences]
|
||||
results = [future.result() for future in result_futures]
|
||||
|
||||
# Verify results are correct
|
||||
assert results == [5, 12, 1, 10]
|
||||
|
||||
|
||||
def test_batch_size_fn_validation():
|
||||
"""Test that batch_size_fn validation works correctly."""
|
||||
from ray.serve.batching import batch
|
||||
|
||||
# Test with non-callable batch_size_fn
|
||||
with pytest.raises(TypeError, match="batch_size_fn must be a callable or None"):
|
||||
|
||||
@batch(batch_size_fn="not_a_function")
|
||||
async def my_batch_handler(items):
|
||||
return items
|
||||
|
||||
|
||||
def test_batch_size_fn_default_behavior(serve_instance):
|
||||
"""Test that default behavior (batch_size_fn=None) still works as expected."""
|
||||
|
||||
@serve.deployment
|
||||
class DefaultBatcher:
|
||||
@serve.batch(max_batch_size=5, batch_wait_timeout_s=0.5)
|
||||
async def handle_batch(self, requests):
|
||||
return [r * 2 for r in requests]
|
||||
|
||||
async def __call__(self, request):
|
||||
return await self.handle_batch(request)
|
||||
|
||||
handle = serve.run(DefaultBatcher.bind())
|
||||
|
||||
# Send 10 requests
|
||||
result_futures = [handle.remote(i) for i in range(10)]
|
||||
results = [future.result() for future in result_futures]
|
||||
|
||||
# Verify all results are correct
|
||||
assert results == [i * 2 for i in range(10)]
|
||||
|
||||
|
||||
def test_batch_size_fn_oversized_item_raises_error(serve_instance):
|
||||
app = FastAPI()
|
||||
|
||||
@serve.deployment
|
||||
@serve.ingress(app)
|
||||
class OversizedItemBatcher:
|
||||
@serve.batch(
|
||||
max_batch_size=10,
|
||||
batch_wait_timeout_s=0.5,
|
||||
batch_size_fn=lambda items: sum(item["size"] for item in items),
|
||||
)
|
||||
async def handle_batch(self, requests: List):
|
||||
return [req["value"] for req in requests]
|
||||
|
||||
@app.post("/")
|
||||
async def f(self, request: Request):
|
||||
body = await request.json()
|
||||
return await self.handle_batch(body)
|
||||
|
||||
serve.run(OversizedItemBatcher.bind())
|
||||
|
||||
# Send a request with size > max_batch_size (15 > 10)
|
||||
# This should return a 500 error with RuntimeError message
|
||||
url = f"{get_application_url(use_localhost=True)}/"
|
||||
response = httpx.post(url, json={"size": 15, "value": "too_large"}, timeout=5)
|
||||
|
||||
assert response.status_code == 500
|
||||
|
||||
|
||||
def test_batch_size_fn_deferred_item_processed(serve_instance):
|
||||
@serve.deployment(max_ongoing_requests=15)
|
||||
class DeferredItemBatcher:
|
||||
def __init__(self):
|
||||
self.batch_sizes = []
|
||||
|
||||
@serve.batch(
|
||||
max_batch_size=10,
|
||||
batch_wait_timeout_s=0.5,
|
||||
batch_size_fn=lambda items: sum(item["size"] for item in items),
|
||||
)
|
||||
async def handle_batch(self, requests: List):
|
||||
# Record actual batch sizes for verification
|
||||
total_size = sum(req["size"] for req in requests)
|
||||
self.batch_sizes.append(total_size)
|
||||
return [req["value"] for req in requests]
|
||||
|
||||
async def __call__(self, request):
|
||||
return await self.handle_batch(request)
|
||||
|
||||
def get_batch_sizes(self):
|
||||
return self.batch_sizes
|
||||
|
||||
handle = serve.run(DeferredItemBatcher.bind())
|
||||
|
||||
# Send requests where some will need to be deferred:
|
||||
# Request 1: size=6 (fits)
|
||||
# Request 2: size=6 (would make total 12 > 10, deferred)
|
||||
# Request 3: size=3 (fits with request 1, total 9)
|
||||
# Request 4: size=4 (would make total 13 > 10, deferred)
|
||||
requests = [
|
||||
{"size": 6, "value": "a"},
|
||||
{"size": 6, "value": "b"},
|
||||
{"size": 3, "value": "c"},
|
||||
{"size": 4, "value": "d"},
|
||||
]
|
||||
|
||||
result_futures = [handle.remote(req) for req in requests]
|
||||
results = [future.result() for future in result_futures]
|
||||
|
||||
# All requests should be processed successfully
|
||||
assert set(results) == {"a", "b", "c", "d"}
|
||||
|
||||
# Verify total size processed equals sum of all request sizes
|
||||
batch_sizes = handle.get_batch_sizes.remote().result()
|
||||
total_processed = sum(batch_sizes)
|
||||
expected_total = sum(req["size"] for req in requests) # 6 + 6 + 3 + 4 = 19
|
||||
assert (
|
||||
total_processed == expected_total
|
||||
), f"Total processed {total_processed} != expected {expected_total}"
|
||||
|
||||
|
||||
def test_batch_size_fn_mixed_normal_and_large_items(serve_instance):
|
||||
@serve.deployment
|
||||
class MixedSizeBatcher:
|
||||
def __init__(self):
|
||||
self.batches_processed = []
|
||||
|
||||
@serve.batch(
|
||||
max_batch_size=100,
|
||||
batch_wait_timeout_s=0.5,
|
||||
batch_size_fn=lambda items: sum(item["tokens"] for item in items),
|
||||
)
|
||||
async def handle_batch(self, requests: List):
|
||||
batch_info = {
|
||||
"total_tokens": sum(req["tokens"] for req in requests),
|
||||
"num_items": len(requests),
|
||||
}
|
||||
self.batches_processed.append(batch_info)
|
||||
return [f"processed_{req['id']}" for req in requests]
|
||||
|
||||
async def __call__(self, request):
|
||||
return await self.handle_batch(request)
|
||||
|
||||
def get_batches(self):
|
||||
return self.batches_processed
|
||||
|
||||
handle = serve.run(MixedSizeBatcher.bind())
|
||||
|
||||
# Mix of small and larger items
|
||||
requests = [
|
||||
{"id": 1, "tokens": 10}, # Small
|
||||
{"id": 2, "tokens": 20}, # Small
|
||||
{"id": 3, "tokens": 50}, # Medium
|
||||
{"id": 4, "tokens": 15}, # Small
|
||||
{"id": 5, "tokens": 90}, # Large (near limit)
|
||||
{"id": 6, "tokens": 5}, # Small
|
||||
]
|
||||
|
||||
result_futures = [handle.remote(req) for req in requests]
|
||||
results = [future.result() for future in result_futures]
|
||||
|
||||
# All requests should be processed
|
||||
expected_results = [f"processed_{i}" for i in range(1, 7)]
|
||||
assert set(results) == set(expected_results)
|
||||
|
||||
# Verify total tokens processed equals sum of all request tokens
|
||||
batches = handle.get_batches.remote().result()
|
||||
total_tokens_processed = sum(batch["total_tokens"] for batch in batches)
|
||||
expected_total = sum(req["tokens"] for req in requests) # 10+20+50+15+90+5 = 190
|
||||
assert (
|
||||
total_tokens_processed == expected_total
|
||||
), f"Total tokens {total_tokens_processed} != expected {expected_total}"
|
||||
|
||||
# Verify total items processed equals number of requests
|
||||
total_items = sum(batch["num_items"] for batch in batches)
|
||||
assert total_items == len(
|
||||
requests
|
||||
), f"Total items {total_items} != expected {len(requests)}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", "-s", __file__]))
|
||||
Reference in New Issue
Block a user