add indirect execution context access #14954

alangenfeld · 2023-06-26T17:22:50Z

adds a means to get the current op/asset execution context via a function call, setup using a ContextVar. Exposed as OpExecutionContext.get() & AssetExecutionContext.get()

How I Tested These Changes

added test

alangenfeld · 2023-06-26T17:23:01Z

Current dependencies on/for this PR:

master
- PR add indirect execution context access #14954 👈
  - PR [prototype] return AssetMaterialization #14931

This stack of pull requests is managed by Graphite.

github-actions · 2023-07-10T20:34:59Z

Deploy preview for dagit-core-storybook ready!

✅ Preview
https://dagit-core-storybook-dkp1njgus-elementl.vercel.app
https://al-06-26-add-indirect-execution-context-access.core-storybook.dagster-docs.io

Built with commit 9457af6.
This pull request is being automatically deployed with vercel-action

github-actions · 2023-07-10T20:35:02Z

Deploy preview for dagster ready!

Preview available at https://dagster-jhtjrovtv-elementl.vercel.app

Direct link to changed pages:

github-actions · 2023-08-18T14:50:37Z

Deploy preview for dagster-docs ready!

Preview available at https://dagster-docs-d2zx3qsfb-elementl.vercel.app
https://al-06-26-add-indirect-execution-context-access.dagster.dagster-docs.io

Direct link to changed pages:

schrockn · 2023-10-24T14:12:40Z

If we do this, I think it should be a getter of the class (e.g. AssetExecutionContext.get()) so that the type of the context is clear.

alangenfeld · 2023-10-24T15:58:15Z

python_modules/dagster/dagster/_core/execution/context/compute.py

@@ -1300,15 +1305,24 @@ def typed_event_stream_error_message(self) -> Optional[str]:
    def set_requires_typed_event_stream(self, *, error_message: Optional[str] = None) -> None:
        self._step_execution_context.set_requires_typed_event_stream(error_message=error_message)

+    @staticmethod
+    def get_current() -> Optional["OpExecutionContext"]:


should these raise or return None when there is no current context?

IMO we should raise

alangenfeld · 2023-10-24T15:59:02Z

python_modules/dagster/dagster/_core/execution/plan/execute_step.py

+    with time_execution_scope() as timer_result, enter_execution_context(
+        step_context
+    ) as compute_context:


this is hoisted up here to the top of the iterator call stack since context managers + active generators get goofy and if you raise an exception based on a yielded value in a frame above where the context manager is opened it does not get closed

schrockn · 2023-10-24T18:03:53Z

Interested @slopp 's feedback on naming, etc

benpankow · 2023-10-24T20:51:13Z

python_modules/dagster/dagster_tests/core_tests/execution_tests/test_context.py

@@ -405,3 +405,18 @@ def the_op(context: int):
        @asset
        def the_asset(context: int):
            pass
+
+
+def test_get_context():


might be worth testing thread behavior? I guess ContextVar should handle it fine

ya feel pretty good about thread & coroutine concurrency
https://docs.python.org/3/library/contextvars.html

schrockn

Req'ing change based on the in-memory instantiations issue (comment inline)

schrockn · 2023-10-25T09:56:25Z

python_modules/dagster/dagster/_core/execution/context/compute.py

+    asset_ctx = AssetExecutionContext(step_context)
+    op_ctx = OpExecutionContext(step_context)


We are going to switch this from an IS-A to a HAS-A relationship soon.

@alangenfeld It's subtle, but I think we should add a method on AssetExecutionContext now called get_op_execution_context and just that for the this instance. I'm a bit spooked by having two objects in memory for this. cc: @jamiedemaria

The thing that led me this current approach was ensuring that the passed down context was == to the result of _ExecutionContext.get(). It will requires some shenaigans to maintain that equality if we do the proposed shifting.

I'm not sure I'm quite following the issue/proposed solution here. can you add some elaboration?

slopp · 2023-10-25T18:48:34Z

I'm not sure I understand the use case enough to have an opinion on the name

I think the only thing that is potentially confusing as a user is whether this method has to be called... eg

@asset
def my_asset(context: AssetExecutionContext): 
     context.get_current() # what? is this necessary; why or why not?

Does re-using the standard Python get make this more or less clear? 🤷 Context injection is just fairly dagster magical.

Otherwise fine with get_current or get_context or even the fairly verbose get_current_context

comments addressed

schrockn · 2023-10-26T13:59:51Z

I'm not sure I understand the use case enough to have an opinion on the name

I think the only thing that is potentially confusing as a user is whether this method has to be called... eg
@asset
def my_asset(context: AssetExecutionContext): 
     context.get_current() # what? is this necessary; why or why not?
Does re-using the standard Python get make this more or less clear? 🤷 Context injection is just fairly dagster magical.

Otherwise fine with get_current or get_context or even the fairly verbose get_current_context

@slopp the use case here is avoid the need to pass contexts around entirely and instead provide a global accessor.

You would instead be able to write code like the following:

@asset
def my_asset(): 
     context = AssetExecutionContext.get_current()

There are tradeoffs to this approach, but we have certainly gotten this feature request in the past.

schrockn · 2023-10-26T14:01:01Z

python_modules/dagster/dagster/_core/execution/context/compute.py

+    asset_ctx = AssetExecutionContext(step_context)
+    asset_token = _current_asset_execution_context.set(asset_ctx)
+
+    try:
+        if context_annotation is EmptyAnnotation:
+            # if no type hint has been given, default to:
+            # * AssetExecutionContext for sda steps not in graph-backed assets, and asset_checks
+            # * OpExecutionContext for non sda steps
+            # * OpExecutionContext for ops in graph-backed assets
+            if is_asset_check:
+                yield asset_ctx
+            elif is_op_in_graph_asset or not is_sda_step:
+                yield asset_ctx.get_op_execution_context()
+            else:
+                yield asset_ctx
+        elif context_annotation is AssetExecutionContext:
+            yield asset_ctx
+        else:
+            yield asset_ctx.get_op_execution_context()
+    finally:
+        _current_asset_execution_context.reset(asset_token)


👍🏻 thanks much more comfortable with this

alangenfeld · 2023-10-26T21:25:15Z

python_modules/dagster/dagster/_core/execution/context/compute.py

@@ -1300,15 +1306,34 @@ def typed_event_stream_error_message(self) -> Optional[str]:
    def set_requires_typed_event_stream(self, *, error_message: Optional[str] = None) -> None:
        self._step_execution_context.set_requires_typed_event_stream(error_message=error_message)

+    @staticmethod
+    def get() -> "OpExecutionContext":


i updated these to just get now that they raise + slopps feedback

benpankow · 2023-11-06T18:07:50Z

looks good on my end, will leave to others w/ feedback to approve

slopp · 2023-11-06T20:06:22Z

I'm not sure I understand the use case enough to have an opinion on the name
I think the only thing that is potentially confusing as a user is whether this method has to be called... eg
@asset
def my_asset(context: AssetExecutionContext): 
     context.get_current() # what? is this necessary; why or why not?
Does re-using the standard Python get make this more or less clear? 🤷 Context injection is just fairly dagster magical.
Otherwise fine with get_current or get_context or even the fairly verbose get_current_context
@slopp the use case here is avoid the need to pass contexts around entirely and instead provide a global accessor.

You would instead be able to write code like the following:
@asset
def my_asset(): 
     context = AssetExecutionContext.get_current()
There are tradeoffs to this approach, but we have certainly gotten this feature request in the past.

Thanks for clarifying. Yes this new capability makes sense to me. I've seen a few users get very tripped up by the use of context (eg trying to put it as the second argument or trying to use a different name). I think this addition is beneficial. Would we go so far as to say this is the best practice?

Understanding the use case, I still don't have super strong opinions on the names.

schrockn · 2023-11-06T21:51:03Z

Also would like @PedramNavid and @tacastillo's thoughts on this. It solves a problem that some users have but at the cost of yet-another-way of doing things.

alangenfeld · 2023-11-06T22:37:23Z

From what I understand @benpankow has an internal use case for this functionality (being able to swap in a insights compatible resource without changing the API).

In the current form of the PR the new static get is not yet marked @public.

If we feel like we are close to resolution I am fine getting the buy in to mark this @public, but otherwise we may want to separate out the end user facing exposure of this functionality to prevent blocking progress of work.

tacastillo · 2023-11-06T23:13:46Z

Oh wow this is beautiful.

I'm a much bigger fan of pulling the contexts into the scope. get_dagster_logger afaik was the only thing that also used that pattern, but there's not much reason why not to do this with the whole context as it'd solve the same issues that get_dagster_logger did.

I used to like the magic arguments, but it's often a stumbling point for users. And I want to say we might've had an early churn because they've said something along the lines of:

Dagster's too complicated. Look at all these magic arguments in order to make something work.

There are handfuls of users that need this. I often get asked:

Hey, I have this huge legacy script with a bunch of nested functions, and I wanted to access the partition key/run ID, etc. from within it. Do I really need to pass the context in to each nested function?

Once this isn't experimental, my vote would be to appoint this as the best practice and to move away from magic arguments unless needed (sounds like there's a use case for keeping them?) I'd say the benefits of it outweigh the overall cost of the addition of it.

PedramNavid · 2023-11-07T01:08:17Z

Generally supportive! I assume same experience with regards to type-ahead/completions? So long as we maintain support for existing context magic for a while I think it's fine.

How would this work with functions called from an asset? You would still need to pass context right?

Or would this work?

def my_wrapper_func(x): 
    if AssetExecutionContext.some_method():
        do_something_with(x)

@asset
def my_asset():
    x = 123
    my_wrapper_func(x)

schrockn · 2023-11-07T13:39:51Z

@PedramNavid @tacastillo The tradeoff I'd like you to consider is 1) should we allow both methods and accept that there are "two ways of doing things" 2) should we drive heavy towards to global accessor and incur switching costs and 3) what is the price we pay for having code depend on global state.

@tacastillo in terms of "my vote would be to appoint this as the best practice and to move away from magic arguments unless needed (sounds like there's a use case for keeping them?)" consider the cases where we have resource initialization or sensors or anywhere else where there is a subtly different context. Should we have N global accessors and error when users call the wrong one? Or consolidate into a single ExecutionContext that covers all cases?

That all being said as @alangenfeld said this solves an immediate problem for Ben and does not make this public, so we can go ahead and land this. Seems like whether or not to make this public is a sep discussion.

github-actions · 2023-11-07T18:13:45Z

Deploy preview for dagit-storybook ready!

✅ Preview
https://dagit-storybook-cwm1gjw8t-elementl.vercel.app
https://al-06-26-add-indirect-execution-context-access.components-storybook.dagster-docs.io

Built with commit 9457af6.
This pull request is being automatically deployed with vercel-action

adds a means to get the current op/asset execution context via a function call, setup using a `ContextVar`. Exposed as `OpExecutionContext.get()` & `AssetExecutionContext.get()` ## How I Tested These Changes added test

alangenfeld mentioned this pull request Jun 26, 2023

[prototype] return AssetMaterialization #14931

Closed

alangenfeld force-pushed the al/06-26-add_indirect_execution_context_access branch from 37be28c to 8b69394 Compare July 10, 2023 20:31

alangenfeld force-pushed the al/06-26-add_indirect_execution_context_access branch from 8b69394 to d0752c8 Compare August 18, 2023 14:46

schrockn mentioned this pull request Sep 10, 2023

RFC: Much simpler AssetExecutionContext #16417

Closed

alangenfeld force-pushed the al/06-26-add_indirect_execution_context_access branch 2 times, most recently from 1ea5b5c to df619f3 Compare October 23, 2023 17:05

alangenfeld requested review from schrockn and benpankow October 24, 2023 14:08

alangenfeld force-pushed the al/06-26-add_indirect_execution_context_access branch from df619f3 to d13f3fc Compare October 24, 2023 14:11

alangenfeld force-pushed the al/06-26-add_indirect_execution_context_access branch 2 times, most recently from fad7e83 to b6ba500 Compare October 24, 2023 15:49

alangenfeld commented Oct 24, 2023

View reviewed changes

benpankow reviewed Oct 24, 2023

View reviewed changes

schrockn previously requested changes Oct 25, 2023

View reviewed changes

alangenfeld force-pushed the al/06-26-add_indirect_execution_context_access branch 2 times, most recently from 5f99d0b to 9889bd2 Compare October 25, 2023 20:02

schrockn reviewed Oct 26, 2023

View reviewed changes

alangenfeld commented Oct 26, 2023

View reviewed changes

alangenfeld requested a review from schrockn November 3, 2023 18:44

schrockn approved these changes Nov 7, 2023

View reviewed changes

add indirect execution context access

9457af6

alangenfeld force-pushed the al/06-26-add_indirect_execution_context_access branch from 9889bd2 to 9457af6 Compare November 7, 2023 18:10

alangenfeld merged commit f500ec9 into master Nov 7, 2023
3 checks passed

alangenfeld deleted the al/06-26-add_indirect_execution_context_access branch November 7, 2023 18:41

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

add indirect execution context access #14954

add indirect execution context access #14954

alangenfeld commented Jun 26, 2023 •

edited

Loading

alangenfeld commented Jun 26, 2023 •

edited

Loading

github-actions bot commented Jul 10, 2023 •

edited

Loading

github-actions bot commented Jul 10, 2023

github-actions bot commented Aug 18, 2023 •

edited

Loading

schrockn commented Oct 24, 2023

alangenfeld Oct 24, 2023

schrockn Oct 24, 2023

alangenfeld Oct 24, 2023

schrockn commented Oct 24, 2023

benpankow Oct 24, 2023

alangenfeld Oct 24, 2023

schrockn left a comment

schrockn Oct 25, 2023

alangenfeld Oct 25, 2023

jamiedemaria Oct 25, 2023

slopp commented Oct 25, 2023 •

edited

Loading

schrockn commented Oct 26, 2023

schrockn Oct 26, 2023

alangenfeld Oct 26, 2023

benpankow commented Nov 6, 2023

slopp commented Nov 6, 2023

schrockn commented Nov 6, 2023

alangenfeld commented Nov 6, 2023

tacastillo commented Nov 6, 2023 •

edited

Loading

PedramNavid commented Nov 7, 2023

schrockn commented Nov 7, 2023

github-actions bot commented Nov 7, 2023

		asset_ctx = AssetExecutionContext(step_context)
		op_ctx = OpExecutionContext(step_context)

add indirect execution context access #14954

add indirect execution context access #14954

Conversation

alangenfeld commented Jun 26, 2023 • edited Loading

How I Tested These Changes

alangenfeld commented Jun 26, 2023 • edited Loading

github-actions bot commented Jul 10, 2023 • edited Loading

github-actions bot commented Jul 10, 2023

github-actions bot commented Aug 18, 2023 • edited Loading

schrockn commented Oct 24, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

schrockn commented Oct 24, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

schrockn 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

slopp commented Oct 25, 2023 • edited Loading

schrockn commented Oct 26, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

benpankow commented Nov 6, 2023

slopp commented Nov 6, 2023

schrockn commented Nov 6, 2023

alangenfeld commented Nov 6, 2023

tacastillo commented Nov 6, 2023 • edited Loading

PedramNavid commented Nov 7, 2023

schrockn commented Nov 7, 2023

github-actions bot commented Nov 7, 2023

alangenfeld commented Jun 26, 2023 •

edited

Loading

alangenfeld commented Jun 26, 2023 •

edited

Loading

github-actions bot commented Jul 10, 2023 •

edited

Loading

github-actions bot commented Aug 18, 2023 •

edited

Loading

slopp commented Oct 25, 2023 •

edited

Loading

tacastillo commented Nov 6, 2023 •

edited

Loading