dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4250259289
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -2299,10 +2903,45 @@ def _execute_task(context: Context, ti:
RuntimeTaskInstance, log: Logger):
assert isinstance(kwargs, dict)
execute = functools.partial(task.resume_execution,
next_method=next_method, next_kwargs=kwargs)
- # Export context in os.environ to make it available for operators to use.
- airflow_context_vars = context_to_airflow_vars(context,
in_env_var_format=True)
- os.environ.update(airflow_context_vars)
+ # Export the context to os.environ for operators that read AIRFLOW_CTX_*
directly. Indexed
+ # sub-tasks skip this: they run concurrently in one process, so the update
would race, and the
+ # parent IterableOperator already exported the same values before they
started.
+ if not isinstance(ti, IndexedTaskInstance):
+ os.environ.update(context_to_airflow_vars(context,
in_env_var_format=True))
+
+ outlet_events = _run_pre_execute(task, context, log)
+ log.info("::endgroup::")
+
+ # An indexed sub-task runs in a worker thread under the parent's
execution_timeout, which the
+ # parent enforces and reported to the supervisor once (see
_run_execute_callable).
+ result = _run_execute_callable(
+ context, execute, task, enforce_timeout=not isinstance(ti,
IndexedTaskInstance)
+ )
+
+ _run_post_execute(task, context, outlet_events, result, log)
+ return result
+
+
+async def _execute_async_task(context: Context, ti: RuntimeTaskInstance, log:
Logger):
+ """Async counterpart of :func:`_execute_task` for
:class:`BaseAsyncOperator` sub-tasks."""
+ task = cast("BaseAsyncOperator", ti.task)
+ # Async tasks cannot be resuming a deferral, so next_method never applies
here.
+ outlet_events = _run_pre_execute(task, context, log)
+
+ ctx = contextvars.copy_context()
+ ctx.run(ExecutorSafeguard.tracker.set, task)
+ # Under the parent's execution_timeout only, as a sync indexed task is:
the operator's own limit
+ # is the same value, started later, so enforcing it here too would never
fire first and would
+ # race the parent's kill when the two coincide.
+ result = await ctx.run(lambda: task.aexecute(context=context))
Review Comment:
Right, `ctx.run()` only created the coroutine in `ctx`; its body ran in the
awaiting task's context, without the tracker. It now runs as a task created
with `context=ctx` (7bff796c30), which matches how a plain async task gets
there: from inside `ctx.run(execute)`, whose context the loop copies into the
task. A new test runs a sync and an async operator under `.iterate()` that each
check `ExecutorSafeguard.tracker` is set to themselves; the async case got
`[False, False]` before the change. Cancellation and exceptions,
`AirflowTaskTimeout` included, still reach the awaiting item unchanged, and a
contextvar the operator sets stays in `ctx` as on the sync path.
Drafted-by: Claude Opus 5.5; reviewed by @dabla before posting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
