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:
Kharec committed 2026-10-13 18:17:53 +02:00
1 parent 8181c6e885
commit 839e9291d6
5 files changed
+313 -69

No files matched your search

+74 -12
View File
@@ -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")
}
}