feat: use subprocess to fork colocated workers

[garrett4wade]

Description

This PR adds an option to create colocated workers via subprocess, instead of hosting them in the same process.

This is primarily because that the python multi-threading performance is too low. If we allocate eval rollout and rollout into the same worker, there will be two threads executing the workflow, leading to congestion.

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 significantly enhances the system's ability to manage colocated workers by introducing a fork mechanism that spawns new worker processes via subprocess calls. This change is primarily driven by the need to overcome Python's Global Interpreter Lock (GIL) limitations, which can hinder performance when multiple threads execute within the same process, especially for I/O-bound or CPU-bound tasks. By allowing colocated workers to run in separate processes, the system gains improved performance, better resource isolation, and more robust error handling, without requiring entirely separate resource allocations for each worker.

Highlights

• New Forking Mechanism for Colocated Workers: Introduced a fork option in SchedulingStrategy that allows colocated workers to run in separate subprocesses rather than sharing the same process. This addresses Python's Global Interpreter Lock (GIL) limitations for multi-threaded performance.
• Scheduler Integration: Implemented subprocess based forking for local and Slurm schedulers, and modified the Ray scheduler to support forked workers by allocating fractional GPU resources and ensuring they share placement groups.
• New RPC Endpoint: Added a new /fork RPC endpoint to allow existing worker processes to spawn new child worker processes, which inherit the parent's environment.
• Worker Lifecycle Management: Updated worker lifecycle management across local, Ray, and Slurm schedulers to correctly handle forked workers, including creation, health checks, and cleanup, by differentiating between processes owned by the scheduler and those forked by other workers.
• Simplified Placement Group Management: Removed the shared_placement_group concept from the Job dataclass and related controllers, simplifying how placement groups are managed.
• Comprehensive Testing and Documentation: Added extensive unit tests for the new forking functionality in local_scheduler and ray_scheduler, and updated CLI documentation to reflect the new fork parameter.

🧠 New Feature in Public Preview: You can now enable Memory 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 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 introduces a valuable feature to fork co-located workers into separate processes, addressing potential performance bottlenecks from Python's GIL. The implementation across the local, ray, and slurm schedulers is largely consistent and well-tested. My review focuses on a few critical issues I've identified concerning asynchronous execution patterns and subprocess management. These could lead to deadlocks or runtime errors, and my suggestions aim to improve the robustness and stability of this new feature.

[garrett4wade]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a valuable feature to fork colocated workers into separate processes, addressing potential performance issues with Python's multi-threading. The implementation spans across local, slurm, and ray schedulers and is generally well-executed. The changes are extensive, including new RPC endpoints and updated scheduler logic to handle the forking behavior. I've identified a few areas for improvement, mainly concerning code maintainability in the RPC server and a minor bug in a new test for the Ray scheduler. Overall, this is a solid contribution that enhances the framework's flexibility and performance.

[garrett4wade]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a significant feature to allow forking colocated workers into separate processes, addressing performance issues related to Python's GIL. The implementation is comprehensive, touching upon API definitions, scheduler logic for local, Ray, and Slurm environments, and introducing new utility functions. The changes are well-structured, and the addition of new integration tests for the forking mechanism is a great step towards ensuring robustness. The refactoring of async task execution and log command generation into shared utilities also improves the overall code quality.

[rchardx]

Guuuud