[tinker] Fix per-tenant blocking on forward/forward_backward calls#1920
[tinker] Fix per-tenant blocking on forward/forward_backward calls#1920kailash109 wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces a per_model option to process_batch_requests in skyrl/tinker/engine.py, allowing batch requests to be grouped and processed by model ID so that futures can be completed as soon as each model's sub-batch finishes. The review feedback suggests refining the type annotation of the grouped dictionary for better static analysis and enhancing the logging within the batch processing loop to include the specific model_id for improved observability.
| if error_results: | ||
| self._complete_futures(error_results) | ||
| if per_model: | ||
| grouped: dict[str, dict] = defaultdict(dict) |
There was a problem hiding this comment.
The type annotation for grouped is currently dict[str, dict], which is a bit too generic. Providing a more precise type annotation like dict[str, dict[str, tuple[str, BaseModel]]] will improve code maintainability and allow static analysis tools (such as mypy or pyright) to correctly type-check the groups list and ensure compatibility with the processor function signature.
| grouped: dict[str, dict] = defaultdict(dict) | |
| grouped: dict[str, dict[str, tuple[str, BaseModel]]] = defaultdict(dict) |
| for group in groups: | ||
| with log_timing(f"process_batch_requests({name}, n={len(group)})"): |
There was a problem hiding this comment.
In a multi-tenant environment where multiple models are processed, logging the timing of each model's sub-batch with its specific model_id would greatly improve observability and debugging. We can extract the model_id from the first item in each non-empty group and include it in the log_timing message.
for group in groups:
model_id = next(iter(group.values()))[0] if group else "unknown"
with log_timing(f"process_batch_requests({name}, model={model_id}, n={len(group)})"):ec71a8b to
350dd89
Compare
Overview
forward_backward/forwardfutures are currently completed only after entire cross-model batch finishes, so each client blocks on all other tenant's sub-batches even though backend executes per-model sub-batches serially anyways.Add per_model mode to
process_batch_requests(enabled for forward and forward_backward), valid requests grouped by model_id and each model's futures are completed as soon as its own sub-batch returns, meaning clients unblock after their own batch finishes and can pipeline next request while other tenants' sub-batches are still running.GPU execution order is unchanged, and _sleep_inference_engines() calls across per-model sub-batches are no-ops (vllm early return if already sleeping).
Validation: Tinker cpu tests pass