Fix mapped tasks partial arguments when DAG default args are provided by hussein-awala · Pull Request #29913 · apache/airflow

hussein-awala · 2023-03-04T01:15:58Z

Currently the BaseOperator overrides the mapped tasks partial arguments by DAG default args when they are provided regardless the value of partial arguments, but by definition, Airflow should uses these default args only when the operator argument is not provided.

This PR fills the partial_kwargs dictionary with the arguments and kwargs of the partial method, and then just replaces the None values with the default arguments values.

…s by dag default values

potiuk · 2023-03-04T23:59:44Z

cc: @uranusjr . This definitely needs your insight :)

airflow/models/baseoperator.py

… comments

…pport None as accepted value

uranusjr · 2023-03-13T05:20:11Z

May be related: #29366

airflow/models/baseoperator.py

uranusjr

I added a commit to tweak the merging part. Should not affect the functionality, I think, but please take another look to make sure.

airflow/models/baseoperator.py

This should improve iteration a bit, I think.

…#29913) * Add a failing test to make it pass * use partial_kwargs when they are provide and override only None values by dag default values * update the test and check if the values are filled in the right order * fix overriding retry_delay with default value when it is equal to 0 * add missing default value for inlets and outlets * set partial_kwargs dict type to dict[str, Any] and remove type ignore comments * create a dict for default values and use NotSet instead of None to support None as accepted value * update partial typing by removing None type from some args and set NotSet for all args * Tweak kwarg merging slightly This should improve iteration a bit, I think. * Fix unit tests --------- Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com> (cherry picked from commit f01051a)

…apache#29913) * Add a failing test to make it pass * use partial_kwargs when they are provide and override only None values by dag default values * update the test and check if the values are filled in the right order * fix overriding retry_delay with default value when it is equal to 0 * add missing default value for inlets and outlets * set partial_kwargs dict type to dict[str, Any] and remove type ignore comments * create a dict for default values and use NotSet instead of None to support None as accepted value * update partial typing by removing None type from some args and set NotSet for all args * Tweak kwarg merging slightly This should improve iteration a bit, I think. * Fix unit tests --------- Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com>

mpgreg · 2023-08-03T13:35:23Z

Is there any reason this fix wouldn't work with branch decorator also?

I'm on Astronomer Runtime 8.8.0 based on Airflow 2.6.3+astro.2 and branch to dynamic tasks

@task.branch
    def check_object_count(doc_count:dict) -> str:
        response = blah
        
        if response 
            return None
        else:
            return ["get_md_docs", "get_rst_docs", "get_code_samples", "get_stack_overflow"]

@task(trigger_rule='none_failed')
    def get_md_docs(source:dict):
         return blah
@task(trigger_rule='none_failed')
    def get_rst_docs(source:dict):
         return blah
@task(trigger_rule='none_failed')
    def get_code_samples(source:dict):
         return blah
@task(trigger_rule='none_failed')
    def get_stack_overflow(source:dict):
         return blah

md_docs = get_md_docs.partial().expand(source=markdown_docs_sources)
rst_docs = get_rst_docs.partial().expand(source=rst_docs_sources)
code_samples = get_code_samples.partial().expand(source=code_samples_sources)
stackoverflow_md = get_stack_overflow.partial().expand(tag=stackoverflow_tags)

All downstream mapped tasks fail with:

Traceback (most recent call last):
  File "/usr/local/lib/python3.10/site-packages/airflow/models/taskinstance.py", line 1407, in _run_raw_task
    self._execute_task_with_callbacks(context, test_mode)
  File "/usr/local/lib/python3.10/site-packages/airflow/models/taskinstance.py", line 1531, in _execute_task_with_callbacks
    task_orig = self.render_templates(context=context)
  File "/usr/local/lib/python3.10/site-packages/airflow/models/taskinstance.py", line 2179, in render_templates
    original_task.render_template_fields(context)
  File "/usr/local/lib/python3.10/site-packages/airflow/models/mappedoperator.py", line 692, in render_template_fields
    unmapped_task = self.unmap(mapped_kwargs)
  File "/usr/local/lib/python3.10/site-packages/airflow/models/mappedoperator.py", line 615, in unmap
    op = self.operator_class(**kwargs, _airflow_from_mapped=True)
  File "/usr/local/lib/python3.10/site-packages/airflow/models/baseoperator.py", line 429, in apply_defaults
    result = func(self, **kwargs, default_args=default_args)
  File "/usr/local/lib/python3.10/site-packages/airflow/decorators/python.py", line 49, in __init__
    super().__init__(
  File "/usr/local/lib/python3.10/site-packages/airflow/models/baseoperator.py", line 429, in apply_defaults
    result = func(self, **kwargs, default_args=default_args)
  File "/usr/local/lib/python3.10/site-packages/airflow/decorators/base.py", line 212, in __init__
    super().__init__(task_id=task_id, **kwargs_to_upstream, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/airflow/models/baseoperator.py", line 429, in apply_defaults
    result = func(self, **kwargs, default_args=default_args)
  File "/usr/local/lib/python3.10/site-packages/airflow/operators/python.py", line 166, in __init__
    super().__init__(**kwargs)
  File "/usr/local/lib/python3.10/site-packages/airflow/models/baseoperator.py", line 429, in apply_defaults
    result = func(self, **kwargs, default_args=default_args)
  File "/usr/local/lib/python3.10/site-packages/airflow/models/baseoperator.py", line 788, in __init__
    raise AirflowException(
airflow.exceptions.AirflowException: Invalid arguments were passed to _PythonDecoratedOperator (task_id: get_md_docs__1). Invalid arguments were:
**kwargs: {'postgres_conn_id': 'postgres_myid'}

mpgreg · 2023-08-03T13:36:34Z

If I remove the default_args it runs okay.

hussein-awala added 3 commits March 4, 2023 00:50

Add a failing test to make it pass

acf78e0

use partial_kwargs when they are provide and override only None value…

934d128

…s by dag default values

update the test and check if the values are filled in the right order

389c6a9

hussein-awala requested a review from uranusjr as a code owner March 4, 2023 01:15

hussein-awala added 2 commits March 4, 2023 03:06

fix overriding retry_delay with default value when it is equal to 0

7725d63

add missing default value for inlets and outlets

5e3b5ff

hussein-awala added this to the Airflow 2.5.2 milestone Mar 4, 2023

uranusjr reviewed Mar 10, 2023

View reviewed changes

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

uranusjr reviewed Mar 10, 2023

View reviewed changes

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

ephraimbuddy modified the milestones: Airflow 2.5.2, Airflow 2.5.3 Mar 10, 2023

hussein-awala added 2 commits March 11, 2023 02:44

set partial_kwargs dict type to dict[str, Any] and remove type ignore…

ad1f7b8

… comments

create a dict for default values and use NotSet instead of None to su…

cb0980a

…pport None as accepted value

uranusjr mentioned this pull request Mar 13, 2023

default_args feature incompatible with Dynamic Task Mapping #29366

Closed

2 tasks

uranusjr reviewed Mar 13, 2023

View reviewed changes

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

vchiapaikeo mentioned this pull request Mar 20, 2023

Support default_args in dynamic task mapping #30056

Closed

pierrejeambrun modified the milestones: Airflow 2.5.3, Airflow 2.5.4 Mar 27, 2023

hussein-awala requested review from ashb, dstandish, eladkal, ephraimbuddy, jedcunningham, jhtimmins, kaxil, o-nikolas, potiuk and ryanahamilton as code owners March 31, 2023 21:11

eladkal removed request for XD-DENG, bbovenzi, dstandish, eladkal, jedcunningham, jhtimmins, kaxil and o-nikolas April 4, 2023 21:50

uranusjr approved these changes Apr 10, 2023

View reviewed changes

hussein-awala commented Apr 10, 2023

View reviewed changes

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

Tweak kwarg merging slightly

cd142b6

This should improve iteration a bit, I think.

uranusjr force-pushed the fix/mapped_tasks_args branch from 819b262 to cd142b6 Compare April 13, 2023 06:46

ephraimbuddy modified the milestones: Airflow 2.5.4, Airflow 2.6.1 Apr 13, 2023

hussein-awala added 2 commits April 14, 2023 02:46

Fix unit tests

0892f25

Merge branch 'main' into fix/mapped_tasks_args

7a59ec7

hussein-awala force-pushed the fix/mapped_tasks_args branch from bb92fe4 to 7a59ec7 Compare April 14, 2023 11:00

potiuk approved these changes Apr 14, 2023

View reviewed changes

potiuk modified the milestones: Airflow 2.6.1, Airflow 2.6.0 Apr 14, 2023

hussein-awala merged commit f01051a into apache:main Apr 14, 2023

ephraimbuddy added the type:bug-fix Changelog: Bug Fixes label Apr 14, 2023

agomez-etsy mentioned this pull request Aug 21, 2023

Dynamic tasks fail when default_args contains custom fields #33600

Closed

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Comments

Fix mapped tasks partial arguments when DAG default args are provided#29913

Fix mapped tasks partial arguments when DAG default args are provided#29913
hussein-awala merged 13 commits intoapache:mainfrom
hussein-awala:fix/mapped_tasks_args

hussein-awala commented Mar 4, 2023

Uh oh!

potiuk commented Mar 4, 2023

Uh oh!

Uh oh!

Uh oh!

uranusjr commented Mar 13, 2023

Uh oh!

Uh oh!

uranusjr left a comment

Uh oh!

Uh oh!

mpgreg commented Aug 3, 2023

Uh oh!

mpgreg commented Aug 3, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Comments

Conversation

hussein-awala commented Mar 4, 2023

Uh oh!

potiuk commented Mar 4, 2023

Uh oh!

Uh oh!

Uh oh!

uranusjr commented Mar 13, 2023

Uh oh!

Uh oh!

uranusjr left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mpgreg commented Aug 3, 2023

Uh oh!

mpgreg commented Aug 3, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants