kaxil commented on code in PR #74349:
URL: https://github.com/apache/airflow/pull/74349#discussion_r4222295977
##########
airflow-core/src/airflow/serialization/definitions/dag.py:
##########
@@ -1290,11 +1299,23 @@ def clear(
if count == 0:
return 0
- clear_task_instances(
- list(tis),
- session,
+ clear_task_instances_for_runs(
Review Comment:
The dry run above returns `tis_result`, but once a gate is among `tis` this
archives more than that: `clear_loop_task_instances` archives every live TI at
a later loop position, including ones `exclude_task_ids` left out, and the
returned count is still `len(tis)`. `airflow tasks clear` builds its
confirmation prompt from the dry run, so users confirm a list that doesn't show
the later passes that get archived. Could the dry run include the archive set
from `select_loop_clear_scope`, and the scope respect `exclude_task_ids`?
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -424,10 +430,50 @@ def clear_task_instances(
:param run_on_latest_version: whether to run on latest serialized DAG and
Bundle version.
A run with no version of its own uses the latest either way, since
there is nothing
else for it to run on; a task instance with no version joins its run's.
+ :param whole_task_keys: ``(dag_id, run_id, task_id)`` of tasks cleared as
a whole. A pre-region
+ mapped expansion of such a task is archived and regenerated in a new
region.
+ :param superseded_ti_ids: executions whose generated work is replaced
rather than retried; they
+ are archived in place instead of receiving a successor.
:meta private:
"""
from airflow.exceptions import AirflowClearRunningTaskException
+ from airflow.models.dagrun import DagRun
+
+ if tis:
+ run_keys = {(ti.dag_id, ti.run_id) for ti in tis}
+ session.scalars(
+ select(DagRun.id)
+ .where(tuple_(DagRun.dag_id, DagRun.run_id).in_(run_keys))
+ .order_by(DagRun.dag_id, DagRun.run_id)
+ .with_for_update()
+ ).all()
+
+ legacy_keys = {
+ (ti.dag_id, ti.run_id, ti.task_id)
+ for ti in tis
+ if ti.region_id == SENTINEL_REGION_ID
+ and ti.region_index >= 0
+ and (ti.dag_id, ti.run_id, ti.task_id) in whole_task_keys
+ }
+ if legacy_keys:
+ existing_keys = {
+ (region.dag_id, region.run_id, region.node_id)
+ for region in session.scalars(
+ select(DynamicRegion).where(
+ tuple_(DynamicRegion.dag_id, DynamicRegion.run_id,
DynamicRegion.node_id).in_(
+ legacy_keys
+ ),
+ DynamicRegion.parent_region_id.is_(None),
+ DynamicRegion.forked_from_region_id.is_(None),
+ )
+ )
+ }
+ session.add_all(
Review Comment:
When the legacy expansion is regenerated in this new region,
`expand_mapped_task` seeds try numbers from
`get_last_try_numbers(region_id=<new region>)`, which never sees the archived
sentinel rows, so the new tries start at 1 again. A top-level map region has an
empty log position and keeps `map_index=region_index`, so the new attempt 1
writes to the same `map_index=i/attempt=1.log` as the archived attempt 1.
That's an ordinary whole-task clear of a mapped task created before the
upgrade, no loops involved. Could the seed also take the max try from the
sentinel region for the same task?
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1901,6 +1903,100 @@ def _emit_duration_stats_for_finished_state(self):
tags={**self.stats_tags, "dag_id": self.dag_id},
)
+ def _reconcile_legacy_expansions(
+ self, *, session: Session, dag_version_id: UUID | None = None
+ ) -> list[TI]:
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+
+ if not self.get_dag().has_dynamic_nodes:
+ return []
+
+ pending = (
+ select(DynamicRegion)
+ .where(
+ DynamicRegion.dag_id == self.dag_id,
+ DynamicRegion.run_id == self.run_id,
+ DynamicRegion.parent_region_id.is_(None),
+ DynamicRegion.forked_from_region_id.is_(None),
+ ~select(TI.id)
+ .where(
+ TI.working_set.is_(True),
+ TI.dag_id == DynamicRegion.dag_id,
+ TI.run_id == DynamicRegion.run_id,
+ TI.region_id == DynamicRegion.id,
+ )
+ .exists(),
+ ~select(TI.id)
+ .where(
+ TI.working_set.is_(True),
+ TI.dag_id == DynamicRegion.dag_id,
+ TI.run_id == DynamicRegion.run_id,
+ TI.task_id == DynamicRegion.node_id,
+ TI.region_id == SENTINEL_REGION_ID,
+ )
+ .exists(),
+ select(TI.id)
+ .where(
+ TI.working_set.is_(None),
+ TI.dag_id == DynamicRegion.dag_id,
+ TI.run_id == DynamicRegion.run_id,
+ TI.task_id == DynamicRegion.node_id,
+ TI.region_id == SENTINEL_REGION_ID,
+ TI.archived_reason == "superseded",
+ )
+ .exists(),
+ )
+ .execution_options(include_all_attempts=True)
+ )
+ regions = session.scalars(pending).all()
+ if not regions:
+ return []
+ session.execute(
+ select(DagRun)
+ .where(DagRun.id == self.id)
+ .with_for_update()
+ .execution_options(populate_existing=True)
Review Comment:
Sessions are `autoflush=False`, so when `clear_task_instances` calls
`dr.verify_integrity` on the latest-version path, the new `clear_number`,
`queued_at`, `state` and `start_date` on that run are still unflushed, and this
`populate_existing` lock reloads them from the row. On a whole-task clear of a
pre-region mapped task this query does match, because the clear has already
flushed the new top-level region and archived the sentinel rows. So a finished
run cleared with `run_on_latest_version` (or one with no
`created_dag_version_id`, which always takes that path) stays SUCCESS or FAILED
with cleared TIs that the scheduler never picks up. Locking with
`select(DagRun.id)` like the other lock sites, or flushing before the lock,
would avoid it.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py:
##########
@@ -963,42 +994,162 @@ def _collect_relatives(run_id: str, direction:
Literal["upstream", "downstream"]
*((t, m) for t, m in mapped_tasks_tuples if t not in
normal_task_ids),
]
- task_instances: Sequence[TI]
- if dag_run_id is not None and not (past or future):
- # Use run_id-based clearing when we have a specific dag_run_id and not
using past/future
- task_instances = dag.clear(
- dry_run=True,
- task_ids=task_markers_to_clear,
- run_id=dag_run_id,
- session=session,
- run_on_latest_version=resolved_run_on_latest,
- only_failed=body.only_failed,
- only_running=body.only_running,
- )
- else:
- # Use date-based clearing when no dag_run_id or when past/future is
specified
- task_instances = dag.clear(
+ clear_markers = task_markers_to_clear
+ if loop_aware and task_markers_to_clear is not None:
+ clear_markers = [marker if isinstance(marker, str) else marker[0] for
marker in task_markers_to_clear]
+
+ def select_candidates() -> Sequence[TI]:
+ if body.task_instance_ids is not None:
+ return session.scalars(
+ select(TI).where(
+ TI.dag_id == dag_id,
+ TI.run_id == dag_run_id,
+ TI.working_set.is_(True),
+ TI.id.in_(body.task_instance_ids),
+ )
+ ).all()
+ if dag_run_id is not None and not (past or future):
+ return dag.clear(
+ dry_run=True,
+ task_ids=clear_markers,
+ run_id=dag_run_id,
+ session=session,
+ run_on_latest_version=resolved_run_on_latest,
+ only_failed=body.only_failed and not loop_aware,
+ only_running=body.only_running and not loop_aware,
+ )
+ return dag.clear(
dry_run=True,
- task_ids=task_markers_to_clear,
+ task_ids=clear_markers,
start_date=body.start_date,
end_date=body.end_date,
session=session,
run_on_latest_version=resolved_run_on_latest,
- only_failed=body.only_failed,
- only_running=body.only_running,
+ only_failed=body.only_failed and not loop_aware,
+ only_running=body.only_running and not loop_aware,
)
- if not dry_run:
+ task_instances = select_candidates()
+ if body.task_instance_ids is not None and {ti.id for ti in task_instances}
!= set(body.task_instance_ids):
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND, "Selected task execution is not current
in this DAG run"
+ )
+
+ if task_instances:
+ if not dry_run:
+ run_keys = {(ti.dag_id, ti.run_id) for ti in task_instances}
+ session.scalars(
+ select(DagRun)
+ .where(tuple_(DagRun.dag_id, DagRun.run_id).in_(run_keys))
+ .order_by(DagRun.dag_id, DagRun.run_id)
+ .with_for_update()
+ .execution_options(populate_existing=True)
+ ).all()
+ if body.task_instance_ids is None:
+ task_instances = [ti for ti in select_candidates() if
(ti.dag_id, ti.run_id) in run_keys]
+ task_instances = session.scalars(
+ select(TI)
+ .where(TI.id.in_([ti.id for ti in task_instances]))
+ .execution_options(populate_existing=True)
+ ).all()
+
+ scopes: list[LoopClearScope] = []
+ if loop_aware and task_markers_to_clear is not None:
Review Comment:
`dag.clear` adds the setup and teardown tasks of the selection through
`partial_subset`, but this filter keeps only TIs whose task_id was in the
request, and `select_loop_clear_scope` doesn't add them back. Since
`loop_aware` is true for any Dag that contains a loop, clearing an ordinary
task with a setup/teardown pair in such a Dag now re-runs it without its setup
and teardown. Could the filter keep the ids `dag.clear` added, or the scope
pull in the relevant setup/teardown?
##########
airflow-ctl/src/airflowctl/api/datamodels/generated.py:
##########
@@ -291,6 +302,27 @@ class ClearTaskInstancesBody(BaseModel):
title="Task Group Id",
),
] = None
+ task_instance_ids: Annotated[
Review Comment:
`airflowctl tasks clear --task-instance-ids` comes through as a plain
string: the annotation is the `TaskInstanceIds` root model, which
`_python_type_from_string` maps to `str`, and `_apply_datamodel_defaults` only
splits `task_ids`, so pydantic rejects it with "Input should be a valid list".
`--whole-expansion-ids` has to be a subset of it, so neither is usable from the
CLI. Could both get the same comma/JSON handling as `task_ids`?
##########
airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py:
##########
@@ -57,6 +57,8 @@ class TaskInstanceResponse(BaseModel):
dag_id: str
run_id: str = Field(alias="dag_run_id")
map_index: int
+ region_id: UUID
Review Comment:
With these required, every ordinary TI in the REST API reports the zero UUID
as its `region_id`, while the execution API maps the sentinel to `None` through
`get_public_region`. Could these be `UUID | None` / `int | None` and go through
`get_public_region` before the generated clients ship? Otherwise every client
has to know that the zero UUID means "no region".
##########
airflow-core/src/airflow/ti_deps/deps/loop_archival_dep.py:
##########
@@ -0,0 +1,30 @@
+# 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.
+from __future__ import annotations
+
+from airflow.ti_deps.deps.base_ti_dep import BaseTIDep
+
+
+class LoopArchivalDep(BaseTIDep):
+ """A rerun gate waits for the generated work it replaces to terminate."""
+
+ NAME = "Loop archival"
+ IS_TASK_DEP = True
+
+ def _get_dep_statuses(self, ti, dep_context, session):
+ if ti.waits_for_loop_archival(session=session):
+ yield self._failing_status(reason="Superseded loop executions are
still terminating")
Review Comment:
`_are_premature_tis` still sets only the retry and reschedule flags, so
`update_state` can't tell this wait apart from a dependency that will never be
met. With a body like `prepare >> process` and `process[3]` running, clearing
the gate of iteration 1 can still fail the run with `all_tasks_deadlocked`
before the worker acks the restart.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -2794,10 +2907,18 @@ def _flush_ti_buffer(*, drain: bool = False) -> int:
while len(ti_carry) >= _TI_CHUNK_SIZE:
slice_tis = ti_carry[:_TI_CHUNK_SIZE]
Review Comment:
These slices still cut across runs, and the buffer query has no ordering. A
gate of one run in this slice archives that run's later passes that are still
waiting in `ti_carry`, and the next slice then fails with "An archived task
instance cannot be cleared" and rolls back the whole partition clear. Flushing
whole runs per call would avoid it.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py:
##########
@@ -896,7 +913,7 @@ def post_clear_task_instances(
error_message = f"Dag Run id {dag_run_id} not found in dag
{dag_id}"
raise HTTPException(status.HTTP_404_NOT_FOUND, error_message)
# Get the specific dag version:
- dag = get_dag_for_run(dag_bag, dag_run, session)
+ dag = get_dag_for_run_or_latest_version(dag_bag, dag_run, dag_id,
session)
Review Comment:
This still picks the run's creation version, where `get_dag_for_run` used to
fall back to the latest one for unversioned bundles. On the default dags folder
the scheduler moves a running run's unfinished TIs to the latest version and
`verify_integrity` adds new tasks, so after `a >> b` becomes `a >> b >> c`
mid-run, clearing `b` with downstream misses `c`, and a newly added
`task_group_id` returns 404.
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py:
##########
@@ -2255,6 +2255,40 @@ def test_clear_dag_run_dry_run(self, test_client,
session, body, dag_run_id, exp
)
assert logs == 0
+ @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
+ def test_clear_dag_run_dry_run_reports_loop_coordinates(self, test_client,
dag_maker, session):
+ @task_group
+ def body():
+ EmptyOperator(task_id="work")
+
+ with dag_maker("dry_run_loop", serialized=True):
+ loop = create_loop(body, max_iterations=3)
+ run = dag_maker.create_dagrun(run_id="dry_run_loop_run")
+ region =
session.scalar(select(DynamicRegion).where(DynamicRegion.node_id ==
loop.group_id))
+ later = TaskInstance(
+ task=dag_maker.dag.get_task("body.work"),
+ run_id=run.run_id,
+ dag_version_id=run.created_dag_version_id,
+ region_id=region.id,
+ region_index=1,
+ )
+ later.state = State.SUCCESS
+ session.add(later)
+ session.commit()
+
+ response = test_client.post(
+ f"/dags/{run.dag_id}/dagRuns/{run.run_id}/clear", json={"dry_run":
True, "only_failed": False}
+ )
+
+ assert response.status_code == 200, response.text
+ work = [row for row in response.json()["task_instances"] if
row["task_id"] == "body.work"]
+ assert {row["region_index"] for row in work} == {0, 1}
+ assert {row["map_index"] for row in work} == {-1}
+ assert all(
+ row["loop_iterations"] == [{"loop_id": loop.group_id, "iteration":
row["region_index"]}]
Review Comment:
`TaskInstanceResponse` still has no `loop_iterations` field at this commit,
so this raises `KeyError`. It can move to the PR that adds the field.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -545,13 +806,92 @@ def handle_bulk_create(
}
)
+ def _handle_regional_bulk(
+ self, action: MutationAction, results: BulkActionResponse
+ ) -> tuple[MutationAction, set[tuple[str, str, str, int]], set[tuple[str,
str, str]]]:
+ regional = [
+ entity
+ for entity in action.entities
+ if isinstance(entity, BulkTaskInstanceBody)
+ and (entity.region_id is not None or entity.region_index is not
None)
+ ]
+ deleting = isinstance(action, BulkDeleteAction)
+ specific, whole = self._categorize_entities(
+ action.entities, results, method="DELETE" if deleting else "PUT",
action_name=action.action.value
+ )
+ keys = {key[:3] for key in specific} | whole
+ run_keys = {key[:2] for key in keys}
+ for entity in action.entities:
+ if isinstance(entity, BulkTaskInstanceBody) and
(entity.include_future or entity.include_past):
+ dag_id, run_id, task_id, _ =
self._extract_task_identifiers(entity)
+ if (dag_id, run_id, task_id) in keys:
+ dag = get_latest_version_of_dag(self.dag_bag, dag_id,
self.session)
+ run_keys.update(
+ (dag_id, selected_run)
+ for selected_run in get_run_ids(
+ dag, run_id, entity.include_future,
entity.include_past, session=self.session
+ )
+ )
+ self.session.scalars(
+ select(DagRun)
+ .where(
+ tuple_(DagRun.dag_id, DagRun.run_id).in_(run_keys),
+ )
+ .order_by(DagRun.dag_id, DagRun.run_id)
+ .with_for_update()
+ ).all()
+ for entity in regional:
+ dag_id, run_id, task_id, map_index =
self._extract_task_identifiers(entity)
+ if (dag_id, run_id, task_id) not in keys:
+ continue
+ try:
+ dag, tis, data = _patch_ti_validate_request(
+ dag_id,
+ run_id,
+ task_id,
+ self.dag_bag,
+ entity,
+ self.session,
+ map_index,
+ getattr(action, "update_mask", None),
+ )
+ if deleting:
+ for ti in tis:
+ self.session.delete(ti)
Review Comment:
This still deletes only the live row, so the archived tries of a loop pass
(and the logs and XCom keyed by their attempt ids) stay behind, unlike the
non-regional path that goes through `TI.delete_attempts`.
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -570,11 +620,297 @@ def get_cached_latest_dag_version(dag_id: str) ->
DagVersion | None:
if dr.created_dag_version_id:
_pin_versionless_tis_to_run_version(dr,
dr.created_dag_version_id, session)
for ti in tis:
- ti.context_carrier = new_task_run_carrier(ti.dag_run.context_carrier)
+ if ti.working_set is True:
+ ti.context_carrier =
new_task_run_carrier(ti.dag_run.context_carrier)
session.flush()
return cleared
[email protected](frozen=True)
+class LoopClearScope:
+ """Live executions to retry and generated executions to archive."""
+
+ retry_ids: frozenset[UUID]
+ archive_ids: frozenset[UUID]
+
+
+def select_loop_clear_scope(
+ selected: Collection[TaskInstance],
+ *,
+ whole_expansion_ids: Collection[UUID] = (),
+ upstream: bool = False,
+ downstream: bool = True,
+ later_loop_iterations: bool = True,
+ session: Session,
+) -> LoopClearScope:
+ """Select loop clear executions; mutation callers must hold the DagRun
lock."""
+ from airflow.models.dagbag import DBDagBag
+ from airflow.models.task_coordinates import TaskCoordinateResolver,
enclosing_loop
+ from airflow.serialization.definitions.mappedoperator import
get_mapped_ti_count
+
+ if not selected:
+ return LoopClearScope(frozenset(), frozenset())
+ runs = {(ti.dag_id, ti.run_id) for ti in selected}
+ if len(runs) != 1:
+ raise ValueError("Loop clear selection must belong to one DagRun")
+ dag_id, run_id = runs.pop()
+ live = {
+ ti.id: ti
+ for ti in session.scalars(
+ select(TaskInstance)
+ .where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == run_id,
+ TaskInstance.working_set.is_(True),
+ )
+ .execution_options(populate_existing=True)
+ )
+ }
+ selected_ids = {ti.id for ti in selected}
+ if not selected_ids <= live.keys():
+ raise ValueError("Loop clear selection contains executions that are no
longer live")
+ if not set(whole_expansion_ids) <= selected_ids:
+ raise ValueError("Whole-task selection requires an explicitly selected
execution")
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+ regions = load_region_ancestry(
+ {ti.region_id for ti in live.values()}, dag_id=dag_id, run_id=run_id,
session=session
+ )
+ retry = {ti_id: live[ti_id] for ti_id in selected_ids}
+ for ti in tuple(retry.values()):
+ task = resolver.get_task(dag_id, run_id, ti.task_id,
dag_version_id=ti.dag_version_id)
+ group = enclosing_loop(task)
+ if group is not None and loop_position(regions, ti.region_id,
ti.region_index, group.node_id) is None:
+ raise ValueError("Loop clear selection requires an execution
inside its pinned loop")
+ if ti.id in whole_expansion_ids:
+ if not task.get_needs_expansion():
+ raise ValueError("Whole-expansion selection requires a mapped
task")
+ retry.update(
+ (other.id, other)
+ for other in resolver.resolve(dag_id=dag_id, run_id=run_id,
task_id=ti.task_id, caller=ti)
+ )
+ for is_upstream, enabled in ((True, upstream), (False, downstream)):
+ if not enabled:
+ continue
+ for ti in tuple(retry.values()):
+ task = resolver.get_task(dag_id, run_id, ti.task_id,
dag_version_id=ti.dag_version_id)
+ contexts = resolver.producer_contexts(ti)
+ count = (
+ get_mapped_ti_count(task, run_id, session=session,
producer_contexts=contexts)
+ if task.get_needs_expansion() and ti.region_index >= 0
+ else None
+ )
+ for relative in task.get_flat_relatives(upstream=is_upstream):
Review Comment:
This still walks relatives once per selected TI, with `producer_contexts`,
`get_mapped_ti_count` and a `resolve` each time, all under the DagRun lock. Any
Dag that contains a loop takes this path, so clearing a 1000-wide mapped task
with downstream there is thousands of queries returning the same rows.
Expanding relatives once per `(task_id, region_id, dag_version_id)` when the
task and relative share no mapped group would collapse most of it.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1468,6 +1468,8 @@ def recalculate(self) -> _UnfinishedStates:
@provide_session
def task_instance_scheduling_decisions(self, *, session: Session =
NEW_SESSION) -> TISchedulingDecision:
tis = self.get_task_instances(session=session, state=State.task_states)
+ reconciled = self._reconcile_legacy_expansions(session=session)
Review Comment:
The `has_dynamic_nodes` guard helps, but any Dag with a mapped task still
runs this three-subquery `dynamic_region` probe on every scheduling pass of
every run, and it only finds work after a whole-task clear of a pre-region
expansion. Could it run from `clear_task_instances` and `complete_restart`,
where that state is created, with `verify_integrity` as the backstop?
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -1543,10 +1543,36 @@ def process_executor_events(
return len(event_buffer)
# Check state of finished tasks
+ candidates = session.execute(
+ select(TI.id, TI.dag_id, TI.run_id).where(
+ TI.working_set.is_(True), TI.id.in_([key.id for key in
tis_with_right_state])
+ )
+ ).all()
+ run_query = (
+ select(DagRun.dag_id, DagRun.run_id)
+ .where(
+ tuple_(DagRun.dag_id, DagRun.run_id).in_(
+ {(dag_id, run_id) for _, dag_id, run_id in candidates}
+ )
+ )
+ .order_by(DagRun.dag_id, DagRun.run_id)
+ )
+ locked_runs = set(
+ session.execute(with_row_locks(run_query, of=DagRun,
session=session, skip_locked=True))
Review Comment:
This takes a `skip_locked` DagRun lock for every run with a pending executor
event, for every Dag, and puts events for runs it can't lock back in the
buffer. In HA, a run that another scheduler is examining
(`get_running_dag_runs_to_examine` holds the same row lock) gets its task
success/failure handling pushed to a later loop, and the heartbeat scan below
does the same. Should this be limited to runs with regional TIs, or is the
delay meant to apply to every Dag?
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py:
##########
@@ -197,33 +221,110 @@ def _patch_ti_validate_request(
session: SessionDep,
map_index: int | None = -1,
update_mask: list[str] | None = None,
+ *,
+ lock: bool = True,
) -> tuple[SerializedDAG, list[TI], dict]:
+ _validate_region_selection(body)
dag = get_latest_version_of_dag(dag_bag, dag_id, session)
- if not dag.has_task(task_id):
+ if lock:
+ _lock_patch_runs(dag, dag_run_id, body, session)
+ if body.region_id is None and not dag.has_task(task_id):
raise HTTPException(status.HTTP_404_NOT_FOUND, f"Task '{task_id}' not
found in Dag '{dag_id}'")
query = (
select(TI)
.where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id ==
task_id)
.options(joinedload(TI.rendered_task_instance_fields))
)
- if map_index is not None:
- query = query.where(TI.map_index == map_index)
- else:
- query = query.order_by(TI.map_index)
-
+ if body.region_id is None:
+ resolver = TaskCoordinateResolver(dag_bag, session)
+ for version in session.scalars(
Review Comment:
This still returns a null `dag_version_id` for TIs from runs that predate
Dag versioning, and `get_task(..., dag_version_id=None)` raises a bare
`ValueError("Pinned Dag ... not found")` when the run has no
`created_dag_version_id` either. Nothing on the PATCH or bulk path catches it,
so a note or state change on such a TI is a 500. Filtering with
`TI.dag_version_id.is_not(None)` here would skip them.
--
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]
