uranusjr commented on a change in pull request #21614:
URL: https://github.com/apache/airflow/pull/21614#discussion_r808589851
##########
File path: airflow/models/dagrun.py
##########
@@ -649,19 +651,41 @@ def task_instance_scheduling_decisions(self, session:
Session = NEW_SESSION) ->
def _get_ready_tis(
self,
- scheduleable_tasks: List[TI],
+ scheduleable_tis: List[TI],
finished_tis: List[TI],
session: Session,
) -> Tuple[List[TI], bool]:
old_states = {}
ready_tis: List[TI] = []
changed_tis = False
- if not scheduleable_tasks:
+ if not scheduleable_tis:
return ready_tis, changed_tis
+ # If we expand TIs, we need a new list so that we iterate over them
too. (We can't alter
+ # `scheduleable_tis` in place and have the `for` loop pick them up
+ expanded_tis: List[TI] = []
+
# Check dependencies
- for st in scheduleable_tasks:
+ for st in itertools.chain(scheduleable_tis, expanded_tis):
Review comment:
Maybe call this `sti` instead to signify it’s a ti, not a t(ask)? (or
`schedulable_ti`)
##########
File path: airflow/models/dagrun.py
##########
@@ -649,19 +651,41 @@ def task_instance_scheduling_decisions(self, session:
Session = NEW_SESSION) ->
def _get_ready_tis(
self,
- scheduleable_tasks: List[TI],
+ scheduleable_tis: List[TI],
finished_tis: List[TI],
session: Session,
) -> Tuple[List[TI], bool]:
old_states = {}
ready_tis: List[TI] = []
changed_tis = False
- if not scheduleable_tasks:
+ if not scheduleable_tis:
return ready_tis, changed_tis
+ # If we expand TIs, we need a new list so that we iterate over them
too. (We can't alter
+ # `scheduleable_tis` in place and have the `for` loop pick them up
+ expanded_tis: List[TI] = []
+
# Check dependencies
- for st in scheduleable_tasks:
+ for st in itertools.chain(scheduleable_tis, expanded_tis):
+
+ # Expansion of last resort! This is ideally handled in the
mini-scheduler in LocalTaskJob, but if
+ # for any reason it wasn't, we need to expand it now
+ if st.map_index < 0 and st.task.is_mapped:
+ # HACK. This needs a better way, one that copes with multiple
upstreams!
+ for ti in finished_tis:
+ if st.task_id in ti.task.downstream_task_ids:
+ upstream = ti
+
+ assert isinstance(st.task, MappedOperator)
+ new_tis = st.task.expand_mapped_task(upstream,
session=session)
+ assert new_tis[0] is st
+ # Add the new TIs to the list to be checked
+ for new_ti in new_tis[1:]:
+ new_ti.task = st.task
Review comment:
Perhaps `expand_mapped_task` should take care of assigning `TI.task`?
--
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]
