Skip to content

fix(device): enforce ResourceQuota in Fit() for all backends #2397

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Nitish08-08 wants to merge 1 commit into from
Closed

fix(device): enforce ResourceQuota in Fit() for all backends #2397

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions pkg/device/amd/device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -312,6 +312,11 @@ func (amddevice *AMDDevices) Fit(devices []*device.DeviceUsage, request device.C
if memReq <= 0 && dev.Totalmem > 0 {
memReq = dev.Totalmem
}
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
}
Comment on lines +315 to +319

@coderabbitai coderabbitai Bot Aug 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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 dev.Totalmem-dev.Usedmem < memReq {
reason[common.CardInsufficientMemory]++
klog.V(5).InfoS(common.CardInsufficientMemory, "pod", klog.KObj(pod), "device", dev.ID, "device total memory", dev.Totalmem, "device used memory", dev.Usedmem, "request memory", memReq)
Expand Down
5 changes: 5 additions & 0 deletions pkg/device/ascend/device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -512,6 +512,11 @@ func (npu *Devices) Fit(devices []*device.DeviceUsage, request device.ContainerD
//This incurs an issue
memreq = dev.Totalmem * k.MemPercentagereq / 100
}
if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(memreq), int64(k.Coresreq), k.Type, npu.GetResourceNames()) {
reason[common.ResourceQuotaNotFit]++
klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", memreq, "coresreq", k.Coresreq)
continue
}
if dev.Totalmem-dev.Usedmem < memreq {
reason[common.CardInsufficientMemory]++
klog.V(5).InfoS(common.CardInsufficientMemory, "pod", klog.KObj(pod), "device", dev.ID, "device index", i, "device total memory", dev.Totalmem, "device used memory", dev.Usedmem, "request memory", memreq)
Expand Down
6 changes: 6 additions & 0 deletions pkg/device/awsneuron/device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -428,6 +428,12 @@ func (neuron *AWSNeuronDevices) Fit(devices []*device.DeviceUsage, request devic
continue
}

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
}

Comment on lines +431 to +436

@coderabbitai coderabbitai Bot Aug 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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 countMaskAvailable(dev.Totalcore)-countMaskAvailable(dev.Usedcores) < k.Coresreq {
reason[common.CardInsufficientCore]++
klog.V(5).InfoS(common.CardInsufficientCore, "pod", klog.KObj(pod), "device", dev.ID, "device index", i, "device total core", dev.Totalcore, "device used core", dev.Usedcores, "request cores", k.Coresreq)
Expand Down
5 changes: 5 additions & 0 deletions pkg/device/biren/device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,11 @@ func (br *BirenDevices) Fit(devices []*device.DeviceUsage, request device.Contai
klog.V(5).InfoS(common.ExclusiveDeviceAllocateConflict, "pod", klog.KObj(pod), "device", dev.ID, "device index", i, "used", dev.Used)
continue
}
if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(k.Memreq), int64(k.Coresreq), BirenDevice, br.GetResourceNames()) {
reason[common.ResourceQuotaNotFit]++
klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", k.Memreq, "coresreq", k.Coresreq)
continue
}
if k.Nums > 0 {
klog.V(5).InfoS("find fit device", "pod", klog.KObj(pod), "device", dev.ID)
k.Nums--
Expand Down
5 changes: 5 additions & 0 deletions pkg/device/cambricon/device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -378,6 +378,11 @@ func (cam *CambriconDevices) Fit(devices []*device.DeviceUsage, request device.C
//This incurs an issue
memreq = dev.Totalmem * k.MemPercentagereq / 100
}
if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(memreq), int64(k.Coresreq), CambriconMLUDevice, cam.GetResourceNames()) {
reason[common.ResourceQuotaNotFit]++
klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", memreq, "coresreq", k.Coresreq)
continue
}
if dev.Totalmem-dev.Usedmem < memreq {
reason[common.CardInsufficientMemory]++
klog.V(5).InfoS(common.CardInsufficientMemory, "pod", klog.KObj(pod), "device", dev.ID, "device index", i, "device total memory", dev.Totalmem, "device used memory", dev.Usedmem, "request memory", memreq)
Expand Down
5 changes: 5 additions & 0 deletions pkg/device/enflame/device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -431,6 +431,11 @@ func (enf *EnflameDevices) Fit(devices []*device.DeviceUsage, request device.Con
klog.V(5).InfoS(common.ExclusiveDeviceAllocateConflict, "pod", klog.KObj(pod), "device", dev.ID, "device index", i, "used", dev.Used)
continue
}
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
}
Comment on lines +434 to +438

@coderabbitai coderabbitai Bot Aug 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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.go

Repository: 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 -S

Repository: 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 -S

Repository: 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 dev.Totalmem-dev.Usedmem < profileMemoryMiB {
reason[common.CardInsufficientMemory]++
klog.V(5).InfoS(common.CardInsufficientMemory, "pod", klog.KObj(pod), "device", dev.ID, "device index", i, "device total memory", dev.Totalmem, "device used memory", dev.Usedmem, "request memory", profileMemoryMiB)
Expand Down
6 changes: 6 additions & 0 deletions pkg/device/enflame/gcu.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,12 @@ func (gcuDev *GCUDevices) Fit(devices []*device.DeviceUsage, request device.Cont
continue
}

if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(k.Memreq), int64(k.Coresreq), EnflameGCUDevice, gcuDev.GetResourceNames()) {
reason[common.ResourceQuotaNotFit]++
klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", k.Memreq, "coresreq", k.Coresreq)
continue
}

if k.Nums > 0 {
klog.V(5).InfoS("find fit device", "pod", klog.KObj(pod), "device", dev.ID)
k.Nums--
Expand Down
5 changes: 5 additions & 0 deletions pkg/device/hygon/device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,11 @@ func (dcu *DCUDevices) Fit(devices []*device.DeviceUsage, request device.Contain
//This incurs an issue
memreq = dev.Totalmem * k.MemPercentagereq / 100
}
if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(memreq), int64(k.Coresreq), HygonDCUDevice, dcu.GetResourceNames()) {
reason[common.ResourceQuotaNotFit]++
klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", memreq, "coresreq", k.Coresreq)
continue
}
if dev.Totalmem-dev.Usedmem < memreq {
reason[common.CardInsufficientMemory]++
klog.V(5).InfoS(common.CardInsufficientMemory, "pod", klog.KObj(pod), "device", dev.ID, "device index", i, "device total memory", dev.Totalmem, "device used memory", dev.Usedmem, "request memory", memreq)
Expand Down
5 changes: 5 additions & 0 deletions pkg/device/iluvatar/device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -317,6 +317,11 @@ func (ilu *IluvatarDevices) Fit(devices []*device.DeviceUsage, request device.Co
//This incurs an issue
memreq = dev.Totalmem * k.MemPercentagereq / 100
}
if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(memreq), int64(k.Coresreq), k.Type, ilu.GetResourceNames()) {
reason[common.ResourceQuotaNotFit]++
klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", memreq, "coresreq", k.Coresreq)
continue
}
if dev.Totalmem-dev.Usedmem < memreq {
reason[common.CardInsufficientMemory]++
klog.V(5).InfoS(common.CardInsufficientMemory, "pod", klog.KObj(pod), "device", dev.ID, "device index", i, "device total memory", dev.Totalmem, "device used memory", dev.Usedmem, "request memory", memreq)
Expand Down
5 changes: 5 additions & 0 deletions pkg/device/kunlun/device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,6 +205,11 @@ func (kl *KunlunDevices) Fit(devices []*device.DeviceUsage, request device.Conta
}
}
}
if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(request.Memreq), int64(request.Coresreq), request.Type, kl.GetResourceNames()) {
reason[common.ResourceQuotaNotFit]++
klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", request.Memreq, "coresreq", request.Coresreq)
return false, tmpDevs, common.GenReason(reason, len(devices))
}
return true, tmpDevs, ""
}

Expand Down
5 changes: 5 additions & 0 deletions pkg/device/kunlun/vdevice.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -271,6 +271,11 @@ func (dev *KunlunVDevices) Fit(devices []*device.DeviceUsage, request device.Con
}
}
}
if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(request.Memreq), int64(request.Coresreq), request.Type, dev.GetResourceNames()) {
reason[common.ResourceQuotaNotFit]++
klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", request.Memreq, "coresreq", request.Coresreq)
return false, tmpDevs, common.GenReason(reason, len(devices))
}
return true, tmpDevs, ""
}

Expand Down
5 changes: 5 additions & 0 deletions pkg/device/metax/device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,11 @@ func (mat *MetaxDevices) Fit(devices []*device.DeviceUsage, request device.Conta
//This incurs an issue
memreq = dev.Totalmem * k.MemPercentagereq / 100
}
if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(memreq), int64(k.Coresreq), MetaxGPUDevice, mat.GetResourceNames()) {
reason[common.ResourceQuotaNotFit]++
klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", memreq, "coresreq", k.Coresreq)
continue
}
if dev.Totalmem-dev.Usedmem < memreq {
reason[common.CardInsufficientMemory]++
klog.V(5).InfoS(common.CardInsufficientMemory, "pod", klog.KObj(pod), "device", dev.ID, "device index", i, "device total memory", dev.Totalmem, "device used memory", dev.Usedmem, "request memory", memreq)
Expand Down
6 changes: 6 additions & 0 deletions pkg/device/metax/sdevice.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -371,6 +371,12 @@ func (mats *MetaxSDevices) Fit(devices []*device.DeviceUsage, request device.Con
memreq = dev.Totalmem * request.MemPercentagereq / 100
}

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
}

Comment on lines +374 to +379

@coderabbitai coderabbitai Bot Aug 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

if dev.Totalmem-dev.Usedmem < memreq {
reason[common.CardInsufficientMemory]++
klog.V(5).InfoS(common.CardInsufficientMemory, "pod", klog.KObj(pod), "device", dev.ID, "device index", i, "device total memory", dev.Totalmem, "device used memory", dev.Usedmem, "request memory", memreq)
Expand Down
5 changes: 5 additions & 0 deletions pkg/device/mthreads/device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -346,6 +346,11 @@ func (mth *MthreadsDevices) Fit(devices []*device.DeviceUsage, request device.Co
//This incurs an issue
memreq = dev.Totalmem * k.MemPercentagereq / 100
}
if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(memreq), int64(k.Coresreq), MthreadsGPUDevice, mth.GetResourceNames()) {
reason[common.ResourceQuotaNotFit]++
klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", memreq, "coresreq", k.Coresreq)
continue
}
if dev.Totalmem-dev.Usedmem < memreq {
reason[common.CardInsufficientMemory]++
klog.V(5).InfoS(common.CardInsufficientMemory, "pod", klog.KObj(pod), "device", dev.ID, "device index", i, "device total memory", dev.Totalmem, "device used memory", dev.Usedmem, "request memory", memreq)
Expand Down
23 changes: 1 addition & 22 deletions pkg/device/nvidia/device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -728,27 +728,6 @@ func (dev *NvidiaGPUDevices) AddResourceUsage(pod *corev1.Pod, n *device.DeviceU
return nil
}

func fitQuota(tmpDevs map[string]device.ContainerDevices, allocated *device.PodDevices, ns string, memreq int64, coresreq int64) bool {
mem := memreq
core := coresreq
for _, val := range tmpDevs[NvidiaGPUDevice] {
mem += int64(val.Usedmem)
core += int64(val.Usedcores)
}
if allocated != nil {
if podSingleDevice, exists := (*allocated)[NvidiaGPUDevice]; exists {
for _, containerDevices := range podSingleDevice {
for _, val := range containerDevices {
mem += int64(val.Usedmem)
core += int64(val.Usedcores)
}
}
}
}
klog.V(4).Infoln("Allocating...", mem, "cores", core)
return device.GetLocalCache().FitQuota(ns, mem, MemoryFactor, core, NvidiaGPUDevice)
}

func (nv *NvidiaGPUDevices) Fit(devices []*device.DeviceUsage, request device.ContainerDeviceRequest, pod *corev1.Pod, nodeInfo *device.NodeInfo, allocated *device.PodDevices) (bool, map[string]device.ContainerDevices, string) {
k := request
originReq := k.Nums
Expand Down Expand Up @@ -812,7 +791,7 @@ func (nv *NvidiaGPUDevices) Fit(devices []*device.DeviceUsage, request device.Co
//This incurs an issue
memreq = dev.Totalmem * k.MemPercentagereq / 100
}
if !fitQuota(tmpDevs, allocated, pod.Namespace, int64(memreq), int64(k.Coresreq)) {
if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(memreq), int64(k.Coresreq), NvidiaGPUDevice, nv.GetResourceNames()) {
reason[common.ResourceQuotaNotFit]++
klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", memreq, "coresreq", k.Coresreq)
continue
Expand Down
25 changes: 25 additions & 0 deletions pkg/device/quota.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,31 @@ func (q *QuotaManager) FitQuota(ns string, memreq int64, memoryFactor int32, cor
return true
}

// FitQuotaForDevice checks whether a pod's total device request (current round
// plus previously allocated containers) fits the namespace ResourceQuota for the
// given device type. Every backend should call this from its Fit() method to
// enforce quota at scheduling time.
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)
Comment on lines +95 to +113

@coderabbitai coderabbitai Bot Aug 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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\(' pkg

Repository: 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.go

Repository: 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.go

Repository: 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.
true

Repository: 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.

}

func countPodDevices(podDev PodDevices) map[string]int64 {
res := make(map[string]int64)
for deviceName, podSingle := range podDev {
Expand Down
78 changes: 78 additions & 0 deletions pkg/device/quota_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -282,3 +282,81 @@ func TestDelQuotaNonLimitsKey(t *testing.T) {
t.Errorf("DelQuota: expected memory limit 100 after deleting non-limits quota, got %d", (*qm.Quotas[ns])[memName].Limit)
}
}

func TestFitQuotaForDevice(t *testing.T) {
initTest()
qm := NewQuotaManager()
ns := "testns"
deviceName := "NVIDIA"
memName := "nvidia.com/gpumem"
coreName := "nvidia.com/gpucore"
rns := ResourceNames{
ResourceMemoryName: memName,
ResourceCoreName: coreName,
MemoryFactor: 1,
}

// Set up quota: limit 1000 mem, 100 cores
qm.Quotas[ns] = &DeviceQuota{
memName: &Quota{Used: 800, Limit: 1000},
coreName: &Quota{Used: 50, Limit: 100},
}

// Should fit: 800 used + 100 requested = 900 <= 1000
tmpDevs := map[string]ContainerDevices{}
if !FitQuotaForDevice(tmpDevs, nil, ns, 100, 10, deviceName, rns) {
t.Error("FitQuotaForDevice should return true when within limits")
}

// Should not fit: 800 used + 300 requested = 1100 > 1000
if FitQuotaForDevice(tmpDevs, nil, ns, 300, 10, deviceName, rns) {
t.Error("FitQuotaForDevice should return false when memory exceeds limit")
}

// Should not fit cores: 50 used + 60 requested = 110 > 100
if FitQuotaForDevice(tmpDevs, nil, ns, 10, 60, deviceName, rns) {
t.Error("FitQuotaForDevice should return false when cores exceed limit")
}

// Should account for devices in tmpDevs
tmpDevs2 := map[string]ContainerDevices{
deviceName: {
{Usedmem: 100, Usedcores: 20},
},
}
// 800 + 100 (tmpDevs) + 50 (request) = 950 <= 1000
if !FitQuotaForDevice(tmpDevs2, nil, ns, 50, 10, deviceName, rns) {
t.Error("FitQuotaForDevice should return true accounting for tmpDevs")
}
// 800 + 100 (tmpDevs) + 200 (request) = 1100 > 1000
if FitQuotaForDevice(tmpDevs2, nil, ns, 200, 10, deviceName, rns) {
t.Error("FitQuotaForDevice should return false accounting for tmpDevs")
}

// Should account for allocated devices
allocated := PodDevices{
deviceName: PodSingleDevice{
{
{Usedmem: 50, Usedcores: 5},
},
},
}
// 800 + 50 (allocated) + 50 (request) = 900 <= 1000
if !FitQuotaForDevice(tmpDevs, &allocated, ns, 50, 10, deviceName, rns) {
t.Error("FitQuotaForDevice should return true accounting for allocated")
}
// 800 + 50 (allocated) + 200 (request) = 1050 > 1000
if FitQuotaForDevice(tmpDevs, &allocated, ns, 200, 10, deviceName, rns) {
t.Error("FitQuotaForDevice should return false accounting for allocated")
}

// Should fit if namespace not present
if !FitQuotaForDevice(tmpDevs, nil, "otherns", 5000, 100, deviceName, rns) {
t.Error("FitQuotaForDevice should return true if namespace not present")
}

// Should fit if device not in quota
if !FitQuotaForDevice(tmpDevs, nil, ns, 5000, 100, "unknown-device", rns) {
t.Error("FitQuotaForDevice should return true if device not in quota")
}
}
5 changes: 5 additions & 0 deletions pkg/device/vastai/device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -281,6 +281,11 @@ func (va *VastaiDevices) Fit(devices []*device.DeviceUsage, request device.Conta
continue
}
}
if !device.FitQuotaForDevice(tmpDevs, allocated, pod.Namespace, int64(k.Memreq), int64(k.Coresreq), VastaiDevice, va.GetResourceNames()) {
reason[common.ResourceQuotaNotFit]++
klog.V(3).InfoS(common.ResourceQuotaNotFit, "pod", pod.Name, "memreq", k.Memreq, "coresreq", k.Coresreq)
continue
}
if k.Nums > 0 {
klog.V(5).InfoS("find fit device", "pod", klog.KObj(pod), "device", dev.ID)
if !dieMode {
Expand Down
Loading