diff --git a/Makefile b/Makefile index cdc2bbfed..ed6e6fdde 100644 --- a/Makefile +++ b/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)) diff --git a/images/gpu/cuda-tests/Dockerfile b/images/gpu/cuda-tests/Dockerfile index 207ab2724..d335c691b 100644 --- a/images/gpu/cuda-tests/Dockerfile +++ b/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"] diff --git a/images/gpu/ollama/Dockerfile b/images/gpu/ollama/Dockerfile new file mode 100644 index 000000000..7b6a22fbe --- /dev/null +++ b/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" \ +' diff --git a/pkg/test/dockerutil/container.go b/pkg/test/dockerutil/container.go index 0b0991f28..c31a357af 100644 --- a/pkg/test/dockerutil/container.go +++ b/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, diff --git a/pkg/test/dockerutil/gpu.go b/pkg/test/dockerutil/gpu.go index a116cec86..76c8de32f 100644 --- a/pkg/test/dockerutil/gpu.go +++ b/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{ diff --git a/test/gpu/BUILD b/test/gpu/BUILD index 34056fd6a..52188b456 100644 --- a/test/gpu/BUILD +++ b/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", + ], +) diff --git a/test/gpu/ollama/BUILD b/test/gpu/ollama/BUILD new file mode 100644 index 000000000..05f4200f3 --- /dev/null +++ b/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", + ], +) diff --git a/test/gpu/ollama/ollama.go b/test/gpu/ollama/ollama.go new file mode 100644 index 000000000..0a816c1a1 --- /dev/null +++ b/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 ", r.data.Response) + } + return "" + } + 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) +} diff --git a/test/gpu/textgen_test.go b/test/gpu/textgen_test.go new file mode 100644 index 000000000..9cceb2386 --- /dev/null +++ b/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) + }) +}