data-designer 0.9.2__tar.gz → 0.9.3__tar.gz

{data_designer-0.9.2 → data_designer-0.9.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: data-designer
-Version: 0.9.2
+Version: 0.9.3
 Summary: General framework for synthetic data generation
 License-Expression: Apache-2.0
 Classifier: Development Status :: 4 - Beta
@@ -15,8 +15,8 @@ Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development
 Requires-Python: >=3.10
-Requires-Dist: data-designer-config==0.9.2
-Requires-Dist: data-designer-engine==0.9.2
+Requires-Dist: data-designer-config==0.9.3
+Requires-Dist: data-designer-engine==0.9.3
 Requires-Dist: huggingface-hub<2,>=1.0.1
 Requires-Dist: opentelemetry-api<1.44,>=1.43
 Requires-Dist: opentelemetry-exporter-prometheus<0.65,>=0.64b0

{data_designer-0.9.2 → data_designer-0.9.3}/src/data_designer/integrations/huggingface/client.py

@@ -119,7 +119,7 @@ class HuggingFaceHubClient:
             URL to the uploaded dataset
 
         Raises:
-            HuggingFaceUploadError: If validation fails or upload encounters errors
+            HuggingFaceHubClientUploadError: If validation fails or upload encounters errors
         """
         logger.info(f"🤗 Uploading dataset to Hugging Face Hub: {repo_id}")
 
@@ -162,7 +162,7 @@ class HuggingFaceHubClient:
             private: Whether to create private repo
 
         Raises:
-            HuggingFaceUploadError: If repository creation fails
+            HuggingFaceHubClientUploadError: If repository creation fails
         """
         logger.info(f"{LOG_INDENT}{RandomEmoji.working()} Checking if repository exists...")
         try:
@@ -203,7 +203,7 @@ class HuggingFaceHubClient:
             parquet_folder: Path to folder containing parquet files
 
         Raises:
-            HuggingFaceUploadError: If upload fails
+            HuggingFaceHubClientUploadError: If upload fails
         """
         logger.info(f"{LOG_INDENT}{RandomEmoji.loading()} Uploading main dataset files...")
         try:
@@ -225,7 +225,7 @@ class HuggingFaceHubClient:
             images_folder: Path to images folder
 
         Raises:
-            HuggingFaceUploadError: If upload fails
+            HuggingFaceHubClientUploadError: If upload fails
         """
         if not images_folder.exists():
             return
@@ -255,7 +255,7 @@ class HuggingFaceHubClient:
             processors_folder: Path to folder containing processor output directories
 
         Raises:
-            HuggingFaceUploadError: If upload fails
+            HuggingFaceHubClientUploadError: If upload fails
         """
         if not processors_folder.exists():
             return
@@ -290,7 +290,7 @@ class HuggingFaceHubClient:
             builder_config_path: Path to builder_config.json file
 
         Raises:
-            HuggingFaceUploadError: If upload fails
+            HuggingFaceHubClientUploadError: If upload fails
         """
         logger.info(f"{LOG_INDENT}{RandomEmoji.loading()} Uploading configuration files...")
 
@@ -345,7 +345,7 @@ class HuggingFaceHubClient:
             tags: Additional custom tags for the dataset
 
         Raises:
-            HuggingFaceUploadError: If dataset card generation or upload fails
+            HuggingFaceHubClientUploadError: If dataset card generation or upload fails
         """
         try:
             with open(metadata_path) as f:
@@ -455,7 +455,7 @@ class HuggingFaceHubClient:
             base_dataset_path: Path to dataset directory
 
         Raises:
-            HuggingFaceUploadError: If directory structure is invalid
+            HuggingFaceHubClientUploadError: If directory structure is invalid
         """
         if not base_dataset_path.exists():
             raise HuggingFaceHubClientUploadError(f"Dataset path does not exist: {base_dataset_path}")

{data_designer-0.9.2 → data_designer-0.9.3}/src/data_designer/interface/__init__.py

@@ -22,19 +22,41 @@ if TYPE_CHECKING:
         DataDesignerWorkflowError,
     )
     from data_designer.interface.results import DatasetCreationResults  # noqa: F401
+    from data_designer.interface.workflow_metadata import (  # noqa: F401
+        CompletedWorkflowStageMetadata,
+        FailedWorkflowStageMetadata,
+        RunningWorkflowStageMetadata,
+        SkippedWorkflowStageMetadata,
+        WorkflowMetadata,
+        WorkflowStageMetadata,
+        validate_workflow_stage_metadata,
+    )
 
 _LAZY_IMPORTS: dict[str, tuple[str, str]] = {
     "CompositeWorkflow": ("data_designer.interface.composite_workflow", "CompositeWorkflow"),
     "CompositeWorkflowResults": ("data_designer.interface.composite_workflow", "CompositeWorkflowResults"),
+    "CompletedWorkflowStageMetadata": (
+        "data_designer.interface.workflow_metadata",
+        "CompletedWorkflowStageMetadata",
+    ),
     "DataDesigner": ("data_designer.interface.data_designer", "DataDesigner"),
     "DataDesignerEarlyShutdownError": ("data_designer.interface.errors", "DataDesignerEarlyShutdownError"),
     "DataDesignerGenerationError": ("data_designer.interface.errors", "DataDesignerGenerationError"),
     "DataDesignerProfilingError": ("data_designer.interface.errors", "DataDesignerProfilingError"),
     "DataDesignerWorkflowError": ("data_designer.interface.errors", "DataDesignerWorkflowError"),
     "DatasetCreationResults": ("data_designer.interface.results", "DatasetCreationResults"),
+    "FailedWorkflowStageMetadata": ("data_designer.interface.workflow_metadata", "FailedWorkflowStageMetadata"),
     "ResumeMode": ("data_designer.config.run_config", "ResumeMode"),
+    "RunningWorkflowStageMetadata": ("data_designer.interface.workflow_metadata", "RunningWorkflowStageMetadata"),
     "SkippedStageResult": ("data_designer.interface.composite_workflow", "SkippedStageResult"),
     "SkippedStageStatus": ("data_designer.interface.composite_workflow", "SkippedStageStatus"),
+    "SkippedWorkflowStageMetadata": ("data_designer.interface.workflow_metadata", "SkippedWorkflowStageMetadata"),
+    "WorkflowMetadata": ("data_designer.interface.workflow_metadata", "WorkflowMetadata"),
+    "WorkflowStageMetadata": ("data_designer.interface.workflow_metadata", "WorkflowStageMetadata"),
+    "validate_workflow_stage_metadata": (
+        "data_designer.interface.workflow_metadata",
+        "validate_workflow_stage_metadata",
+    ),
 }
 
 __all__ = list(_LAZY_IMPORTS.keys())

{data_designer-0.9.2 → data_designer-0.9.3}/src/data_designer/interface/composite_workflow.py

@@ -40,6 +40,7 @@ from data_designer.interface.results import (
     _export_jsonl,
     _export_parquet,
 )
+from data_designer.interface.workflow_metadata import WorkflowMetadata, validate_workflow_stage_metadata
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -219,8 +220,19 @@ class CompositeWorkflow:
         _validate_dir_name(name, "stage name")
         if any(stage.name == name for stage in self._stages):
             raise DataDesignerWorkflowError(f"Stage name {name!r} is already used in workflow {self.name!r}.")
-        if num_records is not None and num_records < 1:
-            raise DataDesignerWorkflowError("Stage num_records must be at least 1.")
+        if num_records is not None:
+            if not isinstance(num_records, int) or isinstance(num_records, bool):
+                raise DataDesignerWorkflowError("Stage num_records must be an integer.")
+            if num_records < 1:
+                raise DataDesignerWorkflowError("Stage num_records must be at least 1.")
+        if on_success_version is not None and not isinstance(on_success_version, str):
+            raise DataDesignerWorkflowError("Stage on_success_version must be a string.")
+        if not isinstance(allow_empty, bool):
+            raise DataDesignerWorkflowError("Stage allow_empty must be a boolean.")
+        if not isinstance(sampling_strategy, SamplingStrategy):
+            raise DataDesignerWorkflowError("Stage sampling_strategy must be a SamplingStrategy.")
+        if selection_strategy is not None and not isinstance(selection_strategy, (IndexRange, PartitionBlock)):
+            raise DataDesignerWorkflowError("Stage selection_strategy must be an IndexRange or PartitionBlock.")
         _validate_stage_output(output)
         output_processors = output_processors or []
         _validate_distinct_output_processors(config_builder, output_processors)
@@ -276,6 +288,7 @@ class CompositeWorkflow:
         workflow_path.mkdir(parents=True, exist_ok=True)
         prior_metadata = _read_prior_workflow_metadata(workflow_path, self.name, resume)
         metadata: dict[str, Any] = {
+            **(prior_metadata or {}),
             "name": self.name,
             "library_version": get_library_version(),
             "stages": [],
@@ -388,6 +401,10 @@ class CompositeWorkflow:
                     f"Cannot resume workflow {self.name!r}: stage {stage.name!r} is not reusable."
                 )
 
+            if stage_resume == ResumeMode.ALWAYS and prior_stage_metadata is not None:
+                prior_stage = validate_workflow_stage_metadata(prior_stage_metadata)
+                stage_metadata.update(prior_stage.model_extra or {})
+
             if stage_resume == ResumeMode.NEVER and stage_path.exists():
                 shutil.rmtree(stage_path)
 
@@ -598,6 +615,22 @@ def _read_prior_workflow_metadata(
         raise DataDesignerWorkflowError(
             f"Cannot resume workflow {workflow_name!r}: workflow metadata has invalid shape."
         )
+    try:
+        validated_metadata = WorkflowMetadata.model_validate(metadata)
+    except ValidationError as exc:
+        if resume != ResumeMode.ALWAYS:
+            error = exc.errors(include_url=False, include_input=False)[0]
+            location = ".".join(str(part) for part in error["loc"])
+            logger.warning(
+                "Workflow metadata for %r has invalid field %s (%s); starting fresh.",
+                workflow_name,
+                location,
+                error["msg"],
+            )
+            return None
+        raise DataDesignerWorkflowError(
+            f"Cannot resume workflow {workflow_name!r}: workflow metadata has invalid shape."
+        ) from exc
     if metadata.get("name") != workflow_name:
         if resume != ResumeMode.ALWAYS:
             logger.warning("Workflow metadata for %r has a different name; starting fresh.", workflow_name)
@@ -605,7 +638,7 @@ def _read_prior_workflow_metadata(
         raise DataDesignerWorkflowError(
             f"Cannot resume workflow {workflow_name!r}: workflow metadata name does not match."
         )
-    return metadata
+    return validated_metadata.model_dump(mode="json", exclude_unset=True)
 
 
 def _get_prior_stage_metadata(
@@ -870,9 +903,13 @@ def _parquet_files(path: Path) -> list[Path]:
 def _write_workflow_metadata(workflow_path: Path, metadata: dict[str, Any]) -> None:
     path = workflow_path / WORKFLOW_METADATA_FILENAME
     tmp_path = path.with_name(f"{path.name}.tmp.{os.getpid()}.{uuid.uuid4().hex}")
+    try:
+        validated_metadata = WorkflowMetadata.model_validate(metadata)
+    except ValidationError as exc:
+        raise DataDesignerWorkflowError("Workflow metadata has invalid shape.") from exc
     try:
         with tmp_path.open("w", encoding="utf-8") as f:
-            json.dump(metadata, f, indent=2, sort_keys=True)
+            json.dump(validated_metadata.model_dump(mode="json", exclude_unset=True), f, indent=2, sort_keys=True)
             f.flush()
             os.fsync(f.fileno())
         os.replace(tmp_path, path)

{data_designer-0.9.2 → data_designer-0.9.3}/src/data_designer/interface/data_designer.py

@@ -354,11 +354,10 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
             )
 
         try:
-            profiler_column_configs = config_builder.get_column_configs() or builder.data_designer_config.columns
             profiler = self._create_dataset_profiler(
                 config_builder,
                 resource_provider,
-                column_configs=profiler_column_configs,
+                column_configs=builder.data_designer_config.columns,
             )
             analysis = profiler.profile_dataset(num_records, dataset_for_profiler)
         except Exception as e:
@@ -457,7 +456,11 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
             dataset_for_profiler = processed_dataset
 
         try:
-            profiler = self._create_dataset_profiler(config_builder, resource_provider)
+            profiler = self._create_dataset_profiler(
+                config_builder,
+                resource_provider,
+                column_configs=builder.data_designer_config.columns,
+            )
             analysis = profiler.profile_dataset(num_records, dataset_for_profiler)
         except Exception as e:
             raise DataDesignerProfilingError(f"🛑 Error profiling preview dataset: {e}") from e

data_designer-0.9.3/src/data_designer/interface/workflow_metadata.py

@@ -0,0 +1,97 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Annotated, Any, Literal, TypeAlias
+
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
+
+
+class _WorkflowStageMetadataBase(BaseModel):
+    """Common metadata for a composite workflow stage."""
+
+    model_config = ConfigDict(extra="allow")
+
+    status: Literal["running", "failed", "completed", "completed_empty", "skipped_empty_upstream"]
+    index: int
+    name: str
+    stage_dir: str
+    depends_on: list[str]
+    allow_empty: bool
+    on_success_version: str | None
+    output_processors: list[dict[str, Any]]
+    output: str
+    sampling_strategy: str
+    selection_strategy: dict[str, Any] | None
+
+
+class _StartedWorkflowStageMetadata(_WorkflowStageMetadataBase):
+    fingerprint: str
+    num_records_requested: int
+    seeded_from_stage: str | None
+    seed_path: str | None
+    config: dict[str, Any]
+
+
+class RunningWorkflowStageMetadata(_StartedWorkflowStageMetadata):
+    """Metadata for a running workflow stage."""
+
+    status: Literal["running"]
+
+
+class FailedWorkflowStageMetadata(_StartedWorkflowStageMetadata):
+    """Metadata for a failed workflow stage."""
+
+    status: Literal["failed"]
+    duration_sec: float | None = None
+
+
+class CompletedWorkflowStageMetadata(_StartedWorkflowStageMetadata):
+    """Metadata for a completed workflow stage."""
+
+    status: Literal["completed", "completed_empty"]
+    num_records_actual: int
+    output_records: int
+    output_seed_path: str
+    callback_output_path: str | None
+    stage_output_override_path: str | None = None
+    output_processor_output_path: str | None
+    duration_sec: float
+
+
+class SkippedWorkflowStageMetadata(_WorkflowStageMetadataBase):
+    """Metadata for a stage skipped after an empty upstream stage."""
+
+    status: Literal["skipped_empty_upstream"]
+    upstream_stage: str
+
+
+WorkflowStageMetadata: TypeAlias = Annotated[
+    RunningWorkflowStageMetadata
+    | FailedWorkflowStageMetadata
+    | CompletedWorkflowStageMetadata
+    | SkippedWorkflowStageMetadata,
+    Field(discriminator="status"),
+]
+
+_WORKFLOW_STAGE_METADATA_ADAPTER = TypeAdapter(WorkflowStageMetadata)
+
+
+def validate_workflow_stage_metadata(value: object) -> WorkflowStageMetadata:
+    """Validate the stage union and return its concrete status model.
+
+    Use the concrete status models for runtime type checks.
+    """
+
+    return _WORKFLOW_STAGE_METADATA_ADAPTER.validate_python(value)
+
+
+class WorkflowMetadata(BaseModel):
+    """Metadata persisted for a composite workflow run."""
+
+    model_config = ConfigDict(extra="allow")
+
+    name: str
+    library_version: str
+    stages: list[WorkflowStageMetadata]

data_designer-0.9.3/tests/docs/test_chart_qa_recipe.py

@@ -0,0 +1,136 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+from types import ModuleType
+from unittest.mock import MagicMock
+
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+RECIPE_PATH = REPO_ROOT / "docs" / "assets" / "recipes" / "chart_qa" / "chart_qa.py"
+
+
+def _load_recipe() -> ModuleType:
+    spec = importlib.util.spec_from_file_location("chart_qa_recipe", RECIPE_PATH)
+    assert spec is not None
+    assert spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.fixture(scope="module")
+def recipe() -> ModuleType:
+    return _load_recipe()
+
+
+@pytest.fixture
+def seed_path(tmp_path: Path) -> Path:
+    path = tmp_path / "chart_seed.jsonl"
+    path.write_text('{"image_url": "https://example.com/chart.png"}\n', encoding="utf-8")
+    return path
+
+
+@pytest.mark.parametrize("prompt_version", ["level1", "level2"])
+def test_chart_prompt_versions_match_pipeline_config(recipe: ModuleType, seed_path: Path, prompt_version: str) -> None:
+    config = recipe.build_config(seed_path=str(seed_path), prompt_version=prompt_version).build()
+    columns = {column.name: column for column in config.columns}
+    models = {model.alias: model for model in config.model_configs}
+
+    assert list(columns) == ["question_type", "question", "answer", "quality_check"]
+    assert columns["question_type"].params.values == list(recipe.QUESTION_TYPES)
+    assert columns["question_type"].params.weights == list(recipe.QUESTION_TYPE_WEIGHTS)
+    assert columns["question"].model_alias == "chart-vlm-question"
+    assert columns["answer"].model_alias == "chart-vlm"
+    assert columns["quality_check"].model_alias == "chart-vlm-quality"
+    assert columns["answer"].extract_reasoning_content is True
+    assert columns["question"].multi_modal_context[0].column_name == "image_url"
+    assert columns["question"].multi_modal_context[0].data_type == "url"
+    assert columns["question"].multi_modal_context[0].image_format is None
+    assert columns["answer"].multi_modal_context == columns["question"].multi_modal_context
+    assert columns["quality_check"].multi_modal_context == columns["question"].multi_modal_context
+    assert "{{ answer__reasoning_content }}" in columns["quality_check"].prompt
+    assert config.seed_config.sampling_strategy == recipe.dd.SamplingStrategy.ORDERED
+    assert list(models) == ["chart-vlm", "chart-vlm-question", "chart-vlm-quality"]
+
+    for model in models.values():
+        inference = model.inference_parameters
+        assert model.model == recipe.DEFAULT_VLM_MODEL
+        assert model.provider == recipe.VLLM_PROVIDER_NAME
+        assert model.skip_health_check is False
+        assert inference.max_tokens == 50000
+        assert inference.timeout == 1200
+        assert inference.temperature == 0.6
+        assert inference.top_p == 0.95
+        assert inference.max_parallel_requests == 20
+        assert inference.extra_body == recipe.PRODUCTION_SAMPLING
+
+
+@pytest.mark.parametrize(
+    ("prompt_version", "raw_prompts"),
+    [
+        (
+            "level1",
+            ("LEVEL1_QUESTION_PROMPT", "LEVEL1_ANSWER_PROMPT", "LEVEL1_QUALITY_PROMPT"),
+        ),
+        (
+            "level2",
+            (
+                "LEVEL2_QUESTION_PROMPT",
+                "LEVEL2_ANSWER_PROMPT",
+                "LEVEL2_QUALITY_PROMPT",
+            ),
+        ),
+    ],
+)
+def test_prompt_pack_matches_runtime_whitespace_handling(
+    recipe: ModuleType, prompt_version: str, raw_prompts: tuple[str, str, str]
+) -> None:
+    prompt_pack = recipe.get_prompt_pack(prompt_version)
+
+    assert tuple(prompt_pack) == tuple(getattr(recipe, name).rstrip() for name in raw_prompts)
+    assert all(not prompt.endswith(("\n", "\r", " ", "\t")) for prompt in prompt_pack)
+
+    if prompt_version == "level2":
+        assert "DIFFICULTY FLOOR (required)" in prompt_pack.question
+        assert "NO-ATTEMPT HANDLING" in prompt_pack.answer
+
+
+def test_unknown_prompt_version_is_rejected(recipe: ModuleType) -> None:
+    with pytest.raises(ValueError, match="Unknown prompt version"):
+        recipe.get_prompt_pack("level3")
+
+
+def test_create_dataset_uses_production_provider_settings(
+    recipe: ModuleType, seed_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    config_builder = recipe.build_config(seed_path=str(seed_path))
+    dataset_results = MagicMock()
+    data_designer = MagicMock()
+    data_designer.create.return_value = dataset_results
+    data_designer_class = MagicMock(return_value=data_designer)
+    monkeypatch.setattr(recipe, "DataDesigner", data_designer_class)
+
+    result = recipe.create_dataset(
+        config_builder,
+        num_records=7,
+        vllm_endpoint="http://localhost:8000/v1",
+        prompt_version="level1",
+        artifact_path="artifacts",
+    )
+
+    assert result is dataset_results
+    provider = data_designer_class.call_args.kwargs["model_providers"][0]
+    assert provider.name == "vllm-vl"
+    assert provider.endpoint == "http://localhost:8000/v1"
+    assert provider.provider_type == "openai"
+    assert provider.api_key is None
+    data_designer.create.assert_called_once_with(
+        config_builder,
+        num_records=7,
+        dataset_name="chart_qa_level1",
+    )

{data_designer-0.9.2 → data_designer-0.9.3}/tests/interface/test_composite_workflow.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 import json
 import shutil
 from pathlib import Path
+from typing import Any
 from unittest.mock import MagicMock
 
 import pytest
@@ -22,6 +23,7 @@ from data_designer.config.seed_source import LocalFileSeedSource
 from data_designer.config.seed_source_dataframe import DataFrameSeedSource
 from data_designer.engine.secret_resolver import PlaintextResolver
 from data_designer.engine.storage.artifact_storage import ArtifactStorage, BatchStage, ResumeMode
+from data_designer.interface import WorkflowMetadata
 from data_designer.interface.composite_workflow import SkippedStageResult, SkippedStageStatus
 from data_designer.interface.data_designer import DataDesigner
 from data_designer.interface.errors import DataDesignerWorkflowError
@@ -194,6 +196,9 @@ def test_composite_workflow_runs_linear_stages_with_disk_handoff(
     assert data_designer.artifact_path == stub_artifact_path
 
     metadata = _load_workflow_metadata(stub_artifact_path, "linear-chain")
+    metadata_path = stub_artifact_path / "linear-chain" / "workflow-metadata.json"
+    validated_metadata = WorkflowMetadata.model_validate_json(metadata_path.read_text())
+    assert validated_metadata.model_dump(mode="json", exclude_unset=True) == metadata
     assert [stage["status"] for stage in metadata["stages"]] == ["completed", "completed"]
     assert metadata["stages"][1]["seeded_from_stage"] == "base"
     assert metadata["stages"][1]["depends_on"] == ["base"]
@@ -364,6 +369,33 @@ def test_composite_workflow_rejects_invalid_stage_outputs(
         workflow.add_stage("base", _category_builder(stub_model_configs), output=output)
 
 
+@pytest.mark.parametrize(
+    ("kwargs", "match"),
+    [
+        ({"num_records": "3"}, "num_records must be an integer"),
+        ({"num_records": True}, "num_records must be an integer"),
+        ({"num_records": 0}, "num_records must be at least 1"),
+        ({"on_success_version": 1}, "on_success_version must be a string"),
+        ({"allow_empty": "yes"}, "allow_empty must be a boolean"),
+        ({"sampling_strategy": "ordered"}, "sampling_strategy must be a SamplingStrategy"),
+        ({"selection_strategy": {}}, "selection_strategy must be an IndexRange or PartitionBlock"),
+    ],
+)
+def test_composite_workflow_rejects_invalid_stage_metadata_before_artifacts(
+    stub_artifact_path: Path,
+    stub_model_providers: list[ModelProvider],
+    stub_model_configs: list[ModelConfig],
+    kwargs: dict[str, Any],
+    match: str,
+) -> None:
+    workflow = _data_designer(stub_artifact_path, stub_model_providers).compose_workflow(name="invalid-version")
+
+    with pytest.raises(DataDesignerWorkflowError, match=match):
+        workflow.add_stage("base", _category_builder(stub_model_configs), **kwargs)
+
+    assert not stub_artifact_path.exists()
+
+
 def test_composite_workflow_rejects_unknown_processor_stage_output(
     stub_artifact_path: Path,
     stub_model_providers: list[ModelProvider],
@@ -568,7 +600,7 @@ def test_composite_workflow_stage_output_override_path_must_exist(
     assert create_mock.call_count == 0
 
 
-def test_composite_workflow_resume_if_possible_skips_completed_stages(
+def test_composite_workflow_resume_if_possible_skips_legacy_completed_stages(
     stub_artifact_path: Path,
     stub_model_providers: list[ModelProvider],
     stub_model_configs: list[ModelConfig],
@@ -580,6 +612,13 @@ def test_composite_workflow_resume_if_possible_skips_completed_stages(
     workflow.add_stage("base", _category_builder(stub_model_configs), num_records=3)
     workflow.add_stage("copy", _copy_builder(stub_model_configs))
     workflow.run()
+    metadata_path = stub_artifact_path / "resume-skip" / "workflow-metadata.json"
+    metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+    metadata["workflow_extension"] = True
+    for stage in metadata["stages"]:
+        stage.pop("stage_output_override_path")
+        stage["stage_extension"] = {"value": stage["index"]}
+    metadata_path.write_text(json.dumps(metadata), encoding="utf-8")
     create_mock.reset_mock()
 
     resumed = data_designer.compose_workflow(name="resume-skip")
@@ -590,6 +629,10 @@ def test_composite_workflow_resume_if_possible_skips_completed_stages(
     assert create_mock.call_count == 0
     assert results.count_records() == 3
     assert results.load_dataset()["category"].tolist() == ["alpha", "alpha", "alpha"]
+    resumed_metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+    assert resumed_metadata["workflow_extension"] is True
+    assert [stage["stage_extension"] for stage in resumed_metadata["stages"]] == [{"value": 0}, {"value": 1}]
+    assert all("stage_output_override_path" not in stage for stage in resumed_metadata["stages"])
 
 
 def test_composite_workflow_resume_if_possible_skips_stage_with_output_processors(
@@ -825,6 +868,30 @@ def test_composite_workflow_resume_if_possible_invalid_metadata_shape_starts_fre
     assert [call.kwargs["dataset_name"] for call in create_mock.call_args_list] == ["stage-0-base", "stage-1-copy"]
 
 
+def test_composite_workflow_resume_if_possible_invalid_stage_metadata_starts_fresh(
+    stub_artifact_path: Path,
+    stub_model_providers: list[ModelProvider],
+    stub_model_configs: list[ModelConfig],
+    stub_dataset_profiler_results,
+) -> None:
+    data_designer = _data_designer(stub_artifact_path, stub_model_providers)
+    create_mock = _patch_create(data_designer, stub_dataset_profiler_results)
+    workflow = data_designer.compose_workflow(name="resume-invalid-stage")
+    workflow.add_stage("base", _category_builder(stub_model_configs), num_records=2)
+    workflow.run()
+    metadata_path = stub_artifact_path / "resume-invalid-stage" / "workflow-metadata.json"
+    metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+    metadata["stages"][0]["status"] = "unknown"
+    metadata_path.write_text(json.dumps(metadata), encoding="utf-8")
+    create_mock.reset_mock()
+
+    resumed = data_designer.compose_workflow(name="resume-invalid-stage")
+    resumed.add_stage("base", _category_builder(stub_model_configs), num_records=2)
+    resumed.run(resume=ResumeMode.IF_POSSIBLE)
+
+    assert [call.kwargs["dataset_name"] for call in create_mock.call_args_list] == ["stage-0-base"]
+
+
 @pytest.mark.parametrize("status", ["running", "failed"])
 def test_composite_workflow_resume_if_possible_delegates_matching_resumable_stage(
     stub_artifact_path: Path,
@@ -842,8 +909,20 @@ def test_composite_workflow_resume_if_possible_delegates_matching_resumable_stag
     metadata_path = stub_artifact_path / "resume-partial" / "workflow-metadata.json"
     metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
     _mark_stage_resumable(metadata, 0, status)
+    # duration_sec is schema-owned by failed metadata, but is an extension on running metadata.
+    if status == "failed":
+        metadata["stages"][0]["duration_sec"] = 123.0
+    metadata["stages"][0]["stage_extension"] = {"value": status}
     metadata_path.write_text(json.dumps(metadata), encoding="utf-8")
     create_mock.reset_mock()
+    create_side_effect = create_mock.side_effect
+    metadata_at_create: list[dict[str, Any]] = []
+
+    def capture_create(*args: Any, **kwargs: Any) -> DatasetCreationResults:
+        metadata_at_create.append(json.loads(metadata_path.read_text(encoding="utf-8")))
+        return create_side_effect(*args, **kwargs)
+
+    create_mock.side_effect = capture_create
 
     resumed = data_designer.compose_workflow(name="resume-partial")
     resumed.add_stage("base", _category_builder(stub_model_configs), num_records=2)
@@ -852,6 +931,10 @@ def test_composite_workflow_resume_if_possible_delegates_matching_resumable_stag
 
     assert [call.kwargs["dataset_name"] for call in create_mock.call_args_list] == ["stage-0-base", "stage-1-copy"]
     assert [call.kwargs["resume"] for call in create_mock.call_args_list] == [ResumeMode.ALWAYS, ResumeMode.NEVER]
+    assert metadata_at_create[0]["stages"][0]["stage_extension"] == {"value": status}
+    assert "duration_sec" not in metadata_at_create[0]["stages"][0]
+    resumed_metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+    assert resumed_metadata["stages"][0]["stage_extension"] == {"value": status}
 
 
 def test_composite_workflow_resume_always_reruns_descendants_after_partial_stage(
@@ -936,6 +1019,29 @@ def test_composite_workflow_resume_always_rejects_invalid_metadata_shape(
         resumed.run(resume=ResumeMode.ALWAYS)
 
 
+def test_composite_workflow_resume_always_rejects_invalid_stage_metadata(
+    stub_artifact_path: Path,
+    stub_model_providers: list[ModelProvider],
+    stub_model_configs: list[ModelConfig],
+    stub_dataset_profiler_results,
+) -> None:
+    data_designer = _data_designer(stub_artifact_path, stub_model_providers)
+    _patch_create(data_designer, stub_dataset_profiler_results)
+    workflow = data_designer.compose_workflow(name="resume-invalid-stage-always")
+    workflow.add_stage("base", _category_builder(stub_model_configs), num_records=2)
+    workflow.run()
+    metadata_path = stub_artifact_path / "resume-invalid-stage-always" / "workflow-metadata.json"
+    metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+    metadata["stages"][0]["status"] = "unknown"
+    metadata_path.write_text(json.dumps(metadata), encoding="utf-8")
+
+    resumed = data_designer.compose_workflow(name="resume-invalid-stage-always")
+    resumed.add_stage("base", _category_builder(stub_model_configs), num_records=2)
+
+    with pytest.raises(DataDesignerWorkflowError, match="workflow metadata has invalid shape"):
+        resumed.run(resume=ResumeMode.ALWAYS)
+
+
 def test_composite_workflow_resume_always_rejects_changed_stage(
     stub_artifact_path: Path,
     stub_model_providers: list[ModelProvider],