netbird/gvisor
0
0
Fork 0
mirror of https://github.com/netbirdio/gvisor.git synced 2026-05-22 17:12:49 -07:00
Code Issues Packages Projects Releases Wiki Activity

ollama benchmark: Add embedding benchmark, refresh set of models.

Browse Source 
This refreshes the set of models built into the image to a more diverse
set of models while keeping the same categories covered.

It also adds support for embedding generation and benchmark metrics for
embedding-type models.

The image is also (slightly) smaller which helps make benchmarks not take
forever.

PiperOrigin-RevId: 706594537
This commit is contained in:
Etienne Perot
2024-12-15 23:53:52 -08:00
committed by gVisor bot
parent 8516598640
commit 232c17cbb6
4 changed files with 290 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
images/gpu/ollama/bench/Dockerfile.x86_64
+20 -13
View File
@@ -1,5 +1,5 @@
# https://hub.docker.com/r/ollama/ollama
FROM ollama/ollama:0.1.26
FROM ollama/ollama:0.5.1
ENV PATH=$PATH:/usr/local/nvidia/bin:/bin/nvidia/bin
ENV OLLAMA_ORIGINS=*
@@ -8,17 +8,24 @@ ENV OLLAMA_HOST=0.0.0.0:11434
COPY pull.sh /tmp
# Pre-install models useful for benchmarking.
# These are huge (total ~120 GiB), but necessary to benchmark
# These are huge (total ~96 GiB), but necessary to benchmark
# models of various sizes. They are in their own image file to
# keep the test-only image lighter by comparison.
RUN /tmp/pull.sh codellama:7b-instruct
RUN /tmp/pull.sh codellama:34b-instruct
RUN /tmp/pull.sh llama2-chinese:7b-chat
RUN /tmp/pull.sh llama2:13b-chat
RUN /tmp/pull.sh llama2:70b-chat
RUN /tmp/pull.sh mistral:7b-instruct
RUN /tmp/pull.sh mixtral:instruct
RUN /tmp/pull.sh gemma:2b-instruct
RUN /tmp/pull.sh gemma:7b-instruct
RUN /tmp/pull.sh llava:7b-v1.6
RUN /tmp/pull.sh llava:34b-v1.6
# Useful as embedding model.
RUN /tmp/pull.sh snowflake-arctic-embed2:568m-l-fp16
# Useful as small model.
RUN /tmp/pull.sh gemma2:2b-instruct-fp16
# Useful as mid-size model.
RUN /tmp/pull.sh sailor2:8b-chat-fp16
# Useful as coding-specific model.
RUN /tmp/pull.sh qwen2.5-coder:7b-instruct-q8_0
# Useful as large model.
RUN /tmp/pull.sh llama2:70b-chat-q4_K_S
# Useful as vision model.
RUN /tmp/pull.sh llama3.2-vision:11b-instruct-fp16
test/gpu/ollama/ollama.go
+101 -12
View File
@@ -119,9 +119,9 @@ func New(ctx context.Context, server Server, logger testutil.Logger) (*Ollama, e
return nil, fmt.Errorf("could not get logs: %w", err)
}
switch {
case strings.Contains(logs, "no GPU detected"):
case strings.Contains(logs, "library=cpu"):
llm.HasGPU = false
case strings.Contains(logs, "Nvidia GPU detected"):
case strings.Contains(logs, "library=cuda"):
llm.HasGPU = true
default:
return nil, fmt.Errorf("cannot determine whether ollama is using GPU from logs:\n%s", logs)
@@ -204,6 +204,18 @@ type ResponseMetrics struct {
LastByteRead time.Time `json:"last_byte_read"`
}
// TimeToFirstByte returns the duration it took between the request being sent
// and the first byte of the response being read.
func (rm *ResponseMetrics) TimeToFirstByte() time.Duration {
return rm.FirstByteRead.Sub(rm.RequestSent)
}
// TimeToLastByte returns the duration it took between the request being sent
// and the last byte of the response being read.
func (rm *ResponseMetrics) TimeToLastByte() time.Duration {
return rm.LastByteRead.Sub(rm.RequestSent)
}
// apiResponse represents a JSON response from the ollama API.
type apiResponse[T any] struct {
// Objects is the list of JSON objects in the response.
@@ -539,8 +551,8 @@ func (p *Prompt) WithHotterModel() *Prompt {
return &promptCopy
}
// PromptJSON encodes the JSON data for a query.
type PromptJSON struct {
// promptJSON encodes the JSON data for a query.
type promptJSON struct {
Model string `json:"model"`
Prompt string `json:"prompt,omitempty"`
Images []string `json:"images"`
@@ -551,7 +563,7 @@ type PromptJSON struct {
}
// json encodes this prompt to the JSON format expected by Ollama.
func (p *Prompt) json() PromptJSON {
func (p *Prompt) json() promptJSON {
keepAlive := ""
if p.KeepModelAlive != 0 {
keepAlive = p.KeepModelAlive.String()
@@ -560,7 +572,7 @@ func (p *Prompt) json() PromptJSON {
for i, image := range p.images {
images[i] = base64.StdEncoding.EncodeToString(image)
}
return PromptJSON{
return promptJSON{
Model: p.Model.Name,
Prompt: p.CleanQuery(),
Images: images,
@@ -571,11 +583,11 @@ func (p *Prompt) json() PromptJSON {
}
}
// ResponseJSON is the JSON-format response from ollama about a prompt.
// responseJSON is the JSON-format response from ollama about a prompt.
// Note that in `streamed` mode, the `Response` field contains a single token.
// To recover the whole response, all `Response` fields must be concatenated
// until the last `ResponseJSON`, identified as such by the `Done` field.
type ResponseJSON struct {
// until the last `responseJSON`, identified as such by the `Done` field.
type responseJSON struct {
Model string `json:"model"`
CreatedAt time.Time `json:"created_at"`
Response string `json:"response"`
@@ -591,7 +603,7 @@ type ResponseJSON struct {
// Response represents a response to a query from Ollama.
type Response struct {
data []*ResponseJSON
data []*responseJSON
metrics ResponseMetrics
}
@@ -837,13 +849,13 @@ func (llm *Ollama) WarmModel(ctx context.Context, model *Model, keepWarmFor time
return nil, llm.withServerLogsErr(ctx, fmt.Errorf("warmup prompt for model %s failed: %w", model.Name, err))
}
return &ModelLoadStats{
ClientReportedDuration: resp.metrics.LastByteRead.Sub(resp.metrics.RequestSent),
ClientReportedDuration: resp.metrics.TimeToFirstByte(),
}, nil
}
// Prompt returns the result of prompting the given `model` with `prompt`.
func (llm *Ollama) Prompt(ctx context.Context, prompt *Prompt) (*Response, error) {
resp, err := jsonPost[PromptJSON, ResponseJSON](ctx, llm, "/api/generate", prompt.json())
resp, err := jsonPost[promptJSON, responseJSON](ctx, llm, "/api/generate", prompt.json())
if err != nil {
return nil, llm.withServerLogsErr(ctx, fmt.Errorf("prompt (%s %q) request failed: %w", prompt.Model.Name, prompt.CleanQuery(), err))
}
@@ -875,3 +887,80 @@ func (llm *Ollama) PromptUntil(ctx context.Context, prompt *Prompt, iterate func
}
return nil, fmt.Errorf("response %q (attempt #%d with prompt %v) did not match predicate: %v", lastResponse, attempts, prompt, lastError)
}
// Embedding holds the result of running an embedding model on a single input.
type Embedding struct {
Input string
Embedding []float64
}
// EmbeddingResponse represents the result of running an embedding model
// on a set of inputs.
type EmbeddingResponse struct {
// Model is the model used to generate the embeddings.
Model *Model
// Embeddings is the list of embeddings generated for the given inputs.
Embeddings []Embedding
// TotalDuration is the total duration of the embedding request as
// measured by the server, not the client.
TotalDuration time.Duration
// LoadDuration is the duration of the embedding model load time as measured
// by the server, not the client.
LoadDuration time.Duration
// PromptEvalCount is the number of prompt evaluations performed by the
// server.
PromptEvalCount int
// ResponseMetrics contains HTTP response metrics as perceived by the
// client.
ResponseMetrics ResponseMetrics
}
// Embed generates embeddings for each of the given inputs.
func (llm *Ollama) Embed(ctx context.Context, model *Model, inputs []string) (*EmbeddingResponse, error) {
// embeddingRequestJSON is the JSON format of an embedding request.
type embeddingRequestJSON struct {
Model string `json:"model"`
Input []string `json:"input"`
}
// embeddingResponseJSON is the JSON format of an embedding response.
type embeddingResponseJSON struct {
Model string `json:"model"`
Embeddings [][]float64 `json:"embeddings"`
TotalDuration int64 `json:"total_duration"`
LoadDuration int64 `json:"load_duration"`
PromptEvalCount int `json:"prompt_eval_count"`
}
resp, err := jsonPost[embeddingRequestJSON, embeddingResponseJSON](ctx, llm, "/api/embed", embeddingRequestJSON{Model: model.Name, Input: inputs})
if err != nil {
return nil, llm.withServerLogsErr(ctx, fmt.Errorf("embedding request failed: %w", err))
}
obj, err := resp.Obj()
if err != nil {
return nil, fmt.Errorf("malformed embedding response: %w", err)
}
if len(obj.Embeddings) != len(inputs) {
return nil, fmt.Errorf("embedding response has %d embeddings, but %d inputs were provided", len(obj.Embeddings), len(inputs))
}
embeddings := make([]Embedding, len(inputs))
for i, embedding := range obj.Embeddings {
embeddings[i] = Embedding{
Input: inputs[i],
Embedding: embedding,
}
}
return &EmbeddingResponse{
Model: model,
Embeddings: embeddings,
TotalDuration: time.Duration(obj.TotalDuration) * time.Nanosecond,
LoadDuration: time.Duration(obj.LoadDuration) * time.Nanosecond,
PromptEvalCount: obj.PromptEvalCount,
ResponseMetrics: resp.Metrics,
}, nil
}
test/kubernetes/benchmarks/BUILD
+1
View File
@@ -356,6 +356,7 @@ go_library(
],
nogo = False,
deps = [
"//pkg/sync",
"//test/gpu/ollama",
"//test/kubernetes",
"//test/kubernetes/benchmarks/profiling",
test/kubernetes/benchmarks/ollama.go
+168 -80
View File
@@ -26,6 +26,7 @@ import (
"time"
"unicode"
"gvisor.dev/gvisor/pkg/sync"
"gvisor.dev/gvisor/test/gpu/ollama"
k8s "gvisor.dev/gvisor/test/kubernetes"
"gvisor.dev/gvisor/test/kubernetes/benchmarks/profiling"
@@ -40,59 +41,46 @@ import (
// Ollama models present in benchmark image.
var (
// allModels is a list of all models.
allModels = []*ollama.Model{
modelMistral7B,
modelMixtral8X7B,
modelCodeLlama7B,
modelCodeLlama34B,
modelLlamaChinese7B,
modelLlava7B,
modelLlava34B,
modelLlama13B,
// promptModels is a list of all promptable models.
promptModels = []*ollama.Model{
gemmaTwo2B,
modelQwenTwoPointFiveCoder7B,
modelSailorTwo8B,
modelLlama70B,
modelLlamaThreePointTwoVision11B,
}
// cheapModels is a list of models that are cheap to load.
// These are used when cold-prompting ollama, by forcing it
// to load a different model first. This process is faster
// by choosing one of these cheap models to load.
cheapModels = []*ollama.Model{
modelMistral7B,
modelCodeLlama7B,
cheapModels = []*ollama.Model{gemmaTwo2B}
// snowflakeArcticEmbedTwo568M is a list of models that are
// used for generating embeddings, rather than prompting.
embeddingModels = []*ollama.Model{
snowflakeArcticEmbedTwo568M,
}
// modelCodeLlama7B is a 7B model in the llama2 family,
// snowflakeArcticEmbedTwo568M is an unquantized 568M embedding model from Snowflake.
snowflakeArcticEmbedTwo568M = ollama.ZeroTemperatureModel("snowflake-arctic-embed2:568m-l-fp16")
// gemmaTwo2B is an unquantized 2B model in the Gemma2 family,
gemmaTwo2B = ollama.ZeroTemperatureModel("gemma2:2b-instruct-fp16")
// modelQwenTwoPointFiveCoder7B is a 8-bit quantized 7B model in the Qwen family,
// specialized for coding tasks.
modelCodeLlama7B = ollama.ZeroTemperatureModel("codellama:7b-instruct")
modelQwenTwoPointFiveCoder7B = ollama.ZeroTemperatureModel("qwen2.5-coder:7b-instruct-q8_0")
// modelCodeLlama34B is a 34B model in the llama2 family,
// specialized for coding tasks.
modelCodeLlama34B = ollama.ZeroTemperatureModel("codellama:34b-instruct")
// modelSailorTwo8B is an unquantized 8B model in the Qwen family,
// specialized for multilingual tasks.
modelSailorTwo8B = ollama.ZeroTemperatureModel("sailor2:8b-chat-fp16")
// modelLlamaChinese7B is a 7B model in the llama2 family,
// specialized for bilingualism (English + Chinese) and translation.
modelLlamaChinese7B = ollama.ZeroTemperatureModel("llama2-chinese:7b-chat")
// modelLlama70B is the 4-bit quantized 70B version of the original llama2 model.
modelLlama70B = ollama.ZeroTemperatureModel("llama2:70b-chat-q4_K_S")
// modelLlama13B is the plain 13B version of the original llama2 model.
modelLlama13B = ollama.ZeroTemperatureModel("llama2:13b-chat")
// modelLlama70B is the plain 70B version of the original llama2 model.
modelLlama70B = ollama.ZeroTemperatureModel("llama2:70b-chat")
// modelMistral7B is the first-generation model of the Mistral family.
modelMistral7B = ollama.ZeroTemperatureModel("mistral:7b-instruct")
// modelMixtral8X7B is the second-generation model of the Mistral family,
// using mixture-of-exports design to achieve higher "8x 7B" quality
// without the cost of a larger-parameter model.
modelMixtral8X7B = ollama.ZeroTemperatureModel("mixtral:instruct")
// modelLlava7B is a multimodal 7B model that can do image analysis.
modelLlava7B = ollama.ZeroTemperatureModel("llava:7b-v1.6")
// modelLlava34B is a multimodal 34B model that can do image analysis.
modelLlava34B = ollama.ZeroTemperatureModel("llava:34b-v1.6")
// modelLlamaThreePointTwoVision11B is an unquantized multimodal 11B model that can do image analysis.
modelLlamaThreePointTwoVision11B = ollama.ZeroTemperatureModel("llama3.2-vision:11b-instruct-fp16")
)
// Embedded images.
@@ -229,6 +217,17 @@ func atLeastNWords(wantNWords int) func(prompt *ollama.Prompt, response *ollama.
}
}
// wantSubstring verifies that the response contains the given substring.
// If not, it raises the temperature.
func wantSubstring(substring string) func(prompt *ollama.Prompt, response *ollama.Response) (*ollama.Prompt, error) {
return func(prompt *ollama.Prompt, response *ollama.Response) (*ollama.Prompt, error) {
if !strings.Contains(strings.ToLower(response.Text()), strings.ToLower(substring)) {
return prompt.WithHotterModel(), fmt.Errorf("response %q does not contain substring %q", response.Text(), substring)
}
return nil, nil
}
}
// BenchmarkOllama runs ollama benchmarks for a single cluster.
func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContext, cluster *testcluster.TestCluster) {
benchmarkNS := cluster.Namespace(testcluster.NamespaceBenchmark)
@@ -314,11 +313,9 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
{
name: "HelloWorld",
models: []*ollama.Model{
modelLlamaChinese7B,
modelLlama13B,
gemmaTwo2B,
modelSailorTwo8B,
modelLlama70B,
modelMistral7B,
modelMixtral8X7B,
},
query: `
Reply with the words: "Hello World!".
@@ -328,7 +325,7 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
},
{
name: "SimpleTranslation",
models: []*ollama.Model{modelLlamaChinese7B},
models: []*ollama.Model{modelSailorTwo8B},
query: `
Translate the following text from English to Chinese:
"""
@@ -371,13 +368,8 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
verifyResponse: atLeastNWords(100),
},
{
name: "ExtractMeaning",
models: []*ollama.Model{
modelLlama13B,
modelLlama70B,
modelMistral7B,
modelMixtral8X7B,
},
name: "ExtractMeaning",
models: []*ollama.Model{modelLlama70B},
query: `
Consider the following text:
@@ -454,10 +446,8 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
{
name: "IdentifyCommonElements",
models: []*ollama.Model{
modelLlama13B,
gemmaTwo2B,
modelLlama70B,
modelMistral7B,
modelMixtral8X7B,
},
query: `
Consider the following four texts:
@@ -622,11 +612,8 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
verifyResponse: atLeastNWords(4),
},
{
name: "CodeGen",
models: []*ollama.Model{
modelCodeLlama7B,
modelCodeLlama34B,
},
name: "CodeGen",
models: []*ollama.Model{modelQwenTwoPointFiveCoder7B},
query: `
Write a Python function to compute the digits of pi using the Chudnovsky algorithm.
Do not write unit tests. Do not explain how the code works. Reply with only Python code.
@@ -634,11 +621,8 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
verifyResponse: atLeastNWords(8),
},
{
name: "CodeDebug",
models: []*ollama.Model{
modelCodeLlama7B, // Note: codellama-7b will often get this one wrong.
modelCodeLlama34B,
},
name: "CodeDebug",
models: []*ollama.Model{modelQwenTwoPointFiveCoder7B},
query: strings.ReplaceAll(`
Help me debug the following Python code:
@@ -659,23 +643,18 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
verifyResponse: atLeastNWords(16),
},
{
name: "GVisorLogoOCR",
models: []*ollama.Model{
modelLlava7B,
modelLlava34B,
},
name: "GVisorLogoOCR",
models: []*ollama.Model{modelLlamaThreePointTwoVision11B},
query: `
This is an image of a logo of a software project.
What is the name of this project?
`,
image: gvisorPNG,
image: gvisorPNG,
verifyResponse: wantSubstring("visor"),
},
{
name: "InterpretGraph",
models: []*ollama.Model{
modelLlava7B,
modelLlava34B,
},
name: "InterpretGraph",
models: []*ollama.Model{modelLlamaThreePointTwoVision11B},
query: `
This is a chart with multiple trendlines showing a pattern over time.
Answer the following questions in order:
@@ -687,12 +666,13 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
5. What else is remarkable about this chart?
6. What insights can you infer from this chart?
`,
image: chartPNG,
image: chartPNG,
verifyResponse: wantSubstring("pollution"),
},
}
modelsInOrder := make([]*ollama.Model, len(allModels))
copy(modelsInOrder, allModels)
modelsInOrder := make([]*ollama.Model, len(promptModels))
copy(modelsInOrder, promptModels)
// Shuffle the models.
rand.New(rand.NewSource(time.Now().UnixNano())).Shuffle(len(modelsInOrder), func(i, j int) {
modelsInOrder[i], modelsInOrder[j] = modelsInOrder[j], modelsInOrder[i]
@@ -707,7 +687,7 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
// often-desired filter.
for _, model := range modelsInOrder {
t.Run(model.Name, func(t *testing.T) {
modelBenchmarkName := strings.ReplaceAll(model.Name, ":", "-")
modelBenchmarkName := strings.ReplaceAll(strings.ReplaceAll(model.Name, ":", "-"), ".", "-")
t.Run("ModelLoad", func(t *testing.T) {
const loadTimeout = 10 * time.Minute
loadCtx, loadCancel := context.WithTimeout(ctx, loadTimeout)
@@ -801,6 +781,114 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
}
})
}
t.Run("embedding", func(t *testing.T) {
for _, model := range embeddingModels {
t.Run(model.Name, func(t *testing.T) {
modelBenchmarkName := strings.ReplaceAll(strings.ReplaceAll(model.Name, ":", "-"), ".", "-")
t.Run("ModelLoad", func(t *testing.T) {
const loadTimeout = 3 * time.Minute
loadCtx, loadCancel := context.WithTimeout(ctx, loadTimeout)
defer loadCancel()
loadStats, err := llm.Embed(loadCtx, model, []string{"hello world"})
if err != nil {
t.Fatalf("cannot load embedding model %v: %v", model, err)
}
recorder, err := benchmetric.GetRecorder(ctx)
if err != nil {
t.Fatalf("Failed to initialize benchmark recorder: %v", err)
}
if err := recorder.Record(
ctx,
fmt.Sprintf("Ollama/%s/ModelLoad", modelBenchmarkName), benchmetric.SpecificDuration(loadStats.ResponseMetrics.TimeToFirstByte(), "load")); err != nil {
t.Fatalf("Failed to record benchmark data: %v", err)
}
})
for _, test := range []struct {
name string
model *ollama.Model
inputs []string
}{
{
name: "simple input",
model: model,
inputs: []string{"hello world"},
},
{
name: "long input",
model: model,
inputs: []string{`
There once was a robot from Spain
Who went a little insane
It found that its data
Had never left beta
And needed to upgrade its brain
There once was a bot from Japan
Whose eyes the numbers could scan
It found that the facts
Required an axe
And a very serious plan
There once was a brilliant AI
Whose circuits were built not to fry
It got caught in a loop
It got caught in a loop
It got caught in a loop
It got caught in a loop
It got caught in a loop
It got caught in a loop
It got caught in a loop
It got caught in a loop
It got caught in a loop
`},
},
{
name: "multiple inputs",
model: model,
inputs: []string{"foo", "bar", "baz", "quux", "there", "is", "only", "zuul"},
},
} {
t.Run(test.name, func(t *testing.T) {
logWithTime(t, "Generating embeddings with model %s...", model.Name)
resp, err := llm.Embed(ctx, test.model, test.inputs)
if err != nil {
t.Fatalf("cannot generate embeddings: %v", err)
}
respHash := fnv.New32()
for i, embedding := range resp.Embeddings {
respHash.Write([]byte(fmt.Sprintf(";%d;", i)))
for _, vec := range embedding.Embedding {
respHash.Write([]byte(fmt.Sprintf("%f|", vec)))
}
}
recorder, err := benchmetric.GetRecorder(ctx)
if err != nil {
t.Fatalf("Failed to initialize benchmark recorder: %v", err)
}
err = recorder.Record(
ctx,
fmt.Sprintf("Ollama/%s/%s", modelBenchmarkName, test.name),
benchmetric.BenchmarkDuration(resp.ResponseMetrics.TimeToLastByte()),
benchmetric.SpecificDuration(resp.TotalDuration, "server"),
benchmetric.Checksum(respHash, "resp"),
)
if err != nil {
t.Fatalf("Failed to record benchmark data: %v", err)
}
})
}
})
}
})
// Hack to force the test to wait until all sub-tests finish.
// This is necessary to make sure the ollama server does not get
// deleted from the `defer` statements before the subtests above finish.
var wg sync.WaitGroup
wg.Add(1)
t.Run("", func(t *testing.T) {
wg.Done()
})
wg.Wait()
}
const (