fix: probe CUDA context on correct device in set_task_pid() - #228
fix: probe CUDA context on correct device in set_task_pid()#228maverick123123 wants to merge 1 commit into
Conversation
set_task_pid() polls NVML processes on the first device where nvml_to_cuda_map(i) >= 0, but the probe context was always created on hardcoded CUDA device 0 via cuDevicePrimaryCtxRetain(&pctx, 0). When CUDA_VISIBLE_DEVICES does not start with the pod's device 0 (e.g., CVD=1 or CVD=1,0), the polled NVML device and the probe context device are different physical GPUs. getextrapid() never finds the new PID, causing a spurious 'host pid is error!' and subsequent Device 0 OOM due to incorrect memory accounting. Fix: use the cudaDev found in the NVML enumeration loop for both cuDevicePrimaryCtxRetain and cuDevicePrimaryCtxRelease, ensuring the probe context runs on the same device being polled. Fixes Project-HAMi#225
5785fb8 to
8efff49
Compare
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: maverick123123 The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
📝 WalkthroughWalkthrough
ChangesCUDA probe device alignment
Estimated code review effort: 2 (Simple) | ~5 minutes Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/utils.c (2)
135-173: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRelease the retained primary context on every post-retain exit path.
After
cuDevicePrimaryCtxRetain(&pctx,probeDev), the polling failures at line 146 and thehostpid == 0path at line 161 return beforecuDevicePrimaryCtxRelease(probeDev), leaving the retained primary-context reference unreleased. Route these exits through cleanup that releasespctxbefore returning.🤖 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 `@src/utils.c` around lines 135 - 173, Ensure every exit after cuDevicePrimaryCtxRetain in the surrounding function releases the retained primary context. Update the polling failure branch in the nvmlDeviceGetComputeRunningProcesses loop and the hostpid == 0 branch to route through cleanup that calls cuDevicePrimaryCtxRelease(probeDev) before returning, while preserving each branch’s existing return status.
114-135: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winDo not fall back to CUDA device 0 when no mapping succeeds.
probeDevstays0if everynvml_to_cuda_map()call fails, socuDevicePrimaryCtxRetain(&pctx, probeDev)selects CUDA device 0 unintentionally. Start the selector as invalid and release only after a probe device was selected; this avoids per-device accounting changes on a device that was never successfully mapped.🤖 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 `@src/utils.c` around lines 114 - 135, Update the probe-device selection around probeDev and cuDevicePrimaryCtxRetain: initialize probeDev to an invalid sentinel, and only retain the primary context after the loop when a valid mapping was found. Preserve the existing process-accounting behavior for successful mappings, while avoiding any fallback to CUDA device 0 when all nvml_to_cuda_map calls fail.
🤖 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.
Outside diff comments:
In `@src/utils.c`:
- Around line 135-173: Ensure every exit after cuDevicePrimaryCtxRetain in the
surrounding function releases the retained primary context. Update the polling
failure branch in the nvmlDeviceGetComputeRunningProcesses loop and the hostpid
== 0 branch to route through cleanup that calls
cuDevicePrimaryCtxRelease(probeDev) before returning, while preserving each
branch’s existing return status.
- Around line 114-135: Update the probe-device selection around probeDev and
cuDevicePrimaryCtxRetain: initialize probeDev to an invalid sentinel, and only
retain the primary context after the loop when a valid mapping was found.
Preserve the existing process-accounting behavior for successful mappings, while
avoiding any fallback to CUDA device 0 when all nvml_to_cuda_map calls fail.
|
Thanks for your pull request. Before we can look at it, you'll need to add a 'DCO signoff' to your commits. 📝 Please follow instructions in the contributing guide to update your commits with the DCO Full details of the Developer Certificate of Origin can be found at developercertificate.org. The list of commits missing DCO signoff:
DetailsInstructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. I understand the commands that are listed here. |
1 participant
Problem
set_task_pid()insrc/utils.cpolls NVML processes on the first physical device wherenvml_to_cuda_map(i) >= 0, but the probe context is always created on hardcoded CUDA device 0 viacuDevicePrimaryCtxRetain(&pctx, 0).When
CUDA_VISIBLE_DEVICESdoes not start with the pod's device 0 (e.g.,CVD=1orCVD=1,0), the polled NVML device and the probe context device are different physical GPUs.getextrapid()never finds the new PID, causing a spurious "host pid is error!" and subsequent fake Device 0 OOM due to incorrect memory accounting.Fix
Save the
cudaDevfound in the NVML enumeration loop and use it for bothcuDevicePrimaryCtxRetainandcuDevicePrimaryCtxRelease, ensuring the probe context runs on the same device being polled.Testing
Verified on an 8x RTX 3090 machine with all CVD combinations.
(Closed in favor of #230)