[Serve] Ray Serve with multiplexing with batching

[manickavela29]

Description

PR introduces multiplexed model serving with batching support to Ray Serve, enabling efficient hosting of multiple models on a single replica with automatic batching capabilities. This feature significantly improves resource utilization and reduces deployment costs for multi-model scenarios.

Performance Improvements
Asynchronous cleanup: Non-blocking cleanup of completed requests
Batching-aware routing: Prioritizes replicas with pending requests for the same model
Efficient request tracking: Optimized data structures for tracking pending requests by model

Related issues

1. [Serve] model multiplexing and batching does not work together #56633
2. [Serve] group requests by model_id in Model Multiplexing #50695

Additional information

Sample Code

@serve.deployment
class MultiModelService:
    def __init__(self):
        # Initialize multiplexer with batching
        self.multiplexer = ModelMultiplexer(
            model_load_func=self.load_model,
            max_num_models_per_replica=3,
            enable_batching=True,
            max_batch_size=8,
            batch_wait_timeout_s=0.01)
    
    async def load_model(self, model_id: str):
        # Your model loading logic
        return load_your_model(model_id)
    
    async def __call__(self, request):
        model_id = request.get("model_id")
        input_data = request.get("input")
        return await self.multiplexer.predict(input_data, model_id)

Test files
test_multiplex_batching.py - Core functionality tests
test_multiplex_batching_router.py - Request routing tests
test_multiplex_batching_utils.py - Utilities and fixtures

[gemini-code-assist]

Code Review

This pull request introduces batching capabilities to the multiplexing feature in Ray Serve, which is a valuable enhancement for performance. The changes include updates to the multiplexed decorator, the _ModelMultiplexWrapper, and the request router to support batching-aware routing. The implementation is mostly solid, but there are a few key issues to address. The batching-aware routing logic in the request router appears incomplete, as it identifies opportunities for batching but doesn't act on them. Additionally, there's a potential performance issue with how completed requests are cleaned up. The new tests are comprehensive, but one of them has flawed logic that doesn't correctly test the batching mechanism, and there are some leftover debugging statements that should be removed. Overall, this is a great feature addition, and with these fixes, it will be a strong contribution.

[cursor]

Bug: Batching Optimization Fails due to Incorrect Future Check

The batching-aware routing logic is ineffective. The _get_pending_requests_for_model method filters out completed requests, but the subsequent batching-friendly replica selection logic incorrectly checks for pending_req.future.done(). This condition is always false, preventing the batching optimization from executing.

 

[cursor]

Bug: Misuse Breaks Lazy Batch Queue Wrapper API

The _LazyBatchQueueWrapper is used incorrectly by directly accessing its internal _queue attribute and its queue attribute. This bypasses the wrapper's lazy initialization and violates its intended API, which can lead to AttributeError or other runtime errors when handling batched requests or during shutdown.

Additional Locations (2)

• python/ray/serve/multiplex.py#L215-L216
• python/ray/serve/multiplex.py#L233-L234