Skip to content

fix: async engine side-effect column propagation and collision resolution #509

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
andreatnvidia merged 7 commits into from
Apr 13, 2026
Merged

fix: async engine side-effect column propagation and collision resolution #509

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
b88508d
fix: async engine side-effect column propagation and collision resolu…
andreatnvidia Apr 8, 2026
c378684
fix: separate side-effect columns from completion tracking in async s…
andreatnvidia Apr 9, 2026
ece8f8c
fix: warn on side-effect collision and clarify scheduler column maps
andreatnvidia Apr 10, 2026
818ef66
merge: resolve conflict with main in test_async_scheduler.py
andreatnvidia Apr 13, 2026
d47abcb
Merge branch 'main' into andreatgretel/fix/async-side-effect-columns
andreatnvidia Apr 13, 2026
604aa21
fix: raise ConfigCompilationError on duplicate side-effect producers
andreatnvidia Apr 13, 2026
ba84679
fix: reject duplicate side-effect producers in sync DAG path
andreatnvidia Apr 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 22 additions & 10 deletions packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,12 +124,24 @@ def __init__(
self._disable_early_shutdown = disable_early_shutdown
self._early_shutdown = False

# Multi-column dedup: group output columns by generator identity
# Multi-column dedup: group output columns by generator identity.
# _instance_to_columns holds only real (graph-registered) columns and
# is used for completion tracking. _instance_to_write_columns extends
# that with side-effect columns for buffer writes only.
instance_to_columns: dict[int, list[str]] = {}
for col, gen in generators.items():
instance_to_columns.setdefault(id(gen), []).append(col)
self._instance_to_columns = instance_to_columns

seen_cols: set[str] = {col for col in generators}
instance_to_write_columns: dict[int, list[str]] = {k: list(v) for k, v in instance_to_columns.items()}
for col, gen in generators.items():
for side_effect_col in getattr(gen.config, "side_effect_columns", []):
if side_effect_col not in seen_cols:
instance_to_write_columns.setdefault(id(gen), []).append(side_effect_col)
seen_cols.add(side_effect_col)
Comment thread
andreatnvidia marked this conversation as resolved.
Resolved
self._instance_to_write_columns = instance_to_write_columns
Comment thread
andreatnvidia marked this conversation as resolved.
Resolved

# Stateful generator tracking: instance_id → asyncio.Lock
self._stateful_locks: dict[int, asyncio.Lock] = {}
for col, gen in generators.items():
Expand Down Expand Up @@ -754,10 +766,10 @@ async def _run_from_scratch(self, task: Task, generator: ColumnGenerator) -> Any
else:
result_df = await generator.agenerate(lazy.pd.DataFrame())

# Write results to buffer
# Write results to buffer (include side-effect columns)
if self._buffer_manager is not None:
output_cols = self._instance_to_columns.get(id(generator), [task.column])
for col in output_cols:
write_cols = self._instance_to_write_columns.get(id(generator), [task.column])
for col in write_cols:
if col in result_df.columns:
values = result_df[col].tolist()
self._buffer_manager.update_batch(task.row_group, col, values)
Expand All @@ -780,10 +792,10 @@ async def _run_cell(self, task: Task, generator: ColumnGenerator) -> Any:

result = await generator.agenerate(row_data)

# Write back to buffer
# Write back to buffer (include side-effect columns)
if self._buffer_manager is not None and not self._tracker.is_dropped(task.row_group, task.row_index):
output_cols = self._instance_to_columns.get(id(generator), [task.column])
for col in output_cols:
write_cols = self._instance_to_write_columns.get(id(generator), [task.column])
for col in write_cols:
if col in result:
self._buffer_manager.update_cell(task.row_group, task.row_index, col, result[col])

Expand All @@ -804,9 +816,9 @@ async def _run_batch(self, task: Task, generator: ColumnGenerator) -> Any:

result_df = await generator.agenerate(batch_df)

# Merge result columns back to buffer
# Merge result columns back to buffer (include side-effect columns)
if self._buffer_manager is not None:
output_cols = self._instance_to_columns.get(id(generator), [task.column])
write_cols = self._instance_to_write_columns.get(id(generator), [task.column])
active_rows = rg_size - len(pre_dropped)
if len(result_df) != active_rows:
raise ValueError(
Expand All @@ -819,7 +831,7 @@ async def _run_batch(self, task: Task, generator: ColumnGenerator) -> Any:
continue
# Skip writing to rows dropped by concurrent tasks during the await
if not self._buffer_manager.is_dropped(task.row_group, ri):
for col in output_cols:
for col in write_cols:
if col in result_df.columns:
self._buffer_manager.update_cell(task.row_group, ri, col, result_df.iloc[result_idx][col])
result_idx += 1
Expand Down
11 changes: 9 additions & 2 deletions ...s/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,8 +105,15 @@ def add_edge(self, upstream: str, downstream: str) -> None:
self._downstream.setdefault(upstream, set()).add(downstream)

def set_side_effect(self, side_effect_col: str, producer: str) -> None:
"""Map a side-effect column name to its producing column."""
self._side_effect_map[side_effect_col] = producer
"""Map a side-effect column name to its producing column.

First-writer-wins: if a side-effect column is already mapped to a
different producer, the earlier mapping is kept. This matches sync
engine semantics where earlier consumers see the first producer's
value.
"""
if side_effect_col not in self._side_effect_map:
self._side_effect_map[side_effect_col] = producer

@nabinchha nabinchha Apr 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isn't a side effect column 1:1 with its producer? In which scenario would multiple producers point to the same side effect column?

@andreatnvidia andreatnvidia Apr 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right that in a correctly configured pipeline it should be 1:1. The collision scenario comes from Anonymizer's detection workflow - it has branching pipeline paths where different stages declare the same side-effect columns (e.g. _merged_tagged_text is declared by both prepare_validation_inputs and merge_and_build_candidates), but only one path is active per run. With last-writer-wins the graph resolves the wrong producer and creates a false cycle.

That said, this is an unusual pattern and silently ignoring the second registration isn't great. I'll add a warning log on collision so misconfigurations don't go unnoticed.

@nabinchha nabinchha Apr 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in this example wouldn't the side effect columns be prepare_validation_inputs _merged_tagged_text and merge_and_build_candidates_merged_tagged_text? It's still 1:1 right?

@andreatnvidia andreatnvidia Apr 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I looked into this - the collision appears to be real. Side-effect column names are bare strings with no namespacing. In Anonymizer's detection workflow (custom_columns.py), both merge_and_build_candidates and prepare_validation_inputs declare the exact same side-effect columns:

# merge_and_build_candidates (line 79)
side_effect_columns=[COL_MERGED_TAGGED_TEXT, COL_VALIDATION_CANDIDATES]

# prepare_validation_inputs (line 125)
side_effect_columns=[COL_MERGED_TAGGED_TEXT, COL_VALIDATION_CANDIDATES]

Both resolve to the same literal strings ("_merged_tagged_text", "_validation_candidates"). As far as I can tell, this is intentional - the two generators are alternative producers for the same downstream consumer depending on which pipeline path runs.

I considered whether namespacing (e.g. producer_name._merged_tagged_text) could solve this, but the downstream column config wouldn't know which producer ran, so it can't pick the right prefixed name. You'd likely need a union/alias mechanism in the DAG, which feels like a lot of complexity for a pattern that already works. The @custom_column_generator API is also public, so naming changes could break external plugins.

Open to other ideas, but first-writer-wins + warning seems like the right pragmatic fix here.

@andreatnvidia andreatnvidia Apr 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After our offline discussion, I agree with you - multiple producers for the same side-effect column is an anti-pattern. The Anonymizer case is actually an overwrite pattern within the same pipeline (not alternative paths), which works in the sync engine by accident.

I've updated set_side_effect() to raise ConfigCompilationError on duplicate producers instead of silently picking a winner (604aa21). We'll follow up on the Anonymizer side with distinct column names per stage.


def resolve_side_effect(self, column: str) -> str:
"""Resolve a column name through the side-effect map.
Expand Down
44 changes: 44 additions & 0 deletions packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1430,3 +1430,47 @@ async def test_scheduler_rg_semaphore_deadlock_with_transient_failures() -> None

assert tracker.is_row_group_complete(0, 2, ["seed", "col"])
assert tracker.is_row_group_complete(1, 2, ["seed", "col"])


def test_side_effect_columns_separated_from_completion_tracking() -> None:
"""Side-effect columns must appear in _instance_to_write_columns (buffer writes)
but NOT in _instance_to_columns (completion tracking), because they are not
registered in the execution graph and would cause KeyError in CompletionTracker.
"""
graph = ExecutionGraph()
graph.add_column("seed", GenerationStrategy.FULL_COLUMN)
graph.add_column("primary", GenerationStrategy.CELL_BY_CELL)
graph.add_edge(upstream="seed", downstream="primary")

row_groups = [(0, 2)]
tracker = CompletionTracker.with_graph(graph, row_groups)

provider = _mock_provider()
seed_gen = MockSeedGenerator(config=_expr_config("seed"), resource_provider=provider)
cell_gen = MockCellGenerator(config=_expr_config("primary"), resource_provider=provider)
# Replace the config with a mock that reports side-effect columns.
mock_config = MagicMock()
mock_config.side_effect_columns = ["side_a", "side_b"]
object.__setattr__(cell_gen, "_config", mock_config)

generators: dict[str, ColumnGenerator] = {"seed": seed_gen, "primary": cell_gen}

scheduler = AsyncTaskScheduler(
generators=generators,
graph=graph,
tracker=tracker,
row_groups=row_groups,
)

cell_id = id(cell_gen)

# Completion tracking dict: only real columns
assert "side_a" not in scheduler._instance_to_columns.get(cell_id, [])
assert "side_b" not in scheduler._instance_to_columns.get(cell_id, [])
assert "primary" in scheduler._instance_to_columns.get(cell_id, [])

# Buffer write dict: includes side-effect columns
write_cols = scheduler._instance_to_write_columns.get(cell_id, [])
assert "primary" in write_cols
assert "side_a" in write_cols
assert "side_b" in write_cols
18 changes: 18 additions & 0 deletions packages/data-designer-engine/tests/engine/dataset_builders/utils/test_execution_graph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,24 @@ def test_side_effect_name_collision_prefers_real_column() -> None:
assert graph.get_downstream_columns("summary") == set()


def test_side_effect_collision_first_writer_wins() -> None:
"""When two producers declare the same side-effect, the first registration wins."""
graph = ExecutionGraph()
graph.add_column("producer_a", GenerationStrategy.CELL_BY_CELL)
graph.add_column("producer_b", GenerationStrategy.CELL_BY_CELL)
graph.add_column("consumer", GenerationStrategy.CELL_BY_CELL)

graph.set_side_effect("shared_se", "producer_a")
graph.set_side_effect("shared_se", "producer_b") # should be ignored

assert graph.resolve_side_effect("shared_se") == "producer_a"

# consumer depends on shared_se -> should resolve to producer_a, not producer_b
resolved = graph.resolve_side_effect("shared_se")
graph.add_edge(upstream=resolved, downstream="consumer")
assert graph.get_upstream_columns("consumer") == {"producer_a"}


# -- Validation tests -------------------------------------------------------


Expand Down
Loading