Add bounded concurrent pipeline execution
This commit is contained in:
355
internal/framework/pipeline/runner_concurrency_test.go
Normal file
355
internal/framework/pipeline/runner_concurrency_test.go
Normal file
@@ -0,0 +1,355 @@
|
||||
package pipeline
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
||||
"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 := typedTestDocument()
|
||||
chunks := make([]source.Chunk, chunkCount)
|
||||
for i := range chunks {
|
||||
chunks[i] = source.Chunk{ID: fmt.Sprintf("chunk-%d", i+1), SourceID: doc.ID, Index: i, Ref: doc.Units[0].Ref, Content: []byte(fmt.Sprintf(`{"chunk":%d}`, i+1)), MediaType: "application/json", Units: []source.SourceUnit{doc.Units[0]}}
|
||||
}
|
||||
prepared.chunker = &typedTestChunker{key: "typed/chunk", chunks: chunks}
|
||||
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 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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user