fix: probe CUDA context on correct device in set_task_pid() - #229
fix: probe CUDA context on correct device in set_task_pid()#229maverick123123 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 Signed-off-by: maverick123123 <yuming.wu@dynamia.ai>
|
[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
ChangesProbe context alignment
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
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 |
There was a problem hiding this comment.
Actionable comments posted: 1
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 context on every exit path.
cuDevicePrimaryCtxRetain()adds the current process to GPU-memory accounting forprobeDevand setsctx_activate[probeDev]. Returning after later failures bypassescuDevicePrimaryCtxRelease(probeDev), leaving activation and memory-accounting state stuck until process exit. Route post-retain failures through a cleanup label that always releasesprobeDev.🤖 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 path after cuDevicePrimaryCtxRetain succeeds releases probeDev. In the surrounding process-query flow, replace direct returns for NVML failures and host-pid errors with a cleanup path that calls cuDevicePrimaryCtxRelease(probeDev), while preserving the returned error codes; also route normal completion through the same cleanup label.
114-135: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick winReject the no-mapped-device fallback instead of probing CUDA device 0.
If every
nvml_to_cuda_map(i)returns-1,probeDevstays0, so the code retains device 0 even though no NVML device was selected. The scan then returns “host pid is error!” while the wrong device’s primary context/accounting state has already been retained. InitializeprobeDevto an invalid value and return beforecuDevicePrimaryCtxRetainwhen no mapping succeeds.🤖 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 device scan around nvml_to_cuda_map and cuDevicePrimaryCtxRetain to initialize probeDev as invalid and detect when no NVML device mapping succeeds. Return the existing host-PID error result before retaining a CUDA primary context when probeDev remains invalid; retain the current successful-device behavior otherwise.
🤖 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 `@src/utils.c`:
- Around line 127-135: Fix the reported CppLint whitespace issues in the shown
process-query block: add spaces around the == operator in the do-while condition
and after the comma in the mergepid call. Keep the surrounding logic unchanged.
---
Outside diff comments:
In `@src/utils.c`:
- Around line 135-173: Ensure every exit path after cuDevicePrimaryCtxRetain
succeeds releases probeDev. In the surrounding process-query flow, replace
direct returns for NVML failures and host-pid errors with a cleanup path that
calls cuDevicePrimaryCtxRelease(probeDev), while preserving the returned error
codes; also route normal completion through the same cleanup label.
- Around line 114-135: Update the device scan around nvml_to_cuda_map and
cuDevicePrimaryCtxRetain to initialize probeDev as invalid and detect when no
NVML device mapping succeeds. Return the existing host-PID error result before
retaining a CUDA primary context when probeDev remains invalid; retain the
current successful-device behavior otherwise.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
| }while(res==NVML_ERROR_INSUFFICIENT_SIZE); | ||
| mergepid(&previous,&merged_num,(nvmlProcessInfo_t1 *)tmp_pids_on_device,pre_pids_on_device); | ||
| probeDev = cudaDev; | ||
| break; | ||
| } | ||
| previous = merged_num; | ||
| merged_num = 0; | ||
| memset(tmp_pids_on_device,0,sizeof(nvmlProcessInfo_v1_t)*SHARED_REGION_MAX_PROCESS_NUM); | ||
| CHECK_CU_RESULT(cuDevicePrimaryCtxRetain(&pctx,0)); | ||
| CHECK_CU_RESULT(cuDevicePrimaryCtxRetain(&pctx,probeDev)); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix the CppLint whitespace failures.
Use spaces around == at Line 127 and after the comma at Line 135; these currently fail the reported cpplint checks.
Proposed fix
- }while(res==NVML_ERROR_INSUFFICIENT_SIZE);
+ } while (res == NVML_ERROR_INSUFFICIENT_SIZE);
...
- CHECK_CU_RESULT(cuDevicePrimaryCtxRetain(&pctx,probeDev));
+ CHECK_CU_RESULT(cuDevicePrimaryCtxRetain(&pctx, probeDev));📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| }while(res==NVML_ERROR_INSUFFICIENT_SIZE); | |
| mergepid(&previous,&merged_num,(nvmlProcessInfo_t1 *)tmp_pids_on_device,pre_pids_on_device); | |
| probeDev = cudaDev; | |
| break; | |
| } | |
| previous = merged_num; | |
| merged_num = 0; | |
| memset(tmp_pids_on_device,0,sizeof(nvmlProcessInfo_v1_t)*SHARED_REGION_MAX_PROCESS_NUM); | |
| CHECK_CU_RESULT(cuDevicePrimaryCtxRetain(&pctx,0)); | |
| CHECK_CU_RESULT(cuDevicePrimaryCtxRetain(&pctx,probeDev)); | |
| } while (res == NVML_ERROR_INSUFFICIENT_SIZE); | |
| mergepid(&previous,&merged_num,(nvmlProcessInfo_t1 *)tmp_pids_on_device,pre_pids_on_device); | |
| probeDev = cudaDev; | |
| break; | |
| } | |
| previous = merged_num; | |
| merged_num = 0; | |
| memset(tmp_pids_on_device,0,sizeof(nvmlProcessInfo_v1_t)*SHARED_REGION_MAX_PROCESS_NUM); | |
| CHECK_CU_RESULT(cuDevicePrimaryCtxRetain(&pctx, probeDev)); |
🧰 Tools
🪛 Cppcheck (2.21.0)
[warning] 135-135: If resource allocation fails, then there is a possible null pointer dereference
(nullPointerOutOfResources)
🪛 GitHub Check: cpplint
[failure] 135-135:
[cpplint] reported by reviewdog 🐶
Missing space after , [whitespace/comma] [3]
Raw Output:
src/utils.c:135: Missing space after , [whitespace/comma] [3]
[failure] 127-127:
[cpplint] reported by reviewdog 🐶
Missing spaces around == [whitespace/operators] [3]
Raw Output:
src/utils.c:127: Missing spaces around == [whitespace/operators] [3]
🤖 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 127 - 135, Fix the reported CppLint whitespace
issues in the shown process-query block: add spaces around the == operator in
the do-while condition and after the comma in the mergepid call. Keep the
surrounding logic unchanged.
Source: Linters/SAST tools
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)