numGPU returns a nondeterministic value when a container requests multiple GPU-like resources

[aarushitandon0]

What happened?

numGPU in pkg/runtime/runtime.go (master @ 9f13c8c) determines the GPU count by scanning the resource list for the first key whose lowercased name contains the substring gpu:

for resName, resQ := range resourcePerNode {
    if strings.Contains(strings.ToLower(resName.String()), "gpu") {
        return int(resQ.Value())
    }
}

corev1.ResourceList is a map, so iteration order is randomized. When a container requests two or more resources whose names contain gpu, the value returned varies between calls, with no error, no log line, and no event.

Two problems compound here. The result is unstable, and most of the colliding keys are not device counts at all. They are memory sizes, percentages, and core fractions, so the helper can return a quantity in completely different units.

Minimal reproduction (HAMi)

resources:
  limits:
    nvidia.com/gpu: 1
    nvidia.com/gpumem: 3000

GetNumGPUPerNode returns either 1 or 3000, unpredictably.

Measured against the real function at 9f13c8c, 1000 calls per case, on Go 1.26.0:

HAMi: gpu=1 + gpumem=3000                       returned 1      in  883/1000 runs
                                                returned 3000   in  117/1000 runs

Volcano: vgpu-number=1 + memory=3000 + cores=50 returned 1      in  747/1000 runs
                                                returned 50     in  119/1000 runs
                                                returned 3000   in  134/1000 runs

Control: nvidia.com/gpu=8 alone                 returned 8      in 1000/1000 runs

The odds are not even, and the reason matters. A map this size lives in a single group, Go places entries in insertion order, and iteration begins at a random slot and wraps. So for n matching keys the first-inserted wins (8-n+1)/8 of the time and each of the others wins 1/8. That predicts 87.5% / 12.5% for the two-key HAMi case and 75% / 12.5% / 12.5% for the three-key Volcano case, which is what the measurements show.

Which key is favoured therefore depends on insertion order, not on chance alone. Inserting gpumem before gpu inverts the result to 87% wrong:

gpu declared first      gpu(1): 87%   gpumem(3000): 12%
gpumem declared first   gpu(1): 12%   gpumem(3000): 87%

That is the awkward part. A cluster can look healthy indefinitely, and CI can stay green, while a user whose resources happen to be ordered the other way is wrong almost every reconcile. The single-key control case is fully deterministic, which is why no existing test can catch any of this.

This is standard usage, not an exotic setup

• Volcano vGPU: volcano.sh/vgpu-number, vgpu-memory, vgpu-cores. Three matching keys, so three different values are reachable from the same spec. Volcano is commonly deployed alongside Kubeflow Trainer as the gang scheduler, so this sits on a mainline path for those users.
• HAMi: nvidia.com/gpu, gpumem, gpumem-percentage, gpucores.
• Alibaba: aliyun.com/gpu-mem, gpu-count.
• Intel: every key under the gpu.intel.com/ domain matches, including i915, tiles, and millicores.

Nothing constrains the input either. ResourcesPerNode is a *corev1.ResourceRequirements carrying only +optional, with no CEL rules and no webhook validation on resource names, so users can supply any keys they like.

Affected call sites

GetNumGPUPerNode has five callers and no unit test coverage. Every existing GPU test literal happens to have exactly one substring-matching key per ResourceList, so no current test can observe the problem.

• Torch: selects numProcPerNode for torchrun.
• TorchTune: reached from the plugin's Validate, which means webhook validation and JobSet generation iterate the map independently and can disagree about the same TrainJob.
• MPI
• XGBoost: the value is plumbed into DMLC_NUM_WORKER. A wrong worker count would be expected to leave the rabit tracker waiting on peers that never arrive. I have confirmed the plumbing but not the tracker's blocking semantics.
• Flux

Why ordering or summing alone would not fix it

Ordering the map deterministically would return a stable but still wrong value, for example reading gpucores: 50 as 50 GPUs. Summing is worse: Volcano's documented vgpu-number: 1 plus vgpu-memory: 3000 would yield 3001. The substring heuristic itself is the part that needs replacing.

Worth noting that the sibling helper getNumCPUPerNode, in the same file and serving the same purpose, already does this correctly using typed exact-key accessors (res.Requests.Cpu()). The fuzzy match looks like a workaround for ResourceList having no accessor for extended resources rather than a deliberate design choice.

Proposed fix

Consult an ordered allowlist of resource names that genuinely represent a count of GPU devices, then fall back to a deterministic, filtered substring scan so that custom vendor names continue to resolve as they do today.

// gpuDeviceCountResources are the extended resources whose quantity is a count of
// GPU devices, in the order they are consulted. Resources that express GPU memory,
// core fractions, tiles or percentages are deliberately excluded: they match the
// "gpu" substring but are not device counts.
var gpuDeviceCountResources = []corev1.ResourceName{
	"nvidia.com/gpu",
	"nvidia.com/gpu.shared",
	"amd.com/gpu",
	"gpu.intel.com/i915",
	"gpu.intel.com/xe",
	"volcano.sh/vgpu-number",
	"aliyun.com/gpu-count",
}

// nonDeviceCountHints mark GPU-adjacent resources that quantify something other
// than a number of devices.
var nonDeviceCountHints = []string{"mem", "core", "percentage", "tile"}

func numGPU(resourcePerNode corev1.ResourceList) int {
	// 1. Known device-count resources win, in a fixed precedence order.
	for _, name := range gpuDeviceCountResources {
		if q, ok := resourcePerNode[name]; ok {
			return int(q.Value())
		}
	}
	// 2. Fall back to a deterministic, filtered scan so that custom vendor names
	//    (for example example.com/gpu) keep working.
	names := make([]string, 0, len(resourcePerNode))
	for name := range resourcePerNode {
		names = append(names, name.String())
	}
	sort.Strings(names)
	for _, name := range names {
		lower := strings.ToLower(name)
		if !strings.Contains(lower, "gpu") {
			continue
		}
		if slices.ContainsFunc(nonDeviceCountHints, func(h string) bool {
			return strings.Contains(lower, h)
		}) {
			continue
		}
		return int(resourcePerNode[corev1.ResourceName(name)].Value())
	}
	return 0
}

How this resolves the cases above:

Request	Result
nvidia.com/gpu: 1 + gpumem: 3000	1, via the allowlist
vgpu-number: 1 + vgpu-memory: 3000 + vgpu-cores: 50	1, via the allowlist
aliyun.com/gpu-mem + gpu-count: 2	2, via the allowlist
gpu.intel.com/i915: 1 + tiles: 2 + millicores: 500	1, via the allowlist
example.com/gpu: 4	4, via the fallback, unchanged from today
nvidia.com/gpumem: 3000 alone	0, since no device count was declared

Determinism is structural rather than incidental: the allowlist is a slice with fixed precedence, and the fallback sorts its keys, so neither path depends on map ordering. Heterogeneous vendors in a single container resolve by list order rather than by chance.

Tests would go in pkg/runtime/runtime_test.go, which currently has no TestGetNumGPUPerNode at all, using the existing table-driven style. Against the current implementation the multi-key cases do not fail reliably, they flake at the rates measured above, which is itself the symptom. With the fix they become deterministic and assert an exact count.

One deliberate exclusion: nvidia.com/mig-* resources contain no gpu substring, so MIG-only containers already report 0 today. That is a separate pre-existing false negative rather than part of this bug, and I would rather raise it on its own than widen the scope here.

I am happy to take this one and open the PR.

Related work

• feat(api): KEP-2782: Dynamic Resource Allocation support for TrainJob #3540 / KEP-2782 (DRA) describes this same contains "gpu" heuristic and propagates it to DRA device request names, but sums across all matches rather than taking the first. That would leave the same heuristic living on two surfaces with two different multi-match policies, one of them order dependent. Fixing the root cause before that lands seems better than fixing it twice.
• fix(framework): derive process counts from the resolved node resources #3974 touches all five call sites for a different bug (raw override versus merged resources). numGPU itself is untouched, so the two are complementary rather than overlapping. Note that it will shift call site line numbers.
• feat(framework): define deterministic execution order for runtime plugins #3957 is a sibling case of Go map iteration nondeterminism in the plugin registry.

The helper was introduced in #2653 and later adopted unchanged by #2832, #3188, and #3200.

What did you expect to happen?

GetNumGPUPerNode should return a stable, correct GPU device count regardless of map iteration order, and regardless of which additional GPU-adjacent resources the user requests alongside the device count.

In the HAMi reproduction above, it should return 1 on every call, never 3000.

Environment

Kubeflow Trainer: master @ 9f13c8c

Go: 1.26.0

This is a defect in the controller's resource parsing logic and is reproducible from the source at that commit, independent of cluster version.

[aarushitandon0]

/kind bug

[aarushitandon0]

/assign

[aarushitandon0]

/assign