From 47c30e5b0acd67d3d784d495d45e48da2a20cab5 Mon Sep 17 00:00:00 2001 From: wenxuanW Date: Thu, 2 Apr 2026 09:41:46 -0700 Subject: [PATCH] Add node taint config --- components/kubelet/action.pb.go | 100 +++++++++++------- components/kubelet/action.proto | 1 + .../assets/20-flex-node-node-config.conf | 2 +- .../kubelet/v20260301/kubelet_service.go | 2 + pkg/bootstrapper/components.go | 6 +- pkg/config/copy.go | 3 + pkg/config/copy_test.go | 10 ++ pkg/config/structs.go | 6 +- 8 files changed, 85 insertions(+), 45 deletions(-) diff --git a/components/kubelet/action.pb.go b/components/kubelet/action.pb.go index 75d645b6..d2e7157b 100644 --- a/components/kubelet/action.pb.go +++ b/components/kubelet/action.pb.go @@ -862,13 +862,11 @@ func (x *NodeAuthInfo) ClearBootstrapTokenCredential() { } } -const ( - NodeAuthInfo_AuthInfo_not_set_case case_NodeAuthInfo_AuthInfo = 0 - NodeAuthInfo_ArcCredential_case case_NodeAuthInfo_AuthInfo = 1 - NodeAuthInfo_MsiCredential_case case_NodeAuthInfo_AuthInfo = 2 - NodeAuthInfo_ServicePrincipalCredential_case case_NodeAuthInfo_AuthInfo = 3 - NodeAuthInfo_BootstrapTokenCredential_case case_NodeAuthInfo_AuthInfo = 4 -) +const NodeAuthInfo_AuthInfo_not_set_case case_NodeAuthInfo_AuthInfo = 0 +const NodeAuthInfo_ArcCredential_case case_NodeAuthInfo_AuthInfo = 1 +const NodeAuthInfo_MsiCredential_case case_NodeAuthInfo_AuthInfo = 2 +const NodeAuthInfo_ServicePrincipalCredential_case case_NodeAuthInfo_AuthInfo = 3 +const NodeAuthInfo_BootstrapTokenCredential_case case_NodeAuthInfo_AuthInfo = 4 func (x *NodeAuthInfo) WhichAuthInfo() case_NodeAuthInfo_AuthInfo { if x == nil { @@ -1171,13 +1169,14 @@ func (b0 KubeletConfig_builder) Build() *KubeletConfig { } type StartKubeletServiceSpec struct { - state protoimpl.MessageState `protogen:"opaque.v1"` - xxx_hidden_ControlPlane *ControlPlane `protobuf:"bytes,1,opt,name=control_plane,json=controlPlane"` - xxx_hidden_NodeAuthInfo *NodeAuthInfo `protobuf:"bytes,2,opt,name=node_auth_info,json=nodeAuthInfo"` - xxx_hidden_NodeLabels map[string]string `protobuf:"bytes,3,rep,name=node_labels,json=nodeLabels" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` - xxx_hidden_KubeletConfig *KubeletConfig `protobuf:"bytes,4,opt,name=kubelet_config,json=kubeletConfig"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + state protoimpl.MessageState `protogen:"opaque.v1"` + xxx_hidden_ControlPlane *ControlPlane `protobuf:"bytes,1,opt,name=control_plane,json=controlPlane"` + xxx_hidden_NodeAuthInfo *NodeAuthInfo `protobuf:"bytes,2,opt,name=node_auth_info,json=nodeAuthInfo"` + xxx_hidden_NodeLabels map[string]string `protobuf:"bytes,3,rep,name=node_labels,json=nodeLabels" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + xxx_hidden_KubeletConfig *KubeletConfig `protobuf:"bytes,4,opt,name=kubelet_config,json=kubeletConfig"` + xxx_hidden_RegisterWithTaints []string `protobuf:"bytes,5,rep,name=register_with_taints,json=registerWithTaints"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *StartKubeletServiceSpec) Reset() { @@ -1233,6 +1232,13 @@ func (x *StartKubeletServiceSpec) GetKubeletConfig() *KubeletConfig { return nil } +func (x *StartKubeletServiceSpec) GetRegisterWithTaints() []string { + if x != nil { + return x.xxx_hidden_RegisterWithTaints + } + return nil +} + func (x *StartKubeletServiceSpec) SetControlPlane(v *ControlPlane) { x.xxx_hidden_ControlPlane = v } @@ -1249,6 +1255,10 @@ func (x *StartKubeletServiceSpec) SetKubeletConfig(v *KubeletConfig) { x.xxx_hidden_KubeletConfig = v } +func (x *StartKubeletServiceSpec) SetRegisterWithTaints(v []string) { + x.xxx_hidden_RegisterWithTaints = v +} + func (x *StartKubeletServiceSpec) HasControlPlane() bool { if x == nil { return false @@ -1289,6 +1299,9 @@ type StartKubeletServiceSpec_builder struct { NodeAuthInfo *NodeAuthInfo NodeLabels map[string]string KubeletConfig *KubeletConfig + // Taints to register the node with via --register-with-taints. + // Each entry must use the kubelet taint format: "key=value:Effect" or "key:Effect". + RegisterWithTaints []string } func (b0 StartKubeletServiceSpec_builder) Build() *StartKubeletServiceSpec { @@ -1299,6 +1312,7 @@ func (b0 StartKubeletServiceSpec_builder) Build() *StartKubeletServiceSpec { x.xxx_hidden_NodeAuthInfo = b.NodeAuthInfo x.xxx_hidden_NodeLabels = b.NodeLabels x.xxx_hidden_KubeletConfig = b.KubeletConfig + x.xxx_hidden_RegisterWithTaints = b.RegisterWithTaints return m0 } @@ -1335,6 +1349,7 @@ func (x *StartKubeletServiceStatus) ProtoReflect() protoreflect.Message { type StartKubeletServiceStatus_builder struct { _ [0]func() // Prevents comparability and use of unkeyed literals for the builder. + } func (b0 StartKubeletServiceStatus_builder) Build() *StartKubeletServiceStatus { @@ -1495,6 +1510,7 @@ func (x *ResetKubeletSpec) ProtoReflect() protoreflect.Message { type ResetKubeletSpec_builder struct { _ [0]func() // Prevents comparability and use of unkeyed literals for the builder. + } func (b0 ResetKubeletSpec_builder) Build() *ResetKubeletSpec { @@ -1537,6 +1553,7 @@ func (x *ResetKubeletStatus) ProtoReflect() protoreflect.Message { type ResetKubeletStatus_builder struct { _ [0]func() // Prevents comparability and use of unkeyed literals for the builder. + } func (b0 ResetKubeletStatus_builder) Build() *ResetKubeletStatus { @@ -1697,6 +1714,7 @@ func (x *StopKubeletServiceSpec) ProtoReflect() protoreflect.Message { type StopKubeletServiceSpec_builder struct { _ [0]func() // Prevents comparability and use of unkeyed literals for the builder. + } func (b0 StopKubeletServiceSpec_builder) Build() *StopKubeletServiceSpec { @@ -1739,6 +1757,7 @@ func (x *StopKubeletServiceStatus) ProtoReflect() protoreflect.Message { type StopKubeletServiceStatus_builder struct { _ [0]func() // Prevents comparability and use of unkeyed literals for the builder. + } func (b0 StopKubeletServiceStatus_builder) Build() *StopKubeletServiceStatus { @@ -1792,13 +1811,14 @@ const file_components_kubelet_action_proto_rawDesc = "" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\x1a?\n" + "\x11EvictionHardEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + - "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\xb3\x03\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\xe5\x03\n" + "\x17StartKubeletServiceSpec\x12N\n" + "\rcontrol_plane\x18\x01 \x01(\v2).aks.flex.components.kubelet.ControlPlaneR\fcontrolPlane\x12O\n" + "\x0enode_auth_info\x18\x02 \x01(\v2).aks.flex.components.kubelet.NodeAuthInfoR\fnodeAuthInfo\x12e\n" + "\vnode_labels\x18\x03 \x03(\v2D.aks.flex.components.kubelet.StartKubeletServiceSpec.NodeLabelsEntryR\n" + "nodeLabels\x12Q\n" + - "\x0ekubelet_config\x18\x04 \x01(\v2*.aks.flex.components.kubelet.KubeletConfigR\rkubeletConfig\x1a=\n" + + "\x0ekubelet_config\x18\x04 \x01(\v2*.aks.flex.components.kubelet.KubeletConfigR\rkubeletConfig\x120\n" + + "\x14register_with_taints\x18\x05 \x03(\tR\x12registerWithTaints\x1a=\n" + "\x0fNodeLabelsEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\x1b\n" + @@ -1816,31 +1836,29 @@ const file_components_kubelet_action_proto_rawDesc = "" + "\x16StopKubeletServiceSpec\"\x1a\n" + "\x18StopKubeletServiceStatusB1Z/github.com/Azure/AKSFlexNode/components/kubeletb\beditionsp\xe9\a" -var ( - file_components_kubelet_action_proto_msgTypes = make([]protoimpl.MessageInfo, 19) - file_components_kubelet_action_proto_goTypes = []any{ - (*StartKubeletService)(nil), // 0: aks.flex.components.kubelet.StartKubeletService - (*KubeletArcCredential)(nil), // 1: aks.flex.components.kubelet.KubeletArcCredential - (*KubeletMSICredential)(nil), // 2: aks.flex.components.kubelet.KubeletMSICredential - (*KubeletServicePrincipalCredential)(nil), // 3: aks.flex.components.kubelet.KubeletServicePrincipalCredential - (*KubeletBootstrapTokenCredential)(nil), // 4: aks.flex.components.kubelet.KubeletBootstrapTokenCredential - (*ControlPlane)(nil), // 5: aks.flex.components.kubelet.ControlPlane - (*NodeAuthInfo)(nil), // 6: aks.flex.components.kubelet.NodeAuthInfo - (*KubeletConfig)(nil), // 7: aks.flex.components.kubelet.KubeletConfig - (*StartKubeletServiceSpec)(nil), // 8: aks.flex.components.kubelet.StartKubeletServiceSpec - (*StartKubeletServiceStatus)(nil), // 9: aks.flex.components.kubelet.StartKubeletServiceStatus - (*ResetKubelet)(nil), // 10: aks.flex.components.kubelet.ResetKubelet - (*ResetKubeletSpec)(nil), // 11: aks.flex.components.kubelet.ResetKubeletSpec - (*ResetKubeletStatus)(nil), // 12: aks.flex.components.kubelet.ResetKubeletStatus - (*StopKubeletService)(nil), // 13: aks.flex.components.kubelet.StopKubeletService - (*StopKubeletServiceSpec)(nil), // 14: aks.flex.components.kubelet.StopKubeletServiceSpec - (*StopKubeletServiceStatus)(nil), // 15: aks.flex.components.kubelet.StopKubeletServiceStatus - nil, // 16: aks.flex.components.kubelet.KubeletConfig.KubeReservedEntry - nil, // 17: aks.flex.components.kubelet.KubeletConfig.EvictionHardEntry - nil, // 18: aks.flex.components.kubelet.StartKubeletServiceSpec.NodeLabelsEntry - (*api.Metadata)(nil), // 19: aks.flex.components.api.Metadata - } -) +var file_components_kubelet_action_proto_msgTypes = make([]protoimpl.MessageInfo, 19) +var file_components_kubelet_action_proto_goTypes = []any{ + (*StartKubeletService)(nil), // 0: aks.flex.components.kubelet.StartKubeletService + (*KubeletArcCredential)(nil), // 1: aks.flex.components.kubelet.KubeletArcCredential + (*KubeletMSICredential)(nil), // 2: aks.flex.components.kubelet.KubeletMSICredential + (*KubeletServicePrincipalCredential)(nil), // 3: aks.flex.components.kubelet.KubeletServicePrincipalCredential + (*KubeletBootstrapTokenCredential)(nil), // 4: aks.flex.components.kubelet.KubeletBootstrapTokenCredential + (*ControlPlane)(nil), // 5: aks.flex.components.kubelet.ControlPlane + (*NodeAuthInfo)(nil), // 6: aks.flex.components.kubelet.NodeAuthInfo + (*KubeletConfig)(nil), // 7: aks.flex.components.kubelet.KubeletConfig + (*StartKubeletServiceSpec)(nil), // 8: aks.flex.components.kubelet.StartKubeletServiceSpec + (*StartKubeletServiceStatus)(nil), // 9: aks.flex.components.kubelet.StartKubeletServiceStatus + (*ResetKubelet)(nil), // 10: aks.flex.components.kubelet.ResetKubelet + (*ResetKubeletSpec)(nil), // 11: aks.flex.components.kubelet.ResetKubeletSpec + (*ResetKubeletStatus)(nil), // 12: aks.flex.components.kubelet.ResetKubeletStatus + (*StopKubeletService)(nil), // 13: aks.flex.components.kubelet.StopKubeletService + (*StopKubeletServiceSpec)(nil), // 14: aks.flex.components.kubelet.StopKubeletServiceSpec + (*StopKubeletServiceStatus)(nil), // 15: aks.flex.components.kubelet.StopKubeletServiceStatus + nil, // 16: aks.flex.components.kubelet.KubeletConfig.KubeReservedEntry + nil, // 17: aks.flex.components.kubelet.KubeletConfig.EvictionHardEntry + nil, // 18: aks.flex.components.kubelet.StartKubeletServiceSpec.NodeLabelsEntry + (*api.Metadata)(nil), // 19: aks.flex.components.api.Metadata +} var file_components_kubelet_action_proto_depIdxs = []int32{ 19, // 0: aks.flex.components.kubelet.StartKubeletService.metadata:type_name -> aks.flex.components.api.Metadata 8, // 1: aks.flex.components.kubelet.StartKubeletService.spec:type_name -> aks.flex.components.kubelet.StartKubeletServiceSpec diff --git a/components/kubelet/action.proto b/components/kubelet/action.proto index 15667588..3412dbc2 100644 --- a/components/kubelet/action.proto +++ b/components/kubelet/action.proto @@ -69,6 +69,7 @@ message StartKubeletServiceSpec { NodeAuthInfo node_auth_info = 2; map node_labels = 3; KubeletConfig kubelet_config = 4; + repeated string register_with_taints = 5; } message StartKubeletServiceStatus { diff --git a/components/kubelet/v20260301/assets/20-flex-node-node-config.conf b/components/kubelet/v20260301/assets/20-flex-node-node-config.conf index 5565d59a..d0d4cb6f 100644 --- a/components/kubelet/v20260301/assets/20-flex-node-node-config.conf +++ b/components/kubelet/v20260301/assets/20-flex-node-node-config.conf @@ -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}}" diff --git a/components/kubelet/v20260301/kubelet_service.go b/components/kubelet/v20260301/kubelet_service.go index 65b2911c..9c39570e 100644 --- a/components/kubelet/v20260301/kubelet_service.go +++ b/components/kubelet/v20260301/kubelet_service.go @@ -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(), diff --git a/pkg/bootstrapper/components.go b/pkg/bootstrapper/components.go index 6c6f98a4..d87b1e62 100644 --- a/pkg/bootstrapper/components.go +++ b/pkg/bootstrapper/components.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, diff --git a/pkg/config/copy.go b/pkg/config/copy.go index 67ba1584..45c735d6 100644 --- a/pkg/config/copy.go +++ b/pkg/config/copy.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) diff --git a/pkg/config/copy_test.go b/pkg/config/copy_test.go index ec4999b0..aa613545 100644 --- a/pkg/config/copy_test.go +++ b/pkg/config/copy_test.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") diff --git a/pkg/config/structs.go b/pkg/config/structs.go index e2b21907..824238b6 100644 --- a/pkg/config/structs.go +++ b/pkg/config/structs.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.