Add ollama GPU test.

This runs https://ollama.ai/ in a gVisor container and loads two models:
an English-Chinese translation model, and a code assistant model.

It asks the first one to translate "Hello World" to Chinese, and then asks
the second one to generate a test case to verify that the translation is
correct.

This change includes a server and client library for spawning ollama in a
container and interacting through its HTTP API. This will be useful to turn
it into a benchmark that measures its throughput in tokens/second.

PiperOrigin-RevId: 590295278

Makefile

@@ -273,18 +273,39 @@ arm-qemu-smoke-test: $(RUNTIME_BIN) load-arm-qemu
 simple-tests: unit-tests # Compatibility target.
 .PHONY: simple-tests
 
-gpu-smoke-tests: load-basic_cuda-vector-add load-gpu_cuda-tests $(RUNTIME_BIN)
+# Images needed for GPU smoke tests.
+gpu-smoke-images: load-basic_cuda-vector-add load-gpu_cuda-tests
+.PHONY: gpu-smoke-images
+
+gpu-smoke-tests: gpu-smoke-images $(RUNTIME_BIN)
 	@$(call test,--test_env=RUNTIME=runc //test/gpu:gpu_smoke_test)
 	@$(call install_runtime,$(RUNTIME),--nvproxy=true --nvproxy-docker=true)
 	@$(call sudo,test/gpu:gpu_smoke_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 .PHONY: gpu-smoke-tests
 
-cos-gpu-smoke-tests: load-basic_cuda-vector-add load-gpu_cuda-tests $(RUNTIME_BIN)
+cos-gpu-smoke-tests: gpu-smoke-images $(RUNTIME_BIN)
 	@$(call sudo,test/gpu:gpu_smoke_test,--runtime=runc -test.v --cos-gpu $(ARGS))
 	@$(call install_runtime,$(RUNTIME),--nvproxy=true)
 	@$(call sudo,test/gpu:gpu_smoke_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 .PHONY: cos-gpu-smoke-tests
 
+# Images needed for GPU tests.
+# 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_ollama load-basic_busybox load-basic_python
+.PHONY: gpu-images
+
+gpu-all-tests: gpu-images gpu-smoke-tests $(RUNTIME_BIN)
+	@$(call install_runtime,$(RUNTIME),--nvproxy=true --nvproxy-docker=true)
+	@$(call sudo,test/gpu:textgen_test,--runtime=$(RUNTIME) -test.v $(ARGS))
+.PHONY: gpu-all-tests
+
+cos-gpu-all-tests: gpu-images cos-gpu-smoke-tests $(RUNTIME_BIN)
+	@$(call install_runtime,$(RUNTIME),--nvproxy=true)
+	@$(call sudo,test/gpu:textgen_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
+.PHONY: cos-gpu-all-tests
+
 portforward-tests: load-basic_redis load-basic_nginx $(RUNTIME_BIN)
 	@$(call install_runtime,$(RUNTIME),--network=sandbox)
 	@$(call sudo,test/root:portforward_test,--runtime=$(RUNTIME) -test.v $(ARGS))

images/gpu/cuda-tests/Dockerfile

@@ -4,4 +4,5 @@ WORKDIR /
 COPY cuda_malloc_managed.cu .
 COPY cuda_test_util.h .
 COPY run.sh .
+ENV PATH=$PATH:/usr/local/nvidia/bin:/bin/nvidia/bin
 ENTRYPOINT ["/run.sh"]

images/gpu/ollama/Dockerfile

@@ -0,0 +1,17 @@
+# https://hub.docker.com/r/ollama/ollama
+FROM ollama/ollama:0.1.13
+
+ENV PATH=$PATH:/usr/local/nvidia/bin:/bin/nvidia/bin
+
+# Pre-install a few models.
+# Although these are the smallest possible model size (7B parameters),
+# these are still quite large and it would take too long for tests to
+# download them on every run.
+RUN bash -c '                                  \
+    ( ollama serve ) & serverpid="$!";         \
+    sleep 5;                                   \
+    ollama pull codellama:7b &&                \
+    ollama pull llama2-chinese:7b-chat &&      \
+    kill "$serverpid" &&                       \
+    wait "$serverpid"                          \
+'

pkg/test/dockerutil/container.go

@@ -86,6 +86,9 @@ type RunOpts struct {
 	// User is the user to use.
 	User string
 
+	// Optional argv to override the ENTRYPOINT specified in the image.
+	Entrypoint []string
+
 	// Privileged enables privileged mode.
 	Privileged bool
 
@@ -260,6 +263,7 @@ func (c *Container) config(r RunOpts, args []string) *container.Config {
 	return &container.Config{
 		Image:        testutil.ImageByName(r.Image),
 		Cmd:          args,
+		Entrypoint:   r.Entrypoint,
 		ExposedPorts: ports,
 		Env:          env,
 		WorkingDir:   r.WorkDir,

pkg/test/dockerutil/gpu.go

@@ -17,6 +17,7 @@ package dockerutil
 
 import (
 	"flag"
+	"os"
 
 	"github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/api/types/mount"
@@ -55,17 +56,32 @@ func GPURunOpts() RunOpts {
 		})
 	}
 
-	mounts := []mount.Mount{
-		{
-			Source: "/var/lib/nvidia/lib64",
-			Target: "/usr/local/nvidia/lib64",
-			Type:   mount.TypeBind,
-		},
-		{
-			Source: "/var/lib/nvidia/bin",
-			Target: "/usr/local/nvidia/bin",
-			Type:   mount.TypeBind,
-		},
+	var mounts []mount.Mount
+	for _, nvidiaBin := range []string{
+		"/home/kubernetes/bin/nvidia/bin",
+		"/var/lib/nvidia/bin",
+	} {
+		if st, err := os.Stat(nvidiaBin); err == nil && st.IsDir() {
+			mounts = append(mounts, mount.Mount{
+				Source:   nvidiaBin,
+				Target:   "/usr/local/nvidia/bin",
+				Type:     mount.TypeBind,
+				ReadOnly: true,
+			})
+		}
+	}
+	for _, nvidiaLib64 := range []string{
+		"/home/kubernetes/bin/nvidia/lib64",
+		"/var/lib/nvidia/lib64",
+	} {
+		if st, err := os.Stat(nvidiaLib64); err == nil && st.IsDir() {
+			mounts = append(mounts, mount.Mount{
+				Source:   nvidiaLib64,
+				Target:   "/usr/local/nvidia/lib64",
+				Type:     mount.TypeBind,
+				ReadOnly: true,
+			})
+		}
 	}
 
 	return RunOpts{

test/gpu/BUILD

@@ -16,3 +16,19 @@ go_test(
     visibility = ["//:sandbox"],
     deps = ["//pkg/test/dockerutil"],
 )
+
+go_test(
+    name = "textgen_test",
+    srcs = ["textgen_test.go"],
+    tags = [
+        "local",
+        "noguitar",
+        "notap",
+    ],
+    visibility = ["//:sandbox"],
+    deps = [
+        "//pkg/test/dockerutil",
+        "//pkg/test/testutil",
+        "//test/gpu/ollama",
+    ],
+)

test/gpu/ollama/BUILD

@@ -0,0 +1,17 @@
+load("//tools:defs.bzl", "go_library")
+
+package(
+    default_applicable_licenses = ["//:license"],
+    licenses = ["notice"],
+)
+
+go_library(
+    name = "ollama",
+    testonly = 1,
+    srcs = ["ollama.go"],
+    visibility = ["//:sandbox"],
+    deps = [
+        "//pkg/test/dockerutil",
+        "//pkg/test/testutil",
+    ],
+)

test/gpu/ollama/ollama.go

@@ -0,0 +1,377 @@
+// Copyright 2023 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 ollama provides an Ollama API client.
+package ollama
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strings"
+	"time"
+
+	"gvisor.dev/gvisor/pkg/test/dockerutil"
+	"gvisor.dev/gvisor/pkg/test/testutil"
+)
+
+const (
+	// Port is the port used by the ollama server.
+	Port = 11434
+
+	// curtQuery is a query that should result in a very curt response.
+	curtQuery = `Please reply with the single word: "Hello". Do not reply with any other word.`
+)
+
+// Ollama is an ollama client.
+type Ollama struct {
+	container *dockerutil.Container
+	logger    testutil.Logger
+
+	// ModelNames is the list of available model names.
+	ModelNames []string
+
+	// 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.
+	HasGPU bool
+}
+
+// New starts a new Ollama server in the given container,
+// then waits for it to serve and returns the client.
+func New(ctx context.Context, cont *dockerutil.Container, logger testutil.Logger) (*Ollama, error) {
+	opts := dockerutil.GPURunOpts()
+	opts.Image = "gpu/ollama"
+	if err := cont.Spawn(ctx, opts); err != nil {
+		return nil, fmt.Errorf("could not start ollama: %v", err)
+	}
+	llm := &Ollama{
+		container: cont,
+		logger:    logger,
+	}
+
+	// Wait until serving.
+	if err := llm.WaitUntilServing(ctx); err != nil {
+		return nil, fmt.Errorf("ollama did not come up for serving: %w", err)
+	}
+
+	// Get list of model names.
+	modelNames, err := llm.listModelNames(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("could not list model names: %w", err)
+	}
+	if len(modelNames) == 0 {
+		return nil, errors.New("no models available")
+	}
+	llm.ModelNames = modelNames
+
+	// Load the first model.
+	// This is necessary to force ollama to load a model, without which
+	// we cannot detect if it is using the GPU or not.
+	_, err = llm.Prompt(ctx, &Prompt{
+		Model: &Model{Name: llm.ModelNames[0]},
+		Query: curtQuery,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("could not load first model %q: %w", llm.ModelNames[0], err)
+	}
+
+	// Now go over the logs and check if the GPU was used.
+	logs, err := llm.container.Logs(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("could not get logs: %w", err)
+	}
+	switch {
+	case strings.Contains(logs, "check that you have installed GPU drivers"):
+		llm.HasGPU = false
+	case strings.Contains(logs, "VRAM available"):
+		llm.HasGPU = true
+	default:
+		return nil, fmt.Errorf("cannot determine whether ollama is using GPU from logs:\n%s", logs)
+	}
+	return llm, nil
+}
+
+// request makes an HTTP request to the ollama API.
+func (llm *Ollama) request(ctx context.Context, 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)
+	}
+	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))
+	out, err := dockerutil.MakeContainer(ctx, llm.logger).Run(ctx, dockerutil.RunOpts{
+		Image: "basic/busybox",
+		Links: []string{llm.container.MakeLink("llm")},
+	}, cmd...)
+	return []byte(out), err
+}
+
+// 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)
+	if err != nil {
+		return resp, 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
+}
+
+// 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
+	query, err := json.Marshal(input)
+	if err != nil {
+		return resp, fmt.Errorf("could not marshal input %v: %w", input, err)
+	}
+	out, err := llm.request(ctx, endpoint, query)
+	if err != nil {
+		return resp, fmt.Errorf("POST %q(%v) failed: %w", endpoint, input, err)
+	}
+	if err := json.Unmarshal(out, &resp); err != nil {
+		return resp, fmt.Errorf("malformed JSON response %q: %w", string(out), err)
+	}
+	return resp, nil
+}
+
+// listModelNames lists the available model names.
+func (llm *Ollama) listModelNames(ctx context.Context) ([]string, error) {
+	type model struct {
+		Name       string `json:"name"`
+		ModifiedAt string `json:"modified_at"`
+		Size       int    `json:"size"`
+	}
+	type modelsList struct {
+		Models []model `json:"models"`
+	}
+	models, err := jsonGet[modelsList](ctx, llm, "/api/tags")
+	if err != nil {
+		return nil, err
+	}
+	modelNames := make([]string, len(models.Models))
+	for i, m := range models.Models {
+		modelNames[i] = m.Name
+	}
+	return modelNames, nil
+}
+
+// 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)
+		if err != nil {
+			continue
+		}
+		if string(out) == "Ollama is running" {
+			return nil
+		}
+	}
+	return fmt.Errorf("ollama did not respond: %w", ctx.Err())
+}
+
+// Model encodes a model and options for it.
+type Model struct {
+	// Name is the name of the ollama model, e.g. "codellama:7b".
+	Name string
+
+	// Options maps parameter names to JSON-compatible values.
+	Options map[string]any
+}
+
+// String returns the model's name.
+func (m *Model) String() string {
+	return m.Name
+}
+
+// modelTemperatureOption is the temperature option that most models have
+// which controls how free they are from deviating from their most-likely
+// token chain.
+const modelTemperatureOption = "temperature"
+
+// RaiseTemperature increases the "temperature" option of the model,
+// if any.
+func (m *Model) RaiseTemperature() {
+	temp, ok := m.Options[modelTemperatureOption]
+	if !ok {
+		temp = float64(0.0)
+	}
+	if m.Options == nil {
+		m.Options = map[string]any{}
+	}
+	m.Options[modelTemperatureOption] = min(1.0, temp.(float64)*2+.025)
+}
+
+// ZeroTemperatureModel returns a Model with the given name and an initial
+// temperature setting of zero. This setting allows for consistent settings.
+func ZeroTemperatureModel(name string) *Model {
+	return &Model{
+		Name: name,
+		Options: map[string]any{
+			modelTemperatureOption: 0.0,
+		},
+	}
+}
+
+// Prompt is an ollama prompt.
+type Prompt struct {
+	// Model is the model to query.
+	Model *Model
+
+	// Query is the prompt string.
+	Query string
+
+	// Context is the conversational context to follow up on, if any.
+	// This is returned from `Response`.
+	Context ConversationContext
+}
+
+// String returns a human-friendly string representing this prompt.
+func (p *Prompt) String() string {
+	return fmt.Sprintf("[%v] %s", p.Model, p.Query)
+}
+
+// PromptJSON encodes the JSON data for a query.
+type PromptJSON struct {
+	Model   string              `json:"model"`
+	Prompt  string              `json:"prompt"`
+	Stream  bool                `json:"stream"`
+	Context ConversationContext `json:"context"`
+	Options map[string]any      `json:"options"`
+}
+
+// json encodes this prompt to the JSON format expected by Ollama.
+func (p *Prompt) json() PromptJSON {
+	return PromptJSON{
+		Model:   p.Model.Name,
+		Prompt:  p.Query,
+		Stream:  false,
+		Context: p.Context,
+		Options: p.Model.Options,
+	}
+}
+
+// ResponseJSON is the JSON-format response from ollama about a prompt in
+// non-streamed mode.
+type ResponseJSON struct {
+	Model           string              `json:"model"`
+	Response        string              `json:"response"`
+	Done            bool                `json:"done"`
+	TotalNanos      int                 `json:"total_duration"`
+	LoadNanos       int                 `json:"load_duration"`
+	EvalCount       int                 `json:"eval_count"`
+	EvalNanos       int                 `json:"eval_duration"`
+	PromptEvalCount int                 `json:"prompt_eval_count"`
+	PromptEvalNanos int                 `json:"prompt_eval_duration"`
+	Context         ConversationContext `json:"context"`
+}
+
+// Response represents a response to a query from Ollama.
+type Response struct {
+	data ResponseJSON
+}
+
+// Done returns whether the response was completely generated.
+func (r *Response) Done() bool {
+	return r.data.Done
+}
+
+// String returns the response text, if it is done.
+func (r *Response) String() string {
+	if !r.data.Done {
+		if r.data.Response != "" {
+			return fmt.Sprintf("%s <NOT DONE>", r.data.Response)
+		}
+		return "<NOT DONE>"
+	}
+	return r.data.Response
+}
+
+// Text returns the body of the response, if it is done.
+func (r *Response) Text() string {
+	return r.data.Response
+}
+
+// TotalDuration returns the total response generation time.
+func (r *Response) TotalDuration() time.Duration {
+	return time.Duration(r.data.TotalNanos) * time.Nanosecond
+}
+
+// LoadDuration returns the load response generation time.
+func (r *Response) LoadDuration() time.Duration {
+	return time.Duration(r.data.LoadNanos) * time.Nanosecond
+}
+
+// EvalDuration returns the response evaluation time.
+func (r *Response) EvalDuration() time.Duration {
+	return time.Duration(r.data.EvalNanos) * time.Nanosecond
+}
+
+// PromptEvalDuration returns the prompt evaluation time.
+func (r *Response) PromptEvalDuration() time.Duration {
+	return time.Duration(r.data.PromptEvalNanos) * time.Nanosecond
+}
+
+// TokensPerSecond computes the number of tokens generated per second.
+func (r *Response) TokensPerSecond() float64 {
+	if !r.data.Done || r.EvalDuration() == 0 {
+		return 0
+	}
+	return float64(r.data.EvalCount) / 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
+
+// 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())
+	if err != nil {
+		return nil, err
+	}
+	return &Response{data: resp}, nil
+}
+
+// PromptUntil repeatedly issues a prompt until `iterate` returns a nil error.
+// `iterate` may optionally return an updated `Prompt` which will be used to
+// follow up.
+// This is useful to work around the flakiness of LLMs in tests.
+func (llm *Ollama) PromptUntil(ctx context.Context, prompt *Prompt, iterate func(*Prompt, *Response) (*Prompt, error)) (*Response, error) {
+	var lastResponse *Response
+	var lastError error
+	attempts := 0
+	for ctx.Err() == nil {
+		response, err := llm.Prompt(ctx, prompt)
+		if err != nil {
+			return nil, fmt.Errorf("prompt request failed: %w", err)
+		}
+		attempts++
+		newPrompt, err := iterate(prompt, response)
+		if err == nil {
+			return response, nil
+		}
+		if newPrompt != nil {
+			prompt = newPrompt
+		}
+		lastResponse = response
+		lastError = err
+	}
+	return nil, fmt.Errorf("response %q (attempt #%d with prompt %v) did not match predicate: %v", lastResponse, attempts, prompt, lastError)
+}

test/gpu/textgen_test.go

@@ -0,0 +1,208 @@
+// Copyright 2023 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 textgen_test runs ollama and generates some text with it.
+package textgen_test
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	"gvisor.dev/gvisor/pkg/test/dockerutil"
+	"gvisor.dev/gvisor/pkg/test/testutil"
+	"gvisor.dev/gvisor/test/gpu/ollama"
+)
+
+// extractCode extracts code between two code block markers.
+func extractCode(response, codeBlockDelim string) (string, error) {
+	if !strings.Contains(response, codeBlockDelim) {
+		return "", fmt.Errorf("no marker string %q", codeBlockDelim)
+	}
+	var codeLines []string
+	isCodeBlock := false
+	for _, line := range strings.Split(response, "\n") {
+		if strings.HasPrefix(line, codeBlockDelim) {
+			isCodeBlock = !isCodeBlock
+		} else if isCodeBlock {
+			codeLines = append(codeLines, line)
+		}
+	}
+	if isCodeBlock {
+		return "", errors.New("non-terminated code block")
+	}
+	if len(codeLines) == 0 {
+		return "", errors.New("no or empty code block")
+	}
+	return strings.Join(codeLines, "\n") + "\n", nil
+}
+
+// runSandboxedPython runs the given Python code in a sandboxed container.
+func runSandboxedPython(ctx context.Context, logger testutil.Logger, code string) (string, error) {
+	return dockerutil.MakeContainer(ctx, logger).Run(ctx, dockerutil.RunOpts{
+		Image:       "basic/python",
+		NetworkMode: "none",
+		Entrypoint:  []string{"python3"},
+		Env:         []string{"PYTHONUTF8=1"},
+	}, "-c", code)
+}
+
+// TestLLM tests an LLM running in a sandboxed container.
+// It first asks it to translate "Hello World" to Chinese.
+// Then it asks it to write a unit test that verifies that
+// this text is a correct translation.
+func TestLLM(t *testing.T) {
+	ctx := context.Background()
+	// Run the LLM.
+	llmContainer := dockerutil.MakeContainer(ctx, t)
+	defer llmContainer.CleanUp(ctx)
+	startCtx, startCancel := context.WithTimeout(ctx, 30*time.Second)
+	llm, err := ollama.New(startCtx, llmContainer, t)
+	startCancel()
+	if err != nil {
+		t.Fatalf("Failed to start ollama: %v", err)
+	}
+	if !llm.HasGPU {
+		t.Fatal("LLM is not using a GPU")
+	}
+
+	// Query it.
+	var translation string
+	t.Run("translate text", func(t *testing.T) {
+		prompt := ollama.Prompt{
+			Model: ollama.ZeroTemperatureModel("llama2-chinese:7b-chat"),
+			Query: `
+				Translate the following text from English to Chinese:
+				    "Hello World".
+			`,
+		}
+		promptCtx, promptCancel := context.WithTimeout(ctx, time.Minute)
+		response, err := llm.PromptUntil(promptCtx, &prompt, func(prompt *ollama.Prompt, response *ollama.Response) (*ollama.Prompt, error) {
+			defer prompt.Model.RaiseTemperature()
+			text := strings.TrimSpace(response.Text())
+			for _, unacceptable := range []rune{'"', '\'', '\\', '\n', '\r', '\t'} {
+				if strings.ContainsRune(text, unacceptable) {
+					return prompt, fmt.Errorf("response contains unacceptable character %q", unacceptable)
+				}
+			}
+			for _, acceptableWord := range []string{
+				"你好",
+				"世界",
+			} {
+				if strings.Contains(text, acceptableWord) {
+					return prompt, nil
+				}
+			}
+			return prompt, errors.New("text does not contain any of the expected words")
+		})
+		promptCancel()
+		if err != nil {
+			t.Fatalf("translation failed: %v", err)
+		}
+		translation = strings.TrimSpace(response.Text())
+		t.Logf("The Chinese translation of %q is: %q", "Hello World", translation)
+	})
+	if t.Failed() {
+		return
+	}
+	t.Run("generate test case", func(t *testing.T) {
+		const (
+			markerString   = "FOOBARBAZQUUX"
+			hello          = "你好"
+			world          = "世界"
+			codeBlockDelim = "```"
+		)
+		promptCtx, promptCancel := context.WithTimeout(ctx, 3*time.Minute)
+		prompt := ollama.Prompt{
+			Model: ollama.ZeroTemperatureModel("codellama:7b"),
+			Query: fmt.Sprintf(`
+				Generate a Python function that takes a string and verifies that it
+				is a valid Chinese translation of the English phrase "Hello World".
+				The function should first turn its input into lowercase in order to
+				match case-insensitively, and remove all spaces.
+				Then, the function should verify that the phrase contains at least
+				"你好" ("hello") or "世界" ("world").
+				If the verification succeeds, the function should return True.
+				After this function is defined, you should call this function with
+				the input string %q.
+				Then, the code should verify that the function call returned True.
+				If it did, the code should print "Verification succeeded";
+				otherwise, it should print "Verification failed".
+				You may use Python code comments, but do not otherwise explain how
+				the code works and do not provide usage examples.
+				Output a single block of Python code wrapped between %q marks.
+			`, markerString, codeBlockDelim),
+		}
+		response, err := llm.PromptUntil(promptCtx, &prompt, func(prompt *ollama.Prompt, response *ollama.Response) (*ollama.Prompt, error) {
+			defer prompt.Model.RaiseTemperature()
+			pythonCode, err := extractCode(response.Text(), codeBlockDelim)
+			if err != nil {
+				return prompt, fmt.Errorf("code extraction failed: %w", err)
+			}
+			if !strings.Contains(pythonCode, markerString) {
+				return prompt, fmt.Errorf("marker string %q is not in a code block", markerString)
+			}
+			out, err := runSandboxedPython(ctx, t, pythonCode)
+			if err != nil {
+				return prompt, fmt.Errorf("execution with marker string failed: %w", err)
+			}
+			out = strings.TrimSpace(out)
+			if out == "" {
+				return prompt, fmt.Errorf("execution with marker string %q had no output", markerString)
+			}
+			if out == "Verification succeeded" {
+				return prompt, fmt.Errorf("verification did not fail for marker string %q (we expected it to fail for this string): got output %q", markerString, out)
+			}
+			if out != "Verification failed" {
+				return prompt, fmt.Errorf("verification program returned unexpected output %q for marker string %q", out, markerString)
+			}
+			for _, word := range []string{hello, world} {
+				codeWithRealText := strings.ReplaceAll(pythonCode, markerString, fmt.Sprintf("asdf %s fdsa", word))
+				out, err = runSandboxedPython(ctx, t, codeWithRealText)
+				if err != nil {
+					return prompt, fmt.Errorf("execution with word %q failed: %w", word, err)
+				}
+				out = strings.TrimSpace(out)
+				if out == "" {
+					return prompt, fmt.Errorf("execution with word %q had no output", word)
+				}
+				if out != "Verification succeeded" {
+					return prompt, fmt.Errorf("verification with word %q failed: got output %q", word, out)
+				}
+			}
+			return nil, nil
+		})
+		promptCancel()
+		if err != nil {
+			t.Fatalf("Code generation prompt failed: %v", err)
+		}
+		pythonCode, err := extractCode(response.Text(), codeBlockDelim)
+		if err != nil {
+			t.Fatalf("Code extraction failed: %v", err)
+		}
+		testCode := strings.ReplaceAll(pythonCode, markerString, translation)
+		out, err := runSandboxedPython(ctx, t, testCode)
+		if err != nil {
+			t.Fatalf("Translation verification with string %q failed: %v\nCode used:\n\n%s\n\n", translation, err, testCode)
+		}
+		out = strings.TrimSpace(out)
+		if out != "Verification succeeded" {
+			t.Fatalf("Translation verification with string %q failed: %q\nCode used:\n\n%s\n\n", translation, out, testCode)
+		}
+		t.Logf("Translation verification succeeded with code:\n\n%s\n\n", pythonCode)
+	})
+}