Stream generator nodes decorated with functools.wraps by charan-rathore · Pull Request #5798 · kedro-org/kedro

charan-rathore · 2026-09-29T06:22:12Z

Description

Fixes #5796. Kedro checks whether a node is a generator function before streaming its output. A decorator made with functools.wraps hides that function from inspect.isgeneratorfunction, so the runner saved the generator object once instead of saving each chunk. The two runner checks now unwrap the callable for this decision. This also makes async loading and saving reject a decorated generator the same way as an undecorated one. Other iterable objects remain ordinary outputs.

Development notes

Added tests for chunk-by-chunk saves and the async rejection. Updated the generator-node documentation and upcoming release note. pytest tests/runner/test_run_generator_node.py tests/runner/test_task.py -q -o addopts='': 20 passed. The runner suite passed with one unrelated CPU-count-sensitive test excluded (306 passed, one skipped, ten deselected); without exclusion, 315 passed, one skipped, one unrelated failure. Ruff import/format checks pass.

Instinct assisted with the change and tests; Charan reviewed before opening.

Developer Certificate of Origin

The commit is signed off as Charan Rathore.

Checklist

Read the contributing guidelines
Linked to a relevant GitHub issue
Signed off the commit with DCO
Opened as draft if still in progress (select when submitting)
Updated documentation
Added RELEASE.md entry
Added tests

Signed-off-by: Charan Rathore <[email protected]>

Add test for wrapped generator function in async context Signed-off-by: Charan Rathore <[email protected]>

Added a test to ensure that a wrapped generator node streams chunks correctly without collapsing into a single save. Signed-off-by: Charan Rathore <[email protected]>

Clarify usage of generator nodes and their decorators in Kedro. Signed-off-by: Charan Rathore <[email protected]>

Stream outputs from generator nodes decorated with `functools.wraps` and reject these nodes when asynchronous loading and saving is enabled. Signed-off-by: Charan Rathore <[email protected]>

# Conflicts: # RELEASE.md

charan-rathore · 2026-10-04T09:03:52Z

Sorry, this was closed by mistake, reopening.

Signed-off-by: Merel Theisen <[email protected]>

merelcht

Thanks @charan-rathore for tackling this. The fix works for the reported case, but I have one blocking concern about a behaviour regression, plus a few smaller points inline.

merelcht · 2026-10-05T15:33:37Z

        # business objects (e.g. ``mne.Epochs``) are passed to
        # ``catalog.save`` unchanged. See kedro-org/kedro#5412.
-        if outputs and inspect.isgeneratorfunction(node.func):
+        if outputs and inspect.isgeneratorfunction(inspect.unwrap(node.func)):


This now decides whether to stream from the declared function, not from what the node actually returns, which silently changes behaviour for existing nodes. A @wraps decorator that consumes the generator (e.g. return list(f(*args, **kwargs))) worked before: the list was saved once. With this change the runner iterates the list and calls catalog.save once per element, which for datasets that overwrite on save keeps only the last element. The same applies to a @wraps(gen) wrapper returning a str, DataFrame or dict (the runner would iterate characters, column labels or keys). The documented caveat is the only safeguard.

Node._outputs_to_dictionary already uses inspect.isgenerator on the actual return value. Could we stream only when the unwrapped function is a generator function and the outputs are generators, and otherwise keep the old behaviour?

merelcht · 2026-10-05T15:33:37Z

    def execute(self) -> Node:
-        if self.is_async and inspect.isgeneratorfunction(self.node.func):
+        if self.is_async and inspect.isgeneratorfunction(
+            inspect.unwrap(self.node.func)


The same unwrap logic is now duplicated at both call sites, and it has two gaps:

inspect.unwrap raises ValueError: wrapper loop for cyclic __wrapped__ chains, which would surface as an opaque error here where main was fine.

functools.partial(wrapped) is still not detected, because unwrap doesn't follow partial.func.

A small private helper (e.g. _is_generator_node_func) that handles partial, unwrap and the cycle fallback would cover both call sites.

merelcht · 2026-10-05T15:33:38Z

        assert fake_dataset.save.call_count == 10
        assert fake_dataset.save.call_args_list == expected

+    def test_wrapped_generator_node_streams_chunks(self, mocker, catalog):


Nice that this fails on main. Could we add a few more cases?

A decorator that consumes the generator (list(...)), asserting it is still saved once. This would have caught the regression above.

partial(wrapped).

A class-based decorator using update_wrapper.

A multi-output (tuple/dict) wrapped generator.

ParallelRunner, which pickles node.func, so it needs a module-level decorated function rather than a local closure.

merelcht · 2026-10-05T15:33:38Z

 kedro new --starter=pandas-iris --checkout=0.18.14
 ```

+A generator node wrapped by a decorator using `functools.wraps` is streamed in the same way as the original generator function. The decorator must return the generator rather than consuming it. Asynchronous data loading and saving does not support generator nodes, decorated or not.


This paragraph sits at the end of "Set up the project", straight after the legacy kedro new --starter=pandas-iris --checkout=0.18.14 block, so it reads as a non-sequitur and falls under the 0.18.14-only context. Could it move to the generator sections below, perhaps under its own heading such as "Use decorators with generator nodes"?

I'd also tighten the wording: "consuming" is jargon, and "does not support" doesn't say what happens. For example:

If you decorate a generator function with a decorator that uses `functools.wraps`, Kedro still streams its output to the catalog in the same way as the undecorated function. The decorator must return the generator produced by the wrapped function without iterating over it, for example by calling `list()` on it. !!! note Generator nodes, whether decorated or not, cannot be run with asynchronous data loading and saving. Kedro raises a `ValueError` if you use `--async` with a generator node.

merelcht

Requesting changes because of the behaviour regression described in my inline comment on kedro/runner/task.py:178: a @wraps decorator that consumes the generator would now be streamed element by element instead of saved once. The other inline comments (shared helper, extra tests, release note, docs placement) are smaller and can be addressed together.

Signed-off-by: Charan Rathore <[email protected]>

charan-rathore · 2026-10-05T15:48:19Z

Thanks, you are right about the regression. Pushed 222664c.

Streaming now needs both things: the unwrapped function is a generator function, and what the node actually returned is iterators (so a wrapper that returns list(...), a str, a DataFrame or a dict is saved once, as before). I check for iterators and not only generators because unzip can return map objects for multi-output nodes, which an existing test covers.
A small private helper _is_generator_node_func in kedro/runner/task.py follows partial.func and __wrapped__, and stops on a cyclic __wrapped__ chain instead of raising. Both call sites use it. The upfront async check is back to the original one. A wrapped generator that is actually streaming now raises the same ValueError in the async path, and a consuming wrapper still works with async.
New tests: consuming wrapper saved once (also with async and with ParallelRunner), partial(wrapped), a class-based decorator using update_wrapper, wrapped tuple and dict multi-output, and a wrapped generator under ParallelRunner with module-level functions. The consuming-wrapper and partial tests fail on the previous commit.
Moved the docs paragraph to its own section, "Use decorators with generator nodes", at the end of the generator section, with your wording and the async note. Updated the release note.

I ran tests/runner and tests/pipeline (702 passed). 3 tests fail in my environment, and the same 3 fail without my change: two test_only_missing_outputs_with_dataset_factory cases (the pickle dataset is not installed here) and one thread runner max_workers case. I ran ruff format and mypy on task.py. I did not run the rest of the suite.

Prepared with AI assistance; Charan Rathore is directing and responsible for this PR.

charan-rathore · 2026-10-07T08:00:23Z

Gentle nudge - the unwrap logic is consolidated and the extra generator cases are in. Anything else you'd like changed before this can move forward?

merelcht

Thanks @charan-rathore this looks good to me now 👍

I'll ask the team for a second review.

SajidAlamQB

Thanks @charan-rathore this looks good to me, PR is blocked due to coverage so please add the necessary tests. Happy to approve once done.

CI coverage requires every line: the cyclic-chain exit in _is_generator_node_func (kedro/runner/task.py:351) was the one uncovered line. The test loops two functions' __wrapped__ attributes into each other and asserts the walk terminates and reports the function is not a generator. Signed-off-by: Charan Rathore <[email protected]>

…1/2) Signed-off-by: Charan Rathore <[email protected]>

CI coverage requires every line: the cyclic-chain exit in _is_generator_node_func (kedro/runner/task.py:351) was the one uncovered line. The test loops two functions' __wrapped__ attributes into each other and asserts the walk terminates and reports the function is not a generator. Signed-off-by: Charan Rathore <[email protected]>

charan-rathore · 2026-10-08T12:27:44Z

Thanks Sajid. The one uncovered line was the cyclic wrapped chain exit in _is_generator_node_func (kedro/runner/task.py:351), so I added test_cyclic_wrapped_chain_is_not_a_generator to tests/runner/test_task.py: two functions whose wrapped attributes point at each other, asserting the walk terminates instead of looping and reports the function is not a generator.

Verified locally: tests/runner/test_task.py is 11/11, and a targeted coverage run of that file shows line 351 covered now. (The commit landed as a small chain via the web UI - my push access was playing up - but the net diff is just the import and the new test.)

ankatiyar

Thanks @charan-rathore

charan-rathore added 5 commits September 29, 2026 11:49

Unwrap decorated callables for generator node checks

3330ffe

Signed-off-by: Charan Rathore <[email protected]>

Test generator checks through functools.wraps

6187b26

Add test for wrapped generator function in async context Signed-off-by: Charan Rathore <[email protected]>

Test streaming for wraps-decorated generator nodes

96e460f

Added a test to ensure that a wrapped generator node streams chunks correctly without collapsing into a single save. Signed-off-by: Charan Rathore <[email protected]>

Document wraps-decorated generator nodes

e991df3

Clarify usage of generator nodes and their decorators in Kedro. Signed-off-by: Charan Rathore <[email protected]>

Add release note for generator node streaming fix

9c96fc2

Stream outputs from generator nodes decorated with `functools.wraps` and reject these nodes when asynchronous loading and saving is enabled. Signed-off-by: Charan Rathore <[email protected]>

charan-rathore requested review from AliceCima10, Huongg and merelcht as code owners September 29, 2026 06:22

MohammadHijjawi97 mentioned this pull request Sep 29, 2026

Generator node wrapped with a functools.wraps decorator is saved as a generator object instead of being streamed #5796

Closed

merelcht requested review from ankatiyar and removed request for AliceCima10 and Huongg September 30, 2026 10:20

Merge branch 'upmain' into pr

fbcf46f

# Conflicts: # RELEASE.md

charan-rathore closed this Oct 4, 2026

charan-rathore reopened this Oct 4, 2026

Merge branch 'main' into fix-wraps-generator-nodes-5796

2764d61

Signed-off-by: Merel Theisen <[email protected]>

merelcht reviewed Oct 5, 2026

View reviewed changes

merelcht requested changes Oct 5, 2026

View reviewed changes

Stream wrapped generator nodes only when they return generators

222664c

Signed-off-by: Charan Rathore <[email protected]>

Merge branch 'main' into fix-wraps-generator-nodes-5796

5bdf2a6

merelcht approved these changes Oct 7, 2026

View reviewed changes

SajidAlamQB self-requested a review October 8, 2026 10:50

SajidAlamQB reviewed Oct 8, 2026

View reviewed changes

charan-rathore added 3 commits October 8, 2026 17:55

Replace test_task.py with coverage for the cyclic __wrapped__ guard (…

8bab112

…1/2) Signed-off-by: Charan Rathore <[email protected]>

ankatiyar approved these changes Oct 8, 2026

View reviewed changes

merelcht and others added 2 commits October 12, 2026 10:27

Merge branch 'main' into fix-wraps-generator-nodes-5796

71b7fe0

Merge branch 'main' into fix-wraps-generator-nodes-5796

0bc762c

ankatiyar merged commit 761a612 into kedro-org:main Oct 12, 2026
28 checks passed

charan-rathore deleted the fix-wraps-generator-nodes-5796 branch October 12, 2026 10:04

Conversation

charan-rathore commented Sep 29, 2026

Description

Development notes

Developer Certificate of Origin

Checklist

Uh oh!

charan-rathore commented Oct 4, 2026

Uh oh!

merelcht left a comment

Choose a reason for hiding this comment

Uh oh!

merelcht Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

merelcht Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

merelcht Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

merelcht Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

merelcht left a comment

Choose a reason for hiding this comment

Uh oh!

charan-rathore commented Oct 5, 2026

Uh oh!

charan-rathore commented Oct 7, 2026

Uh oh!

merelcht left a comment

Choose a reason for hiding this comment

Uh oh!

SajidAlamQB left a comment

Choose a reason for hiding this comment

Uh oh!

charan-rathore commented Oct 8, 2026

Uh oh!

ankatiyar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants