fix(mlx): relax Metal context-store timeout by default#7141
Conversation
![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 configures the Apple Silicon Metal context-store timeout by setting the AGX_RELAX_CDM_CTXSTORE_TIMEOUT environment variable to "1" on Darwin arm64 systems, while preserving any existing user-defined values. It also introduces a comprehensive test suite to verify this behavior and ensure the environment variable is set before MLX initialization. There are no review comments, so I have no feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Lyxot
force-pushed
the
fix/mlx-watchdog-timeout
branch
from
9221fdd to
0874960
Compare
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because no GitHub Actions runner was available. Make sure your repository has a runner available to run Copilot's review, or add a copilot-setup-steps.yml file specifying one with the runs-on attribute. See the docs for more details.
Sets a safer default for Metal’s context-store timeout on Apple Silicon so MLX workloads are less likely to hit the watchdog, while keeping any user-provided environment value authoritative.
Changes:
- Default
AGX_RELAX_CDM_CTXSTORE_TIMEOUT=1on macOS arm64 viaos.environ.setdefault. - Gate the behavior to Apple Silicon only, and keep explicit env overrides (including empty string /
"0") intact.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0874960980
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Lyxot
force-pushed
the
fix/mlx-watchdog-timeout
branch
from
0874960 to
4aa3955
Compare
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4aa39558c9
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if platform.system() == "Darwin" and platform.machine() == "arm64": | ||
| os.environ.setdefault("AGX_RELAX_CDM_CTXSTORE_TIMEOUT", "1") |
There was a problem hiding this comment.
On Apple Silicon users who run the normal Python API (from unsloth import FastLanguageModel) never import studio.backend.utils.hardware before unsloth/__init__.py calls _is_mlx_available() and imports unsloth_zoo.mlx (unsloth/init.py:57-76), so this setdefault is not applied before MLX can initialize. That leaves regular non-Studio MLX training/inference with the old watchdog behavior; the default needs to live in a bootstrap that the unsloth package imports before MLX detection, or be shared by both.
Useful? React with 👍 / 👎.
Summary
AGX_RELAX_CDM_CTXSTORE_TIMEOUTto1on Apple Silicon before MLX can initialize in training and export paths.0and an empty string, withos.environ.setdefault.Why
Sustained MLX workloads can hit Metal's context-store/interactivity watchdog while a display is active, as reported in ml-explore/mlx#3267. The driver setting helped the reported M2 case, but it is private and does not guarantee protection on every chip or workload.
The setting must exist before MLX or Metal initializes. The shared hardware module covers Studio and CLI training hardware detection, and the export backend now imports that hardware bootstrap before importing
unsloth, covering both subprocess and direct export-backend imports.Behavior and risk
This is an Apple-Silicon-only, process-scoped default. It does not modify shell profiles or system configuration, and an explicitly supplied environment value always remains authoritative.
Because the variable is undocumented and has not protected every reported system, this PR treats it as a best-effort mitigation rather than a universal watchdog fix.
Validation