Add sidecar GPU/CPU memory+utilization monitor for HF PTQ#2000
Add sidecar GPU/CPU memory+utilization monitor for HF PTQ#2000Fridah-nv wants to merge 4 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
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:
📝 WalkthroughWalkthroughAdds a standalone GPU and CPU memory monitor with CSV and summary output, integrates it optionally into the HuggingFace PTQ example, adds ChangesMemory monitoring
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant User
participant huggingface_example
participant mem_monitor
participant hf_ptq
participant SAVE_PATH
User->>huggingface_example: set MEM_MONITOR=1
huggingface_example->>mem_monitor: launch monitoring wrapper
mem_monitor->>hf_ptq: execute quantization
mem_monitor->>SAVE_PATH: write CSV trace and peak summary
mem_monitor-->>huggingface_example: propagate child exit status
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2000 +/- ##
==========================================
+ Coverage 69.45% 75.73% +6.27%
==========================================
Files 518 518
Lines 58544 58544
==========================================
+ Hits 40664 44336 +3672
+ Misses 17880 14208 -3672
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:
|
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: Full review of all 4 changed files (examples/hf_ptq/scripts/mem_monitor.py, huggingface_example.sh, requirements.txt, tests/examples/hf_ptq/test_mem_monitor.py). This is a self-contained sidecar utility under examples/ — it does not touch modelopt/ core, mode registration, config schema, or export paths, so the Mode/State and Export categories do not apply.
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1
The one suggestion is a doc/behavior mismatch: the module docstring and PR body say the trace "appends a CSV timeseries," but --out is opened with mode "w", which truncates. Non-blocking — reword to "overwrites" (or switch to true append with a header guard).
Things I verified as correct:
- Exit-code propagation: for a signal-killed child
returncode == -signum, so128 - (-signum) == 128 + signum, matching shell convention; normal codes pass through unchanged. set -esafety in the shell wrapper: whenMEM_MONITOR != 1theMEM_MON_PREFIXarray is empty and"${MEM_MON_PREFIX[@]}"expands to nothing, leaving the originalpython hf_ptq.py …invocation and its behavior fully unchanged (opt-in, default off).nvidia-ml-pyis the correct distribution providing the importedpynvmlmodule (not the deprecatedpynvmlpackage).- Edge cases in
_resolve_gpu_indices(UUID/MIG tokens,all/none),_cell(util0vs missingNone), NVML→smi fallback, and process-tree RSS aggregation are handled and covered by the CPU-only tests.
Risk: Low. New opt-in example tooling with no changes to library code or default behavior; test coverage is adequate for the CPU path.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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: 4
🧹 Nitpick comments (2)
tests/examples/hf_ptq/test_mem_monitor.py (1)
91-92: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winBound subprocess test execution.
A regression in monitor termination can hang all three end-to-end tests indefinitely. Set a default timeout while preserving per-test overrides.
Proposed fix
def _run(args, **kwargs): + kwargs.setdefault("timeout", 30) return subprocess.run([sys.executable, str(_SCRIPT), *args], **kwargs)As per coding guidelines, “Respect the per-test timeout.”
🤖 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/examples/hf_ptq/test_mem_monitor.py` around lines 91 - 92, Update the _run helper to provide a finite default timeout to subprocess.run, while allowing callers’ per-test timeout values in kwargs to override that default. Preserve all existing command construction and subprocess options.Source: Coding guidelines
examples/hf_ptq/scripts/mem_monitor.py (1)
106-125: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winLocal
import pynvmllacks the required justification comment.As per coding guidelines,
`**/*.py`: Keep imports at the top of Python source and test files; use local imports only for justified circular dependencies, optional dependencies, or unusually heavy imports, with a brief explanatory comment.This import is a justified optional-dependency case (NVML may be absent, falls back tonvidia-smi), but has no comment stating that.🔧 Proposed fix
def _init_nvml(self, indices: list[int] | None) -> bool: try: + # Optional dependency: pynvml (nvidia-ml-py) may not be installed; + # fall back to `nvidia-smi` parsing via _init_smi on any failure. import pynvml🤖 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 `@examples/hf_ptq/scripts/mem_monitor.py` around lines 106 - 125, Add a brief explanatory comment immediately before the local pynvml import in _init_nvml, documenting that it is intentionally local because NVML/pynvml is an optional dependency and the monitor falls back to nvidia-smi when unavailable. Keep the existing import placement and fallback behavior unchanged.Source: Coding guidelines
🤖 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 `@examples/hf_ptq/scripts/mem_monitor.py`:
- Around line 291-293: Guard the psutil.Process initialization in the monitoring
setup around target_pid so stale or already-exited PIDs are handled cleanly
instead of propagating psutil.NoSuchProcess. Preserve the child-process path and
existing behavior for valid PIDs, while emitting the established user-facing
error and exiting through the script’s normal failure flow.
- Around line 268-284: Update _write_summary to ensure the parent directory of
the provided summary path exists before calling Path(path).write_text, while
preserving the existing no-path behavior. This must allow the later exit-status
handling around the wrapped command to run so the child’s return code is
propagated even when args.summary targets a missing directory.
- Around line 270-274: Gate the peak GPU memory output in the metrics summary
using mem_acc.seen, matching the existing util_acc.seen check. Update the loop
over metrics.gpu so peak_gpu{i}_used_mb is emitted only when a valid memory
sample was observed, while preserving the existing formatting and utilization
behavior.
- Around line 352-354: Update the child cleanup block to wait for the terminated
process with a finite timeout, then escalate by killing it if the timeout
expires and wait for termination to complete. Preserve the existing cleanup flow
so _write_summary and exit-code propagation still execute even when the child
ignores SIGTERM.
---
Nitpick comments:
In `@examples/hf_ptq/scripts/mem_monitor.py`:
- Around line 106-125: Add a brief explanatory comment immediately before the
local pynvml import in _init_nvml, documenting that it is intentionally local
because NVML/pynvml is an optional dependency and the monitor falls back to
nvidia-smi when unavailable. Keep the existing import placement and fallback
behavior unchanged.
In `@tests/examples/hf_ptq/test_mem_monitor.py`:
- Around line 91-92: Update the _run helper to provide a finite default timeout
to subprocess.run, while allowing callers’ per-test timeout values in kwargs to
override that default. Preserve all existing command construction and subprocess
options.
🪄 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: 8ba3a711-1e7f-4b0e-8809-7d30c4c8195c
📒 Files selected for processing (4)
examples/hf_ptq/requirements.txtexamples/hf_ptq/scripts/huggingface_example.shexamples/hf_ptq/scripts/mem_monitor.pytests/examples/hf_ptq/test_mem_monitor.py
| compressed-tensors | ||
| fire | ||
| flash-attn>=2.6.0 | ||
| nvidia-ml-py |
There was a problem hiding this comment.
nvidia-ml-py is already a generic dependency in pyproject.toml so should not be listed here
There was a problem hiding this comment.
Thanks for the catch, removed nvidia-ml-py from examples/hf_ptq/requirements.txt
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-4-8) — DM the bot to share feedback.
Adds a standalone, opt-in sidecar GPU/CPU memory+utilization monitor (examples/hf_ptq/scripts/mem_monitor.py) for HF PTQ runs, wired into huggingface_example.sh behind MEM_MONITOR=1 (default off). +542/-1, 4 files. The code is clean and well-documented; license headers match the canonical LICENSE_HEADER; no prompt-injection attempts in the PR content. Correctness spot-checks pass: the shell empty-array expansion "${MEM_MON_PREFIX[@]}" python hf_ptq.py is a genuine no-op when off; the signal exit-code math (128 - rc where rc = -N → 128 + N) is correct; the NVML→nvidia-smi fallback via or short-circuits properly; wrap-mode child cleanup (terminate()/wait()) on signal is handled. CPU-only unit + subprocess (wrap/standalone) tests are meaningful.
Flagging for a quick owner look because:
- CHANGELOG not updated for this new feature (the PR checklist leaves it unmarked). Minor for opt-in tooling with no behavior change, but per repo policy new features get a CHANGELOG entry — confirm whether intended.
- GPU sampling paths (NVML + nvidia-smi) are never exercised in CI — tests only cover the disabled (
--gpus none) path, soGpuSampler.sample(),_query_smi()parsing, and the physical-index resolution are unverified by automation. Understandable given CPU-only CI, but worth noting. CUDA_VISIBLE_DEVICES→ NVML physical-index mismatch: the shell passes--gpus "${CUDA_VISIBLE_DEVICES-all}", butCUDA_VISIBLE_DEVICEScan hold logically-remapped indices that don't match NVML/PCI physical indices unlessCUDA_DEVICE_ORDER=PCI_BUS_IDis set. In that case the monitor silently reports the wrong physical GPUs. The docstring documents this caveat, but the shell wiring doesn't set/enforceCUDA_DEVICE_ORDER, so the sidecar could mis-report memory against the wrong device without any warning.
- requirements.txt: drop nvidia-ml-py (already a core dep in pyproject.toml);
keep psutil, which is not.
- _write_summary: create the summary parent dir before writing so a missing
directory can't crash the run and swallow the child's exit code; gate
peak_gpu{i} on .seen so a never-sampled GPU is omitted instead of reported 0.0.
- Guard psutil.Process against a stale/exited --pid (clean error + exit 1).
- Escalate child.terminate() to SIGKILL after a 10s timeout so an unresponsive
child cannot hang the monitor.
- Docs: 'writes (overwriting any existing file)' instead of 'appends'; add the
optional-dependency justification comment for the local pynvml import.
- Tests: default 30s subprocess timeout in the _run helper.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
|
I think this tool should be put into tools folder, it can be used anywhere else. |
shengliangxu
left a comment
shengliangxu
left a comment
There was a problem hiding this comment.
And could you paste an example output CSV?
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Could you take a look at https://github.com/NVIDIA/Model-Optimizer/blob/main/modelopt/torch/utils/memory_monitor.py?
Per PR #2000 review (shengliangxu + reuse feedback), consolidate onto the existing in-process GPUMemoryMonitor (already called by hf_ptq.py / hf_sa.py) instead of a separate sidecar. GPUMemoryMonitor now also samples GPU utilization and CPU (system + process-tree RSS + util), tracks peak/mean, and optionally writes mem_trace.csv + mem_peak.txt when given output_dir. hf_ptq enables file output via MEM_MONITOR=1 (to a sibling dir, not the checkpoint). Adds psutil to core deps; removes examples/hf_ptq/scripts/mem_monitor.py and the shell wrapper. Hardening from the high-effort code review: - launch_memory_monitor never raises (catch Exception) so the diagnostic can't crash the workload; NVML init wrapped -> degrades to CPU-only if unavailable. - Separate try blocks so an unsupported utilization query (MIG/vGPU) no longer drops the GPU memory sample. - mem_peak.txt refreshed every tick so a peak report survives SIGTERM/OOM. - Scope GPUs to CUDA_VISIBLE_DEVICES so shared nodes don't fold other tenants. - Idempotent stop(); no double nvmlShutdown; no NVML init leak on zero devices. - Summary labels main-process CPU accurately. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
|
| # out of the export path so it doesn't pollute the checkpoint). | ||
| launch_memory_monitor( | ||
| output_dir=f"{args.export_path}_mem_monitor" | ||
| if os.environ.get("MEM_MONITOR") == "1" |
There was a problem hiding this comment.
Lets name it MODELOPT_... env variable
| if os.environ.get("MEM_MONITOR") == "1" | |
| if os.environ.get("MODELOPT_MEM_MONITOR") == "1" |
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: 2
🧹 Nitpick comments (3)
tests/unit/torch/utils/test_memory_monitor.py (2)
125-130: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert that the launched monitor is actually running.
The test only checks
monitor is not None, so it would pass iflaunch_memory_monitor()returned a stopped monitor. Assertmonitor.is_runningbefore callingstop()(and optionally verify that its resolved device list is empty).🤖 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/unit/torch/utils/test_memory_monitor.py` around lines 125 - 130, Strengthen test_launch_returns_running_monitor_without_gpu by asserting monitor.is_running immediately after launch_memory_monitor() returns, before stopping it; optionally also assert the monitor’s resolved device list is empty for the no-GPU case.Source: Path instructions
29-38: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winCover NVML initialization failure.
_patch_nvmlalways makesnvmlInit()succeed, so the CPU-only tests do not exercise theNVMLErrorfallback when NVML is unavailable. Add a test that makesnvmlInitraise and verifies that a usable CPU-only monitor is returned.As per path instructions, tests must exercise the behavior they claim to validate; this includes the documented NVML-unavailable fallback.
🤖 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/unit/torch/utils/test_memory_monitor.py` around lines 29 - 38, Extend the tests around _patch_nvml to add an NVML-unavailable case where nvmlInit raises NVMLError, then verify the monitor factory returns a usable CPU-only monitor. Keep the existing successful-NVML tests unchanged and assert the fallback behavior through the public monitor creation API.Source: Path instructions
modelopt/torch/utils/memory_monitor.py (1)
35-56: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
__all__to declare the module's public API.The module exposes
GPUMemoryMonitorandlaunch_memory_monitoras its intended public surface, but no__all__is defined.As per coding guidelines, "Define each module's public API with
__all__ = [...]."♻️ Proposed addition
GB = 1024**3 + +__all__ = ["GPUMemoryMonitor", "launch_memory_monitor"]🤖 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 `@modelopt/torch/utils/memory_monitor.py` around lines 35 - 56, Define the module’s public API with an __all__ declaration listing GPUMemoryMonitor and launch_memory_monitor, while leaving the existing imports and implementation unchanged.Source: Coding guidelines
🤖 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 `@modelopt/torch/utils/memory_monitor.py`:
- Around line 115-124: Update _resolve_devices so it distinguishes an unset
CUDA_VISIBLE_DEVICES variable from an explicitly empty value: return all devices
only when the environment variable is absent, and return no devices when it is
present but empty. Preserve the existing parsing and filtering behavior for
non-empty values.
In `@tests/unit/torch/utils/test_memory_monitor.py`:
- Around line 41-47: Update the test helper _run and the affected idempotency
test to remove fixed time.sleep calls and use deterministic synchronization
instead. After starting the monitor, poll a bounded condition such as
monitor.cpu_mem._count > 0 before stopping it; omit waiting entirely where the
idempotency test does not require a sample, while preserving bounded test
execution.
---
Nitpick comments:
In `@modelopt/torch/utils/memory_monitor.py`:
- Around line 35-56: Define the module’s public API with an __all__ declaration
listing GPUMemoryMonitor and launch_memory_monitor, while leaving the existing
imports and implementation unchanged.
In `@tests/unit/torch/utils/test_memory_monitor.py`:
- Around line 125-130: Strengthen
test_launch_returns_running_monitor_without_gpu by asserting monitor.is_running
immediately after launch_memory_monitor() returns, before stopping it;
optionally also assert the monitor’s resolved device list is empty for the
no-GPU case.
- Around line 29-38: Extend the tests around _patch_nvml to add an
NVML-unavailable case where nvmlInit raises NVMLError, then verify the monitor
factory returns a usable CPU-only monitor. Keep the existing successful-NVML
tests unchanged and assert the fallback behavior through the public monitor
creation API.
🪄 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: d419acac-59cb-4d04-89de-3002b8747979
📒 Files selected for processing (4)
examples/hf_ptq/hf_ptq.pymodelopt/torch/utils/memory_monitor.pypyproject.tomltests/unit/torch/utils/test_memory_monitor.py
|
/claude review |
kevalmorabia97
left a comment
kevalmorabia97
left a comment
There was a problem hiding this comment.
approving as codeowner for pyproject.toml
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: Full review of all 4 changed files in this PR (modelopt/torch/utils/memory_monitor.py, examples/hf_ptq/hf_ptq.py, pyproject.toml, tests/unit/torch/utils/test_memory_monitor.py). Note: a plain git diff origin/main on this shallow checkout surfaces unrelated commits because the merge base is absent; gh pr diff (the authoritative changed-file list) confines the PR to these 4 files.
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1
Most impactful finding (non-blocking): _resolve_devices maps CUDA_VISIBLE_DEVICES integer values directly to NVML physical indices. These orderings only coincide under CUDA_DEVICE_ORDER=PCI_BUS_ID; with CUDA's default FASTEST_FIRST ordering the monitor can attribute another tenant's GPU to this run without warning — which undercuts the stated 'don't fold in other tenants' GPUs' goal on shared nodes. Suggested either a docstring caveat or UUID-based matching (which also handles the MIG/UUID form that currently falls back to monitoring all devices).
Verified as correct:
_Statpeak/mean aggregation; mean guards against zero count.- NVML degradation:
nvmlInit/count failures caught, monitor continues CPU-only;_nvml_readygatesnvmlShutdown. - Per-device sampling correctly records GPU memory even when
nvmlDeviceGetUtilizationRatesraises (MIG/vGPU), via separate try blocks — covered by a dedicated test. stop()is idempotent (_stoppedguard) so the atexit + explicit-call double invocation runsnvmlShutdownonce.- CSV:
DictWriter(restval="")tolerates rows missing the util column; fieldnames derive from the same_deviceslist used for sampling. _write_summaryrefreshed each tick so a peak report survives SIGTERM/OOM beforestop().launch_memory_monitornever raises; both existing callers (hf_ptq.py,hf_sa.py) ignore the return value, and the new GPU-less behavior (returns a running CPU-only monitor rather thanNone) is backward-compatible for them.- Not a modelopt public-API/mode/state/export change:
memory_monitoris a diagnostic util, nomodelopt_stateinteraction;output_dirdefaults toNoneso default behavior is unchanged.
CodeRabbit already raised __all__, the empty-vs-unset CUDA_VISIBLE_DEVICES distinction, and test hardening (timeouts, NVML-failure coverage, is_running assertion) — not duplicated here.
Risk: Low. Opt-in diagnostic enhancement to an existing in-process utility; no changes to optimization flows, default behavior, or checkpoint state.
abbd92c to
77dab2e
Compare
I think the two solve different problems, so I'd like to keep the sidecar. Also there's thin overlap between them so we might not want to unify them. Where the sidecar wins (and why it matters for this task):
Where the in-process monitor wins:
In this case, the in-process monitor is more for debugging/optimizing memory, side car monitor more for benchmarking like a black-box |
Currently it's only consumed by PTQ workflow, do you think a non-PTQ workflow will benefit from it? |
Adds examples/hf_ptq/scripts/mem_monitor.py, a standalone cross-process sidecar that samples device-level GPU (NVML, nvidia-smi fallback) and CPU (psutil) memory + utilization into a CSV trace and peak/mean summary. It binds to a workload via wrap mode (-- <cmd>, propagates the child exit code) or standalone (--pid/--duration/signal), keeping profiling out of hf_ptq.py. Opt-in from huggingface_example.sh via MEM_MONITOR=1 (default off, no behavior change). Adds psutil + nvidia-ml-py deps and CPU-only unit tests. Part of OMNIML-4947 (single-GPU layerwise PTQ memory validation). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
- requirements.txt: drop nvidia-ml-py (already a core dep in pyproject.toml);
keep psutil, which is not.
- _write_summary: create the summary parent dir before writing so a missing
directory can't crash the run and swallow the child's exit code; gate
peak_gpu{i} on .seen so a never-sampled GPU is omitted instead of reported 0.0.
- Guard psutil.Process against a stale/exited --pid (clean error + exit 1).
- Escalate child.terminate() to SIGKILL after a 10s timeout so an unresponsive
child cannot hang the monitor.
- Docs: 'writes (overwriting any existing file)' instead of 'appends'; add the
optional-dependency justification comment for the local pynvml import.
- Tests: default 30s subprocess timeout in the _run helper.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
77dab2e to
11d95e6
Compare
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
🧹 Nitpick comments (1)
tests/examples/hf_ptq/test_mem_monitor.py (1)
91-169: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider a regression test for the stale-
--pidguard.The previously-fixed
psutil.NoSuchProcessguard around--pid(mem_monitor.py lines 296-301) has no corresponding subprocess test here (e.g., passing an already-exited PID and asserting a clean non-zero exit instead of a traceback). Given this was a real bug fixed mid-review, a small regression test would guard against reintroduction.🤖 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/examples/hf_ptq/test_mem_monitor.py` around lines 91 - 169, The tests lack coverage for the stale --pid handling in the mem-monitor entry point. Add a subprocess regression test near test_standalone_writes_csv_and_summary that launches a short-lived process, waits for it to exit, invokes _run with --pid set to that process ID, and asserts a clean non-zero return without a traceback in stderr.
🤖 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 `@examples/hf_ptq/scripts/mem_monitor.py`:
- Around line 142-159: Harden the SMI sampling path by adding a timeout to
_query_smi’s subprocess.check_output call and handling its subprocess failures
in sample() like the NVML fallback. On timeout or command failure, return the
existing unavailable/empty sample result so the monitoring loop continues and
cleanup and summary handling remain reachable.
---
Nitpick comments:
In `@tests/examples/hf_ptq/test_mem_monitor.py`:
- Around line 91-169: The tests lack coverage for the stale --pid handling in
the mem-monitor entry point. Add a subprocess regression test near
test_standalone_writes_csv_and_summary that launches a short-lived process,
waits for it to exit, invokes _run with --pid set to that process ID, and
asserts a clean non-zero return without a traceback in stderr.
🪄 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: 88e1d5eb-64cd-45ba-95c8-af557b9388ab
📒 Files selected for processing (4)
examples/hf_ptq/requirements.txtexamples/hf_ptq/scripts/huggingface_example.shexamples/hf_ptq/scripts/mem_monitor.pytests/examples/hf_ptq/test_mem_monitor.py
| def _query_smi() -> list[tuple[int, int, int | None]]: | ||
| out = subprocess.check_output( | ||
| [ | ||
| "nvidia-smi", | ||
| "--query-gpu=index,memory.used,utilization.gpu", | ||
| "--format=csv,noheader,nounits", | ||
| ], | ||
| text=True, | ||
| ) | ||
| rows = [] | ||
| for line in out.strip().splitlines(): | ||
| index, used_mb, util = (part.strip() for part in line.split(",")) | ||
| try: | ||
| util_pct: int | None = int(util) | ||
| except ValueError: | ||
| util_pct = None # e.g. "[N/A]" on MIG / unsupported devices | ||
| rows.append((int(index), int(used_mb) * MB, util_pct)) | ||
| return rows |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
sample()'s smi branch can crash the whole monitor on a single failed nvidia-smi call.
Unlike the NVML branch (lines 163-172), which catches NVMLError per-device and degrades gracefully, the smi branch at line 174 calls self._query_smi() with no exception handling. _query_smi()'s subprocess.check_output (lines 143-150) also has no timeout. A transient nvidia-smi failure (nonzero exit, driver reset, ECC error) raises CalledProcessError uncaught inside the sampling loop (main() line 335), killing the monitor process entirely — and since the exception happens before the child.terminate()/_write_summary() cleanup path, the wrapped workload is orphaned and no summary/exit-code propagation occurs. A hang in nvidia-smi (no timeout) blocks the loop indefinitely instead. This directly undermines the "avoids crashing workloads if monitoring fails" design goal for the smi fallback backend.
🔧 Proposed fix
`@staticmethod`
def _query_smi() -> list[tuple[int, int, int | None]]:
out = subprocess.check_output(
[
"nvidia-smi",
"--query-gpu=index,memory.used,utilization.gpu",
"--format=csv,noheader,nounits",
],
text=True,
+ timeout=5,
) if self._backend == "smi":
- return {i: (used, util) for i, used, util in self._query_smi() if i in self._wanted}
+ try:
+ return {i: (used, util) for i, used, util in self._query_smi() if i in self._wanted}
+ except Exception:
+ return {}Also applies to: 161-175
🧰 Tools
🪛 ast-grep (0.44.1)
[error] 142-149: Command coming from incoming request
Context: subprocess.check_output(
[
"nvidia-smi",
"--query-gpu=index,memory.used,utilization.gpu",
"--format=csv,noheader,nounits",
],
text=True,
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').
(subprocess-from-request)
[error] 142-149: Avoid command injection
Context: subprocess.check_output(
[
"nvidia-smi",
"--query-gpu=index,memory.used,utilization.gpu",
"--format=csv,noheader,nounits",
],
text=True,
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').
(command-injection-python)
🤖 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 `@examples/hf_ptq/scripts/mem_monitor.py` around lines 142 - 159, Harden the
SMI sampling path by adding a timeout to _query_smi’s subprocess.check_output
call and handling its subprocess failures in sample() like the NVML fallback. On
timeout or command failure, return the existing unavailable/empty sample result
so the monitoring loop continues and cleanup and summary handling remain
reachable.
Extend the sidecar to also sample GPU power draw (nvmlDeviceGetPowerUsage) and
temperature (nvmlDeviceGetTemperature), and CPU total + free memory
(psutil.virtual_memory().total/.available). New CSV columns gpu{i}_power_w,
gpu{i}_temp_c, sys_cpu_total_mb, sys_cpu_free_mb; summary gains peak power/temp,
total (once) and min free. The nvidia-smi fallback query is extended with
power.draw,temperature.gpu, and each GPU metric is now read in its own guarded
block so an unsupported one (MIG/vGPU) no longer drops the others.
Live-validated on GPU 2,3 (power ~26/31 W, temp ~40/39 C) and CPU-side locally.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
|
/claude review |
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 `@examples/hf_ptq/scripts/mem_monitor.py`:
- Around line 309-326: Update _write_summary to emit peak_sys_cpu_used_mb only
when metrics.sys_used has been sampled and mean_sys_cpu_util_pct only when
metrics.sys_util has been sampled, matching the existing .seen guards used for
other summary metrics. Preserve the current formatting and omit both lines when
no samples were collected.
🪄 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: d7be1ff2-1003-4b54-b309-f43af58f5340
📒 Files selected for processing (2)
examples/hf_ptq/scripts/mem_monitor.pytests/examples/hf_ptq/test_mem_monitor.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/examples/hf_ptq/test_mem_monitor.py
| vm = psutil.virtual_memory() | ||
| sys_util = psutil.cpu_percent(None) | ||
| if proc is None: | ||
| return CpuStat(vm.total, vm.used, vm.available, sys_util, None, None) |
There was a problem hiding this comment.
[SUGGESTION] The sys_free field (surfaced as the sys_cpu_free_mb CSV column and min_sys_cpu_free_mb summary line) is populated from psutil.virtual_memory().available, not .free. These are meaningfully different: available is the kernel's estimate of allocatable memory including reclaimable page cache, whereas free is truly-unused RAM.
Using available is the correct choice for a headroom/budget monitor — .free would badly understate remaining capacity on a machine with a warm page cache. The only imprecision is the label. Consider renaming the column/metric to sys_cpu_avail_mb / min_sys_cpu_avail_mb (or noting "available" in the docstring on line 212) so readers of the CSV don't interpret it as raw free. Non-blocking.
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: Full review of the current sidecar-approach diff — all 4 changed files (examples/hf_ptq/scripts/mem_monitor.py, huggingface_example.sh, requirements.txt, tests/examples/hf_ptq/test_mem_monitor.py). This is self-contained example tooling under examples/; it does not touch modelopt/ core, mode registration, config schema, or export paths, so the Mode/State and Export categories don't apply.
Findings this round: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1
The one new suggestion is a labeling nit on the latest commit's CPU additions: sys_cpu_free_mb / min_sys_cpu_free_mb are populated from virtual_memory().available (the correct choice for a headroom monitor), not .free — worth renaming or documenting so CSV readers don't misread it. Non-blocking.
Verified as correct (including the latest power/temp/free additions):
- GPU power (
mW → W, /1000) and temperature reads are each wrapped in their ownsuppress(NVMLError), so an unsupported metric on MIG/vGPU doesn't drop the others — matches the memory/util pattern. nvidia-smi[N/A]fields parse toNonevia_to_numberand are emitted as blank CSV cells / skipped in the summary (_cell+.seen/.min is not Noneguards); util0is correctly distinguished from missingNone.- Exit-code propagation (
128 - rcfor a signal-killed child),set -e-safe empty-array shell expansion whenMEM_MONITOR != 1, NVML→smi fallback, and process-tree RSS aggregation all still hold.
Already addressed / not duplicated: append-vs-overwrite doc wording (now "overwriting"), local pynvml optional-dependency comment, subprocess test timeout, --pid NoSuchProcess guard, _write_summary parent-dir mkdir, _query_smi subprocess timeout, and the CUDA_VISIBLE_DEVICES→NVML physical-index caveat (raised by @cjluo-nv, documented in the module docstring). The tools/-folder placement and unify-with-memory_monitor.py discussions are design calls for the owners, not code correctness.
Risk: Low. Opt-in (MEM_MONITOR=1, default off) diagnostic tooling with no changes to optimization flows, default behavior, or checkpoint state; CPU path has adequate subprocess + unit coverage.
- --gpus now uses nargs=+ so both space-separated (--gpus 0 1 2 3) and CSV (--gpus 2,3) work; previously the natural space-separated form crashed argparse. Backward compatible with the CSV callers (huggingface_example.sh). - Gate peak_sys_cpu_used_mb / mean_sys_cpu_util_pct on .seen, matching the other summary lines (per CodeRabbit). - Document that sys_free is virtual_memory().available (reclaimable-inclusive), the right headroom measure, not raw .free (per Claude review). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
4 participants
What does this PR do?
Type of change: New feature
Extends the existing in-process
GPUMemoryMonitor(modelopt/torch/utils/memory_monitor.py) — already launched byhf_ptq.pyandhf_sa.py— to also capture CPU and utilization, instead of adding a separate tool.GPUMemoryMonitornow samples, per interval, in its background thread:CUDA_VISIBLE_DEVICESwhen set so shared nodes don't fold in other tenants' GPUs;output_dir, it writesmem_trace.csv(per-sample timeseries) and continuously refreshesmem_peak.txt, so a peak report survives a SIGTERM/OOM kill.hf_ptq.pyenables file output viaMEM_MONITOR=1, writing to a sibling"{export_path}_mem_monitor"dir (kept out of the checkpoint).launch_memory_monitornever raises (an opt-in diagnostic can't crash the workload).psutilto core dependencies; removes the standaloneexamples/hf_ptq/scripts/mem_monitor.py.Usage
Testing
tests/unit/torch/utils/test_memory_monitor.py, CPU-only, NVML monkeypatched, 7 tests):_Statpeak/mean; CPU-only trace + summary; GPU columns/peak;CUDA_VISIBLE_DEVICESscoping; GPU memory still recorded when utilization is unsupported (MIG/vGPU); idempotentstop()(singlenvmlShutdown);launch_memory_monitorreturns a running CPU-only monitor when no GPU is present.launch_memory_monitor()/hf_sa.pyare exercised unchanged (defaultoutput_dir=None).Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).output_dirdefaults toNone; existinglaunch_memory_monitor()callers keep working (the default summary now also includes CPU). Nomodeloptpublic API removed.CONTRIBUTING.md: ✅ — no copied code; addspsutiltopyproject.tomlcore dependencies (nvidia-ml-pywas already core).tests/unit/torch/utils/test_memory_monitor.py.Additional Information
Summary by CodeRabbit
psutildependency for monitoring.