fix: harden workflow chaining concurrency

Author: andreatgretel

packages/data-designer-engine/src/data_designer/engine/models/factory.py

@@ -13,6 +13,7 @@
 if TYPE_CHECKING:
     from data_designer.config.run_config import RunConfig
     from data_designer.engine.mcp.registry import MCPRegistry
+    from data_designer.engine.models.clients.throttle_manager import ThrottleManager
     from data_designer.engine.models.registry import ModelRegistry
 
 
@@ -24,6 +25,7 @@ def create_model_registry(
     mcp_registry: MCPRegistry | None = None,
     client_concurrency_mode: ClientConcurrencyMode = ClientConcurrencyMode.SYNC,
     run_config: RunConfig | None = None,
+    throttle_manager: ThrottleManager | None = None,
 ) -> ModelRegistry:
     """Factory function for creating a ModelRegistry instance.
 
@@ -43,6 +45,8 @@ def create_model_registry(
         run_config: Optional runtime configuration.  The nested
             ``run_config.throttle`` (a ``ThrottleConfig``) is forwarded to the
             ``ThrottleManager`` constructor.
+        throttle_manager: Optional shared throttle manager. When omitted, a new
+            manager is created for this registry.
 
     Returns:
         A configured ModelRegistry instance.
@@ -54,7 +58,8 @@ def create_model_registry(
     from data_designer.engine.models.facade import ModelFacade
     from data_designer.engine.models.registry import ModelRegistry
 
-    throttle_manager = ThrottleManager((run_config or RunConfig()).throttle)
+    if throttle_manager is None:
+        throttle_manager = ThrottleManager((run_config or RunConfig()).throttle)
 
     def model_facade_factory(
         model_config: ModelConfig,

packages/data-designer-engine/src/data_designer/engine/resources/resource_provider.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import os
+from typing import TYPE_CHECKING
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.dataset_metadata import DatasetMetadata
@@ -26,6 +27,9 @@
 from data_designer.engine.secret_resolver import SecretResolver
 from data_designer.engine.storage.artifact_storage import ArtifactStorage
 
+if TYPE_CHECKING:
+    from data_designer.engine.models.clients.throttle_manager import ThrottleManager
+
 
 class ResourceType(StrEnum):
     PERSON_READER = "person_reader"
@@ -91,6 +95,7 @@ def create_resource_provider(
     mcp_providers: list[MCPProviderT] | None = None,
     tool_configs: list[ToolConfig] | None = None,
     client_concurrency_mode: ClientConcurrencyMode | None = None,
+    throttle_manager: ThrottleManager | None = None,
 ) -> ResourceProvider:
     """Factory function for creating a ResourceProvider instance.
 
@@ -111,6 +116,7 @@ def create_resource_provider(
         run_config: Optional runtime configuration.
         mcp_providers: Optional list of MCP provider configurations.
         tool_configs: Optional list of tool configurations.
+        throttle_manager: Optional shared throttle manager for model clients.
 
     Returns:
         A configured ResourceProvider instance.
@@ -158,6 +164,7 @@ def create_resource_provider(
             mcp_registry=mcp_registry,
             client_concurrency_mode=client_concurrency_mode,
             run_config=effective_run_config,
+            throttle_manager=throttle_manager,
         ),
         person_reader=person_reader,
         mcp_registry=mcp_registry,

packages/data-designer-engine/src/data_designer/engine/resources/seed_reader.py

@@ -6,6 +6,7 @@
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable, Sequence
+from copy import copy
 from dataclasses import dataclass
 from fnmatch import fnmatchcase
 from pathlib import Path, PurePosixPath
@@ -673,7 +674,7 @@ def add_reader(self, reader: SeedReader) -> Self:
         return self
 
     def get_reader(self, seed_dataset_source: SeedSource, secret_resolver: SecretResolver) -> SeedReader:
-        reader = self._get_reader_for_source(seed_dataset_source)
+        reader = copy(self._get_reader_for_source(seed_dataset_source))
         reader.attach(seed_dataset_source, secret_resolver)
         return reader
 

packages/data-designer-engine/tests/engine/resources/test_seed_reader.py

@@ -412,7 +412,9 @@ def test_get_reader_basic():
 
     reader = registry.get_reader(local_seed_config, PlaintextResolver())
 
-    assert reader == df_reader
+    assert isinstance(reader, DataFrameSeedReader)
+    assert reader is not df_reader
+    assert reader.source is local_seed_config
 
 
 def test_get_reader_missing():

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

@@ -5,10 +5,10 @@
 
 import hashlib
 import json
+import logging
 import shutil
 import time
-from collections.abc import Callable, Iterator
-from contextlib import contextmanager
+from collections.abc import Callable, ItemsView, Iterator, KeysView
 from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -24,9 +24,14 @@
 from data_designer.interface.results import DatasetCreationResults
 
 if TYPE_CHECKING:
+    import pandas as pd
+
+    from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
     from data_designer.interface.data_designer import DataDesigner
 
 
+logger = logging.getLogger(__name__)
+
 OnSuccessCallback = Callable[[Path], Path | str]
 
 
@@ -50,13 +55,21 @@ class SkippedStageResult:
 
 
 class CompositeWorkflowResults:
+    """Results for a composite workflow run.
+
+    Per-stage entries are the original ``DataDesigner.create()`` results. If a
+    stage uses ``on_success``, metadata and downstream seeding use the callback
+    output path while the stage result still points at the stage's generated
+    dataset.
+    """
+
     def __init__(
         self,
         *,
         name: str,
         stage_results: dict[str, DatasetCreationResults | SkippedStageResult],
         final_stage_name: str,
-    ):
+    ) -> None:
         self.name = name
         self.stage_results = stage_results
         self.final_stage_name = final_stage_name
@@ -67,10 +80,10 @@ def __getitem__(self, stage_name: str) -> DatasetCreationResults | SkippedStageR
     def __iter__(self) -> Iterator[str]:
         return iter(self.stage_results)
 
-    def keys(self):
+    def keys(self) -> KeysView[str]:
         return self.stage_results.keys()
 
-    def items(self):
+    def items(self) -> ItemsView[str, DatasetCreationResults | SkippedStageResult]:
         return self.stage_results.items()
 
     @property
@@ -80,24 +93,24 @@ def final_result(self) -> DatasetCreationResults:
             raise DataDesignerWorkflowError(f"Final stage {self.final_stage_name!r} was skipped: {result.status}.")
         return result
 
-    def load_dataset(self):
+    def load_dataset(self) -> pd.DataFrame:
         return self.final_result.load_dataset()
 
-    def load_analysis(self):
+    def load_analysis(self) -> DatasetProfilerResults:
         return self.final_result.load_analysis()
 
     def count_records(self) -> int:
         return self.final_result.count_records()
 
-    def export(self, *args, **kwargs):
+    def export(self, *args: Any, **kwargs: Any) -> Path:
         return self.final_result.export(*args, **kwargs)
 
-    def push_to_hub(self, *args, **kwargs):
+    def push_to_hub(self, *args: Any, **kwargs: Any) -> str:
         return self.final_result.push_to_hub(*args, **kwargs)
 
 
 class CompositeWorkflow:
-    def __init__(self, *, name: str, data_designer: DataDesigner):
+    def __init__(self, *, name: str, data_designer: DataDesigner) -> None:
         _validate_dir_name(name, "workflow name")
         self.name = name
         self._data_designer = data_designer
@@ -123,7 +136,7 @@ def add_stage(
         self._stages.append(
             _WorkflowStage(
                 name=name,
-                config_builder=config_builder,
+                config_builder=_clone_config_builder(config_builder),
                 depends_on=(self._stages[-1].name,) if self._stages else (),
                 num_records=num_records,
                 on_success=on_success,
@@ -136,6 +149,7 @@ def add_stage(
         return self
 
     def run(self) -> CompositeWorkflowResults:
+        """Run all stages from scratch, replacing deterministic stage directories."""
         if not self._stages:
             raise DataDesignerWorkflowError(f"Workflow {self.name!r} has no stages.")
 
@@ -174,6 +188,12 @@ def run(self) -> CompositeWorkflowResults:
 
             stage_builder = _clone_config_builder(stage.config_builder)
             if previous_seed_path is not None:
+                if stage_builder.get_seed_config() is not None:
+                    logger.warning(
+                        "Stage %r has a seed dataset; workflow will seed it from upstream stage %r.",
+                        stage.name,
+                        previous_stage_name,
+                    )
                 stage_builder.with_seed_dataset(
                     _local_seed_source_from_path(previous_seed_path),
                     sampling_strategy=stage.sampling_strategy,
@@ -206,12 +226,12 @@ def run(self) -> CompositeWorkflowResults:
 
             start_time = time.monotonic()
             try:
-                with _temporary_artifact_path(self._data_designer, workflow_path):
-                    result = self._data_designer.create(
-                        stage_builder,
-                        num_records=num_records,
-                        dataset_name=stage_dir_name,
-                    )
+                result = self._data_designer.create(
+                    stage_builder,
+                    num_records=num_records,
+                    dataset_name=stage_dir_name,
+                    artifact_path=workflow_path,
+                )
                 actual_records = result.count_records()
                 output_seed_path = result.artifact_storage.final_dataset_path
                 callback_output_path = None
@@ -263,16 +283,6 @@ def _clone_config_builder(config_builder: DataDesignerConfigBuilder) -> DataDesi
     return DataDesignerConfigBuilder.from_config(BuilderConfig(data_designer=config_builder.build()))
 
 
-@contextmanager
-def _temporary_artifact_path(data_designer: DataDesigner, artifact_path: Path):
-    original_artifact_path = data_designer._artifact_path
-    data_designer._artifact_path = artifact_path
-    try:
-        yield
-    finally:
-        data_designer._artifact_path = original_artifact_path
-
-
 def _stage_dir_name(index: int, name: str) -> str:
     return f"stage-{index}-{name}"