Fix host PID detection when CUDA_VISIBLE_DEVICES doesn't start at device 0

[ihnokim]

Fixes #225 (closed by the lifecycle policy while the bug is still present on main).

Changes?

• set_task_pid() now pairs the probe context and the polled NVML device by UUID: the probe context is created on CUDA device 0, and the NVML device whose UUID matches (cuDeviceGetUuid vs nvmlDeviceGetUUID) is used for both process snapshots.
• If no NVML device matches the UUID, it logs a warning and falls back to polling NVML device 0 (the previous behavior).
• Error paths now release the retained probe context; previously it leaked for the lifetime of the process whenever detection failed (the leak observed during the fix: probe CUDA context on correct device in set_task_pid() #230 testing).
• Not in scope: the nvml_preInit() map reset itself (see below). This PR only makes set_task_pid() independent of the map; other consumers of the map are unchanged.

Why We Need?

• Whenever CUDA_VISIBLE_DEVICES does not start with the container's device 0, host-PID detection fails deterministically (host pid is error!), per-process accounting is merged into the wrong slot, and legal allocations trip spurious per-device OOMs. This is the failure analyzed in set_task_pid: deterministic "host pid is error" when CUDA_VISIBLE_DEVICES does not start with the pod's device 0 (probe context and poll target mismatch), leading to spurious Device 0 OOM #225.
• Direct cause: the probe context is created with cuDevicePrimaryCtxRetain(&pctx, 0), which is CUDA device 0, while the before/after NVML snapshots are taken on the first NVML device that has a valid entry in cuda_to_nvml_map. With such a CUDA_VISIBLE_DEVICES, those are different physical GPUs, so the probe PID never appears in the snapshot diff.
• Root cause of why the map cannot be used here: nvml_preInit() resets cuda_to_nvml_map_array to identity, and it is triggered via pthread_once by the nvmlInit() call at the top of set_task_pid() itself. By the time the polling loop consults the map, it is an identity map regardless of what parse_cuda_visible_env() computed earlier in postInit().
• This is also why the approach in fix: probe CUDA context on correct device in set_task_pid() #230 (keep the probe on the cudaDev found through the map) did not resolve the failure when we tested it: with an identity map that cudaDev is always 0, which is behaviorally identical to the hardcoded value. It is consistent with the current processes num = 0 0 runs reported in the fix: probe CUDA context on correct device in set_task_pid() #230 thread with that patch applied. The unresolved SET_TASK_PID FAILED case on driver 580.65.06 from that thread matches this root cause as well, but we could not verify on that exact driver.
• The UUID pairing holds regardless of the state of the map or of the initialization order.

Tests?

Environment: Kubernetes + HAMi v2.7.1 (hami-scheduler), pod with nvidia.com/gpu: 2, nvidia.com/gpumem: 5800; node: 2x NVIDIA A100-PCIE-40GB, driver 580.173.02. Cases A-F are the reproduction matrix from #225, 3 repetitions each.

• Full path (torch 2.8.0+cu128), v2.7.1 code base:

build	A 0	B 1	C 1,0	D 1,0 + set_device(1)	E 0,1 + set_device(1)	F 0,1
v2.7.1 official binary	ok 3/3	host pid error 3/3	host pid error 3/3	host pid error 3/3	ok 3/3	ok 3/3
v2.7.1 built from source (control)	ok 3/3	host pid error 3/3	host pid error 3/3	host pid error 3/3	ok 3/3	ok 3/3
v2.7.1 + this fix	ok 3/3	ok 3/3	ok 3/3	ok 3/3	ok 3/3	ok 3/3

• Driver-API probe (ctypes cuInit, no CUDA runtime involved), main:

build	A	B	C	D	E	F
main unpatched	ok 3/3	host pid error 3/3	host pid error 3/3	host pid error 3/3	ok 3/3	ok 3/3
main + this fix	ok 3/3	ok 3/3	ok 3/3	ok 3/3	ok 3/3	ok 3/3

• Case C verbose with the fix: current processes num = 1 2 -> hostPid=581730 (v2.7.1 build) / hostPid=1150676 (main build)
• Leak probe (holder process with CVD=1,0, 15 s): no lingering probe context during or after the run
• A clean clone of this branch builds in nvidia/cuda:13.3.0-cudnn-devel-ubi8 (the make build-in-docker image)
• Why two validation paths: on our nodes, libvgpu built from current main (patched or not) fails CUDA runtime entry-point negotiation for every torch build we tried (cudart 12.6 / 12.8 / 13.0, all ending in Found no NVIDIA driver), while the official v2.7.1 binary works and direct driver-API calls also work. That regression is unrelated to this change (it reproduces on unpatched main) and we will report it separately. Because of it, the full-workload validation uses a v2.7.1 backport of this patch, and main is validated through the driver-API probe path. The probe path was cross-checked against the official v2.7.1 binary, where it reproduces the same B/C/D failures as the torch matrix.

[hami-robot]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by: ihnokim
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:

• OWNERS

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

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 029d269f-e1db-4e51-acf8-419ce70ea63e

📥 Commits

Reviewing files that changed from the base of the PR and between cc30e87 and 1b3d92e.

📒 Files selected for processing (1)

• src/utils.c

Included review availability: Your plan provides up to 2 included reviews per hour; 0 remain after this review.

---

📝 Walkthrough

Walkthrough

set_task_pid() now matches the CUDA probe device to its NVML device by UUID, polls the resolved device, falls back to NVML device 0 when needed, and releases the retained primary context through a shared cleanup path.

Changes

Probe Device Alignment

Layer / File(s)	Summary
Resolve and poll the matching device src/utils.c	The code matches CUDA and NVML device UUIDs, falls back to NVML device 0 when no match exists, and uses the resolved device for process snapshots.
Unify context cleanup src/utils.c	Error and success paths preserve the result and release the retained primary context through one cleanup path. Obsolete unified lock state and helpers are removed.

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

Merge Risk: ⚪ Minimal · up to 1b3d9

The change is localized and no actionable merge-blocking risk remains after normal checks and review.

Suggested labels: enhancement

Suggested reviewers: archlitchi, chaunceyjiang, officialgamer91023

Poem

A rabbit matched each UUID bright
NVML polls the proper site
The probe context follows too
Cleanup runs on errors through
Device zero waits as fallback light

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Out of Scope Changes check	⚠️ Warning	The device-matching and context-cleanup changes are in scope for issue #225. However, removing the unified lock state and the exported try_lock_unified_lock() and try_unlock_unified_lock() functions i…	Remove the unrelated unified-lock changes from this pull request, or document their dependency on the host-PID fix and provide evidence that removing these public functions is safe.
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 1 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The changes match issue #225. set_task_pid() now matches the NVML polling device to CUDA device 0 by UUID, uses the matched device for snapshots and context retention, provides a fallback, and release…
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main bug fix: correcting host PID detection when CUDA_VISIBLE_DEVICES does not begin with device 0.

Full details: Linked Issues check

Explanation

The changes match issue #225. set_task_pid() now matches the NVML polling device to CUDA device 0 by UUID, uses the matched device for snapshots and context retention, provides a fallback, and releases retained contexts on cleanup paths. These changes address the host-PID failure and incorrect device accounting.

Full details: Out of Scope Changes check

Explanation

The device-matching and context-cleanup changes are in scope for issue #225. However, removing the unified lock state and the exported try_lock_unified_lock() and try_unlock_unified_lock() functions is not covered by the linked issue and may be unrelated scope.

• Fix all pre-merge checks with AI

✨ 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

• X
• Mastodon
• Reddit
• LinkedIn

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

[mesutoezdil]

This is being closed because it does not comply with the contribution guidelines.