feat: Add DRA-aware GPU accounting by MohanKumar21 · Pull Request #1378 · llm-d/llm-d-workload-variant-autoscaler

MohanKumar21 · 2026-07-04T21:44:24Z

Fixes #

Proposed Changes

🎁 Add DRA-aware GPU accounting: WVA now counts accelerators requested through Kubernetes Dynamic Resource Allocation (DRA) in addition to standard extended resources (e.g. nvidia.com/gpu).
- New internal/resources helpers resolve device counts from pod.spec.resourceClaims — direct ResourceClaim references, ResourceClaimTemplate references (including generated claims from pod.status.resourceClaimStatuses), exactCount allocation mode (with the Kubernetes default of 1 when count is omitted), and firstAvailable subrequests (max across alternatives). Allocated claims use status.allocation.devices.results, which also covers allocationMode: All.
- Container-level DRA extended resources (deviceclass.resource.kubernetes.io/<device-class>) are summed alongside vendor GPU resources.
🎁 Cluster GPU usage discovery (K8sWithGpuOperator.DiscoverUsage) now includes DRA device requests from running pods when computing per-accelerator usage.
🎁 The saturation engine (both BuildVariantStates and the v2 analysis path) resolves gpusPerReplica via a new DRA-aware helper (utils.GetDRAAwareGPUsPerReplica) that handles leader/worker (LWS) pod templates and falls back to the legacy pod-template count if DRA objects cannot be resolved.
🧹 Register resource.k8s.io/v1 types in the scheme and extend the manager ClusterRole with read access (get/list/watch) to deviceclasses, resourceclaims, resourceclaimtemplates, and resourceslices. Lookups degrade gracefully (count 0) on clusters without the DRA API (NoMatchError) or when claims are not found.

Pre-review Checklist

E2E tests for any new behavior — unit tests cover the DRA counting logic (internal/resources/resource_test.go, internal/discovery/k8s_with_gpu_operator_test.go); no e2e added since the kind-emulator environment does not provision a DRA driver.
Docs PR for any user-facing impact
Proposal PR for any new enhancement or change to existing behavior

Release Note

The workload-variant-autoscaler now accounts for GPUs requested via Kubernetes Dynamic Resource Allocation (DRA) — ResourceClaims, ResourceClaimTemplates, and deviceclass.resource.kubernetes.io/* extended resources — when computing per-replica GPU counts and cluster GPU usage. The manager ClusterRole gains read-only access to resource.k8s.io resources. Clusters without the DRA API are unaffected.

Docs

ev-shindin

Thanks @MohanKumar21! Nice feature please rebase PR and address review comments.

Signed-off-by: MohanKumar21! <mohanmrm20@gmail.com>

ev-shindin

Thanks @MohanKumar21 ! Please address two correctness gaps ans several provide missing test-cases.

ev-shindin · 2026-07-21T14:19:48Z

+	// GetLeaderPodTemplateSpec falls back to the worker template. Size counts the
+	// leader, so total = leaderGPUs + (Size-1)*workerGPUs is correct in both the
+	// leaderless (N*workerGPUs) and distinct-leader cases.
+	total := leaderGPUs + (int(groupSize)-1)*workerGPUs


The shared-claim dedup added to DiscoverUsage (the seenClaims map) is not applied on this capacity path. GetPodTemplateGPUs (called at L383 and L401) takes no seen map, so if a scale-target pod template references a literal ResourceClaimName (a single shared claim) rather than a ResourceClaimTemplateName, that one physical claim is counted in the leader term and again in (groupSize-1)*workerGPUs. Thread a shared seen map through the two GetPodTemplateGPUs calls, or explicitly treat ResourceClaimName in a scale-target template as unsupported/ignored and document why.

ev-shindin · 2026-07-21T14:21:19Z

+
+// CountDRADeviceRequests returns the count for DRA requests that declare an
+// exact count. The Kubernetes default for ExactCount with omitted count is one.
+func CountDRADeviceRequests(requests []resourcev1.DeviceRequest) int {


CountDRADeviceRequests (and the allocation-results path in getResourceClaimDeviceCount L169 / getResourceClaimTemplateDeviceCount L189) sum every device regardless of DeviceClassName. The standard path (GetContainersGPUs) filters against constants.VendorResources; here a claim mixing a GPU with a NIC/FPGA — or a non-GPU claim entirely — inflates the GPU total that drives scaling. Filter by a known GPU device-class allow-list, or document the "every DRA claim on a WVA-managed pod is GPU-only" assumption and why it holds.

ev-shindin · 2026-07-21T14:23:29Z

 				containerGPUs += int(qty.Value())
 			}
 		}
+		containerGPUs += resources.GetContainersDRAExtendedResources([]corev1.Container{container})


The regular-container equivalent (L279) is covered by TestGetPodGPURequests_DRAExtendedResource, but no test exercises an init container requesting a deviceclass.resource.kubernetes.io/* resource, leaving this line and its max(initMax, regularTotal) interaction uncovered.

ev-shindin · 2026-07-21T14:24:29Z

+	}
+
+	total := 0
+	for _, podClaim := range spec.ResourceClaims {


The loop accumulates total and returns the partial value with the error (L143/L153/L161); DiscoverUsage adds draCount even when err != nil. No test has a pod with >1 spec.resourceClaims where the first resolves and a later one fails. Add a two-claim pod (first → 2, second → Get fails) asserting 2, err != nil.

ev-shindin · 2026-07-21T14:26:12Z

+	assert.Equal(t, 2, result["NVIDIA-H100-SXM5-80GB"])
+}
+
+func TestDiscoverUsage_DRAErrorDegradesPerPod(t *testing.T) {


The pod has no non-DRA GPU request, so the result is 0 regardless of whether the DRA-error-degrades path runs. Give the pod a real nvidia.com/gpu request so the expected result is non-zero - the test then fails if a DRA error aborts the whole pod's accounting instead of just the DRA piece.

ev-shindin · 2026-07-21T14:27:09Z

+	}
+}
+
+func TestGetDRAAwareGPUsPerReplica(t *testing.T) {


All table cases use plain nvidia.com/gpu templates; only the fallback branch (...FallsBackOnError, L111) touches real DRA objects. Add a case where leader/worker templates reference a ResourceClaimTemplateName resolving in the fake client, e.g. leader 1 GPU + worker template 2 devices + size 3 -> expect 1 + 2*2 = 5.

ev-shindin · 2026-07-21T14:28:20Z

+// GetPodTemplateGPUs returns the accelerator count declared by a pod template,
+// including standard extended GPU resources and exact-count DRA claims.
+func GetPodTemplateGPUs(ctx context.Context, c client.Reader, namespace string, template *corev1.PodTemplateSpec) (int, error) {
+	if template == nil {


The nil guards in GetPodTemplateGPUs (L68), GetPodDRADeviceCount (L93), and GetPodSpecDRADeviceCount (L106) sit on a per-reconcile / per-pod hot path; a refactor that dereferences before the check would panic uncaught. Add three one-line tests asserting each returns 0, nil for nil input.

ev-shindin · 2026-07-21T14:29:42Z

+	require.True(t, apierrors.IsForbidden(err))
+}
+
+func TestCountDRADeviceRequests(t *testing.T) {


Every table case passes a single-element requests slice, so the total += accumulation across multiple DeviceRequests in one claim is unverified. Add a case with two requests (Exactly count=2 + Exactly count=1) asserting the sum (3).

ev-shindin requested changes Jul 19, 2026

View reviewed changes

ev-shindin assigned MohanKumar21 Jul 20, 2026

MohanKumar21 added 2 commits July 20, 2026 18:26

feat: Add DRA-aware GPU accounting

d39c3cc

Signed-off-by: MohanKumar21! <mohanmrm20@gmail.com>

Address review comments

f887463

Signed-off-by: MohanKumar21! <mohanmrm20@gmail.com>

MohanKumar21 force-pushed the ksmkumar/dra-support branch from 0740c2e to f887463 Compare July 20, 2026 13:32

ev-shindin reviewed Jul 21, 2026

View reviewed changes

Uh oh!

Conversation

MohanKumar21 commented Jul 4, 2026

Proposed Changes

Pre-review Checklist

Uh oh!

ev-shindin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ev-shindin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants