wehub-resource-sync
752 lines
26 KiB
Python
752 lines
26 KiB
Python
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__]))
|