feat(launcher): auto-set OMP_NUM_THREADS based on cpus_per_task by rchardx · Pull Request #859 · areal-project/AReaL

rchardx · 2026-01-27T03:53:27Z

Description

Dynamically compute thread control environment variables (OMP_NUM_THREADS, MKL_NUM_THREADS, etc.) based on allocated CPU cores to prevent thread explosion in multi-process training scenarios.

Changes:

Add get_thread_env_vars() function in launcher utils
Update SchedulingSpec.cpu default from 4 to 8
Apply thread env vars in Local, Ray, and Slurm launchers
Support user override via SchedulingSpec.env_vars or os.environ

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 · 2026-01-27T03:53:45Z

Summary of Changes

Hello @rchardx, 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 performance and resource management by implementing dynamic control over thread-related environment variables. By automatically configuring settings like OMP_NUM_THREADS based on the number of allocated CPU cores, the changes aim to mitigate issues like thread explosion in multi-process training. This ensures more efficient utilization of computational resources and a more stable execution environment across different deployment launchers.

Highlights

Dynamic Thread Control: Introduced a new utility function, get_thread_env_vars, to dynamically compute and set thread control environment variables (e.g., OMP_NUM_THREADS, MKL_NUM_THREADS) based on the allocated CPU cores per task. This aims to prevent thread explosion in multi-process training scenarios.
Launcher Integration: The dynamic thread environment variable setting has been integrated across Local, Ray, and Slurm launchers, ensuring that server and trainer submissions automatically receive optimized thread configurations.
Default CPU Allocation Update: The default value for SchedulingSpec.cpu has been increased from 4 to 8, reflecting a common requirement for CPU cores per GPU and aligning with the new thread control logic.
User Override Support: The system supports user overrides for thread environment variables, allowing configurations specified in SchedulingSpec.env_vars or existing os.environ to take precedence over the dynamically computed values.

🧠 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.

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

The pull request introduces a valuable feature to dynamically manage thread control environment variables, preventing thread explosion in multi-process training. The implementation correctly prioritizes user-defined environment variables and provides a sensible fallback. The update of the default cpu value in SchedulingSpec and its corresponding documentation is also well-aligned with the goal of better resource utilization. However, there is a significant opportunity to improve code maintainability by refactoring duplicated logic related to extracting scheduling specifications and environment variables across the launcher files.

Copilot

Pull request overview

This PR introduces automatic configuration of thread control environment variables (OMP_NUM_THREADS, MKL_NUM_THREADS, etc.) to prevent thread explosion in multi-process training scenarios. However, there are critical bugs in how the thread counts are calculated across different launchers.

Changes:

Added get_thread_env_vars() utility function to compute thread environment variables with priority-based override support
Changed SchedulingSpec.cpu default from 4 to 8 to match the new default thread count
Integrated thread environment variable setting in Local, Ray, and Slurm launchers for both training and inference workloads

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
areal/utils/launcher.py	Added `get_thread_env_vars()` function to compute thread control environment variables based on CPU allocation
areal/api/cli_args.py	Updated `SchedulingSpec.cpu` default value from 4 to 8
docs/cli_reference.md	Updated documentation to reflect new default CPU value of 8
areal/launcher/slurm.py	Integrated thread environment variables for trainer and LLM servers (contains bugs)
areal/launcher/ray.py	Integrated thread environment variables for trainer and LLM servers (contains bugs)
areal/launcher/local.py	Integrated thread environment variables for trainer and LLM servers (contains bugs)

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

garrett4wade

THREAD_ENV_VARS has not been set for controllers. We should also set them in @areal/scheduler/

Dynamically compute thread control environment variables (OMP_NUM_THREADS, MKL_NUM_THREADS, etc.) based on allocated CPU cores to prevent thread explosion in multi-process training scenarios. Changes: - Add get_thread_env_vars() function in launcher utils - Update SchedulingSpec.cpu default from 4 to 8 - Apply thread env vars in Local, Ray, and Slurm launchers - Support user override via SchedulingSpec.env_vars or os.environ

garrett4wade

LGTM. Waiting for CI.

…l-project#859) Dynamically compute thread control environment variables (OMP_NUM_THREADS, MKL_NUM_THREADS, etc.) based on allocated CPU cores to prevent thread explosion in multi-process training scenarios. Changes: - Add get_thread_env_vars() function in launcher utils - Update SchedulingSpec.cpu default from 4 to 8 - Apply thread env vars in Local, Ray, and Slurm launchers - Support user override via SchedulingSpec.env_vars or os.environ

rchardx requested review from Copilot, garrett4wade and nuzant January 27, 2026 03:53

Copilot started reviewing on behalf of rchardx January 27, 2026 03:53 View session

gemini-code-assist Bot reviewed Jan 27, 2026

View reviewed changes

Comment thread areal/utils/launcher.py

Comment thread areal/utils/launcher.py

Comment thread areal/launcher/local.py Outdated

Comment thread areal/launcher/local.py Outdated

Copilot AI reviewed Jan 27, 2026

View reviewed changes

Comment thread areal/launcher/slurm.py Outdated

Comment thread areal/launcher/slurm.py

Comment thread areal/launcher/ray.py

Comment thread areal/launcher/local.py

areal-project deleted a comment from Saingsophea Jan 27, 2026

rchardx force-pushed the rchardx/thread branch 2 times, most recently from 5665111 to d766066 Compare January 27, 2026 05:52

rchardx added the safe-to-test Ready to run unit-tests in a PR. label Jan 27, 2026

rchardx temporarily deployed to AReaL-unittests January 27, 2026 06:26 — with GitHub Actions Inactive

garrett4wade reviewed Jan 27, 2026

View reviewed changes

Comment thread areal/launcher/vllm_server.py Outdated

Comment thread areal/launcher/sglang_server.py Outdated

rchardx force-pushed the rchardx/thread branch from d766066 to c07b914 Compare January 27, 2026 08:39

rchardx added safe-to-test Ready to run unit-tests in a PR. and removed safe-to-test Ready to run unit-tests in a PR. labels Jan 27, 2026

rchardx temporarily deployed to AReaL-unittests January 27, 2026 09:09 — with GitHub Actions Inactive

garrett4wade approved these changes Jan 27, 2026

View reviewed changes

garrett4wade merged commit 80ced5f into main Jan 27, 2026
7 checks passed

garrett4wade deleted the rchardx/thread branch January 27, 2026 11:25

Saingsophea approved these changes Jan 28, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat(launcher): auto-set OMP_NUM_THREADS based on cpus_per_task#859

feat(launcher): auto-set OMP_NUM_THREADS based on cpus_per_task#859
garrett4wade merged 1 commit into
mainfrom
rchardx/thread

rchardx commented Jan 27, 2026

Uh oh!

gemini-code-assist Bot commented Jan 27, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

garrett4wade left a comment

Uh oh!

Uh oh!

Uh oh!

garrett4wade left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

rchardx commented Jan 27, 2026

Description

Type of Change

Checklist

Uh oh!

gemini-code-assist Bot commented Jan 27, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

garrett4wade left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

garrett4wade left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants