diff --git a/pkg/downloader/download_plan.go b/pkg/downloader/download_plan.go index 84f0a21a2ab0..bc007ed005c5 100644 --- a/pkg/downloader/download_plan.go +++ b/pkg/downloader/download_plan.go @@ -1,6 +1,10 @@ package downloader -import "context" +import ( + "context" + + "github.com/mudler/xlog" +) // FileTask describes one download operation and an optional post-download // hook that runs after the bytes are present on disk. Callers keep any @@ -26,7 +30,7 @@ func DownloadFilesWithContext(ctx context.Context, tasks []FileTask, status func } taskOpts := append([]DownloadOption{}, opts...) taskOpts = append(taskOpts, task.Options...) - if err := task.URI.DownloadFileWithContext(ctx, task.Destination, task.SHA256, task.FileIndex, task.TotalFiles, status, taskOpts...); err != nil { + if err := downloadTaskWithRetry(ctx, task, status, taskOpts); err != nil { return err } if task.AfterDownload != nil { @@ -37,3 +41,30 @@ func DownloadFilesWithContext(ctx context.Context, tasks []FileTask, status func } return nil } + +// downloadTaskWithRetry fetches one file, retrying transient failures. Without +// this, a single cancelled stream anywhere in a large multi-file repo threw +// away every file already downloaded, and the .partial resume machinery was +// unreachable because nothing ever made a second attempt. +func downloadTaskWithRetry(ctx context.Context, task FileTask, status func(string, string, string, float64), opts []DownloadOption) error { + var err error + for attempt := 1; ; attempt++ { + err = task.URI.DownloadFileWithContext(ctx, task.Destination, task.SHA256, task.FileIndex, task.TotalFiles, status, opts...) + if err == nil { + return nil + } + if attempt >= DownloadRetryAttempts || !IsRetryable(ctx, err) { + return err + } + xlog.Warn("download failed, retrying", + "uri", string(task.URI), + "destination", task.Destination, + "attempt", attempt, + "maxAttempts", DownloadRetryAttempts, + "error", err, + ) + if waitErr := waitBeforeRetry(ctx, attempt); waitErr != nil { + return waitErr + } + } +} diff --git a/pkg/downloader/retry.go b/pkg/downloader/retry.go new file mode 100644 index 000000000000..3db9312460a2 --- /dev/null +++ b/pkg/downloader/retry.go @@ -0,0 +1,92 @@ +package downloader + +import ( + "context" + "errors" + "io" + "time" +) + +// ErrTransientDownload marks a download failure that a later attempt has a +// real chance of getting past: the peer cancelled the stream, the connection +// dropped, the transfer stalled, or the server answered 5xx. Everything else +// is treated as permanent, because retrying it either loops on a condition +// that will never change (404, checksum mismatch, a full disk) or burns +// bandwidth re-fetching gigabytes for nothing. +var ErrTransientDownload = errors.New("transient download failure") + +// DownloadRetryAttempts bounds how many times a single file is attempted +// before the whole plan gives up. These are multi-GB transfers, so the budget +// is deliberately small: the resume path means a retry is cheap in bytes, but +// an aggressive retry loop against a genuinely broken remote is its own +// outage. Overridable so tests (and operators on very flaky links) can adjust. +var DownloadRetryAttempts = 3 + +// DownloadRetryBaseDelay is the first backoff interval; each further retry +// doubles it. Backoff exists to let a momentarily overloaded remote recover, +// not to wait out a long outage, hence the small base and the low attempt cap. +var DownloadRetryBaseDelay = 2 * time.Second + +// transientError wraps an error while keeping its message intact, so the +// caller-facing diagnostics stay exactly as informative as before and only the +// retry classification changes. +type transientError struct{ err error } + +func (e *transientError) Error() string { return e.err.Error() } + +func (e *transientError) Unwrap() error { return e.err } + +func (e *transientError) Is(target error) bool { return target == ErrTransientDownload } + +// asTransient marks err retryable. A nil error stays nil so call sites can +// wrap unconditionally. +func asTransient(err error) error { + if err == nil { + return nil + } + return &transientError{err: err} +} + +// IsRetryable reports whether another attempt is worth making. A cancelled +// caller is never retried through, regardless of how the failure was +// classified: the caller has already given up. +func IsRetryable(ctx context.Context, err error) bool { + if err == nil || ctx.Err() != nil { + return false + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) || errors.Is(err, ErrUserCancelled) { + return false + } + return errors.Is(err, ErrTransientDownload) +} + +// readErrorRecorder remembers the last non-EOF error the source returned. +// io.Copy folds read and write failures into one return value; comparing the +// copy error against the recorded one is what lets the caller say which side +// of the transfer actually failed. +type readErrorRecorder struct { + r io.Reader + err error +} + +func (t *readErrorRecorder) Read(p []byte) (int, error) { + n, err := t.r.Read(p) + if err != nil && !errors.Is(err, io.EOF) { + t.err = err + } + return n, err +} + +// waitBeforeRetry sleeps for the backoff interval of the given attempt +// (1-based), returning the context error if the caller gives up while waiting. +func waitBeforeRetry(ctx context.Context, attempt int) error { + delay := DownloadRetryBaseDelay << (attempt - 1) + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} diff --git a/pkg/downloader/retry_test.go b/pkg/downloader/retry_test.go new file mode 100644 index 000000000000..5fca5644037d --- /dev/null +++ b/pkg/downloader/retry_test.go @@ -0,0 +1,231 @@ +package downloader_test + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strconv" + "strings" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/downloader" +) + +// flakyRangeServer serves payload, honours Range requests, and aborts the +// connection part-way through the body for the first failUntil GET requests. +// Aborting mid-body is how a peer-cancelled HTTP/2 stream or a dropped +// connection presents to the client: the read fails, the write never does. +func flakyRangeServer(payload []byte, failUntil int) (*httptest.Server, func() []string) { + var mu sync.Mutex + var seenRanges []string + gets := 0 + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + w.Header().Set("Accept-Ranges", "bytes") + w.WriteHeader(http.StatusOK) + return + } + + rangeHeader := r.Header.Get("Range") + mu.Lock() + gets++ + attempt := gets + seenRanges = append(seenRanges, rangeHeader) + mu.Unlock() + + start := 0 + if strings.HasPrefix(rangeHeader, "bytes=") { + n, err := strconv.Atoi(strings.TrimSuffix(strings.TrimPrefix(rangeHeader, "bytes="), "-")) + Expect(err).ToNot(HaveOccurred()) + start = n + } + body := payload[start:] + + w.Header().Set("Content-Length", strconv.Itoa(len(body))) + if start > 0 { + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, len(payload)-1, len(payload))) + w.WriteHeader(http.StatusPartialContent) + } else { + w.WriteHeader(http.StatusOK) + } + + if attempt <= failUntil { + // Send a prefix, then kill the connection: the client's Read fails + // with an unexpected EOF while the local file write succeeded. + _, _ = w.Write(body[:len(body)/2]) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + panic(http.ErrAbortHandler) + } + _, _ = w.Write(body) + })) + + return srv, func() []string { + mu.Lock() + defer mu.Unlock() + return append([]string(nil), seenRanges...) + } +} + +var _ = Describe("download failure diagnostics", func() { + var ( + payload []byte + payloadSHA string + destDir string + destPath string + noProgress = func(string, string, string, float64) {} + ) + + BeforeEach(func() { + payload = make([]byte, 65536) + _, err := rand.Read(payload) + Expect(err).ToNot(HaveOccurred()) + sum := sha256.Sum256(payload) + payloadSHA = fmt.Sprintf("%x", sum[:]) + + destDir = GinkgoT().TempDir() + destPath = filepath.Join(destDir, "model.gguf") + }) + + // Defect 1: io.Copy folds read and write failures into a single error, and + // the download path labelled every one of them "failed to write file". + // A peer-cancelled stream then reads as a filesystem problem, which is + // exactly the wrong place to look. + It("does not report a failed response read as a write failure", func() { + server, _ := flakyRangeServer(payload, 1) + DeferCleanup(server.Close) + + err := downloader.URI(server.URL).DownloadFile(destPath, payloadSHA, 1, 1, noProgress) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).ToNot(ContainSubstring("failed to write file"), + "the local write succeeded; only the response body read failed") + Expect(err.Error()).To(ContainSubstring(server.URL), + "a read-side failure should name the source it was reading from") + }) + + // The partial, not the final blob path, is the file actually being written. + It("names the partial file when the local write is the side that failed", func() { + // A directory in place of the partial makes every write fail while the + // response body stays healthy. + Expect(os.MkdirAll(destPath+downloader.PartialFileSuffix, 0750)).To(Succeed()) + + server, _ := flakyRangeServer(payload, 0) + DeferCleanup(server.Close) + + err := downloader.URI(server.URL).DownloadFile(destPath, payloadSHA, 1, 1, noProgress) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(destPath + downloader.PartialFileSuffix)) + }) +}) + +var _ = Describe("DownloadFilesWithContext retries", func() { + var ( + payload []byte + payloadSHA string + destPath string + ) + + BeforeEach(func() { + payload = make([]byte, 65536) + _, err := rand.Read(payload) + Expect(err).ToNot(HaveOccurred()) + sum := sha256.Sum256(payload) + payloadSHA = fmt.Sprintf("%x", sum[:]) + + destPath = filepath.Join(GinkgoT().TempDir(), "model.gguf") + + original := downloader.DownloadRetryBaseDelay + downloader.DownloadRetryBaseDelay = time.Millisecond + DeferCleanup(func() { downloader.DownloadRetryBaseDelay = original }) + }) + + // Defect 2: one cancelled stream half-way through a multi-file plan threw + // away every file already fetched. The .partial resume machinery existed + // but nothing retried, so it was unreachable. + It("retries a transient stream failure and resumes from the partial", func() { + server, ranges := flakyRangeServer(payload, 2) + DeferCleanup(server.Close) + + err := downloader.DownloadFilesWithContext(context.Background(), []downloader.FileTask{{ + URI: downloader.URI(server.URL), + Destination: destPath, + SHA256: payloadSHA, + FileIndex: 1, + TotalFiles: 1, + }}, nil) + Expect(err).ToNot(HaveOccurred()) + + got, err := os.ReadFile(destPath) + Expect(err).ToNot(HaveOccurred()) + Expect(got).To(Equal(payload)) + + seen := ranges() + Expect(seen).To(HaveLen(3), "expected two failed attempts and one success, got %v", seen) + Expect(seen[0]).To(BeEmpty()) + for _, r := range seen[1:] { + Expect(r).To(HavePrefix("bytes="), "retries must resume, not restart: %v", seen) + Expect(r).ToNot(Equal("bytes=0-"), "retries must resume, not restart: %v", seen) + } + }) + + It("does not retry a permanent failure", func() { + attempts := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempts++ + w.WriteHeader(http.StatusNotFound) + })) + DeferCleanup(server.Close) + + err := downloader.DownloadFilesWithContext(context.Background(), []downloader.FileTask{{ + URI: downloader.URI(server.URL), + Destination: destPath, + FileIndex: 1, + TotalFiles: 1, + }}, nil) + Expect(err).To(HaveOccurred()) + Expect(attempts).To(Equal(1), "a 404 is permanent and must not be retried") + }) + + It("does not retry once the caller's context is cancelled", func() { + ctx, cancel := context.WithCancel(context.Background()) + attempts := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + w.Header().Set("Accept-Ranges", "bytes") + w.WriteHeader(http.StatusOK) + return + } + attempts++ + cancel() + w.Header().Set("Content-Length", strconv.Itoa(len(payload))) + w.WriteHeader(http.StatusOK) + _, _ = w.Write(payload[:len(payload)/2]) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + panic(http.ErrAbortHandler) + })) + DeferCleanup(server.Close) + + err := downloader.DownloadFilesWithContext(ctx, []downloader.FileTask{{ + URI: downloader.URI(server.URL), + Destination: destPath, + SHA256: payloadSHA, + FileIndex: 1, + TotalFiles: 1, + }}, nil) + Expect(err).To(HaveOccurred()) + Expect(attempts).To(Equal(1), "a cancelled caller must not be retried through") + }) +}) diff --git a/pkg/downloader/uri.go b/pkg/downloader/uri.go index 6a730e0e0244..ef1b6b5f0f6c 100644 --- a/pkg/downloader/uri.go +++ b/pkg/downloader/uri.go @@ -673,21 +673,33 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string } return ctx.Err() } - return fmt.Errorf("failed to download file %q: %v", filePath, err) + // The transport failed before the response was established (reset + // connection, refused dial, TLS hiccup). Nothing about it is + // specific to this URL, so another attempt may well succeed. + return asTransient(fmt.Errorf("failed to download file %q: %v", filePath, err)) } //defer resp.Body.Close() if startPos > 0 && resp.StatusCode != http.StatusPartialContent { _ = resp.Body.Close() _ = removePartialFile(tmpFilePath) - return fmt.Errorf( + // The partial has just been discarded, so a further attempt starts + // clean and no longer needs the server to honour the range. + return asTransient(fmt.Errorf( "resume request for %q returned status %d instead of 206", filePath, resp.StatusCode, - ) + )) } if resp.StatusCode >= 400 { - return fmt.Errorf("failed to download url %q, invalid status code %d", url, resp.StatusCode) + err := fmt.Errorf("failed to download url %q, invalid status code %d", url, resp.StatusCode) + // 5xx and 429 describe the server's current state, not the request; + // every other 4xx (missing file, bad auth) is settled and retrying + // it only delays the real error. + if resp.StatusCode >= 500 || resp.StatusCode == http.StatusTooManyRequests { + return asTransient(err) + } + return err } source = resp.Body // Guard against a silently-stalled stream: a dropped TCP connection @@ -732,7 +744,12 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string ctx: ctx, } - _, err = xio.Copy(ctx, io.MultiWriter(outFile, progress), source) + // io.Copy reports read and write failures indistinguishably, so the source + // is wrapped to record which side actually broke. Labelling a peer-cancelled + // HTTP/2 stream "failed to write file" once sent an incident investigation + // after filesystem permissions while the disk was perfectly healthy. + tracked := &readErrorRecorder{r: source} + _, err = xio.Copy(ctx, io.MultiWriter(outFile, progress), tracked) if err != nil { // Detect cancellation via the context (a cause-cancelled read surfaces // the cause, not context.Canceled). Keep the .partial for resume, @@ -745,7 +762,18 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string } return ctx.Err() } - return fmt.Errorf("failed to write file %q: %v", filePath, err) + if readErr := tracked.err; readErr != nil && errors.Is(err, readErr) { + // The source died mid-transfer (peer cancelled the stream, the + // connection dropped, the stall guard fired). The bytes already on + // disk are valid, so the .partial is kept and the failure is + // retryable from where it stopped. + return asTransient(fmt.Errorf("failed to read %q while downloading to %q: %v", url, tmpFilePath, readErr)) + } + // A genuine local write failure: no space, bad permissions, a broken + // mount. Retrying writes the same bytes to the same broken target, so + // this stays permanent. Name the partial, which is the file actually + // being written, rather than the final blob path. + return fmt.Errorf("failed to write file %q: %v", tmpFilePath, err) } // Check for cancellation before finalizing. Keep the .partial for resume