baalajimaestro/ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

update download

Browse source
This commit is contained in:
Michael Yang 2023-10-09 10:14:49 -07:00
parent b599946b74
commit a2055a1e93
1 changed files with 10 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
server/download.go
View file

@ -30,6 +30,7 @@ type blobDownload struct {
Total int64
Completed atomic.Int64
done bool
Parts []*blobDownloadPart
@ -124,7 +125,7 @@ func (b *blobDownload) Run(ctx context.Context, requestURL *url.URL, opts *Regis
file.Truncate(b.Total)
g, ctx := errgroup.WithContext(ctx)
g, _ := errgroup.WithContext(ctx)
// TODO(mxyng): download concurrency should be configurable
g.SetLimit(64)
for i := range b.Parts {
@ -168,7 +169,12 @@ func (b *blobDownload) Run(ctx context.Context, requestURL *url.URL, opts *Regis
}
}
return os.Rename(file.Name(), b.Name)
if err := os.Rename(file.Name(), b.Name); err != nil {
return err
}
b.done = true
return nil
}
func (b *blobDownload) downloadChunk(ctx context.Context, requestURL *url.URL, w io.Writer, part *blobDownloadPart, opts *RegistryOptions) error {
@ -267,11 +273,8 @@ func (b *blobDownload) Wait(ctx context.Context, fn func(api.ProgressResponse))
Completed: b.Completed.Load(),
})
if b.Completed.Load() >= b.Total {
// wait for the file to get renamed
if _, err := os.Stat(b.Name); err == nil {
return nil
}
if b.done {
return nil
}
}
}