add support for bouncing third-party gpu client pods during mig config apply

Signed-off-by: Tariq Ibrahim <tibrahim@nvidia.com>

Author: tariq1890

pkg/mig/reconfigure/reconfigure.go

@@ -27,6 +27,8 @@ import (
 	"time"
 
 	log "github.com/sirupsen/logrus"
+	corev1 "k8s.io/api/core/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/wait"
@@ -42,6 +44,7 @@ const (
 	dcgmExporterDeployLabel        = "nvidia.com/gpu.deploy.dcgm-exporter"
 	dcgmDeployLabel                = "nvidia.com/gpu.deploy.dcgm"
 	nvsmDeployLabel                = "nvidia.com/gpu.deploy.nvsm"
+	gpuClientDeployLabel           = "nvidia.com/gpu.deploy.client"
 
 	migStateSuccess = "success"
 	migStateFailed  = "failed"
@@ -93,6 +96,7 @@ type Reconfigure struct {
 	dcgmExporterDeployed  string
 	dcgmDeployed          string
 	nvsmDeployed          string
+	gpuClientDeployed     string
 	currentState          string
 	migModeChangeRequired bool
 	stoppedServices       []string
@@ -255,13 +259,16 @@ func (r *Reconfigure) getCurrentNodeLabels() error {
 	r.dcgmExporterDeployed = labels[dcgmExporterDeployLabel]
 	r.dcgmDeployed = labels[dcgmDeployLabel]
 	r.nvsmDeployed = labels[nvsmDeployLabel]
+	r.gpuClientDeployed = labels[gpuClientDeployLabel]
 	r.currentState = labels[migConfigStateLabel]
 
 	log.Infof("Current value of '%s=%s'\n", devicePluginDeployLabel, r.pluginDeployed)
 	log.Infof("Current value of '%s=%s'\n", gpuFeatureDiscoveryDeployLabel, r.gfdDeployed)
 	log.Infof("Current value of '%s=%s'\n", dcgmExporterDeployLabel, r.dcgmExporterDeployed)
 	log.Infof("Current value of '%s=%s'\n", dcgmDeployLabel, r.dcgmDeployed)
 	log.Infof("Current value of '%s=%s'\n", nvsmDeployLabel, r.nvsmDeployed)
+	log.Infof("Current value of '%s=%s'\n", gpuClientDeployLabel, r.gpuClientDeployed)
+
 	log.Infof("Current value of '%s=%s'\n", migConfigStateLabel, r.currentState)
 
 	return nil
@@ -370,6 +377,7 @@ func (r *Reconfigure) shutdownKubernetesGPUClients() error {
 		dcgmExporterDeployLabel:        r.maybeSetPaused(r.dcgmExporterDeployed),
 		dcgmDeployLabel:                r.maybeSetPaused(r.dcgmDeployed),
 		nvsmDeployLabel:                r.maybeSetPaused(r.nvsmDeployed),
+		gpuClientDeployLabel:           r.maybeSetPaused(r.gpuClientDeployed),
 	}
 
 	return r.setNodeLabels(labels)
@@ -399,6 +407,11 @@ func (r *Reconfigure) waitForPodsToBeDeleted() error {
 		return fmt.Errorf("dcgm pod did not shutdown: %w", err)
 	}
 
+	log.Infof("Waiting for any pods with parent nodeSelector key %s to shutdown", gpuClientDeployLabel)
+	if err := r.waitForPodsWithParentNodeSelector(gpuClientDeployLabel, timeout); err != nil {
+		return fmt.Errorf("third-party gpu client pods did not shutdown: %w", err)
+	}
+
 	log.Info("Removing the cuda-validator pod")
 	if err := r.deletePod("app=nvidia-cuda-validator"); err != nil {
 		return fmt.Errorf("failed to delete cuda-validator pod: %w", err)
@@ -507,6 +520,7 @@ func (r *Reconfigure) restartKubernetesGPUClients() error {
 		dcgmExporterDeployLabel:        r.maybeSetTrue(r.dcgmExporterDeployed),
 		dcgmDeployLabel:                r.maybeSetTrue(r.dcgmDeployed),
 		nvsmDeployLabel:                r.maybeSetTrue(r.nvsmDeployed),
+		gpuClientDeployLabel:           r.maybeSetTrue(r.gpuClientDeployed),
 	}
 
 	return r.setNodeLabels(labels)
@@ -538,6 +552,7 @@ func (r *Reconfigure) setState(state string) error {
 			dcgmExporterDeployLabel:        r.maybeSetTrue(r.dcgmExporterDeployed),
 			dcgmDeployLabel:                r.maybeSetTrue(r.dcgmDeployed),
 			nvsmDeployLabel:                r.maybeSetTrue(r.nvsmDeployed),
+			gpuClientDeployLabel:           r.maybeSetTrue(r.gpuClientDeployed),
 		}
 
 		if err := r.setNodeLabels(labels); err != nil {
@@ -621,6 +636,69 @@ func (r *Reconfigure) deletePod(labelSelector string) error {
 	return nil
 }
 
+// waitForPodsWithParentNodeSelector waits for all pods on the given node whose parent controller
+// has the specified key in its nodeSelector to terminate.
+//
+// Owner controller lookups are cached by UID across poll iterations. We don't expect the owners to actually change
+// during the lifetime of the wait, so caching the owners list optimises for kube API call volumes.
+func (r *Reconfigure) waitForPodsWithParentNodeSelector(nodeSelectorKey string, timeout time.Duration) error {
+	// ownerMatchCache maps owner UID → whether its nodeSelector contains nodeSelectorKey.
+	// Owner controllers are stable during the wait; only the pods themselves are terminating.
+	ownerMatchCache := make(map[types.UID]bool)
+
+	return wait.PollUntilContextTimeout(r.ctx, 5*time.Second, timeout, true, func(ctx context.Context) (bool, error) {
+		podList, err := r.clientset.CoreV1().Pods(corev1.NamespaceAll).List(ctx, metav1.ListOptions{
+			FieldSelector: "spec.nodeName=" + r.opts.NodeName,
+		})
+		if err != nil {
+			return false, fmt.Errorf("failed to list pods on node %s: %w", r.opts.NodeName, err)
+		}
+
+		matchCount := 0
+		for _, pod := range podList.Items {
+			for _, ownerRef := range pod.OwnerReferences {
+				matches, ok := ownerMatchCache[ownerRef.UID]
+				if !ok {
+					nodeSelector, err := r.getOwnerNodeSelector(ctx, pod.Namespace, ownerRef)
+					if err != nil {
+						log.Warnf("transient error checking parent nodeSelector for pod %s/%s, will retry: %v", pod.Namespace, pod.Name, err)
+						return false, nil
+					}
+					_, matches = nodeSelector[nodeSelectorKey]
+					ownerMatchCache[ownerRef.UID] = matches
+				}
+				if matches {
+					matchCount++
+					break
+				}
+			}
+		}
+
+		if matchCount > 0 {
+			log.Infof("Waiting for %d pod(s) with parent nodeSelector key %s to terminate", matchCount, nodeSelectorKey)
+		}
+		return matchCount == 0, nil
+	})
+}
+
+// getOwnerNodeSelector returns the pod-template nodeSelector for the given owner reference.
+// For ReplicaSets, it traces back to the owning Deployment when present.
+func (r *Reconfigure) getOwnerNodeSelector(ctx context.Context, namespace string, ownerRef metav1.OwnerReference) (map[string]string, error) {
+	switch ownerRef.Kind {
+	case "DaemonSet":
+		ds, err := r.clientset.AppsV1().DaemonSets(namespace).Get(ctx, ownerRef.Name, metav1.GetOptions{})
+		if apierrors.IsNotFound(err) {
+			return nil, nil
+		}
+		if err != nil {
+			return nil, fmt.Errorf("failed to get DaemonSet %s/%s: %w", namespace, ownerRef.Name, err)
+		}
+		return ds.Spec.Template.Spec.NodeSelector, nil
+	default:
+		return nil, nil
+	}
+}
+
 func (r *Reconfigure) runNvidiaSMI() error {
 	if r.opts.DriverRootCtrPath == r.opts.DevRootCtrPath {
 		cmd := exec.Command("chroot", r.opts.DriverRootCtrPath, "nvidia-smi")