New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Allow ExternalTaskSensor to wait for taskgroup #14640

Closed

xinbinhuang wants to merge 6 commits into apache:main from xinbinhuang:taskgroup-sensor

Contributor

xinbinhuang commented Mar 6, 2021 •

edited

Loading

closes: #14563

This PR enables ExternalTaskSensor to also wait for the external task_group.

The implementation is to retrieve the external DAG from the DagBag and then check if the TaskGroup exists. If so, query and wait for the states of all tasks within that TaskGroup during the poking cycle.

^ Add meaningful description above

Read the Pull Request Guidelines for more information.
In case of fundamental code change, Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in UPDATING.md.

xinbinhuang requested review from ashb, kaxil, turbaszek and XD-DENG as code owners

March 6, 2021 09:27

xinbinhuang marked this pull request as draft

March 6, 2021 09:29

xinbinhuang changed the title ~~Allow ExternalSensor to wait for taskgroup~~ Allow ExternalTaskSensor to wait for taskgroup

xinbinhuang commented

View reviewed changes

airflow/sensors/external_task.py Outdated Show resolved Hide resolved

airflow/sensors/external_task.py Outdated

    
                              .scalar()

                          )

                          ) / len(external_task_group_task_ids)

Contributor Author

xinbinhuang Mar 6, 2021 •

edited

Loading

This is to keep the poking check behavior return count_allowed == len(dttm_filter)

xinbinhuang commented

View reviewed changes

airflow/sensors/external_task.py

    
            @@ -134,20 +146,23 @@ def __init__(
          
                      self.execution_delta = execution_delta

                      self.execution_date_fn = execution_date_fn

                      self.external_dag_id = external_dag_id

                      self.external_task_group_id = external_task_group_id

                      self.external_task_id = external_task_id

                      self.check_existence = check_existence

Contributor Author

xinbinhuang Mar 6, 2021 •

edited

Loading

self.check_existence = check_existence is False by default, which maybe make sense for external_dag or external_task. But external_task_group has to check and get an existing dag in order to get the list of task_ids.

https://github.com/apache/airflow/blob/fce49402461ee4e7a5f6ffd18cee3121f3496a39/airflow/sensors/external_task.py#L174-L180

I wonder if we can change the default to True or even have check_existence enabled required? This can give more useful errors if the external task/dag does not exist as well as having a consistent behavior as external_task_group. Also, what would be use case to have a Sensor waiting for an object that doesn't exist until it times out?

github-actions bot commented Mar 6, 2021

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Backport packages$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

xinbinhuang force-pushed the taskgroup-sensor branch from d620a7a to e2dcab4 Compare

March 6, 2021 18:37

xinbinhuang commented

View reviewed changes

airflow/sensors/external_task.py Outdated Show resolved Hide resolved

kaxil reviewed

View reviewed changes

airflow/models/dag.py Outdated Show resolved Hide resolved

xinbinhuang force-pushed the taskgroup-sensor branch from e2dcab4 to 77c9eb6 Compare

March 30, 2021 02:37

xinbinhuang marked this pull request as ready for review

March 30, 2021 07:09

xinbinhuang commented

View reviewed changes

airflow/sensors/external_task.py Outdated

    
            @@ -164,18 +184,23 @@ def poke(self, context, session=None):
          
                      if self.failed_states:

                          count_failed = self.get_count(dttm_filter, session, self.failed_states)

                      if count_failed == len(dttm_filter):

                      if count_failed > 0:

Contributor Author

xinbinhuang Mar 30, 2021 •

edited

Loading

Here I am making the assumption that as long as there is at least one external task failure, then we will want to fail the sensor. Though this changes the original behavior, I think this will be a better.

Member

turbaszek Apr 6, 2021

Should we add this comment in code?

Contributor Author

xinbinhuang Apr 7, 2021 •

edited

Loading

I don't think it's necessary more than an entry in the UPDATING.md. I think the only situation where you will have multiple counts is when the execution_date_fn returns more than one execution date to wait for. However, the original behavior will get you into a weird state when only part of the TIs fail, i.e. one fail and one succeeds, resulting in time out. IMHO, I think this's more like a bug than intended behavior. WDYT?

xinbinhuang commented

View reviewed changes

airflow/sensors/external_task.py

Comment on lines +263 to +301

    
                  def get_external_task_group_task_ids(self, session):

                      """Return task ids for the external TaskGroup"""

                      refreshed_dag_info = DagBag(read_dags_from_db=True).get_dag(self.external_dag_id, session)

                      task_group: Optional["TaskGroup"] = refreshed_dag_info.task_group_dict.get(

                          self.external_task_group_id

                      )

                      if not task_group:

                          raise AirflowException(

                              f"The external task group {self.external_task_group_id} in "

                              f"DAG {self.external_dag_id} does not exist."

                          )

                      task_ids = [task.task_id for task in task_group]

                      return task_ids

Contributor Author

xinbinhuang Mar 30, 2021 •

edited

Loading

The main piece that you retrieve a list of tasks for a TaskGroup. I believe that read_dags_from_db=True is safe to use here because serialized dag is enabled by default in 2.0.

Contributor

yuqian90 Apr 16, 2021

The existing task execution code is creating DagBag on its own instead of reading serialized dags from db. For example this line is creating a DagBag. I think we should do the same here. It's important for tasks to get the latest view of the dag during execution.

https://github.com/apache/airflow/blob/f1edc220d3f9cb050016d23246a682276bd09eee/airflow/sensors/external_task.py#L213

github-actions bot commented Mar 30, 2021

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

xinbinhuang force-pushed the taskgroup-sensor branch from 9d27356 to 914bf4e Compare

March 31, 2021 01:12

github-actions bot commented Mar 31, 2021

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

github-actions bot commented Mar 31, 2021

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

github-actions bot commented Mar 31, 2021

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

github-actions bot commented Mar 31, 2021

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

xinbinhuang force-pushed the taskgroup-sensor branch 2 times, most recently from 4e72ea9 to fce4940 Compare

April 3, 2021 16:19

xinbinhuang requested a review from kaxil

April 3, 2021 16:25

Contributor Author

xinbinhuang commented Apr 3, 2021

Test fails on K8S image build job, and I think it's not relevant to this PR

Member

potiuk commented Apr 4, 2021

Just fixed the K8S problem in #15182 - can you please rebase.

xinbinhuang force-pushed the taskgroup-sensor branch from fce4940 to 1d0bf9e Compare

April 4, 2021 02:34

turbaszek reviewed

View reviewed changes

airflow/models/dag.py Outdated Show resolved Hide resolved

xinbinhuang force-pushed the taskgroup-sensor branch from 2e92f65 to 21a2d07 Compare

April 7, 2021 16:09

xinbinhuang removed the stale label

xinbinhuang reopened this

xinbinhuang and others added 5 commits

September 3, 2021 05:20


          Allow ExternalSensor to wait for taskgroup

aace321


          Use public method

6f74efb

Co-authored-by: Kaxil Naik <[email protected]>


          Remove unittest.TestCase

2bfe78e


          Fix test: clear dag_runs for test_external_dag_sensor

81c50e0


          Test external task group sensor

c3495aa

fixup! Test external task group sensor

fixup! fixup! Test external task group sensor

fixup! fixup! fixup! Test external task group sensor

fixup! fixup! fixup! fixup! Test external task group sensor

xinbinhuang force-pushed the taskgroup-sensor branch from f1edc22 to c3495aa Compare

September 3, 2021 05:41


          Update airflow/models/dag.py

c31e8c9

Co-authored-by: Tomek Urbaszek <[email protected]>

kaxil modified the milestones: Airflow 2.2, Airflow 2.3

Member

kaxil commented Sep 14, 2021

Can you fix the conflicts please @xinbinhuang

jedcunningham modified the milestones: Airflow 2.2.0, Airflow 2.3.0

github-actions bot commented Nov 28, 2021

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.

github-actions bot added the stale label

github-actions bot closed this

Contributor

eladkal commented Dec 19, 2021

@xinbinhuang will you have time to complete it?

Contributor Author

xinbinhuang commented Dec 19, 2021

@xinbinhuang will you have time to complete it?

Thanks for the nudge! Will try to wrap it up before the holidays hit.

Contributor

eladkal commented Dec 19, 2021 •

edited

Loading

Thanks for the nudge! Will try to wrap it up before the holidays hit.

Great :) re-opening so it won't be missed

eladkal reopened this

eladkal removed the stale label

github-actions bot commented Feb 3, 2022

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.

github-actions bot added the stale label

github-actions bot closed this

Contributor

russellpierce commented May 27, 2022

Darn, it's a shame this didn't get through - it is exactly what I was looking for!

Member

potiuk commented Jun 1, 2022

Feel free to open a PR and contribute it on your own. You willl just have to make sure to follow it up and implement it to the quality that we expect @russellpierce - Airflow is created by > 2000 contributors like you, so if you need something, implementing it yourself is the fastest way to get things done @russellpierce.

pateash mentioned this pull request

Allowing task_group Sensor in ExternalTaskSensor #24902

Merged

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

turbaszek turbaszek left review comments

yuqian90 yuqian90 requested changes

ashb Awaiting requested review from ashb ashb is a code owner

XD-DENG Awaiting requested review from XD-DENG XD-DENG is a code owner

kaxil Awaiting requested review from kaxil

Labels

area:core-operators stale