Fix llama.cpp download progress updates
- commit
- c420b4eadde626812a24f71028e9bbed078d1403
- author
- patriceckhart · · 2026-08-06 05:49 UTC
- parent
- 4548953
- signature
- unsigned
2 files · +72 -2
M packages/provider/llamacpp.go +16 -2
| @@ -47,6 +47,10 @@ type LlamaCPPModel struct { | |||
| 47 | 47 | Size int64 `json:"size,omitempty"` | |
| 48 | 48 | FileType string `json:"ftype,omitempty"` | |
| 49 | 49 | } `json:"meta,omitempty"` | |
| 50 | + | // Progress is emitted at the model's top level by current llama.cpp | |
| 51 | + | // routers. Status.Progress supports versions that follow the documented | |
| 52 | + | // nested shape. | |
| 53 | + | Progress map[string]LlamaCPPBytes `json:"progress,omitempty"` | |
| 50 | 54 | } | |
| 51 | 55 | ||
| 52 | 56 | // LlamaCPPBytes describes byte progress for one downloaded file. | |
| @@ -204,7 +208,13 @@ func (c *LlamaCPPClient) watch(ctx context.Context, events chan<- llamaCPPEvent) | |||
| 204 | 208 | if c.APIKey != "" { | |
| 205 | 209 | req.Header.Set("authorization", "Bearer "+c.APIKey) | |
| 206 | 210 | } | |
| 207 | - | resp, err := c.HTTP.Do(req) | |
| 211 | + | // http.Client.Timeout covers the complete response body lifetime, which is | |
| 212 | + | // appropriate for ordinary management requests but would terminate this | |
| 213 | + | // long-lived SSE stream. Preserve the configured transport and other client | |
| 214 | + | // behavior while letting ctx control the stream lifetime. | |
| 215 | + | streamClient := *c.HTTP | |
| 216 | + | streamClient.Timeout = 0 | |
| 217 | + | resp, err := streamClient.Do(req) | |
| 208 | 218 | if err != nil { | |
| 209 | 219 | return err | |
| 210 | 220 | } | |
| @@ -408,7 +418,11 @@ func (c *LlamaCPPClient) DownloadAndWait(ctx context.Context, model string, upda | |||
| 408 | 418 | } | |
| 409 | 419 | if entry.Status.Value == "downloading" { | |
| 410 | 420 | sawDownload = true | |
| 411 | - | if progress, ok := downloadProgress(entry.Status.Progress); ok { | |
| 421 | + | files := entry.Status.Progress | |
| 422 | + | if len(files) == 0 { | |
| 423 | + | files = entry.Progress | |
| 424 | + | } | |
| 425 | + | if progress, ok := downloadProgress(files); ok { | |
| 412 | 426 | update(progress) | |
| 413 | 427 | } | |
| 414 | 428 | } else if sawDownload || polls >= 2 { | |
| @@ -47,6 +47,10 @@ type LlamaCPPModel struct { | |||||
| 47 | Size int64 `json:"size,omitempty"` | 47 | Size int64 `json:"size,omitempty"` | ||
| 48 | FileType string `json:"ftype,omitempty"` | 48 | FileType string `json:"ftype,omitempty"` | ||
| 49 | } `json:"meta,omitempty"` | 49 | } `json:"meta,omitempty"` | ||
| 50 | + | // Progress is emitted at the model's top level by current llama.cpp | |||
| 51 | + | // routers. Status.Progress supports versions that follow the documented | |||
| 52 | + | // nested shape. | |||
| 53 | + | Progress map[string]LlamaCPPBytes `json:"progress,omitempty"` | |||
| 50 | } | 54 | } | ||
| 51 | 55 | ||||
| 52 | // LlamaCPPBytes describes byte progress for one downloaded file. | 56 | // LlamaCPPBytes describes byte progress for one downloaded file. | ||
| @@ -204,7 +208,13 @@ func (c *LlamaCPPClient) watch(ctx context.Context, events chan<- llamaCPPEvent) | |||||
| 204 | if c.APIKey != "" { | 208 | if c.APIKey != "" { | ||
| 205 | req.Header.Set("authorization", "Bearer "+c.APIKey) | 209 | req.Header.Set("authorization", "Bearer "+c.APIKey) | ||
| 206 | } | 210 | } | ||
| 207 | - | resp, err := c.HTTP.Do(req) | 211 | + | // http.Client.Timeout covers the complete response body lifetime, which is |
| 212 | + | // appropriate for ordinary management requests but would terminate this | |||
| 213 | + | // long-lived SSE stream. Preserve the configured transport and other client | |||
| 214 | + | // behavior while letting ctx control the stream lifetime. | |||
| 215 | + | streamClient := *c.HTTP | |||
| 216 | + | streamClient.Timeout = 0 | |||
| 217 | + | resp, err := streamClient.Do(req) | |||
| 208 | if err != nil { | 218 | if err != nil { | ||
| 209 | return err | 219 | return err | ||
| 210 | } | 220 | } | ||
| @@ -408,7 +418,11 @@ func (c *LlamaCPPClient) DownloadAndWait(ctx context.Context, model string, upda | |||||
| 408 | } | 418 | } | ||
| 409 | if entry.Status.Value == "downloading" { | 419 | if entry.Status.Value == "downloading" { | ||
| 410 | sawDownload = true | 420 | sawDownload = true | ||
| 411 | - | if progress, ok := downloadProgress(entry.Status.Progress); ok { | 421 | + | files := entry.Status.Progress |
| 422 | + | if len(files) == 0 { | |||
| 423 | + | files = entry.Progress | |||
| 424 | + | } | |||
| 425 | + | if progress, ok := downloadProgress(files); ok { | |||
| 412 | update(progress) | 426 | update(progress) | ||
| 413 | } | 427 | } | ||
| 414 | } else if sawDownload || polls >= 2 { | 428 | } else if sawDownload || polls >= 2 { | ||
M packages/provider/llamacpp_test.go +56 -0
| @@ -32,6 +32,62 @@ func TestLlamaCPPLoadProgressStages(t *testing.T) { | |||
| 32 | 32 | } | |
| 33 | 33 | } | |
| 34 | 34 | ||
| 35 | + | func TestLlamaCPPWatchOutlivesRequestTimeout(t *testing.T) { | |
| 36 | + | requestStarted := make(chan struct{}) | |
| 37 | + | releaseEvent := make(chan struct{}) | |
| 38 | + | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| 39 | + | w.Header().Set("content-type", "text/event-stream") | |
| 40 | + | w.WriteHeader(http.StatusOK) | |
| 41 | + | w.(http.Flusher).Flush() | |
| 42 | + | close(requestStarted) | |
| 43 | + | select { | |
| 44 | + | case <-releaseEvent: | |
| 45 | + | fmt.Fprint(w, "data: {\"model\":\"test\",\"event\":\"download_progress\",\"data\":{}}\n\n") | |
| 46 | + | w.(http.Flusher).Flush() | |
| 47 | + | case <-r.Context().Done(): | |
| 48 | + | } | |
| 49 | + | })) | |
| 50 | + | defer server.Close() | |
| 51 | + | ||
| 52 | + | client, err := NewLlamaCPPClient(server.URL, "") | |
| 53 | + | if err != nil { | |
| 54 | + | t.Fatal(err) | |
| 55 | + | } | |
| 56 | + | client.HTTP.Timeout = 20 * time.Millisecond | |
| 57 | + | ctx, cancel := context.WithCancel(context.Background()) | |
| 58 | + | defer cancel() | |
| 59 | + | events := make(chan llamaCPPEvent, 1) | |
| 60 | + | watchDone := make(chan error, 1) | |
| 61 | + | go func() { watchDone <- client.watch(ctx, events) }() | |
| 62 | + | <-requestStarted | |
| 63 | + | ||
| 64 | + | select { | |
| 65 | + | case err := <-watchDone: | |
| 66 | + | t.Fatalf("watch stopped at ordinary request timeout: %v", err) | |
| 67 | + | case <-time.After(3 * client.HTTP.Timeout): | |
| 68 | + | } | |
| 69 | + | close(releaseEvent) | |
| 70 | + | select { | |
| 71 | + | case event := <-events: | |
| 72 | + | if event.Model != "test" || event.Event != "download_progress" { | |
| 73 | + | t.Fatalf("event = %#v", event) | |
| 74 | + | } | |
| 75 | + | case <-time.After(time.Second): | |
| 76 | + | t.Fatal("timed out waiting for delayed SSE event") | |
| 77 | + | } | |
| 78 | + | } | |
| 79 | + | ||
| 80 | + | func TestLlamaCPPModelAcceptsTopLevelDownloadProgress(t *testing.T) { | |
| 81 | + | var model LlamaCPPModel | |
| 82 | + | if err := json.Unmarshal([]byte(`{"id":"test","status":{"value":"downloading"},"progress":{"model.gguf":{"done":512,"total":1024}}}`), &model); err != nil { | |
| 83 | + | t.Fatal(err) | |
| 84 | + | } | |
| 85 | + | progress, ok := downloadProgress(model.Progress) | |
| 86 | + | if !ok || progress.Ratio != 0.5 { | |
| 87 | + | t.Fatalf("progress = %#v, ok = %v", progress, ok) | |
| 88 | + | } | |
| 89 | + | } | |
| 90 | + | ||
| 35 | 91 | func TestLlamaCPPClientLoadAndDownloadProgress(t *testing.T) { | |
| 36 | 92 | var mu sync.Mutex | |
| 37 | 93 | status := "unloaded" | |
| @@ -32,6 +32,62 @@ func TestLlamaCPPLoadProgressStages(t *testing.T) { | |||||
| 32 | } | 32 | } | ||
| 33 | } | 33 | } | ||
| 34 | 34 | ||||
| 35 | + | func TestLlamaCPPWatchOutlivesRequestTimeout(t *testing.T) { | |||
| 36 | + | requestStarted := make(chan struct{}) | |||
| 37 | + | releaseEvent := make(chan struct{}) | |||
| 38 | + | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |||
| 39 | + | w.Header().Set("content-type", "text/event-stream") | |||
| 40 | + | w.WriteHeader(http.StatusOK) | |||
| 41 | + | w.(http.Flusher).Flush() | |||
| 42 | + | close(requestStarted) | |||
| 43 | + | select { | |||
| 44 | + | case <-releaseEvent: | |||
| 45 | + | fmt.Fprint(w, "data: {\"model\":\"test\",\"event\":\"download_progress\",\"data\":{}}\n\n") | |||
| 46 | + | w.(http.Flusher).Flush() | |||
| 47 | + | case <-r.Context().Done(): | |||
| 48 | + | } | |||
| 49 | + | })) | |||
| 50 | + | defer server.Close() | |||
| 51 | + | ||||
| 52 | + | client, err := NewLlamaCPPClient(server.URL, "") | |||
| 53 | + | if err != nil { | |||
| 54 | + | t.Fatal(err) | |||
| 55 | + | } | |||
| 56 | + | client.HTTP.Timeout = 20 * time.Millisecond | |||
| 57 | + | ctx, cancel := context.WithCancel(context.Background()) | |||
| 58 | + | defer cancel() | |||
| 59 | + | events := make(chan llamaCPPEvent, 1) | |||
| 60 | + | watchDone := make(chan error, 1) | |||
| 61 | + | go func() { watchDone <- client.watch(ctx, events) }() | |||
| 62 | + | <-requestStarted | |||
| 63 | + | ||||
| 64 | + | select { | |||
| 65 | + | case err := <-watchDone: | |||
| 66 | + | t.Fatalf("watch stopped at ordinary request timeout: %v", err) | |||
| 67 | + | case <-time.After(3 * client.HTTP.Timeout): | |||
| 68 | + | } | |||
| 69 | + | close(releaseEvent) | |||
| 70 | + | select { | |||
| 71 | + | case event := <-events: | |||
| 72 | + | if event.Model != "test" || event.Event != "download_progress" { | |||
| 73 | + | t.Fatalf("event = %#v", event) | |||
| 74 | + | } | |||
| 75 | + | case <-time.After(time.Second): | |||
| 76 | + | t.Fatal("timed out waiting for delayed SSE event") | |||
| 77 | + | } | |||
| 78 | + | } | |||
| 79 | + | ||||
| 80 | + | func TestLlamaCPPModelAcceptsTopLevelDownloadProgress(t *testing.T) { | |||
| 81 | + | var model LlamaCPPModel | |||
| 82 | + | if err := json.Unmarshal([]byte(`{"id":"test","status":{"value":"downloading"},"progress":{"model.gguf":{"done":512,"total":1024}}}`), &model); err != nil { | |||
| 83 | + | t.Fatal(err) | |||
| 84 | + | } | |||
| 85 | + | progress, ok := downloadProgress(model.Progress) | |||
| 86 | + | if !ok || progress.Ratio != 0.5 { | |||
| 87 | + | t.Fatalf("progress = %#v, ok = %v", progress, ok) | |||
| 88 | + | } | |||
| 89 | + | } | |||
| 90 | + | ||||
| 35 | func TestLlamaCPPClientLoadAndDownloadProgress(t *testing.T) { | 91 | func TestLlamaCPPClientLoadAndDownloadProgress(t *testing.T) { | ||
| 36 | var mu sync.Mutex | 92 | var mu sync.Mutex | ||
| 37 | status := "unloaded" | 93 | status := "unloaded" | ||