kaxil commented on code in PR #74345:
URL: https://github.com/apache/airflow/pull/74345#discussion_r4222238009
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -1337,11 +1429,12 @@ async def get_previous_task_instance(
:param state: If provided, filters by TaskInstance state
"""
query = (
- select(TI)
+ select(TI, public_map_index_expression(TI))
+ .where(TI.working_set.is_(True))
.join(DR, (TI.dag_id == DR.dag_id) & (TI.run_id == DR.run_id))
.options(contains_eager(TI.dag_run).load_only(DR.logical_date))
- .where(TI.dag_id == dag_id, TI.task_id == task_id, TI.map_index ==
map_index)
- .order_by(DR.logical_date.desc())
+ .where(TI.dag_id == dag_id, TI.task_id == task_id,
*build_coordinate_filters(TI, map_index=map_index))
+ .order_by(DR.logical_date.desc(), TI.region_index.desc())
Review Comment:
For a mapped task inside a loop, every pass has the same logical date and
the map position lives in `region_index`, so pass 0 index 0 and pass 1 index 0
tie on both sort keys and `LIMIT 1` returns whichever row the database hands
back. It only works for unmapped tasks because `region_index` is the iteration
there. I think this needs to order by the resolved loop pass, the same way
`/count` and `/states` pick the latest one.
##########
providers/edge3/src/airflow/providers/edge3/worker_api/routes/logs.py:
##########
@@ -41,6 +44,21 @@
@provide_session
def _logfile_path(task: TaskInstanceKey, *, session=NEW_SESSION) -> str:
"""Elaborate the (relative) path and filename to expect from task
execution."""
+ ti_id = session.scalar(
Review Comment:
This still keys on the five-field `TaskInstanceKey`. For a mapped task
inside a loop, pass 0 index 0 and pass 1 index 0 store the same `map_index` and
both start at try 1, so both `EdgeJobModel` rows match this filter (no `ORDER
BY`, so either can come back) and `@cache` keeps the first answer for the life
of the API server. Pass 1's log chunks then land in pass 0's file. The worker
can only send those five fields today (`logs_logfile_path` in
`cli/api_client.py`), so I think it needs to send the TI id and key the cache
on that.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -1257,10 +1302,45 @@ async def get_previous_successful_dagrun(
return PrevSuccessfulDagRunResponse.model_validate(dag_run)
+def _find_superseded_ids(rows: Iterable[tuple[TI, int]], resolver:
TaskCoordinateResolver) -> set[UUID]:
+ """
+ Find live rows that a later loop pass of the same task slot supersedes.
+
+ A task inside a loop keeps every pass live, so callers that name a task
without a pass get its latest
+ one. Two live rows that are not told apart by their loop pass leave the
slot ambiguous.
+ """
+ slots: dict[tuple[str, str, int], dict[UUID, TI]] = defaultdict(dict)
Review Comment:
Grouping by public slot picks the latest pass per slot, not per task. If
pass 0 maps over two items and pass 1 over one (the loop docs say the next
iteration can map a different collection), pass 0's index 1 is alone in its
group and stays live, so `/states` reports `task_1` from pass 0 next to
`task_0` from pass 1, and `/count` counts both. Could this find the latest pass
per `(run_id, task_id)` first and then keep only that pass's slots? That would
also let `get_loop_pass` load the region ancestry once per run instead of one
`load_region_ancestry` query per row.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -532,27 +655,41 @@ def delete_xcom(
token=CurrentTIToken,
):
"""Delete a single XCom Value."""
- owner = _find_writer_id(
- token.id, dag_id=dag_id, run_id=run_id, task_id=task_id,
map_index=map_index, session=session
+ owner = _get_writer_id(
Review Comment:
`DELETE /xcoms` used to be a 200 no-op when no task instance matched: 3.3.2
deletes by coordinates, and the base of this PR skips when `_find_writer_id`
returns `None`. `_get_writer_id` now raises 404, and `client.xcoms.delete`
doesn't catch errors, so `XCom.delete(...)` for a task removed from the run, or
for a mapped task with the default `map_index`, now fails the calling task.
Keeping `_find_writer_id` here with `if owner is not None` would keep the 409
for ambiguous passes and the old behaviour otherwise.
##########
airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py:
##########
@@ -295,6 +298,8 @@ class TaskInstance(BaseModel):
try_number: int
dag_version_id: uuid.UUID
map_index: int = -1
+ region_id: uuid.UUID | None = None
Review Comment:
`/previous` and breadcrumbs now map the sentinel to `None`, but this field
still carries the zero UUID for every ordinary task: `task_instance_to_runtime`
validates straight from the ORM row, and
`test_execute_task_carries_persisted_legacy_coordinates` asserts `UUID(int=0)`.
The Task SDK maps `None` back to the sentinel in `model_post_init`, so it
doesn't mind, but a Go or TS client reading this spec still sees two spellings
of "not in a region". Could `task_instance_to_runtime` apply
`get_public_region` as well?
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -1230,6 +1237,14 @@ def _get_sentry_integration(executor: BaseExecutor) ->
str:
return ""
return sentry_integration
+ log_contexts = prepare_task_log_contexts(
Review Comment:
This runs inside the critical section, after the pool rows are locked. For
any batch with regional TIs, `prepare_task_log_contexts` loads the
`SerializedDagModel` rows and calls `row.dag`, which deserializes on every
access instead of going through `self.scheduler_dag_bag`. Could it take a dag
bag so the scheduler and triggerer pass their cached ones? `dagrun.py:1538` and
`trigger.py:711` build a fresh `DBDagBag()` per call the same way, and
`trigger.py` deserializes the same Dag twice on the FAILED path. I haven't
measured the cost.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -1273,49 +1353,61 @@ def get_task_instance_count(
run_ids: Annotated[list[str] | None, Query()] = None,
states: Annotated[list[str] | None, Query()] = None,
) -> int:
- """Get the count of task instances matching the given criteria."""
- query = select(func.count()).select_from(TI).where(TI.dag_id == dag_id)
+ """Get the count of task instances matching the given criteria, counting a
looped task's latest pass."""
+ conditions = [TI.dag_id == dag_id, *build_coordinate_filters(TI,
map_index=map_index)]
if task_ids:
- query = query.where(TI.task_id.in_(task_ids))
-
- if map_index is not None:
- query = query.where(TI.map_index == map_index)
+ conditions.append(TI.task_id.in_(task_ids))
if logical_dates:
- query = query.where(TI.logical_date.in_(logical_dates))
+ conditions.append(TI.logical_date.in_(logical_dates))
if run_ids:
- query = query.where(TI.run_id.in_(run_ids))
+ conditions.append(TI.run_id.in_(run_ids))
if task_group_id:
group_tasks = _get_group_tasks(
dag_id, task_group_id, session, dag_bag, logical_dates, run_ids,
map_index
)
- # Get unique (task_id, map_index) pairs
- task_map_pairs = [(ti.task_id, ti.map_index) for ti in group_tasks]
-
- if not task_map_pairs:
+ if not group_tasks:
# If no task group tasks found, default to checking the task group
ID itself
# This matches the behavior in _get_external_task_group_task_ids
- task_map_pairs = [(task_group_id, -1)]
+ conditions.extend([TI.task_id == task_group_id,
public_map_index_expression(TI) == -1])
+ else:
+ conditions.append(TI.id.in_(ti.id for ti in group_tasks))
- # Update query to use task_id, map_index pairs
- query = query.where(tuple_(TI.task_id,
TI.map_index).in_(task_map_pairs))
+ slot = public_map_index_expression(TI)
+ slots = (
+ select(TI.run_id.label("run_id"), TI.task_id.label("task_id"),
slot.label("slot"))
+ .where(*conditions)
+ .subquery()
+ )
+ shared_slots = session.execute(
Review Comment:
This adds a `GROUP BY ... HAVING` over the `CASE`/`EXISTS` slot expression
to every `/count` call, including Dags with no loops, and `/count` is what
`task_mapping.py:84` and `ExternalTaskSensor` call. Rows in the sentinel region
can't share a slot because of `task_instance_current_key`, so could the slots
subquery filter on `TI.region_id != SENTINEL_REGION_ID` and skip this for
ordinary tasks?
##########
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/template_rendering.py:
##########
@@ -62,14 +70,19 @@ def render_k8s_pod_yaml(task_instance: TaskInstance) ->
dict | None:
# The pod spec using the default template.
pod_template_file = kube_config.pod_template_file
- # Generate command args using shared utility function
- command_args = generate_pod_command_args(task_instance)
+ if AIRFLOW_V_3_4_PLUS:
Review Comment:
I think this branch is a no-op. `generate_pod_command_args` already builds
`ExecuteTask.make(task_instance)` on 3.x, and `workload.key.map_index` comes
from `region_index`, which is a synonym for the ORM `map_index`, so it always
equals `task_instance.map_index`. Could the branch and the conditional import
go?
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -657,7 +668,9 @@ def _emit_task_span(ti, state):
"airflow.task_id": ti.task_id,
"airflow.dag_run.run_id": ti.run_id,
"airflow.task_instance.try_number": ti.try_number,
- "airflow.task_instance.map_index": ti.map_index if
ti.map_index is not None else -1,
+ "airflow.task_instance.map_index": map_index,
+ "airflow.task_instance.region_id": str(ti.region_id),
Review Comment:
`str(ti.region_id)` exports the zero UUID to OTel for every ordinary task.
Could this go through `get_public_region` and set the two region attributes
only when the region isn't `None`?
--
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]
