kacpermuda commented on code in PR #75132:
URL: https://github.com/apache/airflow/pull/75132#discussion_r4335165814
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -371,19 +371,28 @@ def ti_run(
should_retry=_is_eligible_to_retry(previous_state, ti.try_number,
ti.max_tries),
multi_team=conf.getboolean("core", "multi_team"),
)
- resolver = TaskCoordinateResolver(dag_bag, session)
- if loop_context := resolver.loop_context(ti):
- group, index = loop_context
- terminal = resolver.get_task(
- ti.dag_id, ti.run_id, group.terminal_task_id,
dag_version_id=ti.dag_version_id
- )
- context.loop = LoopContext(
- node_id=group.node_id,
- index=index,
- max_iterations=group.max_iterations,
- terminal_task_id=group.terminal_task_id,
- terminal_is_mapped=terminal.get_needs_expansion(),
- )
+ if ti.region_id != SENTINEL_REGION_ID:
+ resolver = TaskCoordinateResolver(dag_bag, session)
+ task = resolver.get_task(ti.dag_id, ti.run_id, ti.task_id,
dag_version_id=ti.dag_version_id)
+ if (expanded_ti_count := resolver.get_expanded_ti_count(ti,
task=task)) is not None:
+ context.expanded_ti_count = expanded_ti_count
+ if (loop := resolver.get_loop_position(ti, task=task)) is not None:
+ terminal = resolver.get_task(
+ ti.dag_id, ti.run_id, loop.group.terminal_task_id,
dag_version_id=ti.dag_version_id
+ )
+ context.loop = LoopContext(
+ node_id=loop.group.node_id,
+ index=loop.index,
+ max_iterations=loop.group.max_iterations,
+ terminal_task_id=loop.group.terminal_task_id,
+ terminal_is_mapped=terminal.get_needs_expansion(),
+ region_id=loop.region_id,
+ trajectory=loop.trajectory,
+ )
+ # Set only for a fork, as the route leaves out unset fields
+ if loop.forked_from_region_id is not None:
+ context.loop.forked_from_region_id =
loop.forked_from_region_id
+ context.loop.resumes_from_index = loop.resumes_from_index
Review Comment:
Done: TaskCoordinateResolver.get_loop_context() now builds the LoopContext
the worker gets, and ti_run just assigns it. It also gives the OpenLineage
provider the exact same object for server-side events, so it replaces the
separate LoopPosition struct.
--
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]
