Use client request context in proxy request (#43)

Canceled or closed HTTP requests from clients will also stop the proxied
HTTP requests to upstreamed servers.

README.md

@@ -5,10 +5,12 @@
 # Introduction
 llama-swap is a light weight, transparent proxy server that provides automatic model swapping to llama.cpp's server.
 
-Written in golang, it is very easy to install (single binary with no dependancies) and configure (single yaml file). Download a pre-built [release](https://github.com/mostlygeek/llama-swap/releases) or built it yourself from source with `make clean all`.
+Written in golang, it is very easy to install (single binary with no dependancies) and configure (single yaml file).
+
+Download a pre-built [release](https://github.com/mostlygeek/llama-swap/releases) or build it yourself from source with `make clean all`.
 
 ## How does it work?
-When a request is made to an OpenAI compatible endpoints, lama-swap will extract the `model` value load the appropriate server configuration to serve it. If a server is already running it will stop it and start a new one. This is where the "swap" part comes in. The upstream server is automatically swapped to the correct one to serve the request.
+When a request is made to an OpenAI compatible endpoint, lama-swap will extract the `model` value and load the appropriate server configuration to serve it. If a server is already running it will stop it and start the correct one. This is where the "swap" part comes in. The upstream server is automatically swapped to the correct one to serve the request.
 
 In the most basic configuration llama-swap handles one model at a time. For more advanced use cases, the `profiles` feature can load multiple models at the same time. You have complete control over how your system resources are used.
 
@@ -26,13 +28,15 @@ Any OpenAI compatible server would work. llama-swap was originally designed for
   - `v1/chat/completions`
   - `v1/embeddings`
   - `v1/rerank`
-  - `v1/audio/speech`
+  - `v1/audio/speech` ([#36](https://github.com/mostlygeek/llama-swap/issues/36))
 - ✅ Multiple GPU support
+- ✅ Docker Support ([#40](https://github.com/mostlygeek/llama-swap/pull/40))
 - ✅ Run multiple models at once with `profiles`
 - ✅ Remote log monitoring at `/log`
 - ✅ Automatic unloading of models from GPUs after timeout
 - ✅ Use any local OpenAI compatible server (llama.cpp, vllm, tabbyAPI, etc)
 - ✅ Direct access to upstream HTTP server via `/upstream/:model_id` ([demo](https://github.com/mostlygeek/llama-swap/pull/31))
+-
 
 ## config.yaml
 
@@ -87,6 +91,20 @@ models:
     cmd: llama-server --port 9999 -m Llama-3.2-1B-Instruct-Q4_K_M.gguf -ngl 0
     unlisted: true
 
+  # Docker Support (Experimental)
+  # see: https://github.com/mostlygeek/llama-swap/pull/40
+  "dockertest":
+    proxy: "http://127.0.0.1:9790"
+
+    # introduced to reliably stop containers
+    cmd_stop: docker stop -t 2 dockertest
+
+    cmd: >
+      docker run --name dockertest
+      --init --rm -p 9790:8080 -v /mnt/nvme/models:/models
+      ghcr.io/ggerganov/llama.cpp:server
+      --model '/models/Qwen2.5-Coder-0.5B-Instruct-Q4_K_M.gguf'
+
 # profiles make it easy to managing multi model (and gpu) configurations.
 #
 # Tips:

config.example.yaml

@@ -53,6 +53,21 @@ models:
       --ctx-size 8192
       --reranking
 
+  # EXPERIMENTAL! Docker Support
+  # see:
+  #  - https://github.com/mostlygeek/llama-swap/pull/40
+  #  - https://github.com/mostlygeek/llama-swap/issues/35
+  "dockertest":
+    proxy: "http://127.0.0.1:9790"
+
+    # use this to reliably stop named containers
+    cmd_stop: docker stop -t 2 dockertest
+
+    cmd: >
+      docker run --name dockertest
+      --init --rm -p 9790:8080 -v /mnt/nvme/models:/models
+      ghcr.io/ggerganov/llama.cpp:server
+      --model '/models/Qwen2.5-Coder-0.5B-Instruct-Q4_K_M.gguf'
 
   "simple":
     # example of setting environment variables

llama-swap.go

@@ -4,6 +4,8 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/gin-gonic/gin"
 	"github.com/mostlygeek/llama-swap/proxy"
@@ -39,6 +41,16 @@ func main() {
 	}
 
 	proxyManager := proxy.New(config)
+
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		<-sigChan
+		fmt.Println("Shutting down llama-swap")
+		proxyManager.StopProcesses()
+		os.Exit(0)
+	}()
+
 	fmt.Println("llama-swap listening on " + *listenStr)
 	if err := proxyManager.Run(*listenStr); err != nil {
 		fmt.Printf("Server error: %v\n", err)

proxy/config.go

@@ -11,6 +11,7 @@ import (
 
 type ModelConfig struct {
 	Cmd           string   `yaml:"cmd"`
+	CmdStop       string   `yaml:"cmd_stop"`
 	Proxy         string   `yaml:"proxy"`
 	Aliases       []string `yaml:"aliases"`
 	Env           []string `yaml:"env"`
@@ -22,6 +23,9 @@ type ModelConfig struct {
 func (m *ModelConfig) SanitizedCommand() ([]string, error) {
 	return SanitizeCommand(m.Cmd)
 }
+func (m *ModelConfig) SanitizeCommandStop() ([]string, error) {
+	return SanitizeCommand(m.CmdStop)
+}
 
 type Config struct {
 	HealthCheckTimeout int                    `yaml:"healthCheckTimeout"`

proxy/config_test.go

@@ -35,6 +35,11 @@ models:
     aliases:
       - "m2"
     checkEndpoint: "/"
+  docker:
+    cmd: docker run -p 9999:8080 --name "my_container"
+    cmd_stop: docker stop my_container
+    proxy: "http://localhost:9999"
+    checkEndpoint: "/health"
 healthCheckTimeout: 15
 profiles:
   test:
@@ -56,6 +61,7 @@ profiles:
 		Models: map[string]ModelConfig{
 			"model1": {
 				Cmd:           "path/to/cmd --arg1 one",
+				CmdStop:       "",
 				Proxy:         "http://localhost:8080",
 				Aliases:       []string{"m1", "model-one"},
 				Env:           []string{"VAR1=value1", "VAR2=value2"},
@@ -63,11 +69,19 @@ profiles:
 			},
 			"model2": {
 				Cmd:           "path/to/cmd --arg1 one",
+				CmdStop:       "",
 				Proxy:         "http://localhost:8081",
 				Aliases:       []string{"m2"},
 				Env:           nil,
 				CheckEndpoint: "/",
 			},
+			"docker": {
+				Cmd:           `docker run -p 9999:8080 --name "my_container"`,
+				CmdStop:       "docker stop my_container",
+				Proxy:         "http://localhost:9999",
+				Env:           nil,
+				CheckEndpoint: "/health",
+			},
 		},
 		HealthCheckTimeout: 15,
 		Profiles: map[string][]string{
@@ -99,6 +113,18 @@ func TestConfig_ModelConfigSanitizedCommand(t *testing.T) {
 	assert.Equal(t, []string{"python", "model1.py", "--arg1", "value1", "--arg2", "value2"}, args)
 }
 
+func TestConfig_ModelConfigSanitizedCommandStop(t *testing.T) {
+	config := &ModelConfig{
+		CmdStop: `docker stop my_container \
+		--arg1 1
+		--arg2 2`,
+	}
+
+	args, err := config.SanitizeCommandStop()
+	assert.NoError(t, err)
+	assert.Equal(t, []string{"docker", "stop", "my_container", "--arg1", "1", "--arg2", "2"}, args)
+}
+
 func TestConfig_FindConfig(t *testing.T) {
 
 	// TODO?

proxy/process.go

@@ -135,6 +135,7 @@ func (p *Process) start() error {
 				if time.Since(p.lastRequestHandled) > maxDuration {
 					fmt.Fprintf(p.logMonitor, "!!! Unloading model %s, TTL of %ds reached.\n", p.ID, p.config.UnloadAfter)
 					p.Stop()
+					return
 				}
 			}
 		}()
@@ -152,12 +153,13 @@ func (p *Process) Stop() {
 	defer p.stateMutex.Unlock()
 
 	if p.state != StateReady {
+		fmt.Fprintf(p.logMonitor, "!!! Stop() called but Process State is not READY\n")
 		return
 	}
 
 	if p.cmd == nil || p.cmd.Process == nil {
 		// this situation should never happen... but if it does just update the state
-		fmt.Fprintf(p.logMonitor, "!!! State is Ready but Command is nil.")
+		fmt.Fprintf(p.logMonitor, "!!! State is Ready but Command is nil.\n")
 		p.state = StateStopped
 		return
 	}
@@ -165,29 +167,59 @@ func (p *Process) Stop() {
 	// Pretty sure this stopping code needs some work for windows and
 	// will be a source of pain in the future.
 
-	p.cmd.Process.Signal(syscall.SIGTERM)
-	sigtermTimeout, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
-
-	sigtermNormal := make(chan error, 1)
-	go func() {
-		sigtermNormal <- p.cmd.Wait()
-	}()
-
-	select {
-	case <-sigtermTimeout.Done():
-		fmt.Fprintf(p.logMonitor, "!!! process for %s timed out waiting to stop\n", p.ID)
-		p.cmd.Process.Kill()
-		p.cmd.Wait()
-	case err := <-sigtermNormal:
+	if p.config.CmdStop != "" {
+		// for issue #35 to do things like `docker stop`
+		args, err := p.config.SanitizeCommandStop()
 		if err != nil {
-			if err.Error() != "wait: no child processes" {
-				// possible that simple-responder for testing is just not
-				// existing right, so suppress those errors.
-				fmt.Fprintf(p.logMonitor, "!!! process for %s stopped with error > %v\n", p.ID, err)
+			fmt.Fprintf(p.logMonitor, "!!! Error sanitizing stop command: %v\n", err)
+
+			// leave the state as it is?
+			return
+		}
+
+		fmt.Fprintf(p.logMonitor, "!!! Running stop command: %s\n", strings.Join(args, " "))
+		cmd := exec.Command(args[0], args[1:]...)
+		cmd.Stdout = p.logMonitor
+		cmd.Stderr = p.logMonitor
+		err = cmd.Start()
+		if err != nil {
+			fmt.Fprintf(p.logMonitor, "!!! Error running stop command: %v\n", err)
+
+			// leave the state as it is?
+			return
+		}
+
+		err = cmd.Wait()
+		if err != nil {
+			fmt.Fprintf(p.logMonitor, "!!! WARNING error waiting for stop command to complete: %v\n", err)
+		}
+	} else {
+		sigtermTimeout, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		sigtermNormal := make(chan error, 1)
+		go func() {
+			sigtermNormal <- p.cmd.Wait()
+		}()
+
+		p.cmd.Process.Signal(syscall.SIGTERM)
+
+		select {
+		case <-sigtermTimeout.Done():
+			fmt.Fprintf(p.logMonitor, "XXX Process for %s timed out waiting to stop, sending SIGKILL to PID: %d\n", p.ID, p.cmd.Process.Pid)
+			p.cmd.Process.Kill()
+			p.cmd.Wait()
+		case err := <-sigtermNormal:
+			if err != nil {
+				if err.Error() != "wait: no child processes" {
+					// possible that simple-responder for testing is just not
+					// existing right, so suppress those errors.
+					fmt.Fprintf(p.logMonitor, "!!! process for %s stopped with error > %v\n", p.ID, err)
+				}
 			}
 		}
 	}
+
 	p.state = StateStopped
 }
 
@@ -294,7 +326,7 @@ func (p *Process) ProxyRequest(w http.ResponseWriter, r *http.Request) {
 
 	proxyTo := p.config.Proxy
 	client := &http.Client{}
-	req, err := http.NewRequest(r.Method, proxyTo+r.URL.String(), r.Body)
+	req, err := http.NewRequestWithContext(r.Context(), r.Method, proxyTo+r.URL.String(), r.Body)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return

proxy/process_test.go

@@ -67,7 +67,6 @@ func TestProcess_BrokenModelConfig(t *testing.T) {
 	assert.Contains(t, w.Body.String(), "unable to start process")
 }
 
-// test that the process unloads after the TTL
 func TestProcess_UnloadAfterTTL(t *testing.T) {
 	if testing.Short() {
 		t.Skip("skipping long auto unload TTL test")
@@ -79,7 +78,7 @@ func TestProcess_UnloadAfterTTL(t *testing.T) {
 	config.UnloadAfter = 3 // seconds
 	assert.Equal(t, 3, config.UnloadAfter)
 
-	process := NewProcess("ttl", 2, config, NewLogMonitorWriter(io.Discard))
+	process := NewProcess("ttl_test", 2, config, NewLogMonitorWriter(io.Discard))
 	defer process.Stop()
 
 	// this should take 4 seconds
@@ -111,6 +110,33 @@ func TestProcess_UnloadAfterTTL(t *testing.T) {
 	assert.Equal(t, StateStopped, process.CurrentState())
 }
 
+func TestProcess_LowTTLValue(t *testing.T) {
+	if true { // change this code to run this ...
+		t.Skip("skipping test, edit process_test.go to run it ")
+	}
+
+	config := getTestSimpleResponderConfig("fast_ttl")
+	assert.Equal(t, 0, config.UnloadAfter)
+	config.UnloadAfter = 1 // second
+	assert.Equal(t, 1, config.UnloadAfter)
+
+	process := NewProcess("ttl", 2, config, NewLogMonitorWriter(os.Stdout))
+	defer process.Stop()
+
+	for i := 0; i < 100; i++ {
+		t.Logf("Waiting before sending request %d", i)
+		time.Sleep(1500 * time.Millisecond)
+
+		expected := fmt.Sprintf("echo=test_%d", i)
+		req := httptest.NewRequest("GET", fmt.Sprintf("/slow-respond?echo=%s&delay=50ms", expected), nil)
+		w := httptest.NewRecorder()
+		process.ProxyRequest(w, req)
+		assert.Equal(t, http.StatusOK, w.Code)
+		assert.Contains(t, w.Body.String(), expected)
+	}
+
+}
+
 // issue #19
 func TestProcess_HTTPRequestsHaveTimeToFinish(t *testing.T) {
 	if testing.Short() {

proxy/proxymanager.go

@@ -69,6 +69,19 @@ func New(config *Config) *ProxyManager {
 		})
 	}
 
+	// see: https://github.com/mostlygeek/llama-swap/issues/42
+	// respond with permissive OPTIONS for any endpoint
+	pm.ginEngine.Use(func(c *gin.Context) {
+		if c.Request.Method == "OPTIONS" {
+			c.Header("Access-Control-Allow-Origin", "*")
+			c.Header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
+			c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization")
+			c.AbortWithStatus(204)
+			return
+		}
+		c.Next()
+	})
+
 	// Set up routes using the Gin engine
 	pm.ginEngine.POST("/v1/chat/completions", pm.proxyOAIHandler)
 	// Support legacy /v1/completions api, see issue #12
@@ -202,6 +215,21 @@ func (pm *ProxyManager) swapModel(requestedModel string) (*Process, error) {
 		return nil, fmt.Errorf("could not find modelID for %s", requestedModel)
 	}
 
+	// check if model is part of the profile
+	if profileName != "" {
+		found := false
+		for _, item := range pm.config.Profiles[profileName] {
+			if item == realModelName {
+				found = true
+				break
+			}
+		}
+
+		if !found {
+			return nil, fmt.Errorf("model %s part of profile %s", realModelName, profileName)
+		}
+	}
+
 	// exit early when already running, otherwise stop everything and swap
 	requestedProcessKey := ProcessKeyName(profileName, realModelName)
 

proxy/proxymanager_test.go

@@ -210,3 +210,47 @@ func TestProxyManager_ListModelsHandler(t *testing.T) {
 	// Ensure all expected models were returned
 	assert.Empty(t, expectedModels, "not all expected models were returned")
 }
+
+func TestProxyManager_ProfileNonMember(t *testing.T) {
+
+	model1 := "path1/model1"
+	model2 := "path2/model2"
+
+	profileMemberName := ProcessKeyName("test", model1)
+	profileNonMemberName := ProcessKeyName("test", model2)
+
+	config := &Config{
+		HealthCheckTimeout: 15,
+		Models: map[string]ModelConfig{
+			model1: getTestSimpleResponderConfig("model1"),
+			model2: getTestSimpleResponderConfig("model2"),
+		},
+		Profiles: map[string][]string{
+			"test": {model1},
+		},
+	}
+
+	proxy := New(config)
+	defer proxy.StopProcesses()
+
+	// actual member of profile
+	{
+		reqBody := fmt.Sprintf(`{"model":"%s"}`, profileMemberName)
+		req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBufferString(reqBody))
+		w := httptest.NewRecorder()
+
+		proxy.HandlerFunc(w, req)
+		assert.Equal(t, http.StatusOK, w.Code)
+		assert.Contains(t, w.Body.String(), "model1")
+	}
+
+	// actual model, but non-member will 404
+	{
+		reqBody := fmt.Sprintf(`{"model":"%s"}`, profileNonMemberName)
+		req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBufferString(reqBody))
+		w := httptest.NewRecorder()
+
+		proxy.HandlerFunc(w, req)
+		assert.Equal(t, http.StatusNotFound, w.Code)
+	}
+}