409 lines
15 KiB
Go
409 lines
15 KiB
Go
package pipeline
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"reflect"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
frameworkllm "gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
|
|
)
|
|
|
|
type countingProvider struct {
|
|
active atomic.Int32
|
|
maximum atomic.Int32
|
|
calls atomic.Int32
|
|
}
|
|
|
|
func (p *countingProvider) CompleteStructured(ctx context.Context, _ contracts.StructuredCompletionRequest, _ any) (contracts.StructuredCompletionResponse, error) {
|
|
p.calls.Add(1)
|
|
current := p.active.Add(1)
|
|
defer p.active.Add(-1)
|
|
for {
|
|
seen := p.maximum.Load()
|
|
if current <= seen || p.maximum.CompareAndSwap(seen, current) {
|
|
break
|
|
}
|
|
}
|
|
select {
|
|
case <-time.After(5 * time.Millisecond):
|
|
return contracts.StructuredCompletionResponse{}, nil
|
|
case <-ctx.Done():
|
|
return contracts.StructuredCompletionResponse{}, ctx.Err()
|
|
}
|
|
}
|
|
|
|
func preparedConcurrentPipeline(t *testing.T, chunkCount int) *PreparedPipeline {
|
|
t.Helper()
|
|
catalog := typedResolutionCatalog(t, completeTypedCatalogOptions())
|
|
resolved, err := ResolvePipeline(typedResolutionProfile(), ResolveOptions{}, catalog)
|
|
if err != nil {
|
|
t.Fatalf("ResolvePipeline() error = %v", err)
|
|
}
|
|
prepared, err := Prepare(resolved, registriesFromModuleCatalog(catalog), ModuleDependencies{})
|
|
if err != nil {
|
|
t.Fatalf("Prepare() error = %v", err)
|
|
}
|
|
doc := typedTestDocumentWithUnits(chunkCount)
|
|
if adapter, ok := prepared.input.(*typedTestInput); ok {
|
|
adapter.doc = doc
|
|
} else {
|
|
t.Fatalf("prepared input = %T, want *typedTestInput", prepared.input)
|
|
}
|
|
prepared.chunker = &typedTestChunker{key: "typed/chunk", plan: typedTestPlan(doc)}
|
|
return prepared
|
|
}
|
|
|
|
func installExtractOperation(prepared *PreparedPipeline, laneIndex int, operation func(context.Context, contracts.TypedExtractionRequest) (erasedTypedResult, error)) {
|
|
prepared.lanes[laneIndex].typed.extract = func(ctx context.Context, _ any, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
return operation(ctx, request)
|
|
}
|
|
}
|
|
|
|
func typedValueForLane(laneIndex int, chunkIndex int) any {
|
|
if laneIndex == 0 {
|
|
return codecNotes{Items: []string{fmt.Sprintf("chunk-%d", chunkIndex)}}
|
|
}
|
|
return codecScore{Value: chunkIndex}
|
|
}
|
|
|
|
func TestRunnerBoundsExtractJobsAndStabilizesReverseCompletion(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 3)
|
|
started := make(chan int, 6)
|
|
release := make([]chan struct{}, 6)
|
|
for i := range release {
|
|
release[i] = make(chan struct{})
|
|
}
|
|
var active atomic.Int32
|
|
var maximum atomic.Int32
|
|
for laneIndex := range prepared.lanes {
|
|
lane := laneIndex
|
|
installExtractOperation(prepared, lane, func(ctx context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
jobIndex := request.Chunk.Index*2 + lane
|
|
current := active.Add(1)
|
|
defer active.Add(-1)
|
|
for {
|
|
seen := maximum.Load()
|
|
if current <= seen || maximum.CompareAndSwap(seen, current) {
|
|
break
|
|
}
|
|
}
|
|
started <- jobIndex
|
|
select {
|
|
case <-release[jobIndex]:
|
|
case <-ctx.Done():
|
|
return erasedTypedResult{}, ctx.Err()
|
|
}
|
|
return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index), Warnings: []contracts.Warning{{Scope: fmt.Sprintf("lane-%d/chunk-%d", lane, request.Chunk.Index), ReasonCode: "observed", Message: "ordered"}}}, nil
|
|
})
|
|
}
|
|
|
|
type outcome struct {
|
|
output RunOutput
|
|
err error
|
|
}
|
|
done := make(chan outcome, 1)
|
|
go func() {
|
|
output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
done <- outcome{output: output, err: err}
|
|
}()
|
|
for pair := 0; pair < 3; pair++ {
|
|
first, second := <-started, <-started
|
|
want := map[int]bool{pair * 2: true, pair*2 + 1: true}
|
|
if !want[first] || !want[second] || first == second {
|
|
t.Fatalf("started jobs = %d, %d; want round-robin pair %d", first, second, pair)
|
|
}
|
|
close(release[second])
|
|
close(release[first])
|
|
}
|
|
result := <-done
|
|
if result.err != nil {
|
|
t.Fatalf("Run() error = %v", result.err)
|
|
}
|
|
if got := maximum.Load(); got != 2 {
|
|
t.Fatalf("maximum concurrent extract jobs = %d, want 2", got)
|
|
}
|
|
wantScopes := []string{"lane-0/chunk-0", "lane-0/chunk-1", "lane-0/chunk-2", "lane-1/chunk-0", "lane-1/chunk-1", "lane-1/chunk-2"}
|
|
gotScopes := make([]string, len(result.output.Warnings))
|
|
for i := range result.output.Warnings {
|
|
gotScopes[i] = result.output.Warnings[i].Scope
|
|
}
|
|
if !reflect.DeepEqual(gotScopes, wantScopes) {
|
|
t.Fatalf("warning order = %#v, want %#v", gotScopes, wantScopes)
|
|
}
|
|
}
|
|
|
|
func TestRunnerStartsLaneContinuationWhileOtherLaneExtractsRemain(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 2)
|
|
releaseOtherLane := make(chan struct{})
|
|
mergeStarted := make(chan struct{}, 1)
|
|
installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
return erasedTypedResult{Value: typedValueForLane(0, request.Chunk.Index)}, nil
|
|
})
|
|
installExtractOperation(prepared, 1, func(ctx context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
select {
|
|
case <-releaseOtherLane:
|
|
return erasedTypedResult{Value: typedValueForLane(1, request.Chunk.Index)}, nil
|
|
case <-ctx.Done():
|
|
return erasedTypedResult{}, ctx.Err()
|
|
}
|
|
})
|
|
originalMerge := prepared.lanes[0].typed.merge
|
|
prepared.lanes[0].typed.merge = func(ctx context.Context, implementation any, request contracts.TypedMergeRequest[any]) (erasedTypedResult, error) {
|
|
select {
|
|
case mergeStarted <- struct{}{}:
|
|
default:
|
|
}
|
|
return originalMerge(ctx, implementation, request)
|
|
}
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
done <- err
|
|
}()
|
|
select {
|
|
case <-mergeStarted:
|
|
close(releaseOtherLane)
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("first lane merge did not start while the second lane remained in extract")
|
|
}
|
|
if err := <-done; err != nil {
|
|
t.Fatalf("Run() error = %v", err)
|
|
}
|
|
}
|
|
|
|
func TestRunnerBoundsConcurrentLaneContinuations(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 1)
|
|
base := prepared.lanes[0]
|
|
lanes := make([]preparedLaneExecutor, 8)
|
|
resolvedLanes := make([]ResolvedArtifactLane, len(lanes))
|
|
publicLanes := make([]PreparedArtifactLane, len(lanes))
|
|
var active atomic.Int32
|
|
var maximum atomic.Int32
|
|
for i := range lanes {
|
|
lane := base
|
|
lane.resolved.ID = fmt.Sprintf("lane-%02d", i)
|
|
lane.typed = &preparedTypedLane{
|
|
extractor: base.typed.extractor,
|
|
merger: base.typed.merger,
|
|
normalizer: base.typed.normalizer,
|
|
extract: func(_ context.Context, _ any, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
return erasedTypedResult{Value: codecNotes{Items: []string{request.Chunk.ID}}}, nil
|
|
},
|
|
merge: func(_ context.Context, _ any, _ contracts.TypedMergeRequest[any]) (erasedTypedResult, error) {
|
|
current := active.Add(1)
|
|
defer active.Add(-1)
|
|
for {
|
|
seen := maximum.Load()
|
|
if current <= seen || maximum.CompareAndSwap(seen, current) {
|
|
break
|
|
}
|
|
}
|
|
time.Sleep(5 * time.Millisecond)
|
|
return erasedTypedResult{Value: codecNotes{}}, nil
|
|
},
|
|
normalize: base.typed.normalize,
|
|
codec: base.typed.codec,
|
|
}
|
|
lanes[i] = lane
|
|
resolvedLanes[i] = lane.resolved
|
|
publicLanes[i] = PreparedArtifactLane{Resolved: lane.resolved}
|
|
}
|
|
prepared.lanes = lanes
|
|
prepared.resolved.Steps[0].ArtifactLanes = resolvedLanes
|
|
prepared.ArtifactLanes = publicLanes
|
|
|
|
output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
if err != nil {
|
|
t.Fatalf("Run() error = %v", err)
|
|
}
|
|
if len(output.NormalizeOutputs) != len(lanes) {
|
|
t.Fatalf("normalize outputs = %d, want %d", len(output.NormalizeOutputs), len(lanes))
|
|
}
|
|
if got := maximum.Load(); got != 2 {
|
|
t.Fatalf("maximum concurrent lane continuations = %d, want 2", got)
|
|
}
|
|
}
|
|
|
|
func TestRunnerSelectsFrameworkErrorByStableLaneOrder(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 1)
|
|
ready := make(chan struct{}, 2)
|
|
release := make(chan struct{})
|
|
for laneIndex := range prepared.lanes {
|
|
lane := laneIndex
|
|
installExtractOperation(prepared, lane, func(context.Context, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
ready <- struct{}{}
|
|
<-release
|
|
return erasedTypedResult{}, fmt.Errorf("lane-%d failure", lane)
|
|
})
|
|
}
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
done <- err
|
|
}()
|
|
<-ready
|
|
<-ready
|
|
close(release)
|
|
if err := <-done; err == nil || !strings.Contains(err.Error(), "lane-0 failure") {
|
|
t.Fatalf("Run() error = %v, want first resolved lane failure", err)
|
|
}
|
|
}
|
|
|
|
func TestRunnerSelectsMergeErrorBeforeEarlierLaneNormalizeError(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 1)
|
|
for laneIndex := range prepared.lanes {
|
|
lane := laneIndex
|
|
installExtractOperation(prepared, lane, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index)}, nil
|
|
})
|
|
}
|
|
ready := make(chan struct{}, 2)
|
|
release := make(chan struct{})
|
|
prepared.lanes[0].typed.normalize = func(context.Context, any, contracts.TypedNormalizeRequest[any]) (erasedTypedResult, error) {
|
|
ready <- struct{}{}
|
|
<-release
|
|
return erasedTypedResult{}, errors.New("earlier lane normalize failure")
|
|
}
|
|
prepared.lanes[1].typed.merge = func(context.Context, any, contracts.TypedMergeRequest[any]) (erasedTypedResult, error) {
|
|
ready <- struct{}{}
|
|
<-release
|
|
return erasedTypedResult{}, errors.New("later lane merge failure")
|
|
}
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
done <- err
|
|
}()
|
|
<-ready
|
|
<-ready
|
|
close(release)
|
|
if err := <-done; err == nil || !strings.Contains(err.Error(), "later lane merge failure") {
|
|
t.Fatalf("Run() error = %v, want merge failure before normalize failure", err)
|
|
}
|
|
}
|
|
|
|
func TestRunnerKeepsSuccessfulChunksWhenAnotherExtractIsRejected(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 2)
|
|
for laneIndex := range prepared.lanes {
|
|
lane := laneIndex
|
|
installExtractOperation(prepared, lane, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index)}, nil
|
|
})
|
|
}
|
|
validator := &prepared.lanes[0].extractValidators.validators[0]
|
|
validator.typedValidate = func(_ context.Context, _ any, target typedValidationTarget) (contracts.ValidationResult, error) {
|
|
if target.chunk != nil && target.chunk.Index == 0 {
|
|
return contracts.ValidationResult{Approved: false, ReasonCode: "expected_rejection", Message: "rejected by test"}, nil
|
|
}
|
|
return contracts.ValidationResult{Approved: true}, nil
|
|
}
|
|
output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
if err != nil {
|
|
t.Fatalf("Run() error = %v", err)
|
|
}
|
|
if len(output.Rejected) != 1 || output.Rejected[0].LaneID != prepared.lanes[0].resolved.ID || output.Rejected[0].ChunkIndex != 0 {
|
|
t.Fatalf("rejections = %#v, want the first lane's first chunk", output.Rejected)
|
|
}
|
|
if len(output.NormalizeOutputs) != 2 {
|
|
t.Fatalf("normalize outputs = %d, want both lanes to continue", len(output.NormalizeOutputs))
|
|
}
|
|
if output.Manifest.ValidationStatus != "rejected" {
|
|
t.Fatalf("validation status = %q, want rejected", output.Manifest.ValidationStatus)
|
|
}
|
|
}
|
|
|
|
func TestRunnerSeparatelyBoundsExtractJobsAndSharedProviderCalls(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 3)
|
|
provider := &countingProvider{}
|
|
scheduler, err := frameworkllm.NewScheduler(2)
|
|
if err != nil {
|
|
t.Fatalf("NewScheduler() error = %v", err)
|
|
}
|
|
client := frameworkllm.NewScheduledClient(provider, scheduler)
|
|
var jobActive atomic.Int32
|
|
var jobMaximum atomic.Int32
|
|
var attempts sync.Map
|
|
for laneIndex := range prepared.lanes {
|
|
lane := laneIndex
|
|
prepared.lanes[lane].resolved.Extract.Retries = 1
|
|
installExtractOperation(prepared, lane, func(ctx context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
current := jobActive.Add(1)
|
|
defer jobActive.Add(-1)
|
|
for {
|
|
seen := jobMaximum.Load()
|
|
if current <= seen || jobMaximum.CompareAndSwap(seen, current) {
|
|
break
|
|
}
|
|
}
|
|
if _, callErr := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: "extract"}, nil); callErr != nil {
|
|
return erasedTypedResult{}, callErr
|
|
}
|
|
key := fmt.Sprintf("%d/%d", lane, request.Chunk.Index)
|
|
countValue, _ := attempts.LoadOrStore(key, &atomic.Int32{})
|
|
count := countValue.(*atomic.Int32).Add(1)
|
|
if lane == 0 && request.Chunk.Index == 0 && count == 1 {
|
|
return erasedTypedResult{}, errors.New("retry extract")
|
|
}
|
|
return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index)}, nil
|
|
})
|
|
validator := &prepared.lanes[lane].extractValidators.validators[0]
|
|
validator.typedValidate = func(ctx context.Context, _ any, _ typedValidationTarget) (contracts.ValidationResult, error) {
|
|
if _, callErr := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: "validate"}, nil); callErr != nil {
|
|
return contracts.ValidationResult{}, callErr
|
|
}
|
|
return contracts.ValidationResult{Approved: true}, nil
|
|
}
|
|
}
|
|
output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 3})
|
|
if err != nil {
|
|
t.Fatalf("Run() error = %v", err)
|
|
}
|
|
if len(output.NormalizeOutputs) != 2 {
|
|
t.Fatalf("normalize outputs = %d, want 2", len(output.NormalizeOutputs))
|
|
}
|
|
if got := jobMaximum.Load(); got != 3 {
|
|
t.Fatalf("maximum extract jobs = %d, want 3", got)
|
|
}
|
|
if got := provider.maximum.Load(); got != 2 {
|
|
t.Fatalf("maximum provider calls = %d, want 2", got)
|
|
}
|
|
if got := provider.calls.Load(); got != 13 {
|
|
t.Fatalf("provider calls = %d, want 13 including retry and validators", got)
|
|
}
|
|
}
|
|
|
|
func TestRunnerReturnsParentCancellationAndStopsQueuedExtracts(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 4)
|
|
started := make(chan struct{}, 8)
|
|
for laneIndex := range prepared.lanes {
|
|
lane := laneIndex
|
|
installExtractOperation(prepared, lane, func(ctx context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
started <- struct{}{}
|
|
<-ctx.Done()
|
|
return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index)}, ctx.Err()
|
|
})
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
_, err := New().Run(ctx, RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
done <- err
|
|
}()
|
|
<-started
|
|
<-started
|
|
cancel()
|
|
if err := <-done; !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("Run() error = %v, want context.Canceled", err)
|
|
}
|
|
if got := len(started); got != 0 {
|
|
t.Fatalf("additional started jobs = %d, want none after cancellation", got)
|
|
}
|
|
}
|