From e9b5b3d765b4ee07adf320346d45cf08e6a80e9e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw?= Date: Wed, 8 Apr 2026 18:47:30 +0200 Subject: [PATCH 1/7] refactor: unify DAG construction by moving topological sort into execution_graph.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Eliminates dag.py and its networkx dependency by moving topologically_sort_column_configs into execution_graph.py as a module-level function. Side-effect resolution is now O(1) via a side_effect_map dict (previously O(nĀ²) linear scan). Kahn's algorithm is reused in-place rather than leaning on networkx.topological_sort. Closes #510 Co-Authored-By: Claude Sonnet 4.6 --- .../dataset_builders/utils/config_compiler.py | 2 +- .../engine/dataset_builders/utils/dag.py | 89 ------------------- .../dataset_builders/utils/execution_graph.py | 62 +++++++++++++ .../engine/dataset_builders/utils/test_dag.py | 2 +- 4 files changed, 64 insertions(+), 91 deletions(-) delete mode 100644 packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/dag.py diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/config_compiler.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/config_compiler.py index 8112d87e8..208fa6d80 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/config_compiler.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/config_compiler.py @@ -11,8 +11,8 @@ SamplerMultiColumnConfig, SeedDatasetMultiColumnConfig, ) -from data_designer.engine.dataset_builders.utils.dag import topologically_sort_column_configs from data_designer.engine.dataset_builders.utils.errors import ConfigCompilationError +from data_designer.engine.dataset_builders.utils.execution_graph import topologically_sort_column_configs def compile_dataset_builder_column_configs(config: DataDesignerConfig) -> list[DatasetBuilderColumnConfigT]: diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/dag.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/dag.py deleted file mode 100644 index 4b3e03670..000000000 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/dag.py +++ /dev/null @@ -1,89 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import logging -from itertools import chain - -import data_designer.lazy_heavy_imports as lazy -from data_designer.config.column_types import ColumnConfigT -from data_designer.engine.column_generators.utils.generator_classification import column_type_used_in_execution_dag -from data_designer.engine.dataset_builders.utils.errors import ConfigCompilationError, DAGCircularDependencyError -from data_designer.logging import LOG_INDENT - -logger = logging.getLogger(__name__) - - -def topologically_sort_column_configs(column_configs: list[ColumnConfigT]) -> list[ColumnConfigT]: - dag = lazy.nx.DiGraph() - - non_dag_column_config_list = [ - col for col in column_configs if not column_type_used_in_execution_dag(col.column_type) - ] - dag_column_config_dict = { - col.name: col for col in column_configs if column_type_used_in_execution_dag(col.column_type) - } - - if len(dag_column_config_dict) == 0: - return non_dag_column_config_list - - side_effect_dict = {n: list(c.side_effect_columns) for n, c in dag_column_config_dict.items()} - all_side_effects = set(chain.from_iterable(side_effect_dict.values())) - - side_effect_to_producer: dict[str, str] = {} - for producer, cols in side_effect_dict.items(): - for col in cols: - existing = side_effect_to_producer.get(col) - if existing is not None and existing != producer: - raise ConfigCompilationError( - f"Side-effect column {col!r} is already produced by {existing!r}; " - f"cannot register a second producer {producer!r}. " - f"Use distinct side-effect column names for each pipeline stage." - ) - side_effect_to_producer[col] = producer - - logger.info("ā›“ļø Sorting column configs into a Directed Acyclic Graph") - for name, col in dag_column_config_dict.items(): - dag.add_node(name) - _add_dependency_edges( - dag, name, list(col.required_columns), dag_column_config_dict, side_effect_dict, all_side_effects, "" - ) - if col.skip is not None: - _add_dependency_edges( - dag, name, col.skip.columns, dag_column_config_dict, side_effect_dict, all_side_effects, "skip.when" - ) - - if not lazy.nx.is_directed_acyclic_graph(dag): - raise DAGCircularDependencyError( - "šŸ›‘ The Data Designer column configurations contain cyclic dependencies. Please " - "inspect the column configurations and ensure they can be sorted without " - "circular references." - ) - - sorted_columns = non_dag_column_config_list - sorted_columns.extend([dag_column_config_dict[n] for n in list(lazy.nx.topological_sort(dag))]) - - return sorted_columns - - -def _add_dependency_edges( - dag: lazy.nx.DiGraph, - name: str, - dep_names: list[str], - dag_column_config_dict: dict[str, ColumnConfigT], - side_effect_dict: dict[str, list[str]], - all_side_effects: set[str], - label: str, -) -> None: - """Add DAG edges from *dep_names* to *name*, resolving through side-effect parents.""" - for dep in dep_names: - if dep in dag_column_config_dict: - logger.debug(f"{LOG_INDENT}šŸ”— `{name}` {label} depends on `{dep}`") - dag.add_edge(dep, name) - elif dep in all_side_effects: - for parent, cols in side_effect_dict.items(): - if dep in cols: - logger.debug(f"{LOG_INDENT}šŸ”— `{name}` {label} depends on `{parent}` via `{dep}`") - dag.add_edge(parent, name) - break diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py index 5cd41dd38..731a30eaf 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py @@ -9,12 +9,15 @@ from typing import TYPE_CHECKING from data_designer.config.column_configs import GenerationStrategy +from data_designer.config.column_types import ColumnConfigT +from data_designer.engine.column_generators.utils.generator_classification import column_type_used_in_execution_dag from data_designer.engine.dataset_builders.multi_column_configs import ( DatasetBuilderColumnConfigT, MultiColumnConfig, ) from data_designer.engine.dataset_builders.utils.errors import ConfigCompilationError, DAGCircularDependencyError from data_designer.engine.dataset_builders.utils.task_model import SliceRef +from data_designer.logging import LOG_INDENT logger = logging.getLogger(__name__) @@ -330,3 +333,62 @@ def to_mermaid(self) -> str: for dep in sorted(self._upstream.get(col, set())): lines.append(f" {dep} --> {col}") return "\n".join(lines) + + +def topologically_sort_column_configs(column_configs: list[ColumnConfigT]) -> list[ColumnConfigT]: + non_dag_cols = [col for col in column_configs if not column_type_used_in_execution_dag(col.column_type)] + dag_col_dict = {col.name: col for col in column_configs if column_type_used_in_execution_dag(col.column_type)} + + if not dag_col_dict: + return non_dag_cols + + # side_effect_col_name -> producing column name + side_effect_map: dict[str, str] = {} + for name, col in dag_col_dict.items(): + for se_col in col.side_effect_columns: + existing = side_effect_map.get(se_col) + if existing is not None and existing != name: + raise ConfigCompilationError( + f"Side-effect column {se_col!r} is already produced by {existing!r}; " + f"cannot register a second producer {name!r}. " + f"Use distinct side-effect column names for each pipeline stage." + ) + side_effect_map[se_col] = name + + def resolve(col_name: str) -> str | None: + if col_name in dag_col_dict: + return col_name + return side_effect_map.get(col_name) + + upstream: dict[str, set[str]] = {name: set() for name in dag_col_dict} + downstream: dict[str, set[str]] = {name: set() for name in dag_col_dict} + + logger.info("ā›“ļø Sorting column configs into a Directed Acyclic Graph") + for name, col in dag_col_dict.items(): + for req in col.required_columns: + resolved = resolve(req) + if resolved is None or resolved == name: + continue + logger.debug(f"{LOG_INDENT}šŸ”— `{name}` depends on `{resolved}`") + upstream[name].add(resolved) + downstream[resolved].add(name) + + in_degree = {name: len(ups) for name, ups in upstream.items()} + queue: deque[str] = deque(name for name, deg in in_degree.items() if deg == 0) + order: list[str] = [] + while queue: + name = queue.popleft() + order.append(name) + for child in downstream.get(name, set()): + in_degree[child] -= 1 + if in_degree[child] == 0: + queue.append(child) + + if len(order) != len(dag_col_dict): + raise DAGCircularDependencyError( + "šŸ›‘ The Data Designer column configurations contain cyclic dependencies. Please " + "inspect the column configurations and ensure they can be sorted without " + "circular references." + ) + + return non_dag_cols + [dag_col_dict[n] for n in order] diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_dag.py b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_dag.py index bbb7aa9c8..dc908a1d4 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_dag.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_dag.py @@ -21,8 +21,8 @@ from data_designer.config.utils.code_lang import CodeLang from data_designer.config.validator_params import CodeValidatorParams from data_designer.engine.dataset_builders.multi_column_configs import SamplerMultiColumnConfig -from data_designer.engine.dataset_builders.utils.dag import topologically_sort_column_configs from data_designer.engine.dataset_builders.utils.errors import ConfigCompilationError, DAGCircularDependencyError +from data_designer.engine.dataset_builders.utils.execution_graph import topologically_sort_column_configs MODEL_ALIAS = "stub-model-alias" From eddf3a3917426b2b7605290880e6960de683ea47 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw?= Date: Wed, 8 Apr 2026 19:09:00 +0200 Subject: [PATCH 2/7] test: relax non-deterministic ordering assertion in test_dag test_judge and test_code_and_depends_on_validation_reasoning_traces have no mutual dependency and reach in-degree 0 simultaneously in Kahn's algorithm. Set iteration order varies with PYTHONHASHSEED, making the strict list assertion flaky. Assert only the topological invariants. Co-Authored-By: Claude Sonnet 4.6 --- .../engine/dataset_builders/utils/test_dag.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_dag.py b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_dag.py index dc908a1d4..d16281fa2 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_dag.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_dag.py @@ -82,14 +82,14 @@ def test_dag_construction(): assert sorted_column_configs[0].column_type == DataDesignerColumnType.SAMPLER - assert [c.name for c in sorted_column_configs[1:]] == [ - "test_code", - "test_validation", - "depends_on_validation", - "test_judge", - "test_code_and_depends_on_validation_reasoning_traces", - "uses_all_the_stuff", - ] + names = [c.name for c in sorted_column_configs[1:]] + assert names[0] == "test_code" + assert names[1] == "test_validation" + assert names[2] == "depends_on_validation" + # test_judge and test_code_and_depends_on_validation_reasoning_traces have no mutual + # dependency, so their relative order is not guaranteed by topological sort. + assert set(names[3:5]) == {"test_judge", "test_code_and_depends_on_validation_reasoning_traces"} + assert names[5] == "uses_all_the_stuff" def test_circular_dependencies(): From 6b2a8c891eb20bf7396265e4187c595e818f996a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw?= Date: Wed, 15 Apr 2026 21:43:30 +0200 Subject: [PATCH 3/7] docs(engine): document intentional skip.columns omission in topologically_sort_column_configs ExecutionGraph.create handles skip.when ordering edges in its own two-pass build; the pre-sort function only needs required_columns to produce a valid ColumnConfigT ordering for config compilation. --- .../engine/dataset_builders/utils/execution_graph.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py index 731a30eaf..da2c3a6ad 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py @@ -363,6 +363,10 @@ def resolve(col_name: str) -> str | None: upstream: dict[str, set[str]] = {name: set() for name in dag_col_dict} downstream: dict[str, set[str]] = {name: set() for name in dag_col_dict} + # Only required_columns edges are added here. skip.columns edges are intentionally + # omitted: ExecutionGraph.create handles them in its own two-pass build, which is + # the authoritative execution-order graph. This function only needs to produce a + # valid topological ordering of ColumnConfigT objects for config compilation. logger.info("ā›“ļø Sorting column configs into a Directed Acyclic Graph") for name, col in dag_col_dict.items(): for req in col.required_columns: From 17e9288c36aa00e37382d4ed826e58af4d128ff3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw?= Date: Fri, 17 Apr 2026 16:08:50 +0200 Subject: [PATCH 4/7] fix(engine): restore skip.columns edges in topologically_sort_column_configs The sync builder executes generators in compile-time sort order (from _column_configs, populated via this function), not ExecutionGraph order. Dropping skip.columns edges caused evaluate_skip_when to hit UndefinedError when the referenced column hadn't been generated yet, silently skipping rows. Also: - Refactor edge building into _add_edge() helper with a label parameter to distinguish "required" from "skip.when" edges in debug output - Rename test_dag.py -> test_topological_sort.py to match the new module location - Add from __future__ import annotations (required by AGENTS.md) - Add test_side_effect_column_ordering covering the side_effect_map.get() path - Add test_skip_when_column_ordering covering the skip.columns edge path --- .../dataset_builders/utils/execution_graph.py | 22 ++++++----- .../{test_dag.py => test_topological_sort.py} | 38 +++++++++++++++++++ 2 files changed, 50 insertions(+), 10 deletions(-) rename packages/data-designer-engine/tests/engine/dataset_builders/utils/{test_dag.py => test_topological_sort.py} (76%) diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py index da2c3a6ad..1ac4c2b26 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py @@ -363,19 +363,21 @@ def resolve(col_name: str) -> str | None: upstream: dict[str, set[str]] = {name: set() for name in dag_col_dict} downstream: dict[str, set[str]] = {name: set() for name in dag_col_dict} - # Only required_columns edges are added here. skip.columns edges are intentionally - # omitted: ExecutionGraph.create handles them in its own two-pass build, which is - # the authoritative execution-order graph. This function only needs to produce a - # valid topological ordering of ColumnConfigT objects for config compilation. + def _add_edge(name: str, dep: str, label: str) -> None: + resolved = resolve(dep) + if resolved is None: + return + logger.debug(f"{LOG_INDENT}šŸ”— `{name}` depends on `{resolved}` [{label}]") + upstream[name].add(resolved) + downstream[resolved].add(name) + logger.info("ā›“ļø Sorting column configs into a Directed Acyclic Graph") for name, col in dag_col_dict.items(): for req in col.required_columns: - resolved = resolve(req) - if resolved is None or resolved == name: - continue - logger.debug(f"{LOG_INDENT}šŸ”— `{name}` depends on `{resolved}`") - upstream[name].add(resolved) - downstream[resolved].add(name) + _add_edge(name, req, "required") + if col.skip is not None: + for skip_col in col.skip.columns: + _add_edge(name, skip_col, "skip.when") in_degree = {name: len(ups) for name, ups in upstream.items()} queue: deque[str] = deque(name for name, deg in in_degree.items() if deg == 0) diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_dag.py b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_topological_sort.py similarity index 76% rename from packages/data-designer-engine/tests/engine/dataset_builders/utils/test_dag.py rename to packages/data-designer-engine/tests/engine/dataset_builders/utils/test_topological_sort.py index d16281fa2..e65ace062 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_dag.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_topological_sort.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + from typing import Any import pytest @@ -135,3 +137,39 @@ def gen_b(row: dict[str, Any]) -> dict[str, Any]: ] with pytest.raises(ConfigCompilationError, match="already produced by"): topologically_sort_column_configs(column_configs) + + +def test_side_effect_column_ordering() -> None: + """A column that depends on a side-effect column is sorted after its producer.""" + + @custom_column_generator(required_columns=["seed"], side_effect_columns=["seed_trace"]) + def gen_with_trace(row: dict[str, Any]) -> dict[str, Any]: + return row + + column_configs = [ + LLMTextColumnConfig(name="seed", prompt="generate seed", model_alias=MODEL_ALIAS), + ExpressionColumnConfig(name="consumer", expr="{{ seed_trace }}"), + CustomColumnConfig(name="producer", generator_function=gen_with_trace), + ] + sorted_configs = topologically_sort_column_configs(column_configs) + names = [c.name for c in sorted_configs] + assert names.index("producer") < names.index("consumer") + + +def test_skip_when_column_ordering() -> None: + """A column with skip.when referencing another DAG column is sorted after that column.""" + from data_designer.config.base import SkipConfig + + column_configs = [ + LLMTextColumnConfig(name="seed", prompt="generate seed", model_alias=MODEL_ALIAS), + LLMTextColumnConfig( + name="gated", + prompt="generate gated", + model_alias=MODEL_ALIAS, + skip=SkipConfig(when="{{ seed == 'bad' }}"), + ), + ] + # gated has no required_columns referencing seed, only a skip.when dependency + sorted_configs = topologically_sort_column_configs(column_configs) + names = [c.name for c in sorted_configs] + assert names.index("seed") < names.index("gated") From b4812448b4d3cd14c376e82168a0b33d9abe8f0d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw?= Date: Fri, 17 Apr 2026 16:11:47 +0200 Subject: [PATCH 5/7] fix(tests): move SkipConfig import to module level in test_topological_sort --- .../engine/dataset_builders/utils/test_topological_sort.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_topological_sort.py b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_topological_sort.py index e65ace062..a3add3fe7 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_topological_sort.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_topological_sort.py @@ -17,6 +17,7 @@ Score, ValidationColumnConfig, ) +from data_designer.config.base import SkipConfig from data_designer.config.column_types import DataDesignerColumnType from data_designer.config.custom_column import custom_column_generator from data_designer.config.sampler_params import SamplerType @@ -158,8 +159,6 @@ def gen_with_trace(row: dict[str, Any]) -> dict[str, Any]: def test_skip_when_column_ordering() -> None: """A column with skip.when referencing another DAG column is sorted after that column.""" - from data_designer.config.base import SkipConfig - column_configs = [ LLMTextColumnConfig(name="seed", prompt="generate seed", model_alias=MODEL_ALIAS), LLMTextColumnConfig( From 2e4843084686add17d4fb688db344d7e665a5dc9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw?= Date: Mon, 20 Apr 2026 21:52:35 +0200 Subject: [PATCH 6/7] refactor(execution-graph): extract nested closures to module-level helpers, fix docs and test style - Extract `resolve`/`_add_edge` nested closures in `topologically_sort_column_configs` to module-level `_resolve_dag_column` and `_add_dag_edge` per STYLEGUIDE.md - Add self-edge guard in `_add_dag_edge` (consistent with `ExecutionGraph.create`) - Update `architecture/dataset-builders.md` to remove stale `dag.py`/NetworkX references - Fix import order in `test_topological_sort.py` (SkipConfig before column_configs) - Add `-> None` return annotations to legacy test functions --- architecture/dataset-builders.md | 8 +-- .../dataset_builders/utils/execution_graph.py | 56 ++++++++++++++----- .../utils/test_topological_sort.py | 6 +- 3 files changed, 46 insertions(+), 24 deletions(-) diff --git a/architecture/dataset-builders.md b/architecture/dataset-builders.md index 0cca1c7d8..d7c56b244 100644 --- a/architecture/dataset-builders.md +++ b/architecture/dataset-builders.md @@ -71,11 +71,7 @@ Both execution modes integrate skip at the same points: - **Sequential**: `_run_full_column_generator` and the fan-out methods (`_fan_out_with_threads`, `_fan_out_with_async`) call `_should_skip_cell` per record. Skipped rows are excluded from the generator input, then merged back with skip metadata preserved. A fast `_column_can_skip` check short-circuits the per-record evaluation when no skip config or propagation applies. - **Async**: `_run_cell` and `_run_batch` in `AsyncTaskScheduler` call `_should_skip_record` / `_apply_skip_to_record` with the same logic. Skipped cells report as skipped (not success) in progress tracking. -DAG edges are added for `skip.when` column references (both in `dag.py` and `ExecutionGraph.create`) so skip-gate columns are generated before the gated column. - -### DAG (Config-Level) - -`dataset_builders/utils/dag.py` provides `topologically_sort_column_configs` ā€” builds a NetworkX graph from `required_columns`, side-effect columns, and `skip.when` references, returns a topological ordering. Used by both execution modes for initial column ordering. +DAG edges are added for `skip.when` column references in `ExecutionGraph.create` so skip-gate columns are generated before the gated column. ### DatasetBatchManager @@ -118,7 +114,7 @@ DatasetBuilder.build() - **Dual execution engines behind one API.** The sequential engine is simpler and easier to debug; the async engine adds row-group parallelism for throughput. Users switch via an environment variable without changing their code. - **DAG-driven ordering** ensures columns with dependencies (e.g., a judge column that depends on a text column) are generated in the correct order, regardless of the order they appear in the config. - **Salvage rounds in async mode** retry failed tasks after all other tasks in a round complete, improving resilience against transient LLM failures without blocking the entire generation. -- **Separate config-level and runtime DAGs.** The config-level DAG (`dag.py`) determines column ordering; the runtime `ExecutionGraph` adds strategy-aware dependency tracking for the async scheduler. +- **Unified DAG construction.** `topologically_sort_column_configs` (in `execution_graph.py`) determines column ordering using Kahn's algorithm; the runtime `ExecutionGraph` adds strategy-aware dependency tracking for the async scheduler. ## Cross-References diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py index 1ac4c2b26..ec4c6e6b5 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py @@ -335,6 +335,45 @@ def to_mermaid(self) -> str: return "\n".join(lines) +def _resolve_dag_column( + col_name: str, + dag_col_dict: dict[str, ColumnConfigT], + side_effect_map: dict[str, str], +) -> str | None: + """Resolve a column name to its DAG producer. + + Returns the column itself if it is a direct DAG column, the producing + column if it is a declared side-effect, or ``None`` if the name is not + known to this DAG (e.g. a seed or sampler column). + """ + if col_name in dag_col_dict: + return col_name + return side_effect_map.get(col_name) + + +def _add_dag_edge( + name: str, + dep: str, + label: str, + dag_col_dict: dict[str, ColumnConfigT], + side_effect_map: dict[str, str], + upstream: dict[str, set[str]], + downstream: dict[str, set[str]], +) -> None: + """Add a dependency edge from *dep*'s producer to *name* if the dep is a known DAG column. + + Self-edges are skipped, consistent with ``ExecutionGraph.create``. + The *label* parameter (``"required"`` or ``"skip.when"``) is included in + the debug log so the source of each edge is visible during tracing. + """ + resolved = _resolve_dag_column(dep, dag_col_dict, side_effect_map) + if resolved is None or resolved == name: + return + logger.debug(f"{LOG_INDENT}šŸ”— `{name}` depends on `{resolved}` [{label}]") + upstream[name].add(resolved) + downstream[resolved].add(name) + + def topologically_sort_column_configs(column_configs: list[ColumnConfigT]) -> list[ColumnConfigT]: non_dag_cols = [col for col in column_configs if not column_type_used_in_execution_dag(col.column_type)] dag_col_dict = {col.name: col for col in column_configs if column_type_used_in_execution_dag(col.column_type)} @@ -355,29 +394,16 @@ def topologically_sort_column_configs(column_configs: list[ColumnConfigT]) -> li ) side_effect_map[se_col] = name - def resolve(col_name: str) -> str | None: - if col_name in dag_col_dict: - return col_name - return side_effect_map.get(col_name) - upstream: dict[str, set[str]] = {name: set() for name in dag_col_dict} downstream: dict[str, set[str]] = {name: set() for name in dag_col_dict} - def _add_edge(name: str, dep: str, label: str) -> None: - resolved = resolve(dep) - if resolved is None: - return - logger.debug(f"{LOG_INDENT}šŸ”— `{name}` depends on `{resolved}` [{label}]") - upstream[name].add(resolved) - downstream[resolved].add(name) - logger.info("ā›“ļø Sorting column configs into a Directed Acyclic Graph") for name, col in dag_col_dict.items(): for req in col.required_columns: - _add_edge(name, req, "required") + _add_dag_edge(name, req, "required", dag_col_dict, side_effect_map, upstream, downstream) if col.skip is not None: for skip_col in col.skip.columns: - _add_edge(name, skip_col, "skip.when") + _add_dag_edge(name, skip_col, "skip.when", dag_col_dict, side_effect_map, upstream, downstream) in_degree = {name: len(ups) for name, ups in upstream.items()} queue: deque[str] = deque(name for name, deg in in_degree.items() if deg == 0) diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_topological_sort.py b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_topological_sort.py index a3add3fe7..6ad9d6504 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_topological_sort.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_topological_sort.py @@ -7,6 +7,7 @@ import pytest +from data_designer.config.base import SkipConfig from data_designer.config.column_configs import ( CustomColumnConfig, ExpressionColumnConfig, @@ -17,7 +18,6 @@ Score, ValidationColumnConfig, ) -from data_designer.config.base import SkipConfig from data_designer.config.column_types import DataDesignerColumnType from data_designer.config.custom_column import custom_column_generator from data_designer.config.sampler_params import SamplerType @@ -30,7 +30,7 @@ MODEL_ALIAS = "stub-model-alias" -def test_dag_construction(): +def test_dag_construction() -> None: column_configs = [] column_configs.append( SamplerMultiColumnConfig( @@ -95,7 +95,7 @@ def test_dag_construction(): assert names[5] == "uses_all_the_stuff" -def test_circular_dependencies(): +def test_circular_dependencies() -> None: column_configs = [] column_configs.append( SamplerMultiColumnConfig( From 5796118d372d33712ea4d7c3d450a577319b8270 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw?= Date: Tue, 21 Apr 2026 20:42:11 +0200 Subject: [PATCH 7/7] refactor(execution-graph): extract shared Kahn helper, inline resolve, fix module-level ordering - Extract `_kahns_topological_sort` shared helper used by both `ExecutionGraph.get_topological_order` and `topologically_sort_column_configs` - Inline `_resolve_dag_column` into `_add_dag_edge` (no other call sites) - Move private helpers after the public function (public-before-private per STYLEGUIDE) - Add docstring to `topologically_sort_column_configs` - Update architecture/dataset-builders.md to mention both sort sites for skip.when edges --- architecture/dataset-builders.md | 2 +- .../dataset_builders/utils/execution_graph.py | 147 +++++++++--------- 2 files changed, 73 insertions(+), 76 deletions(-) diff --git a/architecture/dataset-builders.md b/architecture/dataset-builders.md index d7c56b244..cf06eea17 100644 --- a/architecture/dataset-builders.md +++ b/architecture/dataset-builders.md @@ -71,7 +71,7 @@ Both execution modes integrate skip at the same points: - **Sequential**: `_run_full_column_generator` and the fan-out methods (`_fan_out_with_threads`, `_fan_out_with_async`) call `_should_skip_cell` per record. Skipped rows are excluded from the generator input, then merged back with skip metadata preserved. A fast `_column_can_skip` check short-circuits the per-record evaluation when no skip config or propagation applies. - **Async**: `_run_cell` and `_run_batch` in `AsyncTaskScheduler` call `_should_skip_record` / `_apply_skip_to_record` with the same logic. Skipped cells report as skipped (not success) in progress tracking. -DAG edges are added for `skip.when` column references in `ExecutionGraph.create` so skip-gate columns are generated before the gated column. +DAG edges are added for `skip.when` column references in both `topologically_sort_column_configs` (compile-time sort) and `ExecutionGraph.create` (async runtime) so skip-gate columns are generated before the gated column. ### DatasetBatchManager diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py index ec4c6e6b5..b090cf63d 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py @@ -233,27 +233,12 @@ def get_topological_order(self) -> list[str]: if self._topological_order_cache is not None: return list(self._topological_order_cache) - in_degree: dict[str, int] = {col: 0 for col in self._columns} - for col, deps in self._upstream.items(): - if col in in_degree: - in_degree[col] = len(deps) - - queue = deque(col for col, deg in in_degree.items() if deg == 0) - order: list[str] = [] - while queue: - col = queue.popleft() - order.append(col) - for child in self._downstream.get(col, set()): - if child in in_degree: - in_degree[child] -= 1 - if in_degree[child] == 0: - queue.append(child) - - if len(order) != len(self._columns): - raise DAGCircularDependencyError( - f"The execution graph contains cyclic dependencies. Resolved {len(order)}/{len(self._columns)} columns." - ) - + order = _kahns_topological_sort( + self._columns, + self._upstream, + self._downstream, + "The execution graph contains cyclic dependencies.", + ) self._topological_order_cache = order return list(order) @@ -335,46 +320,17 @@ def to_mermaid(self) -> str: return "\n".join(lines) -def _resolve_dag_column( - col_name: str, - dag_col_dict: dict[str, ColumnConfigT], - side_effect_map: dict[str, str], -) -> str | None: - """Resolve a column name to its DAG producer. - - Returns the column itself if it is a direct DAG column, the producing - column if it is a declared side-effect, or ``None`` if the name is not - known to this DAG (e.g. a seed or sampler column). - """ - if col_name in dag_col_dict: - return col_name - return side_effect_map.get(col_name) - +def topologically_sort_column_configs(column_configs: list[ColumnConfigT]) -> list[ColumnConfigT]: + """Return column configs in dependency order using Kahn's algorithm. -def _add_dag_edge( - name: str, - dep: str, - label: str, - dag_col_dict: dict[str, ColumnConfigT], - side_effect_map: dict[str, str], - upstream: dict[str, set[str]], - downstream: dict[str, set[str]], -) -> None: - """Add a dependency edge from *dep*'s producer to *name* if the dep is a known DAG column. + Non-DAG columns (samplers, seeds) are placed first, followed by DAG columns + sorted by ``required_columns`` and ``skip.when`` edges. Side-effect columns + are resolved to their producing column. - Self-edges are skipped, consistent with ``ExecutionGraph.create``. - The *label* parameter (``"required"`` or ``"skip.when"``) is included in - the debug log so the source of each edge is visible during tracing. + Raises: + ConfigCompilationError: If two columns declare the same side-effect name. + DAGCircularDependencyError: If the dependency graph contains a cycle. """ - resolved = _resolve_dag_column(dep, dag_col_dict, side_effect_map) - if resolved is None or resolved == name: - return - logger.debug(f"{LOG_INDENT}šŸ”— `{name}` depends on `{resolved}` [{label}]") - upstream[name].add(resolved) - downstream[resolved].add(name) - - -def topologically_sort_column_configs(column_configs: list[ColumnConfigT]) -> list[ColumnConfigT]: non_dag_cols = [col for col in column_configs if not column_type_used_in_execution_dag(col.column_type)] dag_col_dict = {col.name: col for col in column_configs if column_type_used_in_execution_dag(col.column_type)} @@ -405,22 +361,63 @@ def topologically_sort_column_configs(column_configs: list[ColumnConfigT]) -> li for skip_col in col.skip.columns: _add_dag_edge(name, skip_col, "skip.when", dag_col_dict, side_effect_map, upstream, downstream) - in_degree = {name: len(ups) for name, ups in upstream.items()} - queue: deque[str] = deque(name for name, deg in in_degree.items() if deg == 0) - order: list[str] = [] - while queue: - name = queue.popleft() - order.append(name) - for child in downstream.get(name, set()): - in_degree[child] -= 1 - if in_degree[child] == 0: - queue.append(child) - - if len(order) != len(dag_col_dict): - raise DAGCircularDependencyError( - "šŸ›‘ The Data Designer column configurations contain cyclic dependencies. Please " - "inspect the column configurations and ensure they can be sorted without " - "circular references." - ) + order = _kahns_topological_sort( + list(dag_col_dict), + upstream, + downstream, + "šŸ›‘ The Data Designer column configurations contain cyclic dependencies. Please " + "inspect the column configurations and ensure they can be sorted without " + "circular references.", + ) return non_dag_cols + [dag_col_dict[n] for n in order] + + +def _add_dag_edge( + name: str, + dep: str, + label: str, + dag_col_dict: dict[str, ColumnConfigT], + side_effect_map: dict[str, str], + upstream: dict[str, set[str]], + downstream: dict[str, set[str]], +) -> None: + """Add a dependency edge from *dep*'s producer to *name* if the dep is a known DAG column. + + Self-edges are skipped, consistent with ``ExecutionGraph.create``. + The *label* parameter (``"required"`` or ``"skip.when"``) is included in + the debug log so the source of each edge is visible during tracing. + """ + resolved = dep if dep in dag_col_dict else side_effect_map.get(dep) + if resolved is None or resolved == name: + return + logger.debug(f"{LOG_INDENT}šŸ”— `{name}` depends on `{resolved}` [{label}]") + upstream[name].add(resolved) + downstream[resolved].add(name) + + +def _kahns_topological_sort( + nodes: list[str], + upstream: dict[str, set[str]], + downstream: dict[str, set[str]], + error_message: str, +) -> list[str]: + """Return a topological ordering of *nodes* using Kahn's algorithm. + + Raises: + DAGCircularDependencyError: If the graph contains a cycle. + """ + in_degree: dict[str, int] = {col: len(upstream.get(col, set())) for col in nodes} + queue: deque[str] = deque(col for col, deg in in_degree.items() if deg == 0) + order: list[str] = [] + while queue: + col = queue.popleft() + order.append(col) + for child in downstream.get(col, set()): + if child in in_degree: + in_degree[child] -= 1 + if in_degree[child] == 0: + queue.append(child) + if len(order) != len(nodes): + raise DAGCircularDependencyError(error_message) + return order