Merge pull request #353 from bjwswang/endpoint

feat: utlize worker's UID as the model registration name
kubeagi · Dec 12, 2023 · aa13ca2 · aa13ca2
2 parents 5362f42 + acd3b86
commit aa13ca2
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 37 deletions.
diff --git a/api/base/v1alpha1/worker.go b/api/base/v1alpha1/worker.go
@@ -17,10 +17,6 @@ limitations under the License.
 package v1alpha1
 
 import (
-	"crypto/sha256"
-	"encoding/hex"
-	"fmt"
-
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
@@ -48,14 +44,8 @@ func (worker Worker) Type() WorkerType {
 }
 
 // MakeRegistrationModelName generates a model name used to register itself into fastchat controller
-// {MODEL_NAME}-{WORKER_NAME}-{WORKER_NAMESPACE}
-// When multiple workers run a same model,its will overwrite the previuos registerd worker address.That's why we use a combined model name for registration.
 func (worker Worker) MakeRegistrationModelName() string {
-	modelRef := worker.Spec.Model
-	if modelRef == nil {
-		return ""
-	}
-	return fmt.Sprintf("%s-%s-%s", modelRef.Name, worker.Name, worker.Namespace)
+	return string(worker.UID)
 }
 
 func (worker Worker) PendingCondition() Condition {
@@ -91,27 +81,11 @@ func (worker Worker) ErrorCondition(msg string) Condition {
 	}
 }
 
-func (worker Worker) generateUniqueName() string {
-	// Create a new SHA-256 hasher
-	hasher := sha256.New()
-
-	// Write the input string to the hasher
-	hasher.Write([]byte(worker.Name))
-
-	// Calculate the hash sum
-	hashSum := hasher.Sum(nil)
-
-	// Convert the hash sum to a hexadecimal string
-	hashString := hex.EncodeToString(hashSum)
-
-	return hashString[:10]
-}
-
 func (worker Worker) BuildEmbedder() *Embedder {
 	return &Embedder{
 		ObjectMeta: metav1.ObjectMeta{
 			Namespace: worker.Namespace,
-			Name:      worker.generateUniqueName() + "-worker",
+			Name:      worker.Name,
 		},
 		Spec: EmbedderSpec{
 			CommonSpec: CommonSpec{
@@ -136,7 +110,7 @@ func (worker Worker) BuildLLM() *LLM {
 	return &LLM{
 		ObjectMeta: metav1.ObjectMeta{
 			Namespace: worker.Namespace,
-			Name:      worker.generateUniqueName() + "-worker",
+			Name:      worker.Name,
 		},
 		Spec: LLMSpec{
 			CommonSpec: CommonSpec{

diff --git a/config/samples/arcadia_v1alpha1_worker_baichuan2-7b.yaml b/config/samples/arcadia_v1alpha1_worker_baichuan2-7b.yaml
@@ -1,23 +1,23 @@
 apiVersion: arcadia.kubeagi.k8s.com.cn/v1alpha1
 kind: Model
 metadata:
-  name: baichuan2-7b
+  name: baichuan2-7b-chat
   namespace: arcadia
 spec:
-  displayName: "baichuan2-7b"
+  displayName: "baichuan2-7b-chat"
   description: "百川智能发布的大语言模型模型,同时提供embedding模型能力"
   types: "llm,embedding"
 ---
 apiVersion: arcadia.kubeagi.k8s.com.cn/v1alpha1
 kind: Worker
 metadata:
-  name: worker-baichuan
+  name: baichuan2-7b-chat
   namespace: arcadia
 spec:
   type: "fastchat"
   model:
     kind: "Models"
-    name: "baichuan2-7b"
+    name: "baichuan2-7b-chat"
   resources:
     limits:
       nvidia.com/gpu: "1" # request 1 GPU
diff --git a/config/samples/arcadia_v1alpha1_worker_bge-large-zh-v1.5.yaml b/config/samples/arcadia_v1alpha1_worker_bge-large-zh-v1.5.yaml
@@ -11,7 +11,7 @@ spec:
 apiVersion: arcadia.kubeagi.k8s.com.cn/v1alpha1
 kind: Worker
 metadata:
-  name: worker-bge
+  name: bge-large-zh
   namespace: arcadia
 spec:
   type: "fastchat"

diff --git a/examples/chat_with_worker/main.go b/examples/chat_with_worker/main.go
@@ -14,7 +14,7 @@ func main() {
 		openai.WithToken("fake"),
 		// update base url to fastchat api server
 		openai.WithBaseURL("http://arcadia-fastchat.172.22.96.167.nip.io/v1"),
-		openai.WithModel("baichuan2-7b"),
+		openai.WithModel("fb219b5f-8f3e-49e1-8d5b-f0c6da481186"),
 	)
 	if err != nil {
 		log.Fatal(err)

diff --git a/pkg/worker/runner.go b/pkg/worker/runner.go
@@ -78,13 +78,13 @@ func (runner *RunnerFastchat) Build(ctx context.Context, model *arcadiav1alpha1.
 			"/bin/bash",
 			"-c",
 			`echo "Run model worker..."
-python3.9 -m fastchat.serve.model_worker --model-names $FASTCHAT_MODEL_NAME-$FASTCHAT_WORKER_NAME-$FASTCHAT_WORKER_NAMESPACE \
+python3.9 -m fastchat.serve.model_worker --model-names $FASTCHAT_REGISTRATION_MODEL_NAME \
 --model-path /data/models/$FASTCHAT_MODEL_NAME --worker-address $FASTCHAT_WORKER_ADDRESS \
 --controller-address $FASTCHAT_CONTROLLER_ADDRESS \
 --host 0.0.0.0 --port 21002`},
 		Env: []corev1.EnvVar{
 			{Name: "FASTCHAT_WORKER_NAMESPACE", Value: runner.w.Namespace},
-			{Name: "FASTCHAT_WORKER_NAME", Value: runner.w.Name},
+			{Name: "FASTCHAT_REGISTRATION_MODEL_NAME", Value: runner.w.MakeRegistrationModelName()},
 			{Name: "FASTCHAT_MODEL_NAME", Value: model.Name},
 			{Name: "FASTCHAT_WORKER_ADDRESS", Value: fmt.Sprintf("http://%s.%s.svc.cluster.local:21002", runner.w.Name+WokerCommonSuffix, runner.w.Namespace)},
 			{Name: "FASTCHAT_CONTROLLER_ADDRESS", Value: gw.Controller},