Skip to content

[tinker] Fix per-tenant blocking on forward/forward_backward calls#1920

Open
kailash109 wants to merge 1 commit into
NovaSky-AI:mainfrom
kailash109:opt-model-batch
Open

[tinker] Fix per-tenant blocking on forward/forward_backward calls#1920
kailash109 wants to merge 1 commit into
NovaSky-AI:mainfrom
kailash109:opt-model-batch

Conversation

@kailash109

@kailash109 kailash109 commented Jul 17, 2026

Copy link
Copy Markdown

Overview

forward_backward / forward futures 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

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread skyrl/tinker/engine.py
if error_results:
self._complete_futures(error_results)
if per_model:
grouped: dict[str, dict] = defaultdict(dict)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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.

Suggested change
grouped: dict[str, dict] = defaultdict(dict)
grouped: dict[str, dict[str, tuple[str, BaseModel]]] = defaultdict(dict)

Comment thread skyrl/tinker/engine.py
Comment on lines +757 to +758
for group in groups:
with log_timing(f"process_batch_requests({name}, n={len(group)})"):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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)})"):

@kailash109 kailash109 mentioned this pull request Jul 17, 2026
39 tasks
@erictang000 erictang000 self-assigned this Jul 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants