Skip to content

fix: probe CUDA context on correct device in set_task_pid() - #229

Closed
maverick123123 wants to merge 1 commit into
Project-HAMi:mainfrom
maverick123123:fix/set-task-pid-probe-device-mismatch
Closed

fix: probe CUDA context on correct device in set_task_pid()#229
maverick123123 wants to merge 1 commit into
Project-HAMi:mainfrom
maverick123123:fix/set-task-pid-probe-device-mismatch

Conversation

@maverick123123

@maverick123123 maverick123123 commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

Problem

set_task_pid() in src/utils.c polls NVML processes on the first physical device where nvml_to_cuda_map(i) >= 0, but the probe context is 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 fake Device 0 OOM due to incorrect memory accounting.

Fix

Save the cudaDev found in the NVML enumeration loop and use it for both cuDevicePrimaryCtxRetain and cuDevicePrimaryCtxRelease, 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)

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>
@hami-robot

hami-robot Bot commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by: maverick123123
Once this PR has been reviewed and has the lgtm label, please assign archlitchi for approval. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Details Needs approval from an approver in each of these files:

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

@hami-robot hami-robot Bot added the size/XS label Jul 23, 2026
@coderabbitai

coderabbitai Bot commented Jul 23, 2026

Copy link
Copy Markdown

Review Change Stack

📝 Walkthrough

Walkthrough

set_task_pid() now tracks the CUDA device selected during NVML probing and uses that device for both primary-context retention and release instead of hard-coding device 0.

Changes

Probe context alignment

Layer / File(s) Summary
Use the mapped device for context lifetime
src/utils.c
set_task_pid() records the first successfully mapped CUDA device and uses it for primary-context retention and release.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

Suggested reviewers: mesutoezdil

Poem

I’m a rabbit tracing devices in flight,
No more probing the wrong one at night.
Retain where we scan,
Release by the same plan—
CUDA hops neatly from mismatch to right.

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly describes the main fix: creating the CUDA probe context on the correct device in set_task_pid().
Linked Issues check ✅ Passed The change matches issue #225 by using the mapped CUDA device for primary context retain/release, aligning probe context with NVML polling.
Out of Scope Changes check ✅ Passed The diff is narrowly scoped to the reported CUDA device mismatch in src/utils.c and introduces no unrelated changes.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai
coderabbitai Bot requested a review from mesutoezdil July 23, 2026 08:22

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 win

Release the retained context on every exit path.

cuDevicePrimaryCtxRetain() adds the current process to GPU-memory accounting for probeDev and sets ctx_activate[probeDev]. Returning after later failures bypasses cuDevicePrimaryCtxRelease(probeDev), leaving activation and memory-accounting state stuck until process exit. Route post-retain failures through a cleanup label that always releases probeDev.

🤖 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 win

Reject the no-mapped-device fallback instead of probing CUDA device 0.

If every nvml_to_cuda_map(i) returns -1, probeDev stays 0, 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. Initialize probeDev to an invalid value and return before cuDevicePrimaryCtxRetain when 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

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 9365ad91-48da-4a32-a413-42abc642e387

📥 Commits

Reviewing files that changed from the base of the PR and between 52f33fc and efcc7a2.

📒 Files selected for processing (1)
  • src/utils.c

Comment thread src/utils.c
Comment on lines +127 to +135
}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));

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 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.

Suggested change
}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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant