add device data in resourceclaim

cmd/dra-example-kubeletplugin/driver.go

@@ -94,9 +94,8 @@ func (d *driver) PrepareResourceClaims(ctx context.Context, claims []*resourceap
 	return result, nil
 }
 
-func (d *driver) prepareResourceClaim(_ context.Context, claim *resourceapi.ResourceClaim) kubeletplugin.PrepareResult {
-	klog.Infof("Preparing claim: UID=%s, Namespace=%s, Name=%s", claim.UID, claim.Namespace, claim.Name)
-	preparedPBs, err := d.state.Prepare(claim)
+func (d *driver) prepareResourceClaim(ctx context.Context, claim *resourceapi.ResourceClaim) kubeletplugin.PrepareResult {
+	preparedPBs, err := d.state.Prepare(ctx, claim)
 	if err != nil {
 		klog.Errorf("Error preparing devices for claim %v: %v", claim.UID, err)
 		return kubeletplugin.PrepareResult{

cmd/dra-example-kubeletplugin/main.go

@@ -52,6 +52,7 @@ type Flags struct {
 	healthcheckPort               int
 	profile                       string
 	driverName                    string
+	gpuDeviceStatus               bool
 }
 
 type Config struct {
@@ -147,6 +148,12 @@ func newApp() *cli.App {
 			Destination: &flags.driverName,
 			EnvVars:     []string{"DRIVER_NAME"},
 		},
+		&cli.BoolFlag{
+			Name:        "gpu-device-status",
+			Usage:       "Enable adding allocated device attributes (e.g., model, uuid, driverVersion) into ResourceClaim.status.devices[].data. Disabled by default.",
+			Destination: &flags.gpuDeviceStatus,
+			EnvVars:     []string{"GPU_DEVICE_STATUS"},
+		},
 	}
 	cliFlags = append(cliFlags, flags.kubeClientConfig.Flags()...)
 	cliFlags = append(cliFlags, flags.loggingConfig.Flags()...)

cmd/dra-example-kubeletplugin/state.go

@@ -17,14 +17,21 @@
 package main
 
 import (
+	"context"
+	encode "encoding/json"
 	"fmt"
 	"slices"
 	"sync"
 
 	resourceapi "k8s.io/api/resource/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/serializer/json"
+	"k8s.io/client-go/util/retry"
+	draclient "k8s.io/dynamic-resource-allocation/client"
 	"k8s.io/dynamic-resource-allocation/resourceslice"
+
+	"k8s.io/klog/v2"
 	drapbv1 "k8s.io/kubelet/pkg/apis/dra/v1beta1"
 	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager"
 
@@ -48,6 +55,8 @@ type DeviceState struct {
 	checkpointManager checkpointmanager.CheckpointManager
 	configDecoder     runtime.Decoder
 	configHandler     profiles.ConfigHandler
+	config            *Config
+	gpuDeviceStatus   bool
 }
 
 func NewDeviceState(config *Config) (*DeviceState, error) {
@@ -103,6 +112,8 @@ func NewDeviceState(config *Config) (*DeviceState, error) {
 		checkpointManager: checkpointManager,
 		configDecoder:     decoder,
 		configHandler:     configHandler,
+		config:            config,
+		gpuDeviceStatus:   config.flags.gpuDeviceStatus,
 	}
 
 	checkpoints, err := state.checkpointManager.ListCheckpoints()
@@ -124,7 +135,7 @@ func NewDeviceState(config *Config) (*DeviceState, error) {
 	return state, nil
 }
 
-func (s *DeviceState) Prepare(claim *resourceapi.ResourceClaim) ([]*drapbv1.Device, error) {
+func (s *DeviceState) Prepare(ctx context.Context, claim *resourceapi.ResourceClaim) ([]*drapbv1.Device, error) {
 	s.Lock()
 	defer s.Unlock()
 
@@ -139,7 +150,8 @@ func (s *DeviceState) Prepare(claim *resourceapi.ResourceClaim) ([]*drapbv1.Devi
 	if preparedClaims[claimUID] != nil {
 		return preparedClaims[claimUID].GetDevices(), nil
 	}
-	preparedDevices, err := s.prepareDevices(claim)
+
+	preparedDevices, err := s.prepareDevices(ctx, claim)
 	if err != nil {
 		return nil, fmt.Errorf("prepare failed: %v", err)
 	}
@@ -190,7 +202,7 @@ func (s *DeviceState) Unprepare(claimUID string) error {
 	return nil
 }
 
-func (s *DeviceState) prepareDevices(claim *resourceapi.ResourceClaim) (profiles.PreparedDevices, error) {
+func (s *DeviceState) prepareDevices(ctx context.Context, claim *resourceapi.ResourceClaim) (profiles.PreparedDevices, error) {
 	if claim.Status.Allocation == nil {
 		return nil, fmt.Errorf("claim not yet allocated")
 	}
@@ -212,6 +224,9 @@ func (s *DeviceState) prepareDevices(claim *resourceapi.ResourceClaim) (profiles
 	// 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)
@@ -223,6 +238,12 @@ func (s *DeviceState) prepareDevices(claim *resourceapi.ResourceClaim) (profiles
 		if _, exists := s.allocatable[result.Device]; !exists {
 			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)
@@ -236,6 +257,13 @@ func (s *DeviceState) prepareDevices(claim *resourceapi.ResourceClaim) (profiles
 	// 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 {
@@ -351,3 +379,61 @@ 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.
+	c := draclient.New(s.config.coreclient)
+	rc := c.ResourceClaims(ns)
+
+	return retry.RetryOnConflict(retry.DefaultRetry, func() error {
+		claim, err := rc.Get(ctx, name, metav1.GetOptions{})
+		if err != nil {
+			return err
+		}
+
+		// copy the object and update only status.devices
+		claim = claim.DeepCopy()
+		claim.Status.Devices = devices
+
+		_, err = rc.UpdateStatus(ctx, claim, metav1.UpdateOptions{})
+		return err
+	})
+}

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

@@ -8,6 +8,9 @@ rules:
 - apiGroups: ["resource.k8s.io"]
   resources: ["resourceclaims"]
   verbs: ["get"]
+- apiGroups: ["resource.k8s.io"]
+  resources: ["resourceclaims/status"]
+  verbs: ["update"]
 - apiGroups: [""]
   resources: ["nodes"]
   verbs: ["get"]

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

@@ -80,6 +80,8 @@ spec:
         # Simulated number of devices the example driver will pretend to have.
         - name: NUM_DEVICES
           value: {{ .Values.kubeletPlugin.numDevices | quote }}
+        - name: GPU_DEVICE_STATUS
+          value: {{ .Values.gpuDeviceStatus | quote }}
         {{- if .Values.kubeletPlugin.containers.plugin.healthcheckPort }}
         - name: HEALTHCHECK_PORT
           value: {{ .Values.kubeletPlugin.containers.plugin.healthcheckPort | quote }}

deployments/helm/dra-example-driver/values.yaml

@@ -18,6 +18,10 @@ deviceProfile: "gpu"
 # value is derived from the deviceProfile.
 driverName: ""
 
+# add allocated device attributes (e.g. model, uuid, driverVersion)
+# into ResourceClaim.status.devices[].data via server-side apply.
+gpuDeviceStatus: false
+
 imagePullSecrets: []
 image:
   repository: registry.k8s.io/dra-example-driver/dra-example-driver

test/e2e/e2e.sh

@@ -86,6 +86,49 @@ function gpu-partition-count-from-logs {
   echo "$logs" | sed -nE "s/^declare -x GPU_DEVICE_${id}_PARTITION_COUNT=\"(.+)\"$/\1/p"
 }
 
+function verify-resourceclaim-device-status() {
+  local ns="$1"
+  echo "=== Verifying ResourceClaim device data in namespace ${ns} ==="
+
+  local claim=""
+  for i in {1..30}; do
+    claim="$(kubectl get resourceclaim -n "${ns}" -o jsonpath='{.items[0].metadata.name}' 2>/dev/null || true)"
+    if [[ -n "${claim}" ]]; then
+      break
+    fi
+    sleep 1
+  done
+
+  if [[ -z "${claim}" ]]; then
+    echo "ERROR: no ResourceClaim found in namespace ${ns}"
+    exit 1
+  fi
+
+  echo "Found ResourceClaim ${ns}/${claim}, checking status.devices[0].data ..."
+
+  local uuid
+  uuid="$(kubectl get resourceclaim "${claim}" -n "${ns}" \
+    -o jsonpath='{.status.devices[0].data.uuid.string}')"
+
+  local driver_version
+  driver_version="$(kubectl get resourceclaim "${claim}" -n "${ns}" \
+    -o jsonpath='{.status.devices[0].data.driverVersion.version}')"
+
+  if [[ -z "${uuid}" ]]; then
+    echo "ERROR: ResourceClaim ${ns}/${claim} is missing .status.devices[0].data.uuid.string"
+    kubectl get resourceclaim "${claim}" -n "${ns}" -o yaml
+    exit 1
+  fi
+
+  if [[ -z "${driver_version}" ]]; then
+    echo "ERROR: ResourceClaim ${ns}/${claim} is missing .status.devices[0].data.driverVersion.version"
+    kubectl get resourceclaim "${claim}" -n "${ns}" -o yaml
+    exit 1
+  fi
+
+  echo "OK: ResourceClaim ${ns}/${claim} has device data (uuid=${uuid}, driverVersion=${driver_version})"
+}
+
 declare -a observed_gpus
 function gpu-already-seen {
   local gpu="$1"
@@ -103,6 +146,9 @@ if [ $gpu_test_1 != 2 ]; then
     exit 1
 fi
 
+# Verify that at least one ResourceClaim in gpu-test1 has device data
+verify-resourceclaim-device-status "gpu-test1"
+
 gpu_test1_pod0_ctr0_logs=$(kubectl logs -n gpu-test1 pod0 -c ctr0)
 gpu_test1_pod0_ctr0_gpus=$(gpus-from-logs "$gpu_test1_pod0_ctr0_logs")
 gpu_test1_pod0_ctr0_gpus_count=$(echo "$gpu_test1_pod0_ctr0_gpus" | wc -w | tr -d ' ')

test/e2e/setup-e2e.sh

@@ -35,5 +35,6 @@ helm upgrade -i \
   --namespace dra-example-driver \
   --set webhook.enabled=true \
   --set kubeletPlugin.numDevices=9 \
+  --set gpuDeviceStatus=true \
   dra-example-driver \
   deployments/helm/dra-example-driver