Fix Prepare() idempotency by preventing reverse state transition and atomically detecting already-prepared claims.

* Avoid overwriting PrepareCompleted with PrepareStarted on subsequent Prepare() calls.
* Move claim prepared state checks into a single updateCheckpoint call for atomicity.
* Ensure checkpoint state does not transition from completed to started state again.

Signed-off-by: Shiva Krishna, Merla <[email protected]>

Author: shivamerla

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

@@ -149,34 +149,40 @@ func (s *DeviceState) Prepare(ctx context.Context, claim *resourceapi.ResourceCl
 
 	claimUID := string(claim.UID)
 
-	checkpoint, err := s.getCheckpoint()
-	if err != nil {
-		return nil, fmt.Errorf("unable to get checkpoint: %w", err)
-	}
+	var preparedDevices PreparedDevices
+	var alreadyPrepared bool
 
-	preparedClaim, exists := checkpoint.V2.PreparedClaims[claimUID]
-	if exists && preparedClaim.CheckpointState == ClaimCheckpointStatePrepareCompleted {
-		// 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(4).Infof("Skip prepare: claim %v found in checkpoint", claimUID)
-		return preparedClaim.PreparedDevices.GetDevices(), nil
-	}
+	// Atomically check if prepare for this claim is already completed, otherwise mark as started.
+	err := s.updateCheckpoint(func(cp *Checkpoint) {
+		if pc, ok := cp.V2.PreparedClaims[claimUID]; ok &&
+			pc.CheckpointState == ClaimCheckpointStatePrepareCompleted {
 
-	err = s.updateCheckpoint(func(checkpoint *Checkpoint) {
-		checkpoint.V2.PreparedClaims[claimUID] = PreparedClaim{
+			alreadyPrepared = true
+			preparedDevices = pc.PreparedDevices
+			return
+		}
+
+		cp.V2.PreparedClaims[claimUID] = PreparedClaim{
 			CheckpointState: ClaimCheckpointStatePrepareStarted,
 			Status:          claim.Status,
 			Name:            claim.Name,
 			Namespace:       claim.Namespace,
 		}
+		klog.V(6).Infof("checkpoint updated for claim %v", claimUID)
 	})
 	if err != nil {
 		return nil, fmt.Errorf("unable to update checkpoint: %w", err)
 	}
-	klog.V(6).Infof("checkpoint updated for claim %v", claimUID)
 
-	preparedDevices, err := s.prepareDevices(ctx, claim)
+	if alreadyPrepared {
+		// 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(4).Infof("skip prepare: claim %v found in checkpoint", claimUID)
+		return preparedDevices.GetDevices(), nil
+	}
+
+	preparedDevices, err = s.prepareDevices(ctx, claim)
 	if err != nil {
 		return nil, fmt.Errorf("prepare devices failed: %w", err)
 	}

cmd/gpu-kubelet-plugin/device_state.go

@@ -150,32 +150,38 @@ func (s *DeviceState) Prepare(ctx context.Context, claim *resourceapi.ResourceCl
 
 	claimUID := string(claim.UID)
 
-	checkpoint, err := s.getCheckpoint()
-	if err != nil {
-		return nil, fmt.Errorf("unable to get checkpoint: %v", err)
-	}
+	var preparedDevices PreparedDevices
+	var alreadyPrepared bool
 
-	err = s.updateCheckpoint(func(checkpoint *Checkpoint) {
-		checkpoint.V2.PreparedClaims[claimUID] = PreparedClaim{
+	// Atomically check if prepare for this claim is already completed, otherwise mark as started.
+	err := s.updateCheckpoint(func(cp *Checkpoint) {
+		if pc, ok := cp.V2.PreparedClaims[claimUID]; ok &&
+			pc.CheckpointState == ClaimCheckpointStatePrepareCompleted {
+
+			alreadyPrepared = true
+			preparedDevices = pc.PreparedDevices
+			return
+		}
+
+		cp.V2.PreparedClaims[claimUID] = PreparedClaim{
 			CheckpointState: ClaimCheckpointStatePrepareStarted,
 			Status:          claim.Status,
 		}
+		klog.V(6).Infof("checkpoint updated for claim %v", claimUID)
 	})
 	if err != nil {
 		return nil, fmt.Errorf("unable to update checkpoint: %w", err)
 	}
-	klog.V(6).Infof("checkpoint updated for claim %v", claimUID)
 
-	preparedClaim, exists := checkpoint.V2.PreparedClaims[claimUID]
-	if exists && preparedClaim.CheckpointState == ClaimCheckpointStatePrepareCompleted {
+	if alreadyPrepared {
 		// 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
+		klog.V(4).Infof("skip prepare: claim %v found in checkpoint", claimUID)
+		return preparedDevices.GetDevices(), nil
 	}
 
-	preparedDevices, err := s.prepareDevices(ctx, claim)
+	preparedDevices, err = s.prepareDevices(ctx, claim)
 	if err != nil {
 		return nil, fmt.Errorf("prepare devices failed: %w", err)
 	}