migrate/arm_patch.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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
package main

import (
	"context"
	"fmt"

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

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

func patchDeploymentARMAffinity(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)

	if HasArm64Preference(newDeployment.Spec.Template.Spec.Affinity) {
		fmt.Printf("workload %s %s/%s already has arm preference, skip the prefer affinity\n",
			workload.Kind, workload.Namespace, workload.Name)
	} else {
		newDeployment.Spec.Template.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution =
			AddArm64Preference(newDeployment.Spec.Template.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution)
	}
	if CheckWorkloadHasARM64Toleration(newDeployment.Spec.Template.Spec.Tolerations) {
		fmt.Printf("workload %s %s/%s already has arm64 toleration, skip it\n",
			workload.Kind, workload.Namespace, workload.Name)
	} else {
		newDeployment.Spec.Template.Spec.Tolerations = AddARM64Toleration(newDeployment.Spec.Template.Spec.Tolerations)
	}

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

func patchStatefulSetARMAffinity(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)

	if HasArm64Preference(newSS.Spec.Template.Spec.Affinity) {
		fmt.Printf("workload %s %s/%s already has arm preference, skip the prefer affinity\n",
			workload.Kind, workload.Namespace, workload.Name)
	} else {
		newSS.Spec.Template.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution =
			AddArm64Preference(newSS.Spec.Template.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution)
	}

	if CheckWorkloadHasARM64Toleration(newSS.Spec.Template.Spec.Tolerations) {
		fmt.Printf("workload %s %s/%s already has arm64 toleration, skip it\n",
			workload.Kind, workload.Namespace, workload.Name)
	} else {
		newSS.Spec.Template.Spec.Tolerations = AddARM64Toleration(newSS.Spec.Template.Spec.Tolerations)
	}

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