diff --git a/pkg/llm/drivers/llm_container/backend_parameters.go b/pkg/llm/drivers/llm_container/backend_parameters.go new file mode 100644 index 0000000000..bb8ead99ae --- /dev/null +++ b/pkg/llm/drivers/llm_container/backend_parameters.go @@ -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 +} diff --git a/pkg/llm/drivers/llm_container/sglang.go b/pkg/llm/drivers/llm_container/sglang.go index bacfad7bcb..036ef55f37 100644 --- a/pkg/llm/drivers/llm_container/sglang.go +++ b/pkg/llm/drivers/llm_container/sglang.go @@ -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{ diff --git a/pkg/llm/drivers/llm_container/vllm.go b/pkg/llm/drivers/llm_container/vllm.go index c67f1743dc..cb7602c8cc 100644 --- a/pkg/llm/drivers/llm_container/vllm.go +++ b/pkg/llm/drivers/llm_container/vllm.go @@ -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",