Merge pull request #834 from dayuy/main

fix: also add the argument `revision` for HuggingFaceRepo
kubeagi · Mar 13, 2024 · 001b207 · 001b207
2 parents 433ac09 + 7995b5b
commit 001b207
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/pkg/worker/runner.go b/pkg/worker/runner.go
@@ -217,13 +217,15 @@ func (runner *RunnerFastchatVLLM) Build(ctx context.Context, model *arcadiav1alp
 		if err := runner.c.Get(ctx, types.NamespacedName{Namespace: *model.Namespace, Name: model.Name}, &m); err != nil {
 			return nil, err
 		}
+		if m.Spec.Revision != "" {
+			extraAgrs += fmt.Sprintf(" --revision %s", m.Spec.Revision)
+		}
 		if m.Spec.HuggingFaceRepo != "" {
 			modelFileDir = m.Spec.HuggingFaceRepo
 		}
 		if m.Spec.ModelScopeRepo != "" {
 			modelFileDir = m.Spec.ModelScopeRepo
 			additionalEnvs = append(additionalEnvs, corev1.EnvVar{Name: "FASTCHAT_USE_MODELSCOPE", Value: "True"}, corev1.EnvVar{Name: "VLLM_USE_MODELSCOPE", Value: "True"})
-			extraAgrs += fmt.Sprintf(" --revision %s", m.Spec.Revision)
 		}
 	}