Make _get_ti compatible with RPC #38570
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
potiuk
approved these changes
Mar 27, 2024
dirrao
approved these changes
Mar 28, 2024
eladkal
approved these changes
Mar 28, 2024
hussein-awala
approved these changes
Mar 28, 2024
uranusjr
reviewed
Mar 29, 2024
dstandish
force-pushed
the
remove-@provide_session-from-_get_ti
branch
from
4f0eff5
to
7010ca5
Compare
dirrao
suggested changes
Mar 30, 2024
dstandish
force-pushed
the
remove-@provide_session-from-_get_ti
branch
from
7010ca5
to
b317209
Compare
dstandish
force-pushed
the
remove-@provide_session-from-_get_ti
branch
2 times, most recently
from
c9809df
to
181c0fe
Compare
dstandish
changed the title
dstandish
commented
Apr 2, 2024
dstandish
commented
Apr 2, 2024
dstandish
force-pushed
the
remove-@provide_session-from-_get_ti
branch
from
181c0fe
to
942d63f
Compare
dstandish
force-pushed
the
remove-@provide_session-from-_get_ti
branch
2 times, most recently
from
f0b1cbd
to
22e9a5f
Compare
dstandish
force-pushed
the
remove-@provide_session-from-_get_ti
branch
from
d88aa3f
to
b6fbc74
Compare
uranusjr
reviewed
Apr 8, 2024
uranusjr
reviewed
Apr 8, 2024
vincbeck
approved these changes
Apr 8, 2024
dstandish
force-pushed
the
remove-@provide_session-from-_get_ti
branch
from
b6fbc74
to
285cb7c
Compare
potiuk
approved these changes
Apr 9, 2024
The only usages of the provided session in this function are calls to other funcs which are also decorated with provide_session, and which are already RPC calls. There does not seem to be a reason that they all need to be in the same session. So we should be able to simply remove this decorator, thus making this function AIP-44 compatible.
dstandish
force-pushed
the
remove-@provide_session-from-_get_ti
branch
from
285cb7c
to
0046643
Compare
utkarsharma2
pushed a commit
to astronomer/airflow
that referenced
this pull request
Apr 22, 2024
This is for AIP-44. I had to pull out the "db access" parts from `_get_ti` and move them to RPC function `_get_ti_db_access`. To make that work, I also had to ensure that "task" objects (a.k.a. instances of AbstractOperator) can properly be roundtripped with BaseSerialization.serialize. Up to now they could not be, and they were "manually" serialized as part of SerializedDAG. This changes a bit the way we serialize task objects and so we had to handle backcompat and update a fair amount of tests.
utkarsharma2
added
changelog:skip
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Ok, taking yet another approach with this.
It turns out there was a non-obvious 4th usage of session in this function, so I went back to trying to make it an RPC call.
The problem is, we don't have a TaskPydantic model, and when you serialize an operator object it doesn't really do it properly because it doesn't "encode" it i.e. it doesn't add the TYPE var so the receiver has no idea what the object is. This, I assume, is because historically an operator was always only serialized as part of another object.
What I did to "fix" this is I make base serialization serialize operator objects "properly", i.e. running through
_encode.This necessitated a change to the way SerDag deserializes its tasks.
The other wonky thing here is that we expect (somewhere in the execution process) the task object on a TI to also have a dag, which, happens to get lost in serialization (and yeah, would be a bit reduntant perhaps to include that when serializing DAG for example... think the DAG ser'd on each task....), so what I do is add it back after the TI comes back through the RPC.... It's a bit of a mess but that's kind of where we're at.
Please give it a look when you can @uranusjr @potiuk and let me know what you think of the approach. I suspect there will be a test failure or two that I have to correct with this change.
Orig:
The only usages of the provided session in this function are calls to other funcs which are also decorated with provide_session, and which are already RPC calls. There does not seem to be a reason that they all need to be in the same session. So we should be able to simply remove this decorator, thus making this function AIP-44 compatible.