Make `airflow dags test` be able to execute Mapped Tasks #21210

ashb · 2022-01-29T16:04:03Z

In order to do this there were a few required:

The BackfillJob needs to know about mapped tasks, both to expand them, and in order to update it's TI tracking
The DebugExecutor, task_command.py, LocalTaskJob needed to "unmap" the mapped task to get the real operator back at the right time.

I was testing this with the following dag:

from airflow import DAG
from airflow.decorators import task
from airflow.operators.python import PythonOperator
import pendulum

@task
def make_list():
    return list(map(lambda a: f'echo "{a!r}"', [1, 2, {'a': 'b'}]))

def consumer(*args):
     print(repr(args))

with DAG(dag_id='maptest', start_date=pendulum.DateTime(2022, 1, 18)) as dag:
    PythonOperator(task_id='consumer', python_callable=consumer).map(op_args=make_list())

It can't "unmap" decorated operators successfully yet, so we're using old-school PythonOperator

We also just pass the whole value to the operator, not just the current mapping value(s)

^ 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.

ashb · 2022-01-29T16:04:28Z

No tests yet -- that's what I'm working on now.

uranusjr · 2022-01-31T16:44:04Z

airflow/executors/debug_executor.py

I wonder if it’d be nicer to also implement unmap() on BaseOperator and have it simply return self.

I made it raise an exception (as it should never be called?) but exist for typing.

I could have it return self and then remove the if task.is_mapped -- WDYT?

My OOP training tells me polymorphism is better than an if clause, so my vote goes to return self. But ultimately I guess either makes sense (depending on how your mental model is of the logic).

airflow/jobs/backfill_job.py

airflow/models/baseoperator.py

airflow/serialization/serialized_objects.py

uranusjr · 2022-02-03T10:02:48Z

airflow/jobs/backfill_job.py

Is this guaranteed (and why?) or just a temporary measure to get pass Mypy?

Temporary for now until we support mapping task groups (the DAG api supports it, but the runtime doesn't yet)

The why: Because expand_mapped_task doesn't exist on BaseOperator, only MappedOperator

Makes sense

github-actions · 2022-02-03T10:17:09Z

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

@task

In order to do this there were two steps required: - The BackfillJob needs to know about mapped tasks, both to expand them, and in order to update it's TI tracking - The DebugExecutor needed to "unmap" the mapped task to get the real operator bacj I was testing this with the following dag: ``` from airflow import DAG from airflow.decorators import task from airflow.operators.python import PythonOperator import pendulum @task def make_list(): return list(map(lambda a: f'echo "{a!r}"', [1, 2, {'a': 'b'}])) def consumer(*args): print(repr(args)) with DAG(dag_id='maptest', start_date=pendulum.DateTime(2022, 1, 18)) as dag: PythonOperator(task_id='consumer', python_callable=consumer).map(op_args=make_list()) ``` It can't "unmap" decorated operators successfully yet, so we're using old-school PythonOperator We also just pass the whole value to the operator, not just the current mapping value(s)

And since TaskGroup is a DAGNode, we don't need to store parent group directly anymore -- it'll already be stored

Co-authored-by: Tzu-ping Chung <[email protected]>

ashb · 2022-02-03T11:36:23Z

Not to self: if you want CI to run a test dag file, it helps if you commit it too.

ashb · 2022-02-03T12:43:29Z

A lot of SSH timeouts, and the MSSQL failures.

uranusjr · 2022-02-03T13:47:30Z

Yeah those SSH things are popping up a lot recently, not sure what had changed 😞

potiuk · 2022-02-03T14:01:38Z

Yeah those SSH things are popping up a lot recently, not sure what had changed disappointed

Yeah. I tried to address it via #21262 and it seemed to help a bit but I believe this is a side effect of some unrelated test, which we should track down and fix.

ashb · 2022-02-03T18:51:40Z

Maybe have had some un-expected luck via #21292 (that, or what ever environmental issue was causing it has vanished)

tests/models/__init__.py

Co-authored-by: Tzu-ping Chung <[email protected]>

ashb · 2022-02-04T14:17:17Z

Right, only SSH failures (which I have a PR open for) so I'm going to merge this now.

boring-cyborg bot added the area:Scheduler including HA (high availability) scheduler label Jan 29, 2022

ashb added the area:dynamic-task-mapping AIP-42 label Jan 29, 2022

uranusjr reviewed Jan 31, 2022

View reviewed changes

ashb force-pushed the run-mapped-tasks-debug-executor branch from 41d5177 to 083a26c Compare February 2, 2022 13:14

ashb marked this pull request as ready for review February 2, 2022 13:22

ashb requested review from XD-DENG, dstandish, jedcunningham and kaxil as code owners February 2, 2022 13:22

uranusjr reviewed Feb 3, 2022

View reviewed changes

airflow/serialization/serialized_objects.py Outdated Show resolved Hide resolved

uranusjr reviewed Feb 3, 2022

View reviewed changes

uranusjr approved these changes Feb 3, 2022

View reviewed changes

github-actions bot added the full tests needed We need to run full set of tests for this PR to merge label Feb 3, 2022

ashb and others added 7 commits February 3, 2022 11:35

Always have a task_group property on DAGNodes

05de448

And since TaskGroup is a DAGNode, we don't need to store parent group directly anymore -- it'll already be stored

Add "integation" tests for running mapped tasks via BackfillJob

d6d62bf

Fix tests

b3278cf

Update airflow/serialization/serialized_objects.py

f0965aa

Co-authored-by: Tzu-ping Chung <[email protected]>

Only show "Map Index" in Backfill report when relevant

633b746

If you want to test a DAG, it helps if you commit it

c915beb

ashb force-pushed the run-mapped-tasks-debug-executor branch from d6b04d7 to c915beb Compare February 3, 2022 11:35

uranusjr reviewed Feb 4, 2022

View reviewed changes

tests/models/__init__.py Outdated Show resolved Hide resolved

Update tests/models/__init__.py

d3f1601

Co-authored-by: Tzu-ping Chung <[email protected]>

ashb merged commit 6fc6edf into apache:main Feb 4, 2022

ashb deleted the run-mapped-tasks-debug-executor branch February 4, 2022 14:24

uranusjr mentioned this pull request Feb 4, 2022

Rewrite decorated task mapping #21328

Merged

ashb mentioned this pull request Feb 15, 2022

Fix Tasks getting stuck in scheduled state #19747

Merged

jedcunningham added the changelog:skip Changes that should be skipped from the changelog (CI, tests, etc..) label Mar 1, 2022

jedcunningham added this to the Airflow 2.3.0 milestone Apr 26, 2022

Make airflow dags test be able to execute Mapped Tasks #21210

Make airflow dags test be able to execute Mapped Tasks #21210

Uh oh!

Conversation

ashb commented Jan 29, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ashb commented Jan 29, 2022

Uh oh!

uranusjr Jan 31, 2022

Choose a reason for hiding this comment

Uh oh!

ashb Feb 2, 2022

Choose a reason for hiding this comment

Uh oh!

uranusjr Feb 2, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

uranusjr Feb 3, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ashb Feb 3, 2022

Choose a reason for hiding this comment

Uh oh!

ashb Feb 3, 2022

Choose a reason for hiding this comment

Uh oh!

uranusjr Feb 3, 2022

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Feb 3, 2022

Uh oh!

ashb commented Feb 3, 2022

Uh oh!

ashb commented Feb 3, 2022

Uh oh!

uranusjr commented Feb 3, 2022

Uh oh!

potiuk commented Feb 3, 2022

Uh oh!

ashb commented Feb 3, 2022

Uh oh!

Uh oh!

ashb commented Feb 4, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Make `airflow dags test` be able to execute Mapped Tasks #21210

Make `airflow dags test` be able to execute Mapped Tasks #21210

ashb commented Jan 29, 2022 •

edited

Loading

uranusjr Feb 2, 2022 •

edited

Loading

uranusjr Feb 3, 2022 •

edited

Loading