You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "Lee-W (via GitHub)" <gi...@apache.org> on 2025/04/30 04:56:34 UTC
[PR] fix(task_group): correctly set upstream_map_indexes for mapped task g… [airflow]
Lee-W opened a new pull request, #49996:
URL: https://github.com/apache/airflow/pull/49996
…roup
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
--
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] Handle MappedTaskGroup map indexes [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on PR #49996:
URL: https://github.com/apache/airflow/pull/49996#issuecomment-2850566968
As this is approved, comments addresssed, I'm going to merge this later today
--
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] Handle MappedTaskGroup map indexes [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #49996:
URL: https://github.com/apache/airflow/pull/49996#discussion_r2071359548
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -247,6 +251,31 @@ def ti_run(
)
+def _get_upstream_map_indexes(
+ request: Request, dag_id: str, task_id: str, ti_map_index
+) -> dict[str, int | list[int] | None] | None:
+ dag = request.app.state.dag_bag.get_dag(dag_id)
+ if not dag:
+ return None
+
+ task = dag.get_task(task_id)
+ upstream_map_indexes = {
+ upstream_task.task_id: (
+ # regular tasks
+ None
+ if upstream_task.task_group is None
+ # tasks in the same mapped task group
+ else ti_map_index
+ if task.task_group == upstream_task.task_group
+ # tasks not in the same mapped task group
+ # the upstream mapped task group should combine the xcom as a list and return it
+ else list(range(upstream_task.task_group.get_parse_time_mapped_ti_count()))
+ )
+ for upstream_task in task.upstream_list
+ }
+ return upstream_map_indexes
Review Comment:
This would be a lot nicer to read if implemented like this instead:
```python
def _get_upstream_map_indexes(task, ti_map_index) -> Iterator[tuple[str, int | list[int] | None]]:
for upstream_task in task.upstream_list:
if upstream_task.task_group is None:
yield upstream_task.task_id, None
elif ...
upstream_map_indexes = (
dict(_get_upstream_map_index(dag.get_task(ti.task_id), ti.map_index)
if dag := request.app.state.dag_bag.get_dag(ti.dag_id)
else None
)
```
--
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] Handle MappedTaskGroup map indexes [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #49996:
URL: https://github.com/apache/airflow/pull/49996#discussion_r2071438137
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -247,6 +251,31 @@ def ti_run(
)
+def _get_upstream_map_indexes(
+ request: Request, dag_id: str, task_id: str, ti_map_index
+) -> dict[str, int | list[int] | None] | None:
+ dag = request.app.state.dag_bag.get_dag(dag_id)
+ if not dag:
+ return None
+
+ task = dag.get_task(task_id)
+ upstream_map_indexes = {
+ upstream_task.task_id: (
+ # regular tasks
+ None
+ if upstream_task.task_group is None
+ # tasks in the same mapped task group
+ else ti_map_index
+ if task.task_group == upstream_task.task_group
+ # tasks not in the same mapped task group
+ # the upstream mapped task group should combine the xcom as a list and return it
+ else list(range(upstream_task.task_group.get_parse_time_mapped_ti_count()))
+ )
+ for upstream_task in task.upstream_list
+ }
+ return upstream_map_indexes
Review Comment:
Just updated. Thanks!
--
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] Handle MappedTaskGroup map indexes [airflow]
Posted by "tirkarthi (via GitHub)" <gi...@apache.org>.
tirkarthi commented on code in PR #49996:
URL: https://github.com/apache/airflow/pull/49996#discussion_r2071634966
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -247,6 +261,26 @@ def ti_run(
)
+def _get_upstream_map_indexes(
+ task: Operator, ti_map_index: int
+) -> Iterator[tuple[str, int | list[int] | None]]:
+ for upstream_task in task.upstream_list:
+ map_indexes: int | list[int] | None
+ if upstream_task.task_group is None:
+ # regular tasks
+ map_indexes = None
+ elif task.task_group == upstream_task.task_group:
+ # tasks in the same mapped task group
+ map_indexes = ti_map_index
+ else:
+ # tasks not in the same mapped task group
+ # the upstream mapped task group should combine the xcom as a list and return it
+ mapped_ti_count: int = upstream_task.task_group.get_parse_time_mapped_ti_count()
Review Comment:
`get_parse_time_mapped_ti_count` is present only in `MappedTaskGroup` and not in `TaskGroup` . I was testing this PR and sample dag like below could cause the traceback since here `TaskGroup` is passed which works in Airflow 2. I was also wondering since this is only for parse time count whether there is similar method for dynamic input like `[1, 2, 3]` returned by another `@task`.
This could be taken as another PR too but just thought of adding it since I found it during testing.
```python
from datetime import datetime
try:
from airflow.sdk import DAG, task, task_group
except ImportError:
from airflow import DAG
from airflow.decorators import task, task_group
with DAG(
dag_id="gh49714_1",
start_date=datetime(2025, 1, 1),
schedule=None,
):
@task_group
def add_one_and_double():
@task
def add_one(value: int):
return value + 1
@task
def double(value: int):
return value * 2
new_values = add_one.expand(value=[1, 2, 3])
return double.expand(value=new_values)
@task
def consumer(values: list[int]):
print(f"New values are: {list(values)}")
results = add_one_and_double()
consumer(results)
```
```
│ /home/karthikeyan/stuff/python/airflow/airflow-core/src/airflow/api_fastapi/execution_api/routes │
│ /task_instances.py:278 in _get_upstream_map_indexes │
│ │
│ 275 │ │ else: │
│ 276 │ │ │ # tasks not in the same mapped task group │
│ 277 │ │ │ # the upstream mapped task group should combine the xcom as a list and retur │
│ ❱ 278 │ │ │ mapped_ti_count: int = upstream_task.task_group.get_parse_time_mapped_ti_cou │
│ 279 │ │ │ map_indexes = list(range(mapped_ti_count)) if mapped_ti_count is not None el │
│ 280 │ │ │
│ 281 │ │ yield upstream_task.task_id, map_indexes │
│ │
│ ╭─────────────────────────────────────────── locals ───────────────────────────────────────────╮ │
│ │ task = <Task(_PythonDecoratedOperator): consumer> │ │
│ │ ti_map_index = -1 │ │
│ │ upstream_task = MappedOperator( │ │
│ │ │ operator_class={ │ │
│ │ │ │ 'ui_color': '#ffefeb', │ │
│ │ │ │ 'task_id': 'add_one_and_double.double', │ │
│ │ │ │ 'task_type': '_PythonDecoratedOperator', │ │
│ │ │ │ 'start_from_trigger': False, │ │
│ │ │ │ 'template_ext': [], │ │
│ │ │ │ 'template_fields': ['templates_dict', 'op_args', 'op_kwargs'], │ │
│ │ │ │ 'template_fields_renderers': { │ │
│ │ │ │ │ 'templates_dict': 'json', │ │
│ │ │ │ │ 'op_args': 'py', │ │
│ │ │ │ │ 'op_kwargs': 'py' │ │
│ │ │ │ }, │ │
│ │ │ │ 'ui_fgcolor': '#000', │ │
│ │ │ │ 'downstream_task_ids': ['consumer'], │ │
│ │ │ │ 'start_trigger_args': None │ │
│ │ │ }, │ │
│ │ │ _is_mapped=True, │ │
│ │ │ expand_input=DictOfListsExpandInput(value={}), │ │
│ │ │ partial_kwargs={ │ │
│ │ │ │ 'is_setup': False, │ │
│ │ │ │ 'is_teardown': False, │ │
│ │ │ │ 'on_failure_fail_dagrun': False, │ │
│ │ │ │ 'op_args': [], │ │
│ │ │ │ 'op_kwargs': {} │ │
│ │ │ }, │ │
│ │ │ task_id='add_one_and_double.double', │ │
│ │ │ params={}, │ │
│ │ │ deps=frozenset({ │ │
│ │ │ │ <TIDep(Mapped dependencies have succeeded)>, │ │
│ │ │ │ <TIDep(Previous Dagrun State)>, │ │
│ │ │ │ <TIDep(Not Previously Skipped)>, │ │
│ │ │ │ <TIDep(Not In Retry Period)>, │ │
│ │ │ │ <TIDep(Trigger Rule)> │ │
│ │ │ }), │ │
│ │ │ operator_extra_links=[], │ │
│ │ │ template_ext=[], │ │
│ │ │ template_fields=['templates_dict', 'op_args', 'op_kwargs'], │ │
│ │ │ template_fields_renderers={ │ │
│ │ │ │ 'templates_dict': 'json', │ │
│ │ │ │ 'op_args': 'py', │ │
│ │ │ │ 'op_kwargs': 'py' │ │
│ │ │ }, │ │
│ │ │ ui_color='#ffefeb', │ │
│ │ │ ui_fgcolor='#000', │ │
│ │ │ _is_empty=False, │ │
│ │ │ _can_skip_downstream=False, │ │
│ │ │ _is_sensor=False, │ │
│ │ │ _task_module='airflow.providers.standard.decorators.python', │ │
│ │ │ _task_type='_PythonDecoratedOperator', │ │
│ │ │ _operator_name='@task', │ │
│ │ │ start_trigger_args=None, │ │
│ │ │ start_from_trigger=False, │ │
│ │ │ _needs_expansion=True, │ │
│ │ │ dag=<SerializedDAG dag_id='gh49714_1' schedule=None #tasks=3>, │ │
│ │ │ task_group=<weakproxy at 0x7f26a48cb1f0 to TaskGroup at 0x7f26a6513590>, │ │
│ │ │ start_date=DateTime(2025, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC')), │ │
│ │ │ end_date=None, │ │
│ │ │ upstream_task_ids={'add_one_and_double.add_one'}, │ │
│ │ │ downstream_task_ids={'consumer'}, │ │
│ │ │ _disallow_kwargs_override=False, │ │
│ │ │ _expand_input_attr='op_kwargs_expand_input' │ │
│ │ ) │ │
│ ╰──────────────────────────────────────────────────────────────────────────────────────────────╯ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
AttributeError: 'TaskGroup' object has no attribute 'get_parse_time_mapped_ti_count'
```
--
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] Handle MappedTaskGroup map indexes [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #49996:
URL: https://github.com/apache/airflow/pull/49996#discussion_r2073133313
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -247,6 +261,26 @@ def ti_run(
)
+def _get_upstream_map_indexes(
+ task: Operator, ti_map_index: int
+) -> Iterator[tuple[str, int | list[int] | None]]:
+ for upstream_task in task.upstream_list:
+ map_indexes: int | list[int] | None
+ if upstream_task.task_group is None:
+ # regular tasks
+ map_indexes = None
+ elif task.task_group == upstream_task.task_group:
+ # tasks in the same mapped task group
+ map_indexes = ti_map_index
+ else:
+ # tasks not in the same mapped task group
+ # the upstream mapped task group should combine the xcom as a list and return it
+ mapped_ti_count: int = upstream_task.task_group.get_parse_time_mapped_ti_count()
Review Comment:
This is fixed in the latest verison
--
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] Handle MappedTaskGroup map indexes [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #49996:
URL: https://github.com/apache/airflow/pull/49996#discussion_r2072766292
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -247,6 +261,26 @@ def ti_run(
)
+def _get_upstream_map_indexes(
+ task: Operator, ti_map_index: int
+) -> Iterator[tuple[str, int | list[int] | None]]:
+ for upstream_task in task.upstream_list:
+ map_indexes: int | list[int] | None
+ if upstream_task.task_group is None:
+ # regular tasks
+ map_indexes = None
+ elif task.task_group == upstream_task.task_group:
+ # tasks in the same mapped task group
+ map_indexes = ti_map_index
+ else:
+ # tasks not in the same mapped task group
+ # the upstream mapped task group should combine the xcom as a list and return it
+ mapped_ti_count: int = upstream_task.task_group.get_parse_time_mapped_ti_count()
Review Comment:
Hmmm... this seems to be a valid issue. Let me take a look and see what I can do. Thanks!
--
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] Handle MappedTaskGroup map indexes [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W merged PR #49996:
URL: https://github.com/apache/airflow/pull/49996
--
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