[https://nvbugs/6437410][fix] fix nemotron weight update test - #16712
[https://nvbugs/6437410][fix] fix nemotron weight update test#16712shuyixiong wants to merge 4 commits into
Conversation
Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>
Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
|
/bot run --disable-fail-fast --stage-list "H100_PCIe-PyTorch-Ray-1, DGX_B200-4_GPUs-PyTorch-Ray-1" |
WalkthroughChangesUpdate-weights test coverage
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py`:
- Around line 796-847: Update test_llm_update_weights_nemotron_h to launch the
pytest command with subprocess.Popen(start_new_session=True), replacing
subprocess.run while preserving its output capture and timeout behavior. In the
subprocess.TimeoutExpired handler, terminate the entire process group with
os.killpg using the child’s process group ID before failing the test, ensuring
Ray/NCCL workers are also stopped.
🪄 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: bf9183ba-ae60-4bd9-9a0e-04b6c54031a4
📒 Files selected for processing (3)
tests/integration/test_lists/waives.txttests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.pytests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
| @pytest.mark.part4 | ||
| @skip_pre_hopper | ||
| def test_llm_update_weights_nemotron_h(mamba_deps): | ||
| """Runs _nemotron_h_body in a spawned subprocess so HF transformers | ||
| sees the mamba-ssm / causal-conv1d fast path installed by the | ||
| mamba_deps fixture. See _nemotron_h_body docstring for why.""" | ||
| ctx = multiprocessing.get_context("spawn") | ||
| queue = ctx.Queue() | ||
| proc = ctx.Process(target=_nemotron_h_subprocess_entry, args=(queue,)) | ||
| proc.start() | ||
| proc.join() | ||
| err = queue.get() if not queue.empty() else None | ||
| if proc.exitcode != 0: | ||
| pytest.fail(f"Subprocess exited with code {proc.exitcode}\n{err or ''}") | ||
| if err is not None: | ||
| pytest.fail(err) | ||
| """Runs the Nemotron-H body in a fresh ``python -m pytest`` subprocess so | ||
| HF transformers re-imports cleanly and picks up the mamba-ssm / | ||
| causal-conv1d fast path installed by the ``mamba_deps`` fixture (a plain | ||
| in-process run would keep the parent's negative import caches; see the | ||
| _nemotron_h_body docstring). Driving it as a subprocess — instead of a | ||
| hand-managed ``multiprocessing`` child — lets ``subprocess.run`` and the | ||
| inner pytest own the process lifecycle: a hang is bounded by ``timeout=``, | ||
| a crash surfaces as a non-zero return code, and the failure detail is the | ||
| inner pytest's own traceback.""" | ||
| # Must stay under the outer pytest ``--timeout`` so a genuine hang (e.g. a | ||
| # Ray/NCCL/CUDA deadlock) is reported here with useful output instead of | ||
| # the whole test being hard-killed at the pytest timeout. | ||
| subprocess_timeout_s = 1800.0 | ||
|
|
||
| node_id = f"{os.path.abspath(__file__)}::test_nemotron_h_body_impl" | ||
| cmd = [ | ||
| sys.executable, | ||
| "-m", | ||
| "pytest", | ||
| node_id, | ||
| "--run-ray", | ||
| "-p", | ||
| "no:cacheprovider", | ||
| "-p", | ||
| "no:xdist", | ||
| "--tb=short", | ||
| "-s", | ||
| "-v", | ||
| ] | ||
| env = {**os.environ, _NEMOTRON_H_BODY_ENV: "1"} | ||
| try: | ||
| result = subprocess.run( | ||
| cmd, | ||
| env=env, | ||
| capture_output=True, | ||
| text=True, | ||
| timeout=subprocess_timeout_s, | ||
| ) | ||
| except subprocess.TimeoutExpired as e: | ||
| out = (e.stdout or "") + (e.stderr or "") | ||
| pytest.fail( | ||
| f"Nemotron-H subprocess did not complete within " | ||
| f"{subprocess_timeout_s:.0f}s (likely hung); terminated.\n{out}" | ||
| ) | ||
| if result.returncode != 0: | ||
| pytest.fail( | ||
| f"Nemotron-H subprocess failed (exit code {result.returncode}).\n" | ||
| f"{result.stdout}\n{result.stderr}" | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n 'nemotron' tests/integration/test_lists/waives.txt
rg -n 'test_llm_update_weights_nemotron_h|test_nemotron_h_body_impl' tests/integration/test_lists/waives.txt tests/integration/test_lists/test-db tests/integration/test_lists/qa -r ''Repository: NVIDIA/TensorRT-LLM
Length of output: 1332
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the test file around the affected lines and the inner body helper.
sed -n '720,920p' tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py
printf '\n--- waiver/test-list search ---\n'
rg -n 'test_llm_update_weights_nemotron_h|test_nemotron_h_body_impl|nemotron_h' tests/integration/test_lists tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py || true
printf '\n--- list files mentioning this unittest path ---\n'
rg -n 'unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py|test_llm_update_weights_multi_gpu.py' tests/integration/test_lists || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 8486
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check whether any integration test-list files were modified in this PR.
git diff --name-only -- tests/integration/test_lists
# Show the exact hunks, if any, for the relevant waiver file and test file.
printf '\n--- waives diff ---\n'
git diff -- tests/integration/test_lists/waives.txt || true
printf '\n--- test file diff summary ---\n'
git diff --stat -- tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 210
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Look for established timeout/process-group patterns in nearby tests.
rg -n 'start_new_session|killpg|TimeoutExpired|subprocess\.run\(' tests/unittest tests/integration | sed -n '1,120p'Repository: NVIDIA/TensorRT-LLM
Length of output: 12968
Kill the whole Ray subprocess tree on timeout. subprocess.run(..., timeout=...) only stops the direct python -m pytest child; Ray/NCCL workers can survive a hang and keep GPU memory pinned. Launch it with Popen(start_new_session=True) and os.killpg(...) the group in the TimeoutExpired path.
- Changed tests:
test_llm_update_weights_nemotron_h,test_nemotron_h_body_impl. - Coverage:
test_llm_update_weights_nemotron_his listed intests/integration/test_lists/test-db/l0_dgx_b200.ymlvia-m "part4";test_nemotron_h_body_implis intentionally unlisted. - Verdict: sufficient.
🧰 Tools
🪛 ast-grep (0.44.1)
[error] 829-835: Command coming from incoming request
Context: subprocess.run(
cmd,
env=env,
capture_output=True,
text=True,
timeout=subprocess_timeout_s,
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').
(subprocess-from-request)
[error] 829-835: Use of unsanitized data to create processes
Context: subprocess.run(
cmd,
env=env,
capture_output=True,
text=True,
timeout=subprocess_timeout_s,
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').
(os-system-unsanitized-data)
🪛 Ruff (0.15.21)
[error] 830-830: subprocess call: check for execution of untrusted input
(S603)
🤖 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/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py`
around lines 796 - 847, Update test_llm_update_weights_nemotron_h to launch the
pytest command with subprocess.Popen(start_new_session=True), replacing
subprocess.run while preserving its output capture and timeout behavior. In the
subprocess.TimeoutExpired handler, terminate the entire process group with
os.killpg using the child’s process group ID before failing the test, ensuring
Ray/NCCL workers are also stopped.
Source: Path instructions
There was a problem hiding this comment.
AI is right here, but we should use the utilities already available:
TensorRT-LLM/tests/integration/defs/trt_test_alternative.py
Lines 276 to 279 in 4c17c85
This is also one reason why we should make these integration tests instead of unittests.
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- existing integration subprocess utility ---'
sed -n '250,300p' tests/integration/defs/trt_test_alternative.py
printf '%s\n' '--- utility usages ---'
rg -n -C 3 'killpg|start_new_session|trt_test_alternative' tests/integration tests/unittest \
-g '*.py' | sed -n '1,220p'Length of output: 21468
@tongyuantongyu, confirmed. The helper is the right solution—and stronger than my original killpg() suggestion:
tests/integration/defs/trt_test_alternative.py::popen()defaults tostart_new_session=True.- Its
cleanup_process_tree()recursively identifies and terminates descendants, addressing Ray workers that may create separate process groups. - It also handles timeout output collection and has cross-platform behavior.
So the appropriate follow-up is to migrate this coverage to the integration framework and use trt_test_alternative there, rather than duplicating subprocess lifecycle logic in this unittest.
✏️ Learnings added
Learnt from: tongyuantongyu
Repo: NVIDIA/TensorRT-LLM PR: 16712
File: tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py:796-847
Timestamp: 2026-07-23T10:22:02.769Z
Learning: `tests/integration/defs/trt_test_alternative.py::popen()` and its `cleanup_process_tree()` are the canonical TensorRT-LLM utilities for subprocess-based integration tests. They recursively clean descendant processes because `os.killpg()` alone is insufficient when child processes, including Ray workers, create their own process groups; use these utilities rather than local process-group cleanup implementations.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
There was a problem hiding this comment.
It's gone — no process tree to clean.
There was a problem hiding this comment.
Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!
|
PR_Github #60897 [ run ] triggered by Bot. Commit: |
|
PR_Github #60897 [ run ] completed with state |
|
|
||
| class RefHFModelWithIPCHandles(RefHFModel): | ||
| def __init__(self, model_dir: str, device_id: int = 0, num_hidden_layers: int = 4): | ||
| def __init__(self, model_dir: str, device_id: int = 0, **model_kwargs): |
There was a problem hiding this comment.
nit: Collapsing the explicit num_hidden_layers: int = 4 into an untyped**model_kwargsmakes the signature less self-documenting — a reader can no longer tell from the signature which kwargs actually matter, and the type annotation is gone. Since only two keys are ever meaningful, a keyword-only signature keeps the honesty the old code had while still supporting hybrid models:
def __init__(self, model_dir: str, device_id: int = 0, *,
num_hidden_layers: int | None = None,
layers_block_type: list[str] | None = None):
| # such models. Catch a silently ignored override loudly here. | ||
| num_hidden_layers = model_kwargs.get("num_hidden_layers") | ||
| if num_hidden_layers is not None: | ||
| assert self.model.config.num_hidden_layers == num_hidden_layers |
There was a problem hiding this comment.
nit: Add self-explaining forthe failure:
assert self.model.config.num_hidden_layers == num_hidden_layers, (
f"num_hidden_layers override silently ignored: "
f"HF loaded {self.model.config.num_hidden_layers}, expected {num_hidden_layers}"
)
| # reference and the mamba SSM / selective-scan path introduces small | ||
| # numerical differences (observed top-20 overlap ~0.89 vs the 0.9 | ||
| # default). | ||
| compare_logits(llm_logits, ref_logits, threshold=0.8) |
There was a problem hiding this comment.
Observed overlap is ~0.89 but the threshold is dropped to 0.8 — could we tighten to 0.85 to keep a comfortable margin over BF16/selective-scan noise while still catching a real regression?
There was a problem hiding this comment.
Tightened to 0.85. Correction: ~0.89 is the mean, not the worst. 5 runs x 4 prompts (GB200, TP=4): mean 0.8909, min 0.8672, max 0.9281.
chzblych
left a comment
chzblych
left a comment
There was a problem hiding this comment.
Approved for the waive list change.
| if err is not None: | ||
| pytest.fail(err) | ||
| """Runs the Nemotron-H body in a fresh ``python -m pytest`` subprocess so | ||
| HF transformers re-imports cleanly and picks up the mamba-ssm / |
There was a problem hiding this comment.
Can we find a sane version that won't mess our environment, put it into https://github.com/NVIDIA/TensorRT-LLM/blob/main/requirements-dev.txt, and remove these subprocess & nested pytest stuffs? If not we'd better make them integration tests. They are too heavy to be unittests in the current shape.
There was a problem hiding this comment.
Put them in the Ray-stage install next to ray[default] rather than requirements-dev.txt, as explicit wheel URLs: a requirements file can't pass --no-build-isolation, and a plain version pin means a ~19 min source build per stage since upstream's newest wheels are DLFW 26.04 while our image is 26.05.
| @pytest.mark.part4 | ||
| @skip_pre_hopper | ||
| def test_llm_update_weights_nemotron_h(mamba_deps): | ||
| """Runs _nemotron_h_body in a spawned subprocess so HF transformers | ||
| sees the mamba-ssm / causal-conv1d fast path installed by the | ||
| mamba_deps fixture. See _nemotron_h_body docstring for why.""" | ||
| ctx = multiprocessing.get_context("spawn") | ||
| queue = ctx.Queue() | ||
| proc = ctx.Process(target=_nemotron_h_subprocess_entry, args=(queue,)) | ||
| proc.start() | ||
| proc.join() | ||
| err = queue.get() if not queue.empty() else None | ||
| if proc.exitcode != 0: | ||
| pytest.fail(f"Subprocess exited with code {proc.exitcode}\n{err or ''}") | ||
| if err is not None: | ||
| pytest.fail(err) | ||
| """Runs the Nemotron-H body in a fresh ``python -m pytest`` subprocess so | ||
| HF transformers re-imports cleanly and picks up the mamba-ssm / | ||
| causal-conv1d fast path installed by the ``mamba_deps`` fixture (a plain | ||
| in-process run would keep the parent's negative import caches; see the | ||
| _nemotron_h_body docstring). Driving it as a subprocess — instead of a | ||
| hand-managed ``multiprocessing`` child — lets ``subprocess.run`` and the | ||
| inner pytest own the process lifecycle: a hang is bounded by ``timeout=``, | ||
| a crash surfaces as a non-zero return code, and the failure detail is the | ||
| inner pytest's own traceback.""" | ||
| # Must stay under the outer pytest ``--timeout`` so a genuine hang (e.g. a | ||
| # Ray/NCCL/CUDA deadlock) is reported here with useful output instead of | ||
| # the whole test being hard-killed at the pytest timeout. | ||
| subprocess_timeout_s = 1800.0 | ||
|
|
||
| node_id = f"{os.path.abspath(__file__)}::test_nemotron_h_body_impl" | ||
| cmd = [ | ||
| sys.executable, | ||
| "-m", | ||
| "pytest", | ||
| node_id, | ||
| "--run-ray", | ||
| "-p", | ||
| "no:cacheprovider", | ||
| "-p", | ||
| "no:xdist", | ||
| "--tb=short", | ||
| "-s", | ||
| "-v", | ||
| ] | ||
| env = {**os.environ, _NEMOTRON_H_BODY_ENV: "1"} | ||
| try: | ||
| result = subprocess.run( | ||
| cmd, | ||
| env=env, | ||
| capture_output=True, | ||
| text=True, | ||
| timeout=subprocess_timeout_s, | ||
| ) | ||
| except subprocess.TimeoutExpired as e: | ||
| out = (e.stdout or "") + (e.stderr or "") | ||
| pytest.fail( | ||
| f"Nemotron-H subprocess did not complete within " | ||
| f"{subprocess_timeout_s:.0f}s (likely hung); terminated.\n{out}" | ||
| ) | ||
| if result.returncode != 0: | ||
| pytest.fail( | ||
| f"Nemotron-H subprocess failed (exit code {result.returncode}).\n" | ||
| f"{result.stdout}\n{result.stderr}" | ||
| ) |
There was a problem hiding this comment.
AI is right here, but we should use the utilities already available:
TensorRT-LLM/tests/integration/defs/trt_test_alternative.py
Lines 276 to 279 in 4c17c85
This is also one reason why we should make these integration tests instead of unittests.
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: This test-only fix is mechanically mergeable and the core approach (truncating hybrid NemotronH via layers_block_type instead of the silently-ignored num_hidden_layers) looks correct, but the timeout error path leaks GPU-holding worker processes and should be fixed before merge.
Concerns
- [MAJOR]
tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py:831- TimeoutExpired handler leaks the Ray/NCCL process group- What is wrong: The subprocess is started with
subprocess.run(...)with no new session / process group. Onsubprocess.TimeoutExpired, only the directpytestchild is terminated; the Ray/NCCL/CUDA workers it forked are orphaned. - How it fails: The 1800s timeout exists precisely to catch a Ray/NCCL/CUDA deadlock. When that hang occurs ->
TimeoutExpiredfires ->pytest.fail(...)is raised, but the worker grandchildren survive and keep holding multi-GPU memory. Subsequent tests on the same node then fail with OOM / device-busy, turning a single hang into a cascade. - Suggested fix: launch in its own process group and kill the group on timeout, e.g.
proc = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, start_new_session=True) try: out, err = proc.communicate(timeout=subprocess_timeout_s) except subprocess.TimeoutExpired: os.killpg(os.getpgid(proc.pid), signal.SIGKILL) out, err = proc.communicate() pytest.fail(f"Nemotron-H subprocess hung >{subprocess_timeout_s:.0f}s; killed.\n{out}\n{err}")
- What is wrong: The subprocess is started with
Minor notes (non-blocking)
tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py:772- threshold lowered 0.9 -> 0.8 while observed overlap is ~0.89; the ~0.09 slack could mask a real regression. Consider ~0.85.tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py:44- dropping thenum_hidden_layers=4default means any caller that relied on it now loads the full checkpoint. Please confirm no other instantiations depend on the old default.tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py:812--p no:xdistcan error if pytest-xdist is not installed in the subprocess env.
QA view
- Test coverage: adequate for a test-only change - it re-enables the previously waived Nemotron-H test; there is no production code to cover. Runtime pass/fail is unverified (CBTS results unavailable per the PR description) and the timeout/leak path is not itself exercised.
- SM coverage: arch-guarded by
@skip_pre_hopper, so it runs on Hopper+ only; the FP8 Nemotron/Qwen case stays waived, so the fp8 variant of this path remains untested. No new arch introduced without a test. - Test code: see the MAJOR (process-group leak) plus the loosened threshold and
-p no:xdistnotes above. - Test time: significant - removing the
part4waiver re-enables a 4-GPU weight-update test on a (7-layer-truncated) 30B model with an 1800s budget. Exact runtime can't be read from the diff. - Needs
/qa-verify: yes - this un-waives a previously failing bug-fix test whose result is unavailable, changes test infrastructure (subprocess-driven pytest + process lifecycle), and has a leaky error path. A human QA should confirm it passes on Hopper+ and leaves no orphaned GPU workers.
Does this actually fix nvbugs/6437410?
Likely yes. The traced failure is that NemotronHConfig derives num_hidden_layers from layers_block_type and silently ignores the direct override, so the old model_kwargs={'num_hidden_layers':7} kept the full 30B model (OOM / unmatched logits). The diff truncates via layers_block_type[:7] on both the HF reference and the LLM, adds an assert to catch the silent-ignore case, loosens the threshold for mamba/selective-scan drift, and re-runs in a fresh python -m pytest subprocess so the mamba-ssm fast path is picked up. That chain addresses the root cause, but it is unverified at runtime.
Possible new issues
- Leaked Ray/NCCL GPU workers on a hang (MAJOR) can break later tests on the same node.
- Removed
num_hidden_layers=4default may cause unshown callers to load full checkpoints. - The 0.8 threshold reduces sensitivity to future numerical regressions.
What I could not verify
- The other instantiations of
RefHFModelWithIPCHandles(not in the diff) and whether any relied on the old default of 4 layers. - Actual test pass/fail and whether the subprocess leaves orphaned processes at runtime - CBTS coverage was not provided.
- Whether pytest-xdist is guaranteed present in the subprocess environment.
Automated review by NVCortex Lite, run by @fredricz-20070104.
… drop nested pytest Signed-off-by: shikicloud <shikiw@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
jenkins/L0_Test.groovy (1)
3811-3813: 🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick winVerify the integrity of the downloaded native wheels.
These commands install native-code wheels from public GitHub during CI. The release assets publish SHA-256 checksums, but this command does not verify them. Add per-architecture checksum validation or mirror the wheels into controlled Artifactory before installation. (github.com)
🤖 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 `@jenkins/L0_Test.groovy` around lines 3811 - 3813, Update the native wheel installation block to verify SHA-256 checksums for both causal-conv1d and mamba_ssm assets before pip installation, using the architecture-specific published checksums for mambaArch. Alternatively, source the wheels from controlled Artifactory, but do not install directly from the public URLs without integrity validation.
🤖 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/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py`:
- Around line 44-51: Update RefHFModelWithIPCHandles.__init__ to annotate its
return as None and replace Optional[int] with int | None and Optional[List[str]]
with list[str] | None, preserving the existing constructor behavior and
parameters.
---
Nitpick comments:
In `@jenkins/L0_Test.groovy`:
- Around line 3811-3813: Update the native wheel installation block to verify
SHA-256 checksums for both causal-conv1d and mamba_ssm assets before pip
installation, using the architecture-specific published checksums for mambaArch.
Alternatively, source the wheels from controlled Artifactory, but do not install
directly from the public URLs without integrity validation.
🪄 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: 21608973-137a-4a75-89a8-c7863f851674
📒 Files selected for processing (4)
jenkins/L0_Test.groovyjenkins/scripts/slurm_install.shtests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.pytests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py
| def __init__( | ||
| self, | ||
| model_dir: str, | ||
| device_id: int = 0, | ||
| *, | ||
| num_hidden_layers: Optional[int] = None, | ||
| layers_block_type: Optional[List[str]] = None, | ||
| ): |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Verify the repository Python target before applying Python 3.10+ syntax.
fd -HI '^(pyproject\.toml|setup\.py|\.python-version|\.tool-versions)$' . -0 |
xargs -0 -r rg -n 'requires-python|python_requires|target-version|python'Repository: NVIDIA/TensorRT-LLM
Length of output: 6821
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py"
printf '%s\n' '--- repository guidance ---'
if [ -f CODING_GUIDELINES.md ]; then
cat -n CODING_GUIDELINES.md
else
printf '%s\n' 'CODING_GUIDELINES.md not found'
fi
printf '%s\n' '--- target file ---'
cat -n "$file"
printf '%s\n' '--- relevant imports and constructor usages ---'
rg -n -C 3 '^(from typing|import typing)|RefHFModelWithIPCHandles|num_hidden_layers|layers_block_type' "$file"
printf '%s\n' '--- changed-file summary ---'
git diff --stat -- "$file"
printf '%s\n' '--- test-list references ---'
rg -n 'test_llm_update_weights|update_weights' tests/integration/test_lists tests 2>/dev/null | head -200Repository: NVIDIA/TensorRT-LLM
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target="tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py"
printf '%s\n' '--- target constructor and nearby code ---'
sed -n '43,95p' "$target"
printf '%s\n' '--- test functions in target ---'
rg -n '^def test_|^async def test_|^class ' "$target"
printf '%s\n' '--- base class definition ---'
rg -n -C 8 'class RefHFModel|def __init__' tests utils tensorrt_llm 2>/dev/null | head -240
printf '%s\n' '--- test-list entries for this module ---'
rg -n -C 3 'test_llm_update_weights|ray_orchestrator.*single_gpu|update_weights' tests/integration/test_lists 2>/dev/null || true
printf '%s\n' '--- typing imports and usages in target ---'
python3 - <<'PY'
from pathlib import Path
import ast
path = Path("tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py")
tree = ast.parse(path.read_text())
imports = {}
for node in tree.body:
if isinstance(node, ast.ImportFrom) and node.module == "typing":
for alias in node.names:
imports[alias.asname or alias.name] = alias.name
used = {node.id for node in ast.walk(tree) if isinstance(node, ast.Name)}
for local_name, imported_name in imports.items():
print(f"{imported_name}: {'used' if local_name in used else 'unused'}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 20284
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target="tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py"
printf '%s\n' '--- target diff ---'
git diff --unified=3 -- "$target"
printf '%s\n' '--- base helper candidates ---'
fd -HI -t f 'torch_ref|ref.*model' tests | head -100
printf '%s\n' '--- base helper definitions ---'
rg -n -C 12 'class RefHFModel|def __init__' tests/unittest tests/utils 2>/dev/null | rg -B 12 -A 12 'RefHFModel|torch_ref' | head -240 || true
printf '%s\n' '--- test-list references ---'
if [ -d tests/integration/test_lists ]; then
rg -n -C 3 'test_llm_update_weights|ray_orchestrator|update_weights' tests/integration/test_lists || true
else
printf '%s\n' 'tests/integration/test_lists not found'
fi
printf '%s\n' '--- typing import usage in target ---'
python3 - <<'PY'
from pathlib import Path
import ast
path = Path("tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py")
tree = ast.parse(path.read_text())
typing_names = {}
for node in tree.body:
if isinstance(node, ast.ImportFrom) and node.module == "typing":
for alias in node.names:
typing_names[alias.asname or alias.name] = alias.name
used = {node.id for node in ast.walk(tree) if isinstance(node, ast.Name)}
for local_name, imported_name in typing_names.items():
print(f"{imported_name}: {'used' if local_name in used else 'unused'}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 21409
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- target test decorators and calls ---'
sed -n '170,360p' tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py
printf '%s\n' '--- cross-file helper and Nemotron-H references ---'
rg -n -C 4 'RefHFModelWithIPCHandles|layers_block_type|num_hidden_layers' \
tests/unittest/_torch/ray_orchestrator tests/unittest/utils/torch_ref.py
printf '%s\n' '--- exact CI list entries ---'
sed -n '238,255p' tests/integration/test_lists/test-db/l0_h100.ymlRepository: NVIDIA/TensorRT-LLM
Length of output: 40864
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- Nemotron-H test function ---'
sed -n '620,715p' tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py
rg -n -C 2 'test.*nemotron|test_llm_update_weights.*nemotron|multi_gpu' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa 2>/dev/null || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 50377
Use Python 3.10 type syntax and annotate the constructor return.
RefHFModelWithIPCHandles.__init__ needs a -> None annotation. Replace Optional[int] and Optional[List[str]] with int | None and list[str] | None.
Test coverage summary
- Changed test functions: none. Changed helper:
RefHFModelWithIPCHandles.__init__. - Existing consumers include the single-GPU update-weight tests and
test_llm_update_weights_nemotron_h. - Test lists:
l0_h100.ymlcovers parts 0–2;l0_dgx_b200.ymlcovers part 4. - Verdict: sufficient.
🤖 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/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py`
around lines 44 - 51, Update RefHFModelWithIPCHandles.__init__ to annotate its
return as None and replace Optional[int] with int | None and Optional[List[str]]
with list[str] | None, preserving the existing constructor behavior and
parameters.
Sources: Coding guidelines, Learnings
Signed-off-by: shikicloud <shikiw@nvidia.com> # Conflicts: # tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py
|
/bot run --disable-fail-fast --stage-list "H100_PCIe-PyTorch-Ray-1, DGX_B200-4_GPUs-PyTorch-Ray-1" |
|
PR_Github #63636 [ run ] triggered by Bot. Commit: |
|
PR_Github #63636 [ run ] completed with state |
|
/bot run |
|
PR_Github #63692 [ run ] triggered by Bot. Commit: |
|
PR_Github #63692 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63919 [ run ] triggered by Bot. Commit: |
|
PR_Github #63919 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63952 [ run ] triggered by Bot. Commit: |
|
PR_Github #63952 [ run ] completed with state
|
7 participants
Dev Engineer Review
layers_block_typefrom the checkpoint configuration.causal-conv1dandmamba_ssmwheels by architecture.part4waiver was removed.QA Engineer Review
Modified test functions:
test_llm_update_weights_nemotron_h.mamba_depsfixture._nemotron_h_body._nemotron_h_subprocess_entry.RefHFModelWithIPCHandles.__init__.Coverage:
tests/integration/test_lists/waives.txtpreviously waived the Nemotronpart4case. That waiver was removed.Verdict: needs follow-up — test-list coverage is present, but CBTS coverage data is unavailable.
Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.