fix/test(moe_ep): self-bootstrap 1-rank process group in dg mega oracle test - #4221
Merged
Merged
Conversation
test_deep_gemm_mega_kernel_matches_torch_reference (added in #3980) fails under plain pytest with "ValueError: ... environment variable RANK expected, but not set": its WORLD_SIZE guard only skips when the variable is set and != 1, so without torchrun it falls through to dist.init_process_group(backend="nccl"), whose default env:// rendezvous requires torchrun's RANK/MASTER_* variables. This is only reachable on CI jobs that combine plain-pytest discovery of tests/moe_ep, a cu13 image with the EP stack (deep_gemm) installed, and a capability-10 GPU — i.e. the B300 cu130 unit-test job, where it currently errors on every run. Fix: when RANK is absent, self-bootstrap a 1-rank NCCL group via an explicit tcp://127.0.0.1:<free-port> init (rank=0, world_size=1) instead of env://. Deliberately avoids os.environ mutation so RANK/MASTER_* don't leak to later tests in the same pytest process; teardown is unchanged (conftest.pytest_sessionfinish destroys the group). The torchrun path is preserved and run_tests.sh still exercises it. Verified on GB200 (single GPU): both launch modes pass with rel_l2=0.0027 vs the torch oracle; the plain-pytest mode was confirmed with RANK/WORLD_SIZE/MASTER_ADDR/MASTER_PORT/LOCAL_RANK explicitly unset. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
mhoqueanik
requested review from
Anerudhan,
Aneureka,
IwakuraRein,
aleozlx,
feih-nv,
jiahanc,
nv-yunzheq,
samuellees and
yzh119
as code owners
July 28, 2026 20:46
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
mhoqueanik
changed the title
Contributor
📝 WalkthroughWalkthroughThe deep-gemm oracle test now supports plain ChangesDeep-GEMM test execution
Estimated code review effort: 2 (Simple) | ~10 minutes Sequence Diagram(s)sequenceDiagram
participant Pytest
participant Socket
participant NCCL
Pytest->>NCCL: Use env:// initialization when RANK is set
alt Plain pytest
Pytest->>Socket: Bind an available local port
Socket-->>Pytest: Return TCP rendezvous address
Pytest->>NCCL: Initialize rank 0, world size 1 via tcp://
end
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
mhoqueanik
enabled auto-merge (squash)
July 28, 2026 20:47
Anerudhan
approved these changes
Jul 28, 2026
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/moe_ep/test_deep_gemm_mega_kernel_vs_reference.py`:
- Around line 244-262: Expand the existing try/finally in the test setup to
begin immediately after the conditional process-group bootstrap, covering
torch.cuda.set_device(), _make_problem(), and symmetric-buffer creation. Track
whether this test initialized the group (including the plain-pytest path), and
in finally destroy it only when owned by this test, while preserving externally
managed torchrun groups.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 299bb828-76f2-4852-9c04-415052150bef
📒 Files selected for processing (2)
tests/moe_ep/run_tests.shtests/moe_ep/test_deep_gemm_mega_kernel_vs_reference.py
Comment on lines
+244
to
+262
| if "RANK" in os.environ: | ||
| # torchrun launch: use its env:// rendezvous. | ||
| dist.init_process_group(backend="nccl", timeout=_PG_TIMEOUT) | ||
| else: | ||
| # Plain-pytest launch (CI unit-test jobs): self-bootstrap a 1-rank | ||
| # group with an explicit tcp:// store instead of mutating the | ||
| # RANK/MASTER_* env vars, which would leak to later tests in the | ||
| # same pytest process. The freed port can theoretically be | ||
| # snatched before torch rebinds it; acceptable for a test. | ||
| with socket.socket() as s: | ||
| s.bind(("127.0.0.1", 0)) | ||
| port = s.getsockname()[1] | ||
| dist.init_process_group( | ||
| backend="nccl", | ||
| init_method=f"tcp://127.0.0.1:{port}", | ||
| rank=0, | ||
| world_size=1, | ||
| timeout=_PG_TIMEOUT, | ||
| ) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Ensure bootstrap failures cannot leak the process group.
The group is initialized before torch.cuda.set_device(), _make_problem(), and symmetric-buffer creation, but the cleanup boundary starts afterward. If any of those operations fails, the plain-pytest process group remains initialized for subsequent tests. Move the existing try/finally to cover all post-bootstrap setup and destroy the group only when this test created it.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/moe_ep/test_deep_gemm_mega_kernel_vs_reference.py` around lines 244 -
262, Expand the existing try/finally in the test setup to begin immediately
after the conditional process-group bootstrap, covering torch.cuda.set_device(),
_make_problem(), and symmetric-buffer creation. Track whether this test
initialized the group (including the plain-pytest path), and in finally destroy
it only when owned by this test, while preserving externally managed torchrun
groups.
nv-yunzheq
approved these changes
Jul 28, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
test_deep_gemm_mega_kernel_matches_torch_reference (added in #3980) fails under plain pytest with "ValueError: ... environment variable RANK expected, but not set": its WORLD_SIZE guard only skips when the variable is set and != 1, so without torchrun it falls through to dist.init_process_group(backend="nccl"), whose default env:// rendezvous requires torchrun's RANK/MASTER_* variables.
This is only reachable on CI jobs that combine plain-pytest discovery of tests/moe_ep, a cu13 image with the EP stack (deep_gemm) installed, and a capability-10 GPU — i.e. the B300 cu130 unit-test job, where it currently errors on every run.
Fix: when RANK is absent, self-bootstrap a 1-rank NCCL group via an explicit tcp://127.0.0.1: init (rank=0, world_size=1) instead of env://. Deliberately avoids os.environ mutation so RANK/MASTER_* don't leak to later tests in the same pytest process; teardown is unchanged (conftest.pytest_sessionfinish destroys the group). The torchrun path is preserved and run_tests.sh still exercises it.
Verified on GB200 (single GPU): both launch modes pass with rel_l2=0.0027 vs the torch oracle; the plain-pytest mode was confirmed with RANK/WORLD_SIZE/MASTER_ADDR/MASTER_PORT/LOCAL_RANK explicitly unset.
AI-assisted (Claude Code).
📌 Description
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit