You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "boraberke (via GitHub)" <gi...@apache.org> on 2024/04/09 20:50:56 UTC
[PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
boraberke opened a new pull request, #38883:
URL: https://github.com/apache/airflow/pull/38883
This PR allows TaskGroup(s) to be passed as `branch` in the operators that use `BranchMixIn`.
With this implementation, if the provided string is a `task_group_id` instead of `task_id`, it will be expanded into the roots of that `Task Group`, as suggested [here](https://github.com/apache/airflow/issues/34526#issuecomment-1732136233).
I would greatly appreciate any feedback and would love to work on it more!
Closes: https://github.com/apache/airflow/issues/34526
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #38883:
URL: https://github.com/apache/airflow/pull/38883#issuecomment-2092370209
There’s a test case needs fixing.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #38883:
URL: https://github.com/apache/airflow/pull/38883#discussion_r1577427771
##########
airflow/operators/branch.py:
##########
@@ -34,8 +36,42 @@ class BranchMixIn(SkipMixin):
def do_branch(self, context: Context, branches_to_execute: str | Iterable[str]) -> str | Iterable[str]:
"""Implement the handling of branching including logging."""
self.log.info("Branch into %s", branches_to_execute)
- self.skip_all_except(context["ti"], branches_to_execute)
- return branches_to_execute
+ branch_task_ids = self.expand_task_group_roots(context["ti"], branches_to_execute)
+ self.skip_all_except(context["ti"], branch_task_ids)
+ return branch_task_ids
+
+ def expand_task_group_roots(
+ self, ti: TaskInstance | TaskInstancePydantic, branches_to_execute: str | Iterable[str]
+ ) -> str | Iterable[str]:
Review Comment:
You don’t need to change the call site. But I think `skip_all_except` should support passing in an iterator; you can do that by changing this line
https://github.com/apache/airflow/blob/e69853ac9e222259bb9963b3b8359c294d84e148/airflow/models/skipmixin.py#L184
to use `branch_task_id_set` instead. This should be trivial and has no consequences.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #38883:
URL: https://github.com/apache/airflow/pull/38883#discussion_r1580467975
##########
airflow/operators/branch.py:
##########
@@ -34,8 +36,42 @@ class BranchMixIn(SkipMixin):
def do_branch(self, context: Context, branches_to_execute: str | Iterable[str]) -> str | Iterable[str]:
"""Implement the handling of branching including logging."""
self.log.info("Branch into %s", branches_to_execute)
- self.skip_all_except(context["ti"], branches_to_execute)
- return branches_to_execute
+ branch_task_ids = self.expand_task_group_roots(context["ti"], branches_to_execute)
+ self.skip_all_except(context["ti"], branch_task_ids)
+ return branch_task_ids
+
+ def expand_task_group_roots(
+ self, ti: TaskInstance | TaskInstancePydantic, branches_to_execute: str | Iterable[str]
+ ) -> str | Iterable[str]:
Review Comment:
It accepts any iterator but does not handle it correctly. `branch_task_ids` is iterated twice in the function, which you can’t do with an iterator. The line above is the second iteration we need to get rid of.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal merged PR #38883:
URL: https://github.com/apache/airflow/pull/38883
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #38883:
URL: https://github.com/apache/airflow/pull/38883#discussion_r1577427771
##########
airflow/operators/branch.py:
##########
@@ -34,8 +36,42 @@ class BranchMixIn(SkipMixin):
def do_branch(self, context: Context, branches_to_execute: str | Iterable[str]) -> str | Iterable[str]:
"""Implement the handling of branching including logging."""
self.log.info("Branch into %s", branches_to_execute)
- self.skip_all_except(context["ti"], branches_to_execute)
- return branches_to_execute
+ branch_task_ids = self.expand_task_group_roots(context["ti"], branches_to_execute)
+ self.skip_all_except(context["ti"], branch_task_ids)
+ return branch_task_ids
+
+ def expand_task_group_roots(
+ self, ti: TaskInstance | TaskInstancePydantic, branches_to_execute: str | Iterable[str]
+ ) -> str | Iterable[str]:
Review Comment:
You don’t need to change the call site; `skip_all_except` supports passing in an iterator.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
Posted by "boraberke (via GitHub)" <gi...@apache.org>.
boraberke commented on code in PR #38883:
URL: https://github.com/apache/airflow/pull/38883#discussion_r1575972506
##########
airflow/operators/branch.py:
##########
@@ -34,8 +36,42 @@ class BranchMixIn(SkipMixin):
def do_branch(self, context: Context, branches_to_execute: str | Iterable[str]) -> str | Iterable[str]:
"""Implement the handling of branching including logging."""
self.log.info("Branch into %s", branches_to_execute)
- self.skip_all_except(context["ti"], branches_to_execute)
- return branches_to_execute
+ branch_task_ids = self.expand_task_group_roots(context["ti"], branches_to_execute)
+ self.skip_all_except(context["ti"], branch_task_ids)
+ return branch_task_ids
+
+ def expand_task_group_roots(
+ self, ti: TaskInstance | TaskInstancePydantic, branches_to_execute: str | Iterable[str]
+ ) -> str | Iterable[str]:
Review Comment:
It does not have to be public, I will prefix it with underscore.
If I change this to a generator, how should I modify this part?
https://github.com/apache/airflow/blob/e69853ac9e222259bb9963b3b8359c294d84e148/airflow/operators/branch.py#L39
Should I convert it into list from the generator something like this?
```python
branch_task_ids = list(self._expand_task_group_roots(context["ti"], branches_to_execute))
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #38883:
URL: https://github.com/apache/airflow/pull/38883#discussion_r1575577739
##########
airflow/operators/branch.py:
##########
@@ -34,8 +36,42 @@ class BranchMixIn(SkipMixin):
def do_branch(self, context: Context, branches_to_execute: str | Iterable[str]) -> str | Iterable[str]:
"""Implement the handling of branching including logging."""
self.log.info("Branch into %s", branches_to_execute)
- self.skip_all_except(context["ti"], branches_to_execute)
- return branches_to_execute
+ branch_task_ids = self.expand_task_group_roots(context["ti"], branches_to_execute)
+ self.skip_all_except(context["ti"], branch_task_ids)
+ return branch_task_ids
+
+ def expand_task_group_roots(
+ self, ti: TaskInstance | TaskInstancePydantic, branches_to_execute: str | Iterable[str]
+ ) -> str | Iterable[str]:
Review Comment:
Does this function need to be public? We should add an underscore prefix unless necessary.
Also it would be better if this is a generator that yields task IDs instead.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
Posted by "boraberke (via GitHub)" <gi...@apache.org>.
boraberke commented on code in PR #38883:
URL: https://github.com/apache/airflow/pull/38883#discussion_r1582425547
##########
airflow/operators/branch.py:
##########
@@ -34,8 +36,42 @@ class BranchMixIn(SkipMixin):
def do_branch(self, context: Context, branches_to_execute: str | Iterable[str]) -> str | Iterable[str]:
"""Implement the handling of branching including logging."""
self.log.info("Branch into %s", branches_to_execute)
- self.skip_all_except(context["ti"], branches_to_execute)
- return branches_to_execute
+ branch_task_ids = self.expand_task_group_roots(context["ti"], branches_to_execute)
+ self.skip_all_except(context["ti"], branch_task_ids)
+ return branch_task_ids
+
+ def expand_task_group_roots(
+ self, ti: TaskInstance | TaskInstancePydantic, branches_to_execute: str | Iterable[str]
+ ) -> str | Iterable[str]:
Review Comment:
Oh, got it now, thanks! Changed it as you recommended.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on PR #38883:
URL: https://github.com/apache/airflow/pull/38883#issuecomment-2095431571
> As an additional question or note:
>
> I am not sure if I should modify the docstring of the following operators:
>
> * BaseBranchOperator
> * BranchPythonOperator
> * BranchPythonVirtualenvOperator
> * BranchExternalPythonOperator
> * BranchDateTimeOperator
> * LatestOnlyOperator
> * BranchDayOfWeekOperator
>
> They all mention single task_id or list of task_ids to follow. However, with this PR, they all start to support "single task_id, single task_group_id, and list of combination of task_ids and task_group_ids".
Lets please
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
Posted by "boraberke (via GitHub)" <gi...@apache.org>.
boraberke commented on code in PR #38883:
URL: https://github.com/apache/airflow/pull/38883#discussion_r1580019503
##########
airflow/operators/branch.py:
##########
@@ -34,8 +36,42 @@ class BranchMixIn(SkipMixin):
def do_branch(self, context: Context, branches_to_execute: str | Iterable[str]) -> str | Iterable[str]:
"""Implement the handling of branching including logging."""
self.log.info("Branch into %s", branches_to_execute)
- self.skip_all_except(context["ti"], branches_to_execute)
- return branches_to_execute
+ branch_task_ids = self.expand_task_group_roots(context["ti"], branches_to_execute)
+ self.skip_all_except(context["ti"], branch_task_ids)
+ return branch_task_ids
+
+ def expand_task_group_roots(
+ self, ti: TaskInstance | TaskInstancePydantic, branches_to_execute: str | Iterable[str]
+ ) -> str | Iterable[str]:
Review Comment:
Not sure if I need to change that, `skip_all_except` already accepts passing in an iterator and with the latest commit with generator, it looks like it works fine. What do you think?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #38883:
URL: https://github.com/apache/airflow/pull/38883#discussion_r1580468363
##########
airflow/operators/branch.py:
##########
@@ -34,8 +36,42 @@ class BranchMixIn(SkipMixin):
def do_branch(self, context: Context, branches_to_execute: str | Iterable[str]) -> str | Iterable[str]:
"""Implement the handling of branching including logging."""
self.log.info("Branch into %s", branches_to_execute)
- self.skip_all_except(context["ti"], branches_to_execute)
- return branches_to_execute
+ branch_task_ids = self.expand_task_group_roots(context["ti"], branches_to_execute)
+ self.skip_all_except(context["ti"], branch_task_ids)
+ return branch_task_ids
+
+ def expand_task_group_roots(
+ self, ti: TaskInstance | TaskInstancePydantic, branches_to_execute: str | Iterable[str]
+ ) -> str | Iterable[str]:
Review Comment:
It appears to “work fine” because the generator quietly yields nothing after it is exhausted, and the line becomes a no-op.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
Posted by "boraberke (via GitHub)" <gi...@apache.org>.
boraberke commented on PR #38883:
URL: https://github.com/apache/airflow/pull/38883#issuecomment-2046029118
As an additional question or note:
I am not sure if I should modify the docstring of the following operators:
- BaseBranchOperator
- BranchPythonOperator
- BranchPythonVirtualenvOperator
- BranchExternalPythonOperator
- BranchDateTimeOperator
- LatestOnlyOperator
- BranchDayOfWeekOperator
They all mention single task_id or list of task_ids to follow. However, with this PR, they all start to support "single task_id, single task_group_id, and list of combination of task_ids and task_group_ids".
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow Task Group Ids to be passed as branches in BranchMixIn [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on PR #38883:
URL: https://github.com/apache/airflow/pull/38883#issuecomment-2053978424
Tests are failing. Can you fix them?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org