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
@@ -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)
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user