fix(scheduler): count shared DRA GPU device once per node - #1931
Merged
enoodle merged 2 commits intoAug 6, 2026
Conversation
Contributor
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds reference-counted NVIDIA DRA GPU accounting so a physical device shared by multiple pods contributes to node usage once, with symmetric removal behavior and regression tests. ChangesShared DRA GPU accounting
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
TensorRaya
force-pushed
the
fix/dra-shared-claim-double-count
branch
from
July 16, 2026 09:17
0565d00 to
c79f723
Compare
Contributor
Author
|
@coderabbitai review |
Contributor
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@CHANGELOG.md`:
- Line 28: Correct the changelog entry’s spelling by changing “ine” to “in” and
“then” to “than”; also make the issue link’s displayed number match the
referenced issue URL, using the correct issue number consistently.
In `@pkg/scheduler/api/node_info/dra_shared_device_info.go`:
- Around line 58-100: Update dedupSharedDRAGpus to read the GPU value before
incrementing DRASharedDeviceRefCount, return immediately when current GPU usage
is zero or less, and cap alreadyCounted so the deduction cannot exceed current.
This prevents reservation tasks from registering DRA devices or producing
negative resourcesToTrack values while preserving normal deduplication for
GPU-contributing tasks.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: fa8f5d74-ec03-455d-81cf-f12eb06a355b
📒 Files selected for processing (6)
CHANGELOG.mdpkg/common/constants/constants.gopkg/scheduler/api/node_info/dra_shared_device_info.gopkg/scheduler/api/node_info/dra_shared_device_info_test.gopkg/scheduler/api/node_info/gpu_sharing_node_info.gopkg/scheduler/api/node_info/node_info.go
TensorRaya
force-pushed
the
fix/dra-shared-claim-double-count
branch
2 times, most recently
from
July 16, 2026 10:16
c18422b to
b7767da
Compare
enoodle
reviewed
Jul 17, 2026
enoodle
left a comment
enoodle
left a comment
Collaborator
There was a problem hiding this comment.
I created 2 allocate action tests to verify this PR. one of them fails because of the mismatch between the test GPU driver name and the one tested here (see comment) But I suspect that the one where one of the pods sharing the DRA claims is still pending will fail after your PR is fixed. Can you take a look at this after you fix the gpu driver detection ?
The tests for reference: main...enoodle:KAI-Scheduler:review-1931-test-cases
TensorRaya
force-pushed
the
fix/dra-shared-claim-double-count
branch
from
July 17, 2026 23:37
b7767da to
00e3d38
Compare
Contributor
Author
|
Thanks for writing the test cases, that was a big help. I cherry-picked your commit into the branch (kept your authorship) and addressed everything:
|
|
Total coverage: 54.3% -> 54.0% (delta -0.30%) Merging this branch will increase overall coverage
Coverage by fileChanged files (no unit tests)
Please note that the "Total", "Covered", and "Missed" counts above refer to code statements instead of lines of code. The value in brackets refers to the test coverage of that file in the old version of the code. Changed unit test files
|
📊 Performance Benchmark ResultsComparing PR (
|
Collaborator
|
Hi @TensorRaya sorry for the delay, |
When one physical GPU is shared by several pods through a single ResourceClaim (status.reservedFor lists more than one consumer, as with MPS or time-slicing), every consuming pod carries the same allocated device. The node accounting added that device to the used vector once per pod, so a 1-GPU node shared by two pods reported used: 2, capacity: 1. IdleVector then went negative and MaxNodeResourcesPredicate rejected the node for every task, even ones requesting no GPU at all, which showed up as a confusing "didn't have enough resources: GPUs" on a node that clearly had one. Keep a per-device reference count on the node, keyed by driver/pool/device from the allocation result, and skip a device that another pod on the node already contributed. Removal mirrors this: the device stays counted until the last consumer leaves. Exclusive claims and non-DRA GPUs are untouched. Reservation pods are a special case worth calling out: addTaskResources zeroes their GPU index before the dedup runs, so counting their shared devices would subtract from zero and drive the used count negative. Both the dedup and release paths now leave a zero-GPU task's accounting alone. Signed-off-by: Raya Solano <raya@mbinf.de>
Signed-off-by: Erez Freiberger <enoodle@gmail.com> Signed-off-by: Raya Solano <raya.solano@mbinf.de>
TensorRaya
force-pushed
the
fix/dra-shared-claim-double-count
branch
from
July 29, 2026 12:40
00e3d38 to
4e9cf1e
Compare
Contributor
Author
|
I did it — rebased on current main, PR is mergeable again. The conflict was in Since #1943 touches the same accounting path, I re-ran the counter-proof instead of trusting a green suite: with Gates run locally against main: |
enoodle
approved these changes
Jul 30, 2026
Collaborator
📝 Changelog fragment recordedThanks! This PR added the changelog fragment(s) below. Pending fragments are folded into |
Collaborator
|
Backport failed for Please cherry-pick the changes locally and resolve any conflicts. git fetch origin v0.16
git worktree add -d .worktree/backport-1931-to-v0.16 origin/v0.16
cd .worktree/backport-1931-to-v0.16
git switch --create backport-1931-to-v0.16
git cherry-pick -x 8dc2871286de9c1ba4b0511a6b018a594ded9895 4e9cf1e3cbe1916955f50857bfffc18c55e53fed |
Collaborator
|
Successfully created backport PR for |
4 tasks
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
4 participants
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.
Description
addTaskResourcesadds each task's DRA GPU count — taken from its allocatedResourceClaimInfo— to the nodeUsedVector. When several pods share onephysical GPU through a single
ResourceClaim(itsstatus.reservedForlistsmore than one consumer, as in DRA time-slicing or MPS), each consuming pod
carries the same allocated device, so the device is counted once per pod. On a
node whose shared device pushes the counted usage above the node's physical GPU
count,
IdleVectorgoes negative and the node is rejected for every task —including GPU-less ones — via
MaxNodeResourcesPredicate, surfacing as amisleading "didn't have enough resources: GPUs" message (e.g. on a 1-GPU node
shared by two pods:
used: 2, capacity: 1, idle-1).Fix: track a per-device reference count on the node, keyed by the allocated
device's
driver/pool/devicetuple.addTaskResourcessubtracts from thetracked GPU vector any device already referenced by another pod on the node, so
a shared physical device contributes to
UsedVectorexactly once;removeTaskResourcesmirrors this, keeping the device counted while any consumerremains and releasing it only when the last one leaves. The reference count
lives on
GpuSharingNodeInfo(cloned with the rest of the node's sharing state).Only NVIDIA GPU DRA devices (
driver == gpu.nvidia.com) are tracked; other DRAdevices are untouched.
Non-DRA and single-consumer DRA GPU accounting is unchanged: the dedup only ever
skips a device that a prior pod on the same node already contributed.
Tests (
pkg/scheduler/api/node_info/dra_shared_device_info_test.go), exercisingthe real
AddTask/RemoveTaskaccounting path:1, idle0(not2/-1).1while any consumer of the shared deviceremains, returns to
0when the last leaves.2),guarding against over-dedup.
Verified the shared-claim tests fail without the fix (reproducing
used: 2,idle
-1) and pass with it; the distinct-devices test stays green either way.Full
./pkg/scheduler/...suite passes.Related Issues
Fixes #1930
Checklist
Breaking Changes
None. Accounting changes only in the previously-broken case of a physical DRA
device shared by more than one pod on a node; exclusive claims and non-DRA GPUs
are counted exactly as before.
Additional Notes
CHANGELOG.md updated under
[Unreleased] > Fixed.Validated with:
go test ./pkg/scheduler/api/node_info/ ./pkg/common/resources/and the fullgo test ./pkg/scheduler/...suite (77 packages ok).gofmtandgo vetcleanon all touched files.
Summary by CodeRabbit