refactor: separate BatchTaskDispatcher from WorkflowExecutor

[garrett4wade]

Description

The main aim of this refactor is to enable code reuse in RolloutController. ``BatchTaskDispatcherextracts the common logic of task submission, waiting, and interactive submission-waiting loops. We can reuse the logic but with a new task that calls remote engines inRolloutController`.

There's no breaking change. Previous functionalities should not be affected.

Expected to be merged after #612

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not
  work as expected)
• Documentation update
• Code refactoring (no functional changes)
• Performance improvement
• Test coverage improvement

Checklist

• I have read the Contributing Guide
• I have run formatting tools (pre-commit or manual)
• I have run relevant unit tests and they pass
• I have added tests for new functionality
• I have updated documentation if needed
• My branch is up to date with main
• This PR introduces breaking changes (if yes, fill out details below)
• If this PR changes documentation, I have built and previewed it locally with
  jb build docs
• No critical issues raised by AI reviewers (/gemini review)

[gemini-code-assist]

Summary of Changes

Hello @garrett4wade, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request undertakes a significant refactoring effort by introducing a new BatchTaskDispatcher class, which extracts and centralizes the common asynchronous task management logic previously embedded within WorkflowExecutor. This change aims to enhance code reuse and improve the overall architecture, particularly benefiting components like RolloutController. Additionally, the StalenessManager has been updated to use a VersionProvider protocol for dynamic version retrieval, and the return types of wait methods across various engine APIs have been standardized to return lists of individual results.

Highlights

• Refactoring WorkflowExecutor: The core logic for managing task submission, waiting, and background threads has been extracted from WorkflowExecutor into a new, reusable BatchTaskDispatcher class. This significantly improves modularity and separation of concerns.
• New BatchTaskDispatcher Class: A new generic class, BatchTaskDispatcher, has been introduced to encapsulate the common patterns of asynchronous task dispatching, including interactive submission-waiting loops. This class is now responsible for the AsyncTaskRunner and its associated background threads.
• StalenessManager Update: The StalenessManager has been enhanced to accept a VersionProvider protocol in its constructor. This allows it to dynamically query the current version, making its capacity calculation more flexible and decoupling it from direct version passing.
• Standardized Wait Method Return Types: The wait methods across InferenceEngine, RemoteInfEngine, and specific engine implementations (SglangRemoteEngine, VllmRemoteEngine) now consistently return list[dict[str, Any] | None]. This change reflects that individual trajectory dictionaries are returned, with None indicating rejected trajectories.
• Prepare Batch Logic Centralization: The interactive submit-and-wait loop previously found within WorkflowExecutor.prepare_batch has been moved into the BatchTaskDispatcher's active_submit_and_wait method, centralizing this common pattern.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request effectively refactors the WorkflowExecutor by extracting a generic BatchTaskDispatcher. This is a great improvement that enhances modularity and will enable code reuse as intended. The separation of concerns makes the WorkflowExecutor much cleaner, especially the prepare_batch method.

I've found a few issues to address:

• A potential bug in the capacity calculation within active_submit_and_wait.
• A redundant variable assignment.
• An incorrect exception type in an except block.

Overall, this is a solid refactoring. Please see my detailed comments below.

[garrett4wade]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a significant and well-executed refactoring by separating the generic BatchTaskDispatcher from the WorkflowExecutor. This change greatly improves modularity, making the task submission and processing logic more reusable. The introduction of the VersionProvider protocol to decouple StalenessManager is also a commendable design choice. The overall architecture is cleaner and more maintainable.

One point of attention is that the change to the wait method's return type is a breaking change to the public API, which should be noted. I've also identified a minor performance improvement opportunity in prepare_batch where workflow resolution can be made more efficient. Overall, this is a high-quality contribution.

[rchardx]

In general LGTM!
I may submit a PR trying to elimite the busy loop in the current workflow implementaion.