fweilun commented on issue #52219:
URL: https://github.com/apache/airflow/issues/52219#issuecomment-3018696381
Hi, I'm currently working on this issue.
`
def skip_all_except(
self,
ti: RuntimeTaskInstanceProtocol,
branch_task_ids: None | str | Iterable[str],
):
// some code in skip_all_except
task = ti.task
dag = ti.task.dag
valid_task_ids = set(dag.task_ids)
invalid_task_ids = branch_task_id_set - valid_task_ids
if invalid_task_ids:
raise AirflowException(
"'branch_task_ids' must contain only valid task_ids. "
f"Invalid tasks found: {invalid_task_ids}."
)
log.info(f"[DEBUG] {task.downstream_list}")
downstream_tasks = _ensure_tasks(task.downstream_list)
if downstream_tasks:
# For a branching workflow that looks like this, when "branch"
does skip_all_except("task1"),
# we intuitively expect both "task1" and "join" to execute even
though strictly speaking,
# "join" is also immediately downstream of "branch" and should
have been skipped. Therefore,
# we need a special case here for such empty branches: Check
downstream tasks of branch_task_ids.
# In case the task to skip is also downstream of
branch_task_ids, we add it to branch_task_ids and
# exclude it from skipping.
#
# branch -----> join
# \ ^
# v /
# task1
#
for branch_task_id in list(branch_task_id_set):
branch_task_id_set.update(dag.get_task(branch_task_id).get_flat_relative_ids(upstream=False))
log.info(f"[DEBUG] downstream_tasks: {downstream_tasks}")
log.info(f"[DEBUG] branch_task_id_set: {branch_task_id_set}")
skip_tasks = [
(t.task_id, ti.map_index) for t in downstream_tasks if
t.task_id not in branch_task_id_set
]
follow_task_ids = [t.task_id for t in downstream_tasks if
t.task_id in branch_task_id_set]
log.info("Skipping tasks %s", skip_tasks)
ti.xcom_push(
key=XCOM_SKIPMIXIN_KEY,
value={XCOM_SKIPMIXIN_FOLLOWED: follow_task_ids},
)
# The following could be applied only for non-mapped tasks,
# as future mapped tasks have not been expanded yet. Such tasks
# have to be handled by NotPreviouslySkippedDep.
self._set_state_to_skipped(skip_tasks, ti.map_index) # type:
ignore[arg-type]
`
--
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]
