Repository navigation
Conversation
…oundaries With pipeline_step_boundary on, the loop used to await the send-ahead and then run metric reduction, trajectory logging, the step log and on_step_end inline before starting step N+1, which froze dispatch/polling for ~2.5 s per boundary. Move that bookkeeping into a background task (reduction via asyncio.to_thread), keep gc.collect inline during mb1 compute, defer on_step_begin(N+1) until on_step_end(N) has run, await pending bookkeeping before the next weight sync and on loop errors. Flag-off path unchanged.
Summary of ChangesHello, I'm Gemini Code Assist1! 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 optimizes the pipelined training boundary by offloading heavy end-of-step bookkeeping tasks to a background thread. By decoupling metric reduction and logging from the main event loop, the system can now proceed to the next training step without waiting for the previous step's post-processing to finish. This change significantly reduces idle time in the pipeline, ensuring that critical stages like dispatch and prepack remain responsive while maintaining strict ordering and error propagation guarantees. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe 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
Customization To customize the 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 Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. 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. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces background bookkeeping for pipelined step boundaries in the RL program orchestrator, allowing step k+1 to start and send its second micro-batch without waiting for step k's metric reduction, logging, and host-memory release. These bookkeeping tasks are offloaded to a background task (_finish_step_bookkeeping) and run off the event loop using asyncio.to_thread. Corresponding unit tests have been added to verify this behavior and ensure callbacks maintain their correct order. The feedback suggests improving type safety by replacing Any with more specific type annotations for the background task parameter in _log_task_failure and the filtered_groups parameter in _finish_step_bookkeeping.
| pass | ||
|
|
||
|
|
||
| def _log_task_failure(task: asyncio.Task[Any]) -> None: |
There was a problem hiding this comment.
According to the repository style guide, strict and specific type hints should be used instead of Any whenever a more specific type can be expressed. Since _pending_bookkeeping is typed as asyncio.Task[None] | None, the task parameter in _log_task_failure can be typed as asyncio.Task[None].
| def _log_task_failure(task: asyncio.Task[Any]) -> None: | |
| def _log_task_failure(task: asyncio.Task[None]) -> None: |
References
- Type hints & Strict Typing: Always annotate function signatures, attributes, and data structures with concrete, specific types rather than Any. Avoid Any whenever a specific type can be expressed. (link)
| policy_training_time: float, | ||
| weight_sync_time: float, | ||
| metrics_fetch_time: float, | ||
| filtered_groups: list[Any], |
There was a problem hiding this comment.
According to the repository style guide, strict and specific type hints should be used instead of Any whenever a more specific type can be expressed. Since _drain_filtered_groups returns list[list[datatypes.TrajectoryItem]], filtered_groups can be typed as list[list[datatypes.TrajectoryItem]] instead of list[Any].
| filtered_groups: list[Any], | |
| filtered_groups: list[list[datatypes.TrajectoryItem]], |
References
- Type hints & Strict Typing: Always annotate function signatures, attributes, and data structures with concrete, specific types rather than Any. Avoid Any whenever a specific type can be expressed. (link)
This branch has not been deployed
2 participants
Stacked on #2877 (
parambole/l3-converge). Addresses review finding F1: withpipeline_step_boundaryon, end-of-step bookkeeping still ran on the event loop and blocked micro-batch 2 of step N+1.What changed (only when
pipeline_step_boundaryis on)_finish_step_bookkeeping):_log_consumed_trajectoriesinasyncio.to_thread;last_step_result, writes the Train step log, and callson_step_end._step += 1, then_release_window(). Dispatch, polling and prepack are no longer frozen for about 2.5 s at each boundary. On at-sph1-s1a, the median for metrics + log was 1.35 s and for GC 1.14 s.gc.collect()stays inline, at the point where the base loop runs it, which is during mb1 compute. Running it in the background holds the GIL and starves the next step's prepack feeds. That cost stages in the tests.on_step_begin(N+1)is deferred untilon_step_end(N)has run (_begin_step).policy_version. It is also awaited on loop errors, soon_step_end(N)runs before the error surfaces, and when training finishes. A failure in bookkeeping fails the run.Tests
BackgroundBookkeepingTest:rl_program_pipelined_boundary_test.py,*_robustness_test.py,rl_program_stage_before_update_test.py,rl_program_test.py,config_wiring_test.py: 653 passed (baseline [DO NOT MERGE] perf(orchestrator): opt-in pipelined step boundary that stages the next micro-batch before the deferred update #2877: 649).Supersedes #2885 / #2886.