jason810496 commented on code in PR #75405:
URL: https://github.com/apache/airflow/pull/75405#discussion_r4355805363
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -2355,52 +2356,102 @@ def _run_execute_callable(
return result
-@detail_span("_execute_task")
-def _execute_task(context: Context, ti: RuntimeTaskInstance, log: Logger):
- """Execute Task (optionally with a Timeout) and push Xcom results."""
- task = ti.task
- execute = task.execute
+class _TaskRunner:
+ """
+ Run a task instance and return its result.
- if ti._ti_context_from_server and (next_method :=
ti._ti_context_from_server.next_method):
- from airflow.sdk.serde import deserialize
+ Only :meth:`run` is meant to be called from outside. It selects how the
task
+ runs and delegates. A stub task goes to its registered handler; a concrete
+ operator to its ``execute`` lifecycle (:meth:`_run_operator`).
+
+ .. note::
+
+ Ideally, the stub-or-not dispatching should live with the coordinator
+ dispatch. A coordinator is picked by queue, then either runs the
handler
+ (if stub) or the (non-stub) task. This requires a larger refactor and
is
+ deferred to a future PR.
+
+ .. seealso:: Language SDK ADR-0002.
+ """
+
+ @classmethod
+ @detail_span("_execute_task")
+ def run(cls, ti: RuntimeTaskInstance, context: Context):
+ """Run the task specified by *ti* with *context*."""
+ if ti.task.is_stub:
+ return cls._run_handler(ti)
+ return cls._run_operator(ti, context)
+
+ @classmethod
+ def _run_operator(cls, ti: RuntimeTaskInstance, context: Context):
+ """Execute a concrete operator."""
+ task = ti.task
+ execute = task.execute
+
+ if ti._ti_context_from_server and (next_method :=
ti._ti_context_from_server.next_method):
+ from airflow.sdk.serde import deserialize
+
+ next_kwargs_data = ti._ti_context_from_server.next_kwargs or {}
+ try:
+ if TYPE_CHECKING:
+ assert isinstance(next_kwargs_data, dict)
+ kwargs = deserialize(next_kwargs_data)
+ except (ImportError, KeyError, AttributeError, TypeError):
+ from airflow.serialization.serialized_objects import
BaseSerialization
+
+ kwargs = BaseSerialization.deserialize(next_kwargs_data)
- next_kwargs_data = ti._ti_context_from_server.next_kwargs or {}
- try:
if TYPE_CHECKING:
- assert isinstance(next_kwargs_data, dict)
- kwargs = deserialize(next_kwargs_data)
- except (ImportError, KeyError, AttributeError, TypeError):
- from airflow.serialization.serialized_objects import
BaseSerialization
+ assert isinstance(kwargs, dict)
+ execute = functools.partial(task.resume_execution,
next_method=next_method, next_kwargs=kwargs)
- kwargs = BaseSerialization.deserialize(next_kwargs_data)
+ # 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)
- if TYPE_CHECKING:
- assert isinstance(kwargs, dict)
- execute = functools.partial(task.resume_execution,
next_method=next_method, next_kwargs=kwargs)
+ outlet_events = context_get_outlet_events(context)
- # 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)
+ if (pre_execute_hook := task._pre_execute_hook) is not None:
+ create_executable_runner(pre_execute_hook, outlet_events,
logger=log).run(context)
+ if getattr(pre_execute_hook := task.pre_execute, "__func__", None) is
not BaseOperator.pre_execute:
+ create_executable_runner(pre_execute_hook, outlet_events,
logger=log).run(context)
- outlet_events = context_get_outlet_events(context)
+ _run_task_state_change_callbacks(task, "on_execute_callback", context,
log)
- if (pre_execute_hook := task._pre_execute_hook) is not None:
- create_executable_runner(pre_execute_hook, outlet_events,
logger=log).run(context)
- if getattr(pre_execute_hook := task.pre_execute, "__func__", None) is not
BaseOperator.pre_execute:
- create_executable_runner(pre_execute_hook, outlet_events,
logger=log).run(context)
+ log.info("::endgroup::")
- _run_task_state_change_callbacks(task, "on_execute_callback", context, log)
+ result = _run_execute_callable(context, execute, task)
- log.info("::endgroup::")
+ if (post_execute_hook := task._post_execute_hook) is not None:
+ create_executable_runner(post_execute_hook, outlet_events,
logger=log).run(context, result)
+ if getattr(post_execute_hook := task.post_execute, "__func__", None)
is not BaseOperator.post_execute:
+ create_executable_runner(post_execute_hook, outlet_events,
logger=log).run(context)
- result = _run_execute_callable(context, execute, task)
+ return result
- if (post_execute_hook := task._post_execute_hook) is not None:
- create_executable_runner(post_execute_hook, outlet_events,
logger=log).run(context, result)
- if getattr(post_execute_hook := task.post_execute, "__func__", None) is
not BaseOperator.post_execute:
- create_executable_runner(post_execute_hook, outlet_events,
logger=log).run(context)
+ @classmethod
+ def _run_handler(cls, ti: RuntimeTaskInstance):
+ """Run a stub task's registered Python task handler."""
+ from airflow.sdk.execution_time.coordinator import
get_coordinator_manager
+ from airflow.sdk.task_handlers import call_handler
- return result
+ coordinator = get_coordinator_manager().for_queue(ti.task.queue)
+ declaration = coordinator.describe_task_handler(ti.dag_id, ti.task_id)
Review Comment:
Would it be better to make the `describe_task_handler` method
`PythonCoordinator` specific instead of changing the base interface of
coordinator? Or remove it entirely?
For the non-python SDK, we actually encapsulate the logic of handling
execution path for the task handler and for the Dag definition within the
`execute_task` method.
https://github.com/apache/airflow/blob/7ac92c98ce63201de7b4bf9a59f0364088e7c447/task-sdk/src/airflow/sdk/coordinators/_subprocess.py#L764-L792
IIUC, the necessary info for importing a handler are
`task_handler_bundle_path`, dag_id and task_id.
There're two other approaches in my mind:
1) define a public attribute for accessing `task_handler_bundle_path` --
this means the child need to reparse the config and init the coordinator again
just for resolving the `task_handler_bundle_path`, but change less.
2) passing the `task_handler_bundle_path` that we resolved at supervisor as
`_AIRFLOW__TASK_HANDLER_BUNDLE_PATH` env to the Python task runner.
--
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]
