fix(checkpoint): don't KeyError on manifest stages without a canonical artifact

_validate_artifacts_for_stage looked up CANONICAL_STAGE_ARTIFACTS[stage]
unconditionally, but the valid stage list comes from the pipeline manifest via
get_pipeline_stages(), which declares stages beyond the 9 canonical ones — e.g.
character-animation adds `character_design`/`rig_plan`. Such a stage passes the
`stage in valid_stages` guard, then raised an unhandled KeyError on the
canonical lookup, so those stages could never be checkpointed (the crash hits
write/read_checkpoint and friends, even for in_progress checkpoints).

Look the canonical artifact up defensively with `.get()` and skip the
required-artifact check when there is none. Canonical stages still require their
artifact when completed.
This commit is contained in:
0xDevNinja
2026-07-07 15:29:03 +05:30
parent bcd8eb6e53
commit 5dcd026ef7
2 changed files with 73 additions and 2 deletions

View File

@@ -106,8 +106,17 @@ def _validate_artifacts_for_stage(
status: str,
artifacts: dict[str, Any],
) -> None:
required_artifact = CANONICAL_STAGE_ARTIFACTS[stage]
if status in {"completed", "awaiting_human"} and required_artifact not in artifacts:
# Valid stages come from the pipeline manifest (get_pipeline_stages), which
# can declare stages beyond the 9 canonical ones (e.g. character-animation's
# `character_design`/`rig_plan`, screen-demo's `real_capture`). Those have no
# canonical artifact, so look it up defensively — a missing entry means the
# stage simply has no required artifact, not a crash.
required_artifact = CANONICAL_STAGE_ARTIFACTS.get(stage)
if (
required_artifact is not None
and status in {"completed", "awaiting_human"}
and required_artifact not in artifacts
):
raise CheckpointValidationError(
f"Stage {stage!r} with status {status!r} must include "
f"canonical artifact {required_artifact!r}"