Merge pull request kubeagi#370 from bjwswang/gql

feat: able to configure worker type when create and update
nkwangleiGIT · Dec 14, 2023 · 3783c14 · 3783c14
2 parents 1891ebf + 5df65ad
commit 3783c14
Show file tree

Hide file tree

Showing 6 changed files with 174 additions and 2 deletions.
diff --git a/api/base/v1alpha1/worker.go b/api/base/v1alpha1/worker.go
@@ -35,6 +35,10 @@ const (
 	LabelWorkerType = Group + "/worker-type"
 )
 
+func DefaultWorkerType() WorkerType {
+	return WorkerTypeFastchatNormal
+}
+
 func (worker Worker) Type() WorkerType {
 	if worker.Spec.Type == "" {
 		// use `fastchat` by default

diff --git a/graphql-server/go-server/graph/generated/generated.go b/graphql-server/go-server/graph/generated/generated.go
diff --git a/graphql-server/go-server/graph/generated/models_gen.go b/graphql-server/go-server/graph/generated/models_gen.go
diff --git a/graphql-server/go-server/graph/schema/worker.gql b/graphql-server/go-server/graph/schema/worker.gql
@@ -19,6 +19,7 @@ query listWorkers($input: ListWorkerInput!){
           status
           message
           updateTimestamp
+          type
           model {
             name
             namespace
@@ -50,6 +51,7 @@ query getWorker($name: String!, $namespace: String!) {
           creator
           displayName
           description
+          type
           status
           message
           updateTimestamp
@@ -82,6 +84,7 @@ mutation createWorker($input: CreateWorkerInput!) {
           creator
           displayName
           description
+          type
           status
           message
           updateTimestamp
@@ -114,6 +117,7 @@ mutation updateWorker($input: UpdateWorkerInput) {
           creator
           displayName
           description
+          type
           status
           message
           updateTimestamp

diff --git a/graphql-server/go-server/graph/schema/worker.graphqls b/graphql-server/go-server/graph/schema/worker.graphqls
@@ -65,6 +65,15 @@ type Worker {
     """更新时间"""
     updateTimestamp: Time
 
+    """
+    Worker类型
+    支持两种类型: 
+    - "fastchat" : fastchat提供的通用的推理服务模式
+    - "fastchat-vllm" : fastchat提供的采用VLLM推理加速的推理服务模式
+    规则: 如果为空，则默认为 "fastchat"
+    """
+    type: String
+
     """
     worker对应的模型
     规则: 相同namespace下的模型名称
@@ -109,6 +118,15 @@ input CreateWorkerInput{
     """模型资源描述"""
     description: String
 
+    """
+    Worker类型
+    支持两种类型: 
+    - "fastchat" : fastchat提供的通用的推理服务模式
+    - "fastchat-vllm" : fastchat提供的采用VLLM推理加速的推理服务模式
+    规则: 如果为空，则默认为 "fastchat"
+    """
+    type: String
+
     """
     worker对应的模型
     规则: 必须指定模型准确的namespace
@@ -140,6 +158,15 @@ input UpdateWorkerInput {
     """模型资源描述"""
     description: String
 
+    """
+    Worker类型
+    支持两种类型: 
+    - "fastchat" : fastchat提供的通用的推理服务模式
+    - "fastchat-vllm" : fastchat提供的采用VLLM推理加速的推理服务模式
+    规则: 如果为空，则不更新；如果type类型与当前类型相同，则不更新
+    """
+    type: String
+
     """
     worker运行所需的资源
     """

diff --git a/graphql-server/go-server/pkg/worker/worker.go b/graphql-server/go-server/pkg/worker/worker.go
@@ -73,6 +73,8 @@ func worker2model(ctx context.Context, c dynamic.Interface, obj *unstructured.Un
 		NvidiaGpu: &nvidiaGPUStr,
 	}
 
+	workerType := string(worker.Type())
+
 	// wrap Worker
 	w := generated.Worker{
 		ID:                &id,
@@ -83,6 +85,7 @@ func worker2model(ctx context.Context, c dynamic.Interface, obj *unstructured.Un
 		Annotations:       graphqlutils.MapStr2Any(obj.GetAnnotations()),
 		DisplayName:       &worker.Spec.DisplayName,
 		Description:       &worker.Spec.Description,
+		Type:              &workerType,
 		Status:            &status,
 		CreationTimestamp: &creationtimestamp,
 		UpdateTimestamp:   &updateTime,
@@ -128,6 +131,12 @@ func CreateWorker(ctx context.Context, c dynamic.Interface, input generated.Crea
 		}
 	}
 
+	// Use `fastchat` as the default worker type
+	workerType := v1alpha1.DefaultWorkerType()
+	if input.Type != nil {
+		workerType = v1alpha1.WorkerType(*input.Type)
+	}
+
 	worker := v1alpha1.Worker{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      input.Name,
@@ -142,6 +151,7 @@ func CreateWorker(ctx context.Context, c dynamic.Interface, input generated.Crea
 				DisplayName: displayName,
 				Description: description,
 			},
+			Type: workerType,
 			Model: &v1alpha1.TypedObjectReference{
 				Name:      input.Model.Name,
 				Namespace: &modelNs,
@@ -196,6 +206,13 @@ func UpdateWorker(ctx context.Context, c dynamic.Interface, input *generated.Upd
 		worker.Spec.Description = *input.Description
 	}
 
+	// worker type
+	if input.Type != nil {
+		if worker.Type() != v1alpha1.WorkerType(*input.Type) {
+			worker.Spec.Type = v1alpha1.WorkerType(*input.Type)
+		}
+	}
+
 	// resources
 	if input.Resources != nil {
 		// cpu & memory