Integrate chunk plan caching into the runner

This commit is contained in:
2026-07-18 00:20:56 +00:00
parent ebd449d847
commit 51a36efb6b
15 changed files with 590 additions and 441 deletions

View File

@@ -65,54 +65,6 @@ func (r *WorkspaceRecorder) SourceFailed(moduleKey string, err error) error {
return r.writeManifest("source/manifest.json", coreworkspace.SourceManifest{StageManifest: manifest})
}
func (r *WorkspaceRecorder) ChunkRunning(moduleKey string, sourceDigest string) error {
manifest := r.newStageManifest(coreworkspace.StageChunk, coreworkspace.StatusRunning)
manifest.ModuleKey = moduleKey
manifest.DependencyFingerprints = workspaceFingerprints(digestFingerprints("source_document", sourceDigest))
manifest.StartedAt = timePtr(r.timestamp())
return r.writeManifest("chunk/manifest.json", coreworkspace.ChunkManifest{StageManifest: manifest})
}
func (r *WorkspaceRecorder) ChunkSucceeded(moduleKey string, sourceDigest string, chunks []source.Chunk, warnings []contracts.Warning) error {
outputDigests, err := chunkOutputDigests(chunks)
if err != nil {
return fmt.Errorf("digest chunk checkpoint output: %w", err)
}
payload := chunksEnvelope{Chunks: chunkEnvelopes(chunks), Warnings: cloneWarnings(warnings)}
if err := r.writePayload("chunk/chunks.json", payload); err != nil {
return err
}
manifest := r.newStageManifest(coreworkspace.StageChunk, coreworkspace.StatusSucceeded)
manifest.ModuleKey = moduleKey
manifest.DependencyFingerprints = workspaceFingerprints(digestFingerprints("source_document", sourceDigest))
manifest.OutputDigests = workspaceFingerprints(outputDigests)
manifest.ValidationStatus = validationStatusString(warnings, nil)
manifest.CompletedAt = timePtr(r.timestamp())
return r.writeManifest("chunk/manifest.json", coreworkspace.ChunkManifest{
StageManifest: manifest,
ChunkCount: len(chunks),
})
}
func (r *WorkspaceRecorder) ChunkRejected(moduleKey string, sourceDigest string, rejected contracts.RejectedOutput) error {
manifest := r.newStageManifest(coreworkspace.StageChunk, coreworkspace.StatusSucceededWithRejections)
manifest.ModuleKey = moduleKey
manifest.DependencyFingerprints = workspaceFingerprints(digestFingerprints("source_document", sourceDigest))
manifest.ValidationStatus = "rejected"
manifest.Rejections = rejectionSummaries([]contracts.RejectedOutput{rejected})
manifest.CompletedAt = timePtr(r.timestamp())
return r.writeManifest("chunk/manifest.json", coreworkspace.ChunkManifest{StageManifest: manifest})
}
func (r *WorkspaceRecorder) ChunkFailed(moduleKey string, sourceDigest string, err error) error {
manifest := r.newStageManifest(coreworkspace.StageChunk, coreworkspace.StatusFailed)
manifest.ModuleKey = moduleKey
manifest.DependencyFingerprints = workspaceFingerprints(digestFingerprints("source_document", sourceDigest))
manifest.CompletedAt = timePtr(r.timestamp())
manifest.Metadata = errorMetadata(err)
return r.writeManifest("chunk/manifest.json", coreworkspace.ChunkManifest{StageManifest: manifest})
}
func (r *WorkspaceRecorder) ExtractRunning(laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint) error {
manifest := r.laneManifest(coreworkspace.StageExtract, coreworkspace.StatusRunning, laneID, moduleKey, dependencies)
manifest.StartedAt = timePtr(r.timestamp())
@@ -245,23 +197,6 @@ type sourceDocumentEnvelope struct {
Document source.SourceDocument `json:"document"`
}
type chunksEnvelope struct {
Chunks []chunkEnvelope `json:"chunks"`
Warnings []contracts.Warning `json:"warnings,omitempty"`
}
type chunkEnvelope struct {
ID string `json:"id"`
SourceID string `json:"source_id"`
Index int `json:"index"`
Ref source.SourceRef `json:"ref"`
Content binaryEnvelope `json:"content"`
Units []source.SourceUnit `json:"units,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
Annotations source.ChunkAnnotations `json:"annotations,omitempty"`
PlanAnnotations source.ChunkAnnotations `json:"plan_annotations,omitempty"`
}
type binaryEnvelope struct {
ContentBase64 string `json:"content_base64,omitempty"`
ContentDigest string `json:"content_digest,omitempty"`
@@ -315,27 +250,6 @@ func artifactOutputDigests(outputs []pipeline.CheckpointArtifact) []pipeline.Che
return normalizeFingerprints(values)
}
func chunkEnvelopes(chunks []source.Chunk) []chunkEnvelope {
if len(chunks) == 0 {
return nil
}
out := make([]chunkEnvelope, 0, len(chunks))
for _, chunk := range chunks {
out = append(out, chunkEnvelope{
ID: chunk.ID,
SourceID: chunk.SourceID,
Index: chunk.Index,
Ref: chunk.Ref,
Content: binaryEnvelopeFromContent(chunk.Content, chunk.MediaType, chunk.Metadata, nil),
Units: cloneSourceUnits(chunk.Units),
Metadata: cloneMetadata(chunk.Metadata),
Annotations: source.CloneChunkAnnotations(chunk.Annotations),
PlanAnnotations: source.CloneChunkAnnotations(chunk.PlanAnnotations),
})
}
return out
}
func binaryEnvelopeFromContent(content []byte, mediaType string, metadata map[string]any, warnings []contracts.Warning) binaryEnvelope {
return binaryEnvelope{
ContentBase64: base64.StdEncoding.EncodeToString(content),
@@ -394,21 +308,6 @@ func cloneMetadata(metadata map[string]any) map[string]any {
return out
}
func chunkOutputDigests(chunks []source.Chunk) ([]pipeline.CheckpointFingerprint, error) {
values := make([]pipeline.CheckpointFingerprint, 0, len(chunks))
for _, chunk := range chunks {
digest, err := source.DigestChunk(chunk)
if err != nil {
return nil, fmt.Errorf("chunk %q: %w", chunk.ID, err)
}
values = append(values, pipeline.CheckpointFingerprint{
Name: chunk.ID,
Value: digest,
})
}
return normalizeFingerprints(values), nil
}
func digestFingerprints(name string, digest string) []pipeline.CheckpointFingerprint {
digest = strings.TrimSpace(digest)
if digest == "" {