Finalize bounded typed pipeline implementation

2026-07-17 08:49:08 +00:00
parent adfd3bd052
commit 3013ee044d
11 changed files with 369 additions and 857 deletions
--- a/internal/framework/pipeline/runner_concurrency_test.go
+++ b/internal/framework/pipeline/runner_concurrency_test.go
@@ -178,6 +178,59 @@ func TestRunnerStartsLaneContinuationWhileOtherLaneExtractsRemain(t *testing.T)
 	}
 }

+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.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)
--- a/internal/framework/pipeline/runner_concurrent.go
+++ b/internal/framework/pipeline/runner_concurrent.go
@@ -98,6 +98,7 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints Ch
 	jobs := make(chan extractJob, workerCount)
 	results := make(chan extractJobResult, workerCount)
 	completions := make(chan laneCompletion, len(states))
+	continuations := make(chan *laneExtractState, workerCount)

 	var workers sync.WaitGroup
 	for i := 0; i < workerCount; i++ {
@@ -130,26 +131,46 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints Ch
 		}
 	}()
 	go func() { workers.Wait(); close(results) }()
+	var continuationWorkers sync.WaitGroup
+	for i := 0; i < workerCount; i++ {
+		continuationWorkers.Add(1)
+		go func() {
+			defer continuationWorkers.Done()
+			for state := range continuations {
+				if err := ctx.Err(); err != nil {
+					completions <- laneCompletion{index: state.index, err: err}
+					continue
+				}
+				laneOutput, err := r.continueLane(ctx, input, checkpoints, loader, doc, sourceInput, sessionID, chunks, state)
+				completions <- laneCompletion{index: state.index, output: laneOutput, err: err}
+			}
+		}()
+	}

 	completedOutputs := make([]RunOutput, len(states))
 	var runErrors []orderedRunError
+	var pendingContinuations []*laneExtractState
 	launched, completed := 0, 0
-	launch := func(state *laneExtractState) {
-		launched++
-		go func() {
-			laneOutput, err := r.continueLane(ctx, input, checkpoints, loader, doc, sourceInput, sessionID, chunks, state)
-			completions <- laneCompletion{index: state.index, output: laneOutput, err: err}
-		}()
-	}
 	for _, state := range states {
 		if state.decision.Reused {
-			launch(state)
+			pendingContinuations = append(pendingContinuations, state)
 		}
 	}

 	resultChannel := results
-	for resultChannel != nil || completed < launched {
+	for resultChannel != nil || len(pendingContinuations) > 0 || completed < launched {
+		var continuationChannel chan<- *laneExtractState
+		var nextContinuation *laneExtractState
+		if len(pendingContinuations) > 0 && ctx.Err() == nil {
+			continuationChannel = continuations
+			nextContinuation = pendingContinuations[0]
+		} else if ctx.Err() != nil {
+			pendingContinuations = nil
+		}
 		select {
+		case continuationChannel <- nextContinuation:
+			pendingContinuations = pendingContinuations[1:]
+			launched++
 		case result, ok := <-resultChannel:
 			if !ok {
 				resultChannel = nil
@@ -171,7 +192,7 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints Ch
 					runErrors = append(runErrors, orderedRunError{stage: 0, lane: state.index, chunk: len(chunks), err: err})
 					cancel()
 				} else {
-					launch(state)
+					pendingContinuations = append(pendingContinuations, state)
 				}
 			}
 		case completion := <-completions:
@@ -183,6 +204,8 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints Ch
 			}
 		}
 	}
+	close(continuations)
+	continuationWorkers.Wait()
 	for i := range completedOutputs {
 		mergeLaneOutput(&output, completedOutputs[i])
 	}
--- a/internal/modules/integration/concurrent_runner_test.go
+++ b/internal/modules/integration/concurrent_runner_test.go
@@ -0,0 +1,177 @@
+package integration_test
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"gitea.maximumdirect.net/eric/notarius/internal/core/config"
+	"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"
+	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
+	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
+)
+
+const (
+	concurrentChunkerKey   = "test/concurrent-chunks"
+	concurrentExtractorKey = "test/concurrent-extractor"
+	concurrentValidatorKey = "test/concurrent-validator"
+)
+
+type integrationConcurrencyTracker struct {
+	extractActive   atomic.Int32
+	extractMaximum  atomic.Int32
+	providerActive  atomic.Int32
+	providerMaximum atomic.Int32
+	providerCalls   atomic.Int32
+	validatorCalls  atomic.Int32
+}
+
+func updateMaximum(maximum *atomic.Int32, current int32) {
+	for {
+		seen := maximum.Load()
+		if current <= seen || maximum.CompareAndSwap(seen, current) {
+			return
+		}
+	}
+}
+
+type instrumentedProvider struct {
+	tracker *integrationConcurrencyTracker
+}
+
+func (p instrumentedProvider) CompleteStructured(ctx context.Context, _ contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
+	p.tracker.providerCalls.Add(1)
+	current := p.tracker.providerActive.Add(1)
+	defer p.tracker.providerActive.Add(-1)
+	updateMaximum(&p.tracker.providerMaximum, current)
+	select {
+	case <-time.After(5 * time.Millisecond):
+	case <-ctx.Done():
+		return contracts.StructuredCompletionResponse{}, ctx.Err()
+	}
+	content := []byte(`{"approved":true}`)
+	if out != nil {
+		if err := json.Unmarshal(content, out); err != nil {
+			return contracts.StructuredCompletionResponse{}, err
+		}
+	}
+	return contracts.StructuredCompletionResponse{Content: content}, nil
+}
+
+type concurrentChunker struct{}
+
+func (concurrentChunker) Key() string                               { return concurrentChunkerKey }
+func (concurrentChunker) ReferenceSlots() []contracts.ReferenceSlot { return nil }
+func (concurrentChunker) Chunk(_ context.Context, request contracts.ChunkRequest) (contracts.ChunkResult, error) {
+	chunks := make([]source.Chunk, len(request.Source.Units))
+	for i, unit := range request.Source.Units {
+		chunks[i] = source.Chunk{ID: fmt.Sprintf("%s:chunk:%d", request.Source.ID, i), SourceID: request.Source.ID, Index: i, Ref: unit.Ref, Content: []byte(fmt.Sprintf(`{"unit":%d}`, unit.ID)), MediaType: "application/json", Units: []source.SourceUnit{unit}}
+	}
+	return contracts.ChunkResult{Chunks: chunks}, nil
+}
+
+type concurrentExtractor struct {
+	client  contracts.StructuredLLMClient
+	tracker *integrationConcurrencyTracker
+	failed  atomic.Bool
+}
+
+func (*concurrentExtractor) Key() string                               { return concurrentExtractorKey }
+func (*concurrentExtractor) ReferenceSlots() []contracts.ReferenceSlot { return nil }
+func (e *concurrentExtractor) Extract(ctx context.Context, request contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[dnd.SpellList], error) {
+	current := e.tracker.extractActive.Add(1)
+	defer e.tracker.extractActive.Add(-1)
+	updateMaximum(&e.tracker.extractMaximum, current)
+	var response map[string]any
+	if _, err := e.client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: concurrentExtractorKey}, &response); err != nil {
+		return contracts.TypedExtractionResult[dnd.SpellList]{}, err
+	}
+	if e.failed.CompareAndSwap(false, true) {
+		return contracts.TypedExtractionResult[dnd.SpellList]{}, errors.New("retry requested")
+	}
+	return contracts.TypedExtractionResult[dnd.SpellList]{Value: dnd.SpellList{SpellCasts: []dnd.SpellCast{}}}, nil
+}
+
+type concurrentValidator struct {
+	client  contracts.StructuredLLMClient
+	tracker *integrationConcurrencyTracker
+}
+
+func (*concurrentValidator) Name() string { return concurrentValidatorKey }
+func (*concurrentValidator) ExecutionClass() contracts.ExecutionClass {
+	return contracts.ExecutionClassLLMBacked
+}
+func (v *concurrentValidator) Validate(ctx context.Context, _ contracts.TypedValidationRequest[dnd.SpellList]) (contracts.ValidationResult, error) {
+	v.tracker.validatorCalls.Add(1)
+	var response map[string]any
+	if _, err := v.client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: concurrentValidatorKey}, &response); err != nil {
+		return contracts.ValidationResult{}, err
+	}
+	return contracts.ValidationResult{Approved: true}, nil
+}
+
+func TestRunnerIndependentlyBoundsWorkersAndProviderCallsAcrossRegisteredModules(t *testing.T) {
+	tracker := &integrationConcurrencyTracker{}
+	scheduler, err := frameworkllm.NewScheduler(2)
+	if err != nil {
+		t.Fatalf("NewScheduler() error = %v", err)
+	}
+	client := frameworkllm.NewScheduledClient(instrumentedProvider{tracker: tracker}, scheduler)
+	catalog := dndSpellsTestCatalog(t, dndSpellsCatalogSpecs{})
+	if err := catalog.Chunkers.RegisterWithSpec(pipeline.ModuleSpec{Key: concurrentChunkerKey, Stage: pipeline.StageChunk, Requires: []string{"source.transcript"}, Provides: []string{"chunks"}}, func() (contracts.Chunker, error) { return concurrentChunker{}, nil }); err != nil {
+		t.Fatalf("register chunker: %v", err)
+	}
+	validateOptions := func(options map[string]any) error { return pipeline.RejectUnknownOptions(options) }
+	if err := pipeline.RegisterExtractorBuilder(catalog.Extractors, pipeline.ModuleSpec{Key: concurrentExtractorKey, Stage: pipeline.StageExtract, ArtifactKind: dnd.SpellListKind, Requires: []string{"chunks", "source.transcript"}, Provides: []string{"dnd.spell_casts"}}, validateOptions, func(request pipeline.BuildRequest) (contracts.Extractor[dnd.SpellList], error) {
+		return &concurrentExtractor{client: request.Dependencies.LLM, tracker: tracker}, nil
+	}); err != nil {
+		t.Fatalf("register extractor: %v", err)
+	}
+	validators := pipeline.NewValidatorRegistry()
+	catalog.Validators = validators
+	if err := pipeline.RegisterTypedValidatorBuilder(validators, dnd.SpellListKind, pipeline.ValidatorSpec{Key: concurrentValidatorKey, ExecutionClass: contracts.ExecutionClassLLMBacked}, validateOptions, func(request pipeline.BuildRequest) (contracts.TypedValidator[dnd.SpellList], error) {
+		return &concurrentValidator{client: request.Dependencies.LLM, tracker: tracker}, nil
+	}); err != nil {
+		t.Fatalf("register validator: %v", err)
+	}
+
+	cfg := loadDNDSpellsPipelineConfig(t)
+	profile := cfg.Pipelines["dnd-spells-fixture"]
+	profile.Chunk = pipeline.Binding(concurrentChunkerKey)
+	validatorOverride := pipeline.ValidatorOverride{Set: true, Validators: []pipeline.ModuleBinding{pipeline.Binding(concurrentValidatorKey)}}
+	profile.Artifacts = map[string]pipeline.ArtifactLaneProfile{
+		"alpha": {Extract: pipeline.ModuleBinding{Module: concurrentExtractorKey, Retries: 1, Validators: validatorOverride}, Merge: pipeline.Binding(pipeline.DefaultMergeModule), Normalize: pipeline.Binding(pipeline.DefaultNormalizeModule)},
+		"beta":  {Extract: pipeline.ModuleBinding{Module: concurrentExtractorKey, Retries: 1, Validators: validatorOverride}, Merge: pipeline.Binding(pipeline.DefaultMergeModule), Normalize: pipeline.Binding(pipeline.DefaultNormalizeModule)},
+	}
+	cfg.Pipelines["dnd-spells-fixture"] = profile
+	resolved, err := cfg.Resolve(config.ResolveInput{PipelineID: "dnd-spells-fixture", Catalog: catalog})
+	if err != nil {
+		t.Fatalf("Resolve() error = %v", err)
+	}
+	registries := pipeline.Registries{Inputs: catalog.Inputs, Chunkers: catalog.Chunkers, ArtifactCodecs: catalog.ArtifactCodecs, Extractors: catalog.Extractors, Mergers: catalog.Mergers, Normalizers: catalog.Normalizers, Validators: catalog.Validators, ValidatorChains: catalog.ValidatorChains, Outputs: catalog.Outputs}
+	output, err := runPreparedPipeline(t, registries, resolved.ResolvedPipeline, client, pipeline.RunInput{RawInput: readDNDSpellsFixture(t), 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 := tracker.extractMaximum.Load(); got != 3 {
+		t.Fatalf("maximum extract jobs = %d, want 3", got)
+	}
+	if got := tracker.providerMaximum.Load(); got != 2 {
+		t.Fatalf("maximum provider calls = %d, want 2", got)
+	}
+	if got := tracker.validatorCalls.Load(); got == 0 {
+		t.Fatal("LLM-backed validator was not called")
+	}
+	if got := tracker.providerCalls.Load(); got <= tracker.validatorCalls.Load() {
+		t.Fatalf("provider calls = %d, validator calls = %d; want extractor, retry, and validator calls", got, tracker.validatorCalls.Load())
+	}
+}