Logging Improvements (#88)

This change revamps the internal logging architecture to be more flexible and descriptive. Previously all logs from both llama-swap and upstream services were mixed together. This makes it harder to troubleshoot and identify problems. This PR adds these new endpoints: 

- `/logs/stream/proxy` - just llama-swap's logs
- `/logs/stream/upstream` - stdout output from the upstream server

Author: mostlygeek

README.md

@@ -3,8 +3,6 @@
 ![GitHub Actions Workflow Status](https://img.shields.io/github/actions/workflow/status/mostlygeek/llama-swap/go-ci.yml)
 ![GitHub Repo stars](https://img.shields.io/github/stars/mostlygeek/llama-swap)
 
-
-
 # llama-swap
 
 llama-swap is a light weight, transparent proxy server that provides automatic model swapping to llama.cpp's server.
@@ -69,8 +67,8 @@ models:
 # Default (and minimum) is 15 seconds
 healthCheckTimeout: 60
 
-# Write HTTP logs (useful for troubleshooting), defaults to false
-logRequests: true
+# Valid log levels: debug, info (default), warn, error
+logLevel: info
 
 # define valid model values and the upstream server start
 models:
@@ -221,9 +219,15 @@ Of course, CLI access is also supported:
 # sends up to the last 10KB of logs
 curl http://host/logs'
 
-# streams logs
+# streams combined logs
 curl -Ns 'http://host/logs/stream'
 
+# just llama-swap's logs
+curl -Ns 'http://host/logs/stream/proxy'
+
+# just upstream's logs
+curl -Ns 'http://host/logs/stream/upstream'
+
 # stream and filter logs with linux pipes
 curl -Ns http://host/logs/stream | grep 'eval time'
 

config.example.yaml

@@ -1,9 +1,9 @@
 # Seconds to wait for llama.cpp to be available to serve requests
 # Default (and minimum): 15 seconds
-healthCheckTimeout: 15
+healthCheckTimeout: 90
 
-# Log HTTP requests helpful for troubleshoot, defaults to False
-logRequests: true
+# valid log levels: debug, info (default), warn, error
+logLevel: info
 
 models:
   "llama":

proxy/config.go

@@ -27,6 +27,7 @@ func (m *ModelConfig) SanitizedCommand() ([]string, error) {
 type Config struct {
 	HealthCheckTimeout int                    `yaml:"healthCheckTimeout"`
 	LogRequests        bool                   `yaml:"logRequests"`
+	LogLevel           string                 `yaml:"logLevel"`
 	Models             map[string]ModelConfig `yaml:"models"`
 	Profiles           map[string][]string    `yaml:"profiles"`
 

proxy/logMonitor.go

@@ -2,11 +2,21 @@ package proxy
 
 import (
 	"container/ring"
+	"fmt"
 	"io"
 	"os"
 	"sync"
 )
 
+type LogLevel int
+
+const (
+	LevelDebug LogLevel = iota
+	LevelInfo
+	LevelWarn
+	LevelError
+)
+
 type LogMonitor struct {
 	clients  map[chan []byte]bool
 	mu       sync.RWMutex
@@ -15,6 +25,10 @@ type LogMonitor struct {
 
 	// typically this can be os.Stdout
 	stdout io.Writer
+
+	// logging levels
+	level  LogLevel
+	prefix string
 }
 
 func NewLogMonitor() *LogMonitor {
@@ -26,6 +40,8 @@ func NewLogMonitorWriter(stdout io.Writer) *LogMonitor {
 		clients: make(map[chan []byte]bool),
 		buffer:  ring.New(10 * 1024), // keep 10KB of buffered logs
 		stdout:  stdout,
+		level:   LevelInfo,
+		prefix:  "",
 	}
 }
 
@@ -94,3 +110,77 @@ func (w *LogMonitor) broadcast(msg []byte) {
 		}
 	}
 }
+
+func (w *LogMonitor) SetPrefix(prefix string) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	w.prefix = prefix
+}
+
+func (w *LogMonitor) SetLogLevel(level LogLevel) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	w.level = level
+}
+
+func (w *LogMonitor) formatMessage(level string, msg string) []byte {
+	prefix := ""
+	if w.prefix != "" {
+		prefix = fmt.Sprintf("[%s] ", w.prefix)
+	}
+	return []byte(fmt.Sprintf("%s[%s] %s\n", prefix, level, msg))
+}
+
+func (w *LogMonitor) log(level LogLevel, msg string) {
+	if level < w.level {
+		return
+	}
+	w.Write(w.formatMessage(level.String(), msg))
+}
+
+func (w *LogMonitor) Debug(msg string) {
+	w.log(LevelDebug, msg)
+}
+
+func (w *LogMonitor) Info(msg string) {
+	w.log(LevelInfo, msg)
+}
+
+func (w *LogMonitor) Warn(msg string) {
+	w.log(LevelWarn, msg)
+}
+
+func (w *LogMonitor) Error(msg string) {
+	w.log(LevelError, msg)
+}
+
+func (w *LogMonitor) Debugf(format string, args ...interface{}) {
+	w.log(LevelDebug, fmt.Sprintf(format, args...))
+}
+
+func (w *LogMonitor) Infof(format string, args ...interface{}) {
+	w.log(LevelInfo, fmt.Sprintf(format, args...))
+}
+
+func (w *LogMonitor) Warnf(format string, args ...interface{}) {
+	w.log(LevelWarn, fmt.Sprintf(format, args...))
+}
+
+func (w *LogMonitor) Errorf(format string, args ...interface{}) {
+	w.log(LevelError, fmt.Sprintf(format, args...))
+}
+
+func (l LogLevel) String() string {
+	switch l {
+	case LevelDebug:
+		return "DEBUG"
+	case LevelInfo:
+		return "INFO"
+	case LevelWarn:
+		return "WARN"
+	case LevelError:
+		return "ERROR"
+	default:
+		return "UNKNOWN"
+	}
+}

proxy/process.go

@@ -30,10 +30,12 @@ const (
 )
 
 type Process struct {
-	ID         string
-	config     ModelConfig
-	cmd        *exec.Cmd
-	logMonitor *LogMonitor
+	ID     string
+	config ModelConfig
+	cmd    *exec.Cmd
+
+	processLogger *LogMonitor
+	proxyLogger   *LogMonitor
 
 	healthCheckTimeout      int
 	healthCheckLoopInterval time.Duration
@@ -53,13 +55,14 @@ type Process struct {
 	shutdownCancel context.CancelFunc
 }
 
-func NewProcess(ID string, healthCheckTimeout int, config ModelConfig, logMonitor *LogMonitor) *Process {
+func NewProcess(ID string, healthCheckTimeout int, config ModelConfig, processLogger *LogMonitor, proxyLogger *LogMonitor) *Process {
 	ctx, cancel := context.WithCancel(context.Background())
 	return &Process{
 		ID:                      ID,
 		config:                  config,
 		cmd:                     nil,
-		logMonitor:              logMonitor,
+		processLogger:           processLogger,
+		proxyLogger:             proxyLogger,
 		healthCheckTimeout:      healthCheckTimeout,
 		healthCheckLoopInterval: 5 * time.Second, /* default, can not be set by user - used for testing */
 		state:                   StateStopped,
@@ -68,6 +71,11 @@ func NewProcess(ID string, healthCheckTimeout int, config ModelConfig, logMonito
 	}
 }
 
+// LogMonitor returns the log monitor associated with the process.
+func (p *Process) LogMonitor() *LogMonitor {
+	return p.processLogger
+}
+
 // custom error types for swapping state
 var (
 	ErrExpectedStateMismatch  = errors.New("expected state mismatch")
@@ -85,9 +93,11 @@ func (p *Process) swapState(expectedState, newState ProcessState) (ProcessState,
 	}
 
 	if !isValidTransition(p.state, newState) {
+		p.proxyLogger.Warnf("Invalid state transition from %s to %s", p.state, newState)
 		return p.state, ErrInvalidStateTransition
 	}
 
+	p.proxyLogger.Debugf("State transition from %s to %s", expectedState, newState)
 	p.state = newState
 	return p.state, nil
 }
@@ -152,8 +162,8 @@ func (p *Process) start() error {
 	defer p.waitStarting.Done()
 
 	p.cmd = exec.Command(args[0], args[1:]...)
-	p.cmd.Stdout = p.logMonitor
-	p.cmd.Stderr = p.logMonitor
+	p.cmd.Stdout = p.processLogger
+	p.cmd.Stderr = p.processLogger
 	p.cmd.Env = p.config.Env
 
 	err = p.cmd.Start()
@@ -214,15 +224,16 @@ func (p *Process) start() error {
 				return errors.New("health check interrupted due to shutdown")
 			default:
 				if err := p.checkHealthEndpoint(healthURL); err == nil {
+					p.proxyLogger.Infof("Health check passed on %s", healthURL)
 					cancelHealthCheck()
 					break loop
 				} else {
 					if strings.Contains(err.Error(), "connection refused") {
 						endTime, _ := checkDeadline.Deadline()
 						ttl := time.Until(endTime)
-						fmt.Fprintf(p.logMonitor, "!!! Connection refused on %s, ttl %.0fs\n", healthURL, ttl.Seconds())
+						p.proxyLogger.Infof("Connection refused on %s, retrying in %.0fs", healthURL, ttl.Seconds())
 					} else {
-						fmt.Fprintf(p.logMonitor, "!!! Health check error: %v\n", err)
+						p.proxyLogger.Infof("Health check error on %s, %v", healthURL, err)
 					}
 				}
 			}
@@ -246,7 +257,8 @@ func (p *Process) start() error {
 				p.inFlightRequests.Wait()
 
 				if time.Since(p.lastRequestHandled) > maxDuration {
-					fmt.Fprintf(p.logMonitor, "!!! Unloading model %s, TTL of %ds reached.\n", p.ID, p.config.UnloadAfter)
+
+					p.proxyLogger.Infof("Unloading model %s, TTL of %ds reached.", p.ID, p.config.UnloadAfter)
 					p.Stop()
 					return
 				}
@@ -267,15 +279,15 @@ func (p *Process) Stop() {
 
 	// calling Stop() when state is invalid is a no-op
 	if curState, err := p.swapState(StateReady, StateStopping); err != nil {
-		fmt.Fprintf(p.logMonitor, "!!! Info - Stop() Ready -> StateStopping err: %v, current state: %v\n", err, curState)
+		p.proxyLogger.Infof("Stop() Ready -> StateStopping err: %v, current state: %v", err, curState)
 		return
 	}
 
 	// stop the process with a graceful exit timeout
 	p.stopCommand(5 * time.Second)
 
 	if curState, err := p.swapState(StateStopping, StateStopped); err != nil {
-		fmt.Fprintf(p.logMonitor, "!!! Info - Stop() StateStopping -> StateStopped err: %v, current state: %v\n", err, curState)
+		p.proxyLogger.Infof("Stop() StateStopping -> StateStopped err: %v, current state: %v", err, curState)
 	}
 }
 
@@ -300,33 +312,32 @@ func (p *Process) stopCommand(sigtermTTL time.Duration) {
 	}()
 
 	if p.cmd == nil || p.cmd.Process == nil {
-		fmt.Fprintf(p.logMonitor, "!!! process [%s] cmd or cmd.Process is nil", p.ID)
+		p.proxyLogger.Warnf("Process [%s] cmd or cmd.Process is nil", p.ID)
 		return
 	}
 
 	if err := p.terminateProcess(); err != nil {
-		fmt.Fprintf(p.logMonitor, "!!! failed to gracefully terminate process [%s]: %v\n", p.ID, err)
+		p.proxyLogger.Infof("Failed to gracefully terminate process [%s]: %v", p.ID, err)
 	}
 
 	select {
 	case <-sigtermTimeout.Done():
-		fmt.Fprintf(p.logMonitor, "!!! process [%s] timed out waiting to stop, sending KILL signal\n", p.ID)
+		p.proxyLogger.Infof("Process [%s] timed out waiting to stop, sending KILL signal", p.ID)
 		p.cmd.Process.Kill()
 	case err := <-sigtermNormal:
 		if err != nil {
 			if errno, ok := err.(syscall.Errno); ok {
-				fmt.Fprintf(p.logMonitor, "!!! process [%s] errno >> %v\n", p.ID, errno)
+				p.proxyLogger.Errorf("Process [%s] errno >> %v", p.ID, errno)
 			} else if exitError, ok := err.(*exec.ExitError); ok {
 				if strings.Contains(exitError.String(), "signal: terminated") {
-					fmt.Fprintf(p.logMonitor, "!!! process [%s] stopped OK\n", p.ID)
+					p.proxyLogger.Infof("Process [%s] stopped OK", p.ID)
 				} else if strings.Contains(exitError.String(), "signal: interrupt") {
-					fmt.Fprintf(p.logMonitor, "!!! process [%s] interrupted OK\n", p.ID)
+					p.proxyLogger.Infof("Process [%s] interrupted OK", p.ID)
 				} else {
-					fmt.Fprintf(p.logMonitor, "!!! process [%s] ExitError >> %v, exit code: %d\n", p.ID, exitError, exitError.ExitCode())
+					p.proxyLogger.Warnf("Process [%s] ExitError >> %v, exit code: %d", p.ID, exitError, exitError.ExitCode())
 				}
-
 			} else {
-				fmt.Fprintf(p.logMonitor, "!!! process [%s] exited >> %v\n", p.ID, err)
+				p.proxyLogger.Errorf("Process [%s] exited >> %v", p.ID, err)
 			}
 		}
 	}