diff --git a/agent/app/service/tensorrt_llm.go b/agent/app/service/tensorrt_llm.go index 13218b2a8..d924cff98 100644 --- a/agent/app/service/tensorrt_llm.go +++ b/agent/app/service/tensorrt_llm.go @@ -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") diff --git a/agent/utils/env/env.go b/agent/utils/env/env.go index c9efeb2fd..43ca0a5a8 100644 --- a/agent/utils/env/env.go +++ b/agent/utils/env/env.go @@ -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 {