llm: detect llama-server load stalls from output (#16427)

llama-server model loads could time out after the fixed load duration even while tensor-loading progress dots were still being emitted, so track raw runner output activity and use OLLAMA_LOAD_TIMEOUT as a stall deadline.

Fixes #16416
Fixes #16412

llm/llama_server.go

@@ -36,6 +36,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"golang.org/x/sync/semaphore"
@@ -144,6 +145,8 @@ type llamaServerRunner struct {
 	ggml          *ggml.GGML
 	totalLayers   uint64 // maximum offloadable model layers
 	loadStart     time.Time
+	loadActivity  atomic.Int64
+	loadTracking  atomic.Bool
 	rawEmbeddings bool
 
 	sem *semaphore.Weighted
@@ -848,6 +851,7 @@ func (s *llamaServerRunner) startProcess() error {
 	s.done = make(chan struct{})
 	s.doneErr = nil
 	s.loadStart = time.Now()
+	s.startLoadTracking(s.loadStart)
 
 	// Reap subprocess when it exits.
 	go func(cmd *exec.Cmd, done chan struct{}) {
@@ -974,6 +978,51 @@ func (s *llamaServerRunner) hasParsedVRAM() bool {
 	return len(s.vramByDevice) > 0
 }
 
+func (s *llamaServerRunner) startLoadTracking(t time.Time) {
+	if s == nil {
+		return
+	}
+	s.loadTracking.Store(true)
+	s.noteLoadActivity(t)
+}
+
+func (s *llamaServerRunner) stopLoadTracking() {
+	if s == nil {
+		return
+	}
+	s.loadTracking.Store(false)
+}
+
+func (s *llamaServerRunner) noteLoadActivity(t time.Time) {
+	if s == nil || t.IsZero() {
+		return
+	}
+	if !s.loadTracking.Load() {
+		return
+	}
+
+	ns := t.UnixNano()
+	for {
+		prev := s.loadActivity.Load()
+		if ns <= prev {
+			return
+		}
+		if s.loadActivity.CompareAndSwap(prev, ns) {
+			return
+		}
+	}
+}
+
+func (s *llamaServerRunner) lastLoadActivity() time.Time {
+	if s == nil {
+		return time.Time{}
+	}
+	if ns := s.loadActivity.Load(); ns > 0 {
+		return time.Unix(0, ns)
+	}
+	return time.Time{}
+}
+
 // getServerStatus checks llama-server's /health endpoint.
 // llama-server returns {"status":"ok"}, {"status":"loading model"}, or {"status":"error"}.
 func (s *llamaServerRunner) getServerStatus(ctx context.Context) (ServerStatus, error) {
@@ -1010,6 +1059,9 @@ func (s *llamaServerRunner) getServerStatus(ctx context.Context) (ServerStatus,
 	// llama-server returns {"status":"ok"}, {"status":"loading model"}, {"status":"error", ...}
 	var result struct {
 		Status string `json:"status"`
+		Error  *struct {
+			Message string `json:"message"`
+		} `json:"error"`
 	}
 	if err := json.Unmarshal(body, &result); err != nil {
 		return ServerStatusError, fmt.Errorf("health unmarshal: %w", err)
@@ -1023,6 +1075,14 @@ func (s *llamaServerRunner) getServerStatus(ctx context.Context) (ServerStatus,
 	case "no slot available":
 		return ServerStatusNoSlotsAvailable, nil
 	default:
+		if result.Error != nil {
+			switch strings.ToLower(strings.TrimSpace(result.Error.Message)) {
+			case "loading model":
+				return ServerStatusLoadingModel, nil
+			case "no slot available":
+				return ServerStatusNoSlotsAvailable, nil
+			}
+		}
 		return ServerStatusError, fmt.Errorf("llama-server error: %s", string(body))
 	}
 }
@@ -1055,7 +1115,18 @@ func (s *llamaServerRunner) Ping(ctx context.Context) error {
 }
 
 func (s *llamaServerRunner) WaitUntilRunning(ctx context.Context) error {
-	loadDeadline := time.Now().Add(envconfig.LoadTimeout())
+	s.startLoadTracking(time.Now())
+	defer s.stopLoadTracking()
+
+	stallTimeout := envconfig.LoadTimeout()
+	lastActivity := s.lastLoadActivity()
+	if lastActivity.IsZero() {
+		lastActivity = s.loadStart
+	}
+	if lastActivity.IsZero() {
+		lastActivity = time.Now()
+	}
+	loadDeadline := lastActivity.Add(stallTimeout)
 
 	slog.Info("waiting for llama-server to start responding")
 	var lastStatus ServerStatus = -1
@@ -1091,6 +1162,11 @@ func (s *llamaServerRunner) WaitUntilRunning(ctx context.Context) error {
 		default:
 		}
 
+		if activity := s.lastLoadActivity(); activity.After(lastActivity) {
+			lastActivity = activity
+			loadDeadline = lastActivity.Add(stallTimeout)
+		}
+
 		if time.Now().After(loadDeadline) {
 			msg := s.lastErrMsg()
 			return fmt.Errorf("timed out waiting for llama-server to start - %s", msg)
@@ -1109,6 +1185,10 @@ func (s *llamaServerRunner) WaitUntilRunning(ctx context.Context) error {
 		if statusChanged && status != ServerStatusReady {
 			slog.Info("waiting for llama-server to become available", "status", status)
 		}
+		if statusChanged && status == ServerStatusLoadingModel {
+			lastActivity = time.Now()
+			loadDeadline = lastActivity.Add(stallTimeout)
+		}
 
 		switch status {
 		case ServerStatusReady:
@@ -2478,6 +2558,10 @@ func deviceName(backendName string) string {
 
 func (w *memoryParsingWriter) Write(b []byte) (int, error) {
 	if w.runner != nil {
+		if len(b) > 0 && w.runner.loadTracking.Load() {
+			w.runner.noteLoadActivity(time.Now())
+		}
+
 		func() {
 			w.runner.memoryMu.Lock()
 			defer w.runner.memoryMu.Unlock()

llm/llama_server_test.go

@@ -48,6 +48,12 @@ func TestLlamaServerHealthParsing(t *testing.T) {
 			statusCode: 503,
 			wantStatus: ServerStatusLoadingModel,
 		},
+		{
+			name:       "loading error envelope",
+			body:       `{"error":{"message":"Loading model","type":"unavailable_error","code":503}}`,
+			statusCode: 503,
+			wantStatus: ServerStatusLoadingModel,
+		},
 		{
 			name:       "no slots",
 			body:       `{"status":"no slot available"}`,
@@ -797,7 +803,7 @@ func TestLlamaServerWaitUntilRunningWaitsOnRecoverableStartupOOM(t *testing.T) {
 	}
 }
 
-func TestLlamaServerWaitUntilRunningTimesOutWhenLoadExceedsTimeout(t *testing.T) {
+func TestLlamaServerWaitUntilRunningTimesOutWhenLoadStalls(t *testing.T) {
 	t.Setenv("OLLAMA_LOAD_TIMEOUT", "10ms")
 
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -828,6 +834,55 @@ func TestLlamaServerWaitUntilRunningTimesOutWhenLoadExceedsTimeout(t *testing.T)
 	}
 }
 
+func TestLlamaServerWaitUntilRunningExtendsTimeoutOnOutputActivity(t *testing.T) {
+	t.Setenv("OLLAMA_LOAD_TIMEOUT", "20ms")
+
+	var activityCount atomic.Int32
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/health" {
+			t.Errorf("unexpected path: %s", r.URL.Path)
+			return
+		}
+		if activityCount.Load() < 5 {
+			w.WriteHeader(http.StatusServiceUnavailable)
+			fmt.Fprint(w, `{"error":{"message":"Loading model","type":"unavailable_error","code":503}}`)
+			return
+		}
+		fmt.Fprint(w, `{"status":"ok"}`)
+	}))
+	defer srv.Close()
+
+	parts := strings.Split(srv.URL, ":")
+	var portInt int
+	fmt.Sscanf(parts[len(parts)-1], "%d", &portInt)
+
+	runner := &llamaServerRunner{
+		port: portInt,
+		cmd:  fakeRunningCmd(),
+	}
+	runner.output = &memoryParsingWriter{inner: io.Discard, runner: runner}
+
+	done := make(chan struct{})
+	defer close(done)
+	go func() {
+		ticker := time.NewTicker(5 * time.Millisecond)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-done:
+				return
+			case <-ticker.C:
+				activityCount.Add(1)
+				_, _ = runner.output.Write([]byte("."))
+			}
+		}
+	}()
+
+	if err := runner.WaitUntilRunning(t.Context()); err != nil {
+		t.Fatalf("WaitUntilRunning error: %v", err)
+	}
+}
+
 func TestLlamaServerCompletionRequestFormat(t *testing.T) {
 	tests := []struct {
 		name           string
@@ -2425,6 +2480,39 @@ func TestMemoryParsingWriter(t *testing.T) {
 	}
 }
 
+func TestMemoryParsingWriterRecordsOutputActivityWithoutNewline(t *testing.T) {
+	runner := &llamaServerRunner{}
+	w := &memoryParsingWriter{inner: io.Discard, runner: runner}
+
+	runner.startLoadTracking(time.Now())
+	before := time.Now()
+	if _, err := w.Write([]byte("...")); err != nil {
+		t.Fatal(err)
+	}
+	if got := runner.lastLoadActivity(); got.Before(before) {
+		t.Fatalf("lastLoadActivity = %v, want after %v", got, before)
+	}
+}
+
+func TestMemoryParsingWriterIgnoresOutputActivityAfterLoadTrackingStops(t *testing.T) {
+	runner := &llamaServerRunner{}
+	w := &memoryParsingWriter{inner: io.Discard, runner: runner}
+
+	runner.startLoadTracking(time.Now())
+	if _, err := w.Write([]byte(".")); err != nil {
+		t.Fatal(err)
+	}
+	lastActivity := runner.lastLoadActivity()
+
+	runner.stopLoadTracking()
+	if _, err := w.Write([]byte(".")); err != nil {
+		t.Fatal(err)
+	}
+	if got := runner.lastLoadActivity(); !got.Equal(lastActivity) {
+		t.Fatalf("lastLoadActivity changed after tracking stopped: got %v, want %v", got, lastActivity)
+	}
+}
+
 func TestMemoryParsingPerDevice(t *testing.T) {
 	tests := []struct {
 		name       string