Add prepared snapshot ownership benchmarks
This commit is contained in:
237
internal/usecase/prepared_snapshot_benchmark_test.go
Normal file
237
internal/usecase/prepared_snapshot_benchmark_test.go
Normal file
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user