@@ -25,6 +25,7 @@ import (
2525 "time"
2626
2727 "github.com/sirupsen/logrus"
28+ jsonpatch "gopkg.in/evanphx/json-patch.v4"
2829 corev1 "k8s.io/api/core/v1"
2930 metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
3031 "k8s.io/apimachinery/pkg/labels"
@@ -108,22 +109,7 @@ func (c *Client) UpdateNodeLabels(nodeName string, nodeLabels map[string]string)
108109 return fmt .Errorf ("failed to marshal patch: %w" , err )
109110 }
110111
111- backoff := wait.Backoff {
112- Duration : time .Second ,
113- Factor : 2.0 ,
114- Jitter : 0.2 ,
115- Steps : 7 ,
116- }
117-
118- return retry .OnError (backoff , func (err error ) bool {
119- return true
120- }, func () error {
121- _ , err := c .clientset .CoreV1 ().Nodes ().Patch (c .ctx , nodeName , types .StrategicMergePatchType , patchBytes , metav1.PatchOptions {})
122- if err != nil {
123- c .log .Warnf ("Failed to update labels on node %s, retrying: %v" , nodeName , err )
124- }
125- return err
126- })
112+ return c .patchNodeWithRetry (nodeName , patchBytes , types .StrategicMergePatchType )
127113}
128114
129115// GetNodeAnnotationValue returns the annotation value given a node name and annotation key
@@ -164,6 +150,58 @@ func (c *Client) UncordonNode(nodeName string) error {
164150 return drain .RunCordonOrUncordon (drainHelper , node , false )
165151}
166152
153+ // TaintNode applies a taint on a Node marking it as Unschedulable
154+ func (c * Client ) TaintNode (taint corev1.Taint , nodeName string ) error {
155+ c .log .Infof ("Tainting node %s" , nodeName )
156+
157+ node , err := c .clientset .CoreV1 ().Nodes ().Get (c .ctx , nodeName , metav1.GetOptions {})
158+ if err != nil {
159+ return fmt .Errorf ("failed to get node %s: %w" , nodeName , err )
160+ }
161+
162+ original := node .DeepCopy ()
163+
164+ // First, we check if the taint already exists. If yes, we skip the node patch operation
165+ for _ , nodeTaint := range node .Spec .Taints {
166+ if taintsEqual (nodeTaint , taint ) {
167+ c .log .Infof ("Taint has already been applied. Skipping node taint..." )
168+ return nil
169+ }
170+ }
171+
172+ node .Spec .Taints = append (node .Spec .Taints , taint )
173+
174+ return c .applyNodeDiff (original , node )
175+ }
176+
177+ // UntaintNode removes the taint on a Node
178+ func (c * Client ) UntaintNode (taint corev1.Taint , nodeName string ) error {
179+ c .log .Infof ("Untainting node %s" , nodeName )
180+
181+ node , err := c .clientset .CoreV1 ().Nodes ().Get (c .ctx , nodeName , metav1.GetOptions {})
182+ if err != nil {
183+ return fmt .Errorf ("failed to get node %s: %w" , nodeName , err )
184+ }
185+
186+ original := node .DeepCopy ()
187+
188+ found := false
189+ for i , nodeTaint := range node .Spec .Taints {
190+ if taintsEqual (nodeTaint , taint ) {
191+ node .Spec .Taints = append (node .Spec .Taints [:i ], node .Spec .Taints [i + 1 :]... )
192+ found = true
193+ break
194+ }
195+ }
196+
197+ if ! found {
198+ c .log .Infof ("Taint not found. Skipping node untaint..." )
199+ return nil
200+ }
201+
202+ return c .applyNodeDiff (original , node )
203+ }
204+
167205// WaitForPodTermination will wait for the termination of pods matching labels from the selectorMap on the node with the specified namespace.
168206// It will continue to wait until the specified timeout elapses
169207func (c * Client ) WaitForPodTermination (selectorMap map [string ]string , namespace , nodeName string , timeout time.Duration ) error {
@@ -183,6 +221,29 @@ func (c *Client) WaitForPodTermination(selectorMap map[string]string, namespace,
183221 })
184222}
185223
224+ // ListPodsOnNodeWithAnnotation will list all pods matching a specified annotation on a node
225+ func (c * Client ) ListPodsOnNodeWithAnnotation (annotationKey string , nodeName string ) ([]corev1.Pod , error ) {
226+
227+ podList , err := c .clientset .CoreV1 ().Pods (corev1 .NamespaceAll ).List (c .ctx , metav1.ListOptions {
228+ FieldSelector : "spec.nodeName=" + nodeName ,
229+ })
230+ if err != nil {
231+ return nil , err
232+ }
233+ var result []corev1.Pod
234+ for _ , pod := range podList .Items {
235+ if _ , exists := pod .Annotations [annotationKey ]; exists {
236+ result = append (result , pod )
237+ }
238+ }
239+ return result , nil
240+ }
241+
242+ // DeletePod will delete a pid with the name "podName"
243+ func (c * Client ) DeletePod (podName , namespace string ) error {
244+ return c .clientset .CoreV1 ().Pods (namespace ).Delete (c .ctx , podName , metav1.DeleteOptions {})
245+ }
246+
186247// DrainNode drains a Node given a Node name and a set of drain option parameters
187248func (c * Client ) DrainNode (nodeName string , drainOpts DrainOptions ) error {
188249 c .log .Infof ("Draining node %s" , nodeName )
@@ -293,3 +354,46 @@ func gpuPodSpecFilter(pod corev1.Pod) bool {
293354 }
294355 return false
295356}
357+
358+ func (c * Client ) applyNodeDiff (currentNode * corev1.Node , desiredNode * corev1.Node ) error {
359+ currentNodeJSON , err := json .Marshal (currentNode )
360+ if err != nil {
361+ return fmt .Errorf ("failed to marshal node object %v: %w" , currentNode , err )
362+ }
363+ desiredNodeJSON , err := json .Marshal (desiredNode )
364+ if err != nil {
365+ return fmt .Errorf ("failed to marshal node object %v: %w" , desiredNode , err )
366+ }
367+
368+ patch , err := jsonpatch .CreateMergePatch (currentNodeJSON , desiredNodeJSON )
369+ if err != nil {
370+ return fmt .Errorf ("failed to create merge-patch: %w" , err )
371+ }
372+ return c .patchNodeWithRetry (desiredNode .Name , patch , types .MergePatchType )
373+ }
374+
375+ func (c * Client ) patchNodeWithRetry (nodeName string , patch []byte , patchType types.PatchType ) error {
376+
377+ backoff := wait.Backoff {
378+ Duration : time .Second ,
379+ Factor : 2.0 ,
380+ Jitter : 0.2 ,
381+ Steps : 7 ,
382+ }
383+ return retry .OnError (backoff , func (err error ) bool {
384+ return true
385+ }, func () error {
386+ _ , err := c .clientset .CoreV1 ().Nodes ().Patch (c .ctx , nodeName , patchType , patch , metav1.PatchOptions {})
387+ if err != nil {
388+ c .log .Warnf ("Failed to patch node %s, retrying: %v" , nodeName , err )
389+ }
390+ return err
391+ })
392+ }
393+
394+ func taintsEqual (t1 corev1.Taint , t2 corev1.Taint ) bool {
395+ keysEqual := t1 .Key == t2 .Key
396+ valuesEqual := t1 .Value == t2 .Value
397+ effectsEqual := t1 .Effect == t2 .Effect
398+ return keysEqual && valuesEqual && effectsEqual
399+ }
0 commit comments