eric/notarius
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement selective checkpoint recomputation

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-21 21:59:48 +00:00
parent 22d4f29670
commit c437682407
17 changed files with 902 additions and 180 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
internal/cli/recompute_policy_test.go Normal file
View File

@@ -0,0 +1,56 @@
package cli
import (
"testing"
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
)
func TestRecomputePolicyIncludesDependentLanesAndReusablePredecessors(t *testing.T) {
producer := pipeline.ResolvedArtifactLane{StepID: "first", ID: "producer"}
unrelated := pipeline.ResolvedArtifactLane{StepID: "first", ID: "unrelated"}
consumer := pipeline.ResolvedArtifactLane{
StepID: "second", ID: "consumer",
ExtractReferences: pipeline.ResolvedReferenceTarget{Bindings: []pipeline.ReferenceBinding{{Artifact: &pipeline.ArtifactReference{Step: "first", Lane: "producer"}}}},
}
downstream := pipeline.ResolvedArtifactLane{
StepID: "third", ID: "downstream",
ExtractReferences: pipeline.ResolvedReferenceTarget{Bindings: []pipeline.ReferenceBinding{{Artifact: &pipeline.ArtifactReference{Step: "second", Lane: "consumer"}}}},
}
independent := pipeline.ResolvedArtifactLane{StepID: "third", ID: "independent"}
resolved := pipeline.ResolvedPipeline{Steps: []pipeline.ResolvedPipelineStep{
{ID: "first", ArtifactLanes: []pipeline.ResolvedArtifactLane{producer, unrelated}},
{ID: "second", ArtifactLanes: []pipeline.ResolvedArtifactLane{consumer}},
{ID: "third", ArtifactLanes: []pipeline.ResolvedArtifactLane{downstream, independent}},
}}
policy, err := recomputePolicy(resolved, "second")
if err != nil {
t.Fatal(err)
}
if _, ok := policy.ForcedLanes[pipeline.CheckpointLaneKey("second", "consumer")]; !ok {
t.Fatal("selected lane was not forced")
}
if _, ok := policy.ForcedLanes[pipeline.CheckpointLaneKey("third", "downstream")]; !ok {
t.Fatal("transitive dependent lane was not forced")
}
if _, ok := policy.ForcedLanes[pipeline.CheckpointLaneKey("first", "producer")]; ok {
t.Fatal("predecessor was implicitly forced")
}
if _, ok := policy.RequireReusableLanes[pipeline.CheckpointLaneKey("first", "producer")]; !ok {
t.Fatal("required predecessor was not marked reusable")
}
if _, ok := policy.ForcedLanes[pipeline.CheckpointLaneKey("first", "unrelated")]; ok {
t.Fatal("unrelated lane was forced")
}
if _, ok := policy.ForcedLanes[pipeline.CheckpointLaneKey("third", "independent")]; ok {
t.Fatal("unrelated later lane was forced")
}
}
func TestRecomputePolicyRejectsUnknownStep(t *testing.T) {
_, err := recomputePolicy(pipeline.ResolvedPipeline{Steps: []pipeline.ResolvedPipelineStep{{ID: "known"}}}, "missing")
if err == nil {
t.Fatal("unknown step was accepted")
}
}

171
internal/cli/run.go
View File

@@ -28,7 +28,7 @@ import (
const defaultConfigPath = "/usr/local/etc/notarius/config.yml"
const usage = `Usage:
notarius help
notarius run <pipeline-id> --input path/to/source.json [--config path/to/config.yml] [--output-dir path] [--chunk_cache auto|bypass|refresh] [--resume] [--debug [--debug-dir path]] [--only lane-a,lane-b] [--session-id id] [--reference selector=path] [--without-reference selector]
notarius run <pipeline-id> --input path/to/source.json [--config path/to/config.yml] [--output-dir path] [--chunk_cache auto|bypass|refresh] [--resume] [--recompute-step step-id] [--debug [--debug-dir path]] [--only lane-a,lane-b] [--session-id id] [--reference selector=path] [--without-reference selector]
notarius config validate --config path/to/config.yml [--pipeline pipeline-id] [--only lane-a,lane-b]
notarius pipelines list --config path/to/config.yml [--json]
`
@@ -136,6 +136,7 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
debugDir := fs.String("debug-dir", "", "debug bundle directory")
llmProfile := fs.String("llm-profile", "", "LLM profile override")
resume := fs.Bool("resume", false, "reuse compatible recorded checkpoints")
recomputeStep := singleValueFlag{}
chunkCache := chunkCacheFlag{}
sessionID := sessionIDFlag{}
referenceFlags := stringListFlag{}
@@ -144,6 +145,7 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
fs.Var(&chunkCache, "chunk_cache", "chunk plan cache mode: auto, bypass, or refresh")
fs.Var(&referenceFlags, "reference", "reference binding, as slot=path, chunk.slot=path, merge.slot=path, lane.slot=path, lane.extract.slot=path, lane.merge.slot=path, or lane.normalize.slot=path")
fs.Var(&withoutReferenceFlags, "without-reference", "unbind a reference, using the same selector forms as --reference")
fs.Var(&recomputeStep, "recompute-step", "recompute one ordered pipeline step and dependent lanes")
if err := validateRunFlagValues(args); err != nil {
fmt.Fprintf(stderr, "notarius: %v\n", err)
return 2
@@ -190,6 +192,14 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
fmt.Fprintf(stderr, "notarius: %v\n", err)
return 2
}
if recomputeStep.set && strings.TrimSpace(recomputeStep.value) == "" {
fmt.Fprintln(stderr, "notarius: --recompute-step must not be empty")
return 2
}
if recomputeStep.set && len(only) > 0 {
fmt.Fprintln(stderr, "notarius: --recompute-step cannot be combined with --only")
return 2
}
referenceRequests, err := parseReferenceFlags(referenceFlags)
if err != nil {
fmt.Fprintf(stderr, "notarius: %v\n", err)
@@ -223,6 +233,14 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
fmt.Fprintln(stderr, "notarius: --resume requires cache.checkpoints.enabled: true")
return 1
}
if recomputeStep.set && !*resume {
fmt.Fprintln(stderr, "notarius: --recompute-step requires --resume")
return 2
}
if recomputeStep.set && !cfg.Cache.Checkpoints.Enabled {
fmt.Fprintln(stderr, "notarius: --recompute-step requires cache.checkpoints.enabled: true")
return 1
}
startedAt := opts.Now().UTC()
runID, err := opts.RunIDGenerator(startedAt)
@@ -267,6 +285,7 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
OnlyLanes: append([]string(nil), only...),
ChunkCacheOverride: chunkCache.explicitValue(),
Resume: *resume,
RecomputeStep: strings.TrimSpace(recomputeStep.value),
RunID: runID,
StartedAt: startedAt,
}
@@ -309,6 +328,13 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
return failPipelineCommand(stderr, commandState, terminalWriter, err)
}
effective.ResolvedPipeline = materialized
checkpointPolicy := pipeline.CheckpointExecutionPolicy{}
if recomputeStep.set {
checkpointPolicy, err = recomputePolicy(effective.ResolvedPipeline, recomputeStep.value)
if err != nil {
return failPipelineCommand(stderr, commandState, terminalWriter, err)
}
}
invocation.PipelineDigest = effective.ResolvedPipeline.Digest
if err := writeSummary(summary, func() error { return summary.WriteInvocation(invocation) }); err != nil {
return failPipelineCommand(stderr, commandState, terminalWriter, fmt.Errorf("write debug invocation metadata: %w", err))
@@ -358,21 +384,22 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
}
output, err := pipeline.New().Run(ctx, pipeline.RunInput{
Prepared: prepared,
Path: strings.TrimSpace(*inputPath),
RawInput: rawInput,
SessionID: strings.TrimSpace(sessionID.value),
RunID: runID,
StartedAt: startedAt,
LLMProfiles: llmProfiles,
Metadata: runMetadata(effective.Config.Output.Directory, debugPath),
Warnings: referenceWarnings,
ChunkCacheMode: effective.Config.Cache.ChunkPlans.Mode,
ChunkPlans: chunkPlans,
Checkpoints: checkpointRecorder,
Checkpoint: checkpointLoader,
Debug: debugRecorder,
ExtractWorkers: cfg.Concurrency.StageWorkers["extract"],
Prepared: prepared,
Path: strings.TrimSpace(*inputPath),
RawInput: rawInput,
SessionID: strings.TrimSpace(sessionID.value),
RunID: runID,
StartedAt: startedAt,
LLMProfiles: llmProfiles,
Metadata: runMetadata(effective.Config.Output.Directory, debugPath),
Warnings: referenceWarnings,
ChunkCacheMode: effective.Config.Cache.ChunkPlans.Mode,
ChunkPlans: chunkPlans,
Checkpoints: checkpointRecorder,
Checkpoint: checkpointLoader,
CheckpointPolicy: checkpointPolicy,
Debug: debugRecorder,
ExtractWorkers: cfg.Concurrency.StageWorkers["extract"],
})
commandState.observeOutput(output)
if err != nil {
@@ -481,6 +508,95 @@ func checkpointHandlersForRun(
return recorder, loader, nil
}
func recomputePolicy(resolved pipeline.ResolvedPipeline, requestedStep string) (pipeline.CheckpointExecutionPolicy, error) {
requestedStep = strings.TrimSpace(requestedStep)
if requestedStep == "" {
return pipeline.CheckpointExecutionPolicy{}, fmt.Errorf("--recompute-step must not be empty")
}
var selected *pipeline.ResolvedPipelineStep
for index := range resolved.Steps {
if strings.TrimSpace(resolved.Steps[index].ID) == requestedStep {
selected = &resolved.Steps[index]
break
}
}
if selected == nil {
return pipeline.CheckpointExecutionPolicy{}, fmt.Errorf("unknown pipeline step %q", requestedStep)
}
policy := pipeline.CheckpointExecutionPolicy{ForcedLanes: make(map[string]struct{}), RequireReusableLanes: make(map[string]struct{})}
var forced []pipeline.ResolvedArtifactLane
for _, lane := range selected.ArtifactLanes {
lane.StepID = selected.ID
policy.ForcedLanes[pipeline.CheckpointLaneKey(selected.ID, lane.ID)] = struct{}{}
forced = append(forced, lane)
}
lanes := make(map[string]pipeline.ResolvedArtifactLane)
dependents := make(map[string][]pipeline.ResolvedArtifactLane)
for _, step := range resolved.Steps {
for _, lane := range step.ArtifactLanes {
lane.StepID = step.ID
lanes[pipeline.CheckpointLaneKey(step.ID, lane.ID)] = lane
for _, target := range []pipeline.ResolvedReferenceTarget{lane.ExtractReferences, lane.MergeReferences, lane.NormalizeReferences} {
for _, binding := range target.Bindings {
if binding.Artifact != nil {
producerKey := pipeline.CheckpointLaneKey(binding.Artifact.Step, binding.Artifact.Lane)
dependents[producerKey] = append(dependents[producerKey], lane)
}
}
}
}
}
processed := make(map[string]struct{})
queue := append([]pipeline.ResolvedArtifactLane(nil), forced...)
for len(queue) > 0 {
lane := queue[0]
queue = queue[1:]
key := pipeline.CheckpointLaneKey(lane.StepID, lane.ID)
if _, ok := processed[key]; ok {
continue
}
processed[key] = struct{}{}
for _, dependent := range dependents[key] {
dependentKey := pipeline.CheckpointLaneKey(dependent.StepID, dependent.ID)
if _, alreadyForced := policy.ForcedLanes[dependentKey]; alreadyForced {
continue
}
policy.ForcedLanes[dependentKey] = struct{}{}
forced = append(forced, dependent)
queue = append(queue, dependent)
}
}
processed = make(map[string]struct{})
queue = append([]pipeline.ResolvedArtifactLane(nil), forced...)
for len(queue) > 0 {
lane := queue[0]
queue = queue[1:]
key := pipeline.CheckpointLaneKey(lane.StepID, lane.ID)
if _, ok := processed[key]; ok {
continue
}
processed[key] = struct{}{}
for _, target := range []pipeline.ResolvedReferenceTarget{lane.ExtractReferences, lane.MergeReferences, lane.NormalizeReferences} {
for _, binding := range target.Bindings {
if binding.Artifact == nil {
continue
}
producerKey := pipeline.CheckpointLaneKey(binding.Artifact.Step, binding.Artifact.Lane)
if _, forcedAlready := policy.ForcedLanes[producerKey]; !forcedAlready {
policy.RequireReusableLanes[producerKey] = struct{}{}
}
if producer, ok := lanes[producerKey]; ok {
queue = append(queue, producer)
}
}
}
}
return policy, nil
}
func checkpointIdentityFingerprints(values []pipeline.CheckpointFingerprint) []checkpoint.Fingerprint {
if len(values) == 0 {
return nil
@@ -661,7 +777,7 @@ func reorderRunArgs(args []string) []string {
func runFlagTakesValue(arg string) bool {
switch arg {
case "--config", "--input", "--only", "--output-dir", "--debug-dir", "--llm-profile", "--session-id", "--chunk_cache", "--reference", "--without-reference":
case "--config", "--input", "--only", "--output-dir", "--debug-dir", "--llm-profile", "--session-id", "--chunk_cache", "--reference", "--without-reference", "--recompute-step":
return true
default:
return false
@@ -1008,6 +1124,27 @@ type sessionIDFlag struct {
set bool
}
type singleValueFlag struct {
value string
set bool
}
func (flag *singleValueFlag) String() string {
if flag == nil {
return ""
}
return flag.value
}
func (flag *singleValueFlag) Set(value string) error {
if flag.set {
return fmt.Errorf("--recompute-step may be specified only once")
}
flag.value = value
flag.set = true
return nil
}
func (flag *sessionIDFlag) String() string {
if flag == nil {
return ""