diff --git a/llm/llama_server.go b/llm/llama_server.go index 9c0ddcdca..671d7a2cc 100644 --- a/llm/llama_server.go +++ b/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() diff --git a/llm/llama_server_test.go b/llm/llama_server_test.go index a2cf1e7be..11b07147d 100644 --- a/llm/llama_server_test.go +++ b/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