Finalize bounded typed pipeline implementation
This commit is contained in:
177
internal/modules/integration/concurrent_runner_test.go
Normal file
177
internal/modules/integration/concurrent_runner_test.go
Normal file
@@ -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())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user