Skip to content

[ROCm] xdist: select GPU from runner-allocated ROCR slice#793

Open
mminutoli wants to merge 1 commit into
rocm-jaxlib-v0.9.1from
fix/conftest-r0.9.1
Open

[ROCm] xdist: select GPU from runner-allocated ROCR slice#793
mminutoli wants to merge 1 commit into
rocm-jaxlib-v0.9.1from
fix/conftest-r0.9.1

Conversation

@mminutoli

Copy link
Copy Markdown

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

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

@psanal35 psanal35 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.

LGTM~

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants