DRA: refactor checkpointing

Co-authored-by: Kevin Klues <klueska@gmail.com>

Author: bart0sh

pkg/kubelet/cm/dra/checkpoint/v1/apis.go

@@ -0,0 +1,69 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package v1
+
+import (
+	"encoding/json"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager/checksum"
+)
+
+func New(data []ClaimInfoState) *Checkpoint {
+	pc := &Checkpoint{
+		TypeMeta: metav1.TypeMeta{
+			Kind:       CheckpointKind,
+			APIVersion: CheckpointAPIVersion,
+		},
+		Data:     data,
+		Checksum: 0,
+	}
+	return pc
+}
+
+// MarshalCheckpoint marshals checkpoint to JSON
+func (cp *Checkpoint) 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)
+}
+
+// UnmarshalCheckpoint tries to unmarshal passed bytes to checkpoint
+func (cp *Checkpoint) UnmarshalCheckpoint(blob []byte) error {
+	if err := json.Unmarshal(blob, cp); err != nil {
+		return err
+	}
+	return cp.VerifyChecksum()
+}
+
+// VerifyChecksum verifies that current checksum of checkpoint is valid
+func (cp *Checkpoint) 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)
+}

pkg/kubelet/cm/dra/checkpoint/v1/types.go

@@ -0,0 +1,75 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package v1
+
+import (
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/sets"
+	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager/checksum"
+)
+
+const (
+	GroupName            = "checkpoint.dra.kubelet.k8s.io"
+	CheckpointKind       = "DRACheckpoint"
+	CheckpointAPIVersion = GroupName + "/v1"
+)
+
+type ClaimInfoStateList []ClaimInfoState
+
+// +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object
+// +k8s:deepcopy-gen=true
+type Checkpoint struct {
+	metav1.TypeMeta
+	Data     ClaimInfoStateList
+	Checksum checksum.Checksum
+}
+
+// +k8s:deepcopy-gen=true
+type ClaimInfoState struct {
+	// ClaimUID is the UID of a resource claim
+	ClaimUID types.UID
+
+	// ClaimName is the name of a resource claim
+	ClaimName string
+
+	// Namespace is a claim namespace
+	Namespace string
+
+	// PodUIDs is a set of pod UIDs that reference a resource
+	PodUIDs sets.Set[string]
+
+	// DriverState contains information about all drivers which have allocation
+	// results in the claim, even if they don't provide devices for their results.
+	DriverState map[string]DriverState
+}
+
+// DriverState is used to store per-device claim info state in a checkpoint
+// +k8s:deepcopy-gen=true
+type DriverState struct {
+	Devices []Device
+}
+
+// Device is how a DRA driver described an allocated device in a claim
+// to kubelet. RequestName and CDI device IDs are optional.
+// +k8s:deepcopy-gen=true
+type Device struct {
+	PoolName     string
+	DeviceName   string
+	RequestNames []string
+	CDIDeviceIDs []string
+}

pkg/kubelet/cm/dra/state/zz_generated.deepcopy.go renamed to pkg/kubelet/cm/dra/checkpoint/v1/zz_generated.deepcopy.go

@@ -25,15 +25,16 @@ import (
 	resourceapi "k8s.io/api/resource/v1alpha3"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/sets"
-	"k8s.io/kubernetes/pkg/kubelet/cm/dra/state"
+	checkpointapi "k8s.io/kubernetes/pkg/kubelet/cm/dra/checkpoint/v1"
+	state "k8s.io/kubernetes/pkg/kubelet/cm/dra/state"
 	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
 )
 
 // ClaimInfo holds information required
 // to prepare and unprepare a resource claim.
 // +k8s:deepcopy-gen=true
 type ClaimInfo struct {
-	state.ClaimInfoState
+	checkpointapi.ClaimInfoState
 	prepared bool
 }
 
@@ -47,18 +48,18 @@ type claimInfoCache struct {
 // newClaimInfoFromClaim creates a new claim info from a resource claim.
 // It verifies that the kubelet can handle the claim.
 func newClaimInfoFromClaim(claim *resourceapi.ResourceClaim) (*ClaimInfo, error) {
-	claimInfoState := state.ClaimInfoState{
+	claimInfoState := checkpointapi.ClaimInfoState{
 		ClaimUID:    claim.UID,
 		ClaimName:   claim.Name,
 		Namespace:   claim.Namespace,
 		PodUIDs:     sets.New[string](),
-		DriverState: make(map[string]state.DriverState),
+		DriverState: make(map[string]checkpointapi.DriverState),
 	}
 	if claim.Status.Allocation == nil {
 		return nil, errors.New("not allocated")
 	}
 	for _, result := range claim.Status.Allocation.Devices.Results {
-		claimInfoState.DriverState[result.Driver] = state.DriverState{}
+		claimInfoState.DriverState[result.Driver] = checkpointapi.DriverState{}
 	}
 	info := &ClaimInfo{
 		ClaimInfoState: claimInfoState,
@@ -68,7 +69,7 @@ func newClaimInfoFromClaim(claim *resourceapi.ResourceClaim) (*ClaimInfo, error)
 }
 
 // newClaimInfoFromClaim creates a new claim info from a checkpointed claim info state object.
-func newClaimInfoFromState(state *state.ClaimInfoState) *ClaimInfo {
+func newClaimInfoFromState(state *checkpointapi.ClaimInfoState) *ClaimInfo {
 	info := &ClaimInfo{
 		ClaimInfoState: *state.DeepCopy(),
 		prepared:       false,
@@ -77,9 +78,9 @@ func newClaimInfoFromState(state *state.ClaimInfoState) *ClaimInfo {
 }
 
 // setCDIDevices adds a set of CDI devices to the claim info.
-func (info *ClaimInfo) addDevice(driverName string, device state.Device) {
+func (info *ClaimInfo) addDevice(driverName string, device checkpointapi.Device) {
 	if info.DriverState == nil {
-		info.DriverState = make(map[string]state.DriverState)
+		info.DriverState = make(map[string]checkpointapi.DriverState)
 	}
 	driverState := info.DriverState[driverName]
 	driverState.Devices = append(driverState.Devices, device)
@@ -188,7 +189,7 @@ func (cache *claimInfoCache) hasPodReference(UID types.UID) bool {
 
 // syncToCheckpoint syncs the full claim info cache state to a checkpoint.
 func (cache *claimInfoCache) syncToCheckpoint() error {
-	claimInfoStateList := make(state.ClaimInfoStateList, 0, len(cache.claimInfo))
+	claimInfoStateList := make(checkpointapi.ClaimInfoStateList, 0, len(cache.claimInfo))
 	for _, infoClaim := range cache.claimInfo {
 		claimInfoStateList = append(claimInfoStateList, infoClaim.ClaimInfoState)
 	}