feat(llm): auto append backend_parameters in vllm & sglang (#24876)

2026-10-31 02:38:07 +08:00 · 2026-05-22 17:01:54 +08:00
parent ae431939fd
commit 588c38ec74
3 changed files with 222 additions and 10 deletions
@@ -0,0 +1,138 @@
+package llm_container
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+)
+
+type runtimeArg struct {
+	Key   string
+	Value string
+}
+
+func parseBackendParameterArgs(raw string, validate func(string) error) ([]runtimeArg, error) {
+	raw = strings.TrimSpace(raw)
+	if raw == "" {
+		return nil, nil
+	}
+	items := []string{}
+	if err := json.Unmarshal([]byte(raw), &items); err != nil {
+		items = []string{raw}
+	}
+	args := make([]runtimeArg, 0, len(items))
+	errs := make([]string, 0)
+	for _, item := range items {
+		arg, ok, err := parseBackendParameterArg(item)
+		if err != nil {
+			errs = append(errs, err.Error())
+			continue
+		}
+		if !ok {
+			continue
+		}
+		if err := validate(arg.Key); err != nil {
+			errs = append(errs, err.Error())
+			continue
+		}
+		args = append(args, arg)
+	}
+	normalized, err := normalizeRuntimeArgs(args, validate)
+	if err != nil {
+		errs = append(errs, err.Error())
+	}
+	if len(errs) > 0 {
+		return normalized, fmt.Errorf("%s", strings.Join(errs, "; "))
+	}
+	return normalized, nil
+}
+
+func parseBackendParameterArg(item string) (runtimeArg, bool, error) {
+	item = strings.TrimSpace(item)
+	if item == "" {
+		return runtimeArg{}, false, nil
+	}
+	if strings.HasPrefix(item, "--") {
+		item = strings.TrimSpace(strings.TrimPrefix(item, "--"))
+	} else if strings.HasPrefix(item, "-") {
+		return runtimeArg{}, false, fmt.Errorf("invalid backend parameter %q: only long flags are supported", item)
+	}
+	if item == "" {
+		return runtimeArg{}, false, fmt.Errorf("invalid backend parameter: empty flag")
+	}
+
+	key := item
+	value := ""
+	if idx := strings.Index(item, "="); idx >= 0 {
+		key = strings.TrimSpace(item[:idx])
+		value = item[idx+1:]
+	} else if fields := strings.Fields(item); len(fields) > 1 {
+		key = fields[0]
+		value = strings.TrimSpace(item[len(key):])
+	}
+	if strings.HasPrefix(key, "--") {
+		key = strings.TrimPrefix(key, "--")
+	}
+	return runtimeArg{Key: key, Value: value}, true, nil
+}
+
+func normalizeRuntimeArgs(args []runtimeArg, validate func(string) error) ([]runtimeArg, error) {
+	if len(args) == 0 {
+		return nil, nil
+	}
+	out := make([]runtimeArg, 0, len(args))
+	indexByKey := make(map[string]int, len(args))
+	for _, arg := range args {
+		key := strings.TrimSpace(arg.Key)
+		if err := validate(key); err != nil {
+			return nil, err
+		}
+		next := runtimeArg{Key: key, Value: arg.Value}
+		if idx, ok := indexByKey[key]; ok {
+			out[idx] = next
+			continue
+		}
+		indexByKey[key] = len(out)
+		out = append(out, next)
+	}
+	return out, nil
+}
+
+func mergeRuntimeArgs(base, overrides []runtimeArg, validate func(string) error) ([]runtimeArg, error) {
+	out := make([]runtimeArg, 0, len(base)+len(overrides))
+	indexByKey := make(map[string]int, len(base)+len(overrides))
+	appendNormalized := func(items []runtimeArg) error {
+		normalized, err := normalizeRuntimeArgs(items, validate)
+		if err != nil {
+			return err
+		}
+		for _, arg := range normalized {
+			if idx, ok := indexByKey[arg.Key]; ok {
+				out[idx] = arg
+				continue
+			}
+			indexByKey[arg.Key] = len(out)
+			out = append(out, arg)
+		}
+		return nil
+	}
+	if err := appendNormalized(base); err != nil {
+		return nil, err
+	}
+	if err := appendNormalized(overrides); err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func appendRuntimeFlags(flags []string, args []runtimeArg) []string {
+	for _, arg := range args {
+		flagName := "--" + arg.Key
+		if arg.Value == "" {
+			flags = append(flags, flagName)
+			continue
+		}
+		flags = append(flags, fmt.Sprintf("%s %s", flagName, shellQuoteSingle(arg.Value)))
+	}
+	return flags
+}
@@ -219,7 +219,11 @@ func (s *sglang) GetContainerSpec(ctx context.Context, llm *models.SLLM, image *
 	if eff := s.GetEffectiveSpec(llm, sku); eff != nil {
 		effSpec = eff.(*api.LLMSpecSGLang)
 	}
-	startScript := buildSGLangEntrypointScript(len(postOverlays) > 0, tensorParallelSize, effSpec)
+	backendParameters := ""
+	if sku != nil {
+		backendParameters = sku.BackendParameters
+	}
+	startScript := buildSGLangEntrypointScript(len(postOverlays) > 0, tensorParallelSize, backendParameters, effSpec)
 	envs := []*commonapi.ContainerKeyValue{
 		{
 			Key:   "HUGGING_FACE_HUB_CACHE",
@@ -593,7 +597,38 @@ func appendSGLangCustomizedFlags(flags []string, effSpec *api.LLMSpecSGLang) []s
 	return flags
 }

-func buildSGLangServeFlagsWithModelExpr(modelExpr string, servedModelNameExpr string, tensorParallelSize int, effSpec *api.LLMSpecSGLang) []string {
+func sglangCustomizedArgsToRuntime(args []*api.SGLangCustomizedArg) []runtimeArg {
+	if len(args) == 0 {
+		return nil
+	}
+	out := make([]runtimeArg, 0, len(args))
+	for _, arg := range args {
+		if arg == nil {
+			continue
+		}
+		out = append(out, runtimeArg{Key: arg.Key, Value: arg.Value})
+	}
+	return out
+}
+
+func appendSGLangRuntimeFlags(flags []string, backendParameters string, effSpec *api.LLMSpecSGLang) []string {
+	backendArgs, err := parseBackendParameterArgs(backendParameters, validateSGLangArgKey)
+	if err != nil {
+		log.Errorf("parse sglang backend parameters: %v", err)
+	}
+	var customizedArgs []runtimeArg
+	if effSpec != nil {
+		customizedArgs = sglangCustomizedArgsToRuntime(effSpec.CustomizedArgs)
+	}
+	mergedArgs, err := mergeRuntimeArgs(backendArgs, customizedArgs, validateSGLangArgKey)
+	if err != nil {
+		log.Errorf("merge sglang runtime args: %v", err)
+		return flags
+	}
+	return appendRuntimeFlags(flags, mergedArgs)
+}
+
+func buildSGLangServeFlagsWithModelExpr(modelExpr string, servedModelNameExpr string, tensorParallelSize int, backendParameters string, effSpec *api.LLMSpecSGLang) []string {
 	flags := []string{
 		fmt.Sprintf("--model-path %s", modelExpr),
 		fmt.Sprintf("--served-model-name %s", servedModelNameExpr),
@@ -601,20 +636,21 @@ func buildSGLangServeFlagsWithModelExpr(modelExpr string, servedModelNameExpr st
 		fmt.Sprintf("--port %d", api.LLM_SGLANG_DEFAULT_PORT),
 		fmt.Sprintf("--tp-size %d", tensorParallelSize),
 	}
-	return appendSGLangCustomizedFlags(flags, effSpec)
+	return appendSGLangRuntimeFlags(flags, backendParameters, effSpec)
 }

-func buildSGLangServeFlags(modelPath string, tensorParallelSize int, effSpec *api.LLMSpecSGLang) []string {
+func buildSGLangServeFlags(modelPath string, tensorParallelSize int, backendParameters string, effSpec *api.LLMSpecSGLang) []string {
 	modelQuoted := shellQuoteSingle(modelPath)
 	return buildSGLangServeFlagsWithModelExpr(
 		modelQuoted,
 		fmt.Sprintf(`"$(basename %s)"`, modelQuoted),
 		tensorParallelSize,
+		backendParameters,
 		effSpec,
 	)
 }

-func buildSGLangEntrypointScript(hasMountedModels bool, tensorParallelSize int, effSpec *api.LLMSpecSGLang) string {
+func buildSGLangEntrypointScript(hasMountedModels bool, tensorParallelSize int, backendParameters string, effSpec *api.LLMSpecSGLang) string {
 	modelsPath := shellQuoteSingle(api.LLM_SGLANG_MODELS_PATH)
 	if !hasMountedModels {
 		return fmt.Sprintf("mkdir -p %s && exec sleep infinity", modelsPath)
@@ -628,6 +664,7 @@ func buildSGLangEntrypointScript(hasMountedModels bool, tensorParallelSize int,
 		`"$model"`,
 		`"$(basename "$model")"`,
 		tensorParallelSize,
+		backendParameters,
 		effSpec,
 	), " ")
 	return strings.Join([]string{
@@ -150,27 +150,59 @@ func appendVLLMCustomizedFlags(flags []string, effSpec *api.LLMSpecVllm) []strin
 	return flags
 }

-func buildVLLMServeFlagsWithModelExpr(modelExpr string, servedModelNameExpr string, tensorParallelSize int, effSpec *api.LLMSpecVllm) []string {
+func vllmCustomizedArgsToRuntime(args []*api.VllmCustomizedArg) []runtimeArg {
+	if len(args) == 0 {
+		return nil
+	}
+	out := make([]runtimeArg, 0, len(args))
+	for _, arg := range args {
+		if arg == nil {
+			continue
+		}
+		out = append(out, runtimeArg{Key: arg.Key, Value: arg.Value})
+	}
+	return out
+}
+
+func appendVLLMRuntimeFlags(flags []string, backendParameters string, effSpec *api.LLMSpecVllm) []string {
+	backendArgs, err := parseBackendParameterArgs(backendParameters, validateVLLMArgKey)
+	if err != nil {
+		log.Errorf("parse vllm backend parameters: %v", err)
+	}
+	var customizedArgs []runtimeArg
+	if effSpec != nil {
+		customizedArgs = vllmCustomizedArgsToRuntime(effSpec.CustomizedArgs)
+	}
+	mergedArgs, err := mergeRuntimeArgs(backendArgs, customizedArgs, validateVLLMArgKey)
+	if err != nil {
+		log.Errorf("merge vllm runtime args: %v", err)
+		return flags
+	}
+	return appendRuntimeFlags(flags, mergedArgs)
+}
+
+func buildVLLMServeFlagsWithModelExpr(modelExpr string, servedModelNameExpr string, tensorParallelSize int, backendParameters string, effSpec *api.LLMSpecVllm) []string {
 	flags := []string{
 		fmt.Sprintf("--model %s", modelExpr),
 		fmt.Sprintf("--served-model-name %s", servedModelNameExpr),
 		fmt.Sprintf("--port %d", api.LLM_VLLM_DEFAULT_PORT),
 		fmt.Sprintf("--tensor-parallel-size %d", tensorParallelSize),
 	}
-	return appendVLLMCustomizedFlags(flags, effSpec)
+	return appendVLLMRuntimeFlags(flags, backendParameters, effSpec)
 }

-func buildVLLMServeFlags(modelPath string, tensorParallelSize int, effSpec *api.LLMSpecVllm) []string {
+func buildVLLMServeFlags(modelPath string, tensorParallelSize int, backendParameters string, effSpec *api.LLMSpecVllm) []string {
 	modelQuoted := shellQuoteSingle(modelPath)
 	return buildVLLMServeFlagsWithModelExpr(
 		modelQuoted,
 		fmt.Sprintf(`"$(basename %s)"`, modelQuoted),
 		tensorParallelSize,
+		backendParameters,
 		effSpec,
 	)
 }

-func buildVLLMEntrypointScript(hasMountedModels bool, tensorParallelSize int, effSpec *api.LLMSpecVllm) string {
+func buildVLLMEntrypointScript(hasMountedModels bool, tensorParallelSize int, backendParameters string, effSpec *api.LLMSpecVllm) string {
 	modelsPath := shellQuoteSingle(api.LLM_VLLM_MODELS_PATH)
 	if !hasMountedModels {
 		return fmt.Sprintf("mkdir -p %s && exec sleep infinity", modelsPath)
@@ -184,6 +216,7 @@ func buildVLLMEntrypointScript(hasMountedModels bool, tensorParallelSize int, ef
 		`"$model"`,
 		`"$(basename "$model")"`,
 		tensorParallelSize,
+		backendParameters,
 		effSpec,
 	), " ")
 	return strings.Join([]string{
@@ -394,7 +427,11 @@ func (v *vllm) GetContainerSpec(ctx context.Context, llm *models.SLLM, image *mo
 	if eff := v.GetEffectiveSpec(llm, sku); eff != nil {
 		effSpec = eff.(*api.LLMSpecVllm)
 	}
-	startScript := buildVLLMEntrypointScript(len(postOverlays) > 0, tensorParallelSize, effSpec)
+	backendParameters := ""
+	if sku != nil {
+		backendParameters = sku.BackendParameters
+	}
+	startScript := buildVLLMEntrypointScript(len(postOverlays) > 0, tensorParallelSize, backendParameters, effSpec)
 	envs := []*commonapi.ContainerKeyValue{
 		{
 			Key:   "HUGGING_FACE_HUB_CACHE",