From abf21d38b6b3782ea81c272f3c2eec9db123e9a2 Mon Sep 17 00:00:00 2001
From: Devin Wong <wongsiosun@outlook.com>
Date: Mon, 27 Jan 2025 21:44:00 -0800
Subject: [PATCH] fine tuning

---
 aks-node-controller/parser/helper.go      | 24 ++++++++++++++++++-----
 aks-node-controller/parser/helper_test.go |  2 +-
 aks-node-controller/parser/parser.go      |  4 ++--
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/aks-node-controller/parser/helper.go b/aks-node-controller/parser/helper.go
index 66c269b5b1a..c43eda55851 100644
--- a/aks-node-controller/parser/helper.go
+++ b/aks-node-controller/parser/helper.go
@@ -46,13 +46,13 @@ var (
 	containerdConfigTemplateText string
 	//nolint:gochecknoglobals
 	containerdConfigTemplate = template.Must(
-		template.New("containerdconfigforaksnodeconfig").Funcs(getFuncMapForContainerdConfigTemplate()).Parse(containerdConfigTemplateText),
+		template.New("containerdconfig").Funcs(getFuncMapForContainerdConfigTemplate()).Parse(containerdConfigTemplateText),
 	)
 	//go:embed  templates/containerd_no_GPU.toml.gtpl
 	containerdConfigNoGPUTemplateText string
 	//nolint:gochecknoglobals
 	containerdConfigNoGPUTemplate = template.Must(
-		template.New("containerdconfigforaksnodeconfig").Funcs(getFuncMapForContainerdConfigTemplate()).Parse(containerdConfigNoGPUTemplateText),
+		template.New("nogpucontainerdconfig").Funcs(getFuncMapForContainerdConfigTemplate()).Parse(containerdConfigNoGPUTemplateText),
 	)
 )
 
@@ -145,13 +145,13 @@ func getKubenetTemplate() string {
 	return base64.StdEncoding.EncodeToString(kubenetTemplateContent)
 }
 
-// getContainerdConfig returns the base64 encoded containerd config depending on whether the node is with GPU or not.
-func getContainerdConfig(aksnodeconfig *aksnodeconfigv1.Configuration, noGPU bool) string {
+// getContainerdConfigBase64 returns the base64 encoded containerd config depending on whether the node is with GPU or not.
+func getContainerdConfigBase64(aksnodeconfig *aksnodeconfigv1.Configuration) string {
 	if aksnodeconfig == nil {
 		return ""
 	}
 
-	containerdConfig, err := containerdConfigFromAKSNodeConfig(aksnodeconfig, noGPU)
+	containerdConfig, err := containerdConfigFromAKSNodeConfig(aksnodeconfig, false)
 	if err != nil {
 		return fmt.Sprintf("error getting containerd config from node bootstrap variables: %v", err)
 	}
@@ -159,6 +159,20 @@ func getContainerdConfig(aksnodeconfig *aksnodeconfigv1.Configuration, noGPU boo
 	return base64.StdEncoding.EncodeToString([]byte(containerdConfig))
 }
 
+// getNoGPUContainerdConfigBase64 returns the base64 encoded containerd config depending on whether the node is with GPU or not.
+func getNoGPUContainerdConfigBase64(aksnodeconfig *aksnodeconfigv1.Configuration) string {
+	if aksnodeconfig == nil {
+		return ""
+	}
+
+	containerdConfig, err := containerdConfigFromAKSNodeConfig(aksnodeconfig, true)
+	if err != nil {
+		return fmt.Sprintf("error getting No GPU containerd config from node bootstrap variables: %v", err)
+	}
+
+	return base64.StdEncoding.EncodeToString([]byte(containerdConfig))
+}
+
 func containerdConfigFromAKSNodeConfig(aksnodeconfig *aksnodeconfigv1.Configuration, noGPU bool) (string, error) {
 	if aksnodeconfig == nil {
 		return "", fmt.Errorf("AKSNodeConfig is nil")
diff --git a/aks-node-controller/parser/helper_test.go b/aks-node-controller/parser/helper_test.go
index 859ccc153b7..90bc200a828 100644
--- a/aks-node-controller/parser/helper_test.go
+++ b/aks-node-controller/parser/helper_test.go
@@ -348,7 +348,7 @@ oom_score = 0
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			if got := getContainerdConfig(tt.args.aksnodeconfig, tt.args.noGpu); got != tt.want {
+			if got := getContainerdConfigBase64(tt.args.aksnodeconfig); got != tt.want {
 				t.Errorf("getContainerdConfig() = %v, want %v", got, tt.want)
 			}
 		})
diff --git a/aks-node-controller/parser/parser.go b/aks-node-controller/parser/parser.go
index a3fb80ca4f2..e60c1d1335d 100644
--- a/aks-node-controller/parser/parser.go
+++ b/aks-node-controller/parser/parser.go
@@ -158,8 +158,8 @@ func getCSEEnv(config *aksnodeconfigv1.Configuration) map[string]string {
 		"AZURE_ENVIRONMENT_FILEPATH":                     getAzureEnvironmentFilepath(config),
 		"KUBE_CA_CRT":                                    config.GetKubernetesCaCert(),
 		"KUBENET_TEMPLATE":                               getKubenetTemplate(),
-		"CONTAINERD_CONFIG_CONTENT":                      getContainerdConfig(config, false),
-		"CONTAINERD_CONFIG_NO_GPU_CONTENT":               getContainerdConfig(config, true),
+		"CONTAINERD_CONFIG_CONTENT":                      getContainerdConfigBase64(config),
+		"CONTAINERD_CONFIG_NO_GPU_CONTENT":               getNoGPUContainerdConfigBase64(config),
 		"IS_KATA":                                        fmt.Sprintf("%v", config.GetIsKata()),
 		"ARTIFACT_STREAMING_ENABLED":                     fmt.Sprintf("%v", config.GetEnableArtifactStreaming()),
 		"SYSCTL_CONTENT":                                 getSysctlContent(config.GetCustomLinuxOsConfig().GetSysctlConfig()),