migrate/arm_rollback.go at main · cloudpilot-ai/migrate · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
package main

import (
	"context"
	"fmt"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func rollbackWorkloadARMAffinity(selectedWorkloads []Workload) {
	for _, workload := range selectedWorkloads {
		var err error
		switch workload.Kind {
		case WorkloadDeployment:
			err = rollbackDeploymentARMAffinity(&workload)
		case WorkloadStatefulSet:
			err = rollbackStatefulSetARMAffinity(&workload)
		}
		if err != nil {
			fmt.Printf("Failed to rollback %s workload %s/%s, err: %v\n", workload.Kind,
				workload.Namespace, workload.Name, err)
		}
	}
}

func rollbackDeploymentARMAffinity(workload *Workload) error {
	ctx := context.Background()
	deployment, err := kubeClient.AppsV1().Deployments(workload.Namespace).
		Get(ctx, workload.Name, metav1.GetOptions{})
	if err != nil {
		return fmt.Errorf("get deployment: %w", err)
	}

	newDeployment := deployment.DeepCopy()
	newDeployment.Spec.Template.Spec.Affinity = ensurePreferAffinity(newDeployment.Spec.Template.Spec.Affinity)

	newDeployment.Spec.Template.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution =
		RemoveArm64Preference(newDeployment.Spec.Template.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution)
	newDeployment.Spec.Template.Spec.Tolerations = RemoveARM64Toleration(newDeployment.Spec.Template.Spec.Tolerations)

	return patchResource(ctx, deployment, newDeployment, workload.Namespace, workload.Name, workload.Kind)
}

func rollbackStatefulSetARMAffinity(workload *Workload) error {
	ctx := context.Background()

	ss, err := kubeClient.AppsV1().
		StatefulSets(workload.Namespace).
		Get(ctx, workload.Name, metav1.GetOptions{})
	if err != nil {
		return fmt.Errorf("get statefulset: %w", err)
	}

	newSS := ss.DeepCopy()
	newSS.Spec.Template.Spec.Affinity = ensurePreferAffinity(newSS.Spec.Template.Spec.Affinity)

	newSS.Spec.Template.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution =
		RemoveArm64Preference(newSS.Spec.Template.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution)
	newSS.Spec.Template.Spec.Tolerations = RemoveARM64Toleration(newSS.Spec.Template.Spec.Tolerations)

	return patchResource(ctx, ss, newSS, workload.Namespace, workload.Name, workload.Kind)
}