chore: cluster create support specify schedulingpolicy (#483)

Co-authored-by: yipeng1030 <[email protected]>
apecloud · Nov 8, 2024 · 509f77e · 509f77e
1 parent 0c7d968
commit 509f77e
Show file tree

Hide file tree

Showing 14 changed files with 214 additions and 24 deletions.
diff --git a/docs/user_docs/cli/kbcli_cluster_create_apecloud-mysql.md b/docs/user_docs/cli/kbcli_cluster_create_apecloud-mysql.md
@@ -41,7 +41,7 @@ kbcli cluster create apecloud-mysql NAME [flags]
       --storage-class-name string      Storage class name of the data volume
       --tenancy string                 The tenancy of cluster. Legal values [SharedNode, DedicatedNode]. (default "SharedNode")
       --termination-policy string      The termination policy of cluster. Legal values [DoNotTerminate, Halt, Delete, WipeOut]. (default "Delete")
-      --tolerations strings            Tolerations for cluster, such as "key=value:effect, key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
+      --tolerations strings            Tolerations for cluster, such as "key=value:effect,key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
       --topology-keys stringArray      Topology keys for affinity
       --version string                 Cluster version. (default "ac-mysql-8.0.30")
 ```

diff --git a/docs/user_docs/cli/kbcli_cluster_create_elasticsearch.md b/docs/user_docs/cli/kbcli_cluster_create_elasticsearch.md
@@ -39,7 +39,7 @@ kbcli cluster create elasticsearch NAME [flags]
       --storage float                  Storage size, the unit is Gi. Value range [1, 10000]. (default 20)
       --tenancy string                 The tenancy of cluster. Legal values [SharedNode, DedicatedNode]. (default "SharedNode")
       --termination-policy string      The termination policy of cluster. Legal values [DoNotTerminate, Halt, Delete, WipeOut]. (default "Delete")
-      --tolerations strings            Tolerations for cluster, such as "key=value:effect, key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
+      --tolerations strings            Tolerations for cluster, such as "key=value:effect,key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
       --topology-keys stringArray      Topology keys for affinity
       --version string                 The version of ElasticSearch. (default "8.8.2")
 ```

diff --git a/docs/user_docs/cli/kbcli_cluster_create_kafka.md b/docs/user_docs/cli/kbcli_cluster_create_kafka.md
@@ -51,7 +51,7 @@ kbcli cluster create kafka NAME [flags]
       --storage-enable                 Enable storage for Kafka.
       --tenancy string                 The tenancy of cluster. Legal values [SharedNode, DedicatedNode]. (default "SharedNode")
       --termination-policy string      The termination policy of cluster. Legal values [DoNotTerminate, Halt, Delete, WipeOut]. (default "Delete")
-      --tolerations strings            Tolerations for cluster, such as "key=value:effect, key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
+      --tolerations strings            Tolerations for cluster, such as "key=value:effect,key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
       --topology-keys stringArray      Topology keys for affinity
       --version string                 Cluster version. (default "kafka-3.3.2")
 ```

diff --git a/docs/user_docs/cli/kbcli_cluster_create_llm.md b/docs/user_docs/cli/kbcli_cluster_create_llm.md
@@ -42,7 +42,7 @@ kbcli cluster create llm NAME [flags]
       --replicas int                   The number of replicas, for standalone mode, the replicas is 1, for replication mode, the default replicas is 2. Value range [1, 5]. (default 1)
       --tenancy string                 The tenancy of cluster. Legal values [SharedNode, DedicatedNode]. (default "SharedNode")
       --termination-policy string      The termination policy of cluster. Legal values [DoNotTerminate, Halt, Delete, WipeOut]. (default "Delete")
-      --tolerations strings            Tolerations for cluster, such as "key=value:effect, key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
+      --tolerations strings            Tolerations for cluster, such as "key=value:effect,key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
       --topology-keys stringArray      Topology keys for affinity
       --url string                     Model URL, only work for CPU mode
       --version string                 Cluster version.

diff --git a/docs/user_docs/cli/kbcli_cluster_create_mongodb.md b/docs/user_docs/cli/kbcli_cluster_create_mongodb.md
@@ -40,7 +40,7 @@ kbcli cluster create mongodb NAME [flags]
       --storage-class-name string      Storage class name of the data volume
       --tenancy string                 The tenancy of cluster. Legal values [SharedNode, DedicatedNode]. (default "SharedNode")
       --termination-policy string      The termination policy of cluster. Legal values [DoNotTerminate, Halt, Delete, WipeOut]. (default "Delete")
-      --tolerations strings            Tolerations for cluster, such as "key=value:effect, key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
+      --tolerations strings            Tolerations for cluster, such as "key=value:effect,key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
       --topology-keys stringArray      Topology keys for affinity
       --version string                 Cluster version. Legal values [7.0.12, 6.0.16, 5.0.28, 4.4.29, 4.2.24, 4.0.28]. (default "6.0.16")
 ```

diff --git a/docs/user_docs/cli/kbcli_cluster_create_mysql.md b/docs/user_docs/cli/kbcli_cluster_create_mysql.md
@@ -39,7 +39,7 @@ kbcli cluster create mysql NAME [flags]
       --storage float                  Storage size, the unit is Gi. Value range [1, 10000]. (default 20)
       --tenancy string                 The tenancy of cluster. Legal values [SharedNode, DedicatedNode]. (default "SharedNode")
       --termination-policy string      The termination policy of cluster. Legal values [DoNotTerminate, Halt, Delete, WipeOut]. (default "Delete")
-      --tolerations strings            Tolerations for cluster, such as "key=value:effect, key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
+      --tolerations strings            Tolerations for cluster, such as "key=value:effect,key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
       --topology-keys stringArray      Topology keys for affinity
       --version string                 Cluster version, run "kbcli cv list --devel" to see all versions. Legal values [mysql-8.4, mysql-8.0, mysql-5.7]. (default "mysql-8.0")
 ```

diff --git a/docs/user_docs/cli/kbcli_cluster_create_postgresql.md b/docs/user_docs/cli/kbcli_cluster_create_postgresql.md
@@ -40,7 +40,7 @@ kbcli cluster create postgresql NAME [flags]
       --storage-class-name string      Storage class name of the data volume
       --tenancy string                 The tenancy of cluster. Legal values [SharedNode, DedicatedNode]. (default "SharedNode")
       --termination-policy string      The termination policy of cluster. Legal values [DoNotTerminate, Halt, Delete, WipeOut]. (default "Delete")
-      --tolerations strings            Tolerations for cluster, such as "key=value:effect, key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
+      --tolerations strings            Tolerations for cluster, such as "key=value:effect,key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
       --topology-keys stringArray      Topology keys for affinity
       --version string                 service version. (default "15.7.0")
 ```

diff --git a/docs/user_docs/cli/kbcli_cluster_create_qdrant.md b/docs/user_docs/cli/kbcli_cluster_create_qdrant.md
@@ -39,7 +39,7 @@ kbcli cluster create qdrant NAME [flags]
       --storage-class-name string      Storage class name of the data volume
       --tenancy string                 The tenancy of cluster. Legal values [SharedNode, DedicatedNode]. (default "SharedNode")
       --termination-policy string      The termination policy of cluster. Legal values [DoNotTerminate, Halt, Delete, WipeOut]. (default "Delete")
-      --tolerations strings            Tolerations for cluster, such as "key=value:effect, key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
+      --tolerations strings            Tolerations for cluster, such as "key=value:effect,key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
       --topology-keys stringArray      Topology keys for affinity
       --version string                 The version of Qdrant. (default "1.10.0")
 ```

diff --git a/docs/user_docs/cli/kbcli_cluster_create_redis.md b/docs/user_docs/cli/kbcli_cluster_create_redis.md
@@ -47,7 +47,7 @@ kbcli cluster create redis NAME [flags]
       --storage-class-name string         Storage class name of the data volume
       --tenancy string                    The tenancy of cluster. Legal values [SharedNode, DedicatedNode]. (default "SharedNode")
       --termination-policy string         The termination policy of cluster. Legal values [DoNotTerminate, Halt, Delete, WipeOut]. (default "Delete")
-      --tolerations strings               Tolerations for cluster, such as "key=value:effect, key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
+      --tolerations strings               Tolerations for cluster, such as "key=value:effect,key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
       --topology-keys stringArray         Topology keys for affinity
       --twemproxy.cpu float               twemproxy component cpu cores. Value range [0.1, 8]. (default 0.2)
       --twemproxy.enabled                 Whether have twemproxy component, default is false

diff --git a/docs/user_docs/cli/kbcli_cluster_create_xinference.md b/docs/user_docs/cli/kbcli_cluster_create_xinference.md
@@ -40,7 +40,7 @@ kbcli cluster create xinference NAME [flags]
       --shm-size string                shm size (default "64Mi")
       --tenancy string                 The tenancy of cluster. Legal values [SharedNode, DedicatedNode]. (default "SharedNode")
       --termination-policy string      The termination policy of cluster. Legal values [DoNotTerminate, Halt, Delete, WipeOut]. (default "Delete")
-      --tolerations strings            Tolerations for cluster, such as "key=value:effect, key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
+      --tolerations strings            Tolerations for cluster, such as "key=value:effect,key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'
       --topology-keys stringArray      Topology keys for affinity
 ```
 

diff --git a/pkg/cmd/cluster/cluster_test.go b/pkg/cmd/cluster/cluster_test.go
@@ -102,6 +102,51 @@ var _ = Describe("Cluster", func() {
 			Expect(o.Name).ShouldNot(BeEmpty())
 			Expect(o.Run()).Should(Succeed())
 		})
+		It("should apply SharedNode tenancy and Preferred PodAntiAffinity", func() {
+			o, err := NewSubCmdsOptions(createOptions, clusterType)
+			Expect(err).Should(Succeed())
+			o.Tenancy = "SharedNode"
+			o.TopologyKeys = []string{"test-topology1", "test-topology2"}
+			o.NodeLabels = map[string]string{"environment": "test-env", "region": "test-region"}
+			o.TolerationsRaw = []string{"testKey1=testValue1:NoSchedule", "testKey2=testValue2:NoExecute"}
+			o.PodAntiAffinity = "Preferred"
+
+			Expect(o).ShouldNot(BeNil())
+			Expect(o.ChartInfo).ShouldNot(BeNil())
+			o.Format = printer.YAML
+
+			Expect(o.CreateOptions.Complete()).To(Succeed())
+			o.Client = testing.FakeClientSet()
+			fakeDiscovery1, _ := o.Client.Discovery().(*fakediscovery.FakeDiscovery)
+			fakeDiscovery1.FakedServerVersion = &version.Info{Major: "1", Minor: "27", GitVersion: "v1.27.0"}
+			Expect(o.Complete(nil)).To(Succeed())
+			Expect(o.Validate()).To(Succeed())
+			Expect(o.Name).ShouldNot(BeEmpty())
+			Expect(o.Run()).Should(Succeed())
+		})
+
+		It("should apply DedicatedNode tenancy and Required PodAntiAffinity", func() {
+			o, err := NewSubCmdsOptions(createOptions, clusterType)
+			Expect(err).Should(Succeed())
+			o.Tenancy = "DedicatedNode"
+			o.TopologyKeys = []string{"test-region", "test-zone"}
+			o.NodeLabels = map[string]string{"cluster": "test-cluster", "env": "test-production"}
+			o.TolerationsRaw = []string{"testKey3=testValue3:NoSchedule", "testKey4=testValue4:NoExecute"}
+			o.PodAntiAffinity = "Required"
+
+			Expect(o).ShouldNot(BeNil())
+			Expect(o.ChartInfo).ShouldNot(BeNil())
+			o.Format = printer.YAML
+
+			Expect(o.CreateOptions.Complete()).To(Succeed())
+			o.Client = testing.FakeClientSet()
+			fakeDiscovery2, _ := o.Client.Discovery().(*fakediscovery.FakeDiscovery)
+			fakeDiscovery2.FakedServerVersion = &version.Info{Major: "1", Minor: "27", GitVersion: "v1.27.0"}
+			Expect(o.Complete(nil)).To(Succeed())
+			Expect(o.Validate()).To(Succeed())
+			Expect(o.Name).ShouldNot(BeEmpty())
+			Expect(o.Run()).Should(Succeed())
+		})
 	})
 
 	Context("create validate", func() {

diff --git a/pkg/cmd/cluster/create_subcmds.go b/pkg/cmd/cluster/create_subcmds.go
@@ -21,13 +21,16 @@ package cluster
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"os"
 	"regexp"
 
 	"github.com/spf13/cobra"
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	cmdutil "k8s.io/kubectl/pkg/cmd/util"
 
@@ -61,6 +64,7 @@ type CreateSubCmdsOptions struct {
 	NodeLabels     map[string]string `json:"nodeLabels,omitempty"`
 	Tenancy        string            `json:"tenancy"`
 	TolerationsRaw []string          `json:"-"`
+	Tolerations    []corev1.Toleration
 
 	// SkipSchemaValidation is used to skip the schema validation of the helm chart.
 	SkipSchemaValidation bool `json:"-"`
@@ -123,11 +127,10 @@ func buildCreateSubCmds(createOptions *action.CreateOptions) []*cobra.Command {
 		util.CheckErr(addCreateFlags(cmd, o.Factory, o.ChartInfo, t.String()))
 
 		// Schedule policy
-		// TODO: implement them, and check whether the flag has been defined
 		cmd.Flags().StringVar(&o.PodAntiAffinity, "pod-anti-affinity", "Preferred", "Pod anti-affinity type, one of: (Preferred, Required)")
 		cmd.Flags().StringArrayVar(&o.TopologyKeys, "topology-keys", nil, "Topology keys for affinity")
 		cmd.Flags().StringToStringVar(&o.NodeLabels, "node-labels", nil, "Node label selector")
-		cmd.Flags().StringSliceVar(&o.TolerationsRaw, "tolerations", nil, `Tolerations for cluster, such as "key=value:effect, key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'`)
+		cmd.Flags().StringSliceVar(&o.TolerationsRaw, "tolerations", nil, `Tolerations for cluster, such as "key=value:effect,key:effect", for example '"engineType=mongo:NoSchedule", "diskType:NoSchedule"'`)
 		if cmd.Flag("tenancy") == nil {
 			cmd.Flags().StringVar(&o.Tenancy, "tenancy", "SharedNode", "Tenancy options, one of: (SharedNode, DedicatedNode)")
 		}
@@ -196,6 +199,22 @@ func (o *CreateSubCmdsOptions) Complete(cmd *cobra.Command) error {
 	if o.ChartInfo.ClusterDef == "" && len(o.ChartInfo.ComponentDef) == 0 {
 		return fmt.Errorf("cannot find clusterDef in cluster spec or componentDef in componentSpecs or shardingSpecs")
 	}
+
+	// Build tolerations if raw toleration rules are configured
+	if o.TolerationsRaw != nil {
+		tolerationsResult, err := util.BuildTolerations(o.TolerationsRaw)
+		if err != nil {
+			return err
+		}
+		jsonData, err := json.Marshal(tolerationsResult)
+		if err != nil {
+			return err
+		}
+		err = json.Unmarshal(jsonData, &o.Tolerations)
+		if err != nil {
+			return err
+		}
+	}
 	return nil
 }
 
@@ -226,12 +245,36 @@ func (o *CreateSubCmdsOptions) Run() error {
 		return nil, fmt.Errorf("failed to find cluster object from manifests rendered from %s chart", o.ClusterType)
 	}
 
+	clusterObj, err := getClusterObj()
+	if err != nil {
+		return err
+	}
+
+	spec, _ := clusterObj.Object["spec"].(map[string]interface{})
+	if compSpec, ok := spec["componentSpecs"].([]interface{}); ok {
+		for i := range compSpec {
+			comp := compSpec[i].(map[string]interface{})
+			if err := o.applySchedulingPolicyToComponent(comp); err != nil {
+				return err
+			}
+			compSpec[i] = comp
+		}
+	}
+
+	if shardingSpec, ok := spec["shardings"].([]interface{}); ok {
+		for i := range shardingSpec {
+			shard := shardingSpec[i].(map[string]interface{})
+			if compSpec, ok := shard["template"].(map[string]interface{}); ok {
+				if err := o.applySchedulingPolicyToComponent(compSpec); err != nil {
+					return err
+				}
+				shard["template"] = compSpec
+			}
+		}
+	}
+
 	// only edits the cluster object, other dependency objects are created directly
 	if o.EditBeforeCreate {
-		clusterObj, err := getClusterObj()
-		if err != nil {
-			return err
-		}
 		customEdit := action.NewCustomEditOptions(o.Factory, o.IOStreams, "create")
 		if err = customEdit.Run(clusterObj); err != nil {
 			return err
@@ -319,3 +362,17 @@ func (o *CreateSubCmdsOptions) getClusterObj(objs []*objectInfo) (*unstructured.
 	}
 	return nil, fmt.Errorf("failed to find cluster object from manifests rendered from %s chart", o.ClusterType)
 }
+
+func (o *CreateSubCmdsOptions) applySchedulingPolicyToComponent(item map[string]interface{}) error {
+	if compName, ok := item["name"]; ok {
+		schedulingPolicy, needSet := util.BuildSchedulingPolicy(o.Tenancy, o.Name, compName.(string), o.Tolerations, o.NodeLabels, o.PodAntiAffinity, o.TopologyKeys)
+		if needSet {
+			converted, err := runtime.DefaultUnstructuredConverter.ToUnstructured(schedulingPolicy)
+			if err != nil {
+				return err
+			}
+			_ = unstructured.SetNestedField(item, converted, "schedulingPolicy")
+		}
+	}
+	return nil
+}
diff --git a/pkg/cmd/cluster/register_test.go b/pkg/cmd/cluster/register_test.go
@@ -91,16 +91,14 @@ var _ = Describe("cluster register", func() {
 
 	Context("test register cluster chart", func() {
 		var (
-			source              = "https://github.com/apecloud/helm-charts/releases/download/mysql-cluster-1.0.0-alpha.0/mysql-cluster-1.0.0-alpha.0.tgz"
-			engine              = "mysql"
-			apecloudMysqlEngine = "apecloud-mysql"
-			version             = "1.0.0-alpha.0"
-			repo                = types.ClusterChartsRepoURL
+			source  = "https://github.com/apecloud/helm-charts/releases/download/apecloud-mysql-cluster-1.0.0-alpha.0/apecloud-mysql-cluster-1.0.0-alpha.0.tgz"
+			engine  = "apecloud-mysql"
+			version = "1.0.0-alpha.0"
+			repo    = types.ClusterChartsRepoURL
 		)
 
 		AfterEach(func() {
 			cluster.ClearCharts(cluster.ClusterType(engine))
-			cluster.ClearCharts(cluster.ClusterType(apecloudMysqlEngine))
 		})
 
 		It("test register chart by source", func() {
@@ -112,9 +110,9 @@ var _ = Describe("cluster register", func() {
 		})
 
 		It("test register built-in chart and test validate", func() {
-			Expect(RegisterClusterChart(tf, streams, "", apecloudMysqlEngine, version, repo)).Should(Succeed())
+			Expect(RegisterClusterChart(tf, streams, "", engine, version, repo)).Should(Succeed())
 			validatedChart := &cluster.TypeInstance{
-				Name:      cluster.ClusterType(apecloudMysqlEngine),
+				Name:      cluster.ClusterType(engine),
 				URL:       "",
 				Alias:     "",
 				ChartName: filepath.Base(source),