From 71212d503f0061a826bc72c63894315cd6926f7a Mon Sep 17 00:00:00 2001 From: Etienne Perot Date: Thu, 29 Feb 2024 13:08:05 -0800 Subject: [PATCH] `ollama` library: Add more metrics. This change adds new helper methods on the `Response` to get more metrics. It also adds support for forcing the model to be unloaded prior to prompting, which is useful to measure model loading time. PiperOrigin-RevId: 611571326 --- test/gpu/ollama/ollama.go | 188 ++++++++++++++++++++++++++++++++++---- test/gpu/textgen_test.go | 2 + 2 files changed, 172 insertions(+), 18 deletions(-) diff --git a/test/gpu/ollama/ollama.go b/test/gpu/ollama/ollama.go index 214d870e0..1cab3443f 100644 --- a/test/gpu/ollama/ollama.go +++ b/test/gpu/ollama/ollama.go @@ -23,6 +23,8 @@ import ( "errors" "fmt" "io" + "math" + "sort" "strconv" "strings" "time" @@ -50,6 +52,10 @@ type Ollama struct { // ModelNames is the list of available model names. ModelNames []string + // cheapModels is a list of models that are known to be cheap. + // A caller may set this to make forcefully unloading a model quicker. + cheapModels []*Model + // HasGPU is set depending on whether the LLM has GPU access. // ollama supports running both on CPU and GPU, and detects this // by spawning nvidia-smi. @@ -102,9 +108,9 @@ func New(ctx context.Context, server Server, logger testutil.Logger) (*Ollama, e // This may fail during the process of loading the first model, so we keep // iterating for a while. _, err = llm.Prompt(ctx, &Prompt{ - Model: &Model{Name: llm.ModelNames[0]}, - WarmFirst: false, - Query: curtQuery, + Model: &Model{Name: llm.ModelNames[0]}, + Preamble: DoNothing, + Query: curtQuery, }) if err != nil { return nil, fmt.Errorf("could not load first model %q: %w", llm.ModelNames[0], err) @@ -128,6 +134,15 @@ func New(ctx context.Context, server Server, logger testutil.Logger) (*Ollama, e return llm, nil } +// SetCheapModels can be used to inform this Ollama client as to the list of +// models it can use that are known to be cheap. +// This is useful when forcefully unloading models by swapping them with +// another one, to ensure that the one it is being swapped with is small. +// Therefore, there should be at least two models specified here. +func (llm *Ollama) SetCheapModels(cheapModels []*Model) { + llm.cheapModels = cheapModels +} + // dockerServer implements `Server`. It interfaces with an ollama server // running in a local Docker container. type dockerServer struct { @@ -410,6 +425,22 @@ func ZeroTemperatureModel(name string) *Model { } } +// Preamble dictates what to do before actually issuing a prompt. +type Preamble int + +const ( + // DoNothing means to do nothing when issuing a prompt. + DoNothing Preamble = iota + + // WarmUp issues a short query to the same model before + // issuing a prompt. This ensures the model is warm. + WarmUp Preamble = iota + + // Unload issues a short query with a different model before + // issuing a prompt. This ensures the model is cold. + Unload Preamble = iota +) + // Prompt is an ollama prompt. type Prompt struct { // Model is the model to query. @@ -423,10 +454,10 @@ type Prompt struct { // This is returned from `Response`. Context ConversationContext - // WarmFirst ensures the model is already loaded by issuing a small query - // beforehand. This is necessary for benchmarks to be accurate, but is - // unnecessary when just testing. - WarmFirst bool + // Preamble dictates what to do before actually issuing this prompt. + // This can be used to make sure the model is warm, or to make sure + // the model is cold. + Preamble Preamble } // CleanQuery removes common whitespace from query lines, and all @@ -565,6 +596,11 @@ func (r *Response) Done() bool { return r.data[len(r.data)-1].Done } +// NumTokens returns the number of tokens in the response. +func (r *Response) NumTokens() int { + return len(r.data) +} + // String returns the response text, if it is done. func (r *Response) String() string { if len(r.data) == 0 { @@ -596,6 +632,98 @@ func (r *Response) Text() string { return r.String() } +// TimeToFirstToken returns the time it took between the request starting +// and the first token being received by the client. +func (r *Response) TimeToFirstToken() time.Duration { + if !r.Done() { + return -1 + } + return r.metrics.FirstByteRead.Sub(r.metrics.RequestSent) +} + +// TimeToLastToken returns the time it took between the request starting +// and the last token being received by the client. +func (r *Response) TimeToLastToken() time.Duration { + if !r.Done() { + return -1 + } + return r.metrics.LastByteRead.Sub(r.metrics.RequestSent) +} + +// tokenIntervals returns the time between each token generation. +func (r *Response) tokenIntervals() []time.Duration { + if !r.Done() || len(r.data) < 2 { + return nil + } + intervals := make([]time.Duration, len(r.data)-1) + for i := 0; i < len(r.data)-1; i++ { + intervals[i] = r.data[i+1].CreatedAt.Sub(r.data[i].CreatedAt) + } + return intervals +} + +// OutputTokensPerSecond computes the average number of output tokens +// generated per second. +func (r *Response) OutputTokensPerSecond() float64 { + if !r.Done() || r.EvalDuration() == 0 { + return -1 + } + return float64(r.data[len(r.data)-1].EvalCount) / float64(r.EvalDuration().Seconds()) +} + +// TimePerOutputTokenAverage computes the average time to generate an output +// token. +func (r *Response) TimePerOutputTokenAverage() time.Duration { + if !r.Done() { + return -1 + } + intervals := r.tokenIntervals() + var sum time.Duration + for _, interval := range intervals { + sum += interval + } + return sum / time.Duration(len(intervals)) +} + +// TimePerOutputTokenQuantile computes a quantile of the time it takes to +// generate an output token. +func (r *Response) TimePerOutputTokenQuantile(quantile float64) time.Duration { + if quantile < 0.0 || quantile > 1.0 { + panic("quantile must be between 0.0 and 1.0 inclusively") + } + if !r.Done() || r.EvalDuration() == 0 { + return -1 + } + intervals := r.tokenIntervals() + sort.Slice(intervals, func(i, j int) bool { return intervals[i] < intervals[j] }) + return intervals[int(quantile*float64(len(intervals)-1))] +} + +// TokenGenerationStdDev returns the standard deviation of the time between +// token generations. +func (r *Response) TokenGenerationStdDev() time.Duration { + intervals := r.tokenIntervals() + if len(intervals) == 0 { + return -1 + } + if len(intervals) == 1 { + return 0 + } + + var sum time.Duration + for _, interval := range intervals { + sum += interval + } + mean := sum / time.Duration(len(intervals)) + variance := 0.0 + for _, interval := range intervals { + intervalMinusMean := float64((interval - mean).Nanoseconds()) + variance += intervalMinusMean * intervalMinusMean + } + variance = variance / float64(len(intervals)-1) + return time.Duration(math.Sqrt(variance)) * time.Nanosecond +} + // TotalDuration returns the total response generation time. func (r *Response) TotalDuration() time.Duration { if !r.Done() { @@ -628,14 +756,6 @@ func (r *Response) PromptEvalDuration() time.Duration { return time.Duration(r.data[len(r.data)-1].PromptEvalNanos) * time.Nanosecond } -// TokensPerSecond computes the number of tokens generated per second. -func (r *Response) TokensPerSecond() float64 { - if !r.Done() || r.EvalDuration() == 0 { - return 0 - } - return float64(r.data[len(r.data)-1].EvalCount) / float64(r.EvalDuration().Seconds()) -} - // ConversationContext represents a conversational context. // It is returned by a response and may be passed to a follow-up prompt. type ConversationContext []int @@ -658,9 +778,27 @@ func (llm *Ollama) withServerLogsErr(ctx context.Context, err error) error { return fmt.Errorf("%w (server logs are empty)", err) } +// getReplacementModel picks an available model other than `model`. +// It tries to find a one that is marked cheap if possible. +func (llm *Ollama) getReplacementModel(model *Model) (*Model, error) { + for _, cheapModel := range llm.cheapModels { + if cheapModel.Name != model.Name { + return cheapModel, nil + } + } + for _, otherModelName := range llm.ModelNames { + if otherModelName != model.Name { + return ZeroTemperatureModel(otherModelName), nil + } + } + return nil, fmt.Errorf("cannot find a replacement model to load instead of %q (available: %v; cheap: %v)", model.Name, llm.ModelNames, llm.cheapModels) +} + // Prompt returns the result of prompting the given `model` with `prompt`. func (llm *Ollama) Prompt(ctx context.Context, prompt *Prompt) (*Response, error) { - if prompt.WarmFirst { + switch prompt.Preamble { + case DoNothing: // Do nothing. + case WarmUp: warmCtx, warmCancel := context.WithTimeout(ctx, 3*time.Minute) _, err := jsonPost[PromptJSON, ResponseJSON](warmCtx, llm, "/api/generate", (&Prompt{ Model: prompt.Model, @@ -670,6 +808,20 @@ func (llm *Ollama) Prompt(ctx context.Context, prompt *Prompt) (*Response, error if err != nil { return nil, llm.withServerLogsErr(ctx, fmt.Errorf("warmup prompt for model %s failed: %w", prompt.Model.Name, err)) } + case Unload: + replacementModel, err := llm.getReplacementModel(prompt.Model) + if err != nil { + return nil, fmt.Errorf("cannot find a replacement model to load instead of %q to forcefully unload it: %w", prompt.Model.Name, err) + } + unloadCtx, unloadCancel := context.WithTimeout(ctx, 3*time.Minute) + _, err = jsonPost[PromptJSON, ResponseJSON](unloadCtx, llm, "/api/generate", (&Prompt{ + Model: replacementModel, + Query: curtQuery, + }).json()) + unloadCancel() + if err != nil { + return nil, llm.withServerLogsErr(ctx, fmt.Errorf("unload prompt for model %s failed: %w", replacementModel.Name, err)) + } } resp, err := jsonPost[PromptJSON, ResponseJSON](ctx, llm, "/api/generate", prompt.json()) if err != nil { @@ -692,10 +844,10 @@ func (llm *Ollama) PromptUntil(ctx context.Context, prompt *Prompt, iterate func if err != nil { return nil, fmt.Errorf("prompt request failed: %w", err) } - if prompt.WarmFirst && !warmed { + if prompt.Preamble == WarmUp && !warmed { // Future prompts do not need to specify the WarmFirst option. promptCopy := *prompt - promptCopy.WarmFirst = false + promptCopy.Preamble = DoNothing prompt = &promptCopy warmed = true } diff --git a/test/gpu/textgen_test.go b/test/gpu/textgen_test.go index 0ac7cf07a..bb24f7139 100644 --- a/test/gpu/textgen_test.go +++ b/test/gpu/textgen_test.go @@ -89,6 +89,7 @@ func TestLLM(t *testing.T) { Translate the following text from English to Chinese: "Hello World". `, + Preamble: ollama.DoNothing, } promptCtx, promptCancel := context.WithTimeout(ctx, 3*time.Minute) response, err := llm.PromptUntil(promptCtx, &prompt, func(prompt *ollama.Prompt, response *ollama.Response) (*ollama.Prompt, error) { @@ -146,6 +147,7 @@ func TestLLM(t *testing.T) { the code works and do not provide usage examples. Output a single block of Python code wrapped between %q marks. `, markerString, codeBlockDelim), + Preamble: ollama.DoNothing, } response, err := llm.PromptUntil(promptCtx, &prompt, func(prompt *ollama.Prompt, response *ollama.Response) (*ollama.Prompt, error) { defer prompt.Model.RaiseTemperature()