From c324b9b3faac540e07b31a6e06aaab949d517d55 Mon Sep 17 00:00:00 2001 From: Etienne Perot Date: Wed, 28 Feb 2024 17:56:15 -0800 Subject: [PATCH] Add custom HTTP client for interacting with `ollama`. This allows getting more metrics from the output stream. PiperOrigin-RevId: 611289642 --- Makefile | 2 +- images/gpu/ollama/client/BUILD | 11 ++ images/gpu/ollama/client/Dockerfile | 8 ++ images/gpu/ollama/client/client.go | 152 +++++++++++++++++++++++++++ test/gpu/ollama/ollama.go | 154 +++++++++++++++++++++------- 5 files changed, 288 insertions(+), 39 deletions(-) create mode 100644 images/gpu/ollama/client/BUILD create mode 100644 images/gpu/ollama/client/Dockerfile create mode 100644 images/gpu/ollama/client/client.go diff --git a/Makefile b/Makefile index e06af2564..c12919476 100644 --- a/Makefile +++ b/Makefile @@ -298,7 +298,7 @@ cos-gpu-smoke-tests: gpu-smoke-images $(RUNTIME_BIN) # This is a superset of those needed for smoke tests. # It includes non-GPU images that are used as part of GPU tests, # e.g. busybox and python. -gpu-images: gpu-smoke-images load-gpu_pytorch load-gpu_ollama load-basic_busybox load-basic_python +gpu-images: gpu-smoke-images load-gpu_pytorch load-gpu_ollama load-gpu_ollama_client load-basic_busybox load-basic_python .PHONY: gpu-images gpu-all-tests: gpu-images gpu-smoke-tests $(RUNTIME_BIN) diff --git a/images/gpu/ollama/client/BUILD b/images/gpu/ollama/client/BUILD new file mode 100644 index 000000000..c322daa43 --- /dev/null +++ b/images/gpu/ollama/client/BUILD @@ -0,0 +1,11 @@ +load("//tools:defs.bzl", "go_binary") + +package( + default_applicable_licenses = ["//:license"], + licenses = ["notice"], +) + +go_binary( + name = "client", + srcs = ["client.go"], +) diff --git a/images/gpu/ollama/client/Dockerfile b/images/gpu/ollama/client/Dockerfile new file mode 100644 index 000000000..11356724b --- /dev/null +++ b/images/gpu/ollama/client/Dockerfile @@ -0,0 +1,8 @@ +FROM golang:1.22 AS builder + +COPY client.go /client.go +RUN CGO_ENABLED=0 go build -o /httpclient /client.go + +FROM alpine:latest +COPY --from=builder /httpclient /usr/bin/ +CMD ["/usr/bin/httpclient"] diff --git a/images/gpu/ollama/client/client.go b/images/gpu/ollama/client/client.go new file mode 100644 index 000000000..2faa7a168 --- /dev/null +++ b/images/gpu/ollama/client/client.go @@ -0,0 +1,152 @@ +// Copyright 2024 The gVisor Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// A simple `curl`-like HTTP client that prints metrics after the request. +// All of its output is structured to be unambiguous even if stdout/stderr +// is combined, as is the case for Kubernetes logs. +// Useful for communicating with ollama. +package main + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "flag" + "fmt" + "io" + "net/http" + "os" + "sort" + "time" +) + +// Flags. +var ( + url = flag.String("url", "", "HTTP request URL.") + method = flag.String("method", "GET", "HTTP request method (GET or POST).") + postDataBase64 = flag.String("post_base64", "", "HTTP request POST data in base64 format; ignored for GET requests.") + timeout = flag.Duration("timeout", 0, "HTTP request timeout; 0 for no timeout.") +) + +// bufSize is the size of buffers used for HTTP requests and responses. +const bufSize = 1024 * 1024 // 1MiB + +// fatalf crashes the program with a given error message. +func fatalf(format string, values ...any) { + fmt.Fprintf(os.Stderr, "FATAL: "+format+"\n", values...) + os.Exit(1) +} + +// Metrics contains the request metrics to export to JSON. +// This is parsed by the ollama library at `test/gpu/ollama/ollama.go`. +type Metrics struct { + // ProgramStarted is the time when the program started. + ProgramStarted time.Time `json:"program_started"` + // RequestSent is the time when the HTTP request was sent. + RequestSent time.Time `json:"request_sent"` + // ResponseReceived is the time when the HTTP response headers were received. + ResponseReceived time.Time `json:"response_received"` + // FirstByteRead is the time when the first HTTP response body byte was read. + FirstByteRead time.Time `json:"first_byte_read"` + // LastByteRead is the time when the last HTTP response body byte was read. + LastByteRead time.Time `json:"last_byte_read"` +} + +func main() { + var metrics Metrics + metrics.ProgramStarted = time.Now() + flag.Parse() + if *url == "" { + fatalf("--url is required") + } + client := http.Client{ + Transport: &http.Transport{ + MaxIdleConns: 1, + IdleConnTimeout: *timeout, + ReadBufferSize: bufSize, + WriteBufferSize: bufSize, + }, + Timeout: *timeout, + } + var request *http.Request + var err error + switch *method { + case "GET": + request, err = http.NewRequest("GET", *url, nil) + case "POST": + postData, postDataErr := base64.StdEncoding.DecodeString(*postDataBase64) + if postDataErr != nil { + fatalf("cannot decode POST data: %v", postDataErr) + } + request, err = http.NewRequest("POST", *url, bytes.NewBuffer(postData)) + default: + err = fmt.Errorf("unknown method %q", *method) + } + if err != nil { + fatalf("cannot create request: %v", err) + } + readBuf := make([]byte, bufSize) + orderedReqHeaders := make([]string, 0, len(request.Header)) + for k := range request.Header { + orderedReqHeaders = append(orderedReqHeaders, k) + } + sort.Strings(orderedReqHeaders) + for _, k := range orderedReqHeaders { + for _, v := range request.Header[k] { + fmt.Fprintf(os.Stderr, "REQHEADER: %s: %s\n", k, v) + } + } + metrics.RequestSent = time.Now() + resp, err := client.Do(request) + metrics.ResponseReceived = time.Now() + if err != nil { + fatalf("cannot make request: %v", err) + } + gotFirstByte := false + for { + n, err := resp.Body.Read(readBuf) + if n > 0 { + if !gotFirstByte { + metrics.FirstByteRead = time.Now() + gotFirstByte = true + } + fmt.Printf("BODY: %q\n", string(readBuf[:n])) + } + if err == io.EOF { + metrics.LastByteRead = time.Now() + break + } + if err != nil { + fatalf("cannot read response body: %v", err) + } + } + if err := resp.Body.Close(); err != nil { + fatalf("cannot close response body: %v", err) + } + orderedRespHeaders := make([]string, 0, len(resp.Header)) + for k := range resp.Header { + orderedRespHeaders = append(orderedRespHeaders, k) + } + sort.Strings(orderedRespHeaders) + for _, k := range orderedRespHeaders { + for _, v := range resp.Header[k] { + fmt.Fprintf(os.Stderr, "RESPHEADER: %s: %s\n", k, v) + } + } + metricsBytes, err := json.Marshal(&metrics) + if err != nil { + fatalf("cannot marshal metrics: %v", err) + } + fmt.Fprintf(os.Stderr, "STATS: %s\n", string(metricsBytes)) +} diff --git a/test/gpu/ollama/ollama.go b/test/gpu/ollama/ollama.go index 0f186982b..2f724164b 100644 --- a/test/gpu/ollama/ollama.go +++ b/test/gpu/ollama/ollama.go @@ -16,10 +16,13 @@ package ollama import ( + "bytes" "context" + "encoding/base64" "encoding/json" "errors" "fmt" + "strconv" "strings" "time" @@ -54,9 +57,13 @@ type Ollama struct { // Server performs requests against an ollama server. type Server interface { - // HTTPRequest performs an HTTP request against the ollama server. - // The request is a GET request if postData == nil, otherwise POST. - HTTPRequest(ctx context.Context, endpoint string, postData []byte) ([]byte, error) + // InstrumentedRequest performs an instrumented HTTP request against the + // ollama server, using the `gpu/ollama_client` ollama image. + // `argvFn` takes in a `protocol://host:port` string and returns a + // command-line to use for making an instrumented HTTP request against the + // ollama server. + // InstrumentedRequest should return the logs from the request container. + InstrumentedRequest(ctx context.Context, argvFn func(hostPort string) []string) ([]byte, error) // Logs retrieves logs from the server. Logs(ctx context.Context) (string, error) @@ -144,16 +151,13 @@ func NewDocker(ctx context.Context, cont *dockerutil.Container, logger testutil. return New(ctx, ds, logger) } -// HTTPRequest implements `Server.HTTPRequest`. -func (ds *dockerServer) HTTPRequest(ctx context.Context, endpoint string, data []byte) ([]byte, error) { - cmd := []string{"wget", "-qO-"} - if data != nil { - cmd = append(cmd, "--post-data", string(data)) - } - cmd = append(cmd, fmt.Sprintf("http://llm:%d%s", Port, endpoint)) +// InstrumentedRequest implements `Server.InstrumentedRequest`. +func (ds *dockerServer) InstrumentedRequest(ctx context.Context, argvFn func(hostPort string) []string) ([]byte, error) { + const ollamaHost = "llm" + cmd := argvFn(fmt.Sprintf("http://%s:%d", ollamaHost, Port)) out, err := dockerutil.MakeContainer(ctx, ds.logger).Run(ctx, dockerutil.RunOpts{ - Image: "basic/busybox", - Links: []string{ds.container.MakeLink("llm")}, + Image: "gpu/ollama/client", + Links: []string{ds.container.MakeLink(ollamaHost)}, }, cmd...) if err != nil { if out != "" { @@ -169,42 +173,115 @@ func (ds *dockerServer) Logs(ctx context.Context) (string, error) { return ds.container.Logs(ctx) } -// request makes an HTTP request to the ollama API. -func (llm *Ollama) request(ctx context.Context, endpoint string, data []byte) ([]byte, error) { +// ResponseMetrics are HTTP request metrics from an ollama API query. +// These is the same JSON struct as defined in +// `images/gpu/ollama/client/client.go`. +type ResponseMetrics struct { + // ProgramStarted is the time when the program started. + ProgramStarted time.Time `json:"program_started"` + // RequestSent is the time when the HTTP request was sent. + RequestSent time.Time `json:"request_sent"` + // ResponseReceived is the time when the HTTP response headers were received. + ResponseReceived time.Time `json:"response_received"` + // FirstByteRead is the time when the first HTTP response body byte was read. + FirstByteRead time.Time `json:"first_byte_read"` + // LastByteRead is the time when the last HTTP response body byte was read. + LastByteRead time.Time `json:"last_byte_read"` +} + +type apiResponse[T any] struct { + Response T + Metrics ResponseMetrics +} + +func makeAPIResponse[T any](rawResponse []byte) (*apiResponse[T], error) { + var respBytes bytes.Buffer + var resp apiResponse[T] + for _, line := range strings.Split(string(rawResponse), "\n") { + line = strings.TrimSpace(line) + if line == "" { + continue + } + colonIndex := strings.Index(line, ":") + if colonIndex == -1 { + return nil, fmt.Errorf("malformed line: %q", line) + } + data := strings.TrimSpace(line[colonIndex+1:]) + switch line[:colonIndex] { + case "FATAL": + return nil, fmt.Errorf("request failed: %s", data) + case "REQHEADER", "RESPHEADER": + // Do nothing with these. + case "BODY": + unquoted, err := strconv.Unquote(data) + if err != nil { + return nil, fmt.Errorf("malformed body line: %q", data) + } + respBytes.WriteString(unquoted) + case "STATS": + if err := json.Unmarshal([]byte(data), &resp.Metrics); err != nil { + return nil, fmt.Errorf("malformed stats line: %q", data) + } + default: + return nil, fmt.Errorf("malformed line: %q", line) + } + } + if respBytes.Len() == 0 { + return nil, fmt.Errorf("empty response") + } + if err := json.Unmarshal(respBytes.Bytes(), &resp.Response); err != nil { + return nil, fmt.Errorf("malformed JSON response %q: %w", string(respBytes.Bytes()), err) + } + return &resp, nil +} + +// instrumentedRequest makes an HTTP request to the ollama API. +// It returns the raw bytestream from the instrumented request logs. +func (llm *Ollama) instrumentedRequest(ctx context.Context, method, endpoint string, data []byte) ([]byte, error) { if endpoint != "" && !strings.HasPrefix(endpoint, "/") { return nil, fmt.Errorf("endpoint must be empty or start with '/', got %q", endpoint) } - return llm.server.HTTPRequest(ctx, endpoint, data) + argvFn := func(hostPort string) []string { + argv := []string{ + "httpclient", + fmt.Sprintf("--method=%s", method), + fmt.Sprintf("--url=%s%s", hostPort, endpoint), + } + if data != nil { + argv = append(argv, fmt.Sprintf("--post_base64=%s", base64.StdEncoding.EncodeToString(data))) + } + if ctxDeadline, hasDeadline := ctx.Deadline(); hasDeadline { + argv = append(argv, fmt.Sprintf("--timeout=%v", time.Until(ctxDeadline))) + } + return argv + } + rawResponse, err := llm.server.InstrumentedRequest(ctx, argvFn) + if err != nil { + return nil, fmt.Errorf("%s: %w", endpoint, err) + } + return rawResponse, nil } // jsonGet performs a JSON HTTP GET request. -func jsonGet[Out any](ctx context.Context, llm *Ollama, endpoint string) (Out, error) { - var resp Out - out, err := llm.request(ctx, endpoint, nil) +func jsonGet[Out any](ctx context.Context, llm *Ollama, endpoint string) (*apiResponse[Out], error) { + out, err := llm.instrumentedRequest(ctx, "GET", endpoint, nil) if err != nil { - return resp, fmt.Errorf("GET %q failed: %w", endpoint, err) + return nil, fmt.Errorf("GET %q failed: %w", endpoint, err) } - if err := json.Unmarshal(out, &resp); err != nil { - return resp, fmt.Errorf("malformed JSON response %q: %w", string(out), err) - } - return resp, nil + return makeAPIResponse[Out](out) } // jsonPost performs a JSON HTTP POST request. -func jsonPost[In, Out any](ctx context.Context, llm *Ollama, endpoint string, input In) (Out, error) { - var resp Out +func jsonPost[In, Out any](ctx context.Context, llm *Ollama, endpoint string, input In) (*apiResponse[Out], error) { query, err := json.Marshal(input) if err != nil { - return resp, fmt.Errorf("could not marshal input %v: %w", input, err) + return nil, fmt.Errorf("could not marshal input %v: %w", input, err) } - out, err := llm.request(ctx, endpoint, query) + out, err := llm.instrumentedRequest(ctx, "POST", endpoint, query) if err != nil { - return resp, fmt.Errorf("POST %q %v failed: %w", endpoint, string(query), err) + return nil, fmt.Errorf("POST %q %v failed: %w", endpoint, string(query), err) } - if err := json.Unmarshal(out, &resp); err != nil { - return resp, fmt.Errorf("malformed JSON response %q: %w", string(out), err) - } - return resp, nil + return makeAPIResponse[Out](out) } // listModelNames lists the available model names. @@ -221,8 +298,8 @@ func (llm *Ollama) listModelNames(ctx context.Context) ([]string, error) { if err != nil { return nil, err } - modelNames := make([]string, len(models.Models)) - for i, m := range models.Models { + modelNames := make([]string, len(models.Response.Models)) + for i, m := range models.Response.Models { modelNames[i] = m.Name } return modelNames, nil @@ -231,11 +308,11 @@ func (llm *Ollama) listModelNames(ctx context.Context) ([]string, error) { // WaitUntilServing waits until ollama is serving, or the context expires. func (llm *Ollama) WaitUntilServing(ctx context.Context) error { for ctx.Err() == nil { - out, err := llm.request(ctx, "/", nil) + out, err := llm.instrumentedRequest(ctx, "GET", "/", nil) if err != nil { continue } - if string(out) == "Ollama is running" { + if strings.Contains(string(out), "Ollama is running") { return nil } } @@ -435,7 +512,8 @@ type ResponseJSON struct { // Response represents a response to a query from Ollama. type Response struct { - data ResponseJSON + data ResponseJSON + metrics ResponseMetrics } // Done returns whether the response was completely generated. @@ -526,7 +604,7 @@ func (llm *Ollama) Prompt(ctx context.Context, prompt *Prompt) (*Response, error if err != nil { return nil, llm.withServerLogsErr(ctx, fmt.Errorf("prompt (%s %q) request failed: %w", prompt.Model.Name, prompt.CleanQuery(), err)) } - return &Response{data: resp}, nil + return &Response{data: resp.Response, metrics: resp.Metrics}, nil } // PromptUntil repeatedly issues a prompt until `iterate` returns a nil error.