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 2 commits into
Conversation
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>
|
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. |
📝 WalkthroughWalkthroughAdds an optional memory-monitoring sidecar for HuggingFace PTQ runs, collecting CPU/GPU metrics into CSV and summary files. The change adds ChangesPTQ memory monitoring
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant huggingface_example as huggingface_example.sh
participant mem_monitor as mem_monitor.py
participant hf_ptq as hf_ptq.py
participant outputs as CSV and summary files
huggingface_example->>mem_monitor: invoke when MEM_MONITOR=1
mem_monitor->>hf_ptq: launch wrapped PTQ command
mem_monitor->>mem_monitor: sample CPU and GPU metrics
mem_monitor->>outputs: write trace and peak summary
mem_monitor-->>huggingface_example: return child exit status
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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 78.38% 66.70% -11.69%
===========================================
Files 518 518
Lines 58269 58269
===========================================
- Hits 45674 38868 -6806
- Misses 12595 19401 +6806
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?
4 participants
What does this PR do?
Type of change: New example
Adds an opt-in sidecar that monitors GPU and CPU memory + utilization during an HF PTQ run, without instrumenting
hf_ptq.pyitself. It supports the OMNIML-4947 goal of verifying single-GPU layerwise PTQ stays within tight GPU/CPU memory budgets.examples/hf_ptq/scripts/mem_monitor.py— a standalone, cross-process sampler. It reads GPU memory/utilization at the device level via NVML (falling back tonvidia-smi, and degrading to CPU-only if no driver is reachable) and CPU system + process-tree RSS/utilization viapsutil, writing a CSV timeseries (overwriting any existing file) plus a peak/mean summary. It binds to a workload two ways: wrap mode (-- <cmd>, launches and tracks the child's process tree and propagates its exit code) or standalone (--pid/--duration/ signal).examples/hf_ptq/scripts/huggingface_example.sh— opt-in viaMEM_MONITOR=1(default off → byte-for-byte unchanged behavior); when enabled it wraps the existinghf_ptq.pyinvocation and writesmem_trace.csv+mem_peak.txtunderSAVE_PATH.examples/hf_ptq/requirements.txt— addspsutil(nvidia-ml-pyis already a core dependency inpyproject.toml).tests/examples/hf_ptq/test_mem_monitor.py— CPU-only unit tests.Usage
Testing
pytest tests/examples/hf_ptq/test_mem_monitor.py— 8 passing; covers--gpusparsing (incl. UUID/MIG tokens that must not crash), the disabled-GPU path, CPU/process-tree sampling, CSV+summary output, and wrap-mode exit-code propagation. Subprocess tests are bounded by a default timeout.gpu{i}_used_mb/util columns populate through a workload's allocate → hold → free lifecycle, andmem_peak.txtis written.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.).MEM_MONITOR=1, default off; no changes tohf_ptq.pyor anymodelopt/code.CONTRIBUTING.md: ✅ — no copied code; addspsutilto the example'srequirements.txt(nvidia-ml-pywas already a core dependency).tests/examples/hf_ptq/test_mem_monitor.py.Additional Information
Summary by CodeRabbit
New Features
MEM_MONITOR=1).psutilto enable CPU and process memory tracking.Tests