fix: address workflow review nits

Author: andreatgretel

packages/data-designer/src/data_designer/interface/composite_workflow.py

@@ -107,6 +107,7 @@ def items(self) -> ItemsView[str, DatasetCreationResults | SkippedStageResult]:
 
     @property
     def final_result(self) -> DatasetCreationResults:
+        """Return the final stage result, or raise if it was skipped."""
         return self._require_final_result()
 
     def _require_final_result(self) -> DatasetCreationResults:
@@ -118,33 +119,41 @@ def _require_final_result(self) -> DatasetCreationResults:
         return result
 
     def load_dataset(self) -> pd.DataFrame:
+        """Load the selected output from the final workflow stage."""
         self._require_final_result()
         return self.load_stage_output(self.final_stage_name)
 
     def load_analysis(self) -> DatasetProfilerResults:
+        """Load analysis from the final stage result."""
         return self.final_result.load_analysis()
 
     def count_records(self) -> int:
+        """Count records in the selected output from the final workflow stage."""
         self._require_final_result()
         return self.count_stage_output_records(self.final_stage_name)
 
     def get_stage_output_path(self, stage_name: str) -> Path:
+        """Return the selected output path handed downstream for a stage."""
         result = self.stage_results[stage_name]
         if isinstance(result, SkippedStageResult):
             raise DataDesignerWorkflowError(f"Stage {stage_name!r} was skipped: {result.status.value}.")
         return self._stage_output_paths.get(stage_name, result.artifact_storage.final_dataset_path)
 
     def load_stage_output(self, stage_name: str) -> pd.DataFrame:
+        """Load the selected output handed downstream for a stage."""
         return _load_parquet_dataset(self.get_stage_output_path(stage_name))
 
     def count_stage_output_records(self, stage_name: str) -> int:
+        """Count records in the selected output handed downstream for a stage."""
         return _count_parquet_records(self.get_stage_output_path(stage_name))
 
     def export(self, path: Path | str, *, format: ExportFormat | None = None) -> Path:
+        """Export the selected output from the final workflow stage."""
         self._require_final_result()
         return _export_parquet_dataset(self.get_stage_output_path(self.final_stage_name), Path(path), format=format)
 
     def push_to_hub(self, *args: Any, **kwargs: Any) -> str:
+        """Push the final stage result to Hugging Face Hub when no output override is selected."""
         final_result = self.final_result
         if self.get_stage_output_path(self.final_stage_name) != final_result.artifact_storage.final_dataset_path:
             raise DataDesignerWorkflowError(
@@ -155,7 +164,10 @@ def push_to_hub(self, *args: Any, **kwargs: Any) -> str:
 
 
 class CompositeWorkflow:
+    """Experimental linear workflow for chaining Data Designer stages."""
+
     def __init__(self, *, name: str, data_designer: DataDesigner) -> None:
+        """Create a workflow bound to a parent Data Designer instance."""
         _validate_dir_name(name, "workflow name")
         self.name = name
         self._data_designer = data_designer
@@ -210,11 +222,16 @@ def add_stage(
         return self
 
     def run(self) -> CompositeWorkflowResults:
-        """Run all stages from scratch, replacing deterministic stage directories."""
+        """Run all stages from scratch.
+
+        Each stage writes a deterministic artifact directory under the parent
+        Data Designer artifact path. Downstream stages are seeded from the
+        selected output of the previous stage.
+        """
         if not self._stages:
             raise DataDesignerWorkflowError(f"Workflow {self.name!r} has no stages.")
 
-        workflow_path = self._data_designer._artifact_path / self.name
+        workflow_path = self._data_designer.artifact_path / self.name
         workflow_path.mkdir(parents=True, exist_ok=True)
         metadata: dict[str, Any] = {
             "name": self.name,

packages/data-designer/src/data_designer/interface/data_designer.py

@@ -215,6 +215,11 @@ def info(self) -> InterfaceInfo:
         """
         return self._get_interface_info(self._model_providers)
 
+    @property
+    def artifact_path(self) -> Path:
+        """Directory where Data Designer writes artifacts by default."""
+        return self._artifact_path
+
     def list_mcp_tool_names(self, mcp_provider_name: str, *, timeout_sec: float = 10.0) -> list[str]:
         """Connect to a configured MCP provider and return the names of its available tools.
 
@@ -498,6 +503,17 @@ def preview(
         )
 
     def compose_workflow(self, *, name: str) -> CompositeWorkflow:
+        """Create an experimental composite workflow.
+
+        Workflow chaining is experimental and its API, metadata schema, and
+        artifact layout may change in future releases.
+
+        Args:
+            name: Workflow name used for the artifact directory.
+
+        Returns:
+            A composite workflow that can run named stages in sequence.
+        """
         return CompositeWorkflow(name=name, data_designer=self)
 
     def _log_jinja_rendering_engine_mode(self) -> None:

packages/data-designer/tests/interface/test_acreate.py

@@ -97,8 +97,8 @@ def fake_create(
             started_count += 1
             if started_count == 2:
                 both_started.set()
-        assert both_started.wait(2)
-        assert release.wait(2)
+        assert both_started.wait(5)
+        assert release.wait(5)
         return _creation_result(config_builder, stub_dataset_profiler_results)
 
     data_designer.create = MagicMock(side_effect=fake_create)
@@ -108,7 +108,7 @@ def fake_create(
     left_task = asyncio.create_task(data_designer.acreate(left, num_records=1, dataset_name="left"))
     right_task = asyncio.create_task(data_designer.acreate(right, num_records=1, dataset_name="right"))
     try:
-        assert await asyncio.to_thread(both_started.wait, 2)
+        assert await asyncio.to_thread(both_started.wait, 5)
     finally:
         release.set()
 

packages/data-designer/tests/interface/test_composite_workflow.py

@@ -70,7 +70,7 @@ def fake_create(
         dataset_name: str,
         **kwargs,
     ) -> DatasetCreationResults:
-        artifact_path = Path(kwargs.pop("artifact_path", data_designer._artifact_path))
+        artifact_path = Path(kwargs.pop("artifact_path", data_designer.artifact_path))
         del kwargs
         df = lazy.pd.DataFrame({"category": ["alpha"] * num_records, "category_copy": ["alpha"] * num_records})
         return _result_from_df(
@@ -177,7 +177,7 @@ def test_composite_workflow_runs_linear_stages_with_disk_handoff(
     assert "category_copy" in final_df.columns
     assert (stub_artifact_path / "linear-chain" / "stage-0-base").is_dir()
     assert (stub_artifact_path / "linear-chain" / "stage-1-copy").is_dir()
-    assert data_designer._artifact_path == stub_artifact_path
+    assert data_designer.artifact_path == stub_artifact_path
 
     metadata = _load_workflow_metadata(stub_artifact_path, "linear-chain")
     assert [stage["status"] for stage in metadata["stages"]] == ["completed", "completed"]
@@ -880,7 +880,7 @@ def fake_create(
         dataset_name: str,
         **kwargs,
     ) -> DatasetCreationResults:
-        artifact_path = Path(kwargs.pop("artifact_path", data_designer._artifact_path))
+        artifact_path = Path(kwargs.pop("artifact_path", data_designer.artifact_path))
         del num_records, kwargs
         value = "first" if dataset_name == "stage-0-first" else "final"
         return _result_from_df(
@@ -921,7 +921,7 @@ def fake_create(
         dataset_name: str,
         **kwargs,
     ) -> DatasetCreationResults:
-        artifact_path = Path(kwargs.pop("artifact_path", data_designer._artifact_path))
+        artifact_path = Path(kwargs.pop("artifact_path", data_designer.artifact_path))
         del num_records, kwargs
         return _result_from_df(
             artifact_path,

packages/data-designer/tests/interface/test_data_designer.py

@@ -491,7 +491,8 @@ def test_init_with_string_path(stub_artifact_path, stub_model_providers):
     """Test DataDesigner accepts string paths."""
     designer = DataDesigner(artifact_path=str(stub_artifact_path), model_providers=stub_model_providers)
     assert designer is not None
-    assert isinstance(designer._artifact_path, Path)
+    assert isinstance(designer.artifact_path, Path)
+    assert designer.artifact_path == stub_artifact_path
 
 
 def test_init_with_path_object(stub_artifact_path, stub_model_providers):