Add node taint config (#137)

Author: wenxuan0923

components/kubelet/action.pb.go

@@ -69,6 +69,7 @@ message StartKubeletServiceSpec {
   NodeAuthInfo node_auth_info = 2;
   map<string, string> node_labels = 3;
   KubeletConfig kubelet_config = 4;
+  repeated string register_with_taints = 5;
 }
 
 message StartKubeletServiceStatus {

components/kubelet/action.proto

@@ -4,5 +4,5 @@
 # You may create a higher-numbered drop-in (e.g. 90-custom-node-config.conf)
 # to append or override individual env vars.
 [Service]
-Environment="KUBELET_NODE_CONFIG_ARGS=--node-labels={{.NodeLabels}} --v={{.Verbosity}} --client-ca-file={{.ClientCAFile}}{{range .ClusterDNS}} --cluster-dns={{.}}{{end}}"
+Environment="KUBELET_NODE_CONFIG_ARGS=--node-labels={{.NodeLabels}} --v={{.Verbosity}} --client-ca-file={{.ClientCAFile}}{{range .ClusterDNS}} --cluster-dns={{.}}{{end}}{{if .RegisterWithTaints}} --register-with-taints={{.RegisterWithTaints}}{{end}}"
 Environment="KUBELET_TUNING_ARGS=--eviction-hard={{.EvictionHard}} --kube-reserved={{.KubeReserved}} --image-gc-high-threshold={{.ImageGCHighThreshold}} --image-gc-low-threshold={{.ImageGCLowThreshold}} --max-pods={{.MaxPods}}"

components/kubelet/v20260301/assets/20-flex-node-node-config.conf

@@ -3,6 +3,7 @@ package v20260301
 import (
 	"bytes"
 	"context"
+	"strings"
 
 	"github.com/Azure/AKSFlexNode/components/kubelet"
 	"github.com/Azure/AKSFlexNode/pkg/config"
@@ -74,6 +75,7 @@ func (s *startKubeletServiceAction) ensureSystemdUnit(
 		"Verbosity":            kubeletConfig.GetVerbosity(),
 		"ClientCAFile":         apiServerClientCAPath, // prepared in ensureAPIServerCA
 		"ClusterDNS":           kubeletConfig.GetClusterDns(),
+		"RegisterWithTaints":   strings.Join(spec.GetRegisterWithTaints(), ","),
 		"EvictionHard":         mapPairsToString(kubeletConfig.GetEvictionHard(), "<", ","),
 		"KubeReserved":         mapPairsToString(kubeletConfig.GetKubeReserved(), "=", ","),
 		"ImageGCHighThreshold": kubeletConfig.GetImageGcHighThreshold(),

components/kubelet/v20260301/kubelet_service.go

@@ -286,8 +286,9 @@ var startKubelet resolveActionFunc[*kubelet.StartKubeletService] = func(
 			Server:                   ptr(cfg.Node.Kubelet.ServerURL),
 			CertificateAuthorityData: caData,
 		}.Build(),
-		NodeAuthInfo: nodeAuthInfo.Build(),
-		NodeLabels:   maps.Clone(cfg.Node.Labels),
+		NodeAuthInfo:       nodeAuthInfo.Build(),
+		NodeLabels:         maps.Clone(cfg.Node.Labels),
+		RegisterWithTaints: append([]string(nil), cfg.Node.Taints...),
 		KubeletConfig: kubelet.KubeletConfig_builder{
 			KubeReserved:         maps.Clone(cfg.Node.Kubelet.KubeReserved),
 			EvictionHard:         maps.Clone(cfg.Node.Kubelet.EvictionHard),
@@ -358,6 +359,7 @@ var ensureMachine resolveActionFunc[*aksmachine.EnsureMachine] = func(
 		KubernetesVersion: ptr(cfg.Kubernetes.Version),
 		MaxPods:           ptr(int32(cfg.Node.MaxPods)),
 		NodeLabels:        maps.Clone(cfg.Node.Labels),
+		NodeTaints:        append([]string(nil), cfg.Node.Taints...),
 		KubeletConfig:     kubeletCfg,
 		Enabled:           ptr(cfg.IsDriftDetectionAndRemediationEnabled()),
 		AzureCredential:   azureCred,

pkg/bootstrapper/components.go

@@ -45,6 +45,9 @@ func (cfg *Config) DeepCopy() *Config {
 
 	// Copy node-level maps.
 	out.Node.Labels = cloneStringMap(cfg.Node.Labels)
+	if cfg.Node.Taints != nil {
+		out.Node.Taints = append([]string(nil), cfg.Node.Taints...)
+	}
 	out.Node.Kubelet.KubeReserved = cloneStringMap(cfg.Node.Kubelet.KubeReserved)
 	out.Node.Kubelet.EvictionHard = cloneStringMap(cfg.Node.Kubelet.EvictionHard)
 

pkg/config/copy.go

@@ -52,6 +52,7 @@ func TestConfigDeepCopy_DoesNotSharePointersOrMaps(t *testing.T) {
 		Agent: AgentConfig{EnableDriftDetectionAndRemediation: &falseVal},
 		Node: NodeConfig{
 			Labels: map[string]string{"l": "1"},
+			Taints: []string{"dedicated=infra:NoSchedule"},
 			Kubelet: KubeletConfig{
 				KubeReserved: map[string]string{"cpu": "100m"},
 				EvictionHard: map[string]string{"memory.available": "200Mi"},
@@ -103,6 +104,15 @@ func TestConfigDeepCopy_DoesNotSharePointersOrMaps(t *testing.T) {
 		t.Fatalf("Node.Labels shared; orig=%q, want %q", cfg.Node.Labels["l"], "orig")
 	}
 
+	// Taints slice should not be shared.
+	if len(copy.Node.Taints) != 1 || copy.Node.Taints[0] != "dedicated=infra:NoSchedule" {
+		t.Fatalf("Node.Taints copy=%v, want [dedicated=infra:NoSchedule]", copy.Node.Taints)
+	}
+	cfg.Node.Taints[0] = "mutated:NoExecute"
+	if copy.Node.Taints[0] != "dedicated=infra:NoSchedule" {
+		t.Fatalf("Node.Taints shares backing array; copy=%q, want %q", copy.Node.Taints[0], "dedicated=infra:NoSchedule")
+	}
+
 	cfg.Node.Kubelet.KubeReserved["cpu"] = "200m"
 	if copy.Node.Kubelet.KubeReserved["cpu"] != "100m" {
 		t.Fatalf("KubeReserved shared; copy=%q, want %q", copy.Node.Kubelet.KubeReserved["cpu"], "100m")

pkg/config/copy_test.go

@@ -119,7 +119,11 @@ type ContainerdConfig struct {
 type NodeConfig struct {
 	MaxPods int               `json:"maxPods"`
 	Labels  map[string]string `json:"labels"`
-	Kubelet KubeletConfig     `json:"kubelet"`
+	// Taints to apply at node registration time via --register-with-taints.
+	// Each entry must use the kubelet taint format: "key=value:Effect" or "key:Effect"
+	// (e.g. "dedicated=infra:NoSchedule", "gpu:NoExecute").
+	Taints  []string      `json:"taints,omitempty"`
+	Kubelet KubeletConfig `json:"kubelet"`
 }
 
 // KubeletConfig holds kubelet-specific configuration settings.