[AIP-31] Implement XComArg model to functionally pass output from one operator to the next #8652

jonathanshir · 2020-04-30T18:15:25Z

This PR is a part of [AIP-31], and closes the following issues:
closes #8055, closes #8053, closes #8052 .
The goal of the XComArg model is to create a communication token to be passed between operators. In this fashion, we can wire the output of one operator to the next in a functional and "comfortable" (code-writing-wise) manner.
Collaborated with: @turbaszek @casassg @evgenyshulman

Make sure to mark the boxes below before creating PR: [x]

Description above provides context of the change
Unit tests coverage for changes (not needed for documentation changes)
Target Github ISSUE in description if exists
Commits follow "How to write a good git commit message"
Relevant documentation is updated including usage instructions.
I will engage committers as explained in Contribution Workflow Example.

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.
Read the Pull Request Guidelines for more information.

Added rendering of XComArg in base operator Added xcom arg test, todo multiple keys

Add output property to BaseOperator closes: apache#8055

adding support for runtime sub-item access adding support for usage templated_args (legacy)

Made XCom unit tests run faster

airflow/models/baseoperator.py

airflow/models/xcom_arg.py

scripts/ci/pylint_todo.txt

tests/models/test_xcom_arg.py

airflow/models/xcom_arg.py

turbaszek · 2020-05-05T09:44:37Z

airflow/models/baseoperator.py

@@ -1064,10 +1069,23 @@ def _set_relatives(self,
                       task_or_task_list: Union['BaseOperator', List['BaseOperator']],


I think we need to adjust this type annotation. But it will be hard due to circular imports.

turbaszek · 2020-05-05T09:49:15Z

airflow/models/xcom_arg.py

+        if self.key is not None:
+            xcom_pull_kwargs.append(f"key='{self.key}'")


Is it possible that self.key will be None? The default value is XCOM_RETURN_KEY and as far as I understand the key has to be provided. If it's so, we should have a check for that in __init__.

Key is none is only possible if the base class is overridden (MyXComArg inherits XComArg but doesn't provide a key at all, for example).
I think this is good functionality in case your xcom system doesn't have a key necessarily

Wait, didn't we make XCom overwritable instead? I guess you can overwrite it and set it on the output property, but not sure if I understand the use for that?

Not sure exactly what you meant cas, can you elaborate on "overwritable instead"?

How is XComArg overridden? XComArg is not easy to be overwritten currently. Can you show an example use case where XComArg is overridden by a subclass?

XCom system itself supports key=None, so our goal with this was to allow that kind of usage, so if XCom backend uses key=None we already have support for it in XComArg.

sgtm. Maybe we should still print that key value is None. But not blocking for me.

Co-authored-by: Tomek Urbaszek <turbaszek@gmail.com>

airflow/models/xcom_arg.py

turbaszek · 2020-05-05T15:37:20Z

It looks good to me. @kaxil @casassg @dimberman WDYT?

dimberman

LGTM! Great job @jonathanshir!

kaxil · 2020-05-05T19:55:10Z

airflow/models/baseoperator.py

@@ -1121,6 +1139,12 @@ def set_upstream(self, task_or_task_list: Union['BaseOperator', List['BaseOperat
        """
        self._set_relatives(task_or_task_list, upstream=True)

+    @property
+    def output(self):
+        """Returns default XComArg for the operator"""


Can we add a bit more detail, something like Output from the previous task using XComArg or something similar so that is simpler to understand. WDYT ?

"XCom representation of output received from this operator"?

"Returns reference to XCom pushed by current operator"?

Either of them is fine :)

We haven't updated this yet :)

Merging this PR ignoring this change (we can update it in a related follow-up PR)

Addressed in #8805

airflow/models/xcom_arg.py

tests/models/test_xcom_arg.py

casassg

overall lgtm, added a comment on the self.key being null

Comments, as well as single assert statement in XComArg tests. Co-authored-by: Tomek Urbaszek <turbaszek@gmail.com> Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

kaxil

Almost there, just one comment from me and I think couple of question from @casassg and @turbaszek

turbaszek · 2020-05-07T22:39:47Z

@jonathanshir it seems that some of the tests are flaky or something:

File "/opt/airflow/tests/models/test_xcom_arg.py", line 130 in test_xcom_pass_to_op

In general, we are not running DAGs on CI. So I would suggest either mark the TestXComArgRuntime with system marker or remove/mock this test.

boring-cyborg · 2020-05-09T22:52:13Z

Awesome work, congrats on your first merged pull request!

kaxil · 2020-05-10T16:57:51Z

Tentatively added it to 1.10.11 milestone, will double-check when cherry-picking though to see if it for sure doesn't break backward-compatibility.

ashb · 2020-05-19T13:05:27Z

So this .output wasn't in the AIP proposal.

Why have we introduced it?

turbaszek · 2020-05-19T13:08:15Z

So this .output wasn't in the AIP proposal.

@ashb #8055 , main idea was to be able to mix functional and non-functional ops in one DAG.

ashb · 2020-05-19T13:11:00Z

So this .output wasn't in the AIP proposal.

@ashb #8055 , main idea was to be able to mix functional and non-functional ops in one DAG.

Will the example from the AIP still (eventually) work?

with DAG(
    'send_server_ip', default_args=default_args, schedule_interval=None
) as dag:

  # Using default connection as it's set to httpbin.org by default
  get_ip = SimpleHttpOperator(
      task_id='get_ip', endpoint='get', method='GET', xcom_push=True
  )

  @dag.task(add_context=True)
  def email_subject_generator(context: dict, raw_json: str) -> str:
    external_ip = json.loads(raw_json)['origin']
    return f"Server connected from {external_ip}"

  @task
  def email_body_generator(raw_json: str) -> str:
    return """
    Seems like today your server executing Airflow is connected from the external IP {origin}<br>
    <br>
    Have a nice day!
    """.format(**json.loads(raw_json))

  send_email = EmailOperator(
      task_id='send_email',
      to="example@example.com",
      subject='',
      html_content=''
  )

  ip_info = get_ip()
  subject = email_subject_generator(ip_info)
  body = email_body_generator(ip_info)
  send_email(subject=subject, html_content=body)

turbaszek · 2020-05-19T13:24:56Z

Hm, I think not. I think it should be:

with DAG(
    'send_server_ip', default_args=default_args, schedule_interval=None
) as dag:

  # Using default connection as it's set to httpbin.org by default
  get_ip = SimpleHttpOperator(
      task_id='get_ip', endpoint='get', method='GET', xcom_push=True
  )

  @dag.task(add_context=True)
  def email_subject_generator(context: dict, raw_json: str) -> str:
    external_ip = json.loads(raw_json)['origin']
    return f"Server connected from {external_ip}"

  @task
  def email_body_generator(raw_json: str) -> str:
    return """
    Seems like today your server executing Airflow is connected from the external IP {origin}<br>
    <br>
    Have a nice day!
    """.format(**json.loads(raw_json))

  ip_info = get_ip()
  subject = email_subject_generator(ip_info)
  body = email_body_generator(ip_info)

  send_email = EmailOperator(
      task_id='send_email',
      to="example@example.com",
      subject=subject,
      html_content=body
  )

CC: @casassg

ashb · 2020-05-19T14:35:17Z

Okay, but the email_subject_generator(ip_info) (vs email_subject_generator(ip_info.output) as it has to be right now) will be fixed? Do we already have an issue for that?

casassg · 2020-05-19T14:55:11Z

Main idea is that calling regular operators may be a bit confusing vs setting XComArg on initialization. We can discuss this further w the @task PR coming soon that also brings some more end to end examples of what this would look like. Even though its a change from the initial AIP I think its for the best. Happy to write down the change and ammend/propose it in the original doc. It was mainly an alternative that appeared while we were implementing.

turbaszek · 2020-05-19T14:56:34Z

Okay, but the email_subject_generator(ip_info) (vs email_subject_generator(ip_info.output) as it has to be right now)

It can be email_subject_generator(ip_info). The .output is only for operators defined in "old way" (not decorated with @task).

Calling a function decorated with @task returns XComArg. Calling an operator returns an instance of this operator. So the .output is shorthand for op.output == XComArg(op). But as @casassg says, we can adjust it once we have @task because then things will be simpler.

jonathanshir and others added 8 commits April 30, 2020 20:35

Basic XComArg implementation

af13770

Added rendering of XComArg in base operator Added xcom arg test, todo multiple keys

Support chaining between XcomArg and BaseOperator

692b08b

Add output property to BaseOperator closes: apache#8055

Add operator/keys getters

0edd479

fixing support for XArgCom at BaseOperator

68527f1

adding support for runtime sub-item access adding support for usage templated_args (legacy)

Removed redundancies from xcom arg model and tests

524c661

Adjust doscstrings

c2494dd

Fixed operator property usage

c6ff5dc

Made XCom unit tests run faster

adding support for xcomarg['key_for_xcom'] accordingly to current aip-31

30015ed

boring-cyborg bot added the area:dev-tools label Apr 30, 2020

jonathanshir changed the title ~~[AIP-31] Implement XComArg model to be passed between operators~~ [AIP-31] Implement XComArg model to functionally pass output from one operator to the next Apr 30, 2020

casassg mentioned this pull request May 1, 2020

[8057] [AIP-31] Add @task decorator databand-ai/airflow#5

Closed

13 tasks

kaxil reviewed May 2, 2020

View reviewed changes

airflow/models/baseoperator.py Show resolved Hide resolved

kaxil reviewed May 2, 2020

View reviewed changes

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

kaxil reviewed May 2, 2020

View reviewed changes

scripts/ci/pylint_todo.txt Outdated Show resolved Hide resolved

kaxil reviewed May 2, 2020

View reviewed changes

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

kaxil reviewed May 2, 2020

View reviewed changes

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

jonathanshir added 2 commits May 3, 2020 13:27

Improved docstrings, improved asserts in unit tests

992a406

Improved documentation

be95dca

turbaszek added the AIP-31 Task Flow API for nicer DAG definition label May 4, 2020

turbaszek reviewed May 4, 2020

View reviewed changes

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

turbaszek reviewed May 5, 2020

View reviewed changes

Update xcom_arg.py usage documentation

3e98dee

Co-authored-by: Tomek Urbaszek <turbaszek@gmail.com>

turbaszek reviewed May 5, 2020

View reviewed changes

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

turbaszek reviewed May 5, 2020

View reviewed changes

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

turbaszek reviewed May 5, 2020

View reviewed changes

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

jonathanshir marked this pull request as ready for review May 5, 2020 16:17

dimberman approved these changes May 5, 2020

View reviewed changes

kaxil reviewed May 5, 2020

View reviewed changes

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

kaxil reviewed May 5, 2020

View reviewed changes

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

casassg approved these changes May 6, 2020

View reviewed changes

Apply documentation and test suggestions from code review

93a5376

Comments, as well as single assert statement in XComArg tests. Co-authored-by: Tomek Urbaszek <turbaszek@gmail.com> Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

turbaszek approved these changes May 6, 2020

View reviewed changes

kaxil reviewed May 6, 2020

View reviewed changes

turbaszek force-pushed the feature/aip31_xcom_arg_model branch from e8c652e to 5c00ca8 Compare May 8, 2020 10:35

Use DebugExecutor in TestXComArgRuntime

955b45f

turbaszek force-pushed the feature/aip31_xcom_arg_model branch from 5c00ca8 to 955b45f Compare May 8, 2020 13:36

turbaszek requested a review from kaxil May 8, 2020 15:53

kaxil merged commit bc19778 into apache:master May 9, 2020

kaxil added this to the Airflow 1.10.11 milestone May 10, 2020

casassg mentioned this pull request May 21, 2020

[AIRFLOW-8057] [AIP-31] Add @task decorator #8962

Merged

6 tasks

kaxil modified the milestones: Airflow 1.10.11, Airflow 2.0.0 Jun 22, 2020

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[AIP-31] Implement XComArg model to functionally pass output from one operator to the next #8652

[AIP-31] Implement XComArg model to functionally pass output from one operator to the next #8652

jonathanshir commented Apr 30, 2020 •

edited

turbaszek May 5, 2020

turbaszek May 5, 2020

jonathanshir May 5, 2020

casassg May 5, 2020

jonathanshir May 6, 2020

casassg May 6, 2020

jonathanshir May 7, 2020

casassg May 7, 2020

turbaszek commented May 5, 2020

dimberman left a comment

kaxil May 5, 2020

jonathanshir May 6, 2020

casassg May 6, 2020

kaxil May 6, 2020

kaxil May 9, 2020

kaxil May 9, 2020

turbaszek May 10, 2020

casassg left a comment

kaxil left a comment

turbaszek commented May 7, 2020 •

edited

boring-cyborg bot commented May 9, 2020

kaxil commented May 10, 2020

ashb commented May 19, 2020

turbaszek commented May 19, 2020 •

edited

ashb commented May 19, 2020 •

edited

turbaszek commented May 19, 2020

ashb commented May 19, 2020

casassg commented May 19, 2020

turbaszek commented May 19, 2020

		@@ -1064,10 +1069,23 @@ def _set_relatives(self,
		task_or_task_list: Union['BaseOperator', List['BaseOperator']],

		if self.key is not None:
		xcom_pull_kwargs.append(f"key='{self.key}'")

[AIP-31] Implement XComArg model to functionally pass output from one operator to the next #8652

[AIP-31] Implement XComArg model to functionally pass output from one operator to the next #8652

Conversation

jonathanshir commented Apr 30, 2020 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

turbaszek commented May 5, 2020

dimberman left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

casassg left a comment

Choose a reason for hiding this comment

kaxil left a comment

Choose a reason for hiding this comment

turbaszek commented May 7, 2020 • edited

boring-cyborg bot commented May 9, 2020

kaxil commented May 10, 2020

ashb commented May 19, 2020

turbaszek commented May 19, 2020 • edited

ashb commented May 19, 2020 • edited

turbaszek commented May 19, 2020

ashb commented May 19, 2020

casassg commented May 19, 2020

turbaszek commented May 19, 2020

jonathanshir commented Apr 30, 2020 •

edited

turbaszek commented May 7, 2020 •

edited

turbaszek commented May 19, 2020 •

edited

ashb commented May 19, 2020 •

edited