Skip to content

Commit

Permalink
Support multiline cmds in YAML configuration
Browse files Browse the repository at this point in the history
Add support for multiline `cmd` configurations allowing for nicer looking configuration YAML files.
  • Loading branch information
mostlygeek committed Oct 20, 2024
1 parent 6cf0962 commit be82d1a
Show file tree
Hide file tree
Showing 6 changed files with 185 additions and 18 deletions.
32 changes: 17 additions & 15 deletions config.example.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
10 changes: 9 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
6 changes: 6 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
@@ -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=
Expand Down
22 changes: 22 additions & 0 deletions proxy/config.go
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
package proxy

import (
"fmt"
"os"
"strings"

"gopkg.in/yaml.v3"
)
Expand All @@ -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"`
Expand Down Expand Up @@ -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
}
126 changes: 126 additions & 0 deletions proxy/config_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
7 changes: 5 additions & 2 deletions proxy/manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down

0 comments on commit be82d1a

Please sign in to comment.