Implement selective checkpoint recomputation

This commit is contained in:
2026-07-21 21:59:48 +00:00
parent 22d4f29670
commit c437682407
17 changed files with 902 additions and 180 deletions

View File

@@ -36,12 +36,15 @@ type finalizedExtractResults struct {
decision CheckpointDecision
}
func loadExtract(loader CheckpointLoader, laneID, moduleKey string, deps []CheckpointFingerprint) (ExtractCheckpoint, CheckpointDecision) {
func loadExtract(loader CheckpointLoader, stepID, laneID, moduleKey string, deps []CheckpointFingerprint) (ExtractCheckpoint, CheckpointDecision) {
if stepAware, ok := loader.(StepCheckpointLoader); ok {
return stepAware.ExtractForStep(stepID, laneID, moduleKey, deps)
}
return loader.Extract(laneID, moduleKey, deps)
}
func recordExtract(recorder CheckpointRecorder, laneID, moduleKey string, deps []CheckpointFingerprint, outputs []CheckpointArtifact, rejected []contracts.RejectedOutput, warnings []contracts.Warning) error {
return recorder.ExtractSucceeded(laneID, moduleKey, deps, outputs, rejected, warnings)
func recordExtract(recorder CheckpointRecorder, stepID, laneID, moduleKey string, deps []CheckpointFingerprint, outputs []CheckpointArtifact, rejected []contracts.RejectedOutput, warnings []contracts.Warning) error {
return checkpointExtractSucceeded(recorder, stepID, laneID, moduleKey, deps, outputs, rejected, warnings)
}
type extractJob struct {
@@ -87,10 +90,10 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, step PreparedP
return output, err
}
if !state.decision.Reused {
if err := checkpoints.ExtractRunning(prepared.resolved.ID, prepared.resolved.Extract.Module, state.deps); err != nil {
if err := checkpointExtractRunning(checkpoints, input.stepID, 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)
}
} else if err := finalizeLaneExtract(checkpoints, state); err != nil {
} else if err := finalizeLaneExtract(checkpoints, input.stepID, state); err != nil {
return output, err
}
states[i] = state
@@ -188,13 +191,13 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, step PreparedP
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)
_ = 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, state); 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()
@@ -233,27 +236,31 @@ func prepareLaneExtract(input RunInput, loader CheckpointLoader, doc *source.Sou
extractReferences := operationReferenceSet(input, lane.ExtractReferences)
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, lane.ID, lane.Extract.Module, state.deps)
cp, decision := loadExtract(loader, input.stepID, lane.ID, lane.Extract.Module, state.deps)
decision = forceCheckpointDecision(input.CheckpointPolicy, input.stepID, lane.ID, decision)
if err := requireReusableCheckpoint(input.CheckpointPolicy, input.stepID, lane.ID, decision); err != nil {
return nil, err
}
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()}
if _, _, decodeErr := decodeCanonicalCheckpointArtifact(typed.codec, stored); decodeErr != nil {
decision = CheckpointDecision{Category: "executed", ReasonCode: "artifact_not_canonical", Detail: "stored extract artifact failed canonical codec validation", Reason: "extract artifact checkpoint is not canonical"}
break
}
}
}
if err := requireReusableCheckpoint(input.CheckpointPolicy, input.stepID, lane.ID, decision); err != nil {
return nil, err
}
state.decision = decision
if decision.Reused {
state.remaining = 0
for _, stored := range cp.Outputs {
value, decodeErr := decodeCheckpointArtifact(typed.codec, stored)
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, decodeErr = hydrateCheckpointArtifact(typed.codec, stored, value)
if decodeErr != nil {
return nil, fmt.Errorf("hydrate extract checkpoint for lane %q: %w", lane.ID, decodeErr)
}
stored = hydrated
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
@@ -330,7 +337,7 @@ func (r *Runner) runExtractJob(ctx context.Context, input RunInput, doc *source.
return result
}
func finalizeLaneExtract(checkpoints CheckpointRecorder, state *laneExtractState) error {
func finalizeLaneExtract(checkpoints CheckpointRecorder, stepID string, state *laneExtractState) error {
lane := state.prepared.resolved
indexes := make([]int, 0, len(state.results))
for index := range state.results {
@@ -351,7 +358,7 @@ func finalizeLaneExtract(checkpoints CheckpointRecorder, state *laneExtractState
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 {
if err := recordExtract(checkpoints, stepID, 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)
}
}
@@ -370,7 +377,7 @@ func (r *Runner) continueLane(ctx context.Context, input RunInput, checkpoints C
}
local.Warnings = append(local.Warnings, cloneWarnings(results.warnings)...)
local.Rejected = append(local.Rejected, cloneRejectedOutputs(results.rejected)...)
recordCheckpointEvent(&local, loader, string(StageExtract), lane.ID, lane.Extract.Module, results.decision)
recordCheckpointEvent(&local, loader, string(StageExtract), input.stepID, lane.ID, lane.Extract.Module, results.decision)
if err := writeDebugTimed(input.Debug, path.Join("extract", debugPathComponent(lane.ID), "input.json"), debugTimedEnvelope{Stage: string(StageExtract), StepID: input.stepID, LaneID: lane.ID, ModuleKey: lane.Extract.Module, StartedAt: time.Now().UTC(), Payload: map[string]any{"reused": results.decision.Reused, "decision": results.decision, "source": debugSourceDocumentEnvelope(doc), "chunks": debugSourceChunkEnvelopes(chunks), "options": redactSensitiveMap(lane.Extract.Options), "metadata": redactSensitiveMap(input.Metadata)}}); err != nil {
return local, &laneRunError{stage: StageExtract, err: err}
}