aeroyorch commented on code in PR #65932:
URL: https://github.com/apache/airflow/pull/65932#discussion_r4018116580
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -2202,6 +2202,17 @@ def schedule_tis(
debug_try_number_check = self.log.isEnabledFor(logging.DEBUG)
expected_try_number_by_ti_id: dict[UUID, tuple[int, int, str | None]]
= {}
for ti in schedulable_tis:
+ if ti.state == TaskInstanceState.UP_FOR_RETRY:
+ if TYPE_CHECKING:
+ assert ti.task
+ # Weight strategies and the mutation hook must see the
upcoming try. The updates below
+ # still own the increment, so try_number is restored.
+ try_number = ti.try_number
+ ti.try_number = try_number + 1
Review Comment:
Thanks for digging into this! I agree that the intermediate (B, N) state
does not make much sense. `schedule_tis` only sets the new `try_number` in the
bulk UPDATE, which makes sense to avoid race conditionss between several
schedulers in HA. So removing this intermediate state probably needs a deeper
change.
In the meantime, while digging into the sqlalchemy docs, I found that
[set_committed_value](https://docs.sqlalchemy.org/en/20/orm/session_api.html#sqlalchemy.orm.attributes.set_committed_value)
could make the temporary `try_number` safe here as it sets the value without
marking it as changed, so a flush inside a hook or a strategy cannot write it
to the database.
I am not sure whether you would prefer to use it in this PR for now, or to
wait until the (B, N) state is removed.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
