From 11ecb923bab6c724f003814e6c52d5cad2459d24 Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Fri, 9 Oct 2026 00:56:29 +0000 Subject: [PATCH] Add prepared snapshot ownership benchmarks --- .../prepared_snapshot_benchmark_test.go | 237 ++++++++++++++++++ 1 file changed, 237 insertions(+) create mode 100644 internal/usecase/prepared_snapshot_benchmark_test.go diff --git a/internal/usecase/prepared_snapshot_benchmark_test.go b/internal/usecase/prepared_snapshot_benchmark_test.go new file mode 100644 index 0000000..de8fbb7 --- /dev/null +++ b/internal/usecase/prepared_snapshot_benchmark_test.go @@ -0,0 +1,237 @@ +package usecase + +import ( + "context" + "fmt" + "reflect" + "runtime" + "testing" + + "gitea.maximumdirect.net/eric/promptkit/internal/domain" + "gitea.maximumdirect.net/eric/promptkit/internal/validate" +) + +var preparedSnapshotHandleSink *PreparedExecution + +const preparedSnapshotRetainedBatchSize = 4 + +type preparedSnapshotBenchmarkClient struct{} + +func (preparedSnapshotBenchmarkClient) Generate( + context.Context, + domain.GenerateRequest, +) (*domain.GenerateResponse, error) { + return &domain.GenerateResponse{Content: "ok"}, nil +} + +type preparedSnapshotBenchmarkValidation struct{ schema any } + +func (validation *preparedSnapshotBenchmarkValidation) Validate( + context.Context, + *domain.Artifact, +) (domain.ValidationResult, error) { + return domain.ValidationResult{Status: domain.ValidationPassed, IsValid: true}, nil +} + +func (validation *preparedSnapshotBenchmarkValidation) SchemaDocument() any { + return validation.schema +} + +type preparedSnapshotBenchmarkValidator struct{ schema any } + +func (validator preparedSnapshotBenchmarkValidator) Validate( + context.Context, + *domain.Artifact, + domain.OutputContract, +) (domain.ValidationResult, error) { + return domain.ValidationResult{Status: domain.ValidationPassed, IsValid: true}, nil +} + +func (validator preparedSnapshotBenchmarkValidator) PrepareValidation( + context.Context, + domain.OutputContract, +) (validate.PreparedValidation, error) { + return &preparedSnapshotBenchmarkValidation{schema: validator.schema}, nil +} + +func preparedSnapshotBenchmarkSchema(properties int) map[string]any { + fields := make(map[string]any, properties) + for index := range properties { + fields[fmt.Sprintf("property-%03d", index)] = map[string]any{"type": "string"} + } + return map[string]any{ + "type": "object", "properties": fields, "additionalProperties": false, + } +} + +func preparedSnapshotBenchmarkExtras(entries int) map[string]any { + providers := make(map[string]any, entries) + for index := range entries { + providers[fmt.Sprintf("provider-%03d", index)] = map[string]any{ + "enabled": true, + "labels": []any{fmt.Sprintf("tier-%d", index%4), fmt.Sprintf("region-%d", index%3)}, + "limits": map[string]any{"burst": index + 1, "period": "minute"}, + "routes": []any{map[string]any{"weight": index + 1, "fallback": true}}, + } + } + return map[string]any{"provider_options": providers} +} + +func newPreparedSnapshotBenchmarkRunner( + mutable bool, +) (*Runner, domain.RunRequest, any) { + messageCount := 1 + extraEntries := 0 + schemaProperties := 1 + if mutable { + messageCount = 8 + extraEntries = 64 + schemaProperties = 256 + } + + messages := make([]domain.RenderedMessage, messageCount) + for index := range messages { + messages[index] = domain.RenderedMessage{ + Role: domain.RoleUser, Content: fmt.Sprintf("rendered message %d", index), + CacheControl: &domain.CacheControl{Type: domain.CacheControlEphemeral, TTL: "1h"}, + } + } + renderer := &fakeRenderer{rendered: &domain.RenderedPrompt{ + SessionID: "snapshot-session", Messages: messages, + }} + + profile := defaultExecutionProfile() + if extraEntries > 0 { + profile.ResolvedTarget.ExtraParams = preparedSnapshotBenchmarkExtras(extraEntries) + } + profileRepository := &fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": profile}} + definition := promptDef(domain.FormatJSON, domain.ValidationJSONSchema, 0) + definition.Validation.SchemaPath = "snapshot-schema.json" + schema := preparedSnapshotBenchmarkSchema(schemaProperties) + validator := preparedSnapshotBenchmarkValidator{schema: schema} + runner := NewRunner( + &fakePromptRepo{def: definition}, + profileRepository, + staticArtifactReader{}, + renderer, + preparedSnapshotBenchmarkClient{}, + validator, + nil, + ) + request := domain.RunRequest{ + PromptID: "p", PromptVersion: "1", ProfileID: "exec", + Inputs: singleInputRef(), Vars: map[string]string{"name": "Ada"}, + AppendedMessages: []domain.RenderedMessage{{ + Role: domain.RoleAssistant, Content: "appended message", + CacheControl: &domain.CacheControl{Type: domain.CacheControlEphemeral, TTL: "1h"}, + }}, + } + return runner, request, schema +} + +func checkPreparedSnapshotBenchmarkFixture(b *testing.B, runner *Runner, request domain.RunRequest, schema any) { + b.Helper() + prepared, err := runner.PrepareExecution(context.Background(), request) + if err != nil { + b.Fatal(err) + } + details := prepared.Details() + if details == nil || details.PromptID != request.PromptID || len(details.Messages) < 2 || + details.Messages[0].Content != "rendered message 0" || details.Messages[len(details.Messages)-1].Content != "appended message" || + details.InputHashes["transcript"] == "" || details.StructuredOutput == nil || details.StructuredOutput.JSONSchema == nil { + prepared.Discard() + b.Fatalf("prepared snapshot fixture is incomplete: %#v", details) + } + if got := details.StructuredOutput.JSONSchema.Schema; !reflect.DeepEqual(got, schema) { + prepared.Discard() + b.Fatalf("prepared schema differs from fixture: %#v", got) + } + result, err := runner.RunPrepared(context.Background(), prepared) + if err != nil || result == nil || result.RawOutput != "ok" || !result.Validation.IsValid { + b.Fatalf("completed fixture=(%+v, %v)", result, err) + } + completedDetails := prepared.Details() + if completedDetails == nil || completedDetails.Messages[0].Content != "rendered message 0" || + !reflect.DeepEqual(completedDetails.StructuredOutput.JSONSchema.Schema, schema) { + b.Fatalf("completed handle did not retain its details: %#v", completedDetails) + } +} + +func BenchmarkPreparedSnapshot(b *testing.B) { + for _, fixture := range []struct { + name string + mutable bool + }{ + {name: "small"}, + {name: "mutable_payload", mutable: true}, + } { + b.Run(fixture.name, func(b *testing.B) { + runner, request, schema := newPreparedSnapshotBenchmarkRunner(fixture.mutable) + checkPreparedSnapshotBenchmarkFixture(b, runner, request, schema) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + prepared, err := runner.PrepareExecution(context.Background(), request) + if err != nil { + b.Fatal(err) + } + prepared.Discard() + preparedSnapshotHandleSink = prepared + } + b.StopTimer() + preparedSnapshotHandleSink = nil + }) + } +} + +func BenchmarkRetainedPreparedHandles(b *testing.B) { + for _, state := range []string{"ready", "completed", "discarded"} { + b.Run(state, func(b *testing.B) { + runner, request, schema := newPreparedSnapshotBenchmarkRunner(true) + checkPreparedSnapshotBenchmarkFixture(b, runner, request, schema) + retained := make([]*PreparedExecution, preparedSnapshotRetainedBatchSize) + runtime.GC() + var setupMemory runtime.MemStats + runtime.ReadMemStats(&setupMemory) + runtime.KeepAlive(runner) + runtime.KeepAlive(request) + runtime.KeepAlive(retained) + + retainHandle := func() *PreparedExecution { + prepared, err := runner.PrepareExecution(context.Background(), request) + if err != nil { + b.Fatal(err) + } + switch state { + case "completed": + result, err := runner.RunPrepared(context.Background(), prepared) + if err != nil || result == nil || result.RawOutput != "ok" || !result.Validation.IsValid { + b.Fatalf("complete retained handle=(%+v, %v)", result, err) + } + case "discarded": + prepared.Discard() + } + return prepared + } + + b.ReportAllocs() + b.ResetTimer() + for iteration := 0; iteration < b.N; iteration++ { + retained[iteration%len(retained)] = retainHandle() + } + b.StopTimer() + for index, prepared := range retained { + if prepared == nil { + retained[index] = retainHandle() + } + } + runtime.GC() + var retainedMemory runtime.MemStats + runtime.ReadMemStats(&retainedMemory) + b.ReportMetric(float64(int64(retainedMemory.HeapAlloc)-int64(setupMemory.HeapAlloc)), "retained-B") + runtime.KeepAlive(runner) + runtime.KeepAlive(request) + runtime.KeepAlive(retained) + }) + } +}