@@ -30,6 +30,8 @@ import (
3030 ackclient "github.com/alibabacloud-go/cs-20151215/v5/client"
3131 "github.com/alibabacloud-go/tea/tea"
3232 "github.com/patrickmn/go-cache"
33+ "github.com/samber/lo"
34+ corev1 "k8s.io/api/core/v1"
3335 "sigs.k8s.io/controller-runtime/pkg/log"
3436 karpv1 "sigs.k8s.io/karpenter/pkg/apis/v1"
3537
@@ -39,7 +41,7 @@ import (
3941const defaultNodeLabel = "k8s.aliyun.com=true"
4042
4143type Provider interface {
42- GetNodeRegisterScript (context.Context , map [ string ] string , * v1alpha1.KubeletConfiguration ) (string , error )
44+ GetNodeRegisterScript (context.Context , string , * karpv1. NodeClaim , * v1alpha1.KubeletConfiguration ) (string , error )
4345 GetClusterCNI (context.Context ) (string , error )
4446 LivenessProbe (* http.Request ) error
4547}
@@ -104,10 +106,12 @@ func (p *DefaultProvider) GetClusterCNI(_ context.Context) (string, error) {
104106}
105107
106108func (p * DefaultProvider ) GetNodeRegisterScript (ctx context.Context ,
107- labels map [string ]string ,
109+ capacityType string ,
110+ nodeClaim * karpv1.NodeClaim ,
108111 kubeletCfg * v1alpha1.KubeletConfiguration ) (string , error ) {
112+ labels := lo .Assign (nodeClaim .Labels , map [string ]string {karpv1 .CapacityTypeLabelKey : capacityType })
109113 if cachedScript , ok := p .cache .Get (p .clusterID ); ok {
110- return p .resolveUserData (cachedScript .(string ), labels , kubeletCfg ), nil
114+ return p .resolveUserData (cachedScript .(string ), labels , nodeClaim , kubeletCfg ), nil
111115 }
112116
113117 reqPara := & ackclient.DescribeClusterAttachScriptsRequest {
@@ -126,14 +130,14 @@ func (p *DefaultProvider) GetNodeRegisterScript(ctx context.Context,
126130 }
127131
128132 p .cache .SetDefault (p .clusterID , s )
129- return p .resolveUserData (s , labels , kubeletCfg ), nil
133+ return p .resolveUserData (s , labels , nodeClaim , kubeletCfg ), nil
130134}
131135
132136func (p * DefaultProvider ) resolveUserData (respStr string ,
133137 labels map [string ]string ,
138+ nodeClaim * karpv1.NodeClaim ,
134139 kubeletCfg * v1alpha1.KubeletConfiguration ) string {
135140 cleanupStr := strings .ReplaceAll (respStr , "\r \n " , "" )
136-
137141 // TODO: now, the following code is quite ugly, make it clean in the future
138142 // Add labels
139143 labelsFormated := fmt .Sprintf ("%s,ack.aliyun.com=%s" , defaultNodeLabel , p .clusterID )
@@ -148,8 +152,19 @@ func (p *DefaultProvider) resolveUserData(respStr string,
148152 updatedCommand = fmt .Sprintf ("%s --node-config %s" , updatedCommand , cfg )
149153
150154 // Add taints
151- taint := karpv1 .UnregisteredNoExecuteTaint
152- updatedCommand = fmt .Sprintf ("%s --taints %s" , updatedCommand , taint .ToString ())
155+ taints := lo .Flatten ([][]corev1.Taint {
156+ nodeClaim .Spec .Taints ,
157+ nodeClaim .Spec .StartupTaints ,
158+ })
159+ if ! lo .ContainsBy (taints , func (t corev1.Taint ) bool {
160+ return t .MatchTaint (& karpv1 .UnregisteredNoExecuteTaint )
161+ }) {
162+ taints = append (taints , karpv1 .UnregisteredNoExecuteTaint )
163+ }
164+ updatedCommand = fmt .Sprintf ("%s --taints %s" ,
165+ updatedCommand , strings .Join (lo .Map (taints , func (t corev1.Taint , _ int ) string {
166+ return t .ToString ()
167+ }), "," ))
153168
154169 // Add bash script header
155170 finalScript := fmt .Sprintf ("#!/bin/bash\n \n %s" , updatedCommand )
0 commit comments