mirror of
https://github.com/ollama/ollama.git
synced 2025-05-10 18:06:33 +02:00
348b3e0983
This commit copies (without history) the bmizerany/ollama-go repository with the intention of integrating it into the ollama as a replacement for the pushing, and pulling of models, and management of the cache they are pushed and pulled from. New homes for these packages will be determined as they are integrated and we have a better understanding of proper package boundaries.
107 lines
2.4 KiB
Go
107 lines
2.4 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/ollama/ollama/server/internal/chunks"
|
|
"golang.org/x/sync/errgroup"
|
|
)
|
|
|
|
func BenchmarkDownload(b *testing.B) {
|
|
run := func(fileSize, chunkSize int64) {
|
|
name := fmt.Sprintf("size=%d/chunksize=%d", fileSize, chunkSize)
|
|
b.Run(name, func(b *testing.B) { benchmarkDownload(b, fileSize, chunkSize) })
|
|
}
|
|
|
|
run(100<<20, 8<<20)
|
|
run(100<<20, 16<<20)
|
|
run(100<<20, 32<<20)
|
|
run(100<<20, 64<<20)
|
|
run(100<<20, 128<<20) // 1 chunk
|
|
}
|
|
|
|
func run(ctx context.Context, c *http.Client, chunk chunks.Chunk) error {
|
|
const blobURL = "https://ollama.com/v2/x/x/blobs/sha256-4824460d29f2058aaf6e1118a63a7a197a09bed509f0e7d4e2efb1ee273b447d"
|
|
req, err := http.NewRequestWithContext(ctx, "GET", blobURL, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
req.Header.Set("Range", fmt.Sprintf("bytes=%s", chunk))
|
|
res, err := c.Do(req)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer res.Body.Close()
|
|
|
|
_, err = io.CopyN(io.Discard, res.Body, chunk.Size()) // will io.EOF on short read
|
|
return err
|
|
}
|
|
|
|
var sleepTime atomic.Int64
|
|
|
|
func benchmarkDownload(b *testing.B, fileSize, chunkSize int64) {
|
|
client := &http.Client{
|
|
Transport: func() http.RoundTripper {
|
|
tr := http.DefaultTransport.(*http.Transport).Clone()
|
|
tr.DisableKeepAlives = true
|
|
return tr
|
|
}(),
|
|
}
|
|
defer client.CloseIdleConnections()
|
|
|
|
// warm up the client
|
|
run(context.Background(), client, chunks.New(0, 1<<20))
|
|
|
|
b.SetBytes(fileSize)
|
|
b.ReportAllocs()
|
|
|
|
// Give our CDN a min to breathe between benchmarks.
|
|
time.Sleep(time.Duration(sleepTime.Swap(3)))
|
|
|
|
for b.Loop() {
|
|
g, ctx := errgroup.WithContext(b.Context())
|
|
g.SetLimit(runtime.GOMAXPROCS(0))
|
|
for chunk := range chunks.Of(fileSize, chunkSize) {
|
|
g.Go(func() error { return run(ctx, client, chunk) })
|
|
}
|
|
if err := g.Wait(); err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func BenchmarkWrite(b *testing.B) {
|
|
b.Run("chunksize=1MiB", func(b *testing.B) { benchmarkWrite(b, 1<<20) })
|
|
}
|
|
|
|
func benchmarkWrite(b *testing.B, chunkSize int) {
|
|
b.ReportAllocs()
|
|
|
|
dir := b.TempDir()
|
|
f, err := os.Create(filepath.Join(dir, "write-single"))
|
|
if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
defer f.Close()
|
|
|
|
data := make([]byte, chunkSize)
|
|
b.SetBytes(int64(chunkSize))
|
|
r := bytes.NewReader(data)
|
|
for b.Loop() {
|
|
r.Reset(data)
|
|
_, err := io.Copy(f, r)
|
|
if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
}
|
|
}
|