small

dagster-io · Sep 20, 2023 · 7fa66dd · 7fa66dd
1 parent 6dd716a
commit 7fa66dd
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/python_modules/dagster/dagster/_core/definitions/op_invocation.py b/python_modules/dagster/dagster/_core/definitions/op_invocation.py
@@ -18,6 +18,7 @@
     DagsterInvariantViolationError,
     DagsterTypeCheckDidNotPass,
 )
+from dagster._core.execution.context.invocation import UnboundAssetExecutionContext
 
 from .events import (
     AssetMaterialization,
@@ -147,7 +148,9 @@ def direct_invocation_result(
                 " no context was provided when invoking."
             )
         if len(args) > 0:
-            if args[0] is not None and not isinstance(args[0], UnboundOpExecutionContext):
+            if args[0] is not None and not isinstance(
+                args[0], (UnboundOpExecutionContext, UnboundAssetExecutionContext)
+            ):
                 raise DagsterInvalidInvocationError(
                     f"Decorated function '{compute_fn.name}' has context argument, "
                     "but no context was provided when invoking."

diff --git a/python_modules/dagster/dagster/_core/execution/context/invocation.py b/python_modules/dagster/dagster/_core/execution/context/invocation.py
@@ -811,10 +811,6 @@ def __init__(
             _assets_def=assets_def,
         )
 
-        self._step_execution_context = (
-            self._op_execution_context._step_execution_context  # noqa: SLF001
-        )
-
         self._cm_scope_entered = False
 
     def __enter__(self):