[tinker] Fix per-tenant blocking on forward/forward_backward calls - #1920
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
|
Hey Kailash, thanks so much for putting this up! The speedup looks good on the SkyRL-Train path as we execute per-model sub-batches serially! However, with the JAX backend I believe we use a fused forward/backward that runs sequences from different tenants at once, which causes slowdown if we split the batch by model. Could we gate the per_model flag on this instead of unconditionally setting it in |
can be ignored for now - will be removing batching logic across models for forward/backward in the jax path |
|
@kailash109 can you just resolve the test merge conflict then this should be good to merge |
|
resolved! |
|
@kailash109 Thanks a lot for fixing this, this is great! |
`main` is currently failing pre-commit's black hook: NovaSky-AI#1920 added a test immediately above the `forward_backward_payload` helper introduced by NovaSky-AI#1992 and left one blank line between them instead of two, so `check_code_quality` fails on every open PR rather than just this one. Whitespace only. Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
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