Skip to content

Commit 352b7ec

Browse files
CopilotCopilot
authored
Harden gallery-agent Hugging Face fetches against transient rate limiting (#10187)
* Initial plan * fix: retry HuggingFace trending fetch on transient rate limits * fix: handle body close/write errors in huggingface retry paths --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
1 parent ba70642 commit 352b7ec

3 files changed

Lines changed: 201 additions & 34 deletions

File tree

‎.github/gallery-agent/main.go‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package main
33
import (
44
"context"
55
"encoding/json"
6+
"errors"
67
"fmt"
78
"os"
89
"strconv"
@@ -113,6 +114,17 @@ func main() {
113114
fmt.Println("Searching for trending models on HuggingFace...")
114115
rawModels, err := client.GetTrending(searchTerm, limit)
115116
if err != nil {
117+
if errors.Is(err, hfapi.ErrRateLimited) {
118+
fmt.Printf("HuggingFace API is rate limited after retries, skipping this run: %v\n", err)
119+
writeSummary(AddedModelSummary{
120+
SearchTerm: searchTerm,
121+
TotalFound: 0,
122+
ModelsAdded: 0,
123+
Quantization: quantization,
124+
ProcessingTime: time.Since(startTime).String(),
125+
})
126+
return
127+
}
116128
fmt.Fprintf(os.Stderr, "Error fetching models: %v\n", err)
117129
os.Exit(1)
118130
}
@@ -277,4 +289,3 @@ func truncateString(s string, maxLen int) string {
277289
}
278290
return s[:maxLen] + "..."
279291
}
280-

‎pkg/huggingface-api/client.go‎

Lines changed: 106 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package hfapi
22

33
import (
44
"encoding/json"
5+
"errors"
56
"fmt"
67
"io"
78
"net/http"
@@ -10,6 +11,7 @@ import (
1011
"sort"
1112
"strconv"
1213
"strings"
14+
"time"
1315

1416
"github.com/mudler/LocalAI/pkg/httpclient"
1517
)
@@ -88,57 +90,128 @@ type SearchParams struct {
8890

8991
// Client represents a Hugging Face API client
9092
type Client struct {
91-
baseURL string
92-
client *http.Client
93+
baseURL string
94+
client *http.Client
95+
maxRetries int
96+
retryBackoff time.Duration
97+
maxBackoff time.Duration
98+
sleepFn func(time.Duration)
9399
}
94100

101+
var ErrRateLimited = errors.New("huggingface API rate limited")
102+
95103
// NewClient creates a new Hugging Face API client
96104
func NewClient() *Client {
97105
return &Client{
98-
baseURL: "https://huggingface.co/api/models",
99-
client: httpclient.New(httpclient.WithFollowRedirects()),
106+
baseURL: "https://huggingface.co/api/models",
107+
client: httpclient.New(httpclient.WithFollowRedirects()),
108+
maxRetries: 5,
109+
retryBackoff: 1 * time.Second,
110+
maxBackoff: 30 * time.Second,
111+
sleepFn: time.Sleep,
100112
}
101113
}
102114

103115
// SearchModels searches for models using the Hugging Face API
104116
func (c *Client) SearchModels(params SearchParams) ([]Model, error) {
105-
req, err := http.NewRequest("GET", c.baseURL, nil)
106-
if err != nil {
107-
return nil, fmt.Errorf("failed to create request: %w", err)
108-
}
117+
for attempt := 1; attempt <= c.maxRetries; attempt++ {
118+
req, err := http.NewRequest("GET", c.baseURL, nil)
119+
if err != nil {
120+
return nil, fmt.Errorf("failed to create request: %w", err)
121+
}
109122

110-
// Add query parameters
111-
q := req.URL.Query()
112-
q.Add("sort", params.Sort)
113-
q.Add("direction", fmt.Sprintf("%d", params.Direction))
114-
q.Add("limit", fmt.Sprintf("%d", params.Limit))
115-
q.Add("search", params.Search)
116-
req.URL.RawQuery = q.Encode()
123+
// Add query parameters
124+
q := req.URL.Query()
125+
q.Add("sort", params.Sort)
126+
q.Add("direction", fmt.Sprintf("%d", params.Direction))
127+
q.Add("limit", fmt.Sprintf("%d", params.Limit))
128+
q.Add("search", params.Search)
129+
req.URL.RawQuery = q.Encode()
130+
131+
resp, err := c.client.Do(req)
132+
if err != nil {
133+
if attempt < c.maxRetries {
134+
c.sleepFn(c.exponentialBackoff(attempt))
135+
continue
136+
}
137+
return nil, fmt.Errorf("failed to make request: %w", err)
138+
}
117139

118-
// Make the HTTP request
119-
resp, err := c.client.Do(req)
120-
if err != nil {
121-
return nil, fmt.Errorf("failed to make request: %w", err)
122-
}
123-
defer resp.Body.Close()
140+
if resp.StatusCode != http.StatusOK {
141+
if err := resp.Body.Close(); err != nil {
142+
return nil, fmt.Errorf("failed to close response body: %w", err)
143+
}
144+
if c.isRetryableStatus(resp.StatusCode) && attempt < c.maxRetries {
145+
c.sleepFn(c.retryDelay(resp, attempt))
146+
continue
147+
}
148+
if resp.StatusCode == http.StatusTooManyRequests {
149+
return nil, fmt.Errorf("%w: failed to fetch models. Status code: %d", ErrRateLimited, resp.StatusCode)
150+
}
151+
return nil, fmt.Errorf("failed to fetch models. Status code: %d", resp.StatusCode)
152+
}
124153

125-
if resp.StatusCode != http.StatusOK {
126-
return nil, fmt.Errorf("failed to fetch models. Status code: %d", resp.StatusCode)
127-
}
154+
// Read the response body
155+
body, err := io.ReadAll(resp.Body)
156+
closeErr := resp.Body.Close()
157+
if err != nil {
158+
return nil, fmt.Errorf("failed to read response body: %w", err)
159+
}
160+
if closeErr != nil {
161+
return nil, fmt.Errorf("failed to close response body: %w", closeErr)
162+
}
128163

129-
// Read the response body
130-
body, err := io.ReadAll(resp.Body)
131-
if err != nil {
132-
return nil, fmt.Errorf("failed to read response body: %w", err)
164+
// Parse the JSON response
165+
var models []Model
166+
if err := json.Unmarshal(body, &models); err != nil {
167+
return nil, fmt.Errorf("failed to parse JSON response: %w", err)
168+
}
169+
170+
return models, nil
133171
}
134172

135-
// Parse the JSON response
136-
var models []Model
137-
if err := json.Unmarshal(body, &models); err != nil {
138-
return nil, fmt.Errorf("failed to parse JSON response: %w", err)
173+
return nil, fmt.Errorf("%w: failed to fetch models. Status code: %d", ErrRateLimited, http.StatusTooManyRequests)
174+
}
175+
176+
func (c *Client) isRetryableStatus(code int) bool {
177+
return code == http.StatusTooManyRequests || (code >= http.StatusInternalServerError && code <= http.StatusNetworkAuthenticationRequired)
178+
}
179+
180+
func (c *Client) retryDelay(resp *http.Response, attempt int) time.Duration {
181+
if retryAfter := strings.TrimSpace(resp.Header.Get("Retry-After")); retryAfter != "" {
182+
if seconds, err := strconv.Atoi(retryAfter); err == nil && seconds > 0 {
183+
delay := time.Duration(seconds) * time.Second
184+
if delay > c.maxBackoff {
185+
return c.maxBackoff
186+
}
187+
return delay
188+
}
189+
if at, err := http.ParseTime(retryAfter); err == nil {
190+
delay := time.Until(at)
191+
if delay > 0 {
192+
if delay > c.maxBackoff {
193+
return c.maxBackoff
194+
}
195+
return delay
196+
}
197+
}
139198
}
140199

141-
return models, nil
200+
return c.exponentialBackoff(attempt)
201+
}
202+
203+
func (c *Client) exponentialBackoff(attempt int) time.Duration {
204+
delay := c.retryBackoff
205+
for i := 1; i < attempt; i++ {
206+
delay *= 2
207+
if delay >= c.maxBackoff {
208+
return c.maxBackoff
209+
}
210+
}
211+
if delay > c.maxBackoff {
212+
return c.maxBackoff
213+
}
214+
return delay
142215
}
143216

144217
// GetLatest fetches the latest GGUF models

‎pkg/huggingface-api/client_test.go‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
package hfapi_test
22

33
import (
4+
"errors"
45
"fmt"
56
"net/http"
67
"net/http/httptest"
78
"strings"
9+
"time"
810

911
. "github.com/onsi/ginkgo/v2"
1012
. "github.com/onsi/gomega"
@@ -185,6 +187,87 @@ var _ = Describe("HuggingFace API Client", func() {
185187
Expect(err.Error()).To(ContainSubstring("failed to parse JSON response"))
186188
Expect(models).To(BeNil())
187189
})
190+
191+
It("should retry 429 responses and honor Retry-After", func() {
192+
attempts := 0
193+
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
194+
attempts++
195+
if attempts == 1 {
196+
w.Header().Set("Retry-After", "1")
197+
w.WriteHeader(http.StatusTooManyRequests)
198+
return
199+
}
200+
w.Header().Set("Content-Type", "application/json")
201+
w.WriteHeader(http.StatusOK)
202+
_, err := w.Write([]byte("[]"))
203+
Expect(err).ToNot(HaveOccurred())
204+
}))
205+
client.SetBaseURL(server.URL)
206+
207+
params := hfapi.SearchParams{
208+
Sort: "lastModified",
209+
Direction: -1,
210+
Limit: 30,
211+
Search: "GGUF",
212+
}
213+
214+
start := time.Now()
215+
models, err := client.SearchModels(params)
216+
elapsed := time.Since(start)
217+
218+
Expect(err).ToNot(HaveOccurred())
219+
Expect(models).To(HaveLen(0))
220+
Expect(attempts).To(Equal(2))
221+
Expect(elapsed).To(BeNumerically(">=", 900*time.Millisecond))
222+
})
223+
224+
It("should fail fast on non-retryable 4xx responses", func() {
225+
attempts := 0
226+
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
227+
attempts++
228+
w.WriteHeader(http.StatusBadRequest)
229+
}))
230+
client.SetBaseURL(server.URL)
231+
232+
params := hfapi.SearchParams{
233+
Sort: "lastModified",
234+
Direction: -1,
235+
Limit: 30,
236+
Search: "GGUF",
237+
}
238+
239+
start := time.Now()
240+
models, err := client.SearchModels(params)
241+
elapsed := time.Since(start)
242+
243+
Expect(err).To(HaveOccurred())
244+
Expect(err.Error()).To(ContainSubstring("Status code: 400"))
245+
Expect(models).To(BeNil())
246+
Expect(attempts).To(Equal(1))
247+
Expect(elapsed).To(BeNumerically("<", 500*time.Millisecond))
248+
})
249+
250+
It("should return ErrRateLimited when 429 persists after retries", func() {
251+
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
252+
w.Header().Set("Retry-After", "1")
253+
w.WriteHeader(http.StatusTooManyRequests)
254+
}))
255+
client.SetBaseURL(server.URL)
256+
257+
params := hfapi.SearchParams{
258+
Sort: "trendingScore",
259+
Direction: -1,
260+
Limit: 15,
261+
Search: "GGUF",
262+
}
263+
264+
models, err := client.SearchModels(params)
265+
266+
Expect(err).To(HaveOccurred())
267+
Expect(errors.Is(err, hfapi.ErrRateLimited)).To(BeTrue())
268+
Expect(err.Error()).To(ContainSubstring("Status code: 429"))
269+
Expect(models).To(BeNil())
270+
})
188271
})
189272

190273
Context("when getting latest GGUF models", func() {

0 commit comments

Comments
 (0)