Harden checkpoint reuse and combat validation
This commit is contained in:
@@ -81,9 +81,15 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, step PreparedP
|
||||
output := RunOutput{Manifest: manifestFromPipeline(input)}
|
||||
states, err := initializeLaneStates(input, step, checkpoints, loader, doc, chunks, &output)
|
||||
if err != nil {
|
||||
if mergeErr := mergeTerminalLaneStates(&output, states); mergeErr != nil {
|
||||
return output, errors.Join(err, mergeErr)
|
||||
}
|
||||
return output, err
|
||||
}
|
||||
completedOutputs, runErrors := r.runLaneEngine(parent, input, checkpoints, loader, doc, sourceInput, sessionID, chunks, states)
|
||||
completedOutputs, runErrors := r.runLaneEngine(parent, laneEngineConfig{
|
||||
input: input, checkpoints: checkpoints, loader: loader, doc: doc,
|
||||
sourceInput: sourceInput, sessionID: sessionID, chunks: chunks, states: states,
|
||||
})
|
||||
if err := mergeCompletedLanes(&output, completedOutputs); err != nil {
|
||||
return output, err
|
||||
}
|
||||
@@ -97,183 +103,232 @@ func initializeLaneStates(input RunInput, step PreparedPipelineStep, checkpoints
|
||||
states := make([]*laneExtractState, len(step.lanes))
|
||||
for i, prepared := range step.lanes {
|
||||
if prepared.typed == nil {
|
||||
return nil, fmt.Errorf("typed lane %q executor is not prepared", prepared.resolved.ID)
|
||||
return states, fmt.Errorf("typed lane %q executor is not prepared", prepared.resolved.ID)
|
||||
}
|
||||
if err := setTypedLaneManifestMetadata(output, prepared.resolved.ID, prepared.typed.extractor, prepared.typed.merger, prepared.typed.normalizer); err != nil {
|
||||
return nil, err
|
||||
return states, err
|
||||
}
|
||||
if input.CheckpointPolicy.requiresReusable(input.stepID, prepared.resolved.ID) && !input.CheckpointPolicy.forced(input.stepID, prepared.resolved.ID) {
|
||||
state, err := hydrateRequiredLane(input, loader, doc, i, prepared, output)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
state, err := hydrateRequiredLane(input, loader, doc, i, prepared)
|
||||
states[i] = state
|
||||
if err != nil {
|
||||
return states, err
|
||||
}
|
||||
continue
|
||||
}
|
||||
state, err := prepareLaneExtract(input, loader, doc, chunks, i, prepared, output)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return states, err
|
||||
}
|
||||
if !state.decision.Reused {
|
||||
if err := checkpointExtractRunning(checkpoints, input.stepID, prepared.resolved.ID, prepared.resolved.Extract.Module, state.deps); err != nil {
|
||||
return nil, fmt.Errorf("write extract checkpoint for lane %q: %w", prepared.resolved.ID, err)
|
||||
return states, fmt.Errorf("write extract checkpoint for lane %q: %w", prepared.resolved.ID, err)
|
||||
}
|
||||
} else if err := finalizeLaneExtract(checkpoints, input.stepID, state); err != nil {
|
||||
return nil, err
|
||||
return states, err
|
||||
}
|
||||
states[i] = state
|
||||
}
|
||||
return states, nil
|
||||
}
|
||||
|
||||
func (r *Runner) runLaneEngine(parent context.Context, input RunInput, checkpoints CheckpointRecorder, loader CheckpointLoader, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, chunks []source.Chunk, states []*laneExtractState) ([]RunOutput, []orderedRunError) {
|
||||
workerCount := input.ExtractWorkers
|
||||
type laneEngineConfig struct {
|
||||
input RunInput
|
||||
checkpoints CheckpointRecorder
|
||||
loader CheckpointLoader
|
||||
doc *source.SourceDocument
|
||||
sourceInput contracts.LLMInputMaterial
|
||||
sessionID string
|
||||
chunks []source.Chunk
|
||||
states []*laneExtractState
|
||||
}
|
||||
|
||||
type laneEngine struct {
|
||||
runner *Runner
|
||||
laneEngineConfig
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
workerCount int
|
||||
jobs chan extractJob
|
||||
results chan extractJobResult
|
||||
completions chan laneCompletion
|
||||
continuations chan *laneExtractState
|
||||
continuationWorkers sync.WaitGroup
|
||||
pending []*laneExtractState
|
||||
completedOutputs []RunOutput
|
||||
runErrors []orderedRunError
|
||||
launched int
|
||||
completed int
|
||||
}
|
||||
|
||||
func (r *Runner) runLaneEngine(parent context.Context, config laneEngineConfig) ([]RunOutput, []orderedRunError) {
|
||||
engine := newLaneEngine(r, parent, config)
|
||||
return engine.run()
|
||||
}
|
||||
|
||||
func newLaneEngine(r *Runner, parent context.Context, config laneEngineConfig) *laneEngine {
|
||||
workerCount := config.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))
|
||||
continuations := make(chan *laneExtractState, workerCount)
|
||||
return &laneEngine{
|
||||
runner: r, laneEngineConfig: config, ctx: ctx, cancel: cancel, workerCount: workerCount,
|
||||
jobs: make(chan extractJob, workerCount), results: make(chan extractJobResult, workerCount),
|
||||
completions: make(chan laneCompletion, len(config.states)), continuations: make(chan *laneExtractState, workerCount),
|
||||
completedOutputs: make([]RunOutput, len(config.states)),
|
||||
}
|
||||
}
|
||||
|
||||
func (e *laneEngine) run() ([]RunOutput, []orderedRunError) {
|
||||
defer e.cancel()
|
||||
e.initializeCollection()
|
||||
e.startExtractWorkers()
|
||||
e.startContinuationWorkers()
|
||||
e.collect()
|
||||
close(e.continuations)
|
||||
e.continuationWorkers.Wait()
|
||||
return e.completedOutputs, e.runErrors
|
||||
}
|
||||
|
||||
func (e *laneEngine) initializeCollection() {
|
||||
for _, state := range e.states {
|
||||
if state.terminal {
|
||||
e.completedOutputs[state.index] = state.output
|
||||
} else if state.decision.Reused {
|
||||
e.pending = append(e.pending, state)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (e *laneEngine) startExtractWorkers() {
|
||||
var workers sync.WaitGroup
|
||||
for i := 0; i < workerCount; i++ {
|
||||
for i := 0; i < e.workerCount; i++ {
|
||||
workers.Add(1)
|
||||
go func() {
|
||||
defer workers.Done()
|
||||
for job := range jobs {
|
||||
if ctx.Err() != nil {
|
||||
for job := range e.jobs {
|
||||
if e.ctx.Err() != nil {
|
||||
continue
|
||||
}
|
||||
result := r.runExtractJob(ctx, input, doc, sourceInput, sessionID, job)
|
||||
results <- result
|
||||
e.results <- e.runner.runExtractJob(e.ctx, e.input, e.doc, e.sourceInput, e.sessionID, job)
|
||||
}
|
||||
}()
|
||||
}
|
||||
go func() {
|
||||
defer close(jobs)
|
||||
for chunkIndex := range chunks {
|
||||
for laneIndex := range states {
|
||||
state := states[laneIndex]
|
||||
if state.terminal || state.decision.Reused {
|
||||
continue
|
||||
}
|
||||
select {
|
||||
case jobs <- extractJob{lane: state, chunk: chunks[chunkIndex]}:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
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, runErrors := collectLaneResults(ctx, cancel, input, checkpoints, chunks, states, results, completions, continuations)
|
||||
close(continuations)
|
||||
continuationWorkers.Wait()
|
||||
return completedOutputs, runErrors
|
||||
go e.dispatchExtractJobs()
|
||||
go func() { workers.Wait(); close(e.results) }()
|
||||
}
|
||||
|
||||
func collectLaneResults(ctx context.Context, cancel context.CancelFunc, input RunInput, checkpoints CheckpointRecorder, chunks []source.Chunk, states []*laneExtractState, results <-chan extractJobResult, completions <-chan laneCompletion, continuations chan<- *laneExtractState) ([]RunOutput, []orderedRunError) {
|
||||
completedOutputs := make([]RunOutput, len(states))
|
||||
var runErrors []orderedRunError
|
||||
var pendingContinuations []*laneExtractState
|
||||
launched, completed := 0, 0
|
||||
for _, state := range states {
|
||||
if state.terminal {
|
||||
completedOutputs[state.index] = state.output
|
||||
} else if state.decision.Reused {
|
||||
pendingContinuations = append(pendingContinuations, state)
|
||||
func (e *laneEngine) dispatchExtractJobs() {
|
||||
defer close(e.jobs)
|
||||
for chunkIndex := range e.chunks {
|
||||
for laneIndex := range e.states {
|
||||
state := e.states[laneIndex]
|
||||
if state.terminal || state.decision.Reused {
|
||||
continue
|
||||
}
|
||||
select {
|
||||
case e.jobs <- extractJob{lane: state, chunk: e.chunks[chunkIndex]}:
|
||||
case <-e.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
resultChannel := results
|
||||
for resultChannel != nil || len(pendingContinuations) > 0 || completed < launched {
|
||||
}
|
||||
|
||||
func (e *laneEngine) startContinuationWorkers() {
|
||||
for i := 0; i < e.workerCount; i++ {
|
||||
e.continuationWorkers.Add(1)
|
||||
go func() {
|
||||
defer e.continuationWorkers.Done()
|
||||
for state := range e.continuations {
|
||||
if err := e.ctx.Err(); err != nil {
|
||||
e.completions <- laneCompletion{index: state.index, err: err}
|
||||
continue
|
||||
}
|
||||
laneOutput, err := e.runner.continueLane(e.ctx, e.input, e.checkpoints, e.loader, e.doc, e.sourceInput, e.sessionID, e.chunks, state)
|
||||
e.completions <- laneCompletion{index: state.index, output: laneOutput, err: err}
|
||||
}
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
func (e *laneEngine) collect() {
|
||||
resultChannel := (<-chan extractJobResult)(e.results)
|
||||
for resultChannel != nil || len(e.pending) > 0 || e.completed < e.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
|
||||
if len(e.pending) > 0 && e.ctx.Err() == nil {
|
||||
continuationChannel = e.continuations
|
||||
nextContinuation = e.pending[0]
|
||||
} else if e.ctx.Err() != nil {
|
||||
e.pending = nil
|
||||
}
|
||||
select {
|
||||
case continuationChannel <- nextContinuation:
|
||||
pendingContinuations = pendingContinuations[1:]
|
||||
launched++
|
||||
e.pending = e.pending[1:]
|
||||
e.launched++
|
||||
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})
|
||||
_ = checkpointExtractFailed(checkpoints, input.stepID, 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, input.stepID, state); err != nil {
|
||||
state.failed = true
|
||||
runErrors = append(runErrors, orderedRunError{stage: 0, lane: state.index, chunk: len(chunks), err: err})
|
||||
cancel()
|
||||
} else {
|
||||
pendingContinuations = append(pendingContinuations, 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()
|
||||
}
|
||||
e.handleExtractResult(result)
|
||||
case completion := <-e.completions:
|
||||
e.handleCompletion(completion)
|
||||
}
|
||||
}
|
||||
return completedOutputs, runErrors
|
||||
}
|
||||
|
||||
func hydrateRequiredLane(input RunInput, loader CheckpointLoader, doc *source.SourceDocument, index int, prepared preparedLaneExecutor, output *RunOutput) (*laneExtractState, error) {
|
||||
func (e *laneEngine) handleExtractResult(result extractJobResult) {
|
||||
state := e.states[result.laneIndex]
|
||||
state.remaining--
|
||||
if result.err != nil {
|
||||
state.failed = true
|
||||
e.runErrors = append(e.runErrors, orderedRunError{stage: 0, lane: state.index, chunk: result.chunkIndex, err: result.err})
|
||||
_ = checkpointExtractFailed(e.checkpoints, e.input.stepID, state.prepared.resolved.ID, state.prepared.resolved.Extract.Module, state.deps, result.err)
|
||||
e.cancel()
|
||||
} else {
|
||||
state.results[result.chunkIndex] = result
|
||||
}
|
||||
if state.remaining != 0 || state.failed || e.ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
if err := finalizeLaneExtract(e.checkpoints, e.input.stepID, state); err != nil {
|
||||
state.failed = true
|
||||
e.runErrors = append(e.runErrors, orderedRunError{stage: 0, lane: state.index, chunk: len(e.chunks), err: err})
|
||||
e.cancel()
|
||||
return
|
||||
}
|
||||
e.pending = append(e.pending, state)
|
||||
}
|
||||
|
||||
func (e *laneEngine) handleCompletion(completion laneCompletion) {
|
||||
e.completed++
|
||||
e.completedOutputs[completion.index] = completion.output
|
||||
if completion.err != nil {
|
||||
e.runErrors = append(e.runErrors, classifyLaneError(completion.index, len(e.chunks), completion.err))
|
||||
e.cancel()
|
||||
}
|
||||
}
|
||||
|
||||
func hydrateRequiredLane(input RunInput, loader CheckpointLoader, doc *source.SourceDocument, index int, prepared preparedLaneExecutor) (*laneExtractState, error) {
|
||||
lane, typed := prepared.resolved, prepared.typed
|
||||
local := RunOutput{Manifest: manifestFromPipeline(input)}
|
||||
checkpoint, decision := loader.AcceptedNormalize(input.stepID, lane.ID, lane.Normalize.Module)
|
||||
if decision.Reused {
|
||||
decision = checkpointDecision(CheckpointDecisionReused, CheckpointReasonAcceptedArtifactReused, "accepted normalized artifact is reusable")
|
||||
decision = checkpointDecision(CheckpointDecisionReused, CheckpointReasonAcceptedArtifactReused)
|
||||
if checkpoint.Output.LaneID != lane.ID || checkpoint.Output.ModuleKey != lane.Normalize.Module || checkpoint.Output.SourceID != doc.ID {
|
||||
decision = checkpointDecision(CheckpointDecisionExecuted, CheckpointReasonArtifactPayloadInvalid, "accepted normalized artifact provenance does not match the producer lane")
|
||||
decision = checkpointDecision(CheckpointDecisionExecuted, CheckpointReasonArtifactPayloadInvalid)
|
||||
}
|
||||
}
|
||||
decision, err := resolveCheckpointDecision(&local, loader, input.CheckpointPolicy, StageNormalize, input.stepID, lane.ID, lane.Normalize.Module, decision, typed.codec, []CheckpointArtifact{checkpoint.Output})
|
||||
resolution, err := resolveCheckpointDecision(&local, loader, input.CheckpointPolicy, StageNormalize, input.stepID, lane.ID, lane.Normalize.Module, decision, typed.codec, []CheckpointArtifact{checkpoint.Output})
|
||||
decision = resolution.decision
|
||||
state := &laneExtractState{index: index, prepared: prepared, decision: decision, terminal: true, output: local}
|
||||
if err != nil {
|
||||
if mergeErr := mergeLaneOutput(output, local); mergeErr != nil {
|
||||
return nil, mergeErr
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
_, hydrated, err := decodeCanonicalCheckpointArtifact(typed.codec, checkpoint.Output)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("hydrate accepted normalized artifact for step %q lane %q: %w", input.stepID, lane.ID, err)
|
||||
return state, err
|
||||
}
|
||||
hydrated := resolution.artifacts[0]
|
||||
local.Warnings = append(local.Warnings, cloneWarnings(checkpoint.Warnings)...)
|
||||
local.NormalizeOutputs = append(local.NormalizeOutputs, contracts.SerializedOutput{
|
||||
StepID: input.stepID,
|
||||
@@ -282,7 +337,20 @@ func hydrateRequiredLane(input RunInput, loader CheckpointLoader, doc *source.So
|
||||
SourceID: doc.ID,
|
||||
Artifact: contracts.CloneSerializedArtifact(hydrated.Artifact),
|
||||
})
|
||||
return &laneExtractState{index: index, prepared: prepared, decision: decision, terminal: true, output: local}, nil
|
||||
state.output = local
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func mergeTerminalLaneStates(output *RunOutput, states []*laneExtractState) error {
|
||||
for _, state := range states {
|
||||
if state == nil || !state.terminal {
|
||||
continue
|
||||
}
|
||||
if err := mergeLaneOutput(output, state.output); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func mergeCompletedLanes(output *RunOutput, completedOutputs []RunOutput) error {
|
||||
@@ -304,19 +372,16 @@ func prepareLaneExtract(input RunInput, loader CheckpointLoader, doc *source.Sou
|
||||
deps := append(digestFingerprints("chunks", digest), generatedReferenceDependencies(extractReferences)...)
|
||||
state := &laneExtractState{index: index, prepared: prepared, deps: normalizeCheckpointFingerprints(deps), remaining: len(chunks), results: make(map[int]extractJobResult, len(chunks))}
|
||||
cp, decision := loadExtract(loader, input.stepID, lane.ID, lane.Extract.Module, state.deps)
|
||||
decision, err = resolveCheckpointDecision(output, loader, input.CheckpointPolicy, StageExtract, input.stepID, lane.ID, lane.Extract.Module, decision, typed.codec, cp.Outputs)
|
||||
resolution, err := resolveCheckpointDecision(output, loader, input.CheckpointPolicy, StageExtract, input.stepID, lane.ID, lane.Extract.Module, decision, typed.codec, cp.Outputs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
decision = resolution.decision
|
||||
state.decision = decision
|
||||
if decision.Reused {
|
||||
state.remaining = 0
|
||||
for _, stored := range cp.Outputs {
|
||||
value, hydrated, decodeErr := decodeCanonicalCheckpointArtifact(typed.codec, stored)
|
||||
if decodeErr != nil {
|
||||
return nil, fmt.Errorf("decode extract checkpoint for lane %q: %w", lane.ID, decodeErr)
|
||||
}
|
||||
stored = hydrated
|
||||
for i, stored := range resolution.artifacts {
|
||||
value := resolution.values[i]
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user