Merge pull request #651 from jgehrcke/jp/issue-694

CD daemon: coordinate CD updates on shutdown via mutation cache

Author: jgehrcke

cmd/compute-domain-daemon/computedomain.go

@@ -78,7 +78,6 @@ func NewComputeDomainManager(config *ManagerConfig) *ComputeDomainManager {
 		previousNodes:    []*nvapi.ComputeDomainNode{},
 		updatedNodesChan: make(chan []*nvapi.ComputeDomainNode),
 	}
-	m.podManager = NewPodManager(config, m.Get)
 
 	return m
 }
@@ -112,6 +111,8 @@ func (m *ComputeDomainManager) Start(ctx context.Context) (rerr error) {
 		true,
 	)
 
+	m.podManager = NewPodManager(m.config, m.Get, m.mutationCache)
+
 	_, err = m.informer.AddEventHandler(cache.ResourceEventHandlerFuncs{
 		AddFunc: func(obj any) {
 			m.config.workQueue.Enqueue(obj, m.onAddOrUpdate)
@@ -191,8 +192,9 @@ func (m *ComputeDomainManager) onAddOrUpdate(ctx context.Context, obj any) error
 		return fmt.Errorf("failed to cast to ComputeDomain")
 	}
 
-	// Get the latest ComputeDomain object from the informer cache since we
-	// plan to update it later and always *must* have the latest version.
+	// Get the latest ComputeDomain object from the mutation cache (backed by
+	// the informer cache) since we plan to update it later and always *must*
+	// have the latest version.
 	cd, err := m.Get(string(o.GetUID()))
 	if err != nil {
 		return fmt.Errorf("error getting latest ComputeDomain: %w", err)
@@ -215,8 +217,9 @@ func (m *ComputeDomainManager) onAddOrUpdate(ctx context.Context, obj any) error
 	return nil
 }
 
-// UpdateComputeDomainNodeInfo updates the Nodes field in the ComputeDomain
-// with info about the ComputeDomain daemon running on this node.
+// UpdateComputeDomainNodeInfo updates the Nodes field in the ComputeDomain with
+// info about the ComputeDomain daemon running on this node. Upon success, it
+// reflects the mutation in `m.mutationCache`.
 func (m *ComputeDomainManager) UpdateComputeDomainNodeInfo(ctx context.Context, cd *nvapi.ComputeDomain) (rerr error) {
 	var nodeInfo *nvapi.ComputeDomainNode
 
@@ -268,12 +271,12 @@ func (m *ComputeDomainManager) UpdateComputeDomainNodeInfo(ctx context.Context,
 		newCD.Status.Status = nvapi.ComputeDomainStatusNotReady
 	}
 
-	// Update the status
-	if _, err := m.config.clientsets.Nvidia.ResourceV1beta1().ComputeDomains(newCD.Namespace).UpdateStatus(ctx, newCD, metav1.UpdateOptions{}); err != nil {
-		return fmt.Errorf("error updating nodes in ComputeDomain status: %w", err)
+	// Update status and (upon success) store the latest version of the object
+	// (as returned by the API server) in the mutation cache.
+	newCD, err := m.config.clientsets.Nvidia.ResourceV1beta1().ComputeDomains(newCD.Namespace).UpdateStatus(ctx, newCD, metav1.UpdateOptions{})
+	if err != nil {
+		return fmt.Errorf("error updating ComputeDomain status: %w", err)
 	}
-
-	// Add the updated ComputeDomain to the mutation cache
 	m.mutationCache.Mutation(newCD)
 
 	return nil
@@ -360,15 +363,13 @@ func (m *ComputeDomainManager) GetNodesUpdateChan() chan []*nvapi.ComputeDomainN
 
 // removeNodeFromComputeDomain removes the current node's entry from the ComputeDomain status.
 func (m *ComputeDomainManager) removeNodeFromComputeDomain(ctx context.Context) error {
-	objs := m.informer.GetIndexer().List()
-	if len(objs) == 0 {
-		klog.Infof("No ComputeDomain objects found in informer cache during cleanup")
-		return nil
+	cd, err := m.Get(m.config.computeDomainUUID)
+	if err != nil {
+		return fmt.Errorf("error getting ComputeDomain from mutation cache: %w", err)
 	}
-
-	cd, ok := objs[0].(*nvapi.ComputeDomain)
-	if !ok {
-		return fmt.Errorf("failed to cast object to ComputeDomain")
+	if cd == nil {
+		klog.Infof("No ComputeDomain object found in mutation cache during cleanup")
+		return nil
 	}
 
 	newCD := cd.DeepCopy()
@@ -391,10 +392,13 @@ func (m *ComputeDomainManager) removeNodeFromComputeDomain(ctx context.Context)
 		newCD.Status.Status = nvapi.ComputeDomainStatusNotReady
 	}
 
+	// Update status and (upon success) store the latest version of the object
+	// (as returned by the API server) in the mutation cache.
 	newCD.Status.Nodes = updatedNodes
 	if _, err := m.config.clientsets.Nvidia.ResourceV1beta1().ComputeDomains(newCD.Namespace).UpdateStatus(ctx, newCD, metav1.UpdateOptions{}); err != nil {
 		return fmt.Errorf("error removing node from ComputeDomain status: %w", err)
 	}
+	m.mutationCache.Mutation(newCD)
 
 	klog.Infof("Successfully removed node with IP %s from ComputeDomain %s/%s", m.config.podIP, newCD.Namespace, newCD.Name)
 	return nil

cmd/compute-domain-daemon/podmanager.go

@@ -33,16 +33,17 @@ import (
 
 // PodManager watches for changes to its own pod and updates the CD status accordingly.
 type PodManager struct {
-	config           *ManagerConfig
-	getComputeDomain GetComputeDomainFunc
-	waitGroup        sync.WaitGroup
-	cancelContext    context.CancelFunc
-	podInformer      cache.SharedIndexInformer
-	informerFactory  informers.SharedInformerFactory
+	config                     *ManagerConfig
+	getComputeDomain           GetComputeDomainFunc
+	waitGroup                  sync.WaitGroup
+	cancelContext              context.CancelFunc
+	podInformer                cache.SharedIndexInformer
+	informerFactory            informers.SharedInformerFactory
+	computeDomainMutationCache cache.MutationCache
 }
 
 // NewPodManager creates a new PodManager instance.
-func NewPodManager(config *ManagerConfig, getComputeDomain GetComputeDomainFunc) *PodManager {
+func NewPodManager(config *ManagerConfig, getComputeDomain GetComputeDomainFunc, mutationCache cache.MutationCache) *PodManager {
 	informerFactory := informers.NewSharedInformerFactoryWithOptions(
 		config.clientsets.Core,
 		informerResyncPeriod,
@@ -55,10 +56,11 @@ func NewPodManager(config *ManagerConfig, getComputeDomain GetComputeDomainFunc)
 	podInformer := informerFactory.Core().V1().Pods().Informer()
 
 	p := &PodManager{
-		config:           config,
-		getComputeDomain: getComputeDomain,
-		podInformer:      podInformer,
-		informerFactory:  informerFactory,
+		config:                     config,
+		getComputeDomain:           getComputeDomain,
+		podInformer:                podInformer,
+		informerFactory:            informerFactory,
+		computeDomainMutationCache: mutationCache,
 	}
 
 	return p
@@ -111,6 +113,7 @@ func (pm *PodManager) Stop() error {
 	}
 
 	pm.waitGroup.Wait()
+	klog.Infof("Terminating: pod manager")
 	return nil
 }
 
@@ -186,10 +189,14 @@ func (pm *PodManager) updateNodeStatus(ctx context.Context, status string) error
 	// Update the node status to the new value
 	node.Status = status
 
-	// Update the CD status
-	if _, err := pm.config.clientsets.Nvidia.ResourceV1beta1().ComputeDomains(newCD.Namespace).UpdateStatus(ctx, newCD, metav1.UpdateOptions{}); err != nil {
+	// Update the CD status. Store the CD object returned by the API server in
+	// the mutationCache (it now has a newer ResourceVersion than `newCD` and
+	// hence any further mutations should be performed based on that).
+	newCD, err = pm.config.clientsets.Nvidia.ResourceV1beta1().ComputeDomains(newCD.Namespace).UpdateStatus(ctx, newCD, metav1.UpdateOptions{})
+	if err != nil {
 		return fmt.Errorf("error updating node status in ComputeDomain: %w", err)
 	}
+	pm.computeDomainMutationCache.Mutation(newCD)
 
 	klog.Infof("Successfully updated node %s status to %s", pm.config.nodeName, status)
 	return nil

tests/bats/tests.bats

@@ -143,6 +143,49 @@ log_objects() {
   kubectl delete -f demo/specs/imex/channel-injection-all.yaml
 }
 
+@test "CD daemon shutdown: confirm CD status cleanup" {
+  log_objects
+
+  kubectl apply -f demo/specs/imex/channel-injection.yaml
+  kubectl wait --for=condition=READY pods imex-channel-injection --timeout=100s
+  run kubectl logs imex-channel-injection
+  assert_output --partial "channel0"
+
+  local LOGPATH="${BATS_TEST_TMPDIR}/cd-daemon.log"
+  local PNAME
+  PNAME=$( \
+    kubectl get pods -n nvidia-dra-driver-gpu | \
+    grep imex-channel-injection | \
+    awk '{print $1}'
+  )
+
+  # Expect `nodes` key to be present in CD status.
+  run bats_pipe kubectl get computedomain imex-channel-injection -o json \| jq '.status'
+  assert_output --partial 'nodes'
+
+  echo "attach background log follower to daemon pod: $PNAME"
+  kubectl logs -n nvidia-dra-driver-gpu --follow "$PNAME" > "$LOGPATH" 2>&1 &
+  kubectl delete pods imex-channel-injection
+
+  # Note: the log follower child process terminates when the pod terminates.
+  kubectl wait --for=delete pods imex-channel-injection --timeout=10s
+
+  # Expect `nodes` key to not be be present (single-node CD).
+  run bats_pipe kubectl get computedomain imex-channel-injection -o json \| jq '.status'
+  refute_output --partial 'nodes'
+
+  # Inspect CD daemon log, dump tail for easier debug-on-failure.
+  cat "$LOGPATH" | tail -n 50
+
+  # Explicitly confirm cleanup-on-shutdown behavior by inspecting CD log.
+  cat "$LOGPATH" | grep -e "Successfully updated node .* status to NotReady"
+  cat "$LOGPATH" | grep "Successfully removed node" | \
+    grep "from ComputeDomain default/imex-channel-injection"
+
+  # Delete CD.
+  kubectl delete computedomain imex-channel-injection
+}
+
 @test "NodePrepareResources: catch unknown field in opaque cfg in ResourceClaim" {
   log_objects