kaxil commented on code in PR #74354:
URL: https://github.com/apache/airflow/pull/74354#discussion_r4222355072
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py:
##########
@@ -411,10 +417,29 @@ def _build_ti_summaries(
dag_bag: DBDagBag,
) -> dict[str, Any] | None:
ti_details: dict[str, GridNodeAgg] = {}
- dag_version_id = None
+ task_instances = list(task_instances)
Review Comment:
`list(task_instances)` pulls every TI row of the run into memory for every
Dag, loop or not, which undoes the `yield_per=1000` streaming the caller still
sets at line 571. It's only needed for the first `dag_version_id` and the set
of region ids. Could the version come from a small query before the stream, the
regions from one `DynamicRegion.load_for_run` only when the Dag has loops, and
the aggregation stay a single streaming pass? `region_id` is a selected,
non-nullable column here too, so the `getattr(ti, "region_id", None)` checks at
434 and 455 can go.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -1538,12 +1547,28 @@ async def get_previous_task_instance(
if state:
query = query.where(TI.state == state)
- row = (await session.execute(query.limit(1))).first()
- if row is None:
+ first = session.execute(query.limit(1)).first()
+ if first is None:
return None
- ti, public_index = row
+ if first[0].region_id != SENTINEL_REGION_ID:
+ resolver = TaskCoordinateResolver(dag_bag, session)
+ requester = session.get(TI, token.id)
+ iterations = (
+ resolver.loop_iterations(requester)
+ if requester is not None and (requester.dag_id, requester.task_id)
== (dag_id, task_id)
+ else []
+ )
+ if iterations:
+ with contextlib.closing(
+
session.execute(query.limit(_MAX_PREVIOUS_TIS_SCANNED).execution_options(yield_per=50))
Review Comment:
Each scanned row still goes through `resolver.loop_iterations`, which calls
`load_region_ancestry` for every run it hasn't seen, on the same session while
this `yield_per=50` stream has unread rows. That's a query per earlier run
across up to 500 rows, and I don't think MySQL allows it: `yield_per` uses an
unbuffered server-side cursor there and mysqlclient rejects a new statement
until it's drained. Since the scan is capped anyway, could this `.all()` the
rows, call `resolver.prefetch_regions(...)` and then filter? Pushing
`region_index == iteration` into SQL when the task sits directly in the loop
region would also stop a loop with more than 500 earlier passes from returning
`None` when a match exists.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py:
##########
@@ -48,6 +49,18 @@ def _ui_colors(node) -> dict[str, str]:
}
+def loop_metadata(task_group) -> dict:
+ if not isinstance(task_group, SerializedLoopTaskGroup):
+ return {}
+ gate = task_group.dag.get_task(task_group.gate_task_id)
Review Comment:
This looks up the gate for every loop group, including fixed-count ones that
never use it. With `root=` both the structure route and the grid structure
route call `partial_subset`, and since the gate sits downstream of the body
terminal, filtering to a body task without `include_downstream` drops the gate
while the group survives. `get_task` then raises `TaskNotFound` and
`/ui/structure/structure_data` returns a 500 (traced, not run). Could this read
the gate only when `has_until`, via `task_dict.get(...)`, with a test that
filters a loop Dag by `root`?
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -1504,17 +1504,25 @@ def get_task_instance_count(
return
session.scalar(select(func.count()).select_from(TI).where(*conditions)) or 0
+_MAX_PREVIOUS_TIS_SCANNED = 500
+
+
@router.get("/previous/{dag_id}/{task_id}", status_code=status.HTTP_200_OK)
-async def get_previous_task_instance(
+def get_previous_task_instance(
dag_id: str,
task_id: str,
- session: AsyncSessionDep,
+ session: SessionDep,
+ dag_bag: DagBagDep,
logical_date: Annotated[UtcDateTime | None, Query()] = None,
map_index: Annotated[int, Query()] = -1,
state: Annotated[TaskInstanceState | None, Query()] = None,
+ token: TIToken = CurrentTIToken,
) -> PreviousTIResponse | None:
"""
- Get the previous task instance matching the given criteria, preferring a
looped task's latest pass.
+ Get the previous task instance matching the given criteria.
+
+ When the requesting task instance is of ``task_id``, the previous task
instance sits at the same
Review Comment:
This changes what `ti.get_previous_ti()` returns for a task inside a loop:
the same loop position in an earlier run, where it used to be that run's latest
pass. The commit message only talks about the UI and loops.rst doesn't mention
`get_previous_ti`, so a sentence in each would help users who rely on the old
result.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py:
##########
@@ -621,6 +619,9 @@ def get_task_instances(
"Pass an empty string for the first page, then use ``next_cursor``
from the response. "
"When ``cursor`` is provided, ``offset`` is ignored.",
),
+ loop_id: Annotated[str | None, Query()] = None,
+ iteration: Annotated[int | None, Query(ge=0)] = None,
+ loop_region_id: Annotated[UUID | None, Query()] = None,
Review Comment:
`loop_region_id` is still on the public list endpoint, and `loop_iterations`
is still a list on every TI response. Authoring rejects nested loops and loops
in mapped groups, so a run has at most one invocation and the list at most one
entry. Once a release ships these shapes they're hard to change, so I'd keep
`loop_id` + `iteration` here and a single `loop_iteration` on the response, and
add the plural and invocation forms together with nested-loop support.
##########
airflow-core/src/airflow/ui/src/pages/GroupTaskInstance/Header.tsx:
##########
@@ -20,19 +20,31 @@ import type { ReactNode } from "react";
import { useTranslation } from "react-i18next";
import { AiOutlineGroup } from "react-icons/ai";
+import { FiBookOpen } from "react-icons/fi";
+import { useParams } from "react-router-dom";
import type { LightGridTaskInstanceSummary } from "openapi/requests/types.gen";
import { ClearTaskInstanceButton } from "src/components/Clear";
+import DisplayMarkdownButton from "src/components/DisplayMarkdownButton";
import { HeaderCard } from "src/components/HeaderCard";
import { MarkTaskGroupAsButton } from "src/components/MarkAs";
import Time from "src/components/Time";
+import { useLoopGroupNode } from "src/queries/useIsLoopGroup";
+import { useLoopSummary } from "src/queries/useLoopSummary";
import { formatNumber, useDurationFormat } from "src/utils";
+import { useLoopOutcomeStats } from "./useLoopOutcomeStats";
+
export const Header = ({ taskInstance }: { readonly taskInstance:
LightGridTaskInstanceSummary }) => {
const { i18n, t: translate } = useTranslation();
const { formatElapsed } = useDurationFormat();
+ const { dagId = "", groupId = "", runId = "" } = useParams();
+ const outcomeStats = useLoopOutcomeStats({ dagId, groupId, runId });
+ const { data: loopSummary } = useLoopSummary({ dagId, groupId, runId });
+ const groupNode = useLoopGroupNode(groupId);
+ const docMd = loopSummary === undefined ? groupNode?.doc_md :
loopSummary.doc_md;
Review Comment:
`LoopSummaryResponse.doc_md` doesn't go through the `inspect.cleandoc`
validator `GridNodeResponse.doc_md` has (datamodels/ui/common.py:79-85), and
`@task_group` stores the raw docstring. So a multi-line indented docstring
renders fine from `groupNode` and then turns into an indented code block once
the summary loads. Could the header just read `groupNode?.doc_md` and the field
come off the summary?
##########
airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/grid.py:
##########
@@ -43,3 +46,75 @@ class GridTISummaries(BaseModel):
run_id: str
dag_id: str
task_instances: list[LightGridTaskInstanceSummary]
+
+
+class LoopCriteriaComparison(BaseModel):
+ """Optional comparison details reported by a loop condition."""
+
+ field: str
+ op: str
+ target: Any = None
+ actual: Any = None
+
+
+class LoopIterationSummary(BaseModel):
+ """Execution state for one existing iteration."""
+
+ index: int
+ state: TaskInstanceState | None = None
+ decision: Literal["stop", "continue"] | None = None
Review Comment:
Nothing sets `decision`, `is_tail`, `result` or `criteria`, and
`loop_run_summaries` only produces `cap_reached` / `iteration_failed`, never
`criteria_met` / `not_converged`. So the UI code reading them (the decision
labels and `finalCriteria` in loopUtils.ts, the "loop.target" row, the
`converged` stat in `LoopHistoryChart`) can't render, and the new
loopUtils.test.ts cases build summaries the API can't return.
`parent_iterations` is the same: with nested loops rejected, the ancestor walk
in `services/ui/grid.py` that fills it never appends. I'd drop these fields,
`LoopCriteriaComparison` and the two literals until something produces them.
##########
airflow-core/src/airflow/ui/src/pages/Task/Overview/LoopHistoryChart.tsx:
##########
@@ -0,0 +1,146 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import { Box, HStack, Heading, Stat, useToken } from "@chakra-ui/react";
+import { BarElement, CategoryScale, Chart as ChartJS, Legend, LinearScale,
Tooltip } from "chart.js";
+import annotationPlugin from "chartjs-plugin-annotation";
+import { Bar } from "react-chartjs-2";
+import { useTranslation } from "react-i18next";
+import { useNavigate, useParams } from "react-router-dom";
+
+import { SearchParamsKeys } from "src/constants/searchParams";
+import { useLoopHistory } from "src/queries/useLoopHistory";
+import { getComputedCSSVariableValue } from "src/theme";
+import { median } from "src/utils/median";
+
+ChartJS.register(CategoryScale, LinearScale, BarElement, Legend, Tooltip,
annotationPlugin);
+
+export const LoopHistoryChart = ({ groupId }: { readonly groupId: string }) =>
{
+ const { dagId = "" } = useParams();
+ const { t: translate } = useTranslation("dag");
+ const navigate = useNavigate();
+ const { data } = useLoopHistory({ dagId, groupId });
+ const [successColor, failedColor, mutedColor] = useToken("colors", [
+ "success.solid",
+ "failed.solid",
+ "fg.muted",
+ ]);
+
+ const runs = data?.runs ?? [];
+
+ if (runs.length === 0) {
+ return undefined;
+ }
+
+ const converged = runs.filter((run) => run.reason === "criteria_met");
+ const capHits = runs.filter((run) => run.iterations_ran ===
run.max_iterations);
Review Comment:
`iterations_ran === max_iterations` also counts a run that failed on its
last allowed iteration, or one still running it, since `iterations_ran` counts
any started iteration. `run.status === "ran_to_cap"` is what the backend
already computes for this.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py:
##########
@@ -52,9 +71,12 @@ def add_ti(
end_date: datetime | None,
dag_version_number: int | None,
has_note: bool = False,
+ loop_positions: dict[str, tuple[UUID, int]] | None = None,
) -> None:
"""Merge one task instance row into the summary."""
self.child_states[state] += 1
+ for loop_id, position in (loop_positions or {}).items():
+ self.loop_iterations.setdefault(loop_id, set()).add(position)
Review Comment:
This records a loop position for every TI whatever its state, so
`loop_iterations_count` (the grid tooltip's "Iterations run") counts an
iteration whose TIs exist but haven't started. `iterations_ran` in
`loop_run_summaries` needs a start date or a finished state, so while the next
iteration waits on a pool the tooltip says 2 and the header says 1/5. Could
both use the same rule?
##########
airflow-core/src/airflow/ui/src/pages/GroupTaskInstance/GroupTaskInstances.tsx:
##########
@@ -0,0 +1,35 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import { useParams } from "react-router-dom";
+
+import { useLoopSummary } from "src/queries/useLoopSummary";
+
+import { TaskInstances } from "../TaskInstances";
+import { IterationSelect } from "./LoopIterations";
+
+export const GroupTaskInstances = () => {
+ const { dagId = "", groupId = "", runId = "" } = useParams();
+ const { data: summary } = useLoopSummary({ dagId, groupId, runId });
+
+ return summary === undefined ? (
Review Comment:
`loopGroupId` is only passed once the summary has loaded, so every visit
first fetches and renders the unfiltered list, then refetches filtered. If the
summary request fails (`retry: false`, or a stale `loop_region_id` that 404s),
the page stays on all iterations while the URL says `?iteration=N` and no
selector shows. `IterationSelect` also still writes `loop_region_id` into the
URL on entry (IterationSelect.tsx:63-64), which changes the summary query key
and drops `data`, so this swaps back to the plain list and the outcome stats
flash until the second fetch lands. Could `loopGroupId` come from
`useIsLoopGroup(groupId)`, and the selector leave `loop_region_id` alone unless
the user picks one?
##########
airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py:
##########
@@ -205,3 +234,176 @@ def _find_aggregates(
summary,
)
return
+
+
+def loop_run_summaries(
+ run: DagRun,
+ group_id: str,
+ *,
+ session: Session,
+ dag_bag: DBDagBag,
+ loop_region_id: UUID | None = None,
+) -> list[LoopSummaryResponse]:
+ """Read each invocation of a loop separately, keyed by its fork family."""
+ live = (
+ TaskInstance.dag_id == run.dag_id,
+ TaskInstance.run_id == run.run_id,
+ TaskInstance.working_set.is_(True),
+ )
+ definitions: dict[UUID | None, SerializedLoopTaskGroup | None] = {}
+ for version_id in
session.scalars(select(TaskInstance.dag_version_id).where(*live).distinct()):
+ pinned_id = version_id or run.created_dag_version_id
+ if pinned_id is None:
+ raise HTTPException(404, "Pinned DAG definition not found")
+ dag = dag_bag.get_dag(pinned_id, session=session)
+ definition = dag.task_group.get_task_group_dict().get(group_id) if dag
else None
+ definitions[version_id] = definition if isinstance(definition,
SerializedLoopTaskGroup) else None
+ found = {
+ version_id: definition for version_id, definition in
definitions.items() if definition is not None
+ }
+ task_ids_by_version = {
+ version_id: {task.task_id for task in definition.iter_tasks()}
+ for version_id, definition in found.items()
+ }
+ members = []
+ if task_ids_by_version:
+ member_rows = session.execute(
+ select(
+ TaskInstance.task_id,
+ TaskInstance.region_id,
+ TaskInstance.region_index,
+ TaskInstance.state,
+ TaskInstance.start_date,
+ TaskInstance.end_date,
+ TaskInstance.dag_version_id,
+ )
+ .where(*live,
TaskInstance.task_id.in_(set().union(*task_ids_by_version.values())))
+ .order_by(TaskInstance.task_id, TaskInstance.region_id,
TaskInstance.region_index)
+ )
+ members = [ti for ti in member_rows if ti.task_id in
task_ids_by_version.get(ti.dag_version_id, ())]
+ loop_definitions = {ti.dag_version_id: found[ti.dag_version_id] for ti in
members}
+ group = next(iter(loop_definitions.values()), None)
+ if group is None:
+ if run.created_dag_version_id is None:
+ raise HTTPException(404, "Pinned DAG definition not found")
+ dag = dag_bag.get_dag(run.created_dag_version_id, session=session)
+ definition = dag.task_group.get_task_group_dict().get(group_id) if dag
else None
+ if definition is None:
+ raise HTTPException(404, "Task group not found")
+ if not isinstance(definition, SerializedLoopTaskGroup):
+ raise HTTPException(422, "Task group is not a loop")
+ group = definition
+ regions = load_region_ancestry(
+ {ti.region_id for ti in members}, dag_id=run.dag_id,
run_id=run.run_id, session=session
+ )
+ invocations: dict[UUID, dict[int, list[Any]]] = defaultdict(lambda:
defaultdict(list))
+ for ti in members:
+ if position := loop_position(regions, ti.region_id, ti.region_index,
group_id):
+ family, index = position
+ invocations[family][index].append(ti)
+ families = sorted(invocations, key=str)
+ if loop_region_id is not None and loop_region_id not in invocations:
+ raise HTTPException(404, "Loop invocation not found in this Dag run")
+ selected: list[UUID | None] = (
+ [loop_region_id] if loop_region_id is not None else [*families] if
families else [None]
+ )
+ region_options = []
+ for family in families:
+ parents = []
+ region = regions[family]
+ while region.parent_region_id is not None:
+ parent = regions[region.parent_region_id]
+ parent_group =
group.dag.task_group.get_task_group_dict().get(parent.node_id)
+ if isinstance(parent_group, SerializedLoopTaskGroup):
+ if TYPE_CHECKING:
+ assert region.parent_region_index is not None
+ parents.append(
+ LoopIterationResponse(loop_id=parent.node_id,
iteration=region.parent_region_index)
+ )
+ region = parent
+ region_options.append(
+ LoopInvocationResponse(region_id=family,
parent_iterations=list(reversed(parents)))
+ )
+ summaries = []
+ for selected_family in selected:
+ iterations = invocations[selected_family] if selected_family is not
None else {}
+ gate_rows = [
+ ti
+ for tasks in iterations.values()
+ for ti in tasks
+ if ti.task_id == loop_definitions[ti.dag_version_id].gate_task_id
+ ]
+ if gate_rows:
+ group = loop_definitions[max(gate_rows, key=lambda ti:
ti.region_index).dag_version_id]
+ gate = group.dag.get_task(group.gate_task_id)
+ rows = []
+ failed = None
+ reason_task = None
+ for index, tasks in sorted(iterations.items()):
+ states = [ti.state for ti in tasks]
+ starts = [ti.start_date for ti in tasks if ti.start_date is not
None]
+ ends = [ti.end_date for ti in tasks if ti.end_date is not None]
+ failures = [ti for ti in tasks if ti.state in State.failed_states]
Review Comment:
`State.failed_states` includes `UPSTREAM_FAILED`, so when something upstream
of the loop fails, the summary comes back `failed` with `iteration_failed`,
`failed_at_iteration=0` and `iterations_ran=1` (the predicate at 388-391 counts
it too). The header then says iteration 0 failed in a body task that never ran.
Should an iteration whose tasks only reached `UPSTREAM_FAILED` count as not
run, and the reason only name a task that actually failed?
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_coordinates.py:
##########
@@ -46,6 +58,13 @@ class TaskCoordinateView:
def __getattr__(self, name: str) -> Any:
if name == "map_index":
return self.resolver.public_map_index(self.value)
+ if name == "rendered_map_index" and
self.resolver.public_map_index(self.value) < 0:
Review Comment:
The response side looks right now, but the filter still disagrees with it:
`rendered_map_index_pattern` filters on the `TaskInstance.rendered_map_index`
hybrid expression, which falls back to `cast(region_index, String)`. So
`rendered_map_index_pattern=2` still matches iteration-2 loop rows that come
back with `rendered_map_index: null`.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py:
##########
@@ -555,3 +582,76 @@ def _generate() -> Generator[str, None, None]:
yield GridTISummaries.model_validate(summary).model_dump_json() +
"\n"
return StreamingResponse(content=_generate(),
media_type="application/x-ndjson")
+
+
+@grid_router.get(
+ "/loop/{dag_id}/{run_id}/{group_id}",
+ responses=create_openapi_http_exception_doc([404, 422]),
+ dependencies=[
+ Depends(requires_access_dag(method="GET",
access_entity=DagAccessEntity.TASK_INSTANCE)),
+ Depends(requires_access_dag(method="GET",
access_entity=DagAccessEntity.RUN)),
+ ],
+)
+def get_loop_summary(
+ dag_id: str,
+ run_id: str,
+ group_id: str,
+ dag_bag: DagBagDep,
+ session: SessionDep,
+ loop_region_id: UUID | None = None,
+) -> LoopSummaryResponse:
+ """Summarize one loop invocation from the run's live task instances."""
+ run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id,
DagRun.run_id == run_id))
+ if run is None:
+ raise HTTPException(404, "DAG run not found")
+ return loop_run_summaries(run, group_id, session=session, dag_bag=dag_bag,
loop_region_id=loop_region_id)[
+ 0
+ ]
+
+
+@grid_router.get(
+ "/loop-history/{dag_id}/{group_id}",
+ responses=create_openapi_http_exception_doc([404]),
+ dependencies=[
+ Depends(requires_access_dag(method="GET",
access_entity=DagAccessEntity.TASK_INSTANCE)),
+ Depends(requires_access_dag(method="GET",
access_entity=DagAccessEntity.RUN)),
+ ],
+)
+def get_loop_history(
+ dag_id: str,
+ group_id: str,
+ dag_bag: DagBagDep,
+ session: SessionDep,
+ limit: Annotated[int, Query(ge=1, le=100)] = 14,
+) -> LoopHistoryResponse:
+ """Show loop invocations from recent runs, oldest first."""
+ runs = session.scalars(
+ select(DagRun)
+ .where(DagRun.dag_id == dag_id)
+ .order_by(DagRun.run_after.desc(), DagRun.id.desc())
+ .limit(limit)
+ ).all()
+ history = []
+ for run in reversed(runs):
+ try:
+ summaries = loop_run_summaries(run, group_id, session=session,
dag_bag=dag_bag)
Review Comment:
This still runs `loop_run_summaries` once per run, each with a version
query, the member TI query and a region load, so the default `limit=14` is
around 40 queries and `limit=100` around 300, and every live loop TI of every
run comes back to Python to draw one bar. Could the member and region queries
take `run_id.in_(...)` once and group in Python?
##########
airflow-core/src/airflow/models/task_coordinates.py:
##########
@@ -147,6 +148,46 @@ def adopt_dag(self, dag: SerializedDAG | None) -> None:
if dag is not None and dag.dag_version_id is not None:
self._dags.setdefault(dag.dag_version_id, dag)
+ def prefetch_regions(self, tis: Iterable[TaskCoordinate]) -> None:
+ wanted: dict[tuple[str, str], set[UUID]] = {}
+ for ti in tis:
+ if ti.region_id != SENTINEL_REGION_ID and ti.region_id not in
self._regions.get(
Review Comment:
This still loads ancestry for every non-sentinel TI, and every mapped
expansion has a region, so a Dag with mapping and no loops pays a
`load_region_ancestry` per run on the Task Instances list and Gantt, and
`loop_iterations` then returns `[]` without reading it. Could it skip TIs whose
pinned task has no `enclosing_loop` (the Dags are already in `_dags`)?
--
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]
