feat: add MODEL_PATH environment for tensorrt llm (#10921)

This commit is contained in:
CityFun 2025-11-11 15:41:29 +08:00 committed by GitHub
parent 8e94170fed
commit 51e98623c9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 21 additions and 5 deletions

View file

@ -163,12 +163,11 @@ func handleLLMParams(llm *model.TensorRTLLM, create request.TensorRTLLMCreate) e
delete(serviceValue, "environment")
var environments []interface{}
environments = append(environments, fmt.Sprintf("MODEL_PATH=%s", create.ModelDir))
for _, e := range create.Environments {
environments = append(environments, fmt.Sprintf("%s=%s", e.Key, e.Value))
}
if len(environments) > 0 {
serviceValue["environment"] = environments
}
serviceValue["environment"] = environments
var volumes []interface{}
var defaultVolumes = map[string]string{
@ -255,7 +254,7 @@ func (t TensorRTLLMService) Create(create request.TensorRTLLMCreate) error {
envMap := handleLLMEnv(tensorrtLLM, create)
llmDir := path.Join(global.Dir.TensorRTLLMDir, create.Name)
envPath := path.Join(llmDir, ".env")
if err := env.Write(envMap, envPath); err != nil {
if err := env.WriteWithOrder(envMap, envPath, []string{"MODEL_PATH", "COMMAND"}); err != nil {
return err
}
dockerComposePath := path.Join(llmDir, "docker-compose.yml")
@ -294,7 +293,7 @@ func (t TensorRTLLMService) Update(req request.TensorRTLLMUpdate) error {
tensorrtLLM.Env = envStr
llmDir := path.Join(global.Dir.TensorRTLLMDir, tensorrtLLM.Name)
envPath := path.Join(llmDir, ".env")
if err := env.Write(envMap, envPath); err != nil {
if err := env.WriteWithOrder(envMap, envPath, []string{"MODEL_PATH", "COMMAND"}); err != nil {
return err
}
dockerComposePath := path.Join(llmDir, "docker-compose.yml")

View file

@ -40,6 +40,23 @@ func Marshal(envMap map[string]string) (string, error) {
return strings.Join(lines, "\n"), nil
}
func WriteWithOrder(envMap map[string]string, filename string, orders []string) error {
content, err := MarshalWithOrder(envMap, orders)
if err != nil {
return err
}
file, err := os.Create(filename)
if err != nil {
return err
}
defer file.Close()
_, err = file.WriteString(content + "\n")
if err != nil {
return err
}
return file.Sync()
}
func MarshalWithOrder(envMap map[string]string, orders []string) (string, error) {
lines := make([]string, 0, len(envMap))
for _, k := range orders {