Fix GPU device index resolution and init/shutdown lifecycle races - #8
Open
vancraar wants to merge 4 commits into
Open
Fix GPU device index resolution and init/shutdown lifecycle races#8vancraar wants to merge 4 commits into
vancraar wants to merge 4 commits into
Conversation
gpu_amd_hardware_sampler used the HIP-relative device index directly as the ROCm SMI index for every rsmi_dev_* call. ROCm SMI's own device enumeration isn't filtered by HIP_VISIBLE_DEVICES/ROCR_VISIBLE_DEVICES the way HIP's is, so under a non-default visibility mask the sampler measured the wrong physical GPU. Fix: resolve the actual ROCm SMI index once at construction via a PCI-bus-ID match against every ROCm SMI device, throwing if no match is found instead of silently falling back to the HIP-relative index. The original HIP-relative index is kept in a new hip_device_id_ member for the one HIP call that still needs HIP-space. Also adds the missing find_dependency(MPI) to hwsConfig.cmake.in, alongside the existing ones for CUDA/HIP/ROCm-SMI/Level-Zero.
The constructor's rollback on a failed device resolution kept ROCm SMI's shared instances_/init_finished_ lifecycle keyed off is_first_instance. Under concurrent construction of several gpu_amd_hardware_sampler instances, that could shut ROCm SMI down while a sibling instance was still using it, or leave a sibling stuck forever waiting on init_finished_. Fixed by mirroring the destructor's own "last instance out" check (instances_ reaching zero, and only if ROCm SMI was actually initialized) instead, and by moving rsmi_init() itself into the try block so its own failure is handled the same way. resolve_rsmi_device_id() also silently used the first ROCm SMI device whose PCI bus ID matched. bdfid_to_pci_bus_id() intentionally drops the BDFID's partition bits, so on a partitioned MI-series accelerator several ROCm SMI entries can share one normalized bus ID - silently picking the first would reintroduce the exact silently-wrong-device failure mode this fix exists to close. Now throws on an ambiguous match instead. Also drops the device_id()/hip_device_id()/pci_bus_id() getters and enumerate_all_amd_gpu_pci_bus_ids(): unused on this branch, they only exist to serve the accel<->GPU correlation-hints feature that stays on feature/cray-pm-counters.
gpu_nvidia_hardware_sampler had the same bug just fixed for AMD: nvmlDeviceGetHandleByIndex() was called with the raw CUDA-relative device index. NVML's own device enumeration is not filtered by CUDA_VISIBLE_DEVICES the way CUDA's is, so under a non-default visibility mask the sampler silently measured the wrong physical GPU. Confirmed on real hardware (2x NVIDIA RTX 3090): masking CUDA_VISIBLE_DEVICES to the second GPU and constructing with CUDA-relative index 0 resolved to the first physical GPU before this fix, and correctly to the second after. Fix mirrors gpu_amd_hardware_sampler's resolve_rsmi_device_id(): resolve the true NVML index once at construction via a PCI-bus-ID match against every NVML device, throwing on zero or more than one match instead of guessing. NVIDIA MIG instances are not disambiguated by this - NVML's device enumeration only exposes physical parent GPUs, not MIG device handles, so a MIG-sliced CUDA device still resolves to (and samples) its whole parent GPU; documented as a known limitation rather than silently claimed as handled. The pre-existing instances_/init_finished_ busy-wait (shared by both the AMD and NVIDIA backends, unchanged by either fix) can still deadlock a concurrent waiter if the first instance's rsmi_init()/nvmlInit() call itself fails, or race a destructor's shutdown against a new constructor's init. Both are latent, pre-existing issues unreachable through hws's own API (system_hardware_sampler constructs sequentially) - tracked separately rather than fixed here to keep this change scoped to device index resolution.
gpu_amd_hardware_sampler, gpu_nvidia_hardware_sampler and
gpu_intel_hardware_sampler each used an atomic instances_ counter plus an
atomic init_finished_ flag to lazily call rsmi_init()/nvmlInit()/zeInit()
exactly once, with later constructors busy-waiting on the flag. This has
two bugs, both pre-existing (unrelated to and unchanged by the device
index resolution fixes elsewhere on this branch):
- if the first instance's init call itself throws, init_finished_ never
becomes true, and any other constructor already spinning on
`while (!init_finished_) {}` blocks forever - the failure is never
signaled to waiters.
- a destructor decrementing instances_ to zero and beginning
rsmi_shut_down()/nvmlShutdown() can race a new constructor concurrently
incrementing from zero and calling rsmi_init()/nvmlInit() again - the
atomics serialize the counter, not the actual init/shutdown calls.
Both are only reachable by constructing multiple samplers of the same
backend concurrently from different threads - system_hardware_sampler
itself always constructs sequentially - but they're part of the public
API surface (nothing stops a caller from doing this directly).
Fix: replace the atomic counter/flag pair with a single mutex that's held
for the whole "is this the first/last instance?" decision plus the actual
init/shutdown call, for all three backends. This makes the decision and
the call one atomic step, so:
- a failing init call leaves the shared count untouched and the mutex
unlocked - the next constructor to acquire it simply retries the init
call itself instead of ever spinning on a flag that might never be set.
- a shutdown and a concurrent init can no longer interleave, since both
hold the same mutex around their respective counter check and call.
gpu_intel_hardware_sampler has no equivalent shut down call (Level Zero
has none), so it only needed the simpler "has zeInit() ever succeeded"
half of this - a single mutex-guarded bool instead of a counter.
Verified with a 16-thread x 200-iteration concurrent construct/destruct
stress test against the NVIDIA backend on real hardware (2x RTX 3090):
0 failures, no hangs. The AMD and Intel changes are the identical pattern
with different underlying API calls, reviewed by inspection but not
independently stress-tested on real ROCm/Level-Zero hardware this round.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem 1: ROCm SMI device index resolution (AMD)
gpu_amd_hardware_sampler used the HIP-relative device index directly as
the ROCm SMI (rsmi_dev_*) index. ROCm SMI's own device enumeration is not
filtered by HIP_VISIBLE_DEVICES/ROCR_VISIBLE_DEVICES the way HIP's is, so
under a non-default visibility mask the sampler silently measured the
wrong physical GPU.
Fix: resolve the true ROCm SMI index once at construction via a
PCI-bus-ID match against every ROCm SMI device (resolve_rsmi_device_id()
in src/hws/gpu_amd/hardware_sampler.cpp), throwing on zero or more than
one match instead of guessing.
Problem 2: NVML device index resolution (NVIDIA)
Same bug, same fix, on the NVIDIA backend: gpu_nvidia_hardware_sampler
passed the CUDA-relative index directly to nvmlDeviceGetHandleByIndex(),
whose enumeration NVML doesn't filter by CUDA_VISIBLE_DEVICES either.
Fix mirrors the AMD one: resolve_nvml_device_id() in
src/hws/gpu_nvidia/hardware_sampler.cpp.
Problem 3: init/shutdown lifecycle race (all three GPU backends)
gpu_amd_hardware_sampler, gpu_nvidia_hardware_sampler and
gpu_intel_hardware_sampler each used an atomic counter plus an atomic
flag to lazily call rsmi_init()/nvmlInit()/zeInit() exactly once, with
later constructors busy-waiting on the flag. This has two bugs,
pre-existing and unrelated to problems 1/2 above:
becomes true, and any waiting constructor spins forever.
new constructor concurrently incrementing from zero and calling init
again - the atomics serialize the counter, not the actual calls.
Both are only reachable by constructing multiple samplers of the same
backend concurrently from different threads (system_hardware_sampler
itself always constructs sequentially), but they're part of the public
API surface.
Fix: replace the atomic counter/flag pair with a single mutex held for
the whole check-plus-call in both constructor and destructor, for all
three backends. gpu_intel_hardware_sampler has no shutdown call (Level
Zero has none), so it only needed the simpler "has zeInit() ever
succeeded" half of this.
Known limitations
Partitioned MI-series accelerators (SPX/CPX modes) and NVIDIA MIG
instances: both resolvers throw on an ambiguous PCI-bus-ID match rather
than guessing, but neither resolves a specific partition/instance. Not
exercised on real partitioned/MIG hardware.
Scope note
This was split out of feature/cray-pm-counters, where the AMD fix was
originally developed entangled with an accel<->GPU correlation-hints
feature. That branch still contains its own copy of the AMD fix; once
this PR merges, feature/cray-pm-counters will be rebased onto the new
develop to drop the duplicate before its own PR opens.