Include llama-server output in server error

Signed-off-by: Piotr Stankiewicz <[email protected]>

Author: Piotr Stankiewicz

go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/opencontainers/go-digest v1.0.0
 	github.com/opencontainers/image-spec v1.1.1
 	github.com/sirupsen/logrus v1.9.3
+	github.com/stretchr/testify v1.10.0
 	golang.org/x/sync v0.12.0
 )
 
@@ -20,6 +21,7 @@ require (
 	github.com/containerd/errdefs v1.0.0 // indirect
 	github.com/containerd/log v0.1.0 // indirect
 	github.com/containerd/stargz-snapshotter/estargz v0.16.3 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/distribution/reference v0.6.0 // indirect
 	github.com/docker/cli v27.5.0+incompatible // indirect
 	github.com/docker/distribution v2.8.3+incompatible // indirect
@@ -38,6 +40,7 @@ require (
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/rs/dnscache v0.0.0-20230804202142-fc85eb664529 // indirect
 	github.com/smallnest/ringbuffer v0.0.0-20241116012123-461381446e3d // indirect
 	github.com/vbatts/tar-split v0.11.6 // indirect

pkg/inference/backends/llamacpp/llamacpp.go

@@ -4,18 +4,21 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"io/fs"
 	"net/http"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"runtime"
 	"strconv"
+	"strings"
 
 	"github.com/docker/model-runner/pkg/diskusage"
 	"github.com/docker/model-runner/pkg/inference"
 	"github.com/docker/model-runner/pkg/inference/models"
 	"github.com/docker/model-runner/pkg/logging"
+	"github.com/docker/model-runner/pkg/tailbuffer"
 )
 
 const (
@@ -162,9 +165,11 @@ func (l *llamaCpp) Run(ctx context.Context, socket, model string, mode inference
 		}
 		return llamaCppProcess.Process.Signal(os.Interrupt)
 	}
+	tailBuf := tailbuffer.NewTailBuffer(1024)
 	serverLogStream := l.serverLog.Writer()
+	out := io.MultiWriter(serverLogStream, tailBuf)
 	llamaCppProcess.Stdout = serverLogStream
-	llamaCppProcess.Stderr = serverLogStream
+	llamaCppProcess.Stderr = out
 
 	if err := llamaCppProcess.Start(); err != nil {
 		return fmt.Errorf("unable to start llama.cpp: %w", err)
@@ -174,6 +179,18 @@ func (l *llamaCpp) Run(ctx context.Context, socket, model string, mode inference
 	go func() {
 		llamaCppErr := llamaCppProcess.Wait()
 		serverLogStream.Close()
+
+		errOutput := new(strings.Builder)
+		if _, err := io.Copy(errOutput, tailBuf); err != nil {
+			l.log.Warnf("failed to read server output tail: %w", err)
+		}
+
+		if len(errOutput.String()) != 0 {
+			llamaCppErr = fmt.Errorf("llama.cpp exit status: %w\nwith output: %s", llamaCppErr, errOutput.String())
+		} else {
+			llamaCppErr = fmt.Errorf("llama.cpp exit status: %w", llamaCppErr)
+		}
+
 		llamaCppErrors <- llamaCppErr
 		close(llamaCppErrors)
 		if err := os.Remove(socket); err != nil && !errors.Is(err, fs.ErrNotExist) {

pkg/tailbuffer/tailbuffer.go

@@ -0,0 +1,81 @@
+package tailbuffer
+
+import (
+	"io"
+	"sync"
+)
+
+type tailBuffer struct {
+	lock     sync.Mutex
+	buf      []byte
+	capacity uint
+	size     uint
+	read     uint
+	write    uint
+}
+
+func NewTailBuffer(size uint) io.ReadWriter {
+	return &tailBuffer{
+		buf:      make([]byte, size),
+		capacity: size,
+		size:     0,
+		read:     0,
+		write:    0,
+	}
+}
+
+func (w *tailBuffer) Write(buffer []byte) (int, error) {
+	w.lock.Lock()
+	defer w.lock.Unlock()
+
+	written := 0
+	shouldPushRead := false
+	si := 0
+	if len(buffer) > int(w.capacity) {
+		si = len(buffer) - int(w.capacity)
+	}
+	for _, b := range buffer[si:] {
+		if shouldPushRead {
+			if w.read+1 < w.capacity {
+				w.read += 1
+			} else {
+				w.read = 0
+			}
+		}
+		w.buf[w.write] = b
+		if w.write+1 < w.capacity {
+			w.write += 1
+		} else {
+			w.write = 0
+		}
+		w.size += 1
+		if w.size > w.capacity {
+			w.size = w.capacity
+		}
+		shouldPushRead = w.write == w.read
+		written += 1
+	}
+	return si + written, nil
+}
+
+func (w *tailBuffer) Read(buffer []byte) (int, error) {
+	w.lock.Lock()
+	defer w.lock.Unlock()
+
+	var err error
+	read := uint(0)
+	for read < w.size && int(read) < len(buffer) {
+		buffer[read] = w.buf[w.read]
+		if w.read+1 < w.capacity {
+			w.read += 1
+		} else {
+			w.read = 0
+		}
+		read += 1
+	}
+	w.size -= read
+	if read == 0 {
+		err = io.EOF
+	}
+	return int(read), err
+}

pkg/tailbuffer/tailbuffer_test.go

@@ -0,0 +1,63 @@
+package tailbuffer
+
+import (
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestLogBufferCreation(t *testing.T) {
+	lb := NewTailBuffer(0)
+	require.NotNil(t, lb)
+}
+
+func TestLogBufferWrite(t *testing.T) {
+	lb := NewTailBuffer(1024)
+	n, err := lb.Write([]byte("asdf"))
+	require.NoError(t, err)
+	require.Equal(t, 4, n)
+}
+
+func TestLogBufferReadEmpty(t *testing.T) {
+	lb := NewTailBuffer(4)
+	buf := make([]byte, 4)
+	_, err := lb.Read(buf)
+	require.Error(t, err, io.EOF)
+}
+
+func TestLogBufferWriteRead(t *testing.T) {
+	lb := NewTailBuffer(4)
+	n, err := lb.Write([]byte("asdfg"))
+	require.NoError(t, err)
+	require.Equal(t, 5, n)
+	buf := make([]byte, 4)
+	n, err = lb.Read(buf)
+	require.NoError(t, err)
+	require.Equal(t, 4, n)
+	require.Equal(t, []byte("sdfg"), buf)
+	n, err = lb.Write([]byte("hjklzx"))
+	require.NoError(t, err)
+	require.Equal(t, 6, n)
+	buf = make([]byte, 3)
+	n, err = lb.Read(buf)
+	require.NoError(t, err)
+	require.Equal(t, 3, n)
+	require.Equal(t, []byte("klz"), buf)
+	n, err = lb.Read(buf)
+	require.NoError(t, err)
+	require.Equal(t, 1, n)
+}
+
+func TestLogBufferWriteReadString(t *testing.T) {
+	lb := NewTailBuffer(4)
+	n, err := lb.Write([]byte("asdfg"))
+	require.NoError(t, err)
+	require.Equal(t, 5, n)
+	str := new(strings.Builder)
+	nw, err := io.Copy(str, lb)
+	require.NoError(t, err)
+	require.Equal(t, int64(4), nw)
+	require.Equal(t, "sdfg", str.String())
+}