[branching io manager][fix] Fix loading partitioned parent assets in BranchingIOManager #18491

benpankow · 2023-12-05T00:11:03Z

Summary

Tweak of fix in #17303 which uses the upstream asset partition rather than the current asset partition when using the branching IO manager. Uses context.asset_partition_key rather than the upstream output's partition key, which can sometimes be None when the upstream output is not available.

Test Plan

Adds a new unit test with a non-standard partition mapping, which previously failed.

benpankow · 2023-12-05T00:11:14Z

Current dependencies on/for this PR:

master
- PR [branching io manager][fix] Fix loading partitioned parent assets in BranchingIOManager #18491 👈

This stack of pull requests is managed by Graphite.

smackesey · 2023-12-05T12:59:57Z

I'm confused as to why InputContext.asset_partition_key and InputContext.partition_key differ in the first place. Shouldn't the InputContext.partition_key be scoped to the input being loaded instead of the run? Do we ever use the run value?

danielgafni · 2023-12-05T15:24:42Z

We might not be loading an asset at all. The IOManager should also work for normal jobs. This is because it’s often used as a drop-in replacement for an existing IOManager.

smackesey · 2023-12-05T15:31:33Z

We might not be loading an asset at all. The IOManager should also work for normal jobs. This is because it’s often used as a drop-in replacement for an existing IOManager.

That's true, but the question I'm raising is about InputContext.partition_key. Seems to me this should mean "partition key for the input that this InputContext is defined for". In cases where that key matches the run partition key, then OK. But when it doesn't, then it should be set to the appropriate key.

danielgafni · 2023-12-05T15:42:01Z

Sorry, I don’t see any tests logs or anything. Maybe that’s because I’m on mobile, not sure. Anyway, keeping this use case in mind (working with non-assets) is all what I wanted to point out, and it seems like we are on the same page about this :)

benpankow · 2023-12-05T23:44:37Z

I'm confused as to why InputContext.asset_partition_key and InputContext.partition_key differ in the first place. Shouldn't the InputContext.partition_key be scoped to the input being loaded instead of the run? Do we ever use the run value?

I'm not sure why it's set to be the partition for the asset that's currently being materialized.

It is pretty confusing especially since the two properties would seemingly be the same thing - the input manager should really be agnostic to what's being produced in the actual compute step.

We might not be loading an asset at all. The IOManager should also work for normal jobs. This is because it’s often used as a drop-in replacement for an existing IOManager.

This is a good call-out - right now if there's no upstream asset involved the IO manager will fallback to None partition, so it should continue to work in the op case. Will add a test just to sanity check.

smackesey · 2023-12-06T12:51:23Z

I'm not sure why it's set to be the partition for the asset that's currently being materialized.

It is pretty confusing especially since the two properties would seemingly be the same thing - the input manager should really be agnostic to what's being produced in the actual compute step.

My inclination is to consider this a bug and change it so that partition_key and asset_partition_key are the same. Curious what @alangenfeld thinks of this.

To summarize Alex, InputContext.partition_key currently is the partition key of the run while InputContext.asset_partition_key is the partition_key of the key of the upstream asset partition being loaded. IMO partition_key should always indicate the partition key of the input being loaded, which would be the same as asset_partition_key for asset runs.

alangenfeld

This touches on the complexity around how partitioning works and is modeled that I don't have a great grasp on.

I don't think its OK to make a breaking change the @public partition_key method, until 2.0 given my understanding of our policy. My first reaction was that maybe we could deprecate it and replace it with run_partition_key to more clearly disambiguate, but I don't have confidence that thats the right move.

When does the upstream output partition key differ from the run partition key in the non asset case? From what I understand you can only attach PartitionMapping to assets.

alangenfeld · 2023-12-06T17:26:34Z

python_modules/dagster/dagster/_core/storage/branching/branching_io_manager.py

+            try:
+                partition_key = context.asset_partition_key
+            except CheckError:
+                pass


is there not a better way to do this? would prefer adding a has_ over control flow exceptions

Yeah, updated to handle the multi-partition-key case and this logic is a bit cleaner now.

python_modules/dagster/dagster_tests/storage_tests/branching_io_manager_tests/utils.py

alangenfeld · 2023-12-06T17:54:13Z

python_modules/dagster/dagster/_core/storage/branching/branching_io_manager.py

+                else None
+            )
+            try:
+                partition_key = context.asset_partition_key


what about for other complex partition mapping schemes where asset_partition_key would raise due to there being more than 1 value? https://github.com/dagster-io/dagster/blob/master/python_modules/dagster/dagster/_core/execution/context/input.py#L347-L350

Updated logic & added test case, should be able to handle non-standard (e.g. many-to-one) partition mappings now.

alangenfeld · 2023-12-06T18:02:00Z

python_modules/dagster/dagster/_core/storage/branching/branching_io_manager.py

@@ -69,10 +70,20 @@ def load_input(self, context: InputContext) -> Any:
            return self.branch_io_manager.load_input(context)
        else:
            # we are dealing with an asset input
+            partition_key = (
+                context.upstream_output.partition_key
+                if context.upstream_output and context.upstream_output.has_partition_key


this appears to be not ok
https://github.com/dagster-io/dagster/blob/master/python_modules/dagster/dagster/_core/execution/context/output.py#L363-L368

#7900

smackesey · 2023-12-07T18:12:07Z

When does the upstream output partition key differ from the run partition key in the non asset case? From what I understand you can only attach PartitionMapping to assets.

Pretty sure it doesn't. So it would only change in the asset case. I am supportive of introducing run_partition_key and changing partition_key to reflect the input in 2.0 (with a warning introduced in the meantime, that we can emit only if we are loading an asset).

alangenfeld

this version seems like a more principled fix to me, I assume the cases that danielgafni mentioned above are under test somewhere so this continues to work for them?

alangenfeld · 2023-12-11T18:52:51Z

python_modules/dagster/dagster/_core/storage/branching/branching_io_manager.py

+                )
+                for partition_key in partition_keys
+            ]
+            if all(


careful here on all([]) == True if nothing comes back - may want an explicit check for it

https://github.com/dagster-io/dagster/pull/18491/files/7bee3fd9cdf04c032e6955d5eadd05d5c17b7729#r1422995242 should ensure that the length is always >=1, will add a note though

alangenfeld · 2023-12-11T18:53:41Z

python_modules/dagster/dagster/_core/storage/branching/branching_io_manager.py

-            )
-            if (
-                event_log_entry
+


i think a little comment block explaining why were doing this would be nice, make [1] easier to grok

alangenfeld · 2023-12-11T18:54:16Z

python_modules/dagster/dagster/_core/storage/branching/branching_io_manager.py

+                partition_keys = context.asset_partition_keys
+
+            if len(partition_keys) == 0:
+                partition_keys = [None]


…BranchingIOManager

github-actions · 2023-12-11T23:34:04Z

Deploy preview for dagit-storybook ready!

✅ Preview
https://dagit-storybook-eebygj240-elementl.vercel.app
https://benpankow-fix-branching-io-manager.components-storybook.dagster-docs.io

Built with commit 39e178d.
This pull request is being automatically deployed with vercel-action

github-actions · 2023-12-11T23:34:37Z

Deploy preview for dagit-core-storybook ready!

✅ Preview
https://dagit-core-storybook-8z1knpx0x-elementl.vercel.app
https://benpankow-fix-branching-io-manager.core-storybook.dagster-docs.io

Built with commit 39e178d.
This pull request is being automatically deployed with vercel-action

github-actions · 2023-12-11T23:36:53Z

Deploy preview for dagster-docs ready!

Preview available at https://dagster-docs-39dn3bt9j-elementl.vercel.app
https://benpankow-fix-branching-io-manager.dagster.dagster-docs.io

Direct link to changed pages:

benpankow mentioned this pull request Dec 5, 2023

🐛 fix loading partitioned parent assets in BranchingIOManager #17303

Closed

benpankow requested review from alangenfeld, schrockn and smackesey and removed request for alangenfeld December 5, 2023 00:13

benpankow force-pushed the benpankow/fix-branching-io-manager branch from 13f501c to dc9f19e Compare December 5, 2023 23:49

alangenfeld reviewed Dec 6, 2023

View reviewed changes

benpankow force-pushed the benpankow/fix-branching-io-manager branch from 7531548 to 4b47f02 Compare December 7, 2023 18:09

benpankow requested a review from alangenfeld December 7, 2023 18:09

benpankow force-pushed the benpankow/fix-branching-io-manager branch from c7daeea to 7bee3fd Compare December 7, 2023 18:28

alangenfeld approved these changes Dec 11, 2023

View reviewed changes

Ben and others added 7 commits December 11, 2023 15:30

[branching io manager][fix] Fix loading partitioned parent assets in …

8807040

…BranchingIOManager

add job/op test case

8954875

fix job/op test case

e00cba9

Properly handle the multi-input case

03df00d

add another test

7461306

get non asset partition key too

d02eed2

add more comments

39e178d

benpankow force-pushed the benpankow/fix-branching-io-manager branch from 7bee3fd to 39e178d Compare December 11, 2023 23:31

benpankow merged commit 262853f into master Dec 12, 2023
3 checks passed

benpankow deleted the benpankow/fix-branching-io-manager branch December 12, 2023 00:22

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[branching io manager][fix] Fix loading partitioned parent assets in BranchingIOManager #18491

[branching io manager][fix] Fix loading partitioned parent assets in BranchingIOManager #18491

benpankow commented Dec 5, 2023 •

edited

Loading

benpankow commented Dec 5, 2023

smackesey commented Dec 5, 2023

danielgafni commented Dec 5, 2023 •

edited

Loading

smackesey commented Dec 5, 2023 •

edited

Loading

danielgafni commented Dec 5, 2023

benpankow commented Dec 5, 2023

smackesey commented Dec 6, 2023 •

edited

Loading

alangenfeld left a comment

alangenfeld Dec 6, 2023

benpankow Dec 7, 2023

alangenfeld Dec 6, 2023

benpankow Dec 7, 2023 •

edited

Loading

alangenfeld Dec 6, 2023

smackesey commented Dec 7, 2023

alangenfeld left a comment

alangenfeld Dec 11, 2023

benpankow Dec 11, 2023

alangenfeld Dec 11, 2023

alangenfeld Dec 11, 2023

github-actions bot commented Dec 11, 2023

github-actions bot commented Dec 11, 2023

github-actions bot commented Dec 11, 2023

[branching io manager][fix] Fix loading partitioned parent assets in BranchingIOManager #18491

[branching io manager][fix] Fix loading partitioned parent assets in BranchingIOManager #18491

Conversation

benpankow commented Dec 5, 2023 • edited Loading

Summary

Test Plan

benpankow commented Dec 5, 2023

smackesey commented Dec 5, 2023

danielgafni commented Dec 5, 2023 • edited Loading

smackesey commented Dec 5, 2023 • edited Loading

danielgafni commented Dec 5, 2023

benpankow commented Dec 5, 2023

smackesey commented Dec 6, 2023 • edited Loading

alangenfeld left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

benpankow Dec 7, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

smackesey commented Dec 7, 2023

alangenfeld left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented Dec 11, 2023

github-actions bot commented Dec 11, 2023

github-actions bot commented Dec 11, 2023

benpankow commented Dec 5, 2023 •

edited

Loading

danielgafni commented Dec 5, 2023 •

edited

Loading

smackesey commented Dec 5, 2023 •

edited

Loading

smackesey commented Dec 6, 2023 •

edited

Loading

benpankow Dec 7, 2023 •

edited

Loading