You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "dstandish (via GitHub)" <gi...@apache.org> on 2024/04/25 14:59:37 UTC

[PR] Handle db isolation for mapped operators and task groups [airflow]

dstandish opened a new pull request, #39259:
URL: https://github.com/apache/airflow/pull/39259

   (no comment)


-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1607024896


##########
airflow/models/taskinstance.py:
##########
@@ -1003,25 +1009,40 @@ def _refresh_from_task(
     task_instance_mutation_hook(task_instance)
 
 
+@internal_api_call
+@provide_session
 def _record_task_map_for_downstreams(
-    *, task_instance: TaskInstance | TaskInstancePydantic, task: Operator, value: Any, session: Session
+    *,
+    task_instance: TaskInstance | TaskInstancePydantic,
+    task: Operator,
+    dag: DAG,
+    value: Any,
+    session: Session,
 ) -> None:
     """
     Record the task map for downstream tasks.
 
     :param task_instance: the task instance
     :param task: The task object
+    :param dag: the dag associated with the task
     :param value: The value
     :param session: SQLAlchemy ORM Session
 
     :meta private:
     """
+    # when taking task over RPC, we need to add the dag back
+    if isinstance(task, MappedOperator):
+        if not task.dag:
+            task.dag = dag
+    elif not task._dag:
+        task._dag = dag
+
     if next(task.iter_mapped_dependants(), None) is None:  # No mapped dependants, no need to validate.
         return
     # TODO: We don't push TaskMap for mapped task instances because it's not
-    # currently possible for a downstream to depend on one individual mapped
-    # task instance. This will change when we implement task mapping inside
-    # a mapped task group, and we'll need to further analyze the case.
+    #  currently possible for a downstream to depend on one individual mapped
+    #  task instance. This will change when we implement task mapping inside
+    #  a mapped task group, and we'll need to further analyze the case.

Review Comment:
   Accidental?



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1583394477


##########
airflow/models/taskinstance.py:
##########
@@ -1003,25 +1009,40 @@ def _refresh_from_task(
     task_instance_mutation_hook(task_instance)
 
 
+@internal_api_call
+@provide_session
 def _record_task_map_for_downstreams(
-    *, task_instance: TaskInstance | TaskInstancePydantic, task: Operator, value: Any, session: Session
+    *,
+    task_instance: TaskInstance | TaskInstancePydantic,
+    task: Operator,
+    dag: DAG,
+    value: Any,
+    session: Session,
 ) -> None:
     """
     Record the task map for downstream tasks.
 
     :param task_instance: the task instance
     :param task: The task object
+    :param dag: the dag associated with the task
     :param value: The value
     :param session: SQLAlchemy ORM Session
 
     :meta private:
     """
+    # when taking task over RPC, we need to add the dag back
+    if isinstance(task, MappedOperator):
+        if not task.dag:
+            task.dag = dag
+    elif not task._dag:
+        task._dag = dag

Review Comment:
   the issue @uranusjr is that this _is_ early in the stack when it's a RPC call. the only earlier place we could do it is in the decorator. WDYT?  we could stick it in a private function though and get it out of the way and reuse in module though....
   
   when not a RPC call, this has no effect and is not needed.



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1607061565


##########
airflow/models/taskinstance.py:
##########
@@ -1003,25 +1009,40 @@ def _refresh_from_task(
     task_instance_mutation_hook(task_instance)
 
 
+@internal_api_call
+@provide_session
 def _record_task_map_for_downstreams(
-    *, task_instance: TaskInstance | TaskInstancePydantic, task: Operator, value: Any, session: Session
+    *,
+    task_instance: TaskInstance | TaskInstancePydantic,
+    task: Operator,
+    dag: DAG,
+    value: Any,
+    session: Session,
 ) -> None:
     """
     Record the task map for downstream tasks.
 
     :param task_instance: the task instance
     :param task: The task object
+    :param dag: the dag associated with the task
     :param value: The value
     :param session: SQLAlchemy ORM Session
 
     :meta private:
     """
+    # when taking task over RPC, we need to add the dag back
+    if isinstance(task, MappedOperator):
+        if not task.dag:
+            task.dag = dag
+    elif not task._dag:
+        task._dag = dag

Review Comment:
   fyi @uranusjr this is resolved here ([Use sentinel to elide the dag object on reserialization](https://github.com/apache/airflow/pull/37851/commits/7ba5fc8ac05b2003ef5050f9b13f34e0fdb1ef80)) but i can't make this PR yet because it's depending on too many other PRs to get merged first



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on PR #39259:
URL: https://github.com/apache/airflow/pull/39259#issuecomment-2078024289

   and here @uranusjr ?


-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk merged PR #39259:
URL: https://github.com/apache/airflow/pull/39259


-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1580437417


##########
airflow/models/taskinstance.py:
##########
@@ -3167,6 +3188,8 @@ def render_templates(
         # MappedOperator is useless for template rendering, and we need to be
         # able to access the unmapped task instead.
         original_task.render_template_fields(context, jinja_env)
+        if isinstance(self.task, MappedOperator):
+            self.task = context["ti"].task

Review Comment:
   That’s right because only when it is mappedoperator is ti.task mutated. Otherwise ti.task is the result of rpc call and long story short it can’t be used



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1583394477


##########
airflow/models/taskinstance.py:
##########
@@ -1003,25 +1009,40 @@ def _refresh_from_task(
     task_instance_mutation_hook(task_instance)
 
 
+@internal_api_call
+@provide_session
 def _record_task_map_for_downstreams(
-    *, task_instance: TaskInstance | TaskInstancePydantic, task: Operator, value: Any, session: Session
+    *,
+    task_instance: TaskInstance | TaskInstancePydantic,
+    task: Operator,
+    dag: DAG,
+    value: Any,
+    session: Session,
 ) -> None:
     """
     Record the task map for downstream tasks.
 
     :param task_instance: the task instance
     :param task: The task object
+    :param dag: the dag associated with the task
     :param value: The value
     :param session: SQLAlchemy ORM Session
 
     :meta private:
     """
+    # when taking task over RPC, we need to add the dag back
+    if isinstance(task, MappedOperator):
+        if not task.dag:
+            task.dag = dag
+    elif not task._dag:
+        task._dag = dag

Review Comment:
   the issue @uranusjr is that this _is_ early in the stack when it's a RPC call. the only earlier place we could do it is in the decorator. WDYT?  we could stick it in a private function though and get it out of the way and reuse in module though....



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1580481269


##########
airflow/models/taskinstance.py:
##########
@@ -3167,6 +3188,8 @@ def render_templates(
         # MappedOperator is useless for template rendering, and we need to be
         # able to access the unmapped task instead.
         original_task.render_template_fields(context, jinja_env)
+        if isinstance(self.task, MappedOperator):
+            self.task = context["ti"].task

Review Comment:
   > It makes sense, but `if isinstance(self.task, MappedOperator)` is an awkward condition to check for the case.
   
   yeah, i see what you're saying. e.g. better would be for the code to "tell us" when an unmap has happened.
   
   like when we call 
   ```
   original_task.render_template_fields(context, jinja_env)
   ```
   that could like... return a new task when it creates one. that would certainly make it more obvious what is going on too.



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1579984115


##########
airflow/models/taskinstance.py:
##########
@@ -1003,25 +1009,43 @@ def _refresh_from_task(
     task_instance_mutation_hook(task_instance)
 
 
+@internal_api_call
+@provide_session
 def _record_task_map_for_downstreams(
-    *, task_instance: TaskInstance | TaskInstancePydantic, task: Operator, value: Any, session: Session
+    *,
+    task_instance: TaskInstance | TaskInstancePydantic,
+    task: Operator,
+    dag: DAG,
+    value: Any,
+    session: Session,
 ) -> None:
     """
     Record the task map for downstream tasks.
 
     :param task_instance: the task instance
     :param task: The task object
+    :param dag: the dag associated with the task
     :param value: The value
     :param session: SQLAlchemy ORM Session
 
     :meta private:
     """
+    # if not task._dag:
+    #     task._dag = dag  # required when on RPC server side
+

Review Comment:
   ```suggestion
   ```



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1579978066


##########
airflow/models/taskinstance.py:
##########
@@ -1003,25 +1009,43 @@ def _refresh_from_task(
     task_instance_mutation_hook(task_instance)
 
 
+@internal_api_call
+@provide_session
 def _record_task_map_for_downstreams(
-    *, task_instance: TaskInstance | TaskInstancePydantic, task: Operator, value: Any, session: Session
+    *,
+    task_instance: TaskInstance | TaskInstancePydantic,
+    task: Operator,
+    dag: DAG,
+    value: Any,
+    session: Session,
 ) -> None:
     """
     Record the task map for downstream tasks.
 
     :param task_instance: the task instance
     :param task: The task object
+    :param dag: the dag associated with the task
     :param value: The value
     :param session: SQLAlchemy ORM Session
 
     :meta private:
     """
+    # if not task._dag:

Review Comment:
   left-over? 



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1579983764


##########
airflow/models/taskinstance.py:
##########
@@ -1003,25 +1009,43 @@ def _refresh_from_task(
     task_instance_mutation_hook(task_instance)
 
 
+@internal_api_call
+@provide_session
 def _record_task_map_for_downstreams(
-    *, task_instance: TaskInstance | TaskInstancePydantic, task: Operator, value: Any, session: Session
+    *,
+    task_instance: TaskInstance | TaskInstancePydantic,
+    task: Operator,
+    dag: DAG,
+    value: Any,
+    session: Session,
 ) -> None:
     """
     Record the task map for downstream tasks.
 
     :param task_instance: the task instance
     :param task: The task object
+    :param dag: the dag associated with the task
     :param value: The value
     :param session: SQLAlchemy ORM Session
 
     :meta private:
     """
+    # if not task._dag:

Review Comment:
   correct



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1580432816


##########
airflow/models/taskinstance.py:
##########
@@ -1003,25 +1009,40 @@ def _refresh_from_task(
     task_instance_mutation_hook(task_instance)
 
 
+@internal_api_call
+@provide_session
 def _record_task_map_for_downstreams(
-    *, task_instance: TaskInstance | TaskInstancePydantic, task: Operator, value: Any, session: Session
+    *,
+    task_instance: TaskInstance | TaskInstancePydantic,
+    task: Operator,
+    dag: DAG,
+    value: Any,
+    session: Session,
 ) -> None:
     """
     Record the task map for downstream tasks.
 
     :param task_instance: the task instance
     :param task: The task object
+    :param dag: the dag associated with the task
     :param value: The value
     :param session: SQLAlchemy ORM Session
 
     :meta private:
     """
+    # when taking task over RPC, we need to add the dag back
+    if isinstance(task, MappedOperator):
+        if not task.dag:
+            task.dag = dag
+    elif not task._dag:
+        task._dag = dag

Review Comment:
   Not a fan for this… can we do this earlier in the stack, say when the task is created 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] Handle db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1580450719


##########
airflow/models/taskinstance.py:
##########
@@ -3167,6 +3188,8 @@ def render_templates(
         # MappedOperator is useless for template rendering, and we need to be
         # able to access the unmapped task instead.
         original_task.render_template_fields(context, jinja_env)
+        if isinstance(self.task, MappedOperator):
+            self.task = context["ti"].task

Review Comment:
   that make sense @uranusjr ?
   
   so with normal task, `self.task` is the task that is created locally, and there is no need to override it from the one in context dict.  and if you _did_ that then you'd take a task object that isn't quite complete, essentially because we don't have proper serialization of Task since there's no real Task entity and no TaskPydantic.  But generally it's not a problem because most of the time we don't need to serialize a task object.
   
   in the mappedoperator case though, as we saw last night, the "unmap" operation mutates the ti in the context dict.
   
   if searching for alternatives, we could look at not relying on the context dict for this "unmapping".  e.g. we could forword the "original" ti object to the thing doing the unmapping so we don't need to mutate what's in context.
   
   another option would be, upon receiving a fresh context dict over RPC, we could replace the TIs in the context with the local TIPydantic object -- or something to this effect.  then perhaps we could keep the context["ti"] mutation approach for unmapping.



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1580465386


##########
airflow/models/taskinstance.py:
##########
@@ -3167,6 +3188,8 @@ def render_templates(
         # MappedOperator is useless for template rendering, and we need to be
         # able to access the unmapped task instead.
         original_task.render_template_fields(context, jinja_env)
+        if isinstance(self.task, MappedOperator):
+            self.task = context["ti"].task

Review Comment:
   It makes sense, but `if isinstance(self.task, MappedOperator)` is an awkward condition to check for the case.
   
   > upon receiving a fresh context dict over RPC, we could replace the TIs in the context with the local TIPydantic object
   
   This sounds somewhat promising. Instead of just the ti, we could probably try to replace the entire relationship (including e.g. `dag`) so we can get rid of needing to pass in `dag` separately into `_record_task_map_for_downstreams`.



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1579989027


##########
airflow/models/taskinstance.py:
##########
@@ -3167,6 +3188,8 @@ def render_templates(
         # MappedOperator is useless for template rendering, and we need to be
         # able to access the unmapped task instead.
         original_task.render_template_fields(context, jinja_env)
+        if isinstance(self.task, MappedOperator):
+            self.task = context["ti"].task

Review Comment:
   So, this is an interesting one @potiuk. The way mapped operators are "expanded" or "unmapped"... it happens inside of `MappedOperator.render_template_fields`.  It does so by replacing the `task` attr on the `ti` in the context dictionary, which in the non-db-isolation case mutates what is here `self.task`!  But in db isolation case, the context dict is created via RPC and so the pydantic TI in the context dict is not the same as the PydanticTI that is running....  it's .... quite complicated.  But anyway this here is one way to ensure that the task gets properly unmapped -- we don't here rely on mutating the TI in the context dict.



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1580433108


##########
airflow/models/taskinstance.py:
##########
@@ -3167,6 +3188,8 @@ def render_templates(
         # MappedOperator is useless for template rendering, and we need to be
         # able to access the unmapped task instead.
         original_task.render_template_fields(context, jinja_env)
+        if isinstance(self.task, MappedOperator):
+            self.task = context["ti"].task

Review Comment:
   Does this not work with BaseOperator? The conditional makes this a lot weirder.



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1623787572


##########
airflow/models/taskinstance.py:
##########
@@ -1003,25 +1009,40 @@ def _refresh_from_task(
     task_instance_mutation_hook(task_instance)
 
 
+@internal_api_call
+@provide_session
 def _record_task_map_for_downstreams(
-    *, task_instance: TaskInstance | TaskInstancePydantic, task: Operator, value: Any, session: Session
+    *,
+    task_instance: TaskInstance | TaskInstancePydantic,
+    task: Operator,
+    dag: DAG,
+    value: Any,
+    session: Session,
 ) -> None:
     """
     Record the task map for downstream tasks.
 
     :param task_instance: the task instance
     :param task: The task object
+    :param dag: the dag associated with the task
     :param value: The value
     :param session: SQLAlchemy ORM Session
 
     :meta private:
     """
+    # when taking task over RPC, we need to add the dag back
+    if isinstance(task, MappedOperator):
+        if not task.dag:
+            task.dag = dag
+    elif not task._dag:
+        task._dag = dag
+
     if next(task.iter_mapped_dependants(), None) is None:  # No mapped dependants, no need to validate.
         return
     # TODO: We don't push TaskMap for mapped task instances because it's not
-    # currently possible for a downstream to depend on one individual mapped
-    # task instance. This will change when we implement task mapping inside
-    # a mapped task group, and we'll need to further analyze the case.
+    #  currently possible for a downstream to depend on one individual mapped
+    #  task instance. This will change when we implement task mapping inside
+    #  a mapped task group, and we'll need to further analyze the case.

Review Comment:
   Yes it does and the editor prefers it that way. If you veto just let me know and I’ll revert it otherwise let’s move on



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #39259:
URL: https://github.com/apache/airflow/pull/39259#issuecomment-2162615938

   Just to add a note that eventually we should strive to remove all those cases where MappedOperator and BaseOperator need different behaviour.


-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1580450719


##########
airflow/models/taskinstance.py:
##########
@@ -3167,6 +3188,8 @@ def render_templates(
         # MappedOperator is useless for template rendering, and we need to be
         # able to access the unmapped task instead.
         original_task.render_template_fields(context, jinja_env)
+        if isinstance(self.task, MappedOperator):
+            self.task = context["ti"].task

Review Comment:
   that make sense @uranusjr ?
   
   so with normal task, `self.task` is the task that is created locally, and there is no need to override it from the one in context dict.  and if you _did_ that then you'd take a task object that isn't quite complete, essentially because we don't have proper serialization of Task since there's no real Task entity and no TaskPydantic.  But generally it's not a problem because most of the time we don't need to serialize a task object.
   
   in the mappedoperator case though, as we saw last night, "unmapping" is achieved by mutating the ti in the context dict, and it relies on the assumption that the TI in the context dict is the same object as the one that is created locally and being run, which isn't true when the context comes from RPC.
   
   if searching for alternatives, we could look at not relying on the context dict for this "unmapping".  e.g. we could forword the "original" ti object to the thing doing the unmapping so we don't need to mutate what's in context.
   
   another option would be, upon receiving a fresh context dict over RPC, we could replace the TIs in the context with the local TIPydantic object -- or something to this effect.  then perhaps we could keep the context["ti"] mutation approach for unmapping.



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1580450719


##########
airflow/models/taskinstance.py:
##########
@@ -3167,6 +3188,8 @@ def render_templates(
         # MappedOperator is useless for template rendering, and we need to be
         # able to access the unmapped task instead.
         original_task.render_template_fields(context, jinja_env)
+        if isinstance(self.task, MappedOperator):
+            self.task = context["ti"].task

Review Comment:
   that make sense @uranusjr ?
   
   so with normal task, `self.task` is the task that is created locally, and there is no need to override it from the one in context dict.  and if you _did_ that then you'd take a task object that isn't quite complete, essentially because we don't have proper serialization of Task since there's no real Task entity and no TaskPydantic.  But generally it's not a problem because most of the time we don't need to serialize a task object.
   
   in the mappedoperator case though, as we saw last night, "unmapping" is achieved by mutating the ti in the context dict, and it relies on the assumption that the TI in the context dict is the same object as the one that is created locally and being run, which isn't true when the context comes from RPC.
   
   if searching for alternatives, we could look at not relying on the context dict for this "unmapping".  e.g. we could forword the "original" ti object to the thing doing the unmapping so we don't need to mutate what's in context.
   
   another option would be, upon receiving a fresh context dict over RPC, we could replace the TIs in the context with the local TIPydantic object -- or something to this effect.  then perhaps we could keep the context["ti"] mutation approach for unmapping.
   
   we could also look at changing the way we handle context over RPC. currently it's just a "working" approach but not optimal because there's no laziness.  we could optimize by making each context object an accessor that is an RPC call (and we should do something like this ). and something like that could help here too.



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1623772242


##########
airflow/models/taskinstance.py:
##########
@@ -1003,25 +1009,40 @@ def _refresh_from_task(
     task_instance_mutation_hook(task_instance)
 
 
+@internal_api_call
+@provide_session
 def _record_task_map_for_downstreams(
-    *, task_instance: TaskInstance | TaskInstancePydantic, task: Operator, value: Any, session: Session
+    *,
+    task_instance: TaskInstance | TaskInstancePydantic,
+    task: Operator,
+    dag: DAG,
+    value: Any,
+    session: Session,
 ) -> None:
     """
     Record the task map for downstream tasks.
 
     :param task_instance: the task instance
     :param task: The task object
+    :param dag: the dag associated with the task
     :param value: The value
     :param session: SQLAlchemy ORM Session
 
     :meta private:
     """
+    # when taking task over RPC, we need to add the dag back
+    if isinstance(task, MappedOperator):
+        if not task.dag:
+            task.dag = dag
+    elif not task._dag:
+        task._dag = dag
+
     if next(task.iter_mapped_dependants(), None) is None:  # No mapped dependants, no need to validate.
         return
     # TODO: We don't push TaskMap for mapped task instances because it's not
-    # currently possible for a downstream to depend on one individual mapped
-    # task instance. This will change when we implement task mapping inside
-    # a mapped task group, and we'll need to further analyze the case.
+    #  currently possible for a downstream to depend on one individual mapped
+    #  task instance. This will change when we implement task mapping inside
+    #  a mapped task group, and we'll need to further analyze the case.

Review Comment:
   But this only adds a space to every line of the comment…?



-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on PR #39259:
URL: https://github.com/apache/airflow/pull/39259#issuecomment-2155287687

   @uranusjr you have any concerns here?


-- 
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 db isolation for mapped operators and task groups [airflow]

Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #39259:
URL: https://github.com/apache/airflow/pull/39259#discussion_r1607059883


##########
airflow/models/taskinstance.py:
##########
@@ -1003,25 +1009,40 @@ def _refresh_from_task(
     task_instance_mutation_hook(task_instance)
 
 
+@internal_api_call
+@provide_session
 def _record_task_map_for_downstreams(
-    *, task_instance: TaskInstance | TaskInstancePydantic, task: Operator, value: Any, session: Session
+    *,
+    task_instance: TaskInstance | TaskInstancePydantic,
+    task: Operator,
+    dag: DAG,
+    value: Any,
+    session: Session,
 ) -> None:
     """
     Record the task map for downstream tasks.
 
     :param task_instance: the task instance
     :param task: The task object
+    :param dag: the dag associated with the task
     :param value: The value
     :param session: SQLAlchemy ORM Session
 
     :meta private:
     """
+    # when taking task over RPC, we need to add the dag back
+    if isinstance(task, MappedOperator):
+        if not task.dag:
+            task.dag = dag
+    elif not task._dag:
+        task._dag = dag
+
     if next(task.iter_mapped_dependants(), None) is None:  # No mapped dependants, no need to validate.
         return
     # TODO: We don't push TaskMap for mapped task instances because it's not
-    # currently possible for a downstream to depend on one individual mapped
-    # task instance. This will change when we implement task mapping inside
-    # a mapped task group, and we'll need to further analyze the case.
+    #  currently possible for a downstream to depend on one individual mapped
+    #  task instance. This will change when we implement task mapping inside
+    #  a mapped task group, and we'll need to further analyze the case.

Review Comment:
   no -- adding the indent will make it so all of this text is "part of" the todo (i.e. all show yellow in IDE)  if we don't do this then it looks like separate comment... just a driveby "fix" but i can remove if you like



-- 
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