move build device status to gpu profile

Signed-off-by: Swati Gupta <swatig@nvidia.com>

add model in test assertion

Signed-off-by: Swati Gupta <swatig@nvidia.com>

Author: guptaNswati

cmd/dra-example-kubeletplugin/main.go

@@ -67,7 +67,7 @@ type Config struct {
 
 var validProfiles = map[string]func(flags Flags) profiles.Profile{
 	gpu.ProfileName: func(flags Flags) profiles.Profile {
-		return gpu.NewProfile(flags.nodeName, flags.numDevices, flags.gpuPartitions)
+		return gpu.NewProfile(flags.nodeName, flags.numDevices, flags.gpuPartitions, flags.gpuDeviceStatus)
 	},
 }
 

cmd/dra-example-kubeletplugin/state.go

@@ -18,7 +18,6 @@ package main
 
 import (
 	"context"
-	encode "encoding/json"
 	"fmt"
 	"path/filepath"
 	"slices"
@@ -161,7 +160,7 @@ func (s *DeviceState) Prepare(ctx context.Context, claim *resourceapi.ResourceCl
 	if err != nil {
 		return nil, fmt.Errorf("unable to sync from checkpoint: %v", err)
 	}
-	restoredDevices, err := s.restoreClaimFromCheckpoint(ctx, checkpoint, claim)
+	restoredDevices, err := s.restoreClaimFromCheckpoint(checkpoint, claim)
 	if err != nil {
 		return nil, fmt.Errorf("unable to restore from checkpoint: %v", err)
 	}
@@ -217,7 +216,42 @@ func (s *DeviceState) Unprepare(claimUID types.UID) error {
 // prepareDevices performs one-time setup for the devices allocated to a
 // ResourceClaim before being consumed by a Pod.
 func (s *DeviceState) prepareDevices(ctx context.Context, claim *resourceapi.ResourceClaim) (PreparedDevices, error) {
-	return s.computeDeviceConfig(ctx, claim)
+	preparedDevices, err := s.computeDeviceConfig(claim)
+	if err != nil {
+		return nil, err
+	}
+
+	// Publish per-device status (e.g. uuid, model, driverVersion) into
+	// ResourceClaim.status.devices[].data when the profile implements
+	// [profiles.DeviceStatusBuilder]. This is a side-effect on the API server
+	// and therefore lives in prepareDevices (rather than computeDeviceConfig,
+	// which must be deterministic and side-effect free).
+	builder, ok := s.configHandler.(profiles.DeviceStatusBuilder)
+	if !ok {
+		return preparedDevices, nil
+	}
+
+	var deviceStatuses []resourceapi.AllocatedDeviceStatus
+	for _, result := range claim.Status.Allocation.Devices.Results {
+		if result.Driver != s.driverName {
+			continue
+		}
+		if status := builder.BuildDeviceStatus(s.allocatable, &result); status != nil {
+			deviceStatuses = append(deviceStatuses, *status)
+		}
+	}
+	if len(deviceStatuses) > 0 {
+		klog.FromContext(ctx).Info("Publishing device status to ResourceClaim",
+			"namespace", claim.Namespace, "name", claim.Name, "devices", len(deviceStatuses))
+		if err := s.updateDeviceStatus(ctx, claim.Namespace, claim.Name, deviceStatuses...); err != nil {
+			// A failure to publish status is non-fatal: the device is still
+			// prepared and the claim status will simply be missing the data.
+			klog.FromContext(ctx).Error(err, "Failed to update device status on ResourceClaim",
+				"namespace", claim.Namespace, "name", claim.Name)
+		}
+	}
+
+	return preparedDevices, nil
 }
 
 // unprepareDevices undoes any side-effects produced by
@@ -232,7 +266,7 @@ func (s *DeviceState) unprepareDevices(claimUID types.UID, checkpoint *checkpoin
 // should be deterministic and produce no side-effects. Non-deterministic data or
 // side-effects should be produced by [DeviceState.prepareDevices] directly and
 // recorded in the checkpoint by [DeviceState.addClaimToCheckpoint].
-func (s *DeviceState) computeDeviceConfig(ctx context.Context, claim *resourceapi.ResourceClaim) (PreparedDevices, error) {
+func (s *DeviceState) computeDeviceConfig(claim *resourceapi.ResourceClaim) (PreparedDevices, error) {
 	if claim.Status.Allocation == nil {
 		return nil, fmt.Errorf("claim not yet allocated")
 	}
@@ -254,9 +288,6 @@ func (s *DeviceState) computeDeviceConfig(ctx context.Context, claim *resourceap
 	// the list with len(Requests) == 0 for the lookup below.
 	configs = slices.Insert(configs, 0, &OpaqueDeviceConfig{})
 
-	// build device status
-	var devicesStatus []resourceapi.AllocatedDeviceStatus
-
 	// Look through the configs and figure out which one will be applied to
 	// each device allocation result based on their order of precedence.
 	configResultsMap := make(map[runtime.Object][]*resourceapi.DeviceRequestAllocationResult)
@@ -269,11 +300,6 @@ func (s *DeviceState) computeDeviceConfig(ctx context.Context, claim *resourceap
 			return nil, fmt.Errorf("requested device is not allocatable: %v", result.Device)
 		}
 
-		if s.gpuDeviceStatus {
-			deviceStatus := s.buildDeviceStatus(result)
-			devicesStatus = append(devicesStatus, deviceStatus)
-		}
-
 		for _, c := range slices.Backward(configs) {
 			if len(c.Requests) == 0 || slices.Contains(c.Requests, result.Request) {
 				configResultsMap[c.Config] = append(configResultsMap[c.Config], &result)
@@ -287,13 +313,6 @@ func (s *DeviceState) computeDeviceConfig(ctx context.Context, claim *resourceap
 	// of device allocation results.
 	perDeviceCDIContainerEdits := make(profiles.PerDeviceCDIContainerEdits)
 	for config, results := range configResultsMap {
-		if s.gpuDeviceStatus {
-			klog.Infof("Adding device attribute to claim %s/%s", claim.Namespace, claim.Name)
-			if err := s.updateDeviceStatus(ctx, claim.Namespace, claim.Name, devicesStatus...); err != nil {
-				klog.Warningf("Failed to update device attributes for claim %s/%s: %v", claim.Namespace, claim.Name, err)
-			}
-		}
-
 		// Apply the config to the list of results associated with it.
 		containerEdits, err := s.configHandler.ApplyConfig(config, results)
 		if err != nil {
@@ -344,12 +363,12 @@ func (*DeviceState) removeClaimFromCheckpoint(checkpoint *checkpointapi.Checkpoi
 
 // restoreClaimFromCheckpoint returns the device definitions for devices already prepared
 // for the given claim. If the claim has not yet been prepared, it returns nil.
-func (s *DeviceState) restoreClaimFromCheckpoint(ctx context.Context, checkpoint *checkpointapi.Checkpoint, claim *resourceapi.ResourceClaim) (PreparedDevices, error) {
+func (s *DeviceState) restoreClaimFromCheckpoint(checkpoint *checkpointapi.Checkpoint, claim *resourceapi.ResourceClaim) (PreparedDevices, error) {
 	if slices.ContainsFunc(checkpoint.PreparedClaims, func(c checkpointapi.PreparedClaim) bool { return c.UID == claim.UID }) {
 		// If [DeviceState.addClaimToCheckpoint] associated any other data with
 		// the claim in the checkpoint, then that should be added to the
 		// returned [PreparedDevices] here.
-		return s.computeDeviceConfig(ctx, claim)
+		return s.computeDeviceConfig(claim)
 	}
 	return nil, nil
 }
@@ -453,43 +472,6 @@ func GetOpaqueDeviceConfigs(
 	return resultConfigs, nil
 }
 
-func (s *DeviceState) buildDeviceStatus(res resourceapi.DeviceRequestAllocationResult) resourceapi.AllocatedDeviceStatus {
-	dn := res.Device
-	deviceInfo := make(map[string]resourceapi.DeviceAttribute)
-
-	if d, ok := s.allocatable[dn]; ok {
-		if uuid, ok := d.Attributes["uuid"]; ok {
-			deviceInfo["uuid"] = uuid
-		}
-		if model, ok := d.Attributes["model"]; ok {
-			deviceInfo["model"] = model
-		}
-		if driverVersion, ok := d.Attributes["driverVersion"]; ok {
-			deviceInfo["driverVersion"] = driverVersion
-		}
-	}
-
-	jsonBytes, err := encode.Marshal(deviceInfo)
-	if err != nil {
-		klog.Errorf("Failed to marshal device data: %v", err)
-		jsonBytes = []byte("{}")
-	}
-
-	return resourceapi.AllocatedDeviceStatus{
-		Device: dn,
-		Driver: res.Driver,
-		Pool:   res.Pool,
-		// Data records per-allocation metadata used for monitoring and debugging:
-		//   - Pod→GPU mapping: makes it easier to see which GPU a given pod is using,
-		//     which is not readily available elsewhere.
-		//   - Device attributes (e.g. UUID, model, driverVersion): remain available
-		//     even if the device is later removed from a ResourceSlice (for example,
-		//     because it becomes unhealthy), so past allocations can still be
-		//     correlated with later health or scheduling issues.
-		Data: &runtime.RawExtension{Raw: jsonBytes},
-	}
-}
-
 func (s *DeviceState) updateDeviceStatus(ctx context.Context, ns, name string, devices ...resourceapi.AllocatedDeviceStatus) error {
 	// Converting wrapper to use latest API types,
 	// converts to/from server-supported version.

deployments/helm/dra-example-driver/templates/clusterrole.yaml

@@ -11,6 +11,16 @@ rules:
 - apiGroups: ["resource.k8s.io"]
   resources: ["resourceclaims/status"]
   verbs: ["update"]
+# Kubernetes 1.36+ enforces granular authorization for ResourceClaim status
+# writes via the DRAResourceClaimGranularStatusAuthorization feature gate. As a
+# node-local driver, we need the "associated-node" verbs on
+# "resourceclaims/driver" for our own driver name. This rule is inert on older
+# clusters that have not yet enabled the feature gate.
+# See https://github.com/kubernetes/kubernetes/issues/138149
+- apiGroups: ["resource.k8s.io"]
+  resources: ["resourceclaims/driver"]
+  verbs: ["associated-node:update", "associated-node:patch"]
+  resourceNames: ["{{ include "dra-example-driver.driverName" . }}"]
 - apiGroups: [""]
   resources: ["nodes"]
   verbs: ["get"]

internal/profiles/gpu/gpu.go

@@ -17,6 +17,7 @@
 package gpu
 
 import (
+	"encoding/json"
 	"fmt"
 	"maps"
 	"math/rand"
@@ -27,6 +28,7 @@ import (
 	"k8s.io/apimachinery/pkg/api/resource"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/dynamic-resource-allocation/resourceslice"
+	"k8s.io/klog/v2"
 	"k8s.io/utils/ptr"
 	cdiapi "tags.cncf.io/container-device-interface/pkg/cdi"
 	cdispec "tags.cncf.io/container-device-interface/specs-go"
@@ -38,16 +40,18 @@ import (
 const ProfileName = "gpu"
 
 type Profile struct {
-	nodeName         string
-	numGPUs          int
-	partitionsPerGPU int
+	nodeName           string
+	numGPUs            int
+	partitionsPerGPU   int
+	enableDeviceStatus bool
 }
 
-func NewProfile(nodeName string, numGPUs int, partitionsPerGPU int) Profile {
+func NewProfile(nodeName string, numGPUs int, partitionsPerGPU int, enableDeviceStatus bool) Profile {
 	return Profile{
-		nodeName:         nodeName,
-		numGPUs:          numGPUs,
-		partitionsPerGPU: partitionsPerGPU,
+		nodeName:           nodeName,
+		numGPUs:            numGPUs,
+		partitionsPerGPU:   partitionsPerGPU,
+		enableDeviceStatus: enableDeviceStatus,
 	}
 }
 
@@ -290,3 +294,42 @@ func applyGpuConfig(config *configapi.GpuConfig, results []*resourceapi.DeviceRe
 
 	return perDeviceEdits, nil
 }
+
+// BuildDeviceStatus implements [profiles.DeviceStatusBuilder]. It returns an
+// [resourceapi.AllocatedDeviceStatus] populated with a subset of the device's
+// attributes (uuid, model, driverVersion) to publish into
+// ResourceClaim.status.devices[].data.
+func (p Profile) BuildDeviceStatus(allocatable map[string]resourceapi.Device, result *resourceapi.DeviceRequestAllocationResult) *resourceapi.AllocatedDeviceStatus {
+	if !p.enableDeviceStatus {
+		return nil
+	}
+
+	deviceInfo := make(map[string]resourceapi.DeviceAttribute)
+	if d, ok := allocatable[result.Device]; ok {
+		for _, name := range []resourceapi.QualifiedName{"uuid", "model", "driverVersion"} {
+			if v, ok := d.Attributes[name]; ok {
+				deviceInfo[string(name)] = v
+			}
+		}
+	}
+
+	jsonBytes, err := json.Marshal(deviceInfo)
+	if err != nil {
+		klog.Errorf("Failed to marshal device data for %s: %v", result.Device, err)
+		jsonBytes = []byte("{}")
+	}
+
+	// Data records per-allocation metadata used for monitoring and debugging:
+	//   - Pod->GPU mapping: makes it easier to see which GPU a given pod is
+	//     using, which is not readily available elsewhere.
+	//   - Device attributes (e.g. UUID, model, driverVersion): remain available
+	//     even if the device is later removed from a ResourceSlice (for
+	//     example, because it becomes unhealthy), so past allocations can
+	//     still be correlated with later health or scheduling issues.
+	return &resourceapi.AllocatedDeviceStatus{
+		Device: result.Device,
+		Driver: result.Driver,
+		Pool:   result.Pool,
+		Data:   &runtime.RawExtension{Raw: jsonBytes},
+	}
+}