Merge pull request #394 from klueska/migrate-checkpoint

Migrate checkpoint format from the format in 25.3.0-rc.2 to the latest

Author: klueska

cmd/compute-domain-kubelet-plugin/checkpoint.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"encoding/json"
+	"fmt"
 
 	resourceapi "k8s.io/api/resource/v1beta1"
 	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager/checksum"
@@ -45,7 +46,31 @@ func (cp *Checkpoint) MarshalCheckpoint() ([]byte, error) {
 }
 
 func (cp *Checkpoint) UnmarshalCheckpoint(data []byte) error {
-	return json.Unmarshal(data, cp)
+	// If we can unmarshal to a Checkpoint directly we are done
+	if err := json.Unmarshal(data, cp); err == nil {
+		return nil
+	}
+
+	// Otherwise attempt to unmarshal to a Checkpoint2503RC2
+	// TODO: Remove this one release cycle following the v25.3.0 release
+	var cp2503rc2 Checkpoint2503RC2
+	if err := cp2503rc2.UnmarshalCheckpoint(data); err != nil {
+		return fmt.Errorf("unable to unmarshal as %T or %T", *cp, cp2503rc2)
+	}
+
+	// If that succeeded, verify its checksum...
+	if err := cp2503rc2.VerifyChecksum(); err != nil {
+		return fmt.Errorf("error verifying checksum for %T: %w", cp2503rc2, err)
+	}
+
+	// And then convert it to the v1 checkpoint format and try again
+	cpconv := cp2503rc2.ToV1()
+	data, err := cpconv.MarshalCheckpoint()
+	if err != nil {
+		return fmt.Errorf("error remarshalling %T after conversion from %T: %w", cpconv, cp2503rc2, err)
+	}
+
+	return cp.UnmarshalCheckpoint(data)
 }
 
 func (cp *Checkpoint) VerifyChecksum() error {

cmd/compute-domain-kubelet-plugin/checkpoint_legacy.go

@@ -0,0 +1,143 @@
+package main
+
+import (
+	"encoding/json"
+
+	"k8s.io/dynamic-resource-allocation/kubeletplugin"
+	drapbv1 "k8s.io/kubelet/pkg/apis/dra/v1beta1"
+	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager/checksum"
+)
+
+// Legacy structs and methods to help with conversion as necessary.
+type PreparedDeviceList2503RC2 []PreparedDevice2503RC2
+type PreparedDevices2503RC2 []*PreparedDeviceGroup2503RC2
+type PreparedClaims2503RC2 map[string]PreparedDevices2503RC2
+
+type PreparedDevice2503RC2 struct {
+	Channel *PreparedComputeDomainChannel2503RC2 `json:"channel"`
+	Daemon  *PreparedComputeDomainDaemon2503RC2  `json:"daemon"`
+}
+
+type PreparedComputeDomainChannel2503RC2 struct {
+	Info   *ComputeDomainChannelInfo `json:"info"`
+	Device *drapbv1.Device           `json:"device"`
+}
+
+type PreparedComputeDomainDaemon2503RC2 struct {
+	Info   *ComputeDomainDaemonInfo `json:"info"`
+	Device *drapbv1.Device          `json:"device"`
+}
+
+type PreparedDeviceGroup2503RC2 struct {
+	Devices     PreparedDeviceList2503RC2 `json:"devices"`
+	ConfigState DeviceConfigState         `json:"configState"`
+}
+
+type Checkpoint2503RC2 struct {
+	Checksum checksum.Checksum    `json:"checksum"`
+	V1       *Checkpoint2503RC2V1 `json:"v1,omitempty"`
+}
+
+type Checkpoint2503RC2V1 struct {
+	PreparedClaims PreparedClaims2503RC2 `json:"preparedClaims,omitempty"`
+}
+
+func (cp *Checkpoint2503RC2) MarshalCheckpoint() ([]byte, error) {
+	cp.Checksum = 0
+	out, err := json.Marshal(*cp)
+	if err != nil {
+		return nil, err
+	}
+	cp.Checksum = checksum.New(out)
+	return json.Marshal(*cp)
+}
+
+func (cp *Checkpoint2503RC2) UnmarshalCheckpoint(data []byte) error {
+	return json.Unmarshal(data, cp)
+}
+
+func (cp *Checkpoint2503RC2) VerifyChecksum() error {
+	ck := cp.Checksum
+	cp.Checksum = 0
+	defer func() {
+		cp.Checksum = ck
+	}()
+	out, err := json.Marshal(*cp)
+	if err != nil {
+		return err
+	}
+	return ck.Verify(out)
+}
+
+// ToV1 converts a PreparedDevice2503RC2 to a PreparedDevice.
+func (d *PreparedDevice2503RC2) ToV1() PreparedDevice {
+	device := PreparedDevice{}
+	if d.Channel != nil {
+		device.Channel = d.Channel.ToV1()
+	}
+	if d.Daemon != nil {
+		device.Daemon = d.Daemon.ToV1()
+	}
+	return device
+}
+
+// ToV1 converts a PreparedComputeDomainChannel2503RC2 to a PreparedComputeDomainChannel.
+func (c *PreparedComputeDomainChannel2503RC2) ToV1() *PreparedComputeDomainChannel {
+	channel := &PreparedComputeDomainChannel{}
+	if c.Info != nil {
+		channel.Info = c.Info
+	}
+	if c.Device != nil {
+		channel.Device = &kubeletplugin.Device{
+			Requests:     c.Device.RequestNames,
+			PoolName:     c.Device.PoolName,
+			DeviceName:   c.Device.DeviceName,
+			CDIDeviceIDs: c.Device.CDIDeviceIDs,
+		}
+	}
+	return channel
+}
+
+// ToV1 converts a PreparedComputeDomainDaemon2503RC2 to a PreparedComputeDomainDaemon.
+func (d *PreparedComputeDomainDaemon2503RC2) ToV1() *PreparedComputeDomainDaemon {
+	daemon := &PreparedComputeDomainDaemon{}
+	if d.Info != nil {
+		daemon.Info = d.Info
+	}
+	if d.Device != nil {
+		daemon.Device = &kubeletplugin.Device{
+			Requests:     d.Device.RequestNames,
+			PoolName:     d.Device.PoolName,
+			DeviceName:   d.Device.DeviceName,
+			CDIDeviceIDs: d.Device.CDIDeviceIDs,
+		}
+	}
+	return daemon
+}
+
+// ToV1 converts a PreparedDeviceGroup2503RC2 to a PreparedDeviceGroup.
+func (g *PreparedDeviceGroup2503RC2) ToV1() *PreparedDeviceGroup {
+	group := &PreparedDeviceGroup{
+		Devices:     make(PreparedDeviceList, 0, len(g.Devices)),
+		ConfigState: g.ConfigState,
+	}
+	for _, d := range g.Devices {
+		group.Devices = append(group.Devices, d.ToV1())
+	}
+	return group
+}
+
+// ToV1 converts a Checkpoint2503RC2 to a Checkpoint.
+func (cp *Checkpoint2503RC2) ToV1() *Checkpoint {
+	cpv1 := newCheckpoint()
+	for k, v := range cp.V1.PreparedClaims {
+		pds := make(PreparedDevices, 0, len(v))
+		for _, pd := range v {
+			pds = append(pds, pd.ToV1())
+		}
+		cpv1.V1.PreparedClaims[k] = PreparedClaim{
+			PreparedDevices: pds,
+		}
+	}
+	return cpv1
+}

cmd/compute-domain-kubelet-plugin/device_state.go

@@ -23,6 +23,7 @@ import (
 	"sync"
 
 	resourceapi "k8s.io/api/resource/v1beta1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/dynamic-resource-allocation/kubeletplugin"
 	"k8s.io/klog/v2"
@@ -194,6 +195,27 @@ func (s *DeviceState) Unprepare(ctx context.Context, claimRef kubeletplugin.Name
 		return nil
 	}
 
+	// If pc.Status.Allocation is 'nil', attempt to pull the status from the
+	// API server. This should only ever happen if we have unmarshaled from a
+	// legacy checkpoint format that did not include the Status field.
+	//
+	// TODO: Remove this one release cycle following the v25.3.0 release
+	if pc.Status.Allocation == nil {
+		klog.Infof("PreparedClaim status was unset in Checkpoint for ResourceClaim %s: attempting to pull it from API server", claimRef.String())
+		claim, err := s.config.clientsets.Core.ResourceV1beta1().ResourceClaims(claimRef.Namespace).Get(
+			ctx,
+			claimRef.Name,
+			metav1.GetOptions{})
+
+		if err != nil {
+			return permanentError{fmt.Errorf("failed to fetch ResourceClaim %s: %w", claimRef.String(), err)}
+		}
+		if claim.Status.Allocation == nil {
+			return permanentError{fmt.Errorf("no allocation set in ResourceClaim %s", claim.String())}
+		}
+		pc.Status = claim.Status
+	}
+
 	if err := s.unprepareDevices(ctx, &pc.Status); err != nil {
 		return fmt.Errorf("unprepare devices failed: %w", err)
 	}

go.mod

@@ -22,6 +22,7 @@ require (
 	k8s.io/component-base v0.33.0
 	k8s.io/dynamic-resource-allocation v0.33.0
 	k8s.io/klog/v2 v2.130.1
+	k8s.io/kubelet v0.33.0
 	k8s.io/kubernetes v1.33.0
 	k8s.io/mount-utils v0.33.0
 	k8s.io/utils v0.0.0-20241104100929-3ea5e8cea738
@@ -84,7 +85,6 @@ require (
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	k8s.io/kube-openapi v0.0.0-20250318190949-c8a335a9a2ff // indirect
-	k8s.io/kubelet v0.33.0 // indirect
 	sigs.k8s.io/json v0.0.0-20241010143419-9aa6b5e7a4b3 // indirect
 	sigs.k8s.io/randfill v1.0.0 // indirect
 	sigs.k8s.io/structured-merge-diff/v4 v4.6.0 // indirect