kaxil commented on code in PR #73918:
URL: https://github.com/apache/airflow/pull/73918#discussion_r4139049357
##########
providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py:
##########
@@ -368,6 +380,16 @@ def _purge_jobs(self, session: Session) -> bool:
EdgeLogsModel.task_id == job.task_id,
EdgeLogsModel.map_index == job.map_index,
EdgeLogsModel.try_number == job.try_number,
+ ~select(EdgeJobModel.dag_id)
+ .where(
+ EdgeJobModel.dag_id == job.dag_id,
+ EdgeJobModel.run_id == job.run_id,
+ EdgeJobModel.task_id == job.task_id,
+ EdgeJobModel.map_index == job.map_index,
+ EdgeJobModel.try_number == job.try_number,
+ EdgeJobModel.task_instance_id !=
job.task_instance_id,
Review Comment:
If two jobs at the same coordinates are both purged in the same pass,
neither one deletes the logs. Airflow sessions use `autoflush=False`, so the
`session.delete(job)` at 375 is still pending when the next row's `NOT EXISTS`
runs. Each row sees the other as a live sibling, and after commit both jobs are
gone while their `edge_logs` rows are left with nothing that will ever purge
them. I reproduced it with this statement shape on SQLite (jobs 0, logs 1
left). Adding `session.flush()` after the delete fixed it. `revoke_task` has
the same gap: if A is purged while B exists and B is then revoked (stuck in
queued, say), B's delete at 448 never touches logs. Could the test next to
`test_purging_retired_job_keeps_logs_while_coordinate_sibling_exists` also
cover both siblings expiring together?
##########
providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py:
##########
@@ -368,6 +380,16 @@ def _purge_jobs(self, session: Session) -> bool:
EdgeLogsModel.task_id == job.task_id,
EdgeLogsModel.map_index == job.map_index,
EdgeLogsModel.try_number == job.try_number,
+ ~select(EdgeJobModel.dag_id)
Review Comment:
On Airflow 3.0.x this makes the `edge_logs` delete raise on every purge. 3.0
pins SQLAlchemy 1.4, where an ORM `delete()` defaults to
`synchronize_session="evaluate"`, and the evaluator can't handle the new
`~select(...).exists()` clause: `InvalidRequestError: Could not evaluate
current criteria in Python: "Cannot evaluate SelectStatementGrouping"`. The
subquery is unconditional, so it hits the coordinate-key path on older cores
too, and `sync()` fails on any heartbeat that purges a job. That's what
`TestEdgeExecutor::test_sync` / `test_sync_orphaned_tasks` show in the Compat
3.0.6 job (SQLAlchemy 1.4.54 there).
`.execution_options(synchronize_session=False)` on this statement should fix
it, since nothing in the session holds `EdgeLogsModel` rows that need syncing.
##########
providers/edge3/docs/edge_executor.rst:
##########
@@ -244,3 +244,27 @@ pipeline. The table below documents known compatibility
issues and workarounds:
**For Airflow 3.2 users:** If upgrading to Edge provider >= 3.6.0 breaks
metrics export, either
(1) upgrade Airflow to 3.3+, or (2) downgrade to Edge provider <= 3.5.0 with
the workaround above.
+
+Task-instance UUIDs and upgrades
+-------------------------------
Review Comment:
The underline is one character shorter than the title, so both docs builds
fail with `edge_executor.rst:249: WARNING: Title underline too short.`
##########
providers/edge3/src/airflow/providers/edge3/migrations/versions/0006_5_0_0_add_task_instance_id_to_edge_job.py:
##########
@@ -0,0 +1,71 @@
+# 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.
+"""
+Add task instance identity to Edge jobs.
+
+Revision ID: f2a4b6c8d0e1
+Revises: c6b3c3d093fd
+Create Date: 2026-09-29 00:00:00.000000
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import mysql
+
+revision = "f2a4b6c8d0e1"
+down_revision = "c6b3c3d093fd"
+branch_labels = None
+depends_on = None
+edge3_version = "5.0.0"
Review Comment:
`providers-edge3/5.0.0rc1` is already tagged (a99354726, before this PR) and
its `_REVISION_HEADS_MAP` stops at `c6b3c3d093fd`. If rc1 passes the vote,
5.0.0 ships without this revision. Then `db-manager downgrade --to-version
5.0.0` resolves to `f2a4b6c8d0e1` through the map in db.py and does nothing,
and migrations-ref.rst names the wrong release. Is this going into an rc2? If
not, the version here, the file prefix, the map key and migrations-ref probably
want the next edge3 version. Either way, the upgrade-order requirement (API
downtime, older workers get no new jobs) seems worth a line in the changelog
pointing at the new docs section.
##########
providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py:
##########
@@ -100,6 +102,10 @@ def fetch(
job: EdgeJobModel | None = session.scalar(query)
if not job:
return None
+ if job.task_instance_id and not (worker.sysinfo or
{}).get("supports_task_instance_uuid"):
+ raise HTTPException(
+ status.HTTP_409_CONFLICT, "Upgrade this Edge worker to report
task-instance UUIDs."
Review Comment:
The only reader of this message is a pre-PR worker, and it never sees it.
Its `_make_generic_request` calls `raise_for_status()` without reading the
body, the 409 isn't retried, and `fetch_and_run_job` doesn't catch it, so the
worker just logs "Task exception was never retrieved ... 409,
message='Conflict'" on every poll. By default `_assert_version` already rejects
a mismatched worker, so this only happens when
`minimum_acceptable_edge_version_for_workers` has been relaxed. Could this log
a warning on the API server naming `worker_name` before raising? The docs
paragraph at edge_executor.rst:265 could also say to raise that setting to this
release so old workers drain rather than idle.
##########
providers/edge3/docs/edge_executor.rst:
##########
@@ -244,3 +244,27 @@ pipeline. The table below documents known compatibility
issues and workarounds:
**For Airflow 3.2 users:** If upgrading to Edge provider >= 3.6.0 breaks
metrics export, either
(1) upgrade Airflow to 3.3+, or (2) downgrade to Edge provider <= 3.5.0 with
the workaround above.
+
+Task-instance UUIDs and upgrades
+-------------------------------
+
+On cores that support UUID executor keys, Edge tracks each task attempt by its
UUID.
+Jobs with the same Dag, task, run, map index and try number can coexist without
+sharing completion events. Older supported cores continue to use coordinate
keys.
+
+This requires the Edge job identity migration and updated Edge workers.
Upgrade in
+this order:
+
+1. Stop Edge scheduling and stop the old Edge worker API processes. Old API
processes
Review Comment:
`EdgeExecutor.start()` runs `initdb()` itself when `check_migration()`
fails, so a new scheduler that starts before step 1 migrates the schema and
starts writing UUID rows while old API pods are still serving. In a rolling
deploy that's the default order. Old pods then hand UUID jobs out without the
409 gate. The worker's body-less PATCH that lands on a new pod matches only
`task_instance_id == ""`, finds nothing, and `state()` returns silently at
jobs.py:163. The row stays RESTARTING, which the orphan path never looks at,
until `job_fail_purge`. Could the docs say explicitly that the scheduler
performs this migration on start? And could the state route log a warning when
it misses but a row exists at the same coordinates under another
`task_instance_id`? That's one extra query on the miss path only.
##########
providers/edge3/tests/unit/edge3/executors/test_edge_executor.py:
##########
@@ -886,3 +896,135 @@ def test_queue_workload_unknown_type_raises(self):
with create_session() as session:
with pytest.raises(TypeError, match="Don't know how to queue
workload"):
executor.queue_workload(MagicMock(spec=[]), session=session)
+
+
[email protected](
+ not hasattr(EdgeExecutor, "get_task_key"), reason="UUID executor contract
requires new core"
+)
+class TestUUIDTaskIdentity(_WorkloadFactory):
+ @pytest.fixture(autouse=True)
+ def uuid_executor(self, monkeypatch):
+ monkeypatch.setattr(EdgeExecutor, "supports_task_instance_uuid", True)
+
+ @pytest.mark.parametrize("terminal_state", [TaskInstanceState.SUCCESS,
TaskInstanceState.FAILED])
+ def test_same_coordinate_jobs_report_their_own_uuid(self, terminal_state):
+ first = self._make_execute_task()
+ second = first.model_copy(update={"ti":
first.ti.model_copy(update={"id": uuid4()})})
+ executor = EdgeExecutor()
+ with create_session() as session:
+ executor.queue_workload(first, session=session)
+ executor.queue_workload(second, session=session)
+ session.flush()
+ assert len(session.scalars(select(EdgeJobModel)).all()) == 2
+ assert executor.running == {first.ti.id, second.ti.id}
+ set_job_state(
Review Comment:
This only pins the `task_instance_id` filter in `state()` part of the time.
`state()` uses `session.scalar(...)`, so with the filter removed it just
returns whichever sibling the scan hits first. `first` is inserted first, and
on SQLite/MySQL the PK order sorts on a random uuid4, so the test still passes
about half the time (a quick SQLite probe returned the first-inserted row
106/200 times). `test_legacy_worker_report_cannot_complete_new_uuid_job` never
pins it, because the "" row sorts and inserts first. Parametrizing the report
target over `first`/`second` and asserting the sibling row is still QUEUED
would make it deterministic.
##########
providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py:
##########
@@ -156,18 +159,22 @@ def queue_workload(
key = workload.key
elif isinstance(workload, workloads.ExecuteTask):
task_instance = workload.ti
- key = task_instance.key
+ coordinates = task_instance.key
+ key = self.get_task_key(task_instance) if
self.supports_task_instance_uuid else coordinates
+ task_instance_id = str(task_instance.id) if
self.supports_task_instance_uuid else ""
# Check if job already exists with same dag_id, task_id, run_id,
map_index, try_number
existing_job = session.scalars(
Review Comment:
`existing_job` is already typed `EdgeJobModel | None` from the callback
branch at 133, so rebinding it to a `ScalarResult` here fails MyPy providers
checks (errors at 167 `[assignment]` and 177 `[union-attr]`). A separate name
like `candidates = session.scalars(...)` then `existing_job = next((job for job
in candidates if ...), None)` clears both.
##########
providers/edge3/tests/unit/edge3/migrations/test_task_instance_identity.py:
##########
@@ -0,0 +1,159 @@
+# 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 importlib import import_module
+from io import StringIO
+from pathlib import Path
+from uuid import uuid4
+
+import pytest
+import sqlalchemy as sa
+from alembic.migration import MigrationContext
+from alembic.operations import Operations
+from alembic.script import ScriptDirectory
+
+migration = import_module(
+
"airflow.providers.edge3.migrations.versions.0006_5_0_0_add_task_instance_id_to_edge_job"
+)
+COORDINATES = ["dag_id", "task_id", "run_id", "map_index", "try_number"]
+
+
+def reflected_columns(connection):
+ return [
+ column | {"type": str(column["type"])} for column in
sa.inspect(connection).get_columns("edge_job")
+ ]
+
+
[email protected]
+def legacy_jobs():
+ engine = sa.create_engine("sqlite://")
+ with engine.begin() as connection:
+ metadata = sa.MetaData(naming_convention={"pk": "%(table_name)s_pkey"})
+ context = MigrationContext.configure(connection,
opts={"target_metadata": metadata})
+ scripts = ScriptDirectory(str(Path(migration.__file__).parents[1]))
+ with Operations.context(context):
+ for revision in reversed(list(scripts.walk_revisions(base="base",
head=migration.down_revision))):
+ revision.module.upgrade()
Review Comment:
All five tests that use this fixture error before running, in the Non-DB
providers job as well as Compat 3.0.6/3.1.8 and lowest-deps. Replaying 0003
evaluates `TIMESTAMP()` from `airflow.migrations.db_types`, whose lazy loader
calls `alembic.context.get_bind()`. That needs an `EnvironmentContext`, and
`Operations.context(...)` alone doesn't set one up: `NameError: Can't invoke
function 'get_bind', as the proxy object has not yet been established for the
Alembic 'EnvironmentContext' class`. The loader caches the resolved type in
module globals, so it only passes locally when an earlier test in the process
has already resolved it. That same cache also means a fix that resolves it here
will leak the SQLite type into later tests unless it's reset afterwards.
Running the walk inside `EnvironmentContext(Config(), scripts)` with
`env.configure(connection=...)` should fix it. Building the pre-0006 `edge_job`
table directly instead of replaying 0001-0005 would also avoid it.
##########
providers/edge3/tests/unit/edge3/executors/test_edge_executor.py:
##########
@@ -51,7 +53,8 @@
class TestEdgeExecutor:
@pytest.fixture(autouse=True)
- def setup_test_cases(self):
+ def setup_test_cases(self, monkeypatch):
+ monkeypatch.setattr(EdgeExecutor, "supports_task_instance_uuid", False)
Review Comment:
Pinning the existing classes to `False` means UUID mode, which is the
production default on a core with `get_task_key`, only runs through
`TestUUIDTaskIdentity`. None of those tests call `_update_orphaned_jobs`, so
the new `TaskInstance.id == key` lookup and its legacy command-JSON fallback
aren't covered. Could the orphan tests be parametrized over both modes?
##########
providers/edge3/tests/unit/edge3/models/test_db.py:
##########
@@ -261,7 +261,7 @@ def
test_initdb_stamps_and_upgrades_when_tables_exist_without_version(self, sess
version = conn.execute(text("SELECT version_num FROM
alembic_version_edge3")).scalar()
columns = {col["name"] for col in
inspect(conn).get_columns("edge_worker")}
- assert version == "c6b3c3d093fd"
+ assert version == "f2a4b6c8d0e1"
Review Comment:
This test,
`test_upgradedb_stamps_and_upgrades_when_tables_exist_without_version` and
`test_migration_adds_concurrency_column` fail on every backend. The Postgres
job log shows `3 failed` with `DuplicateColumn: column "task_instance_id" of
relation "edge_job" already exists`, even though the check is green. MySQL does
the same. SQLite on the older alembic in Compat 3.0.6/3.1.8 and lowest-deps
gives `CircularDependencyError` inside 0006's `batch_alter_table`. The tables
come from `edge_metadata.create_all`, so they're already head-shaped, and the
tests stamp back and re-run the chain. 0002 and 0004 survive that because they
check the inspector before adding columns, but 0006 always runs `add_column`
and the PK rebuild. Guarding 0006 the same way (skip when `task_instance_id` is
already there and in the PK) would fix it. So would putting `edge_job` back to
the 5-column shape in these tests' setup, the way they already rebuild
`edge_worker`.
##########
providers/edge3/src/airflow/providers/edge3/cli/worker.py:
##########
@@ -389,6 +389,7 @@ async def _get_sysinfo(self) -> dict[str, str | int | float
| datetime]:
),
"airflow_version": airflow_version,
"edge_provider_version": edge_provider_version,
+ "supports_task_instance_uuid": True,
Review Comment:
Because this lives in sysinfo, it becomes a metric.
`WorkerStateBody.sysinfo` coerces `True` to `1`, and `set_metrics` gauges every
numeric key outside its fixed exclusion set. So each heartbeat now emits
`edge_worker.supports_task_instance_uuid`, and on Airflow < 3.3 also a
per-worker `edge_worker.supports_task_instance_uuid.<worker_name>`. The worker
hover card in the UI shows "Supports_task_instance_uuid: 1" too, and an
`extended_sysinfo` hook can override it because its `update()` runs after this.
A typed `supports_task_instance_uuid: bool = False` on `WorkerQueuesBody`,
which `fetch()` reads directly, would avoid all three, since old workers omit
it. The smaller fix is adding the key to the exclusion set in `edge_worker.py`.
Nothing asserts the flag either: `test_get_sysinfo` doesn't check it, so
dropping it here keeps the suite green while every UUID job 409s.
##########
providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py:
##########
@@ -261,13 +269,17 @@ def _update_orphaned_jobs(self, session: Session) -> bool:
).all()
for job in lifeless_jobs:
- ti = TaskInstance.get_task_instance(
- dag_id=job.dag_id,
- run_id=job.run_id,
- task_id=job.task_id,
- map_index=job.map_index,
- session=session,
- )
+ key = self._job_key(job)
+ if isinstance(key, UUID):
+ ti = session.scalar(select(TaskInstance).where(TaskInstance.id
== key))
Review Comment:
When the attempt was already retried through the Execution API
(`prepare_db_for_next_try` rotated `ti.id`) and the worker dies before its
final state report, this lookup returns None and the job becomes REMOVED.
`_purge_jobs` then keeps the key in `running`: REMOVED is in
`_PURGE_HANDLED_STATES` and the `else` branch never calls `fail()`. So the slot
stays taken until `job_fail_purge` deletes the row, 60 minutes by default. The
scheduler's heartbeat purge doesn't cover it either, because it only selects
RUNNING/RESTARTING TIs and this one is UP_FOR_RETRY. Before, the coordinate
lookup found the UP_FOR_RETRY row and `fail()` freed the slot on the next sync.
Could the not-found case release the slot directly (`self.fail(key)`) instead
of parking the job as REMOVED?
##########
providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py:
##########
@@ -284,23 +296,25 @@ def _update_orphaned_jobs(self, session: Session) -> bool:
return bool(lifeless_jobs)
+ def _job_key(self, job: EdgeJobModel) -> UUID | TaskInstanceKey |
CallbackKey:
+ key = job.key
+ if self.supports_task_instance_uuid and isinstance(key,
TaskInstanceKey):
+ return UUID(job.task_instance_id or
json.loads(job.command)["ti"]["id"])
+ return key
+
def _get_tracked_job_keys(
self, session: Session, states: Sequence[TaskInstanceState]
- ) -> set[TaskInstanceKey | CallbackKey]:
+ ) -> set[UUID | TaskInstanceKey | CallbackKey]:
"""
Read the keys of this team's jobs that are in one of ``states``.
Rows are read without locking on purpose: an edge worker fetches its
next job with
``FOR UPDATE SKIP LOCKED``, so locking the queued rows here would make
it come back empty.
"""
- query = select(
- EdgeJobModel.dag_id,
- EdgeJobModel.task_id,
- EdgeJobModel.run_id,
- EdgeJobModel.try_number,
- EdgeJobModel.map_index,
- ).where(EdgeJobModel.team_name == self.team_name,
EdgeJobModel.state.in_(states))
- return {build_job_key(*row) for row in session.execute(query)}
+ query = select(EdgeJobModel).where(
+ EdgeJobModel.team_name == self.team_name,
EdgeJobModel.state.in_(states)
+ )
+ return {self._job_key(job) for job in session.scalars(query)}
Review Comment:
This used to select five columns and now hydrates full rows, including
`command`, for every QUEUED and handled-state row on every `sync()`. A quick
Postgres probe with 5000 rows and a ~1 KB command measured about 5 ms for the
column tuples vs 18 ms for full entities. Only legacy rows (`task_instance_id
== ""`) need `command`, so selecting the key columns plus
`case((EdgeJobModel.task_instance_id == "", EdgeJobModel.command))`, or
`.options(defer(EdgeJobModel.command))`, keeps the old cost.
##########
providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py:
##########
@@ -138,41 +145,34 @@ def state(
map_index: Annotated[int, WorkerApiDocs.map_index],
state: Annotated[TaskInstanceState, WorkerApiDocs.state],
session: SessionDep,
+ task_instance_id: Annotated[UUID | None, Body(embed=True)] = None,
) -> None:
"""Update the state of a job running on the edge worker."""
- # execute query to catch the queue and check if state toggles to success
or failed
- # otherwise possible that Executor resets orphaned jobs and stats are
exported 2 times
- if state in [TaskInstanceState.SUCCESS, state == TaskInstanceState.FAILED]:
- query = select(EdgeJobModel).where(
- EdgeJobModel.dag_id == dag_id,
- EdgeJobModel.task_id == task_id,
- EdgeJobModel.run_id == run_id,
- EdgeJobModel.map_index == map_index,
- EdgeJobModel.try_number == try_number,
- EdgeJobModel.state == TaskInstanceState.RUNNING,
- )
- job = session.scalar(query)
-
- if job:
- # Edge worker does not backport emitted Airflow metrics, so export
some metrics
- tags = {
- "dag_id": job.dag_id,
- "task_id": job.task_id,
- "queue": job.queue,
- "state": str(state),
- "team_name": job.team_name,
- }
- Stats.incr("edge_worker.ti.finish", tags=prune_dict(tags))
-
- query2 = (
- update(EdgeJobModel)
+ job = session.scalar(
+ select(EdgeJobModel)
.where(
EdgeJobModel.dag_id == dag_id,
EdgeJobModel.task_id == task_id,
EdgeJobModel.run_id == run_id,
EdgeJobModel.map_index == map_index,
EdgeJobModel.try_number == try_number,
+ EdgeJobModel.task_instance_id == (str(task_instance_id) if
task_instance_id else ""),
)
- .values(state=state, last_update=timezone.utcnow())
+ .with_for_update()
)
- session.execute(query2)
+ if job is None:
+ return
+ if job.state == TaskInstanceState.RUNNING and state in (
+ TaskInstanceState.SUCCESS,
+ TaskInstanceState.FAILED,
+ ):
+ tags = {
+ "dag_id": job.dag_id,
+ "task_id": job.task_id,
+ "queue": job.queue,
+ "state": str(state),
Review Comment:
This also changes the metric: the old `state in [SUCCESS, state == FAILED]`
never matched FAILED, so `edge_worker.ti.finish{state=failed}` is new with this
PR. It's the right fix, but dashboards will see a new series. Could it get a
changelog line and a FAILED case in `test_state`?
--
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]
