jax
6dba736e - [ROCm] xdist: select GPU from runner-allocated ROCR slice

Commit
124 days ago
[ROCm] xdist: select GPU from runner-allocated ROCR slice CI runners isolate each job to a subset of the host's physical GPUs by injecting ROCR_VISIBLE_DEVICES (e.g. "0,3,4,5"); these are absolute physical indices and may be non-contiguous. Re-setting ROCR_VISIBLE_DEVICES does not compose across processes, so the previous `xdist_worker_number % num_rocm_devices` selection could land a worker on a GPU owned by another job running on the same host. Pick the device from the inherited ROCR_VISIBLE_DEVICES slice instead, and fall back to the old behaviour when no slice is injected (local dev). (cherry picked from commit 639b5c0d24a25b219116fe38e32da8b396d3e36f)
Author
Committer
Parents
Loading