feat: add workflow stage resume

[andreatgretel]

📋 Summary

Adds stage-level resume support for chained workflows so compatible completed stages can be reused, matching partial stages can continue through the existing single-stage resume path, and downstream stages rerun when upstream outputs change.

🔗 Related Issue

N/A

🔄 Changes

• Add CompositeWorkflow.run(resume=...) with completed-stage reuse and partial-stage delegation.
• Invalidate downstream stages when an upstream stage reruns, changes, or has missing selected/callback output.
• Harden workflow metadata writes and let ResumeMode.IF_POSSIBLE fall back to fresh runs when prior metadata is unusable.
• Add public workflow resume tests for skip, rerun, partial/failed-stage resume, callback-output, output-processor, completed-empty, corrupt metadata, and strict resume behavior.
• Document workflow resume behavior in MkDocs and Fern docs.
• Update the workflow chaining plan with the completed stage-level resume slice.

🧪 Testing

• make test passes
• Unit tests added/updated
• E2E tests added/updated (if applicable)

Ran:

• .venv/bin/ruff format .
• .venv/bin/ruff check packages/data-designer/src/data_designer/interface/composite_workflow.py packages/data-designer/tests/interface/test_composite_workflow.py
• .venv/bin/ruff format --check packages/data-designer/src/data_designer/interface/composite_workflow.py packages/data-designer/tests/interface/test_composite_workflow.py
• .venv/bin/pytest packages/data-designer/tests/interface/test_composite_workflow.py -q - 55 passed, 2 warnings
• .venv/bin/pytest /home/ubuntu/Code/reviews/DataDesigner-747/smoke_test.py -q -s - 2 passed against NVIDIA Build (nvidia/nemotron-3-nano-30b-a3b) and NVIDIA Inference (openai/openai/gpt-5.4-nano) using /home/ubuntu/Code/.env

Note: full .venv/bin/ruff check --fix . currently hits an unrelated existing generated-notebook lint in docs/colab_notebooks/7-nemotron-personas.ipynb (F404).

✅ Checklist

• Follows commit message conventions
• Commits are signed off (DCO)
• Architecture docs updated (if applicable)

[github-actions]

Fern preview: https://nvidia-preview-pr-747.docs.buildwithfern.com/nemo/datadesigner

Fern previews include the docs-website version archive with PR changes synced into latest. Notebook tutorials are rendered without execution outputs in previews.

[greptile-apps]

Greptile Summary

This PR adds stage-level resume support to CompositeWorkflow.run(resume=...). Completed stages with matching fingerprints are reused, partially-run stages delegate to the single-stage resume path, and any stage that runs marks all descendants for forced rerun via force_rerun_downstream.

• Core logic (composite_workflow.py): introduces _read_prior_workflow_metadata, _can_skip_prior_stage, _get_prior_stage_metadata, _stage_result_from_metadata, and supporting helpers for path normalisation; hardens metadata writes with a pid/uuid-namespaced temp file and os.replace for atomicity.
• Tests (test_composite_workflow.py): adds 12 new parametrised tests covering skip, changed-stage rerun, partial/failed-stage resume, missing callback output, moved artifacts, corrupt metadata, and strict-resume rejection.
• Docs (workflow-chaining.mdx, workflow-chaining.md): adds a ## Resume section and updates the current-limits list and plan status.

Confidence Score: 5/5

Safe to merge. The resume logic is well-encapsulated, the existing fresh-run path is unchanged, and the new code is gated entirely behind the new resume parameter which defaults to ResumeMode.NEVER.

The fingerprint comparison, downstream invalidation via force_rerun_downstream, callback-output validation in _can_skip_prior_stage, and the ALWAYS/IF_POSSIBLE branching all behave as intended across the 12 new tests. Path normalisation to relative values is handled symmetrically in write and read helpers. The atomic metadata write is correct. No logic errors were found.

No files require special attention.

Important Files Changed

Filename	Overview
packages/data-designer/src/data_designer/interface/composite_workflow.py	Adds resume orchestration: fingerprint comparison, skip/delegate/rerun branching, relative-path metadata, and atomic temp-file metadata writes. Logic is internally consistent and well-tested.
packages/data-designer/tests/interface/test_composite_workflow.py	Adds 12 new resume tests; _mark_stage_resumable helper correctly strips completion-only fields. All documented scenarios are covered.
fern/versions/latest/pages/concepts/workflow-chaining.mdx	Adds a Resume section and removes the now-obsolete 'Stage-level resume is not implemented yet' limit bullet.
plans/workflow-chaining/workflow-chaining.md	Updates the plan status with what was implemented in this slice; no logic changes.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["workflow.run(resume)"] --> B{resume == NEVER}
    B -- Yes --> C[prior_metadata = None]
    B -- No --> D[_read_prior_workflow_metadata]
    D --> E{readable and name matches}
    E -- No or corrupt --> F{ALWAYS mode}
    F -- Yes --> G[raise WorkflowError]
    F -- No --> C
    E -- Yes --> H[prior_metadata loaded]
    C --> I
    H --> I
    I[For each stage] --> J{skipped_upstream_stage set}
    J -- Yes --> K[status = skipped_empty_upstream]
    K --> I
    J -- No --> L[Compute fingerprint]
    L --> M{prior_matches AND _can_skip_prior_stage}
    M -- Yes --> N[Reuse stage output from metadata]
    N --> O{status == completed_empty}
    O -- Yes --> P[Set skipped_upstream_stage]
    O -- No --> I
    P --> I
    M -- No --> Q{prior_matches AND status running/failed AND path exists}
    Q -- Yes --> R[stage_resume = ALWAYS]
    Q -- No --> S{resume==ALWAYS AND NOT force_rerun_downstream}
    S -- Yes --> T[raise WorkflowError]
    S -- No --> U[stage_resume = NEVER / rmtree]
    R --> V
    U --> V
    V["data_designer.create(stage_resume)"] --> W[force_rerun_downstream = True]
    W --> I
    I --> X[CompositeWorkflowResults]

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A["workflow.run(resume)"] --> B{resume == NEVER}
    B -- Yes --> C[prior_metadata = None]
    B -- No --> D[_read_prior_workflow_metadata]
    D --> E{readable and name matches}
    E -- No or corrupt --> F{ALWAYS mode}
    F -- Yes --> G[raise WorkflowError]
    F -- No --> C
    E -- Yes --> H[prior_metadata loaded]
    C --> I
    H --> I
    I[For each stage] --> J{skipped_upstream_stage set}
    J -- Yes --> K[status = skipped_empty_upstream]
    K --> I
    J -- No --> L[Compute fingerprint]
    L --> M{prior_matches AND _can_skip_prior_stage}
    M -- Yes --> N[Reuse stage output from metadata]
    N --> O{status == completed_empty}
    O -- Yes --> P[Set skipped_upstream_stage]
    O -- No --> I
    P --> I
    M -- No --> Q{prior_matches AND status running/failed AND path exists}
    Q -- Yes --> R[stage_resume = ALWAYS]
    Q -- No --> S{resume==ALWAYS AND NOT force_rerun_downstream}
    S -- Yes --> T[raise WorkflowError]
    S -- No --> U[stage_resume = NEVER / rmtree]
    R --> V
    U --> V
    V["data_designer.create(stage_resume)"] --> W[force_rerun_downstream = True]
    W --> I
    I --> X[CompositeWorkflowResults]

Loading

_{Reviews (5): Last reviewed commit: "merge main into workflow resume branch" | Re-trigger Greptile}

[github-actions]

Review: PR #747 — feat: add workflow stage resume

Summary

Adds stage-level resume to CompositeWorkflow.run(resume=...):

• Reuses compatible completed stages (fingerprint match + output verifiable).
• Delegates partial stages (running / failed) to the existing single-stage resume path via DataDesigner.create(..., resume=ResumeMode.ALWAYS).
• Invalidates downstream stages on first re-run, change, or missing output.
• Adds atomic, fsync'd workflow-metadata.json writes via tmp+os.replace.
• Falls back to fresh runs in ResumeMode.IF_POSSIBLE when prior metadata is missing/corrupt; raises in ResumeMode.ALWAYS.
• Tests for skip / rerun / partial-resume / callback-output / output-processor / completed-empty / corrupt metadata / strict-resume.
• Docs added in both mkdocs and Fern; phase-3 plan updated.

Diff is well-scoped (5 files, +503/−10) and lives entirely in the interface layer — no engine/config touch, no import-direction violations.

Findings

Correctness — Medium

1. output_seed_path is persisted and reused as an absolute path (composite_workflow.py:400, consumed at composite_workflow.py:309). If a user moves their artifact root (rsync, archive, container mount), prior_metadata still matches by name/stage_dir, but output_seed_path points at the old location while _stage_result_from_metadata builds an ArtifactStorage rooted at the new workflow_path. The result is a mixed-state read (old data, new storage object) or a _count_parquet_records raise — silently invalidating reuse in the best case, confusing behavior in the worst. Storing this as a path relative to workflow_path would be more robust. (Same applies to callback_output_path.)

2. ResumeMode.ALWAYS semantics drift once any stage is re-run. The added doc says: "If a stage changed or its selected output is missing, the workflow raises instead of starting fresh." But once an ALWAYS resume successfully resumes a partial upstream stage from checkpoint (stage_resume == ALWAYS, line 328), force_rerun_downstream flips to True and the downstream elif resume == ResumeMode.ALWAYS and not force_rerun_downstream: guard (line 330) is bypassed — so a downstream stage whose fingerprint also differs from prior metadata silently runs fresh instead of raising. That may be the intended behavior, but it deviates from the doc and from the strict reading of the test_composite_workflow_resume_always_rejects_changed_stage contract. Consider tightening the doc and adding a test for "ALWAYS, upstream resumed-from-checkpoint, downstream fingerprint differs".

Correctness — Low

3. _load_stage_analysis swallows every exception (composite_workflow.py:552):

   try:
       return DatasetProfilerResults.model_validate({...})
   except Exception:
       return None

   Bare except Exception masks both ValidationError (the only one you'd plausibly recover from) and unrelated bugs introduced by future schema changes. Narrow it to pydantic.ValidationError so a real bug doesn't decay into a silent analysis=None.

4. stage_metadata.update(prior_stage_metadata) on the skip path (line 308) wholesale-imports every key the prior run wrote, including ones not produced by the current run (config, seed_path, seeded_from_stage, num_records_requested, duration_sec). That's by design here, but means the new metadata file's seeded_from_stage for a re-skipped stage may name a stage from the prior run that no longer exists in the current workflow definition. Fine for inspection; worth a note if later phases lean on those fields.

5. _stage_result_from_metadata returns DatasetMetadata() — empty (line 528). The original run may have collected real DatasetMetadata. Reusing the cached stage exposes a stripped-down DatasetCreationResults to user code. Document or persist+rehydrate.

Style / Nits

6. Unused import in TYPE_CHECKING removal: DatasetProfilerResults was lifted out of TYPE_CHECKING (line 18) because it's now used at module scope by _load_stage_analysis. That's correct, but DatasetProfilerResults is heavy; the project lazy-loads pandas already (lazy.pd). Confirm that data_designer.config.analysis.dataset_profiler doesn't drag in numpy/pandas at import — if it does, this regresses the import-time profile (make perf-import would catch it).

7. tmp_path = path.with_name(f"{path.name}.tmp.{os.getpid()}") (line 706): if two threads in the same process both write workflow metadata for the same workflow (uncommon but possible during async cleanup), they collide on the temp filename. Adding a uuid suffix would be safer; PID alone is not unique within a process. Probably out of scope.

8. force_rerun_downstream is largely redundant with the upstream-fingerprint chain (since each stage's fingerprint folds in upstream_fingerprint). Its real job is to short-circuit the ResumeMode.ALWAYS raise on line 330. A short comment to that effect would help future readers — right now the flag's necessity is non-obvious.

Tests

Coverage is strong. Suggestions only:

• Add a test pinning the ALWAYS-after-partial-upstream-resume behavior (see finding DataDesigner.make_seed_reference_from_file doesn't support paths with multiple parquet partition #2) — whichever way the project decides on, lock it in.
• test_composite_workflow_resume_if_possible_skips_stage_with_output_processors checks main-batch mtime; consider also asserting the output-processors directory wasn't touched (its mtime is the actual cache-hit signal for the output-processor work).
• Worth a regression test for prior metadata whose stages length exceeds the current workflow's stage list — _get_prior_stage_metadata handles index >= len(stages), but the inverse case (current workflow has fewer stages than prior) isn't explicitly covered.

Security

No new attack surface. Metadata is JSON-validated on read and treated as data. The os.fsync + os.replace pattern is the right defense against torn writes from crashes mid-run.

Performance

• One extra _count_parquet_records call per skippable stage (validation in _can_skip_prior_stage + the actual count in run). Cheap (parquet metadata, not row read), but doubling it is unnecessary — could cache the count or skip the validation since the count immediately follows.
• _stage_result_from_metadata instantiates ArtifactStorage with resume=ResumeMode.ALWAYS, which triggers the resolved_dataset_name check. Fine for hits; on misses you'll get an ArtifactStorageError from the validator instead of going through DataDesignerWorkflowError. Consider catching and re-wrapping at the boundary so callers see a single error type per the project's "Errors normalize at boundaries" invariant.

Docs

• docs/concepts/workflow-chaining.md and the Fern mirror are updated symmetrically — good.
• Phase-3 plan update accurately describes the slice and remaining deferred items.
• One precision nit: the ResumeMode.ALWAYS blurb should clarify the "first changed stage raises; downstream of a checkpoint-resumed stage runs fresh" behavior (or close the gap).

Verdict

Solid, well-tested addition. The core logic is correct and the new behavior is opt-in through resume=, so risk to existing callers is minimal. Address #1 (relative paths) before this is relied on for portable artifacts; #2 (ALWAYS doc/test) and #3 (narrow the except) are easy follow-ups. Everything else is taste.

Recommendation: approve with the path-portability and except Exception narrowing addressed (or filed as follow-ups). Other findings are non-blocking polish.