CI: CUDA build/test hygiene + fix Puzzletron Nemotron test failures#1901
Merged
Conversation
Install mamba_ssm/causal-conv1d from their latest PyPI releases instead of git main. mamba main (ed6ce09, #984) bumped its cap to apache-tvm-ffi<=0.1.12, which crashes at import ("attribute '__dict__' of 'type' objects is not writable", tvm_ffi/registry.py:403) and broke importing mamba_ssm, failing all nightly Nemotron puzzletron tests. The latest release still caps at the known-good apache-tvm-ffi<=0.1.9, and pinning to releases stops CI from tracking unreleased upstream dependency changes. Also set TORCH_CUDA_ARCH_LIST=12.0 in the GPU/example/regression workflow container env so CUDA kernels compile only for the runner's RTX PRO 6000 (sm_120) instead of the container image's ~6 archs, cutting JIT/source compile time across all containerized tests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughUpdates CUDA architecture targeting across CI workflows, adjusts nox test environments and CUDA extension loading checks, moves an AWQ-Lite calibration regression into CUDA-only coverage, and adds/removes descriptor caching tests between GPU and unit suites. ChangesCUDA runtime and calibration coverage
Descriptor caching coverage
Estimated code review effort: 2 (Simple) | ~12 minutes 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
noxfile.py (1)
111-120: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winConsider pinning
mamba_ssm/causal-conv1dversions.Switching from git-tracking-
mainto PyPI releases fixes the immediate breakage, but installing unpinned "latest" packages reintroduces the same risk class (an unvetted upstream release breaking CI) that this PR is trying to fix. Pinning to a known-good version would make the CI build reproducible and prevent a future release from silently reintroducing theapache-tvm-ffiincompatibility or similar issues.♻️ Example pin
- # Install the latest *released* sdists (built against the container torch) - "mamba_ssm", - "causal-conv1d", + # Install known-good released sdists (built against the container torch) + "mamba_ssm==2.2.4", + "causal-conv1d==1.5.0.post8",🤖 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 `@noxfile.py` around lines 111 - 120, The package install in the nox session is still using unpinned PyPI releases, which can make CI non-reproducible and reintroduce upstream breakages. Update the install logic in the session that calls session.run for pip install so mamba_ssm and causal-conv1d are pinned to known-good versions instead of “latest” releases. Use the existing nox session block as the place to apply the version pins and keep the build behavior stable.
🤖 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.
Nitpick comments:
In `@noxfile.py`:
- Around line 111-120: The package install in the nox session is still using
unpinned PyPI releases, which can make CI non-reproducible and reintroduce
upstream breakages. Update the install logic in the session that calls
session.run for pip install so mamba_ssm and causal-conv1d are pinned to
known-good versions instead of “latest” releases. Use the existing nox session
block as the place to apply the version pins and keep the build behavior stable.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 97beb32c-02a8-48d6-97d0-304d7976cbaf
📒 Files selected for processing (4)
.github/workflows/_example_tests_runner.yml.github/workflows/gpu_tests.yml.github/workflows/regression_tests.ymlnoxfile.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1901 +/- ##
===========================================
+ Coverage 61.21% 77.69% +16.48%
===========================================
Files 515 515
Lines 57245 57977 +732
===========================================
+ Hits 35043 45047 +10004
+ Misses 22202 12930 -9272
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Set CUDA_VISIBLE_DEVICES="" in tests/unit/conftest.py so unit tests always run on CPU even on a GPU machine: device-agnostic tests take their CPU path, and a test that explicitly needs CUDA now fails loudly (belongs in tests/gpu) instead of being silently skipped in CI's CPU-only container. Move test_awq_lite_uncalibrated_linear_keeps_input_quantizer_enabled (the only such unit test, NVFP4 CUDA-only) to tests/gpu so it actually runs in CI. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
changed the title
test_resolve_descriptor_caches_dynamic_modules needs mamba_ssm (only installed in the GPU lane) and a real Nemotron-H checkpoint, so it can only ever be skipped in the CPU-only unit lane. Move it to tests/gpu where mamba_ssm is available; the two mocked tests in the class stay in unit. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
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/gpu/torch/puzzletron/test_resolve_descriptor_caching.py`:
- Around line 36-42: `decoder_layer_cls()` can return either a single class or a
list, but the diagnostic print in `test_resolve_descriptor_caching` assumes it
is always iterable. Update the test to normalize the result from
`descriptor.decoder_layer_cls()` before printing, so the debug output handles
both `Type[nn.Module]` and `List[Type[nn.Module]]` safely. Keep the `assert
layer_classes` check as-is, and adjust the print logic near
`descriptor`/`layer_classes` to avoid iterating over a single class.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: dbfd3458-a5ad-4750-bb43-7234c3c37cab
📒 Files selected for processing (2)
tests/gpu/torch/puzzletron/test_resolve_descriptor_caching.pytests/unit/torch/puzzletron/test_resolve_descriptor_caching.py
💤 Files with no reviewable changes (1)
- tests/unit/torch/puzzletron/test_resolve_descriptor_caching.py
Set CUDA_VISIBLE_DEVICES="" in the nox unit/partial_unit session env (before the pytest process imports torch, unlike tests/unit/conftest.py which runs after the root conftest already initialized CUDA), so torch.cuda.is_available() is False on GPU dev boxes too. Gate load_cpp_extension on torch.cuda.is_available() so no CUDA extension is JIT-compiled when no GPU is usable, and drop the ineffective CUDA_VISIBLE_DEVICES line from tests/unit/conftest.py. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
force-pushed
the
kmorabia/fix-tests
branch
from
290196a to
cc02ce3
Compare
Contributor
|
kevalmorabia97
changed the title
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
noxfile.py (1)
115-124: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd explicit version caps for
mamba_ssmandcausal-conv1d.Floating released PyPI versions can bring this nightly GPU job back into the same failure mode; cap these to a known-good pair, or document why tracking upstream releases is intentional.
🤖 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 `@noxfile.py` around lines 115 - 124, The pip install step in the session setup is still pulling floating released sdists for mamba_ssm and causal-conv1d, which can reintroduce the GPU failure; update the install list in the session.run block to pin both packages to a known-good compatible version pair, or add a clear inline note explaining why unpinned upstream tracking is intentional. Use the session.run call in noxfile.py as the target location.
🤖 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.
Nitpick comments:
In `@noxfile.py`:
- Around line 115-124: The pip install step in the session setup is still
pulling floating released sdists for mamba_ssm and causal-conv1d, which can
reintroduce the GPU failure; update the install list in the session.run block to
pin both packages to a known-good compatible version pair, or add a clear inline
note explaining why unpinned upstream tracking is intentional. Use the
session.run call in noxfile.py as the target location.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3ca8ab0a-6ec6-4878-870a-f25c415ac8c6
📒 Files selected for processing (2)
modelopt/torch/utils/cpp_extension.pynoxfile.py
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
run_example_command now captures the subprocess output and, on a non-zero exit matching a transient HF/network signature (5xx, rate limit, dropped connection, interrupted download), retries before failing. Non-transient failures still raise CalledProcessError immediately, so real example bugs are not masked. Reduces flaky example/GPU test failures caused by HF Hub server issues. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
changed the title
kevalmorabia97
changed the title
Add a skip_flaky_on_blackwell pytest marker (sm_120) and apply it to test_homogeneous_sharded_state_dict and the test_regular_state_dict[hidden_size=320] case. Both intermittently hit a CUDA illegal memory access on the nemo:26.06 / TE 2.16 / CUDA 13 stack that poisons the CUDA context and cascades 120s timeouts across gpu_megatron. Consolidates the ad-hoc per-site skips into one marker/TODO; still runs on non-Blackwell GPUs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
The VLM Minitron prune test is slow on the 2-GPU nightly runners and intermittently exceeds the default per-test cap; bump it to 360s. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
force-pushed
the
kmorabia/fix-tests
branch
from
961d4e2 to
bd0d970
Compare
talorabr
pushed a commit
that referenced
this pull request
Jul 7, 2026
…1901) ### What does this PR do? Type of change: Bug fix (CI / tests) - **Fix Nemotron nightly failures:** install `mamba_ssm`/`causal-conv1d` from PyPI releases instead of git `main` (avoids the broken `apache-tvm-ffi 0.1.12` that crashes on import). - **Speed up CUDA builds:** set `TORCH_CUDA_ARCH_LIST=12.0` (runner's sm_120) in the GPU/example/regression workflow container env instead of the image's ~6 archs. - **Make unit tests CPU-only:** force CUDA off in the nox `unit` env and skip JIT-compiling CUDA extensions when no GPU is usable; move the two GPU-/`mamba_ssm`-requiring unit tests to `tests/gpu`. - **Harden example tests against HF flakes:** capture subprocess output and retry transient HuggingFace access errors (5xx / rate-limit / connection). - **Skip Blackwell-flaky sharded-state-dict tests:** `test_homogeneous_sharded_state_dict` and `test_regular_state_dict[320]` intermittently hit a CUDA illegal-memory-access on the sm_120 runner that poisons the CUDA context and cascades timeouts; gate them behind a reusable `skip_flaky_on_blackwell` marker (still run on non-Blackwell GPUs). - **Bump slow test timeout:** `test_prune_minitron_vlm` → 360s for the 2-GPU nightly. ### Testing - CI tests on this PR pass (1-gpu) - Manually triggerred 2-gpu test: - GPU: https://github.com/NVIDIA/Model-Optimizer/actions/runs/28774553356 - Examples: https://github.com/NVIDIA/Model-Optimizer/actions/runs/28774556693 - Regression: https://github.com/NVIDIA/Model-Optimizer/actions/runs/28771197586 ### Additional Information - Backward compatible: N/A (CI/tests only) - New dependency: N/A - Changelog: N/A (CI/test infra) 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: talora <talora@nvidia.com>
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
1 participant
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.
What does this PR do?
Type of change: Bug fix (CI / tests)
mamba_ssm/causal-conv1dfrom PyPI releases instead of gitmain(avoids the brokenapache-tvm-ffi 0.1.12that crashes on import).TORCH_CUDA_ARCH_LIST=12.0(runner's sm_120) in the GPU/example/regression workflow container env instead of the image's ~6 archs.unitenv and skip JIT-compiling CUDA extensions when no GPU is usable; move the two GPU-/mamba_ssm-requiring unit tests totests/gpu.test_homogeneous_sharded_state_dictandtest_regular_state_dict[320]intermittently hit a CUDA illegal-memory-access on the sm_120 runner that poisons the CUDA context and cascades timeouts; gate them behind a reusableskip_flaky_on_blackwellmarker (still run on non-Blackwell GPUs).test_prune_minitron_vlm→ 360s for the 2-GPU nightly.Testing
Additional Information
🤖 Generated with Claude Code