feat: resumable download and cancellable load of the translation model
Downloading the 300 MB model restarted from zero after any failure and hung forever on a stalled connection. Download into a .partial file that survives interruptions, resume it with a Range request (falling back to a fresh download when the server cannot), skip files already verified, and abort when nothing arrives for 30 s. A model directory only counts as installed when every file has its pinned size, so a truncated model is offered for download again instead of failing to load. The model load held the translation mutex, so Cancel did nothing for the seconds it takes to read 300 MB. Load without holding the lock, make marian.Load context-aware, and let CancelTranslation wait for the task to wind down so a new one can start at once. The idle release timer re-arms when it fires while a download holds the task slot, instead of never freeing the model.
This commit is contained in:
1 parent
8181c6e885
commit
839e9291d6
5 files changed
+313
-69
No files matched your search
+110
-19
@@ -4,6 +4,7 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"crypto/sha256"
|
"crypto/sha256"
|
||||||
"encoding/hex"
|
"encoding/hex"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
@@ -30,10 +31,12 @@ func Installed(dir string) bool {
|
|||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
|
|
||||||
// Download fetches files into dir. Each file is written next to its
|
// Download fetches files into dir. A file already there with the right
|
||||||
// final name and renamed once its digest is checked, so an interrupted
|
// size and digest is skipped. A file is downloaded into a ".partial" file
|
||||||
// download never leaves a half-written model behind. progress receives
|
// that is renamed only once its digest checks out, so an interrupted
|
||||||
// the bytes downloaded so far and the total.
|
// 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 {
|
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 {
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||||
return err
|
return err
|
||||||
@@ -42,57 +45,145 @@ func Download(ctx context.Context, client *http.Client, dir string, files []Remo
|
|||||||
for _, f := range files {
|
for _, f := range files {
|
||||||
total += f.Size
|
total += f.Size
|
||||||
}
|
}
|
||||||
|
report := func(n int64) {
|
||||||
|
if progress != nil {
|
||||||
|
progress(done+n, total)
|
||||||
|
}
|
||||||
|
}
|
||||||
for _, f := range files {
|
for _, f := range files {
|
||||||
n, err := download(ctx, client, dir, f, func(n int64) {
|
report(0)
|
||||||
if progress != nil {
|
n, err := download(ctx, client, dir, f, report)
|
||||||
progress(done+n, total)
|
|
||||||
}
|
|
||||||
})
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("downloading %s: %w", f.Name, err)
|
return fmt.Errorf("downloading %s: %w", f.Name, err)
|
||||||
}
|
}
|
||||||
done += n
|
done += n
|
||||||
|
report(0)
|
||||||
}
|
}
|
||||||
return nil
|
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) {
|
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)
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, f.URL, nil)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return 0, err
|
return 0, err
|
||||||
}
|
}
|
||||||
|
if off > 0 {
|
||||||
|
req.Header.Set("Range", fmt.Sprintf("bytes=%d-", off))
|
||||||
|
}
|
||||||
resp, err := client.Do(req)
|
resp, err := client.Do(req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return 0, err
|
return 0, err
|
||||||
}
|
}
|
||||||
defer resp.Body.Close()
|
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)
|
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 {
|
if err != nil {
|
||||||
return 0, err
|
return 0, err
|
||||||
}
|
}
|
||||||
defer os.Remove(tmp.Name()) // no-op once renamed
|
|
||||||
hash := sha256.New()
|
hash := sha256.New()
|
||||||
w := &countingWriter{w: io.MultiWriter(tmp, hash), progress: progress}
|
if off > 0 {
|
||||||
_, err = io.Copy(w, io.LimitReader(resp.Body, f.Size+1))
|
// 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 {
|
if cerr := tmp.Close(); err == nil {
|
||||||
err = cerr
|
err = cerr
|
||||||
}
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return 0, err
|
return 0, err // the partial file stays, for the next run
|
||||||
}
|
}
|
||||||
if w.n != f.Size {
|
if got := off + w.n; got != f.Size {
|
||||||
return 0, fmt.Errorf("got %d bytes, expected %d", w.n, f.Size)
|
return 0, fmt.Errorf("got %d bytes, expected %d", got, f.Size)
|
||||||
}
|
}
|
||||||
if sum := hex.EncodeToString(hash.Sum(nil)); sum != f.SHA256 {
|
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")
|
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 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 {
|
type countingWriter struct {
|
||||||
|
|||||||
@@ -127,7 +127,7 @@ func TestSentencePiece(t *testing.T) {
|
|||||||
func TestSafetensors(t *testing.T) {
|
func TestSafetensors(t *testing.T) {
|
||||||
path := filepath.Join(t.TempDir(), "m.safetensors")
|
path := filepath.Join(t.TempDir(), "m.safetensors")
|
||||||
writeSafetensors(t, path, []testTensor{{"w", []int{2, 2}, []float32{1, -2, 3.5, 0}}})
|
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 {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
@@ -141,7 +141,7 @@ func TestSafetensors(t *testing.T) {
|
|||||||
data = append(data, h...)
|
data = append(data, h...)
|
||||||
data = append(data, 0x00, 0x3c, 0x00, 0xc0, 0x80, 0x3f) // 1, -2 (F16), 1 (BF16)
|
data = append(data, 0x00, 0x3c, 0x00, 0xc0, 0x80, 0x3f) // 1, -2 (F16), 1 (BF16)
|
||||||
writeFile(t, path, data)
|
writeFile(t, path, data)
|
||||||
got, err = readSafetensors(path)
|
got, err = readSafetensors(context.Background(), path)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
@@ -151,7 +151,7 @@ func TestSafetensors(t *testing.T) {
|
|||||||
|
|
||||||
h = []byte(`{"a":{"dtype":"F32","shape":[3],"data_offsets":[0,4]}}`)
|
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)...))
|
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")
|
t.Error("shape mismatch accepted")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -254,7 +254,7 @@ func TestTinyModel(t *testing.T) {
|
|||||||
if !Installed(dir) {
|
if !Installed(dir) {
|
||||||
t.Fatal("tiny model not installed")
|
t.Fatal("tiny model not installed")
|
||||||
}
|
}
|
||||||
m, err := Load(dir)
|
m, err := Load(context.Background(), dir)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
@@ -284,7 +284,7 @@ func TestTinyModel(t *testing.T) {
|
|||||||
if Installed(dir) {
|
if Installed(dir) {
|
||||||
t.Error("incomplete model reported installed")
|
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")
|
t.Error("incomplete model loaded")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -297,7 +297,7 @@ func TestRealModel(t *testing.T) {
|
|||||||
if dir == "" {
|
if dir == "" {
|
||||||
t.Skip("PLOTZILLA_MT_MODEL not set")
|
t.Skip("PLOTZILLA_MT_MODEL not set")
|
||||||
}
|
}
|
||||||
m, err := Load(dir)
|
m, err := Load(context.Background(), dir)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
@@ -315,12 +315,23 @@ func TestRealModel(t *testing.T) {
|
|||||||
func TestDownload(t *testing.T) {
|
func TestDownload(t *testing.T) {
|
||||||
content := []byte("model weights")
|
content := []byte("model weights")
|
||||||
sum := sha256.Sum256(content)
|
sum := sha256.Sum256(content)
|
||||||
|
var requests int
|
||||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
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)
|
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()
|
defer srv.Close()
|
||||||
dir := filepath.Join(t.TempDir(), "model")
|
dir := filepath.Join(t.TempDir(), "model")
|
||||||
@@ -334,6 +345,52 @@ func TestDownload(t *testing.T) {
|
|||||||
t.Errorf("downloaded %q, progress %d", got, last)
|
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 := file
|
||||||
bad.Name, bad.SHA256 = "bad.bin", strings.Repeat("0", 64)
|
bad.Name, bad.SHA256 = "bad.bin", strings.Repeat("0", 64)
|
||||||
missing := file
|
missing := file
|
||||||
@@ -348,8 +405,13 @@ func TestDownload(t *testing.T) {
|
|||||||
t.Errorf("%s: a failed download left the file", f.Name)
|
t.Errorf("%s: a failed download left the file", f.Name)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
entries, _ := os.ReadDir(dir)
|
// Only the interrupted download keeps a partial file, to be resumed.
|
||||||
if len(entries) != 1 {
|
for _, f := range []RemoteFile{bad, missing} {
|
||||||
t.Errorf("temporary files left: %v", entries)
|
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")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -79,8 +79,12 @@ type Model struct {
|
|||||||
activation func(float32) float32
|
activation func(float32) float32
|
||||||
}
|
}
|
||||||
|
|
||||||
// Load reads the model in dir.
|
// Load reads the model in dir. The load can take seconds for a large
|
||||||
func Load(dir string) (*Model, error) {
|
// 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
|
var cfg config
|
||||||
data, err := os.ReadFile(filepath.Join(dir, "config.json"))
|
data, err := os.ReadFile(filepath.Join(dir, "config.json"))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -130,7 +134,10 @@ func Load(dir string) (*Model, error) {
|
|||||||
}
|
}
|
||||||
m.unk = m.ids["<unk>"]
|
m.unk = m.ids["<unk>"]
|
||||||
|
|
||||||
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 {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("reading model.safetensors: %w", err)
|
return nil, fmt.Errorf("reading model.safetensors: %w", err)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
package marian
|
package marian
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"encoding/binary"
|
"encoding/binary"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
@@ -27,7 +28,7 @@ type tensorInfo struct {
|
|||||||
// little-endian header length, a JSON header, then the raw data) as
|
// little-endian header length, a JSON header, then the raw data) as
|
||||||
// float32. F32 tensors are read straight into their slice: the file is
|
// 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.
|
// 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)
|
f, err := os.Open(path)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
@@ -54,6 +55,9 @@ func readSafetensors(path string) (map[string]tensor, error) {
|
|||||||
if name == "__metadata__" {
|
if name == "__metadata__" {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
var info tensorInfo
|
var info tensorInfo
|
||||||
if err := json.Unmarshal(raw, &info); err != nil {
|
if err := json.Unmarshal(raw, &info); err != nil {
|
||||||
return nil, fmt.Errorf("tensor %s: %w", name, err)
|
return nil, fmt.Errorf("tensor %s: %w", name, err)
|
||||||
|
|||||||
+114
-34
@@ -3,6 +3,7 @@ package main
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
|
"fmt"
|
||||||
"net/http"
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
@@ -28,8 +29,14 @@ const (
|
|||||||
// modelIdle is how long the model (300 MB of memory) stays loaded
|
// modelIdle is how long the model (300 MB of memory) stays loaded
|
||||||
// after a translation, ready for the next one.
|
// after a translation, ready for the next one.
|
||||||
modelIdle = 5 * time.Minute
|
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{
|
var modelFiles = []marian.RemoteFile{
|
||||||
{Name: "config.json", Size: 1416, SHA256: "b3be13d046d9899d7aab8cf4ed624d9a79f5776038ba793f6b4d2ce3e02192f7"},
|
{Name: "config.json", Size: 1416, SHA256: "b3be13d046d9899d7aab8cf4ed624d9a79f5776038ba793f6b4d2ce3e02192f7"},
|
||||||
{Name: "source.spm", Size: 802397, SHA256: "78d0e717c77053f1c4b856d8661d9cb87c64f083a35418c087b9146300e4f585"},
|
{Name: "source.spm", Size: 802397, SHA256: "78d0e717c77053f1c4b856d8661d9cb87c64f083a35418c087b9146300e4f585"},
|
||||||
@@ -46,6 +53,7 @@ type translation struct {
|
|||||||
lexicon translate.Lexicon
|
lexicon translate.Lexicon
|
||||||
release *time.Timer
|
release *time.Timer
|
||||||
cancel context.CancelFunc
|
cancel context.CancelFunc
|
||||||
|
done chan struct{} // closed when the running task is over
|
||||||
}
|
}
|
||||||
|
|
||||||
// TranslationStatus says whether the model is installed, and how much a
|
// TranslationStatus says whether the model is installed, and how much a
|
||||||
@@ -72,15 +80,29 @@ func modelDirs() []string {
|
|||||||
return dirs
|
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 {
|
func installedModel() string {
|
||||||
for _, d := range modelDirs() {
|
for _, d := range modelDirs() {
|
||||||
if marian.Installed(d) {
|
if modelComplete(d) {
|
||||||
return d
|
return d
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return ""
|
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 {
|
func (a *App) TranslationStatus() TranslationStatus {
|
||||||
var size int64
|
var size int64
|
||||||
for _, f := range modelFiles {
|
for _, f := range modelFiles {
|
||||||
@@ -94,10 +116,11 @@ func (a *App) startTask() (context.Context, error) {
|
|||||||
a.tr.mu.Lock()
|
a.tr.mu.Lock()
|
||||||
defer a.tr.mu.Unlock()
|
defer a.tr.mu.Unlock()
|
||||||
if a.tr.cancel != nil {
|
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)
|
ctx, cancel := context.WithCancel(a.ctx)
|
||||||
a.tr.cancel = cancel
|
a.tr.cancel = cancel
|
||||||
|
a.tr.done = make(chan struct{})
|
||||||
return ctx, nil
|
return ctx, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -107,16 +130,22 @@ func (a *App) endTask() {
|
|||||||
if a.tr.cancel != nil {
|
if a.tr.cancel != nil {
|
||||||
a.tr.cancel()
|
a.tr.cancel()
|
||||||
a.tr.cancel = nil
|
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() {
|
func (a *App) CancelTranslation() {
|
||||||
a.tr.mu.Lock()
|
a.tr.mu.Lock()
|
||||||
defer a.tr.mu.Unlock()
|
cancel, done := a.tr.cancel, a.tr.done
|
||||||
if a.tr.cancel != nil {
|
a.tr.mu.Unlock()
|
||||||
a.tr.cancel()
|
if cancel == nil {
|
||||||
|
return
|
||||||
}
|
}
|
||||||
|
cancel()
|
||||||
|
<-done
|
||||||
}
|
}
|
||||||
|
|
||||||
// progressEmitter sends "translation-progress" events, at most every
|
// 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
|
// 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) {
|
func (a *App) DownloadTranslationModel() (bool, error) {
|
||||||
|
if installedModel() != "" {
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
ctx, err := a.startTask()
|
ctx, err := a.startTask()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return false, err
|
return false, err
|
||||||
@@ -148,11 +181,30 @@ func (a *App) DownloadTranslationModel() (bool, error) {
|
|||||||
f.URL = modelURL + f.Name
|
f.URL = modelURL + f.Name
|
||||||
files[i] = f
|
files[i] = f
|
||||||
}
|
}
|
||||||
err = marian.Download(ctx, &http.Client{}, modelDirs()[0], files, a.progressEmitter())
|
// No total timeout (300 MB can take a while) but abort a stalled
|
||||||
if errors.Is(err, context.Canceled) {
|
// transfer: the watchdog is reset by every progress report.
|
||||||
return false, nil
|
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
|
// 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
|
return nil, err
|
||||||
}
|
}
|
||||||
defer a.endTask()
|
defer a.endTask()
|
||||||
|
if installedModel() == "" {
|
||||||
|
return nil, errors.New("the translation model is not installed")
|
||||||
|
}
|
||||||
wails.EventsEmit(a.ctx, "translation-progress", 0, 0) // loading
|
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 {
|
if err != nil {
|
||||||
return nil, err
|
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
|
// translator returns the model and the dictionaries that help tell French
|
||||||
// from English, loading them the first time.
|
// from English, loading them the first time. The load happens without
|
||||||
func (a *App) translator() (*marian.Model, translate.Lexicon, error) {
|
// 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()
|
a.tr.mu.Lock()
|
||||||
defer a.tr.mu.Unlock()
|
|
||||||
if a.tr.release != nil {
|
if a.tr.release != nil {
|
||||||
a.tr.release.Stop()
|
a.tr.release.Stop()
|
||||||
}
|
}
|
||||||
if a.tr.model == nil {
|
model, lexicon := a.tr.model, a.tr.lexicon
|
||||||
dir := installedModel()
|
a.tr.mu.Unlock()
|
||||||
if dir == "" {
|
if model != nil {
|
||||||
return nil, nil, errors.New("the translation model is not installed")
|
return model, lexicon, nil
|
||||||
}
|
|
||||||
m, err := marian.Load(dir)
|
|
||||||
if err != nil {
|
|
||||||
return nil, nil, err
|
|
||||||
}
|
|
||||||
a.tr.model = m
|
|
||||||
a.tr.lexicon = loadLexicon()
|
|
||||||
}
|
}
|
||||||
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() {
|
func (a *App) releaseLater() {
|
||||||
a.tr.mu.Lock()
|
a.tr.mu.Lock()
|
||||||
defer a.tr.mu.Unlock()
|
defer a.tr.mu.Unlock()
|
||||||
a.tr.release = time.AfterFunc(modelIdle, func() {
|
a.tr.release = time.AfterFunc(modelIdle, a.releaseModel)
|
||||||
a.tr.mu.Lock()
|
}
|
||||||
defer a.tr.mu.Unlock()
|
|
||||||
if a.tr.cancel == nil {
|
// releaseModel frees the model when no task is running. When one is (a
|
||||||
a.tr.model, a.tr.lexicon = nil, nil
|
// 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
|
// lexicon tells French words from English ones with the installed
|
||||||
|
|||||||
Reference in new issue
Block a user