From be82d1a6a0cceb6babdd91e59780d7ecdae0e23d Mon Sep 17 00:00:00 2001 From: Benson Wong Date: Sat, 19 Oct 2024 20:06:59 -0700 Subject: [PATCH] Support multiline cmds in YAML configuration Add support for multiline `cmd` configurations allowing for nicer looking configuration YAML files. --- config.example.yaml | 32 +++++------ go.mod | 10 +++- go.sum | 6 +++ proxy/config.go | 22 ++++++++ proxy/config_test.go | 126 +++++++++++++++++++++++++++++++++++++++++++ proxy/manager.go | 7 ++- 6 files changed, 185 insertions(+), 18 deletions(-) create mode 100644 proxy/config_test.go diff --git a/config.example.yaml b/config.example.yaml index 9fedd0b..fc738e5 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -4,36 +4,38 @@ healthCheckTimeout: 60 models: "llama": - cmd: "models/llama-server-osx --port 8999 -m models/Llama-3.2-1B-Instruct-Q4_K_M.gguf" - proxy: "http://127.0.0.1:8999" + cmd: > + models/llama-server-osx + --port 8999 + -m models/Llama-3.2-1B-Instruct-Q4_K_M.gguf + proxy: http://127.0.0.1:8999 # list of model name aliases this llama.cpp instance can serve aliases: - - "gpt-4o-mini" + - gpt-4o-mini # check this path for a HTTP 200 response for the server to be ready - checkEndpoint: "/health" + checkEndpoint: /health "qwen": - cmd: "models/llama-server-osx --port 8999 -m models/Qwen2.5-1.5B-Instruct-Q4_K_M.gguf" - proxy: "http://127.0.0.1:8999" + cmd: models/llama-server-osx --port 8999 -m models/Qwen2.5-1.5B-Instruct-Q4_K_M.gguf + proxy: http://127.0.0.1:8999 aliases: - - "gpt-3.5-turbo" + - gpt-3.5-turbo "simple": # example of setting environment variables env: - - "CUDA_VISIBLE_DEVICES=0,1" - - "env1=hello" - cmd: "build/simple-responder --port 8999" - proxy: "http://127.0.0.1:8999" + - CUDA_VISIBLE_DEVICES=0,1 + - env1=hello + cmd: build/simple-responder --port 8999 + proxy: http://127.0.0.1:8999 # use "none" to skip check. Caution this may cause some requests to fail # until the upstream server is ready for traffic - checkEndpoint: "none" + checkEndpoint: none # don't use this, just for testing if things are broken "broken": - cmd: "models/llama-server-osx --port 8999 -m models/doesnotexist.gguf" - proxy: "http://127.0.0.1:8999" - + cmd: models/llama-server-osx --port 8999 -m models/doesnotexist.gguf + proxy: http://127.0.0.1:8999 \ No newline at end of file diff --git a/go.mod b/go.mod index 7b25993..56cd702 100644 --- a/go.mod +++ b/go.mod @@ -2,4 +2,12 @@ module github.com/mostlygeek/llama-swap go 1.23.0 -require gopkg.in/yaml.v3 v3.0.1 +require ( + github.com/stretchr/testify v1.9.0 + gopkg.in/yaml.v3 v3.0.1 +) + +require ( + github.com/davecgh/go-spew v1.1.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect +) diff --git a/go.sum b/go.sum index a62c313..60ce688 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,9 @@ +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= diff --git a/proxy/config.go b/proxy/config.go index e3cc213..27e3773 100644 --- a/proxy/config.go +++ b/proxy/config.go @@ -1,7 +1,9 @@ package proxy import ( + "fmt" "os" + "strings" "gopkg.in/yaml.v3" ) @@ -14,6 +16,10 @@ type ModelConfig struct { CheckEndpoint string `yaml:"checkEndpoint"` } +func (m *ModelConfig) SanitizedCommand() ([]string, error) { + return SanitizeCommand(m.Cmd) +} + type Config struct { Models map[string]ModelConfig `yaml:"models"` HealthCheckTimeout int `yaml:"healthCheckTimeout"` @@ -55,3 +61,19 @@ func LoadConfig(path string) (*Config, error) { return &config, nil } + +func SanitizeCommand(cmdStr string) ([]string, error) { + // Remove trailing backslashes + cmdStr = strings.ReplaceAll(cmdStr, "\\ \n", " ") + cmdStr = strings.ReplaceAll(cmdStr, "\\\n", " ") + + // Split the command into arguments + args := strings.Fields(cmdStr) + + // Ensure the command is not empty + if len(args) == 0 { + return nil, fmt.Errorf("empty command") + } + + return args, nil +} diff --git a/proxy/config_test.go b/proxy/config_test.go new file mode 100644 index 0000000..a6e95e7 --- /dev/null +++ b/proxy/config_test.go @@ -0,0 +1,126 @@ +package proxy + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestLoadConfig(t *testing.T) { + // Create a temporary YAML file for testing + tempDir, err := os.MkdirTemp("", "test-config") + if err != nil { + t.Fatalf("Failed to create temporary directory: %v", err) + } + defer os.RemoveAll(tempDir) + + tempFile := filepath.Join(tempDir, "config.yaml") + content := `models: + model1: + cmd: path/to/cmd --arg1 one + proxy: "http://localhost:8080" + aliases: + - "m1" + - "model-one" + env: + - "VAR1=value1" + - "VAR2=value2" + checkEndpoint: "/health" +healthCheckTimeout: 15 +` + + if err := os.WriteFile(tempFile, []byte(content), 0644); err != nil { + t.Fatalf("Failed to write temporary file: %v", err) + } + + // Load the config and verify + config, err := LoadConfig(tempFile) + if err != nil { + t.Fatalf("Failed to load config: %v", err) + } + + expected := &Config{ + Models: map[string]ModelConfig{ + "model1": { + Cmd: "path/to/cmd --arg1 one", + Proxy: "http://localhost:8080", + Aliases: []string{"m1", "model-one"}, + Env: []string{"VAR1=value1", "VAR2=value2"}, + CheckEndpoint: "/health", + }, + }, + HealthCheckTimeout: 15, + } + + assert.Equal(t, expected, config) +} + +func TestModelConfigSanitizedCommand(t *testing.T) { + config := &ModelConfig{ + Cmd: `python model1.py \ + --arg1 value1 \ + --arg2 value2`, + } + + args, err := config.SanitizedCommand() + assert.NoError(t, err) + assert.Equal(t, []string{"python", "model1.py", "--arg1", "value1", "--arg2", "value2"}, args) +} + +func TestFindConfig(t *testing.T) { + config := &Config{ + Models: map[string]ModelConfig{ + "model1": { + Cmd: "python model1.py", + Proxy: "http://localhost:8080", + Aliases: []string{"m1", "model-one"}, + Env: []string{"VAR1=value1", "VAR2=value2"}, + CheckEndpoint: "/health", + }, + "model2": { + Cmd: "python model2.py", + Proxy: "http://localhost:8081", + Aliases: []string{"m2", "model-two"}, + Env: []string{"VAR3=value3", "VAR4=value4"}, + CheckEndpoint: "/status", + }, + }, + HealthCheckTimeout: 10, + } + + // Test finding a model by its name + modelConfig, found := config.FindConfig("model1") + assert.True(t, found) + assert.Equal(t, config.Models["model1"], modelConfig) + + // Test finding a model by its alias + modelConfig, found = config.FindConfig("m1") + assert.True(t, found) + assert.Equal(t, config.Models["model1"], modelConfig) + + // Test finding a model that does not exist + modelConfig, found = config.FindConfig("model3") + assert.False(t, found) + assert.Equal(t, ModelConfig{}, modelConfig) +} + +func TestSanitizeCommand(t *testing.T) { + // Test a simple command + args, err := SanitizeCommand("python model1.py") + assert.NoError(t, err) + assert.Equal(t, []string{"python", "model1.py"}, args) + + // Test a command with spaces and newlines + args, err = SanitizeCommand(`python model1.py \ + --arg1 value1 \ + --arg2 value2`) + assert.NoError(t, err) + assert.Equal(t, []string{"python", "model1.py", "--arg1", "value1", "--arg2", "value2"}, args) + + // Test an empty command + args, err = SanitizeCommand("") + assert.Error(t, err) + assert.Nil(t, args) +} diff --git a/proxy/manager.go b/proxy/manager.go index ad028d6..17a641e 100644 --- a/proxy/manager.go +++ b/proxy/manager.go @@ -65,13 +65,16 @@ func (pm *ProxyManager) swapModel(requestedModel string) error { pm.currentConfig = modelConfig - args := strings.Fields(modelConfig.Cmd) + args, err := modelConfig.SanitizedCommand() + if err != nil { + return fmt.Errorf("unable to get sanitized command: %v", err) + } cmd := exec.Command(args[0], args[1:]...) cmd.Stdout = os.Stdout cmd.Stderr = os.Stderr cmd.Env = modelConfig.Env - err := cmd.Start() + err = cmd.Start() if err != nil { return err }