diff --git a/internal/marian/download.go b/internal/marian/download.go index 85ad758..8d13540 100644 --- a/internal/marian/download.go +++ b/internal/marian/download.go @@ -4,6 +4,7 @@ import ( "context" "crypto/sha256" "encoding/hex" + "errors" "fmt" "io" "net/http" @@ -30,10 +31,12 @@ func Installed(dir string) bool { return true } -// Download fetches files into dir. Each file is written next to its -// final name and renamed once its digest is checked, so an interrupted -// download never leaves a half-written model behind. progress receives -// the bytes downloaded so far and the total. +// Download fetches files into dir. A file already there with the right +// size and digest is skipped. A file is downloaded into a ".partial" file +// that is renamed only once its digest checks out, so an interrupted +// download never leaves a half-written model behind; the partial file is +// kept and the next Download resumes it where it stopped. progress +// receives the bytes downloaded so far and the total. func Download(ctx context.Context, client *http.Client, dir string, files []RemoteFile, progress func(done, total int64)) error { if err := os.MkdirAll(dir, 0o755); err != nil { return err @@ -42,57 +45,145 @@ func Download(ctx context.Context, client *http.Client, dir string, files []Remo for _, f := range files { total += f.Size } + report := func(n int64) { + if progress != nil { + progress(done+n, total) + } + } for _, f := range files { - n, err := download(ctx, client, dir, f, func(n int64) { - if progress != nil { - progress(done+n, total) - } - }) + report(0) + n, err := download(ctx, client, dir, f, report) if err != nil { return fmt.Errorf("downloading %s: %w", f.Name, err) } done += n + report(0) } return nil } +// download fetches one file, resuming the partial download left by an +// interrupted run when the server supports it. func download(ctx context.Context, client *http.Client, dir string, f RemoteFile, progress func(int64)) (int64, error) { + final := filepath.Join(dir, f.Name) + if ok, err := verify(final, f); err != nil { + return 0, err + } else if ok { + return f.Size, nil // downloaded by a previous run + } + partial := filepath.Join(dir, "."+f.Name+".partial") + var off int64 + if st, err := os.Stat(partial); err == nil && st.Size() > 0 && st.Size() <= f.Size { + off = st.Size() + } + if off == f.Size { + // A previous run was interrupted between the digest check and + // the rename. + if ok, _ := verify(partial, f); !ok { + off = 0 // corrupt: start over + } else if err := os.Rename(partial, final); err != nil { + return 0, err + } else { + return f.Size, nil + } + } req, err := http.NewRequestWithContext(ctx, http.MethodGet, f.URL, nil) if err != nil { return 0, err } + if off > 0 { + req.Header.Set("Range", fmt.Sprintf("bytes=%d-", off)) + } resp, err := client.Do(req) if err != nil { return 0, err } defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { + if off > 0 && resp.StatusCode == http.StatusOK { + off = 0 // the server cannot resume: start over + } + want := http.StatusOK + if off > 0 { + want = http.StatusPartialContent + } + if resp.StatusCode != want { + if resp.StatusCode == http.StatusRequestedRangeNotSatisfiable { + os.Remove(partial) // the partial file cannot be resumed + } return 0, fmt.Errorf("server answered %s", resp.Status) } - tmp, err := os.CreateTemp(dir, "."+f.Name+".*") + tmp, err := os.OpenFile(partial, os.O_CREATE|os.O_RDWR, 0o644) if err != nil { return 0, err } - defer os.Remove(tmp.Name()) // no-op once renamed hash := sha256.New() - w := &countingWriter{w: io.MultiWriter(tmp, hash), progress: progress} - _, err = io.Copy(w, io.LimitReader(resp.Body, f.Size+1)) + if off > 0 { + // The digest covers the whole file: seed it with the bytes + // already on disk. + if _, err := io.CopyN(hash, tmp, off); err != nil { + tmp.Close() + return 0, fmt.Errorf("reading the partial download: %w", err) + } + } else if err := tmp.Truncate(0); err != nil { + tmp.Close() + return 0, err + } + if _, err := tmp.Seek(off, io.SeekStart); err != nil { + tmp.Close() + return 0, err + } + w := &countingWriter{w: io.MultiWriter(tmp, hash), progress: func(n int64) { progress(off + n) }} + _, err = io.Copy(w, io.LimitReader(resp.Body, f.Size-off+1)) if cerr := tmp.Close(); err == nil { err = cerr } if err != nil { - return 0, err + return 0, err // the partial file stays, for the next run } - if w.n != f.Size { - return 0, fmt.Errorf("got %d bytes, expected %d", w.n, f.Size) + if got := off + w.n; got != f.Size { + return 0, fmt.Errorf("got %d bytes, expected %d", got, f.Size) } if sum := hex.EncodeToString(hash.Sum(nil)); sum != f.SHA256 { + os.Remove(partial) // corrupt: resuming it would fail again return 0, fmt.Errorf("checksum mismatch") } - if err := os.Chmod(tmp.Name(), 0o644); err != nil { + if err := os.Chmod(partial, 0o644); err != nil { return 0, err } - return w.n, os.Rename(tmp.Name(), filepath.Join(dir, f.Name)) + return f.Size, os.Rename(partial, final) +} + +// verify reports whether path holds the file, with its expected size and +// digest. A missing file is not an error. +func verify(path string, f RemoteFile) (bool, error) { + st, err := os.Stat(path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return false, nil + } + return false, err + } + if st.Size() != f.Size { + return false, nil + } + sum, err := hashFile(path) + if err != nil { + return false, err + } + return sum == f.SHA256, nil +} + +func hashFile(path string) (string, error) { + r, err := os.Open(path) + if err != nil { + return "", err + } + defer r.Close() + h := sha256.New() + if _, err := io.Copy(h, r); err != nil { + return "", err + } + return hex.EncodeToString(h.Sum(nil)), nil } type countingWriter struct { diff --git a/internal/marian/marian_test.go b/internal/marian/marian_test.go index dd5e057..533e1b1 100644 --- a/internal/marian/marian_test.go +++ b/internal/marian/marian_test.go @@ -127,7 +127,7 @@ func TestSentencePiece(t *testing.T) { func TestSafetensors(t *testing.T) { path := filepath.Join(t.TempDir(), "m.safetensors") writeSafetensors(t, path, []testTensor{{"w", []int{2, 2}, []float32{1, -2, 3.5, 0}}}) - got, err := readSafetensors(path) + got, err := readSafetensors(context.Background(), path) if err != nil { t.Fatal(err) } @@ -141,7 +141,7 @@ func TestSafetensors(t *testing.T) { data = append(data, h...) data = append(data, 0x00, 0x3c, 0x00, 0xc0, 0x80, 0x3f) // 1, -2 (F16), 1 (BF16) writeFile(t, path, data) - got, err = readSafetensors(path) + got, err = readSafetensors(context.Background(), path) if err != nil { t.Fatal(err) } @@ -151,7 +151,7 @@ func TestSafetensors(t *testing.T) { h = []byte(`{"a":{"dtype":"F32","shape":[3],"data_offsets":[0,4]}}`) writeFile(t, path, append(binary.LittleEndian.AppendUint64(nil, uint64(len(h))), append(h, 0, 0, 0, 0)...)) - if _, err := readSafetensors(path); err == nil { + if _, err := readSafetensors(context.Background(), path); err == nil { t.Error("shape mismatch accepted") } } @@ -254,7 +254,7 @@ func TestTinyModel(t *testing.T) { if !Installed(dir) { t.Fatal("tiny model not installed") } - m, err := Load(dir) + m, err := Load(context.Background(), dir) if err != nil { t.Fatal(err) } @@ -284,7 +284,7 @@ func TestTinyModel(t *testing.T) { if Installed(dir) { t.Error("incomplete model reported installed") } - if _, err := Load(dir); err == nil { + if _, err := Load(context.Background(), dir); err == nil { t.Error("incomplete model loaded") } } @@ -297,7 +297,7 @@ func TestRealModel(t *testing.T) { if dir == "" { t.Skip("PLOTZILLA_MT_MODEL not set") } - m, err := Load(dir) + m, err := Load(context.Background(), dir) if err != nil { t.Fatal(err) } @@ -315,12 +315,23 @@ func TestRealModel(t *testing.T) { func TestDownload(t *testing.T) { content := []byte("model weights") sum := sha256.Sum256(content) + var requests int srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/missing" { + requests++ + switch r.URL.Path { + case "/missing": http.NotFound(w, r) - return + case "/full": // a server without resume support + w.Write(content) + default: + var off int + if _, err := fmt.Sscanf(r.Header.Get("Range"), "bytes=%d-", &off); err == nil && off > 0 && off < len(content) { + w.WriteHeader(http.StatusPartialContent) + w.Write(content[off:]) + return + } + w.Write(content) } - w.Write(content) })) defer srv.Close() dir := filepath.Join(t.TempDir(), "model") @@ -334,6 +345,52 @@ func TestDownload(t *testing.T) { t.Errorf("downloaded %q, progress %d", got, last) } + // A file already downloaded and verified is skipped. + if err := Download(context.Background(), srv.Client(), dir, []RemoteFile{file}, nil); err != nil { + t.Fatal(err) + } + if requests != 1 { + t.Errorf("a complete file was downloaded again (%d requests)", requests) + } + + // An interrupted download resumes its partial file. + resumed := file + resumed.Name, resumed.URL = "r.bin", srv.URL+"/r" + writeFile(t, filepath.Join(dir, ".r.bin.partial"), content[:6]) + if err := Download(context.Background(), srv.Client(), dir, []RemoteFile{resumed}, nil); err != nil { + t.Fatal(err) + } + if got, _ := os.ReadFile(filepath.Join(dir, "r.bin")); string(got) != string(content) { + t.Errorf("resumed download = %q", got) + } + + // A server that cannot resume answers the whole file: the partial + // download is restarted. + full := file + full.Name, full.URL = "full.bin", srv.URL+"/full" + writeFile(t, filepath.Join(dir, ".full.bin.partial"), content[:6]) + if err := Download(context.Background(), srv.Client(), dir, []RemoteFile{full}, nil); err != nil { + t.Fatal(err) + } + if got, _ := os.ReadFile(filepath.Join(dir, "full.bin")); string(got) != string(content) { + t.Errorf("restarted download = %q", got) + } + + // A corrupt partial file fails the digest check and is removed: the + // next download starts from scratch. + corrupt := file + corrupt.Name, corrupt.URL = "c.bin", srv.URL+"/c" + writeFile(t, filepath.Join(dir, ".c.bin.partial"), []byte("garbag")) + if err := Download(context.Background(), srv.Client(), dir, []RemoteFile{corrupt}, nil); err == nil { + t.Error("corrupt partial file accepted") + } + if _, err := os.Stat(filepath.Join(dir, ".c.bin.partial")); err == nil { + t.Error("corrupt partial file kept") + } + if err := Download(context.Background(), srv.Client(), dir, []RemoteFile{corrupt}, nil); err != nil { + t.Fatal(err) + } + bad := file bad.Name, bad.SHA256 = "bad.bin", strings.Repeat("0", 64) missing := file @@ -348,8 +405,13 @@ func TestDownload(t *testing.T) { t.Errorf("%s: a failed download left the file", f.Name) } } - entries, _ := os.ReadDir(dir) - if len(entries) != 1 { - t.Errorf("temporary files left: %v", entries) + // Only the interrupted download keeps a partial file, to be resumed. + for _, f := range []RemoteFile{bad, missing} { + if _, err := os.Stat(filepath.Join(dir, "."+f.Name+".partial")); err == nil { + t.Errorf("%s: a failed download left a partial file", f.Name) + } + } + if _, err := os.Stat(filepath.Join(dir, ".short.bin.partial")); err != nil { + t.Error("an interrupted download left no partial file to resume") } } diff --git a/internal/marian/model.go b/internal/marian/model.go index ec0c91b..2db4478 100644 --- a/internal/marian/model.go +++ b/internal/marian/model.go @@ -79,8 +79,12 @@ type Model struct { activation func(float32) float32 } -// Load reads the model in dir. -func Load(dir string) (*Model, error) { +// Load reads the model in dir. The load can take seconds for a large +// model: it stops early when ctx is cancelled. +func Load(ctx context.Context, dir string) (*Model, error) { + if err := ctx.Err(); err != nil { + return nil, err + } var cfg config data, err := os.ReadFile(filepath.Join(dir, "config.json")) if err != nil { @@ -130,7 +134,10 @@ func Load(dir string) (*Model, error) { } m.unk = m.ids[""] - tensors, err := readSafetensors(filepath.Join(dir, "model.safetensors")) + if err := ctx.Err(); err != nil { + return nil, err + } + tensors, err := readSafetensors(ctx, filepath.Join(dir, "model.safetensors")) if err != nil { return nil, fmt.Errorf("reading model.safetensors: %w", err) } diff --git a/internal/marian/safetensors.go b/internal/marian/safetensors.go index e12d092..77c124b 100644 --- a/internal/marian/safetensors.go +++ b/internal/marian/safetensors.go @@ -1,6 +1,7 @@ package marian import ( + "context" "encoding/binary" "encoding/json" "fmt" @@ -27,7 +28,7 @@ type tensorInfo struct { // little-endian header length, a JSON header, then the raw data) as // float32. F32 tensors are read straight into their slice: the file is // large and a second copy would double the memory needed to load it. -func readSafetensors(path string) (map[string]tensor, error) { +func readSafetensors(ctx context.Context, path string) (map[string]tensor, error) { f, err := os.Open(path) if err != nil { return nil, err @@ -54,6 +55,9 @@ func readSafetensors(path string) (map[string]tensor, error) { if name == "__metadata__" { continue } + if err := ctx.Err(); err != nil { + return nil, err + } var info tensorInfo if err := json.Unmarshal(raw, &info); err != nil { return nil, fmt.Errorf("tensor %s: %w", name, err) diff --git a/translation.go b/translation.go index ff94740..ad13215 100644 --- a/translation.go +++ b/translation.go @@ -3,6 +3,7 @@ package main import ( "context" "errors" + "fmt" "net/http" "os" "path/filepath" @@ -28,8 +29,14 @@ const ( // modelIdle is how long the model (300 MB of memory) stays loaded // after a translation, ready for the next one. modelIdle = 5 * time.Minute + // stallTimeout is how long a download may go without receiving a byte + // (or its response headers) before it is aborted as stalled. + stallTimeout = 30 * time.Second ) +// errStalled aborts a download that receives nothing for stallTimeout. +var errStalled = errors.New("the download stalled (no data received for 30 seconds)") + var modelFiles = []marian.RemoteFile{ {Name: "config.json", Size: 1416, SHA256: "b3be13d046d9899d7aab8cf4ed624d9a79f5776038ba793f6b4d2ce3e02192f7"}, {Name: "source.spm", Size: 802397, SHA256: "78d0e717c77053f1c4b856d8661d9cb87c64f083a35418c087b9146300e4f585"}, @@ -46,6 +53,7 @@ type translation struct { lexicon translate.Lexicon release *time.Timer cancel context.CancelFunc + done chan struct{} // closed when the running task is over } // TranslationStatus says whether the model is installed, and how much a @@ -72,15 +80,29 @@ func modelDirs() []string { return dirs } +// installedModel returns the first model directory holding every file at +// its pinned size. Checking the digests would mean reading 300 MB on +// every call; the sizes catch truncated or stale files, and the digests +// were checked when the files were downloaded. func installedModel() string { for _, d := range modelDirs() { - if marian.Installed(d) { + if modelComplete(d) { return d } } return "" } +func modelComplete(dir string) bool { + for _, f := range modelFiles { + st, err := os.Stat(filepath.Join(dir, f.Name)) + if err != nil || st.Size() != f.Size { + return false + } + } + return true +} + func (a *App) TranslationStatus() TranslationStatus { var size int64 for _, f := range modelFiles { @@ -94,10 +116,11 @@ func (a *App) startTask() (context.Context, error) { a.tr.mu.Lock() defer a.tr.mu.Unlock() if a.tr.cancel != nil { - return nil, errors.New("a translation is already running") + return nil, errors.New("a translation task is already running") } ctx, cancel := context.WithCancel(a.ctx) a.tr.cancel = cancel + a.tr.done = make(chan struct{}) return ctx, nil } @@ -107,16 +130,22 @@ func (a *App) endTask() { if a.tr.cancel != nil { a.tr.cancel() a.tr.cancel = nil + close(a.tr.done) + a.tr.done = nil } } -// CancelTranslation stops the running download or translation. +// CancelTranslation stops the running download or translation, waiting +// for it to wind down so that a new task can start right away. func (a *App) CancelTranslation() { a.tr.mu.Lock() - defer a.tr.mu.Unlock() - if a.tr.cancel != nil { - a.tr.cancel() + cancel, done := a.tr.cancel, a.tr.done + a.tr.mu.Unlock() + if cancel == nil { + return } + cancel() + <-done } // progressEmitter sends "translation-progress" events, at most every @@ -136,8 +165,12 @@ func (a *App) progressEmitter() func(done, total int64) { } // DownloadTranslationModel downloads the model into the user's data -// directory. It returns false if the user cancelled. +// directory. An interrupted download resumes where it stopped. It returns +// false if the user cancelled. func (a *App) DownloadTranslationModel() (bool, error) { + if installedModel() != "" { + return true, nil + } ctx, err := a.startTask() if err != nil { return false, err @@ -148,11 +181,30 @@ func (a *App) DownloadTranslationModel() (bool, error) { f.URL = modelURL + f.Name files[i] = f } - err = marian.Download(ctx, &http.Client{}, modelDirs()[0], files, a.progressEmitter()) - if errors.Is(err, context.Canceled) { - return false, nil + // No total timeout (300 MB can take a while) but abort a stalled + // transfer: the watchdog is reset by every progress report. + ctx, stop := context.WithCancelCause(ctx) + defer stop(nil) + emit := a.progressEmitter() + stalled := time.AfterFunc(stallTimeout, func() { stop(errStalled) }) + defer stalled.Stop() + transport := http.DefaultTransport.(*http.Transport).Clone() + transport.ResponseHeaderTimeout = stallTimeout + client := &http.Client{Transport: transport} + err = marian.Download(ctx, client, modelDirs()[0], files, func(done, total int64) { + stalled.Reset(stallTimeout) + emit(done, total) + }) + if err != nil { + if errors.Is(context.Cause(ctx), errStalled) { + return false, errStalled + } + if errors.Is(err, context.Canceled) { + return false, nil + } + return false, err } - return err == nil, err + return true, nil } // Translate translates the French text of the elements (flat indexes; all @@ -164,8 +216,14 @@ func (a *App) Translate(doc screenplay.Document, elements []int, titlePage bool) return nil, err } defer a.endTask() + if installedModel() == "" { + return nil, errors.New("the translation model is not installed") + } wails.EventsEmit(a.ctx, "translation-progress", 0, 0) // loading - model, lexicon, err := a.translator() + model, lexicon, err := a.translator(ctx) + if errors.Is(err, context.Canceled) { + return nil, nil + } if err != nil { return nil, err } @@ -189,38 +247,60 @@ func (a *App) Translate(doc screenplay.Document, elements []int, titlePage bool) } // translator returns the model and the dictionaries that help tell French -// from English, loading them the first time. -func (a *App) translator() (*marian.Model, translate.Lexicon, error) { +// from English, loading them the first time. The load happens without +// a.tr.mu: reading 300 MB takes seconds, and CancelTranslation must stay +// responsive. +func (a *App) translator(ctx context.Context) (*marian.Model, translate.Lexicon, error) { a.tr.mu.Lock() - defer a.tr.mu.Unlock() if a.tr.release != nil { a.tr.release.Stop() } - if a.tr.model == nil { - dir := installedModel() - if dir == "" { - return nil, nil, errors.New("the translation model is not installed") - } - m, err := marian.Load(dir) - if err != nil { - return nil, nil, err - } - a.tr.model = m - a.tr.lexicon = loadLexicon() + model, lexicon := a.tr.model, a.tr.lexicon + a.tr.mu.Unlock() + if model != nil { + return model, lexicon, nil } - return a.tr.model, a.tr.lexicon, nil + dir := installedModel() + if dir == "" { + return nil, nil, errors.New("the translation model is not installed") + } + m, err := marian.Load(ctx, dir) + if err != nil { + if ctx.Err() != nil { + return nil, nil, ctx.Err() + } + return nil, nil, fmt.Errorf("loading the translation model: %w", err) + } + lexicon = loadLexicon() + a.tr.mu.Lock() + if ctx.Err() == nil { // cancelled while loading: cache nothing + a.tr.model, a.tr.lexicon = m, lexicon + } + a.tr.mu.Unlock() + if err := ctx.Err(); err != nil { + return nil, nil, err + } + return m, lexicon, nil } +// releaseLater frees the model once it has been unused for modelIdle. func (a *App) releaseLater() { a.tr.mu.Lock() defer a.tr.mu.Unlock() - a.tr.release = time.AfterFunc(modelIdle, func() { - a.tr.mu.Lock() - defer a.tr.mu.Unlock() - if a.tr.cancel == nil { - a.tr.model, a.tr.lexicon = nil, nil - } - }) + a.tr.release = time.AfterFunc(modelIdle, a.releaseModel) +} + +// releaseModel frees the model when no task is running. When one is (a +// download loads nothing but holds the task slot), it tries again after +// another modelIdle. +func (a *App) releaseModel() { + a.tr.mu.Lock() + defer a.tr.mu.Unlock() + if a.tr.cancel != nil { + a.tr.release = time.AfterFunc(modelIdle, a.releaseModel) + return + } + a.tr.model, a.tr.lexicon = nil, nil } // lexicon tells French words from English ones with the installed