Do not GET claim in CD unprepare (use local checkpoint)

Signed-off-by: Dr. Jan-Philip Gehrcke <[email protected]>

Author: jgehrcke

cmd/compute-domain-controller/computedomain.go

@@ -191,7 +191,7 @@ func (m *ComputeDomainManager) RemoveFinalizer(ctx context.Context, uid string)
 	return nil
 }
 
-// AssertWorkloadsCompletes ensures that all workloads asssociated with a ComputeDomain have completed.
+// AssertWorkloadsCompletes ensures that all workloads associated with a ComputeDomain have completed.
 //
 // TODO: We should probably also check to ensure that all ResourceClaims
 // generated from our ResourceClaimTemplate for workloads are gone. Doing

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

@@ -3,6 +3,7 @@ package main
 import (
 	"encoding/json"
 
+	resourceapi "k8s.io/api/resource/v1beta1"
 	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager/checksum"
 )
 
@@ -12,14 +13,22 @@ type Checkpoint struct {
 }
 
 type CheckpointV1 struct {
-	PreparedClaims PreparedClaims `json:"preparedClaims,omitempty"`
+	PreparedClaimsByUID `json:"preparedClaims,omitempty"`
+}
+
+// key: stringified claim UUID
+type PreparedClaimsByUID map[string]PreparedClaim
+
+type PreparedClaim struct {
+	Status          resourceapi.ResourceClaimStatus
+	PreparedDevices PreparedDevices
 }
 
 func newCheckpoint() *Checkpoint {
 	pc := &Checkpoint{
 		Checksum: 0,
 		V1: &CheckpointV1{
-			PreparedClaims: make(PreparedClaims),
+			PreparedClaimsByUID: make(PreparedClaimsByUID),
 		},
 	}
 	return pc

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

@@ -437,7 +437,7 @@ func (m *ComputeDomainManager) periodicCleanup(ctx context.Context) {
 				continue
 			}
 			if err != nil {
-				klog.Errorf("error checking for existenc of directory '%s': %v", m.configFilesRoot, err)
+				klog.Errorf("error checking for existence of directory '%s': %v", m.configFilesRoot, err)
 				continue
 			}
 

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

@@ -115,13 +115,13 @@ func NewDeviceState(ctx context.Context, config *Config) (*DeviceState, error) {
 	}
 
 	for _, c := range checkpoints {
-		if c == DriverPluginCheckpointFile {
+		if c == DriverPluginCheckpointFileBasename {
 			return state, nil
 		}
 	}
 
 	checkpoint := newCheckpoint()
-	if err := state.checkpointManager.CreateCheckpoint(DriverPluginCheckpointFile, checkpoint); err != nil {
+	if err := state.checkpointManager.CreateCheckpoint(DriverPluginCheckpointFileBasename, checkpoint); err != nil {
 		return nil, fmt.Errorf("unable to sync to checkpoint: %v", err)
 	}
 
@@ -135,13 +135,17 @@ func (s *DeviceState) Prepare(ctx context.Context, claim *resourceapi.ResourceCl
 	claimUID := string(claim.UID)
 
 	checkpoint := newCheckpoint()
-	if err := s.checkpointManager.GetCheckpoint(DriverPluginCheckpointFile, checkpoint); err != nil {
-		return nil, fmt.Errorf("unable to sync from checkpoint: %v", err)
+	if err := s.checkpointManager.GetCheckpoint(DriverPluginCheckpointFileBasename, checkpoint); err != nil {
+		return nil, fmt.Errorf("unable to get checkpoint: %w", err)
 	}
-	preparedClaims := checkpoint.V1.PreparedClaims
 
-	if preparedClaims[claimUID] != nil {
-		return preparedClaims[claimUID].GetDevices(), nil
+	preparedClaim, exists := checkpoint.V1.PreparedClaimsByUID[claimUID]
+	if exists {
+		// Make this a noop. Associated device(s) has/ave been prepared by us.
+		// Prepare() must be idempotent, as it may be invoked more than once per
+		// claim (and actual device preparation must happen at most once).
+		klog.V(6).Infof("skip prepare: claim %v found in checkpoint", claimUID)
+		return preparedClaim.PreparedDevices.GetDevices(), nil
 	}
 
 	preparedDevices, err := s.prepareDevices(ctx, claim)
@@ -153,50 +157,65 @@ func (s *DeviceState) Prepare(ctx context.Context, claim *resourceapi.ResourceCl
 		return nil, fmt.Errorf("unable to create CDI spec file for claim: %w", err)
 	}
 
-	preparedClaims[claimUID] = preparedDevices
-	if err := s.checkpointManager.CreateCheckpoint(DriverPluginCheckpointFile, checkpoint); err != nil {
-		return nil, fmt.Errorf("unable to sync to checkpoint: %v", err)
+	// Add ResourceClaimStatus API object to node-local checkpoint: the
+	// 'unprepare' code path must use local state exclusively (ResourceClaim
+	// object might have been deleted from the API server).
+	checkpoint.V1.PreparedClaimsByUID[claimUID] = PreparedClaim{
+		Status:          claim.Status,
+		PreparedDevices: preparedDevices,
+	}
+	if err := s.checkpointManager.CreateCheckpoint(DriverPluginCheckpointFileBasename, checkpoint); err != nil {
+		return nil, fmt.Errorf("unable to create checkpoint: %w", err)
 	}
+	klog.V(6).Infof("checkpoint written for claim %v", claimUID)
 
-	return preparedClaims[claimUID].GetDevices(), nil
+	return preparedDevices.GetDevices(), nil
 }
 
-func (s *DeviceState) Unprepare(ctx context.Context, claim *resourceapi.ResourceClaim) error {
+func (s *DeviceState) Unprepare(ctx context.Context, claimRef kubeletplugin.NamespacedObject) error {
 	s.Lock()
 	defer s.Unlock()
 
-	claimUID := string(claim.UID)
-
-	if err := s.unprepareDevices(ctx, claim); err != nil {
-		return fmt.Errorf("unprepare devices failed: %w", err)
-	}
+	claimUID := string(claimRef.UID)
 
+	// Rely on local checkpoint state for ability to clean up.
 	checkpoint := newCheckpoint()
-	if err := s.checkpointManager.GetCheckpoint(DriverPluginCheckpointFile, checkpoint); err != nil {
-		return fmt.Errorf("unable to sync from checkpoint: %v", err)
+	if err := s.checkpointManager.GetCheckpoint(DriverPluginCheckpointFileBasename, checkpoint); err != nil {
+		return fmt.Errorf("unable to get checkpoint: %w", err)
 	}
-	preparedClaims := checkpoint.V1.PreparedClaims
 
-	if preparedClaims[claimUID] == nil {
+	pc, exists := checkpoint.V1.PreparedClaimsByUID[claimUID]
+	if !exists {
+		// Not an error: if this claim UID is not in the checkpoint then this
+		// device was never prepared or has already been unprepared (assume that
+		// Prepare+Checkpoint are done transactionally). Note that
+		// claimRef.String() contains namespace, name, UID.
+		klog.Infof("unprepare noop: claim not found in checkpoint data: %v", claimRef.String())
 		return nil
 	}
 
+	if err := s.unprepareDevices(ctx, &pc.Status); err != nil {
+		return fmt.Errorf("unprepare devices failed: %w", err)
+	}
+
 	err := s.cdi.DeleteClaimSpecFile(claimUID)
 	if err != nil {
 		return fmt.Errorf("unable to delete CDI spec file for claim: %w", err)
 	}
 
-	delete(preparedClaims, claimUID)
-	if err := s.checkpointManager.CreateCheckpoint(DriverPluginCheckpointFile, checkpoint); err != nil {
-		return fmt.Errorf("unable to sync to checkpoint: %v", err)
+	// Write new checkpoint reflecting that all devices for this claim have been
+	// unprepared (by virtue of removing its UID from all mappings).
+	delete(checkpoint.V1.PreparedClaimsByUID, claimUID)
+	if err := s.checkpointManager.CreateCheckpoint(DriverPluginCheckpointFileBasename, checkpoint); err != nil {
+		return fmt.Errorf("create checkpoint failed: %w", err)
 	}
 
 	return nil
 }
 
 func (s *DeviceState) prepareDevices(ctx context.Context, claim *resourceapi.ResourceClaim) (PreparedDevices, error) {
 	// Generate a mapping of each OpaqueDeviceConfigs to the Device.Results it applies to
-	configResultsMap, err := s.getConfigResultsMap(claim)
+	configResultsMap, err := s.getConfigResultsMap(&claim.Status)
 	if err != nil {
 		return nil, fmt.Errorf("error generating configResultsMap: %w", err)
 	}
@@ -283,9 +302,9 @@ func (s *DeviceState) prepareDevices(ctx context.Context, claim *resourceapi.Res
 	return preparedDevices, nil
 }
 
-func (s *DeviceState) unprepareDevices(ctx context.Context, claim *resourceapi.ResourceClaim) error {
+func (s *DeviceState) unprepareDevices(ctx context.Context, cs *resourceapi.ResourceClaimStatus) error {
 	// Generate a mapping of each OpaqueDeviceConfigs to the Device.Results it applies to
-	configResultsMap, err := s.getConfigResultsMap(claim)
+	configResultsMap, err := s.getConfigResultsMap(cs)
 	if err != nil {
 		return fmt.Errorf("error generating configResultsMap: %w", err)
 	}
@@ -407,12 +426,12 @@ func (s *DeviceState) applyComputeDomainDaemonConfig(ctx context.Context, config
 	return &configState, nil
 }
 
-func (s *DeviceState) getConfigResultsMap(claim *resourceapi.ResourceClaim) (map[runtime.Object][]*resourceapi.DeviceRequestAllocationResult, error) {
+func (s *DeviceState) getConfigResultsMap(rcs *resourceapi.ResourceClaimStatus) (map[runtime.Object][]*resourceapi.DeviceRequestAllocationResult, error) {
 	// Retrieve the full set of device configs for the driver.
 	configs, err := GetOpaqueDeviceConfigs(
 		configapi.Decoder,
 		DriverName,
-		claim.Status.Allocation.Devices.Config,
+		rcs.Allocation.Devices.Config,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("error getting opaque device configs: %v", err)
@@ -433,7 +452,7 @@ func (s *DeviceState) getConfigResultsMap(claim *resourceapi.ResourceClaim) (map
 	// Look through the configs and figure out which one will be applied to
 	// each device allocation result based on their order of precedence and type.
 	configResultsMap := make(map[runtime.Object][]*resourceapi.DeviceRequestAllocationResult)
-	for _, result := range claim.Status.Allocation.Devices.Results {
+	for _, result := range rcs.Allocation.Devices.Results {
 		device, exists := s.allocatable[result.Device]
 		if !exists {
 			return nil, fmt.Errorf("requested device is not allocatable: %v", result.Device)

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

@@ -24,7 +24,6 @@ import (
 	"time"
 
 	resourceapi "k8s.io/api/resource/v1beta1"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 	coreclientset "k8s.io/client-go/kubernetes"
 	"k8s.io/dynamic-resource-allocation/kubeletplugin"
@@ -117,7 +116,7 @@ func (d *driver) Shutdown() error {
 }
 
 func (d *driver) PrepareResourceClaims(ctx context.Context, claims []*resourceapi.ResourceClaim) (map[types.UID]kubeletplugin.PrepareResult, error) {
-	klog.Infof("PrepareResourceClaims called with %d claim(s)", len(claims))
+	klog.V(6).Infof("PrepareResourceClaims called with %d claim(s)", len(claims))
 
 	var wg sync.WaitGroup
 	ctx, cancel := context.WithTimeout(ctx, ErrorRetryMaxTimeout)
@@ -147,7 +146,7 @@ func (d *driver) PrepareResourceClaims(ctx context.Context, claims []*resourceap
 }
 
 func (d *driver) UnprepareResourceClaims(ctx context.Context, claimRefs []kubeletplugin.NamespacedObject) (map[types.UID]error, error) {
-	klog.Infof("UnprepareResourceClaims called with %d claim(s)", len(claimRefs))
+	klog.V(6).Infof("UnprepareResourceClaims called with %d claim(s)", len(claimRefs))
 
 	var wg sync.WaitGroup
 	ctx, cancel := context.WithTimeout(ctx, ErrorRetryMaxTimeout)
@@ -204,27 +203,11 @@ func (d *driver) nodeUnprepareResource(ctx context.Context, claimRef kubeletplug
 	d.Lock()
 	defer d.Unlock()
 
-	// Fetching the resource claim should not be needed (and not be done) in the
-	// unprepare code path. Any state required during unprepare can be stored
-	// via checkpointing.
-	claim, err := d.client.ResourceV1beta1().ResourceClaims(claimRef.Namespace).Get(
-		ctx,
-		claimRef.Name,
-		metav1.GetOptions{})
-
-	if err != nil {
-		return isPermanentError(err), fmt.Errorf("failed to fetch ResourceClaim %s in namespace %s: %w", claimRef.Name, claimRef.Namespace, err)
-	}
-
-	if claim.Status.Allocation == nil {
-		return true, fmt.Errorf("no allocation set in ResourceClaim %s in namespace %s", claim.Name, claim.Namespace)
-	}
-
-	if err := d.state.Unprepare(ctx, claim); err != nil {
-		return isPermanentError(err), fmt.Errorf("error unpreparing devices for claim '%v': %w", claim.UID, err)
+	if err := d.state.Unprepare(ctx, claimRef); err != nil {
+		return isPermanentError(err), fmt.Errorf("error unpreparing devices for claim '%v': %w", claimRef.String(), err)
 	}
 
-	klog.Infof("unprepared devices for claim '%v'", claim.UID)
+	klog.Infof("unprepared devices for claim '%v'", claimRef.String())
 	return true, nil
 }
 

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

@@ -32,9 +32,9 @@ import (
 )
 
 const (
-	DriverName                 = "compute-domain.nvidia.com"
-	DriverPluginPath           = "/var/lib/kubelet/plugins/" + DriverName
-	DriverPluginCheckpointFile = "checkpoint.json"
+	DriverName                         = "compute-domain.nvidia.com"
+	DriverPluginPath                   = "/var/lib/kubelet/plugins/" + DriverName
+	DriverPluginCheckpointFileBasename = "checkpoint.json"
 )
 
 type Flags struct {

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

@@ -22,7 +22,6 @@ import (
 
 type PreparedDeviceList []PreparedDevice
 type PreparedDevices []*PreparedDeviceGroup
-type PreparedClaims map[string]PreparedDevices
 
 type PreparedDevice struct {
 	Channel *PreparedComputeDomainChannel `json:"channel"`

cmd/gpu-kubelet-plugin/checkpoint.go

@@ -3,6 +3,7 @@ package main
 import (
 	"encoding/json"
 
+	resourceapi "k8s.io/api/resource/v1beta1"
 	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager/checksum"
 )
 
@@ -12,14 +13,22 @@ type Checkpoint struct {
 }
 
 type CheckpointV1 struct {
-	PreparedClaims PreparedClaims `json:"preparedClaims,omitempty"`
+	PreparedClaimsByUID `json:"preparedClaims,omitempty"`
+}
+
+// key: stringified claim UUID
+type PreparedClaimsByUID map[string]PreparedClaim
+
+type PreparedClaim struct {
+	Status          resourceapi.ResourceClaimStatus
+	PreparedDevices PreparedDevices
 }
 
 func newCheckpoint() *Checkpoint {
 	pc := &Checkpoint{
 		Checksum: 0,
 		V1: &CheckpointV1{
-			PreparedClaims: make(PreparedClaims),
+			PreparedClaimsByUID: make(PreparedClaimsByUID),
 		},
 	}
 	return pc