add vllm benchmark

The benchmark loads a small model to allow running on many platforms and to keep execution time low for experimentation. It will make sense to add a version to this test that loads a larger model (e.g. llama3-8b).
2026-05-22 17:12:49 -07:00 · 2024-07-09 09:43:26 +02:00
parent fa204cf4cc
commit 2cf6705467
5 changed files with 236 additions and 3 deletions
@@ -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-gpu_ollama_client load-basic_busybox load-basic_python load-gpu_stable-diffusion-xl
+gpu-images: gpu-smoke-images load-gpu_pytorch load-gpu_ollama load-gpu_ollama_client load-basic_busybox load-basic_python load-gpu_stable-diffusion-xl load-gpu_vllm
 .PHONY: gpu-images

 gpu-all-tests: gpu-images gpu-smoke-tests $(RUNTIME_BIN)
@@ -0,0 +1,21 @@
+# v2.45.1
+FROM alpine/git@sha256:4d7fe8d770483993c0cec264d49a573bac49e5239db47a9846572352e72da49c as downloader
+# post checkout hook. checks that command git lfs is available.
+RUN apk add git-lfs && \
+    git lfs install && \
+    GIT_CLONE_PROTECTION_ACTIVE=false git clone https://huggingface.co/facebook/opt-125m /opt-125m && \
+    rm /opt-125m/flax_model.msgpack && \
+    rm /opt-125m/tf_model.h5
+# post checkout hook. checks that command git lfs is available.
+RUN GIT_CLONE_PROTECTION_ACTIVE=false git clone https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered /dataset
+RUN git clone https://github.com/vllm-project/vllm.git /vllm && cd /vllm && git checkout v0.4.2
+
+# v0.4.2
+FROM vllm/vllm-openai@sha256:c63bd9e99c6d5914032b396be38a80fd9640693da62b64b907434d38bf38a8e9
+COPY --from=downloader /vllm/examples/template_chatml.jinja /vllm/examples/template_chatml.jinja
+COPY --from=downloader /vllm/benchmarks /vllm/benchmarks
+COPY --from=downloader /opt-125m /model
+COPY --from=downloader /dataset/ShareGPT_V3_unfiltered_cleaned_split.json /ShareGPT_V3_unfiltered_cleaned_split.json
+
+ENTRYPOINT ["python3"]
+CMD ["-m", "vllm.entrypoints.openai.api_server", "--model", "/model", "--chat-template", "/vllm/examples/template_chatml.jinja"]
@@ -0,0 +1,19 @@
+load("//test/benchmarks:defs.bzl", "benchmark_test")
+
+package(
+    default_applicable_licenses = ["//:license"],
+    licenses = ["notice"],
+)
+
+benchmark_test(
+    name = "vllm_test",
+    srcs = ["vllm_test.go"],
+    visibility = ["//:sandbox"],
+    deps = [
+        "//pkg/test/dockerutil",
+        "//test/benchmarks/harness",
+        "//test/benchmarks/tools",
+        "//test/metricsviz",
+        "@com_github_docker_docker//api/types/mount:go_default_library",
+    ],
+)
@@ -0,0 +1,193 @@
+// 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.
+
+// Package vllm_test holds benchmarks around the vLLM inference engine.
+package vllm_test
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/docker/docker/api/types/mount"
+	"gvisor.dev/gvisor/pkg/test/dockerutil"
+	"gvisor.dev/gvisor/test/benchmarks/harness"
+)
+
+// BenchmarkVLLM runs a vLLM workload.
+func BenchmarkVLLM(b *testing.B) {
+	doVLLMTest(b)
+}
+
+func doVLLMTest(b *testing.B) {
+	serverMachine, err := harness.GetMachine()
+	if err != nil {
+		b.Fatalf("failed to get machine: %v", err)
+	}
+	defer serverMachine.CleanUp()
+
+	ctx := context.Background()
+	serverCtr := serverMachine.GetContainer(ctx, b)
+	defer serverCtr.CleanUp(ctx)
+	if err := harness.DropCaches(serverMachine); err != nil {
+		b.Skipf("failed to drop caches: %v. You probably need root.", err)
+	}
+
+	// Run vllm.
+	runOpts := dockerutil.GPURunOpts()
+	runOpts.Image = "gpu/vllm"
+	runOpts.Env = []string{"PYTHONPATH=$PYTHONPATH:/vllm"}
+
+	if err := serverCtr.Spawn(ctx, runOpts); err != nil {
+		b.Errorf("failed to run container: %v", err)
+	}
+	if out, err := serverCtr.WaitForOutput(ctx, "Uvicorn running on http://0.0.0.0:8000", 10*time.Minute); err != nil {
+		b.Fatalf("failed to start vllm model: %v %s", err, out)
+	}
+
+	b.Run("opt-125", func(b *testing.B) {
+		ctx := context.Background()
+
+		b.ResetTimer()
+		b.StopTimer()
+
+		clientMachine, err := harness.GetMachine()
+		if err != nil {
+			b.Fatalf("failed to get machine: %v", err)
+		}
+		defer clientMachine.CleanUp()
+		clientCtr := clientMachine.GetContainer(ctx, b)
+		defer clientCtr.CleanUp(ctx)
+
+		// store vllm logs here
+		logsDir := b.TempDir()
+
+		b.StartTimer()
+		out, err := clientCtr.Run(ctx, dockerutil.RunOpts{
+			Links: []string{serverCtr.MakeLink("vllmctr")},
+			Image: "gpu/vllm",
+			Env:   []string{"PYTHONPATH=$PYTHONPATH:/vllm"},
+			Mounts: []mount.Mount{
+				// The logs dir is used because vllm only outputs json to a file.
+				{
+					Source: logsDir,
+					Target: "/tmp",
+					Type:   "bind",
+				},
+			},
+		}, "/vllm/benchmarks/benchmark_serving.py", "--num-prompts", fmt.Sprintf("%d", b.N), "--host", "vllmctr", "--model", "/model", "--tokenizer", "/model", "--endpoint", "/v1/completions", "--backend", "openai", "--dataset", "/ShareGPT_V3_unfiltered_cleaned_split.json", "--save-result", "--result-dir", "/tmp")
+		if err != nil {
+			b.Errorf("failed to run container: %v logs: %s", err, out)
+		}
+
+		b.StopTimer()
+
+		metrics, err := parseVLLMJSON(logsDir)
+		if err != nil {
+			b.Errorf("failed to parse vllm output: %v", err)
+		}
+
+		if metrics.Completed == 0 {
+			b.Errorf("did not complete at least one request")
+		}
+
+		for _, err := range metrics.Errors {
+			if err != "" {
+				b.Errorf("unexpected errors: %s", metrics.Errors)
+			}
+		}
+
+		b.ReportMetric(metrics.RequestsPerSec, "requests/sec")
+		b.ReportMetric(metrics.InputToksPerSec, "input-tok/sec")
+		b.ReportMetric(metrics.OutputToksPerSec, "output-tok/sec")
+		b.ReportMetric(metrics.TTFTAvgMS*1000000, "ttft-avg-ns")
+		b.ReportMetric(metrics.TTFTP50MS*1000000, "ttft-p50-ns")
+		b.ReportMetric(metrics.TTFTP99MS*1000000, "ttft-p99-ns")
+		b.ReportMetric(metrics.TPOTAvgMS*1000000, "tpot-avg-ns")
+		b.ReportMetric(metrics.TPOTP50MS*1000000, "tpot-p50-ns")
+		b.ReportMetric(metrics.TPOTP99MS*1000000, "tpot-p99-ns")
+	})
+}
+
+// Modeled after the metrics reported here: https://github.com/vllm-project/vllm/blob/main/benchmarks/benchmark_serving.py#L338-L358
+type metrics struct {
+	Duration          float64 `json:"duration"`
+	Completed         int     `json:"completed"`
+	TotalInputTokens  int     `json:"total_input_tokens"`
+	TotalOutputTokens int     `json:"total_output_tokens"`
+	RequestsPerSec    float64 `json:"request_throughput"`
+	InputToksPerSec   float64 `json:"input_throughput"`
+	OutputToksPerSec  float64 `json:"output_throughput"`
+	TTFTAvgMS         float64 `json:"mean_ttft_ms"`
+	TTFTP50MS         float64 `json:"median_ttft_ms"`
+	TTFTP99MS         float64 `json:"p99_ttft_ms"`
+	TPOTAvgMS         float64 `json:"mean_tpot_ms"`
+	TPOTP50MS         float64 `json:"median_tpot_ms"`
+	TPOTP99MS         float64 `json:"p99_tpot_ms"`
+	// metrics that are available but unused so far.
+	// InputLens         []int   `json:"input_lens"`
+	// OutputLens        []int   `json:"output_lens"`
+	// TTFTs             []float64   `json:"ttfts"`
+	// ITLS              [][]float64 `json:"itls"`
+	// GeneratedTexts    []string `json:"generated_texts"`
+	Errors []string `json:"errors"`
+}
+
+// parseVLLMJSON expects a path that contains only one json file.
+func parseVLLMJSON(path string) (metrics, error) {
+	files, err := os.ReadDir(path)
+	if err != nil {
+		return metrics{}, fmt.Errorf("failed to read directory: %w", err)
+	}
+
+	var jsonPath string
+	// Take the first json file as logs output.
+	// Error if there are multiple.
+	for _, name := range files {
+		if strings.HasSuffix(name.Name(), ".json") {
+			jsonPath = filepath.Join(path, name.Name())
+			continue
+		}
+
+		if jsonPath != "" && strings.HasSuffix(name.Name(), ".json") {
+			return metrics{}, errors.New("found more than one json file, expected 1")
+		}
+	}
+	if jsonPath == "" {
+		return metrics{}, errors.New("no json file found")
+	}
+
+	data, err := os.ReadFile(jsonPath)
+	if err != nil {
+		return metrics{}, fmt.Errorf("failed to read file: %w", err)
+	}
+
+	var vllm metrics
+	if err := json.Unmarshal(data, &vllm); err != nil {
+		return metrics{}, fmt.Errorf("failed to unmarshal data: %w", err)
+	}
+
+	return vllm, nil
+}
+
+func TestMain(m *testing.M) {
+	harness.Init()
+	os.Exit(m.Run())
+}
@@ -43,8 +43,8 @@ endif
 REMOTE_IMAGE_PREFIX ?= us-central1-docker.pkg.dev/gvisor-presubmit/gvisor-presubmit-images
 LOCAL_IMAGE_PREFIX  ?= gvisor.dev/images
 ALL_IMAGES          := $(subst /,_,$(subst images/,,$(shell find images/ -name Dockerfile -o -name Dockerfile.$(ARCH) | xargs -n 1 dirname | uniq)))
-NON_TEST_IMAGES     := gpu/ollama/bench
-TEST_IMAGES         := $(subst /,_,$(subst images/,,$(shell find images/ -name Dockerfile -o -name Dockerfile.$(ARCH) | xargs -n 1 dirname | uniq | grep -v $(NON_TEST_IMAGES))))
+NON_TEST_IMAGES     := gpu/ollama/bench\|gpu/vllm
+TEST_IMAGES         := $(subst /,_,$(subst images/,,$(shell find images/ -name Dockerfile -o -name Dockerfile.$(ARCH) | xargs -n 1 dirname | uniq | grep -v "$(NON_TEST_IMAGES)")))
 SUB_IMAGES          := $(foreach image,$(ALL_IMAGES),$(if $(findstring _,$(image)),$(image),))
 IMAGE_GROUPS        := $(sort $(foreach image,$(SUB_IMAGES),$(firstword $(subst _, ,$(image)))))