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.

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 {

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")
 	}
 }

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["<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 {
 		return nil, fmt.Errorf("reading model.safetensors: %w", err)
 	}

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)