WIP: better versioned checkpoints

Author: nojnhuh

Makefile

@@ -102,15 +102,24 @@ coverage: test
 	cat $(COVERAGE_FILE) | grep -v "_mock.go" > $(COVERAGE_FILE).no-mocks
 	go tool cover -func=$(COVERAGE_FILE).no-mocks
 
-generate: generate-deepcopy
+generate: generate-deepcopy generate-conversion
 
 generate-deepcopy:
 	for api in $(APIS); do \
-		rm -f $(CURDIR)/api/$(VENDOR)/resource/$${api}/zz_generated.deepcopy.go; \
+		rm -f $${api}/zz_generated.deepcopy.go; \
 		controller-gen \
 			object:headerFile=$(CURDIR)/hack/boilerplate.generatego.txt \
-			paths=$(CURDIR)/api/$(VENDOR)/resource/$${api}/ \
-			output:object:dir=$(CURDIR)/api/$(VENDOR)/resource/$${api}; \
+			paths=$${api}/ \
+			output:object:dir=$${api}; \
+	done
+
+generate-conversion:
+	for api in $(APIS); do \
+		rm -f $${api}/zz_generated.conversion.go; \
+		conversion-gen \
+			--go-header-file=$(CURDIR)/hack/boilerplate.generatego.txt \
+			--output-file=zz_generated.conversion.go \
+			$${api}/; \
 	done
 
 setup-e2e:
@@ -128,7 +137,6 @@ teardown-e2e:
 .build-image: docker/Dockerfile.devel
 	if [ x"$(SKIP_IMAGE_BUILD)" = x"" ]; then \
 		$(CONTAINER_TOOL) build \
-			--progress=plain \
 			--build-arg GOLANG_VERSION="$(GOLANG_VERSION)" \
 			--tag $(BUILDIMAGE) \
 			-f $(^) \

cmd/dra-example-kubeletplugin/cdi.go

@@ -26,7 +26,7 @@ import (
 	cdiparser "tags.cncf.io/container-device-interface/pkg/parser"
 	cdispec "tags.cncf.io/container-device-interface/specs-go"
 
-	"sigs.k8s.io/dra-example-driver/internal/profiles"
+	checkpointapi "sigs.k8s.io/dra-example-driver/internal/api/checkpoint"
 )
 
 const cdiCommonDeviceName = "common"
@@ -85,7 +85,7 @@ func (cdi *CDIHandler) CreateCommonSpecFile() error {
 	return cdi.cache.WriteSpec(spec, specName)
 }
 
-func (cdi *CDIHandler) CreateClaimSpecFile(claimUID string, devices profiles.PreparedDevices) error {
+func (cdi *CDIHandler) CreateClaimSpecFile(claimUID string, devices []checkpointapi.PreparedDevice) error {
 	specName := cdiapi.GenerateTransientSpecName(cdi.vendor(), cdi.class, claimUID)
 
 	spec := &cdispec.Spec{

cmd/dra-example-kubeletplugin/checkpoint.go

@@ -17,49 +17,31 @@
 package main
 
 import (
-	"encoding/json"
 	"fmt"
 	"os"
 	"path/filepath"
 
-	"sigs.k8s.io/dra-example-driver/internal/profiles"
-)
-
-type PreparedClaims map[string]profiles.PreparedDevices
-
-type Checkpoint struct {
-	V1 *CheckpointV1 `json:"v1,omitempty"`
-}
-
-type CheckpointV1 struct {
-	PreparedClaims PreparedClaims `json:"preparedClaims,omitempty"`
-}
+	"k8s.io/apimachinery/pkg/runtime/serializer/json"
+	drapbv1 "k8s.io/kubelet/pkg/apis/dra/v1beta1"
+	"k8s.io/utils/ptr"
 
-func newCheckpoint() *Checkpoint {
-	pc := &Checkpoint{
-		V1: &CheckpointV1{},
-	}
-	return pc
-}
+	checkpointapi "sigs.k8s.io/dra-example-driver/internal/api/checkpoint"
+)
 
-func readCheckpoint(path string) (*Checkpoint, error) {
+func readCheckpoint(serializer *json.Serializer, path string) (*checkpointapi.Checkpoint, error) {
 	data, err := os.ReadFile(path)
 	if err != nil {
 		return nil, err
 	}
-	checkpoint := new(Checkpoint)
-	err = json.Unmarshal(data, checkpoint)
+	checkpoint := new(checkpointapi.Checkpoint)
+	_, _, err = serializer.Decode(data, ptr.To(checkpointapi.SchemeGroupVersion.WithKind("Checkpoint")), checkpoint)
 	if err != nil {
 		return nil, fmt.Errorf("unmarshal json from %s: %w", path, err)
 	}
 	return checkpoint, nil
 }
 
-func writeCheckpoint(path string, checkpoint *Checkpoint) (err error) {
-	data, err := json.Marshal(checkpoint)
-	if err != nil {
-		return fmt.Errorf("marshal json: %w", err)
-	}
+func writeCheckpoint(serializer *json.Serializer, path string, checkpoint *checkpointapi.Checkpoint) (err error) {
 	dir := filepath.Dir(path)
 	tmp, err := os.CreateTemp(dir, "tmp-checkpoint-*")
 	if err != nil {
@@ -70,8 +52,9 @@ func writeCheckpoint(path string, checkpoint *Checkpoint) (err error) {
 			err = fmt.Errorf("close temp file: %w", err1)
 		}
 	}()
-	if _, err := tmp.Write(data); err != nil {
-		return fmt.Errorf("write to temp file %s: %w", tmp.Name(), err)
+	err = serializer.Encode(checkpoint, tmp)
+	if err != nil {
+		return fmt.Errorf("serialize json: %w", err)
 	}
 	if err := tmp.Sync(); err != nil {
 		return fmt.Errorf("sync temp file: %w", err)
@@ -82,32 +65,15 @@ func writeCheckpoint(path string, checkpoint *Checkpoint) (err error) {
 	return nil
 }
 
-func (cp *Checkpoint) GetPreparedDevices(claimUID string) profiles.PreparedDevices {
-	if cp.V1 == nil {
-		return nil
+func getPreparedDevices(preparedClaim checkpointapi.PreparedClaim) []*drapbv1.Device {
+	var devices []*drapbv1.Device
+	for _, preparedDevice := range preparedClaim.PreparedDevices {
+		devices = append(devices, &drapbv1.Device{
+			RequestNames: preparedDevice.Device.RequestNames,
+			PoolName:     preparedDevice.Device.PoolName,
+			DeviceName:   preparedDevice.Device.DeviceName,
+			CdiDeviceIds: preparedDevice.Device.CdiDeviceIds,
+		})
 	}
-	if devices, ok := cp.V1.PreparedClaims[claimUID]; ok {
-		return devices
-	}
-	return nil
-}
-
-func (cp *Checkpoint) AddPreparedDevices(claimUID string, pds profiles.PreparedDevices) {
-	if cp.V1 == nil {
-		return
-	}
-
-	if cp.V1.PreparedClaims == nil {
-		cp.V1.PreparedClaims = make(PreparedClaims)
-	}
-
-	cp.V1.PreparedClaims[claimUID] = pds
-}
-
-func (cp *Checkpoint) RemovePreparedDevices(claimUID string) {
-	if cp.V1 == nil {
-		return
-	}
-
-	delete(cp.V1.PreparedClaims, claimUID)
+	return devices
 }

cmd/dra-example-kubeletplugin/state.go

@@ -30,6 +30,8 @@ import (
 	"k8s.io/dynamic-resource-allocation/resourceslice"
 	drapbv1 "k8s.io/kubelet/pkg/apis/dra/v1beta1"
 
+	checkpointapi "sigs.k8s.io/dra-example-driver/internal/api/checkpoint"
+	checkpointv1alpha1 "sigs.k8s.io/dra-example-driver/internal/api/checkpoint/v1alpha1"
 	"sigs.k8s.io/dra-example-driver/internal/profiles"
 )
 
@@ -42,13 +44,14 @@ type OpaqueDeviceConfig struct {
 
 type DeviceState struct {
 	sync.Mutex
-	driverName      string
-	cdi             *CDIHandler
-	driverResources resourceslice.DriverResources
-	allocatable     AllocatableDevices
-	configDecoder   runtime.Decoder
-	configHandler   profiles.ConfigHandler
-	checkpointPath  string
+	driverName           string
+	cdi                  *CDIHandler
+	driverResources      resourceslice.DriverResources
+	allocatable          AllocatableDevices
+	configDecoder        runtime.Decoder
+	configHandler        profiles.ConfigHandler
+	checkpointPath       string
+	checkpointSerializer *json.Serializer
 }
 
 func NewDeviceState(config *Config) (*DeviceState, error) {
@@ -84,6 +87,23 @@ func NewDeviceState(config *Config) (*DeviceState, error) {
 		},
 	)
 
+	checkpointScheme := runtime.NewScheme()
+	sb = runtime.NewSchemeBuilder(
+		checkpointapi.AddToScheme,
+		checkpointv1alpha1.AddToScheme,
+	)
+	if err := sb.AddToScheme(checkpointScheme); err != nil {
+		return nil, fmt.Errorf("create config scheme: %w", err)
+	}
+	checkpointSerializer := json.NewSerializerWithOptions(
+		json.DefaultMetaFactory,
+		checkpointScheme,
+		checkpointScheme,
+		json.SerializerOptions{
+			Pretty: true, Strict: true,
+		},
+	)
+
 	allocatable := make(AllocatableDevices)
 	for _, slice := range driverResources.Pools[config.flags.nodeName].Slices {
 		for _, device := range slice.Devices {
@@ -92,24 +112,25 @@ func NewDeviceState(config *Config) (*DeviceState, error) {
 	}
 
 	state := &DeviceState{
-		driverName:      config.flags.driverName,
-		cdi:             cdi,
-		driverResources: driverResources,
-		allocatable:     allocatable,
-		configDecoder:   decoder,
-		configHandler:   configHandler,
-		checkpointPath:  filepath.Join(config.DriverPluginPath(), DriverPluginCheckpointFile),
+		driverName:           config.flags.driverName,
+		cdi:                  cdi,
+		driverResources:      driverResources,
+		allocatable:          allocatable,
+		configDecoder:        decoder,
+		configHandler:        configHandler,
+		checkpointPath:       filepath.Join(config.DriverPluginPath(), DriverPluginCheckpointFile),
+		checkpointSerializer: checkpointSerializer,
 	}
 
-	_, err = readCheckpoint(state.checkpointPath)
+	_, err = readCheckpoint(state.checkpointSerializer, state.checkpointPath)
 	if err != nil && !errors.Is(err, fs.ErrNotExist) {
 		return nil, fmt.Errorf("failed to read checkpoint: %w", err)
 	}
 	if err == nil {
 		return state, nil
 	}
 
-	if err := writeCheckpoint(state.checkpointPath, newCheckpoint()); err != nil {
+	if err := writeCheckpoint(state.checkpointSerializer, state.checkpointPath, &checkpointapi.Checkpoint{}); err != nil {
 		return nil, fmt.Errorf("unable to sync to checkpoint: %v", err)
 	}
 
@@ -122,30 +143,30 @@ func (s *DeviceState) Prepare(claim *resourceapi.ResourceClaim) ([]*drapbv1.Devi
 
 	claimUID := string(claim.UID)
 
-	checkpoint, err := readCheckpoint(s.checkpointPath)
+	checkpoint, err := readCheckpoint(s.checkpointSerializer, s.checkpointPath)
 	if err != nil {
 		return nil, fmt.Errorf("unable to sync from checkpoint: %v", err)
 	}
 
-	preparedDevices := checkpoint.GetPreparedDevices(claimUID)
-	if preparedDevices != nil {
-		return preparedDevices.GetDevices(), nil
+	preparedClaim, ok := checkpoint.PreparedClaims[claim.UID]
+	if ok {
+		return getPreparedDevices(preparedClaim), nil
 	}
-	preparedDevices, err = s.prepareDevices(claim)
+	preparedClaim, err = s.prepareDevices(claim)
 	if err != nil {
 		return nil, fmt.Errorf("prepare failed: %v", err)
 	}
 
-	if err = s.cdi.CreateClaimSpecFile(claimUID, preparedDevices); err != nil {
+	if err = s.cdi.CreateClaimSpecFile(claimUID, preparedClaim); err != nil {
 		return nil, fmt.Errorf("unable to create CDI spec file for claim: %v", err)
 	}
 
-	checkpoint.AddPreparedDevices(claimUID, preparedDevices)
+	checkpoint.AddPreparedDevices(claimUID, preparedClaim)
 	if err := writeCheckpoint(s.checkpointPath, checkpoint); err != nil {
 		return nil, fmt.Errorf("unable to sync to checkpoint: %v", err)
 	}
 
-	return preparedDevices.GetDevices(), nil
+	return preparedClaim.GetDevices(), nil
 }
 
 func (s *DeviceState) Unprepare(claimUID string) error {
@@ -182,9 +203,9 @@ func (s *DeviceState) Unprepare(claimUID string) error {
 	return nil
 }
 
-func (s *DeviceState) prepareDevices(claim *resourceapi.ResourceClaim) (profiles.PreparedDevices, error) {
+func (s *DeviceState) prepareDevices(claim *resourceapi.ResourceClaim) (checkpointapi.PreparedClaim, error) {
 	if claim.Status.Allocation == nil {
-		return nil, fmt.Errorf("claim not yet allocated")
+		return checkpointapi.PreparedClaim{}, fmt.Errorf("claim not yet allocated")
 	}
 	// Check if any device request has admin access
 	hasAdminAccess := s.checkAdminAccess(claim)
@@ -196,7 +217,7 @@ func (s *DeviceState) prepareDevices(claim *resourceapi.ResourceClaim) (profiles
 		claim.Status.Allocation.Devices.Config,
 	)
 	if err != nil {
-		return nil, fmt.Errorf("error getting opaque device configs: %v", err)
+		return checkpointapi.PreparedClaim{}, fmt.Errorf("error getting opaque device configs: %v", err)
 	}
 
 	// Add the default GPU Config to the front of the config list with the
@@ -213,7 +234,7 @@ func (s *DeviceState) prepareDevices(claim *resourceapi.ResourceClaim) (profiles
 			continue
 		}
 		if _, exists := s.allocatable[result.Device]; !exists {
-			return nil, fmt.Errorf("requested device is not allocatable: %v", result.Device)
+			return checkpointapi.PreparedClaim{}, fmt.Errorf("requested device is not allocatable: %v", result.Device)
 		}
 		for _, c := range slices.Backward(configs) {
 			if len(c.Requests) == 0 || slices.Contains(c.Requests, result.Request) {
@@ -231,7 +252,7 @@ func (s *DeviceState) prepareDevices(claim *resourceapi.ResourceClaim) (profiles
 		// Apply the config to the list of results associated with it.
 		containerEdits, err := s.configHandler.ApplyConfig(config, results)
 		if err != nil {
-			return nil, fmt.Errorf("error applying config: %w", err)
+			return checkpointapi.PreparedClaim{}, fmt.Errorf("error applying config: %w", err)
 		}
 
 		// Merge any new container edits with the overall per device map.
@@ -242,11 +263,11 @@ func (s *DeviceState) prepareDevices(claim *resourceapi.ResourceClaim) (profiles
 
 	// Walk through each config and its associated device allocation results
 	// and construct the list of prepared devices to return.
-	var preparedDevices profiles.PreparedDevices
+	var preparedDevices []checkpointapi.PreparedDevice
 	for _, results := range configResultsMap {
 		for _, result := range results {
-			device := &profiles.PreparedDevice{
-				Device: drapbv1.Device{
+			device := checkpointapi.PreparedDevice{
+				Device: checkpointapi.Device{
 					RequestNames: []string{result.Request},
 					PoolName:     result.Pool,
 					DeviceName:   result.Device,
@@ -259,7 +280,7 @@ func (s *DeviceState) prepareDevices(claim *resourceapi.ResourceClaim) (profiles
 		}
 	}
 
-	return preparedDevices, nil
+	return checkpointapi.PreparedClaim{PreparedDevices: preparedDevices}, nil
 }
 
 func (s *DeviceState) unprepareDevices(claimUID string, devices profiles.PreparedDevices) error {

common.mk

@@ -21,7 +21,7 @@ VERSION  ?=
 vVERSION := v$(VERSION:v%=%)
 
 VENDOR := example.com
-APIS := gpu/v1alpha1
+APIS := $(CURDIR)/api/$(VENDOR)/resource/gpu/v1alpha1 $(CURDIR)/internal/api/checkpoint $(CURDIR)/internal/api/checkpoint/v1alpha1
 
 PLURAL_EXCEPTIONS  = DeviceClassParameters:DeviceClassParameters
 PLURAL_EXCEPTIONS += GpuClaimParameters:GpuClaimParameters

docker/Dockerfile.devel

@@ -20,3 +20,4 @@ RUN go install github.com/gordonklaus/ineffassign@latest && \
 RUN go install github.com/golangci/golangci-lint/cmd/golangci-lint@v1.52.0
 RUN go install sigs.k8s.io/controller-tools/cmd/controller-gen@v0.20.0
 RUN go install k8s.io/code-generator/cmd/client-gen@v0.35.0
+RUN go install k8s.io/code-generator/cmd/conversion-gen@v0.35.0

internal/api/checkpoint/doc.go

@@ -0,0 +1,23 @@
+/*
+Copyright 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.
+*/
+
+// +k8s:deepcopy-gen=package
+
+// Package checkpoint contains internal (unversioned) types for checkpoints.
+// These types are the canonical in-memory representation that the driver
+// programs against. Versioned types (e.g. v1alpha1) are converted to/from these
+// internal types via the scheme.
+package checkpoint