Add bounded concurrent pipeline execution
This commit is contained in:
@@ -202,7 +202,7 @@ func wrapDebugLLMClient(client contracts.StructuredLLMClient, recorder DebugReco
|
||||
if _, ok := client.(*debugLLMClient); ok {
|
||||
return client
|
||||
}
|
||||
return &debugLLMClient{inner: client, recorder: recorder}
|
||||
return &debugLLMClient{inner: client, recorder: synchronizedDebugRecorder(recorder)}
|
||||
}
|
||||
|
||||
// WithDebugLLMRecording decorates a shared LLM client so calls made by
|
||||
|
||||
@@ -50,9 +50,13 @@ type RunInput struct {
|
||||
Checkpoints CheckpointRecorder
|
||||
Checkpoint CheckpointLoader
|
||||
Debug DebugRecorder
|
||||
// ExtractWorkers bounds run-wide extract jobs. Values less than one use a
|
||||
// single worker so direct framework callers retain deterministic behavior.
|
||||
ExtractWorkers int
|
||||
|
||||
pipeline ResolvedPipeline
|
||||
llmClient contracts.StructuredLLMClient
|
||||
pipeline ResolvedPipeline
|
||||
llmClient contracts.StructuredLLMClient
|
||||
extractDecision *CheckpointDecision
|
||||
}
|
||||
|
||||
type RunOutput struct {
|
||||
@@ -87,6 +91,9 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
|
||||
if debugRecorder == nil {
|
||||
debugRecorder = NoopDebugRecorder()
|
||||
}
|
||||
checkpoints = synchronizedCheckpointRecorder(checkpoints)
|
||||
checkpointLoader = synchronizedCheckpointLoader(checkpointLoader)
|
||||
debugRecorder = synchronizedDebugRecorder(debugRecorder)
|
||||
input.Debug = debugRecorder
|
||||
input.llmClient = wrapDebugLLMClient(input.llmClient, debugRecorder)
|
||||
defer func() {
|
||||
@@ -311,10 +318,10 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
|
||||
}
|
||||
|
||||
if chunksAccepted {
|
||||
for _, lane := range input.Prepared.lanes {
|
||||
if err := r.runLane(ctx, input, checkpoints, checkpointLoader, doc, sourceInput, sessionID, canonicalChunks, lane, &output); err != nil {
|
||||
return failOutput(output), err
|
||||
}
|
||||
laneOutput, laneErr := r.runLanes(ctx, input, checkpoints, checkpointLoader, doc, sourceInput, sessionID, canonicalChunks)
|
||||
mergeLaneOutput(&output, laneOutput)
|
||||
if laneErr != nil {
|
||||
return failOutput(output), laneErr
|
||||
}
|
||||
}
|
||||
|
||||
@@ -376,10 +383,6 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
|
||||
return output, nil
|
||||
}
|
||||
|
||||
func (r *Runner) runLane(ctx context.Context, input RunInput, checkpoints CheckpointRecorder, checkpointLoader CheckpointLoader, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, chunks []source.Chunk, prepared preparedLaneExecutor, output *RunOutput) error {
|
||||
return r.runTypedLane(ctx, input, checkpoints, checkpointLoader, doc, sourceInput, sessionID, chunks, prepared, output)
|
||||
}
|
||||
|
||||
func runWithRetry(ctx context.Context, retries int, run func(attempt int) (bool, *contracts.RejectedOutput, error)) (bool, *contracts.RejectedOutput, error) {
|
||||
attempts := retries + 1
|
||||
var last *contracts.RejectedOutput
|
||||
@@ -389,12 +392,12 @@ func runWithRetry(ctx context.Context, retries int, run func(attempt int) (bool,
|
||||
}
|
||||
accepted, rejection, err := run(attempt)
|
||||
if err != nil {
|
||||
if ctxErr := ctx.Err(); ctxErr != nil {
|
||||
return false, nil, ctxErr
|
||||
}
|
||||
if attempt == attempts {
|
||||
return false, nil, fmt.Errorf("failed after %d attempt(s): %w", attempt, err)
|
||||
}
|
||||
if ctxErr := ctx.Err(); ctxErr != nil {
|
||||
return false, nil, ctxErr
|
||||
}
|
||||
continue
|
||||
}
|
||||
if accepted {
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
376
internal/framework/pipeline/runner_concurrent.go
Normal file
376
internal/framework/pipeline/runner_concurrent.go
Normal file
@@ -0,0 +1,376 @@
|
||||
package pipeline
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"path"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||
)
|
||||
|
||||
type laneExtractState struct {
|
||||
index int
|
||||
prepared preparedLaneExecutor
|
||||
deps []CheckpointFingerprint
|
||||
decision CheckpointDecision
|
||||
values []erasedExtractArtifact
|
||||
serialized []CheckpointArtifact
|
||||
warnings []contracts.Warning
|
||||
rejected []contracts.RejectedOutput
|
||||
results map[int]extractJobResult
|
||||
remaining int
|
||||
failed bool
|
||||
}
|
||||
|
||||
type extractJob struct {
|
||||
lane *laneExtractState
|
||||
chunk source.Chunk
|
||||
}
|
||||
|
||||
type extractJobResult struct {
|
||||
laneIndex int
|
||||
chunkIndex int
|
||||
value erasedExtractArtifact
|
||||
serialized CheckpointArtifact
|
||||
warnings []contracts.Warning
|
||||
rejected *contracts.RejectedOutput
|
||||
err error
|
||||
}
|
||||
|
||||
type laneCompletion struct {
|
||||
index int
|
||||
output RunOutput
|
||||
err error
|
||||
}
|
||||
|
||||
type orderedRunError struct {
|
||||
stage int
|
||||
lane int
|
||||
chunk int
|
||||
err error
|
||||
}
|
||||
|
||||
type completedExtractLoader struct {
|
||||
CheckpointLoader
|
||||
laneID string
|
||||
checkpoint ExtractCheckpoint
|
||||
}
|
||||
|
||||
func (l completedExtractLoader) Extract(laneID, _ string, _ []CheckpointFingerprint) (ExtractCheckpoint, CheckpointDecision) {
|
||||
if laneID == l.laneID {
|
||||
return l.checkpoint, CheckpointDecision{Reused: true, Reason: "coordinated extract result"}
|
||||
}
|
||||
return ExtractCheckpoint{}, CheckpointDecision{Reason: "extract result unavailable"}
|
||||
}
|
||||
|
||||
func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints CheckpointRecorder, loader CheckpointLoader, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, chunks []source.Chunk) (RunOutput, error) {
|
||||
output := RunOutput{Manifest: manifestFromPipeline(input)}
|
||||
states := make([]*laneExtractState, len(input.Prepared.lanes))
|
||||
for i, prepared := range input.Prepared.lanes {
|
||||
if prepared.typed == nil {
|
||||
return output, fmt.Errorf("typed lane %q executor is not prepared", prepared.resolved.ID)
|
||||
}
|
||||
setTypedLaneManifestMetadata(&output, prepared.resolved.ID, prepared.typed.extractor, prepared.typed.merger, prepared.typed.normalizer)
|
||||
state, err := prepareLaneExtract(input, loader, doc, chunks, i, prepared)
|
||||
if err != nil {
|
||||
return output, err
|
||||
}
|
||||
if !state.decision.Reused {
|
||||
if err := checkpoints.ExtractRunning(prepared.resolved.ID, prepared.resolved.Extract.Module, state.deps); err != nil {
|
||||
return output, fmt.Errorf("write extract checkpoint for lane %q: %w", prepared.resolved.ID, err)
|
||||
}
|
||||
}
|
||||
states[i] = state
|
||||
}
|
||||
|
||||
workerCount := input.ExtractWorkers
|
||||
if workerCount < 1 {
|
||||
workerCount = 1
|
||||
}
|
||||
ctx, cancel := context.WithCancel(parent)
|
||||
defer cancel()
|
||||
jobs := make(chan extractJob, workerCount)
|
||||
results := make(chan extractJobResult, workerCount)
|
||||
completions := make(chan laneCompletion, len(states))
|
||||
|
||||
var workers sync.WaitGroup
|
||||
for i := 0; i < workerCount; i++ {
|
||||
workers.Add(1)
|
||||
go func() {
|
||||
defer workers.Done()
|
||||
for job := range jobs {
|
||||
if ctx.Err() != nil {
|
||||
continue
|
||||
}
|
||||
result := r.runExtractJob(ctx, input, doc, sourceInput, sessionID, job)
|
||||
results <- result
|
||||
}
|
||||
}()
|
||||
}
|
||||
go func() {
|
||||
defer close(jobs)
|
||||
for chunkIndex := range chunks {
|
||||
for laneIndex := range states {
|
||||
state := states[laneIndex]
|
||||
if state.decision.Reused {
|
||||
continue
|
||||
}
|
||||
select {
|
||||
case jobs <- extractJob{lane: state, chunk: chunks[chunkIndex]}:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
go func() { workers.Wait(); close(results) }()
|
||||
|
||||
completedOutputs := make([]RunOutput, len(states))
|
||||
var runErrors []orderedRunError
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
resultChannel := results
|
||||
for resultChannel != nil || completed < launched {
|
||||
select {
|
||||
case result, ok := <-resultChannel:
|
||||
if !ok {
|
||||
resultChannel = nil
|
||||
continue
|
||||
}
|
||||
state := states[result.laneIndex]
|
||||
state.remaining--
|
||||
if result.err != nil {
|
||||
state.failed = true
|
||||
runErrors = append(runErrors, orderedRunError{stage: 0, lane: state.index, chunk: result.chunkIndex, err: result.err})
|
||||
_ = checkpoints.ExtractFailed(state.prepared.resolved.ID, state.prepared.resolved.Extract.Module, state.deps, result.err)
|
||||
cancel()
|
||||
} else {
|
||||
state.results[result.chunkIndex] = result
|
||||
}
|
||||
if state.remaining == 0 && !state.failed && ctx.Err() == nil {
|
||||
if err := finalizeLaneExtract(checkpoints, state); err != nil {
|
||||
state.failed = true
|
||||
runErrors = append(runErrors, orderedRunError{stage: 0, lane: state.index, chunk: len(chunks), err: err})
|
||||
cancel()
|
||||
} else {
|
||||
launch(state)
|
||||
}
|
||||
}
|
||||
case completion := <-completions:
|
||||
completed++
|
||||
completedOutputs[completion.index] = completion.output
|
||||
if completion.err != nil {
|
||||
runErrors = append(runErrors, classifyLaneError(completion.index, len(chunks), completion.err))
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
}
|
||||
for i := range completedOutputs {
|
||||
mergeLaneOutput(&output, completedOutputs[i])
|
||||
}
|
||||
if err := selectRunError(parent, runErrors); err != nil {
|
||||
return output, err
|
||||
}
|
||||
return output, nil
|
||||
}
|
||||
|
||||
func prepareLaneExtract(input RunInput, loader CheckpointLoader, doc *source.SourceDocument, chunks []source.Chunk, index int, prepared preparedLaneExecutor) (*laneExtractState, error) {
|
||||
lane, typed := prepared.resolved, prepared.typed
|
||||
digest, err := joinedChunkDigest(chunks)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("digest chunks for lane %q: %w", lane.ID, err)
|
||||
}
|
||||
state := &laneExtractState{index: index, prepared: prepared, deps: digestFingerprints("chunks", digest), remaining: len(chunks), results: make(map[int]extractJobResult, len(chunks))}
|
||||
cp, decision := loadExtract(loader, lane.ID, lane.Extract.Module, state.deps)
|
||||
if decision.Reused {
|
||||
for _, stored := range cp.Outputs {
|
||||
if _, decodeErr := decodeCheckpointArtifact(typed.codec, stored); decodeErr != nil {
|
||||
decision = CheckpointDecision{Reason: "extract artifact checkpoint codec is incompatible: " + decodeErr.Error()}
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
state.decision = decision
|
||||
if decision.Reused {
|
||||
state.remaining = 0
|
||||
for _, stored := range cp.Outputs {
|
||||
value, decodeErr := decodeCheckpointArtifact(typed.codec, stored)
|
||||
if decodeErr != nil {
|
||||
return nil, fmt.Errorf("decode extract checkpoint for lane %q: %w", lane.ID, decodeErr)
|
||||
}
|
||||
stored = hydrateCheckpointArtifact(typed.codec, stored, value)
|
||||
artifact := erasedExtractArtifact{LaneID: lane.ID, ExtractorKey: lane.Extract.Module, SourceID: doc.ID, ChunkID: stored.ChunkID, ChunkIndex: stored.ChunkIndex, ChunkRef: stored.ChunkRef, Value: value}
|
||||
if stored.ChunkIndex >= 0 && stored.ChunkIndex < len(chunks) && artifact.ChunkRef == (source.SourceRef{}) {
|
||||
artifact.ChunkRef = chunks[stored.ChunkIndex].Ref
|
||||
}
|
||||
state.values = append(state.values, artifact)
|
||||
state.serialized = append(state.serialized, cloneCheckpointArtifact(stored))
|
||||
}
|
||||
state.warnings, state.rejected = cloneWarnings(cp.Warnings), cloneRejectedOutputs(cp.Rejected)
|
||||
}
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (r *Runner) runExtractJob(ctx context.Context, input RunInput, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, job extractJob) extractJobResult {
|
||||
state, chunk := job.lane, job.chunk
|
||||
lane, typed := state.prepared.resolved, state.prepared.typed
|
||||
result := extractJobResult{laneIndex: state.index, chunkIndex: chunk.Index}
|
||||
var accepted erasedExtractArtifact
|
||||
var serialized CheckpointArtifact
|
||||
var acceptedWarnings []contracts.Warning
|
||||
ok, rejection, err := runWithRetry(ctx, lane.Extract.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
|
||||
started := time.Now().UTC()
|
||||
attemptPath := path.Join("extract", debugPathComponent(lane.ID), fmt.Sprintf("chunk-%06d", chunk.Index+1), fmt.Sprintf("attempt-%02d", attempt))
|
||||
attemptCtx, llmScope := withDebugLLMScope(ctx, attemptPath)
|
||||
extracted, callErr := typed.extract(attemptCtx, typed.extractor, contracts.TypedExtractionRequest{Source: doc, Chunk: &chunk, SourceInput: chunkInputMaterial(sourceInput, chunk), SessionID: sessionID, References: CloneReferenceSet(lane.ExtractReferences.ReferenceSet), LLMProfile: lane.Extract.LLMProfile, Metadata: cloneMetadata(input.Metadata)})
|
||||
if callErr != nil {
|
||||
_ = writeDebugTimed(input.Debug, attemptPath+".json", debugEnvelopeWithLLMCalls(debugTimedEnvelope{Stage: string(StageExtract), LaneID: lane.ID, ModuleKey: lane.Extract.Module, Attempt: attempt, StartedAt: started, Error: callErr.Error()}, llmScope))
|
||||
return false, nil, fmt.Errorf("extract lane %q chunk %q with extractor %q: %w", lane.ID, chunk.ID, lane.Extract.Module, callErr)
|
||||
}
|
||||
artifact := erasedExtractArtifact{LaneID: lane.ID, ExtractorKey: lane.Extract.Module, SourceID: doc.ID, ChunkID: chunk.ID, ChunkIndex: chunk.Index, ChunkRef: chunk.Ref, Value: extracted.Value}
|
||||
warnings, rejected, validateErr := r.validateTypedArtifact(attemptCtx, typed.codec, typedValidationTarget{stage: StageExtract, laneID: lane.ID, moduleKey: lane.Extract.Module, source: doc, sourceID: doc.ID, sourceInput: chunkInputMaterial(sourceInput, chunk), sessionID: sessionID, references: lane.ExtractReferences.ReferenceSet, metadata: input.Metadata, chunk: &chunk, ref: chunk.Ref, value: extracted.Value}, state.prepared.extractValidators, attempt, input.Debug)
|
||||
if validateErr != nil || rejected != nil {
|
||||
return false, rejected, validateErr
|
||||
}
|
||||
stored, encodeErr := checkpointArtifact(typed.codec, artifact.LaneID, artifact.ExtractorKey, artifact.SourceID, artifact.Value)
|
||||
if encodeErr != nil {
|
||||
return false, nil, encodeErr
|
||||
}
|
||||
stored.ChunkID, stored.ChunkIndex, stored.ChunkRef = artifact.ChunkID, artifact.ChunkIndex, artifact.ChunkRef
|
||||
accepted, serialized = artifact, stored
|
||||
acceptedWarnings = append(cloneWarnings(extracted.Warnings), warnings...)
|
||||
if debugErr := writeDebugTimed(input.Debug, attemptPath+".json", debugEnvelopeWithLLMCalls(debugTimedEnvelope{Stage: string(StageExtract), LaneID: lane.ID, ModuleKey: lane.Extract.Module, Attempt: attempt, StartedAt: started, Payload: map[string]any{"output": debugCheckpointArtifact(stored), "warnings": debugWarningEnvelopes(acceptedWarnings)}}, llmScope)); debugErr != nil {
|
||||
return false, nil, debugErr
|
||||
}
|
||||
return true, nil, nil
|
||||
})
|
||||
result.err = err
|
||||
if err == nil && !ok {
|
||||
result.rejected = rejection
|
||||
return result
|
||||
}
|
||||
result.value, result.serialized, result.warnings = accepted, serialized, acceptedWarnings
|
||||
return result
|
||||
}
|
||||
|
||||
func finalizeLaneExtract(checkpoints CheckpointRecorder, state *laneExtractState) error {
|
||||
lane := state.prepared.resolved
|
||||
indexes := make([]int, 0, len(state.results))
|
||||
for index := range state.results {
|
||||
indexes = append(indexes, index)
|
||||
}
|
||||
sort.Ints(indexes)
|
||||
for _, index := range indexes {
|
||||
result := state.results[index]
|
||||
if result.rejected != nil {
|
||||
state.rejected = append(state.rejected, *result.rejected)
|
||||
continue
|
||||
}
|
||||
state.values = append(state.values, result.value)
|
||||
state.serialized = append(state.serialized, result.serialized)
|
||||
state.warnings = append(state.warnings, result.warnings...)
|
||||
}
|
||||
sort.SliceStable(state.values, func(i, j int) bool { return state.values[i].ChunkIndex < state.values[j].ChunkIndex })
|
||||
sort.SliceStable(state.serialized, func(i, j int) bool { return state.serialized[i].ChunkIndex < state.serialized[j].ChunkIndex })
|
||||
sort.SliceStable(state.rejected, func(i, j int) bool { return state.rejected[i].ChunkIndex < state.rejected[j].ChunkIndex })
|
||||
if !state.decision.Reused {
|
||||
if err := recordExtract(checkpoints, lane.ID, lane.Extract.Module, state.deps, state.serialized, state.rejected, state.warnings); err != nil {
|
||||
return fmt.Errorf("write extract checkpoint for lane %q: %w", lane.ID, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *Runner) continueLane(ctx context.Context, input RunInput, checkpoints CheckpointRecorder, loader CheckpointLoader, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, chunks []source.Chunk, state *laneExtractState) (RunOutput, error) {
|
||||
lane := state.prepared.resolved
|
||||
local := RunOutput{Manifest: manifestFromPipeline(input)}
|
||||
checkpoint := ExtractCheckpoint{Outputs: state.serialized, Rejected: state.rejected, Warnings: state.warnings}
|
||||
coordinatedLoader := completedExtractLoader{CheckpointLoader: loader, laneID: lane.ID, checkpoint: checkpoint}
|
||||
input.extractDecision = &state.decision
|
||||
err := r.runTypedLane(ctx, input, checkpoints, coordinatedLoader, doc, sourceInput, sessionID, chunks, state.prepared, &local)
|
||||
return local, err
|
||||
}
|
||||
|
||||
func classifyLaneError(lane, sentinel int, err error) orderedRunError {
|
||||
stage := 1
|
||||
var laneErr *laneRunError
|
||||
if errors.As(err, &laneErr) {
|
||||
switch laneErr.stage {
|
||||
case StageExtract:
|
||||
stage = 0
|
||||
case StageNormalize:
|
||||
stage = 2
|
||||
}
|
||||
} else if strings.Contains(strings.ToLower(err.Error()), "normalize") {
|
||||
stage = 2
|
||||
}
|
||||
return orderedRunError{stage: stage, lane: lane, chunk: sentinel, err: err}
|
||||
}
|
||||
|
||||
func selectRunError(parent context.Context, values []orderedRunError) error {
|
||||
if err := parent.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(values) == 0 {
|
||||
return nil
|
||||
}
|
||||
hasReal := false
|
||||
for _, value := range values {
|
||||
if !errors.Is(value.err, context.Canceled) && !errors.Is(value.err, context.DeadlineExceeded) {
|
||||
hasReal = true
|
||||
break
|
||||
}
|
||||
}
|
||||
filtered := values[:0]
|
||||
for _, value := range values {
|
||||
if hasReal && (errors.Is(value.err, context.Canceled) || errors.Is(value.err, context.DeadlineExceeded)) {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, value)
|
||||
}
|
||||
sort.SliceStable(filtered, func(i, j int) bool {
|
||||
if filtered[i].stage != filtered[j].stage {
|
||||
return filtered[i].stage < filtered[j].stage
|
||||
}
|
||||
if filtered[i].lane != filtered[j].lane {
|
||||
return filtered[i].lane < filtered[j].lane
|
||||
}
|
||||
return filtered[i].chunk < filtered[j].chunk
|
||||
})
|
||||
return filtered[0].err
|
||||
}
|
||||
|
||||
func mergeLaneOutput(dst *RunOutput, src RunOutput) {
|
||||
if dst == nil {
|
||||
return
|
||||
}
|
||||
dst.NormalizeOutputs = append(dst.NormalizeOutputs, cloneSerializedOutputs(src.NormalizeOutputs)...)
|
||||
dst.Rejected = append(dst.Rejected, cloneRejectedOutputs(src.Rejected)...)
|
||||
dst.Warnings = append(dst.Warnings, cloneWarnings(src.Warnings)...)
|
||||
dst.CheckpointEvents = append(dst.CheckpointEvents, src.CheckpointEvents...)
|
||||
for i := range dst.Manifest.ArtifactLanes {
|
||||
for j := range src.Manifest.ArtifactLanes {
|
||||
if dst.Manifest.ArtifactLanes[i].ID == src.Manifest.ArtifactLanes[j].ID && src.Manifest.ArtifactLanes[j].Metadata != nil {
|
||||
dst.Manifest.ArtifactLanes[i].Metadata = cloneMetadata(src.Manifest.ArtifactLanes[j].Metadata)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -118,7 +118,21 @@ func checkpointArtifact(codec artifactCodecEntry, laneID, moduleKey, sourceID st
|
||||
return CheckpointArtifact{LaneID: laneID, ModuleKey: moduleKey, SourceID: sourceID, Artifact: serialized, SchemaDigest: contracts.DigestArtifactSchema(serialized.Schema)}, nil
|
||||
}
|
||||
|
||||
func (r *Runner) runTypedLane(ctx context.Context, input RunInput, checkpoints CheckpointRecorder, loader CheckpointLoader, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, chunks []source.Chunk, prepared preparedLaneExecutor, output *RunOutput) error {
|
||||
type laneRunError struct {
|
||||
stage ModuleStage
|
||||
err error
|
||||
}
|
||||
|
||||
func (e *laneRunError) Error() string { return e.err.Error() }
|
||||
func (e *laneRunError) Unwrap() error { return e.err }
|
||||
|
||||
func (r *Runner) runTypedLane(ctx context.Context, input RunInput, checkpoints CheckpointRecorder, loader CheckpointLoader, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, chunks []source.Chunk, prepared preparedLaneExecutor, output *RunOutput) (err error) {
|
||||
activeStage := StageExtract
|
||||
defer func() {
|
||||
if err != nil {
|
||||
err = &laneRunError{stage: activeStage, err: err}
|
||||
}
|
||||
}()
|
||||
lane, typed := prepared.resolved, prepared.typed
|
||||
if typed == nil {
|
||||
return fmt.Errorf("typed lane %q executor is not prepared", lane.ID)
|
||||
@@ -143,8 +157,12 @@ func (r *Runner) runTypedLane(ctx context.Context, input RunInput, checkpoints C
|
||||
}
|
||||
}
|
||||
}
|
||||
recordCheckpointEvent(output, loader, string(StageExtract), lane.ID, lane.Extract.Module, decision)
|
||||
if err := writeDebugTimed(input.Debug, path.Join("extract", debugPathComponent(lane.ID), "input.json"), debugTimedEnvelope{Stage: string(StageExtract), LaneID: lane.ID, ModuleKey: lane.Extract.Module, StartedAt: time.Now().UTC(), Payload: map[string]any{"reused": decision.Reused, "decision": decision, "source": debugSourceDocumentEnvelope(doc), "chunks": debugSourceChunkEnvelopes(chunks), "options": redactSensitiveMap(lane.Extract.Options), "metadata": redactSensitiveMap(input.Metadata)}}); err != nil {
|
||||
reportedDecision := decision
|
||||
if input.extractDecision != nil {
|
||||
reportedDecision = *input.extractDecision
|
||||
}
|
||||
recordCheckpointEvent(output, loader, string(StageExtract), lane.ID, lane.Extract.Module, reportedDecision)
|
||||
if err := writeDebugTimed(input.Debug, path.Join("extract", debugPathComponent(lane.ID), "input.json"), debugTimedEnvelope{Stage: string(StageExtract), LaneID: lane.ID, ModuleKey: lane.Extract.Module, StartedAt: time.Now().UTC(), Payload: map[string]any{"reused": reportedDecision.Reused, "decision": reportedDecision, "source": debugSourceDocumentEnvelope(doc), "chunks": debugSourceChunkEnvelopes(chunks), "options": redactSensitiveMap(lane.Extract.Options), "metadata": redactSensitiveMap(input.Metadata)}}); err != nil {
|
||||
return err
|
||||
}
|
||||
if decision.Reused {
|
||||
@@ -218,12 +236,13 @@ func (r *Runner) runTypedLane(ctx context.Context, input RunInput, checkpoints C
|
||||
}
|
||||
sort.SliceStable(values, func(i, j int) bool { return values[i].ChunkIndex < values[j].ChunkIndex })
|
||||
sort.SliceStable(serializedExtracts, func(i, j int) bool { return serializedExtracts[i].ChunkIndex < serializedExtracts[j].ChunkIndex })
|
||||
if err := writeDebugTimed(input.Debug, path.Join("extract", debugPathComponent(lane.ID), "output.json"), debugTimedEnvelope{Stage: string(StageExtract), LaneID: lane.ID, ModuleKey: lane.Extract.Module, StartedAt: time.Now().UTC(), Payload: map[string]any{"reused": decision.Reused, "outputs": debugCheckpointArtifacts(serializedExtracts), "rejected": debugRejectedOutputEnvelopes(output.Rejected[rejectedStart:]), "warnings": debugWarningEnvelopes(extractWarnings)}}); err != nil {
|
||||
if err := writeDebugTimed(input.Debug, path.Join("extract", debugPathComponent(lane.ID), "output.json"), debugTimedEnvelope{Stage: string(StageExtract), LaneID: lane.ID, ModuleKey: lane.Extract.Module, StartedAt: time.Now().UTC(), Payload: map[string]any{"reused": reportedDecision.Reused, "outputs": debugCheckpointArtifacts(serializedExtracts), "rejected": debugRejectedOutputEnvelopes(output.Rejected[rejectedStart:]), "warnings": debugWarningEnvelopes(extractWarnings)}}); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(values) == 0 {
|
||||
return nil
|
||||
}
|
||||
activeStage = StageMerge
|
||||
|
||||
mergeInputs := make([]contracts.ExtractArtifact[any], len(values))
|
||||
for i, value := range values {
|
||||
@@ -294,6 +313,7 @@ func (r *Runner) runTypedLane(ctx context.Context, input RunInput, checkpoints C
|
||||
return err
|
||||
}
|
||||
|
||||
activeStage = StageNormalize
|
||||
normalizeDeps := artifactCheckpointDigests([]CheckpointArtifact{serializedMerge})
|
||||
normalizeCP, normalizeDecision := loadNormalize(loader, lane.ID, lane.Normalize.Module, normalizeDeps)
|
||||
if normalizeDecision.Reused {
|
||||
|
||||
154
internal/framework/pipeline/synchronized_collaborators.go
Normal file
154
internal/framework/pipeline/synchronized_collaborators.go
Normal file
@@ -0,0 +1,154 @@
|
||||
package pipeline
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||
)
|
||||
|
||||
type lockedDebugRecorder struct {
|
||||
inner DebugRecorder
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func synchronizedDebugRecorder(inner DebugRecorder) DebugRecorder {
|
||||
if _, ok := inner.(*lockedDebugRecorder); ok {
|
||||
return inner
|
||||
}
|
||||
return &lockedDebugRecorder{inner: inner}
|
||||
}
|
||||
|
||||
// SynchronizedDebugRecorder serializes access when a recorder is shared by
|
||||
// pipeline operations and construction-injected LLM clients.
|
||||
func SynchronizedDebugRecorder(inner DebugRecorder) DebugRecorder {
|
||||
return synchronizedDebugRecorder(inner)
|
||||
}
|
||||
func (r *lockedDebugRecorder) Enabled() bool {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.inner.Enabled()
|
||||
}
|
||||
func (r *lockedDebugRecorder) WriteJSON(name string, payload any) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.inner.WriteJSON(name, payload)
|
||||
}
|
||||
func (r *lockedDebugRecorder) WriteBytes(name string, data []byte) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.inner.WriteBytes(name, data)
|
||||
}
|
||||
|
||||
type lockedCheckpointLoader struct {
|
||||
inner CheckpointLoader
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func synchronizedCheckpointLoader(inner CheckpointLoader) CheckpointLoader {
|
||||
if _, ok := inner.(*lockedCheckpointLoader); ok {
|
||||
return inner
|
||||
}
|
||||
return &lockedCheckpointLoader{inner: inner}
|
||||
}
|
||||
func (l *lockedCheckpointLoader) Enabled() bool {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
return l.inner.Enabled()
|
||||
}
|
||||
func (l *lockedCheckpointLoader) Source(key string) (SourceCheckpoint, CheckpointDecision) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
return l.inner.Source(key)
|
||||
}
|
||||
func (l *lockedCheckpointLoader) Chunk(key, digest string) (ChunkCheckpoint, CheckpointDecision) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
return l.inner.Chunk(key, digest)
|
||||
}
|
||||
func (l *lockedCheckpointLoader) Extract(lane, key string, deps []CheckpointFingerprint) (ExtractCheckpoint, CheckpointDecision) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
return l.inner.Extract(lane, key, deps)
|
||||
}
|
||||
func (l *lockedCheckpointLoader) Merge(lane, key string, deps []CheckpointFingerprint) (MergeCheckpoint, CheckpointDecision) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
return l.inner.Merge(lane, key, deps)
|
||||
}
|
||||
func (l *lockedCheckpointLoader) Normalize(lane, key string, deps []CheckpointFingerprint) (NormalizeCheckpoint, CheckpointDecision) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
return l.inner.Normalize(lane, key, deps)
|
||||
}
|
||||
|
||||
type lockedCheckpointRecorder struct {
|
||||
inner CheckpointRecorder
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func synchronizedCheckpointRecorder(inner CheckpointRecorder) CheckpointRecorder {
|
||||
if _, ok := inner.(*lockedCheckpointRecorder); ok {
|
||||
return inner
|
||||
}
|
||||
return &lockedCheckpointRecorder{inner: inner}
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) call(fn func() error) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return fn()
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) SourceRunning(key string) error {
|
||||
return r.call(func() error { return r.inner.SourceRunning(key) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) SourceSucceeded(key string, doc *source.SourceDocument) error {
|
||||
return r.call(func() error { return r.inner.SourceSucceeded(key, doc) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) SourceFailed(key string, err error) error {
|
||||
return r.call(func() error { return r.inner.SourceFailed(key, err) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) ChunkRunning(key, digest string) error {
|
||||
return r.call(func() error { return r.inner.ChunkRunning(key, digest) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) ChunkSucceeded(key, digest string, chunks []source.Chunk, warnings []contracts.Warning) error {
|
||||
return r.call(func() error { return r.inner.ChunkSucceeded(key, digest, chunks, warnings) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) ChunkRejected(key, digest string, rejected contracts.RejectedOutput) error {
|
||||
return r.call(func() error { return r.inner.ChunkRejected(key, digest, rejected) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) ChunkFailed(key, digest string, err error) error {
|
||||
return r.call(func() error { return r.inner.ChunkFailed(key, digest, err) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) ExtractRunning(lane, key string, deps []CheckpointFingerprint) error {
|
||||
return r.call(func() error { return r.inner.ExtractRunning(lane, key, deps) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) ExtractSucceeded(lane, key string, deps []CheckpointFingerprint, outputs []CheckpointArtifact, rejected []contracts.RejectedOutput, warnings []contracts.Warning) error {
|
||||
return r.call(func() error { return r.inner.ExtractSucceeded(lane, key, deps, outputs, rejected, warnings) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) ExtractFailed(lane, key string, deps []CheckpointFingerprint, err error) error {
|
||||
return r.call(func() error { return r.inner.ExtractFailed(lane, key, deps, err) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) MergeRunning(lane, key string, deps []CheckpointFingerprint) error {
|
||||
return r.call(func() error { return r.inner.MergeRunning(lane, key, deps) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) MergeSucceeded(lane, key string, deps []CheckpointFingerprint, output CheckpointArtifact, warnings []contracts.Warning) error {
|
||||
return r.call(func() error { return r.inner.MergeSucceeded(lane, key, deps, output, warnings) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) MergeRejected(lane, key string, deps []CheckpointFingerprint, rejected contracts.RejectedOutput) error {
|
||||
return r.call(func() error { return r.inner.MergeRejected(lane, key, deps, rejected) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) MergeFailed(lane, key string, deps []CheckpointFingerprint, err error) error {
|
||||
return r.call(func() error { return r.inner.MergeFailed(lane, key, deps, err) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) NormalizeRunning(lane, key string, deps []CheckpointFingerprint) error {
|
||||
return r.call(func() error { return r.inner.NormalizeRunning(lane, key, deps) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) NormalizeSucceeded(lane, key string, deps []CheckpointFingerprint, output CheckpointArtifact, warnings []contracts.Warning) error {
|
||||
return r.call(func() error { return r.inner.NormalizeSucceeded(lane, key, deps, output, warnings) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) NormalizeRejected(lane, key string, deps []CheckpointFingerprint, rejected contracts.RejectedOutput) error {
|
||||
return r.call(func() error { return r.inner.NormalizeRejected(lane, key, deps, rejected) })
|
||||
}
|
||||
func (r *lockedCheckpointRecorder) NormalizeFailed(lane, key string, deps []CheckpointFingerprint, err error) error {
|
||||
return r.call(func() error { return r.inner.NormalizeFailed(lane, key, deps, err) })
|
||||
}
|
||||
Reference in New Issue
Block a user