Implement enough interface for MappedOperator to be baggable #20945

uranusjr · 2022-01-19T09:58:37Z

The tests are changed to actually bag mapped tasks so we can test they are actually compatible.

ashb · 2022-01-19T10:02:50Z

airflow/models/base.py

Anyone else find it odd that Operator is "lower level" than BaseOperator 😁

ashb · 2022-01-19T10:03:46Z

airflow/models/baseoperator.py

Why change from a collection to a sequence? (This might cause knock-on impacts for typing of all other operators)

It’s the other way around, from Sequence to a Collection; Collection is looser than Sequence and should work in subclasses (I think; we can always change this back).

uranusjr · 2022-01-20T06:55:57Z

~~Hmm, #20931 still does not correctly establish task dependencies, we’re missing something.~~

Oh that’s a simple yet difficult-to-spot oversight. Everything should be working now!

uranusjr · 2022-01-20T07:39:51Z

As an unrelated discovery, we can sort of trace what we need to check/change to account for MappedOperator by adding annotation to DAG.add_task() and DAG.add_tasks(). It cascades into a nice chain of errors that alerts us what originally can only contain BaseOperator now may also contain MappedOperator and may need to be modified.

airflow/models/base.py

airflow/models/baseoperator.py

airflow/serialization/serialized_objects.py

ashb · 2022-01-21T15:37:37Z

airflow/models/baseoperator.py

We could choose to make them writable properties that affect the slot in partial_kwargs couldn't we?

@weight_rule.setter def _set_weigh_rule(self, value): self.partial_kwargs['weight_rule'] = value

Hmm yeah theoratically (and maybe also a deleter for Mypy). I’ll investigate.

Uh, Mypy is not smart enough unfortunately python/mypy#4125.

We can’t use the Protocol workaround mentioned in the issue either because BaseOperator has a custom metaclass, while Protocol is also implemented with metaclass, and the two would conflict.

Type ignore it is the other option?

github-actions · 2022-01-21T15:38:34Z

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.

uranusjr · 2022-01-21T18:12:51Z

(Documenting for myself on Monday) This is currently failing with

   tests/models/test_taskinstance.py:2288: 
  _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
  airflow/utils/session.py:71: in wrapper
      return func(*args, session=session, **kwargs)
  airflow/models/taskinstance.py:1621: in run
      session=session,
  airflow/utils/session.py:68: in wrapper
      return func(*args, **kwargs)
  airflow/models/taskinstance.py:1194: in check_and_change_state_before_execution
      dep_context=non_requeueable_dep_context, session=session, verbose=True
  airflow/utils/session.py:68: in wrapper
      return func(*args, **kwargs)
  airflow/models/taskinstance.py:1031: in are_dependencies_met
      for dep_status in self.get_failed_dep_statuses(dep_context=dep_context, session=session):
  airflow/models/taskinstance.py:1052: in get_failed_dep_statuses
      for dep_status in dep.get_dep_statuses(self, session, dep_context):
  airflow/ti_deps/deps/base_ti_dep.py:95: in get_dep_statuses
      yield from self._get_dep_statuses(ti, session, dep_context)
  _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
  
  self = <TIDep(Previous Dagrun State)>
  ti = <[DetachedInstanceError('Instance <TaskInstance at 0x7fefa43f9850> is not bound to a Session; attribute refresh operation cannot proceed') raised in repr()] TaskInstance object at 0x7fefa43f9850>
  session = <sqlalchemy.orm.session.Session object at 0x7fefd5d5f050>
  dep_context = <airflow.ti_deps.dep_context.DepContext object at 0x7fefa4430050>
  
      @provide_session
      def _get_dep_statuses(self, ti, session, dep_context):
          if dep_context.ignore_depends_on_past:
              reason = "The context specified that the state of past DAGs could be ignored."
              yield self._passing_status(reason=reason)
              return
      
  >       if not ti.task.depends_on_past:
  E       AttributeError: 'MappedOperator' object has no attribute 'depends_on_past'
  
  airflow/ti_deps/deps/prev_dagrun_dep.py:41: AttributeError

ashb · 2022-01-22T13:41:48Z

We also need wait_for_downstream in a similar vein for the scheduler - but that can be added in a separate pr

uranusjr · 2022-01-24T07:00:38Z

depends_on_past depends on wait_for_downstream, so I implemented them both.

airflow/models/baseoperator.py

ashb · 2022-01-24T13:51:09Z

Only failure is an actual timing error:

test_heartbeat_failed_fast: assert 0.5497589999999999 < 0.5

Good to go now!

uranusjr requested a review from ashb January 19, 2022 09:58

ashb mentioned this pull request Jan 19, 2022

Set dependencies in MappedOperator via XComArgs #20931

Merged

ashb reviewed Jan 19, 2022

View reviewed changes

airflow/models/base.py Outdated

Copy link

Member

ashb Jan 19, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Anyone else find it odd that Operator is "lower level" than BaseOperator 😁

uranusjr force-pushed the mapped-operator-bag branch 3 times, most recently from c1fd8bc to 20c98a3 Compare January 19, 2022 10:06

ashb reviewed Jan 19, 2022

View reviewed changes

uranusjr force-pushed the mapped-operator-bag branch from 99c78e1 to c4f8d89 Compare January 20, 2022 07:28

uranusjr marked this pull request as ready for review January 20, 2022 07:28

uranusjr requested review from XD-DENG and kaxil as code owners January 20, 2022 07:28

uranusjr requested a review from ashb January 20, 2022 07:28

ashb reviewed Jan 20, 2022

View reviewed changes

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

ashb reviewed Jan 20, 2022

View reviewed changes

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

ashb reviewed Jan 20, 2022

View reviewed changes

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

ashb reviewed Jan 20, 2022

View reviewed changes

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

uranusjr force-pushed the mapped-operator-bag branch from c4f8d89 to 08becd9 Compare January 21, 2022 08:27

uranusjr commented Jan 21, 2022

View reviewed changes

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

ashb mentioned this pull request Jan 21, 2022

Function to expand mapped tasks in to multiple "real" TIs #21019

Merged

ashb reviewed Jan 21, 2022

View reviewed changes

ashb approved these changes Jan 21, 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 Jan 21, 2022

uranusjr force-pushed the mapped-operator-bag branch from 0dece2f to 3cec6d6 Compare January 24, 2022 07:23

ashb reviewed Jan 24, 2022

View reviewed changes

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

ashb approved these changes Jan 24, 2022

View reviewed changes

uranusjr force-pushed the mapped-operator-bag branch from 3cec6d6 to 1d21f86 Compare January 24, 2022 11:47

uranusjr added 11 commits January 24, 2022 20:28

WIP MappedOperator interface parity

4b4a1cc

Remove unneeded getattr

f12ffd9

Fix MappedOperator pool default

d68fb0d

Correctly set relations when circumventing init

8a55a7b

Fix test mocking to point to the correct target

552a982

Declare task_group on BaseOperator for parity

e789c19

Add trigger_rule field to MappedOperator

e887d7a

Fix MappedOperator template_ext serialization

e2952db

Exclude partial_kwargs values from serialization

859dad0

Add wait_for_downstream and depends_on_past

d48f07e

Fix flakiness caused by ti ordering

d56bf4d

uranusjr force-pushed the mapped-operator-bag branch from 27bab1f to d56bf4d Compare January 24, 2022 12:28

uranusjr merged commit ff3bbc3 into apache:main Jan 24, 2022

uranusjr deleted the mapped-operator-bag branch January 24, 2022 13:52

uranusjr mentioned this pull request Jan 25, 2022

Fix last remaining MyPy errors #21020

Merged

ashb added the area:dynamic-task-mapping AIP-42 label Feb 22, 2022

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

Implement enough interface for MappedOperator to be baggable #20945

Implement enough interface for MappedOperator to be baggable #20945

Uh oh!

Conversation

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

Uh oh!

Uh oh!

ashb Jan 19, 2022

Choose a reason for hiding this comment

Uh oh!

ashb Jan 19, 2022

Choose a reason for hiding this comment

Uh oh!

uranusjr Jan 19, 2022

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Uh oh!

uranusjr commented Jan 20, 2022

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ashb Jan 21, 2022

Choose a reason for hiding this comment

Uh oh!

uranusjr Jan 21, 2022

Choose a reason for hiding this comment

Uh oh!

uranusjr Jan 21, 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 Jan 21, 2022

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jan 21, 2022

Uh oh!

uranusjr commented Jan 21, 2022

Uh oh!

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

Uh oh!

Uh oh!

uranusjr commented Jan 24, 2022

Uh oh!

Uh oh!

ashb commented Jan 24, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

uranusjr commented Jan 19, 2022 •

edited

Loading

uranusjr commented Jan 20, 2022 •

edited

Loading

uranusjr Jan 21, 2022 •

edited

Loading

ashb commented Jan 22, 2022 •

edited

Loading