From 8a5636b7941935f765a35bb12efeee32269a9055 Mon Sep 17 00:00:00 2001 From: Arjun Date: Sat, 12 Sep 2026 00:59:05 +0000 Subject: [PATCH] Ignore already cordoned nodes during driver upgrades Signed-off-by: Arjun --- cmd/driver-manager/main.go | 85 ++++- internal/kubernetes/client.go | 222 +++++++++++- internal/kubernetes/client_test.go | 549 +++++++++++++++++++++++++++-- 3 files changed, 811 insertions(+), 45 deletions(-) diff --git a/cmd/driver-manager/main.go b/cmd/driver-manager/main.go index be702193..cab6e93c 100644 --- a/cmd/driver-manager/main.go +++ b/cmd/driver-manager/main.go @@ -75,6 +75,13 @@ const ( nvidiaDRAValidatorDeployLabel = nvidiaDomainPrefix + "/" + "gpu.deploy.dra-validator" nvidiaDRADCGMExporterDeployLabel = nvidiaDomainPrefix + "/" + "gpu.deploy.dcgm-exporter-dra" nvidiaDRADCGMDeployLabel = nvidiaDomainPrefix + "/" + "gpu.deploy.dcgm-dra" + + // nvidiaDriverManagerCordonAnnotation marks a cordon applied by driver-manager, since + // Node.Spec.Unschedulable does not record who set it. + nvidiaDriverManagerCordonAnnotation = nvidiaDomainPrefix + "/" + "gpu-driver-manager.node-cordoned" + // nvidiaUpgradeControllerCordonAnnotation is the ownership marker proposed for the upgrade + // controller, which does not write it yet. + nvidiaUpgradeControllerCordonAnnotation = nvidiaDomainPrefix + "/" + "gpu-driver-upgrade-controller.node-cordoned" ) // Configuration holds all the configuration from environment variables @@ -293,6 +300,8 @@ func newDriverManager(ctx context.Context, cfg *config, components *componentSta func (dm *DriverManager) uninstallDriver() error { dm.log.Info("Starting driver uninstallation process") + dm.clearStaleCordonClaim() + // Check if driver is pre-installed on host if dm.isHostDriver() { dm.log.Info("NVIDIA GPU driver is already pre-installed on the node, disabling the containerized driver") @@ -369,6 +378,12 @@ func (dm *DriverManager) uninstallDriver() error { return nil } + // Wait out an external cordon before evicting anything, so the node keeps running + // untouched while it is in place. + if err := dm.waitForExternalCordonToClear(); err != nil { + return err + } + // Always evict all GPU operator components across a driver restart. The DRA // kubelet-plugin is the exception: it services NodeUnprepareResources for the // claim-holders evicted here (e.g. dra-validator), so it must outlive them and @@ -1225,6 +1240,61 @@ func (dm *DriverManager) isGPUPodEvictionEnabled() bool { return dm.config.enableGPUPodEviction } +// waitForExternalCordonToClear blocks while the node is cordoned by someone other than +// GPU Operator. A GPU-owned cordon is in-progress operator work and does not block. +func (dm *DriverManager) waitForExternalCordonToClear() error { + err := dm.kubeClient.WaitUntilNoExternalCordon(dm.config.nodeName, + nvidiaDriverManagerCordonAnnotation, nvidiaUpgradeControllerCordonAnnotation) + if err != nil { + return fmt.Errorf("failed while waiting for an external cordon on node %s to clear: %w", + dm.config.nodeName, err) + } + return nil +} + +// clearStaleCordonClaim is best effort and must not block the driver install. +func (dm *DriverManager) clearStaleCordonClaim() { + if err := dm.kubeClient.ClearStaleCordonClaim(dm.config.nodeName, nvidiaDriverManagerCordonAnnotation); err != nil { + dm.log.Warnf("Failed to clear stale cordon claim on node %s: %v", dm.config.nodeName, err) + } +} + +// cordonNode cordons the node and records driver-manager as its owner. An existing external +// cordon is waited out rather than taken over. +func (dm *DriverManager) cordonNode() error { + for { + if err := dm.waitForExternalCordonToClear(); err != nil { + return err + } + var acquired bool + err := retryOnAnyError(dm.ctx, dm.log, retryBackoff(defaultCordonRetries), + fmt.Sprintf("cordon node %s", dm.config.nodeName), func() error { + var err error + acquired, err = dm.kubeClient.AcquireNodeCordon(dm.config.nodeName, + nvidiaDriverManagerCordonAnnotation, nvidiaUpgradeControllerCordonAnnotation) + return err + }) + if err != nil { + return err + } + if acquired { + return nil + } + dm.log.Infof("Node %s was recordoned by an external actor before the claim could be recorded, waiting again", dm.config.nodeName) + } +} + +// uncordonNode releases the cordon only when driver-manager applied it; anyone else's +// cordon is left in place. +func (dm *DriverManager) uncordonNode() error { + return retryOnAnyError(dm.ctx, dm.log, retryBackoff(defaultCordonRetries), + fmt.Sprintf("uncordon node %s", dm.config.nodeName), + func() error { + return dm.kubeClient.ReleaseNodeCordon(dm.config.nodeName, + nvidiaDriverManagerCordonAnnotation, nvidiaUpgradeControllerCordonAnnotation) + }) +} + func (dm *DriverManager) nvDrainNode() error { dm.log.Infof("Draining node %s of any GPU pods...", dm.config.nodeName) drainOpts := kube.DrainOptions{ @@ -1250,21 +1320,6 @@ func (dm *DriverManager) isDriverAutoUpgradePolicyEnabled() bool { return false } -// cordonNode and uncordonNode wrap the kube client calls with this command's -// retry policy. The policy lives here rather than in the client so the client -// stays a thin wrapper over the Kubernetes API. -func (dm *DriverManager) cordonNode() error { - return retryOnAnyError(dm.ctx, dm.log, retryBackoff(defaultCordonRetries), - fmt.Sprintf("cordon node %s", dm.config.nodeName), - func() error { return dm.kubeClient.CordonNode(dm.config.nodeName) }) -} - -func (dm *DriverManager) uncordonNode() error { - return retryOnAnyError(dm.ctx, dm.log, retryBackoff(defaultCordonRetries), - fmt.Sprintf("uncordon node %s", dm.config.nodeName), - func() error { return dm.kubeClient.UncordonNode(dm.config.nodeName) }) -} - func (dm *DriverManager) cleanupOnFailure() { dm.log.Info("Performing cleanup on failure") policyEnabled := dm.isDriverAutoUpgradePolicyEnabled() diff --git a/internal/kubernetes/client.go b/internal/kubernetes/client.go index 87022c3f..9a05a553 100644 --- a/internal/kubernetes/client.go +++ b/internal/kubernetes/client.go @@ -50,8 +50,16 @@ const ( computeDomainDRADriverName = "compute-domain." + nvidiaDomainPrefix computeDomainDaemonConfigKind = "ComputeDomainDaemonConfig" + // Presence means a GPU component recorded this cordon; the value is the + // node's prior unschedulable state, not an ownership flag. nodeInitialUnschedulableAnnotationKey = nvidiaDomainPrefix + "/" + "driver-manager.node-initial-state.unschedulable" + // Signals the upgrade controller (k8s-operator-libs) publishes while a node + // is mid-upgrade. + gpuUpgradeInitialUnschedulableAnnotationKey = nvidiaDomainPrefix + "/" + "gpu-driver-upgrade.node-initial-state.unschedulable" + gpuUpgradeWaitForSafeLoadAnnotationKey = nvidiaDomainPrefix + "/" + "gpu-driver-upgrade.driver-wait-for-safe-load" + gpuUpgradeStateLabelKey = nvidiaDomainPrefix + "/" + "gpu-driver-upgrade-state" + kubeClientPollInterval = 5 * time.Second ) @@ -61,6 +69,9 @@ type Client struct { log *logrus.Logger clientset kubernetes.Interface + // pollInterval is the wait between retries for long-running watches such as + // WaitUntilNoExternalCordon. Zero means kubeClientPollInterval. + pollInterval time.Duration } // DrainOptions represents the option parameters that can passed to the drain.Helper struct @@ -150,6 +161,198 @@ func (c *Client) GetNodeAnnotationValue(nodeName, annotation string) (string, er return node.Annotations[annotation], nil } +func annotationTrue(annotations map[string]string, key string) bool { + return annotations[key] == "true" +} + +func annotationPresent(annotations map[string]string, key string) bool { + _, ok := annotations[key] + return ok +} + +func gpuUpgradeInProgress(labels map[string]string) bool { + switch labels[gpuUpgradeStateLabelKey] { + case "cordon-required", "wait-for-jobs-required", "pod-deletion-required", + "drain-required", "node-maintenance-required", "post-maintenance-required", + "pod-restart-required", "validation-required", "uncordon-required", + "upgrade-failed": + return true + default: + return false + } +} + +// gpuCordonClaimed reports whether a GPU Operator component owns the cordon. +// Initial-state keys are presence checks because their values are not ownership. +func gpuCordonClaimed(node *corev1.Node, claimAnnotation, peerClaimAnnotation string) bool { + if node == nil { + return false + } + annotations := node.Annotations + if annotationTrue(annotations, claimAnnotation) || annotationTrue(annotations, peerClaimAnnotation) { + return true + } + if annotationPresent(annotations, nodeInitialUnschedulableAnnotationKey) || + annotationPresent(annotations, gpuUpgradeInitialUnschedulableAnnotationKey) || + annotationPresent(annotations, gpuUpgradeWaitForSafeLoadAnnotationKey) { + return true + } + return gpuUpgradeInProgress(node.Labels) +} + +// gpuPeerOwnsCordon reports whether another GPU component still holds the cordon +// once this component drops its claim. +func gpuPeerOwnsCordon(node *corev1.Node, peerClaimAnnotation string) bool { + if node == nil { + return false + } + annotations := node.Annotations + if annotationTrue(annotations, peerClaimAnnotation) || + annotationPresent(annotations, gpuUpgradeInitialUnschedulableAnnotationKey) || + annotationPresent(annotations, gpuUpgradeWaitForSafeLoadAnnotationKey) { + return true + } + return gpuUpgradeInProgress(node.Labels) +} + +func initialUnschedulableRecording(node *corev1.Node) (string, bool) { + _, recorded := node.Annotations[nodeInitialUnschedulableAnnotationKey] + if recorded && node.Spec.Unschedulable { + return "", false + } + return strconv.FormatBool(node.Spec.Unschedulable), true +} + +// WaitUntilNoExternalCordon blocks until the node is schedulable or a GPU +// component has claimed the cordon, or until c.ctx is cancelled. +func (c *Client) WaitUntilNoExternalCordon(nodeName, claimAnnotation, peerClaimAnnotation string) error { + interval := kubeClientPollInterval + if c.pollInterval > 0 { + interval = c.pollInterval + } + loggedWaiting := false + return wait.PollUntilContextCancel(c.ctx, interval, true, func(ctx context.Context) (bool, error) { + node, err := c.clientset.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{}) + if err != nil { + c.log.Warnf("Failed to inspect cordon on node %s, retrying: %v", nodeName, err) + return false, nil + } + if !node.Spec.Unschedulable || gpuCordonClaimed(node, claimAnnotation, peerClaimAnnotation) { + if loggedWaiting { + c.log.Infof("External cordon on node %s has cleared, continuing driver rotation", nodeName) + } + return true, nil + } + if !loggedWaiting { + c.log.Infof("Node %s is already cordoned by an external actor, waiting until it is uncordoned or a GPU component claims the cordon", nodeName) + loggedWaiting = true + } else { + c.log.Infof("Still waiting for external cordon on node %s to be released", nodeName) + } + return false, nil + }) +} + +// AcquireNodeCordon atomically claims the cordon and marks the Node unschedulable. +// An externally owned cordon is left unchanged and returns false. +func (c *Client) AcquireNodeCordon(nodeName, claimAnnotation, peerClaimAnnotation string) (bool, error) { + acquired := false + err := retry.RetryOnConflict(retry.DefaultBackoff, func() error { + acquired = false + node, err := c.clientset.CoreV1().Nodes().Get(c.ctx, nodeName, metav1.GetOptions{}) + if err != nil { + return fmt.Errorf("failed to get node %s: %w", nodeName, err) + } + + if node.Spec.Unschedulable && !gpuCordonClaimed(node, claimAnnotation, peerClaimAnnotation) { + return nil + } + if node.Spec.Unschedulable && annotationTrue(node.Annotations, claimAnnotation) { + acquired = true + return nil + } + + annotations := map[string]any{ + claimAnnotation: "true", + } + if initialState, write := initialUnschedulableRecording(node); write { + annotations[nodeInitialUnschedulableAnnotationKey] = initialState + } + if err := c.patchNodeSchedulingState(node, true, annotations); err != nil { + return err + } + acquired = true + return nil + }) + if err != nil { + return false, fmt.Errorf("failed to acquire cordon claim on node %s: %w", nodeName, err) + } + return acquired, nil +} + +// ReleaseNodeCordon atomically removes this component's claim, uncordoning the +// Node only when no other GPU component still owns the cordon. +func (c *Client) ReleaseNodeCordon(nodeName, claimAnnotation, peerClaimAnnotation string) error { + err := retry.RetryOnConflict(retry.DefaultBackoff, func() error { + node, err := c.clientset.CoreV1().Nodes().Get(c.ctx, nodeName, metav1.GetOptions{}) + if err != nil { + return fmt.Errorf("failed to get node %s: %w", nodeName, err) + } + if !annotationTrue(node.Annotations, claimAnnotation) { + return nil + } + if !node.Spec.Unschedulable { + return c.clearStaleClaim(node, claimAnnotation) + } + + annotations := map[string]any{ + claimAnnotation: nil, + } + unschedulable := gpuPeerOwnsCordon(node, peerClaimAnnotation) + if !unschedulable { + annotations[nodeInitialUnschedulableAnnotationKey] = nil + } + return c.patchNodeSchedulingState(node, unschedulable, annotations) + }) + if err != nil { + return fmt.Errorf("failed to release cordon claim on node %s: %w", nodeName, err) + } + return nil +} + +// ClearStaleCordonClaim removes a claim left on a schedulable Node, since +// `kubectl uncordon` does not clear annotations. A cordoned Node is unchanged. +func (c *Client) ClearStaleCordonClaim(nodeName, claimAnnotation string) error { + err := retry.RetryOnConflict(retry.DefaultBackoff, func() error { + node, err := c.clientset.CoreV1().Nodes().Get(c.ctx, nodeName, metav1.GetOptions{}) + if err != nil { + return fmt.Errorf("failed to get node %s: %w", nodeName, err) + } + if node.Spec.Unschedulable { + return nil + } + if !annotationPresent(node.Annotations, claimAnnotation) && + !annotationPresent(node.Annotations, nodeInitialUnschedulableAnnotationKey) { + return nil + } + return c.clearStaleClaim(node, claimAnnotation) + }) + if err != nil { + return fmt.Errorf("failed to clear stale cordon claim on node %s: %w", nodeName, err) + } + return nil +} + +// clearStaleClaim drops the claim from a schedulable Node. The resourceVersion +// precondition makes a concurrent cordon conflict instead of being undone. +func (c *Client) clearStaleClaim(node *corev1.Node, claimAnnotation string) error { + c.log.Infof("Node %s is schedulable but carries a cordon claim from an earlier rotation, clearing it", node.Name) + return c.patchNodeSchedulingState(node, false, map[string]any{ + claimAnnotation: nil, + nodeInitialUnschedulableAnnotationKey: nil, + }) +} + // CordonNode atomically records a Node's initial schedulable state and marks it // Unschedulable. An existing recording on an Unschedulable Node is retained so // that restarting driver-manager does not lose the initial state. @@ -159,14 +362,15 @@ func (c *Client) CordonNode(nodeName string) error { return fmt.Errorf("failed to get node %s: %w", nodeName, err) } - _, recorded := node.Annotations[nodeInitialUnschedulableAnnotationKey] - if node.Spec.Unschedulable && recorded { + initialState, write := initialUnschedulableRecording(node) + if node.Spec.Unschedulable && !write { return nil } - initialState := strconv.FormatBool(node.Spec.Unschedulable) c.log.Infof("Cordoning node %s and recording initial state in annotation %s=%s", nodeName, nodeInitialUnschedulableAnnotationKey, initialState) - return c.patchNodeSchedulingState(node, true, initialState) + return c.patchNodeSchedulingState(node, true, map[string]any{ + nodeInitialUnschedulableAnnotationKey: initialState, + }) } // UncordonNode atomically marks a Node as schedulable and removes the initial @@ -184,16 +388,16 @@ func (c *Client) UncordonNode(nodeName string) error { } c.log.Infof("Uncordoning node %s", nodeName) - return c.patchNodeSchedulingState(node, false, nil) + return c.patchNodeSchedulingState(node, false, map[string]any{ + nodeInitialUnschedulableAnnotationKey: nil, + }) } -func (c *Client) patchNodeSchedulingState(node *corev1.Node, unschedulable bool, annotationValue any) error { +func (c *Client) patchNodeSchedulingState(node *corev1.Node, unschedulable bool, annotations map[string]any) error { patch := map[string]any{ "metadata": map[string]any{ "resourceVersion": node.ResourceVersion, - "annotations": map[string]any{ - nodeInitialUnschedulableAnnotationKey: annotationValue, - }, + "annotations": annotations, }, "spec": map[string]any{ "unschedulable": unschedulable, diff --git a/internal/kubernetes/client_test.go b/internal/kubernetes/client_test.go index 406c7fb6..9bd94cae 100644 --- a/internal/kubernetes/client_test.go +++ b/internal/kubernetes/client_test.go @@ -53,6 +53,8 @@ type flakyNodeAPI struct { requests int patches int unschedulable bool + annotations map[string]string + labels map[string]string } func (f *flakyNodeAPI) ServeHTTP(w http.ResponseWriter, r *http.Request) { @@ -83,26 +85,56 @@ func (f *flakyNodeAPI) ServeHTTP(w http.ResponseWriter, r *http.Request) { http.Error(w, err.Error(), http.StatusBadRequest) return } - // A strategic merge patch clears spec.unschedulable with a null rather - // than setting it to false, so key presence is what matters here. - var patch map[string]any - if err := json.Unmarshal(body, &patch); err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - if spec, ok := patch["spec"].(map[string]any); ok { - if value, present := spec["unschedulable"]; present { - cordoned, _ := value.(bool) - f.unschedulable = cordoned + if r.Method == http.MethodPut { + var node corev1.Node + if err := json.Unmarshal(body, &node); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + f.unschedulable = node.Spec.Unschedulable + f.annotations = node.Annotations + } else { + // A strategic merge patch clears spec.unschedulable with a null rather + // than setting it to false, so key presence is what matters here. + var patch map[string]any + if err := json.Unmarshal(body, &patch); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + if spec, ok := patch["spec"].(map[string]any); ok { + if value, present := spec["unschedulable"]; present { + cordoned, _ := value.(bool) + f.unschedulable = cordoned + } + } + if meta, ok := patch["metadata"].(map[string]any); ok { + if ann, ok := meta["annotations"].(map[string]any); ok { + if f.annotations == nil { + f.annotations = make(map[string]string) + } + for key, value := range ann { + if value == nil { + delete(f.annotations, key) + continue + } + if s, ok := value.(string); ok { + f.annotations[key] = s + } + } + } } } } w.Header().Set("Content-Type", "application/json") node := &corev1.Node{ - TypeMeta: metav1.TypeMeta{Kind: "Node", APIVersion: "v1"}, - ObjectMeta: metav1.ObjectMeta{Name: f.nodeName}, - Spec: corev1.NodeSpec{Unschedulable: f.unschedulable}, + TypeMeta: metav1.TypeMeta{Kind: "Node", APIVersion: "v1"}, + ObjectMeta: metav1.ObjectMeta{ + Name: f.nodeName, + Annotations: f.annotations, + Labels: f.labels, + }, + Spec: corev1.NodeSpec{Unschedulable: f.unschedulable}, } if err := json.NewEncoder(w).Encode(node); err != nil { f.requests-- // the request never completed, do not count it @@ -121,6 +153,42 @@ func (f *flakyNodeAPI) cordoned() bool { return f.unschedulable } +func (f *flakyNodeAPI) annotation(key string) string { + f.mu.Lock() + defer f.mu.Unlock() + return f.annotations[key] +} + +func (f *flakyNodeAPI) patchCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.patches +} + +func (f *flakyNodeAPI) setUnschedulable(unschedulable bool) { + f.mu.Lock() + defer f.mu.Unlock() + f.unschedulable = unschedulable +} + +func (f *flakyNodeAPI) setAnnotation(key, value string) { + f.mu.Lock() + defer f.mu.Unlock() + if f.annotations == nil { + f.annotations = make(map[string]string) + } + f.annotations[key] = value +} + +func (f *flakyNodeAPI) setLabel(key, value string) { + f.mu.Lock() + defer f.mu.Unlock() + if f.labels == nil { + f.labels = make(map[string]string) + } + f.labels[key] = value +} + func discardLogger() *logrus.Logger { log := logrus.New() log.SetOutput(io.Discard) @@ -133,18 +201,17 @@ func newHTTPTestClient(t *testing.T, api *flakyNodeAPI) *Client { server := httptest.NewServer(api) t.Cleanup(server.Close) - clientset, err := kubernetes.NewForConfig(&rest.Config{Host: server.URL}) + clientset, err := kubernetes.NewForConfig(&rest.Config{ + Host: server.URL, + ContentConfig: rest.ContentConfig{ + ContentType: runtime.ContentTypeJSON, + }, + }) require.NoError(t, err) return &Client{ctx: context.Background(), log: discardLogger(), clientset: clientset} } -func (f *flakyNodeAPI) patchCount() int { - f.mu.Lock() - defer f.mu.Unlock() - return f.patches -} - // The client makes one attempt per call: the retry policy lives in the // driver-manager command, not here. func TestCordonNodeCordonsTheNode(t *testing.T) { @@ -961,3 +1028,443 @@ func TestDeleteOrEvictPodsIgnoring(t *testing.T) { }) } } + +const ( + testOwnerClaim = "example.com/owner" + testPeerClaim = "example.com/peer" +) + +func TestAcquireNodeCordon(t *testing.T) { + tests := []struct { + name string + unschedulable bool + annotations map[string]string + labels map[string]string + wantAcquired bool + wantCordoned bool + wantOwner string + wantInitialState string + wantPatches int + }{ + { + name: "claims a schedulable node", + wantAcquired: true, + wantCordoned: true, + wantOwner: "true", + wantInitialState: "false", + wantPatches: 1, + }, + { + name: "leaves an external cordon unclaimed", + unschedulable: true, + wantAcquired: false, + wantCordoned: true, + wantPatches: 0, + }, + { + name: "claims a node already recorded by driver-manager", + unschedulable: true, + annotations: map[string]string{nodeInitialUnschedulableAnnotationKey: "false"}, + wantAcquired: true, + wantCordoned: true, + wantOwner: "true", + wantInitialState: "false", + wantPatches: 1, + }, + { + name: "claims a node already claimed by the peer", + unschedulable: true, + annotations: map[string]string{testPeerClaim: "true"}, + wantAcquired: true, + wantCordoned: true, + wantOwner: "true", + wantInitialState: "true", + wantPatches: 1, + }, + { + name: "is a no-op when this component already claimed the cordon", + unschedulable: true, + annotations: map[string]string{ + testOwnerClaim: "true", + nodeInitialUnschedulableAnnotationKey: "false", + }, + wantAcquired: true, + wantCordoned: true, + wantOwner: "true", + wantInitialState: "false", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + api := &flakyNodeAPI{ + nodeName: "gpu-node", + unschedulable: tt.unschedulable, + annotations: tt.annotations, + labels: tt.labels, + } + c := newHTTPTestClient(t, api) + + acquired, err := c.AcquireNodeCordon("gpu-node", testOwnerClaim, testPeerClaim) + require.NoError(t, err) + require.Equal(t, tt.wantAcquired, acquired) + require.Equal(t, tt.wantCordoned, api.cordoned()) + require.Equal(t, tt.wantOwner, api.annotation(testOwnerClaim)) + require.Equal(t, tt.wantInitialState, api.annotation(nodeInitialUnschedulableAnnotationKey)) + require.Equal(t, tt.wantPatches, api.patchCount()) + }) + } +} + +func TestWaitUntilNoExternalCordonReturnsImmediately(t *testing.T) { + tests := []struct { + name string + unschedulable bool + annotations map[string]string + labels map[string]string + }{ + {name: "when the node is schedulable"}, + { + name: "when a GPU component has claimed the cordon", + unschedulable: true, + annotations: map[string]string{testOwnerClaim: "true"}, + }, + { + name: "when driver-manager recorded initial state", + unschedulable: true, + annotations: map[string]string{nodeInitialUnschedulableAnnotationKey: "false"}, + }, + { + name: "when the upgrade controller recorded initial state", + unschedulable: true, + annotations: map[string]string{gpuUpgradeInitialUnschedulableAnnotationKey: "false"}, + }, + { + name: "when the upgrade controller is mid-upgrade", + unschedulable: true, + labels: map[string]string{gpuUpgradeStateLabelKey: "pod-restart-required"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + api := &flakyNodeAPI{ + nodeName: "gpu-node", + unschedulable: tt.unschedulable, + annotations: tt.annotations, + labels: tt.labels, + } + c := newHTTPTestClient(t, api) + c.pollInterval = time.Millisecond + + require.NoError(t, c.WaitUntilNoExternalCordon("gpu-node", testOwnerClaim, testPeerClaim)) + }) + } +} + +func TestWaitUntilNoExternalCordonReturnsOnceCleared(t *testing.T) { + tests := []struct { + name string + mutate func(*flakyNodeAPI) + }{ + { + name: "when the node is uncordoned", + mutate: func(api *flakyNodeAPI) { api.setUnschedulable(false) }, + }, + { + name: "when a GPU claim appears", + mutate: func(api *flakyNodeAPI) { api.setAnnotation(testPeerClaim, "true") }, + }, + { + name: "when the upgrade controller records initial state", + mutate: func(api *flakyNodeAPI) { api.setAnnotation(gpuUpgradeInitialUnschedulableAnnotationKey, "false") }, + }, + { + name: "when the upgrade controller enters an in-progress state", + mutate: func(api *flakyNodeAPI) { api.setLabel(gpuUpgradeStateLabelKey, "pod-restart-required") }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + api := &flakyNodeAPI{nodeName: "gpu-node", unschedulable: true} + c := newHTTPTestClient(t, api) + c.pollInterval = time.Millisecond + + done := make(chan error, 1) + go func() { + done <- c.WaitUntilNoExternalCordon("gpu-node", testOwnerClaim, testPeerClaim) + }() + + time.Sleep(20 * time.Millisecond) + tt.mutate(api) + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for external cordon to clear") + } + }) + } +} + +func TestWaitUntilNoExternalCordonStopsWhenContextCancelled(t *testing.T) { + api := &flakyNodeAPI{nodeName: "gpu-node", unschedulable: true} + c := newHTTPTestClient(t, api) + c.pollInterval = time.Millisecond + + ctx, cancel := context.WithCancel(context.Background()) + c.ctx = ctx + + done := make(chan error, 1) + go func() { + done <- c.WaitUntilNoExternalCordon("gpu-node", testOwnerClaim, testPeerClaim) + }() + + cancel() + + select { + case err := <-done: + require.Error(t, err) + require.ErrorIs(t, err, context.Canceled) + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for cancelled wait to return") + } +} + +func TestWaitUntilNoExternalCordonRetriesGetFailures(t *testing.T) { + api := &flakyNodeAPI{nodeName: "gpu-node", getFailures: 2} + c := newHTTPTestClient(t, api) + c.pollInterval = time.Millisecond + + require.NoError(t, c.WaitUntilNoExternalCordon("gpu-node", testOwnerClaim, testPeerClaim)) +} + +func TestReleaseNodeCordonPreservesPeerClaim(t *testing.T) { + api := &flakyNodeAPI{ + nodeName: "gpu-node", + unschedulable: true, + annotations: map[string]string{ + testOwnerClaim: "true", + testPeerClaim: "true", + nodeInitialUnschedulableAnnotationKey: "false", + }, + } + c := newHTTPTestClient(t, api) + + require.NoError(t, c.ReleaseNodeCordon("gpu-node", testOwnerClaim, testPeerClaim)) + require.True(t, api.cordoned()) + require.Empty(t, api.annotation(testOwnerClaim)) + require.Equal(t, "true", api.annotation(testPeerClaim)) + require.Equal(t, "false", api.annotation(nodeInitialUnschedulableAnnotationKey)) +} + +func TestReleaseNodeCordonUncordonsWhenSoleOwner(t *testing.T) { + api := &flakyNodeAPI{ + nodeName: "gpu-node", + unschedulable: true, + annotations: map[string]string{ + testOwnerClaim: "true", + nodeInitialUnschedulableAnnotationKey: "false", + }, + } + c := newHTTPTestClient(t, api) + + require.NoError(t, c.ReleaseNodeCordon("gpu-node", testOwnerClaim, testPeerClaim)) + require.False(t, api.cordoned()) + require.Empty(t, api.annotation(testOwnerClaim)) + require.Empty(t, api.annotation(nodeInitialUnschedulableAnnotationKey)) +} + +func TestReleaseNodeCordonPreservesUpgradeControllerOwnership(t *testing.T) { + api := &flakyNodeAPI{ + nodeName: "gpu-node", + unschedulable: true, + annotations: map[string]string{ + testOwnerClaim: "true", + gpuUpgradeInitialUnschedulableAnnotationKey: "false", + }, + } + c := newHTTPTestClient(t, api) + + require.NoError(t, c.ReleaseNodeCordon("gpu-node", testOwnerClaim, testPeerClaim)) + require.True(t, api.cordoned()) + require.Empty(t, api.annotation(testOwnerClaim)) + require.Equal(t, "false", api.annotation(gpuUpgradeInitialUnschedulableAnnotationKey)) +} + +func TestReleaseNodeCordonOnSchedulableNodeOnlyDropsStaleClaim(t *testing.T) { + api := &flakyNodeAPI{ + nodeName: "gpu-node", + annotations: map[string]string{ + testOwnerClaim: "true", + testPeerClaim: "true", + nodeInitialUnschedulableAnnotationKey: "false", + }, + } + c := newHTTPTestClient(t, api) + + require.NoError(t, c.ReleaseNodeCordon("gpu-node", testOwnerClaim, testPeerClaim)) + require.False(t, api.cordoned()) + require.Empty(t, api.annotation(testOwnerClaim)) + require.Empty(t, api.annotation(nodeInitialUnschedulableAnnotationKey)) + require.Equal(t, "true", api.annotation(testPeerClaim)) +} + +func TestClearStaleCordonClaim(t *testing.T) { + tests := []struct { + name string + unschedulable bool + annotations map[string]string + wantOwner string + wantInitialState string + wantPatches int + }{ + { + name: "clears a claim left on a schedulable node", + annotations: map[string]string{ + testOwnerClaim: "true", + nodeInitialUnschedulableAnnotationKey: "false", + }, + wantPatches: 1, + }, + { + name: "clears an initial-state recording left on a schedulable node", + annotations: map[string]string{nodeInitialUnschedulableAnnotationKey: "false"}, + wantPatches: 1, + }, + { + name: "is a no-op on a schedulable node without a claim", + }, + { + name: "keeps the claim on a cordoned node", + unschedulable: true, + annotations: map[string]string{ + testOwnerClaim: "true", + nodeInitialUnschedulableAnnotationKey: "false", + }, + wantOwner: "true", + wantInitialState: "false", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + api := &flakyNodeAPI{ + nodeName: "gpu-node", + unschedulable: tt.unschedulable, + annotations: tt.annotations, + } + c := newHTTPTestClient(t, api) + + require.NoError(t, c.ClearStaleCordonClaim("gpu-node", testOwnerClaim)) + require.Equal(t, tt.unschedulable, api.cordoned()) + require.Equal(t, tt.wantOwner, api.annotation(testOwnerClaim)) + require.Equal(t, tt.wantInitialState, api.annotation(nodeInitialUnschedulableAnnotationKey)) + require.Equal(t, tt.wantPatches, api.patchCount()) + }) + } +} + +// A claim left by an interrupted rotation must not let driver-manager release a +// later administrator cordon. +func TestStaleClaimDoesNotReleaseExternalCordon(t *testing.T) { + api := &flakyNodeAPI{ + nodeName: "gpu-node", + annotations: map[string]string{ + testOwnerClaim: "true", + nodeInitialUnschedulableAnnotationKey: "false", + }, + } + c := newHTTPTestClient(t, api) + c.pollInterval = time.Millisecond + + require.NoError(t, c.ClearStaleCordonClaim("gpu-node", testOwnerClaim)) + api.setUnschedulable(true) + + acquired, err := c.AcquireNodeCordon("gpu-node", testOwnerClaim, testPeerClaim) + require.NoError(t, err) + require.False(t, acquired) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + c.ctx = ctx + require.ErrorIs(t, c.WaitUntilNoExternalCordon("gpu-node", testOwnerClaim, testPeerClaim), context.DeadlineExceeded) + + c.ctx = context.Background() + require.NoError(t, c.ReleaseNodeCordon("gpu-node", testOwnerClaim, testPeerClaim)) + require.True(t, api.cordoned()) +} + +func TestAcquireNodeCordonPatchIsAtomic(t *testing.T) { + client, clientset := newTestClient(newTestNode(false, nil)) + + acquired, err := client.AcquireNodeCordon(testNodeName, testOwnerClaim, testPeerClaim) + require.NoError(t, err) + require.True(t, acquired) + + patches := getNodePatches(t, clientset) + require.Len(t, patches, 1) + require.Equal(t, true, patches[0]["spec"].(map[string]any)["unschedulable"]) + metadata := patches[0]["metadata"].(map[string]any) + require.Equal(t, "1", metadata["resourceVersion"]) + annotations := metadata["annotations"].(map[string]any) + require.Equal(t, "true", annotations[testOwnerClaim]) + require.Equal(t, "false", annotations[nodeInitialUnschedulableAnnotationKey]) + + node := getTestNode(t, clientset) + require.True(t, node.Spec.Unschedulable) + require.Equal(t, "true", node.Annotations[testOwnerClaim]) + require.Equal(t, "false", node.Annotations[nodeInitialUnschedulableAnnotationKey]) +} + +func TestGPUCordonClaimed(t *testing.T) { + tests := []struct { + name string + node *corev1.Node + wantClaimed bool + wantPeer bool + }{ + {name: "nil node"}, + { + name: "owner claim", + node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{testOwnerClaim: "true"}}}, + wantClaimed: true, + }, + { + name: "peer claim", + node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{testPeerClaim: "true"}}}, + wantClaimed: true, + wantPeer: true, + }, + { + name: "driver-manager initial state is a claim but not a remaining peer", + node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{nodeInitialUnschedulableAnnotationKey: "false"}}}, + wantClaimed: true, + }, + { + name: "upgrade controller initial state", + node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{gpuUpgradeInitialUnschedulableAnnotationKey: "false"}}}, + wantClaimed: true, + wantPeer: true, + }, + { + name: "upgrade controller mid-upgrade label", + node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{gpuUpgradeStateLabelKey: "pod-restart-required"}}}, + wantClaimed: true, + wantPeer: true, + }, + { + name: "upgrade-done is not a claim", + node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{gpuUpgradeStateLabelKey: "upgrade-done"}}}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.wantClaimed, gpuCordonClaimed(tt.node, testOwnerClaim, testPeerClaim)) + require.Equal(t, tt.wantPeer, gpuPeerOwnsCordon(tt.node, testPeerClaim)) + }) + } +}