fix(device): enforce ResourceQuota in Fit() for all backends - #2397
fix(device): enforce ResourceQuota in Fit() for all backends#2397Nitish08-08 wants to merge 1 commit into
Conversation
FitQuota was only called from nvidia's Fit(), leaving a race window for every other backend. A burst of pods created back-to-back all pass admission with the same Used value and all get scheduled, exceeding the namespace quota. Move the accumulate-then-check logic from nvidia's private fitQuota helper to a shared FitQuotaForDevice function in pkg/device/quota.go. Every backend now calls this from its Fit() method, closing the race. Closes Project-HAMi#2363
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: Nitish08-08 The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
|
Thanks for your pull request. Before we can look at it, you'll need to add a 'DCO signoff' to your commits. 📝 Please follow instructions in the contributing guide to update your commits with the DCO Full details of the Developer Certificate of Origin can be found at developercertificate.org. The list of commits missing DCO signoff:
DetailsInstructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. I understand the commands that are listed here. |
📝 WalkthroughWalkthroughChangesDevice quota enforcement
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (1)
pkg/device/quota_test.go (1)
293-297: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a
MemoryFactortest case.All new cases use
MemoryFactor: 1. Add a case with a factor greater than one. This verifies thatFitQuotaForDeviceforwards scaled memory usage correctly.Based on PR context,
MemoryFactordefines the conversion between recorded memory usage and the ResourceQuota limit.🤖 Prompt for 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. In `@pkg/device/quota_test.go` around lines 293 - 297, Add a test case in the quota test cases around FitQuotaForDevice using a MemoryFactor greater than one, with expected quota values reflecting the converted memory usage. Keep the existing MemoryFactor: 1 cases unchanged and verify that FitQuotaForDevice forwards scaled memory usage correctly.
🤖 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 `@pkg/device/amd/device.go`:
- Around line 315-319: Calculate the converted core usage as coreReq before the
FitQuotaForDevice call in the AMD device allocation flow, using the same
conversion later used for ContainerDevice.Usedcores. Pass coreReq instead of
k.Coresreq to FitQuotaForDevice, and reuse that value when persisting Usedcores
to keep quota accounting consistent.
In `@pkg/device/awsneuron/device.go`:
- Around line 431-436: Update the k.Nums > 1 branch to validate quota before
returning, using the complete selected allocation that includes both temporary
and previously allocated devices. Reuse device.FitQuotaForDevice with the same
resource requirements and AWSNeuronDevice parameters as the single-device path,
and preserve the existing quota failure handling.
In `@pkg/device/enflame/device.go`:
- Around line 434-438: Update the Enflame resource-name construction used by
FitQuotaForDevice so GetResourceNames() includes the MemoryFactor from the
Enflame configuration before the quota check. Preserve the existing quota
arguments and add a regression test verifying GB-configured memory quotas
correctly accept a matching 4-GB profile after conversion.
In `@pkg/device/metax/sdevice.go`:
- Around line 374-379: Update MetaxSDevices.Fit around FitQuotaForDevice to
construct the selected ContainerDevices allocation before quota validation
instead of referencing undeclared tmpDevs. Aggregate the exact selected devices
and validate the memory and core values that will be persisted, using zero core
usage for Online allocations; pass this completed allocation to
FitQuotaForDevice before finalizing the device list.
In `@pkg/device/quota.go`:
- Around line 95-113: Make FitQuotaForDevice perform the quota fit check and
usage reservation atomically instead of only calling GetLocalCache().FitQuota.
Reuse the QuotaManager synchronization and reservation path so concurrent
requests cannot all pass against unchanged usage; the successful check must
record mem and core usage before returning, while rejected requests leave usage
unchanged.
---
Nitpick comments:
In `@pkg/device/quota_test.go`:
- Around line 293-297: Add a test case in the quota test cases around
FitQuotaForDevice using a MemoryFactor greater than one, with expected quota
values reflecting the converted memory usage. Keep the existing MemoryFactor: 1
cases unchanged and verify that FitQuotaForDevice forwards scaled memory usage
correctly.
🪄 Autofix
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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 3fab7099-b7b2-4026-b2b7-da794b249f5a
📒 Files selected for processing (18)
pkg/device/amd/device.gopkg/device/ascend/device.gopkg/device/awsneuron/device.gopkg/device/biren/device.gopkg/device/cambricon/device.gopkg/device/enflame/device.gopkg/device/enflame/gcu.gopkg/device/hygon/device.gopkg/device/iluvatar/device.gopkg/device/kunlun/device.gopkg/device/kunlun/vdevice.gopkg/device/metax/device.gopkg/device/metax/sdevice.gopkg/device/mthreads/device.gopkg/device/nvidia/device.gopkg/device/quota.gopkg/device/quota_test.gopkg/device/vastai/device.go
| if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(memReq), int64(k.Coresreq), AMDDevice, amddevice.GetResourceNames()) { | ||
| reason[common.ResourceQuotaNotFit]++ | ||
| klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", memReq, "coresreq", k.Coresreq) | ||
| continue | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Pass the converted core usage to the quota check.
Line 315 passes k.Coresreq, but Lines 326-358 convert that percentage to coreReq and persist coreReq in ContainerDevice.Usedcores. The quota check can undercount the current allocation when dev.Totalcore is not 100.
Calculate coreReq before this check. Pass coreReq to FitQuotaForDevice.
Based on PR context, quota usage is accumulated from ContainerDevice.Usedcores.
🤖 Prompt for 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.
In `@pkg/device/amd/device.go` around lines 315 - 319, Calculate the converted
core usage as coreReq before the FitQuotaForDevice call in the AMD device
allocation flow, using the same conversion later used for
ContainerDevice.Usedcores. Pass coreReq instead of k.Coresreq to
FitQuotaForDevice, and reuse that value when persisting Usedcores to keep quota
accounting consistent.
| if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(k.Memreq), int64(k.Coresreq), AWSNeuronDevice, neuron.GetResourceNames()) { | ||
| reason[common.ResourceQuotaNotFit]++ | ||
| klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", k.Memreq, "coresreq", k.Coresreq) | ||
| continue | ||
| } | ||
|
|
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Check quota before the multi-device return.
The k.Nums > 1 branch at Lines 369-392 returns before Line 431. Multi-device AWS Neuron requests therefore bypass the new quota validation.
Validate the complete selected allocation before the return. Include temporary and previously allocated devices in that validation.
🤖 Prompt for 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.
In `@pkg/device/awsneuron/device.go` around lines 431 - 436, Update the k.Nums > 1
branch to validate quota before returning, using the complete selected
allocation that includes both temporary and previously allocated devices. Reuse
device.FitQuotaForDevice with the same resource requirements and AWSNeuronDevice
parameters as the single-device path, and preserve the existing quota failure
handling.
| if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(profileMemoryMiB), int64(profileCorePercent), EnflameVGCUDevice, enf.GetResourceNames()) { | ||
| reason[common.ResourceQuotaNotFit]++ | ||
| klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", profileMemoryMiB, "coresreq", profileCorePercent) | ||
| continue | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 5 \
'MemoryFactor|profileMemoryMiB|EnflameResourceNameGCUMemory|FitQuotaForDevice' \
pkg/device/enflame/device.go pkg/device/quota.go pkg/device/quota_test.goRepository: Project-HAMi/HAMi
Length of output: 15058
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '--- quota implementation ---\n'
sed -n '1,150p' pkg/device/quota.go
printf '\n--- Enflame device structs and resource names ---\n'
sed -n '70,110p' pkg/device/enflame/device.go
sed -n '478,488p' pkg/device/enflame/device.go
printf '\n--- FitQuota signatures/usages ---\n'
rg -n 'func \([^)]*\) FitQuota|func FitQuota|MemoryFactor|ResourceNames' pkg/device -SRepository: Project-HAMi/HAMi
Length of output: 246
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- quota implementation ---'
sed -n '1,150p' pkg/device/quota.go
echo
echo '--- Enflame device structs and resource names ---'
sed -n '70,110p' pkg/device/enflame/device.go
sed -n '478,488p' pkg/device/enflame/device.go
echo
echo '--- FitQuota signatures/usages ---'
rg -n 'func \([^)]*\) FitQuota|func FitQuota|MemoryFactor|ResourceNames' pkg/device -SRepository: Project-HAMi/HAMi
Length of output: 17587
Set the Enflame memory conversion factor before enforcing this quota check.
GetResourceNames() returns a zero MemoryFactor, so FitQuotaForDevice() passes usage in MiB into FitQuota() without scaling the configured quota limit. When users configure the Enflame memory quota in GB, a 4-GB profile is checked as 4096 MiB and can exceed a limit like 4. Add MemoryFactor from the Enflame config to GetResourceNames() and add a regression test.
🤖 Prompt for 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.
In `@pkg/device/enflame/device.go` around lines 434 - 438, Update the Enflame
resource-name construction used by FitQuotaForDevice so GetResourceNames()
includes the MemoryFactor from the Enflame configuration before the quota check.
Preserve the existing quota arguments and add a regression test verifying
GB-configured memory quotas correctly accept a matching 4-GB profile after
conversion.
Source: MCP tools
| if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(memreq), int64(request.Coresreq), MetaxSGPUDevice, mats.GetResourceNames()) { | ||
| reason[common.ResourceQuotaNotFit]++ | ||
| klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", memreq, "coresreq", request.Coresreq) | ||
| continue | ||
| } | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | 🏗️ Heavy lift
Build the selected allocation before quota validation.
tmpDevs is not declared in MetaxSDevices.Fit, so this package does not compile.
Do not fix this by adding an empty map. This check runs before the final device list is selected, so an empty map cannot aggregate a multi-device allocation. Build the selected ContainerDevices first, then validate the exact memory and core values that will be persisted. Use zero core usage for Online allocations, which later store coreReq = 0.
🤖 Prompt for 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.
In `@pkg/device/metax/sdevice.go` around lines 374 - 379, Update MetaxSDevices.Fit
around FitQuotaForDevice to construct the selected ContainerDevices allocation
before quota validation instead of referencing undeclared tmpDevs. Aggregate the
exact selected devices and validate the memory and core values that will be
persisted, using zero core usage for Online allocations; pass this completed
allocation to FitQuotaForDevice before finalizing the device list.
| func FitQuotaForDevice(tmpDevs map[string]ContainerDevices, allocated *PodDevices, ns string, memreq int64, coresreq int64, deviceName string, resourceNames ResourceNames) bool { | ||
| mem := memreq | ||
| core := coresreq | ||
| for _, val := range tmpDevs[deviceName] { | ||
| mem += int64(val.Usedmem) | ||
| core += int64(val.Usedcores) | ||
| } | ||
| if allocated != nil { | ||
| if podSingleDevice, exists := (*allocated)[deviceName]; exists { | ||
| for _, containerDevices := range podSingleDevice { | ||
| for _, val := range containerDevices { | ||
| mem += int64(val.Usedmem) | ||
| core += int64(val.Usedcores) | ||
| } | ||
| } | ||
| } | ||
| } | ||
| klog.V(4).Infoln("FitQuotaForDevice: device", deviceName, "mem", mem, "cores", core) | ||
| return GetLocalCache().FitQuota(ns, mem, resourceNames.MemoryFactor, core, deviceName) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline pkg/device/quota.go --items all
rg -n -C 5 '\b(FitQuotaForDevice|FitQuota|AddUsage|RmUsage)\b' pkg
rg -n -C 5 '\.Fit\(|FitQuotaForDevice\(' pkgRepository: Project-HAMi/HAMi
Length of output: 50373
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect quota implementation and relevant scheduler paths.
sed -n '50,190p' pkg/device/quota.go
printf '\n--- scheduler relevant sections ---\n'
sed -n '130,175p' pkg/scheduler/scheduler.go
sed -n '910,1005p' pkg/scheduler/scheduler.go
printf '\n--- webhook quota path ---\n'
sed -n '110,170p' pkg/scheduler/webhook.go
printf '\n--- all direct calls to device quota functions outside tests ---\n'
rg -n -C 3 '\.FitQuotaForDevice\(|FitQuotaForDevice\(|GetLocalCache\(\)\.FitQuota\(' `git ls-files '*_test.go' | sed 's/.*//;q'` 2>/dev/null >/tmp/except_tests.txt || true
rg -n -C 3 '\.FitQuotaForDevice\(|FitQuotaForDevice\(|device\.GetLocalCache\(\)\.FitQuota\(' --glob '!**/*_test.go' pkg
printf '\n--- search for quota related methods and lock patterns ---\n'
rg -n -C 3 'FitQuota|QuotaManager|mutex\.RLock|mutex\.Lock|AddUsage|TakeAndDelete' pkg --glob '!**/*_test.go'Repository: Project-HAMi/HAMi
Length of output: 10814
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Behavioral probe of the QuotaManager state update semantics with an RLock-only check.
# This model matches: FitQuota reads Quota.Used under RLock; AddUsage records usage under Lock;
# multiple Fits can complete before the winning reservation is written.
cat > /tmp/quota_concurrency_model.go <<'EOF'
package main
import (
"context"
"fmt"
"sync"
"time"
)
type Quota struct {
Limit int64
Used int64
}
type QuotaManager struct {
Quotas map[string]Quota
mu sync.RWMutex
}
func fit(q *QuotaManager, ns string, request int64) bool {
q.mu.RLock()
v := q.Quotas[ns]
accept := v.Used+request <= v.Limit
q.mu.RUnlock()
if !accept {
return false
}
time.Sleep(time.Microsecond)
return true
}
func addUsage(q *QuotaManager, ns string, usage int64) {
q.mu.Lock()
q.Quotas[ns].Used += usage
q.mu.Unlock()
}
func main() {
m := &QuotaManager{Quotas: map[string]Quota{"ns": {Limit: 1000}}}
var done sync.WaitGroup
ctx := context.Background()
for i := 0; i < 5; i++ {
done.Add(1)
go func(id int) {
defer done.Done()
select {
case <-time.After(time.Microsecond):
case <-ctx.Done():
return
}
if fit(m, "ns", 300) {
println("fit", id)
addUsage(m, "ns", 300)
}
}(i)
}
done.Wait()
fmt.Printf("final_used=%d\n", m.Quotas["ns"].Used)
}
EOF
go run /tmp/quota_concurrency_model.goRepository: Project-HAMi/HAMi
Length of output: 273
🏁 Script executed:
#!/bin/bash
set -euo pipefail
cat > /tmp/quota_concurrency_model.go <<'EOF'
package main
import (
"fmt"
"sync"
"sync/atomic"
"time"
)
type Quota struct {
Limit int64
Used int64
}
type QuotaManager struct {
Quotas map[string]*Quota
mu sync.RWMutex
}
func fit(q *QuotaManager, ns string, request int64) bool {
q.mu.RLock()
v := q.Quotas[ns]
accept := v.Used+request <= v.Limit
q.mu.RUnlock()
if !accept {
return false
}
time.Sleep(time.Microsecond)
return true
}
func addUsage(q *QuotaManager, ns string, usage int64) {
q.mu.Lock()
defer q.mu.Unlock()
if q.Quotas[ns] == nil {
q.Quotas[ns] = &Quota{Limit: 1000}
}
q.Quotas[ns].Used += usage
}
func main() {
m := &QuotaManager{Quotas: map[string]*Quota{"ns": {Limit: 1000}}}
var attempts, successes atomic.Int64
var wg sync.WaitGroup
for i := 0; i < 50; i++ {
wg.Add(1)
go func() {
defer wg.Done()
if fit(m, "ns", 300) {
values := make(chan bool, 1)
go func() { values <- addUsage(m, "ns", 300) }()
time.Sleep(time.Microsecond * 20)
select {
case <-values:
successes.Add(1)
default:
m.mu.RLock()
used := m.Quotas["ns"].Used
m.mu.RUnlock()
if used == 300 {
successes.Add(1)
}
}
attempts.Add(1)
}
}()
}
wg.Wait()
fmt.Printf("attempts=%d successes=%d final_used=%d\n", attempts.Load(), successes.Load(), m.Quotas["ns"].Used)
}
EOF
go run /tmp/quota_concurrency_model.goRepository: Project-HAMi/HAMi
Length of output: 267
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Remove the probe from output; it is not needed for a deterministic code-base result because
# the source shows FitQuota only reads while AddUsage later mutates separately.
trueRepository: Project-HAMi/HAMi
Length of output: 155
Make quota admission atomic.
FitQuotaForDevice only calls QuotaManager.FitQuota, which reads quota usage under RLock. The usage is added later by QuotaManager.AddUsage under Lock. Concurrent Fit checks can all pass while Used is unchanged, then each increments the same Used value. Merge the fit check with usage reservation, or serialize both steps so one winner records the new usage and later requests see the updated quota.
🤖 Prompt for 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.
In `@pkg/device/quota.go` around lines 95 - 113, Make FitQuotaForDevice perform
the quota fit check and usage reservation atomically instead of only calling
GetLocalCache().FitQuota. Reuse the QuotaManager synchronization and reservation
path so concurrent requests cannot all pass against unchanged usage; the
successful check must record mem and core usage before returning, while rejected
requests leave usage unchanged.
|
There's overlap here with #2377, which is open against the same issue and covers ascend, cambricon, hygon, iluvatar and mthreads. Flagging that early so we don't both spend time on the same five. One concrete thing worth checking in your ascend change. CodeRabbit raised this on #2377 and it applies to the same insertion point here: in topology mode #2377 handles it by skipping the per-card check when Your coverage of amd, awsneuron, biren, enflame, kunlun, metax sgpu and vastai goes further than #2377 does — I left those out because their |
3 participants
Problem
Namespace ResourceQuota enforcement has a race window for every backend
except NVIDIA.
The admission webhook records no device usage — that happens at Filter
time via QuotaManager.AddUsage. So when a burst of pods is created
back-to-back, they all reach the webhook before any has been scheduled,
all read the same Used value, and all pass.
NVIDIA catches this because its Fit() method calls a private fitQuota
helper that re-checks quota after each device assignment. The second
pod in the burst sees the first pod's recorded usage and is correctly
denied when over limit.
Every other backend — cambricon, ascend, hygon, iluvatar, mthreads,
metax, amd, enflame, kunlun, awsneuron, vastai, biren — has no such
check. The whole burst schedules and the namespace ends up over its
limit.
This was identified during review of #2347, which fixed the case where
non-NVIDIA quotas were ignored outright. This is the remaining half.
Reproducer
limits.cambricon.com/mlu.smlu.vmemory sized for ~2 pods.
Solution
Extract the accumulate-then-check logic from nvidia's private fitQuota
helper into a shared FitQuotaForDevice function that any backend can
call. Every backend's Fit() method now invokes this shared function
before assigning a device, closing the race.
The shared function:
and from previously allocated containers in the pod (allocated)
Files changed
How to verify
Closes #2363
Summary by CodeRabbit