This is an automated email from the ASF dual-hosted git repository.
dheerajturaga pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 2f3e60cb949 Support priority_weight in EdgeExecutor (#72057)
2f3e60cb949 is described below
commit 2f3e60cb94901e9f9bc47e3b0aeb7f8290e4e983
Author: PoAn Yang <[email protected]>
AuthorDate: Wed Oct 14 03:18:50 2026 +0800
Support priority_weight in EdgeExecutor (#72057)
* Support priority_weight in EdgeExecutor
Signed-off-by: PoAn Yang <[email protected]>
* Move Edge priority_weight migration to edge3 5.1.0
edge3 5.0.0 was released with f2a4b6c8d0e1 as its migration head, so this
migration ships in the next release. Recording it under 5.0.0 made
`airflow db-manager` resolve `--to-version 5.0.0` and `--from-version 5.0.0`
to the wrong revision.
---------
Signed-off-by: PoAn Yang <[email protected]>
Co-authored-by: Dheeraj Turaga <[email protected]>
---
providers/edge3/docs/changelog.rst | 9 +
providers/edge3/docs/edge_executor.rst | 7 +
providers/edge3/docs/migrations-ref.rst | 4 +-
.../providers/edge3/executors/edge_executor.py | 2 +
.../0007_5_1_0_add_priority_weight_to_edge_job.py | 55 +++++
.../edge3/src/airflow/providers/edge3/models/db.py | 1 +
.../src/airflow/providers/edge3/models/edge_job.py | 5 +-
.../src/airflow/providers/edge3/models/types.py | 20 +-
.../providers/edge3/worker_api/routes/jobs.py | 34 ++--
.../unit/edge3/executors/test_edge_executor.py | 19 +-
providers/edge3/tests/unit/edge3/models/test_db.py | 72 ++++++-
.../unit/edge3/worker_api/routes/test_jobs.py | 222 ++++++++++++++++++++-
12 files changed, 413 insertions(+), 37 deletions(-)
diff --git a/providers/edge3/docs/changelog.rst
b/providers/edge3/docs/changelog.rst
index 70a1ffb9cd8..4d6c9ab0778 100644
--- a/providers/edge3/docs/changelog.rst
+++ b/providers/edge3/docs/changelog.rst
@@ -27,6 +27,15 @@
Changelog
---------
+.. note::
+ Instead of strict queue order, ``EdgeExecutor`` now hands out queued
callbacks first, then tasks
+ by ``priority_weight``, as in the other executors. A steady stream of
higher-priority tasks can
+ hold back lower-priority tasks on the same queue. With the default
+ ``[core] default_task_weight_rule`` of ``downstream``, a task's weight
includes the weights of its
+ downstream tasks. Upstream tasks of large Dags therefore go first even when
no ``priority_weight``
+ is set. Setting the rule to ``absolute`` keeps tasks without their own
``priority_weight`` in
+ queue order among themselves. The rule affects every executor.
+
5.0.0
.....
diff --git a/providers/edge3/docs/edge_executor.rst
b/providers/edge3/docs/edge_executor.rst
index b29dafbe3af..df6e697e734 100644
--- a/providers/edge3/docs/edge_executor.rst
+++ b/providers/edge3/docs/edge_executor.rst
@@ -54,6 +54,13 @@ could take thousands of tasks without a problem), or from an
environment
perspective (you want a worker running from a specific location where required
infrastructure is available).
+Within the queues a worker listens to, jobs are handed out by
+:ref:`apache-airflow:concepts:priority-weight`. Among the queued tasks that
fit in the worker's
+free :ref:`concurrency slots <edge_executor:concurrency_slots>`, the one with
the highest
+``priority_weight`` is picked up first. Tasks of equal weight go in the order
they were queued.
+:ref:`Callbacks <edge_executor:execute_callback>` are picked up before any
task, in the order they
+were queued.
+
When using EdgeExecutor in addition to other executors and EdgeExecutor not
being the default executor
(that is to say the first one in the list of executors), be reminded to also
define EdgeExecutor
as the executor at task or Dag level in addition to the queues you are
targeting.
diff --git a/providers/edge3/docs/migrations-ref.rst
b/providers/edge3/docs/migrations-ref.rst
index 45600d053a2..e4ade8ae266 100644
--- a/providers/edge3/docs/migrations-ref.rst
+++ b/providers/edge3/docs/migrations-ref.rst
@@ -34,7 +34,9 @@ Here's the list of all the Database Migrations that are
executed via when you ru
+-------------------------+------------------+-----------------+----------------------------------------------------------+
| Revision ID | Revises ID | Edge3 Version | Description
|
+=========================+==================+=================+==========================================================+
-| ``f2a4b6c8d0e1`` (head) | ``c6b3c3d093fd`` | ``5.0.0`` | Add task
instance identity to Edge jobs. |
+| ``d5a2f8b41c07`` (head) | ``f2a4b6c8d0e1`` | ``5.1.0`` | Add
priority_weight column to edge_job table. |
++-------------------------+------------------+-----------------+----------------------------------------------------------+
+| ``f2a4b6c8d0e1`` | ``c6b3c3d093fd`` | ``5.0.0`` | Add task
instance identity to Edge jobs. |
+-------------------------+------------------+-----------------+----------------------------------------------------------+
| ``c6b3c3d093fd`` | ``a09c3ee8e1d3`` | ``3.5.0`` | Replace
individual counters with extended JSON based |
| | | | sysinfo.
|
diff --git
a/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py
b/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py
index 11bcda605dd..2d298a16386 100644
--- a/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py
+++ b/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py
@@ -186,6 +186,7 @@ class EdgeExecutor(BaseExecutor):
existing_job.state = TaskInstanceState.QUEUED
existing_job.queue = task_instance.queue
existing_job.concurrency_slots = task_instance.pool_slots
+ existing_job.priority_weight = task_instance.priority_weight
existing_job.command = workload.model_dump_json()
existing_job.team_name = self.team_name
else:
@@ -200,6 +201,7 @@ class EdgeExecutor(BaseExecutor):
state=TaskInstanceState.QUEUED,
queue=task_instance.queue,
concurrency_slots=task_instance.pool_slots,
+ priority_weight=task_instance.priority_weight,
command=workload.model_dump_json(),
team_name=self.team_name,
)
diff --git
a/providers/edge3/src/airflow/providers/edge3/migrations/versions/0007_5_1_0_add_priority_weight_to_edge_job.py
b/providers/edge3/src/airflow/providers/edge3/migrations/versions/0007_5_1_0_add_priority_weight_to_edge_job.py
new file mode 100644
index 00000000000..30aa0a1c579
--- /dev/null
+++
b/providers/edge3/src/airflow/providers/edge3/migrations/versions/0007_5_1_0_add_priority_weight_to_edge_job.py
@@ -0,0 +1,55 @@
+#
+# 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 priority_weight column to edge_job table.
+
+Revision ID: d5a2f8b41c07
+Revises: f2a4b6c8d0e1
+Create Date: 2026-08-25 00:00:00.000000
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "d5a2f8b41c07"
+down_revision = "f2a4b6c8d0e1"
+branch_labels = None
+depends_on = None
+edge3_version = "5.1.0"
+
+
+def upgrade() -> None:
+ with op.batch_alter_table("edge_job", schema=None) as batch_op:
+ batch_op.add_column(
+ sa.Column("priority_weight", sa.Integer(),
server_default=sa.text("1"), nullable=False)
+ )
+ op.drop_index("rj_order", table_name="edge_job")
+ op.create_index(
+ "rj_order", "edge_job", ["state", sa.text("priority_weight DESC"),
"queued_dttm", "queue"]
+ )
+
+
+def downgrade() -> None:
+ op.drop_index("rj_order", table_name="edge_job")
+ op.create_index("rj_order", "edge_job", ["state", "queued_dttm", "queue"])
+ with op.batch_alter_table("edge_job", schema=None) as batch_op:
+ batch_op.drop_column("priority_weight")
diff --git a/providers/edge3/src/airflow/providers/edge3/models/db.py
b/providers/edge3/src/airflow/providers/edge3/models/db.py
index f997afe49b0..315b8397d81 100644
--- a/providers/edge3/src/airflow/providers/edge3/models/db.py
+++ b/providers/edge3/src/airflow/providers/edge3/models/db.py
@@ -48,6 +48,7 @@ _REVISION_HEADS_MAP: dict[str, str] = {
"3.4.0": "a09c3ee8e1d3",
"3.5.0": "c6b3c3d093fd",
"5.0.0": "f2a4b6c8d0e1",
+ "5.1.0": "d5a2f8b41c07",
}
diff --git a/providers/edge3/src/airflow/providers/edge3/models/edge_job.py
b/providers/edge3/src/airflow/providers/edge3/models/edge_job.py
index 45192181003..41c8c2f5c17 100644
--- a/providers/edge3/src/airflow/providers/edge3/models/edge_job.py
+++ b/providers/edge3/src/airflow/providers/edge3/models/edge_job.py
@@ -84,6 +84,7 @@ class EdgeJobModel(Base, LoggingMixin):
state: Mapped[str] = mapped_column(String(20))
queue: Mapped[str] = mapped_column(String(256))
concurrency_slots: Mapped[int] = mapped_column(Integer)
+ priority_weight: Mapped[int] = mapped_column(Integer, nullable=False,
server_default=text("1"))
command: Mapped[str] = mapped_column(String(2048))
queued_dttm: Mapped[datetime | None] = mapped_column(UtcDateTime)
edge_worker: Mapped[str | None] = mapped_column(String(64))
@@ -101,6 +102,7 @@ class EdgeJobModel(Base, LoggingMixin):
queue: str,
concurrency_slots: int,
command: str,
+ priority_weight: int = 1,
queued_dttm: datetime | None = None,
edge_worker: str | None = None,
last_update: datetime | None = None,
@@ -116,6 +118,7 @@ class EdgeJobModel(Base, LoggingMixin):
self.queue = queue
self.concurrency_slots = concurrency_slots
self.command = command
+ self.priority_weight = priority_weight
self.queued_dttm = queued_dttm or timezone.utcnow()
self.edge_worker = edge_worker
self.last_update = last_update
@@ -123,7 +126,7 @@ class EdgeJobModel(Base, LoggingMixin):
self.task_instance_id = task_instance_id
super().__init__()
- __table_args__ = (Index("rj_order", state, queued_dttm, queue),)
+ __table_args__ = (Index("rj_order", state, priority_weight.desc(),
queued_dttm, queue),)
@property
def key(self) -> TaskInstanceKey | CallbackKey:
diff --git a/providers/edge3/src/airflow/providers/edge3/models/types.py
b/providers/edge3/src/airflow/providers/edge3/models/types.py
index dce93f52f82..1eb4e143d2b 100644
--- a/providers/edge3/src/airflow/providers/edge3/models/types.py
+++ b/providers/edge3/src/airflow/providers/edge3/models/types.py
@@ -18,10 +18,14 @@ from __future__ import annotations
from typing import TYPE_CHECKING, TypeAlias, TypeGuard
+from sqlalchemy import and_, literal
+
from airflow.executors.workloads import ExecuteTask
from airflow.providers.edge3.version_compat import AIRFLOW_V_3_3_PLUS
if TYPE_CHECKING:
+ from sqlalchemy.sql.elements import ColumnElement
+
from airflow.executors import workloads
from airflow.executors.workloads import ExecuteCallback
@@ -49,10 +53,11 @@ EXECUTE_CALLBACK_TAG = "ExecuteCallback"
# Dag id, so a row is a callback only when all four fields match.
CALLBACK_JOB_TRY_NUMBER = 0
CALLBACK_JOB_MAP_INDEX = -1
+CALLBACK_JOB_RUN_ID_PREFIX = f"{EXECUTE_CALLBACK_TAG}-"
def build_callback_run_id(callback_id: str) -> str:
- return f"{EXECUTE_CALLBACK_TAG}-{callback_id}"
+ return f"{CALLBACK_JOB_RUN_ID_PREFIX}{callback_id}"
def is_callback_job(dag_id: str, task_id: str, run_id: str, try_number: int,
map_index: int) -> bool:
@@ -63,3 +68,16 @@ def is_callback_job(dag_id: str, task_id: str, run_id: str,
try_number: int, map
and try_number == CALLBACK_JOB_TRY_NUMBER
and map_index == CALLBACK_JOB_MAP_INDEX
)
+
+
+def build_callback_job_filter() -> ColumnElement[bool]:
+ """Return :func:`is_callback_job` as a SQL filter on ``edge_job`` rows."""
+ # Imported here because the edge_job module imports this one.
+ from airflow.providers.edge3.models.edge_job import EdgeJobModel
+
+ return and_(
+ EdgeJobModel.dag_id == EXECUTE_CALLBACK_TAG,
+ EdgeJobModel.run_id == literal(CALLBACK_JOB_RUN_ID_PREFIX) +
EdgeJobModel.task_id,
+ EdgeJobModel.try_number == CALLBACK_JOB_TRY_NUMBER,
+ EdgeJobModel.map_index == CALLBACK_JOB_MAP_INDEX,
+ )
diff --git
a/providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py
b/providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py
index 5d1b791b763..ea1073b88f1 100644
--- a/providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py
+++ b/providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py
@@ -31,6 +31,7 @@ from airflow.executors.workloads import ExecuteTask
from airflow.providers.common.compat.sdk import Stats, timezone
from airflow.providers.edge3.models.edge_job import EdgeJobModel
from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel
+from airflow.providers.edge3.models.types import build_callback_job_filter,
is_callback_job
from airflow.providers.edge3.version_compat import AIRFLOW_V_3_3_PLUS
from airflow.providers.edge3.worker_api.auth import
jwt_token_authorization_rest
from airflow.providers.edge3.worker_api.datamodels import (
@@ -49,13 +50,13 @@ log = logging.getLogger(__name__)
jobs_router = AirflowRouter(tags=["Jobs"], prefix="/jobs")
-def parse_command(command: str, dag_id: str, run_id: str) -> ExecuteTypeBody:
- if AIRFLOW_V_3_3_PLUS:
+def parse_command(
+ command: str, dag_id: str, task_id: str, run_id: str, try_number: int,
map_index: int
+) -> ExecuteTypeBody:
+ if AIRFLOW_V_3_3_PLUS and is_callback_job(dag_id, task_id, run_id,
try_number, map_index):
from airflow.executors.workloads import ExecuteCallback
- from airflow.providers.edge3.models.types import EXECUTE_CALLBACK_TAG
- if dag_id == EXECUTE_CALLBACK_TAG and
run_id.startswith(EXECUTE_CALLBACK_TAG):
- return ExecuteCallback.model_validate_json(command) # type:
ignore[return-value]
+ return ExecuteCallback.model_validate_json(command) # type:
ignore[return-value]
return ExecuteTask.model_validate_json(command)
@@ -88,13 +89,9 @@ def fetch(
if not worker:
raise HTTPException(status.HTTP_404_NOT_FOUND, "Worker not found")
- query = (
- select(EdgeJobModel)
- .where(
- EdgeJobModel.state == TaskInstanceState.QUEUED,
- EdgeJobModel.concurrency_slots <= body.free_concurrency,
- )
- .order_by(EdgeJobModel.queued_dttm)
+ query = select(EdgeJobModel).where(
+ EdgeJobModel.state == TaskInstanceState.QUEUED,
+ EdgeJobModel.concurrency_slots <= body.free_concurrency,
)
if body.queues:
query = query.where(EdgeJobModel.queue.in_(body.queues))
@@ -102,7 +99,14 @@ def fetch(
query = query.where(EdgeJobModel.team_name == worker.team_name)
query = query.limit(1)
query = query.with_for_update(skip_locked=True)
- job: EdgeJobModel | None = session.scalar(query)
+
+ # Callbacks go first, as in other executors (WORKLOAD_TYPE_PRIORITY). A
single query would sort
+ # the whole backlog. dag_id leads the primary key, so the callback query
reads only callbacks.
+ job: EdgeJobModel | None = session.scalar(
+
query.where(build_callback_job_filter()).order_by(EdgeJobModel.queued_dttm)
+ )
+ if job is None:
+ job =
session.scalar(query.order_by(EdgeJobModel.priority_weight.desc(),
EdgeJobModel.queued_dttm))
if not job:
return None
if job.task_instance_id and not body.supports_task_instance_uuid:
@@ -125,7 +129,9 @@ def fetch(
run_id=job.run_id,
map_index=job.map_index,
try_number=job.try_number,
- command=parse_command(job.command, job.dag_id, job.run_id),
+ command=parse_command(
+ job.command, job.dag_id, job.task_id, job.run_id, job.try_number,
job.map_index
+ ),
concurrency_slots=job.concurrency_slots,
task_instance_id=UUID(job.task_instance_id) if job.task_instance_id
else None,
)
diff --git a/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py
b/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py
index df117c00d73..023b248bd90 100644
--- a/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py
+++ b/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py
@@ -844,7 +844,9 @@ class _WorkloadFactory:
session.execute(delete(EdgeJobModel))
session.commit()
- def _make_execute_task(self, task_id: str = "test_task", dag_id: str =
"test_dag") -> ExecuteTask:
+ def _make_execute_task(
+ self, task_id: str = "test_task", dag_id: str = "test_dag",
priority_weight: int = 1
+ ) -> ExecuteTask:
ti = TaskInstanceDTO(
id=uuid4(),
dag_version_id=uuid4(),
@@ -855,7 +857,7 @@ class _WorkloadFactory:
map_index=-1,
pool_slots=1,
queue="default",
- priority_weight=1,
+ priority_weight=priority_weight,
)
return ExecuteTask(
ti=ti,
@@ -1074,6 +1076,19 @@ class TestQueueWorkload(_WorkloadFactory):
assert len(jobs) == 1
assert jobs[0].state == TaskInstanceState.QUEUED
+ def test_queue_workload_stores_task_priority_weight(self):
+ executor = EdgeExecutor()
+
+ with create_session() as session:
+
executor.queue_workload(self._make_execute_task(priority_weight=42),
session=session)
+ with create_session() as session:
+ assert session.scalar(select(EdgeJobModel)).priority_weight == 42
+
+ with create_session() as session:
+
executor.queue_workload(self._make_execute_task(priority_weight=7),
session=session)
+ with create_session() as session:
+ assert session.scalar(select(EdgeJobModel)).priority_weight == 7
+
def test_queue_workload_execute_callback(self):
executor = EdgeExecutor()
workload = self._make_execute_callback()
diff --git a/providers/edge3/tests/unit/edge3/models/test_db.py
b/providers/edge3/tests/unit/edge3/models/test_db.py
index 2405f39fdb4..d44640a7289 100644
--- a/providers/edge3/tests/unit/edge3/models/test_db.py
+++ b/providers/edge3/tests/unit/edge3/models/test_db.py
@@ -22,10 +22,12 @@ from unittest import mock
import pytest
import sqlalchemy as sa
+from alembic import command
from alembic.migration import MigrationContext
from alembic.operations import Operations
from airflow import settings
+from airflow.providers.edge3.models.db import EdgeDBManager
from airflow.providers.edge3.models.edge_base import edge_metadata
from airflow.utils.db_manager import RunDBManager
@@ -244,9 +246,6 @@ class TestEdgeDBManager:
"""Test that initdb runs incremental migrations when tables exist but
alembic version table does not."""
from sqlalchemy import inspect, text
- from airflow import settings
- from airflow.providers.edge3.models.db import EdgeDBManager
-
manager = EdgeDBManager(session)
# Simulate pre-alembic state: tables exist but no version table and no
concurrency column
@@ -281,7 +280,7 @@ class TestEdgeDBManager:
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 == "f2a4b6c8d0e1"
+ assert version in manager.get_script_object().get_heads()
assert "concurrency" in columns
assert "team_name" in columns
@@ -291,9 +290,6 @@ class TestEdgeDBManager:
"""Test upgradedb runs incremental migrations when tables exist but
alembic version table does not."""
from sqlalchemy import inspect, text
- from airflow import settings
- from airflow.providers.edge3.models.db import EdgeDBManager
-
manager = EdgeDBManager(session)
# Simulate pre-alembic state: tables exist but no version table and no
concurrency column
@@ -334,14 +330,10 @@ class TestEdgeDBManager:
def test_migration_adds_concurrency_column(self, session,
legacy_edge_job_table):
"""Test that upgrading from 3.0.0 actually adds the concurrency
column."""
- from alembic import command
from alembic.migration import MigrationContext
from alembic.operations import Operations
from sqlalchemy import inspect
- from airflow import settings
- from airflow.providers.edge3.models.db import EdgeDBManager
-
manager = EdgeDBManager(session)
config = manager.get_alembic_config()
@@ -380,6 +372,64 @@ class TestEdgeDBManager:
assert "concurrency" in columns, "Migration 0002 should have added the
concurrency column"
assert "team_name" in columns, "Migration 0003 should have added the
team_name column"
+ @pytest.fixture
+ def restore_edge_tables(self, session):
+ """
+ Rebuild the edge3 tables at head with the migrations, whatever state
the test left behind.
+
+ From Airflow 3.2 on, the test database also builds these tables with
the migrations.
+ ``create_all()`` would give the columns different collations on MySQL.
+ """
+ yield
+ # The test's session is still in a transaction that read the Alembic
version table. On
+ # PostgreSQL, that read holds a lock, so the DROP on another
connection would wait forever.
+ session.rollback()
+ manager = EdgeDBManager(session)
+ with settings.engine.begin() as conn:
+ manager.drop_tables(conn)
+ manager.upgradedb(use_migration_files=True)
+
+ def test_priority_weight_migration_upgrade_and_downgrade(self, session,
restore_edge_tables):
+ manager = EdgeDBManager(session)
+ config = manager.get_alembic_config()
+ # Before Airflow 3.2, the edge3 tables come from the conftest's
create_all(), which leaves no
+ # Alembic version to downgrade from. Stamping changes nothing when the
head is already recorded.
+ command.stamp(config, "heads")
+ command.downgrade(config, "f2a4b6c8d0e1")
+ # SQLite reflects indexes from a connection's cached schema. The
downgrade ran on another pooled
+ # connection, so idle connections can still hold the old schema until
the pool is disposed.
+ settings.engine.dispose()
+
+ with settings.engine.begin() as conn:
+ inspector = sa.inspect(conn)
+ rj_order = next(idx for idx in inspector.get_indexes("edge_job")
if idx["name"] == "rj_order")
+ assert "priority_weight" not in {col["name"] for col in
inspector.get_columns("edge_job")}
+ assert rj_order["column_names"] == ["state", "queued_dttm",
"queue"]
+ conn.execute(
+ sa.text(
+ "INSERT INTO edge_job (dag_id, task_id, run_id, map_index,
try_number, state, queue, "
+ "concurrency_slots, command) "
+ "VALUES ('dag', 'task', 'run', -1, 1, 'queued', 'default',
1, '{}')"
+ )
+ )
+
+ manager.upgradedb()
+ # The upgrade also ran on another connection, so dispose the pool
again.
+ settings.engine.dispose()
+
+ with settings.engine.connect() as conn:
+ inspector = sa.inspect(conn)
+ rj_order = next(idx for idx in inspector.get_indexes("edge_job")
if idx["name"] == "rj_order")
+ priority_weight = conn.execute(
+ sa.text("SELECT priority_weight FROM edge_job WHERE dag_id =
'dag'")
+ ).scalar_one()
+
+ assert rj_order["column_names"] == ["state", "priority_weight",
"queued_dttm", "queue"]
+ # Of the supported backends, only PostgreSQL reflects the sort order
of index columns.
+ if settings.engine.dialect.name == "postgresql":
+ assert rj_order["column_sorting"] == {"priority_weight": ("desc",)}
+ assert priority_weight == 1
+
def test_drop_tables_handles_missing_tables(self, session):
"""Test that drop_tables handles missing tables gracefully."""
from airflow.providers.edge3.models.db import EdgeDBManager
diff --git a/providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py
b/providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py
index 4a53dad4158..ff4be7d5041 100644
--- a/providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py
+++ b/providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import json
+from datetime import timedelta
from pathlib import Path
from typing import TYPE_CHECKING
from unittest.mock import patch
@@ -27,10 +28,15 @@ from fastapi import HTTPException, status
from sqlalchemy import delete, select
from airflow.executors.workloads import BundleInfo, ExecuteTask
-from airflow.providers.common.compat.sdk import Stats
+from airflow.providers.common.compat.sdk import Stats, timezone
from airflow.providers.edge3.models.edge_job import EdgeJobModel
from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel,
EdgeWorkerState
-from airflow.providers.edge3.models.types import EXECUTE_CALLBACK_TAG
+from airflow.providers.edge3.models.types import (
+ CALLBACK_JOB_MAP_INDEX,
+ CALLBACK_JOB_TRY_NUMBER,
+ EXECUTE_CALLBACK_TAG,
+ build_callback_run_id,
+)
from airflow.providers.edge3.worker_api.datamodels import WorkerQueuesBody
from airflow.providers.edge3.worker_api.routes import jobs
from airflow.providers.edge3.worker_api.routes.jobs import fetch,
parse_command, state
@@ -309,6 +315,188 @@ class TestJobsApiRoutes:
},
)
+ @patch(f"{Stats.__module__}.Stats.incr")
+ def test_fetch_returns_highest_priority_job_first(self, mock_stats_incr,
session: Session):
+ with create_session() as session:
+ session.add(
+ EdgeWorkerModel(
+ worker_name="worker1", state=EdgeWorkerState.IDLE,
queues=[QUEUE], team_name=None
+ )
+ )
+ queued_dttm = timezone.datetime(2026, 1, 1)
+ for task_id, priority_weight, delay in [
+ ("low", 1, 0),
+ ("high", 100, 1),
+ ("medium_late", 10, 3),
+ ("medium_early", 10, 2),
+ ]:
+ session.add(
+ EdgeJobModel(
+ dag_id=DAG_ID,
+ task_id=task_id,
+ run_id=RUN_ID,
+ try_number=1,
+ map_index=-1,
+ state=TaskInstanceState.QUEUED,
+ queue=QUEUE,
+ concurrency_slots=1,
+ command=MOCK_COMMAND_STR,
+ priority_weight=priority_weight,
+ queued_dttm=queued_dttm + timedelta(seconds=delay),
+ )
+ )
+ session.commit()
+
+ body = WorkerQueuesBody(free_concurrency=1, queues=[QUEUE],
team_name=None)
+ fetched = [fetch("worker1", body, session) for _ in range(4)]
+
+ assert [job.task_id for job in fetched if job] == ["high",
"medium_early", "medium_late", "low"]
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="ExecuteCallback
requires Airflow 3.3+")
+ @patch(f"{Stats.__module__}.Stats.incr")
+ def test_fetch_returns_callbacks_before_higher_priority_tasks(self,
mock_stats_incr, session: Session):
+ """Callbacks are fetched in queue order regardless of their own
priority_weight."""
+ queued_dttm = timezone.datetime(2026, 1, 1)
+ callbacks = [
+ ("11111111-1111-1111-1111-111111111111", 100, queued_dttm),
+ ("22222222-2222-2222-2222-222222222222", 1, queued_dttm -
timedelta(seconds=1)),
+ ]
+ with create_session() as session:
+ session.add(
+ EdgeWorkerModel(
+ worker_name="worker1", state=EdgeWorkerState.IDLE,
queues=[QUEUE], team_name=None
+ )
+ )
+ for callback_id, priority_weight, callback_queued_dttm in
callbacks:
+ callback = ExecuteCallback(
+ callback=CallbackDTO(
+ id=callback_id,
+ fetch_method=CallbackFetchMethod.IMPORT_PATH,
+ data={"path": "builtins.dict", "kwargs": {}},
+ ),
+ dag_rel_path=Path("test.py"),
+ bundle_info=BundleInfo(name="test_bundle", version="1.0"),
+ token="test_token",
+ log_path="test.log",
+ )
+ session.add(
+ EdgeJobModel(
+ dag_id=EXECUTE_CALLBACK_TAG,
+ task_id=callback_id,
+ run_id=build_callback_run_id(callback_id),
+ try_number=CALLBACK_JOB_TRY_NUMBER,
+ map_index=CALLBACK_JOB_MAP_INDEX,
+ state=TaskInstanceState.QUEUED,
+ queue=QUEUE,
+ concurrency_slots=1,
+ command=callback.model_dump_json(),
+ priority_weight=priority_weight,
+ queued_dttm=callback_queued_dttm,
+ )
+ )
+ session.add(
+ EdgeJobModel(
+ dag_id=DAG_ID,
+ task_id="high",
+ run_id=RUN_ID,
+ try_number=1,
+ map_index=-1,
+ state=TaskInstanceState.QUEUED,
+ queue=QUEUE,
+ concurrency_slots=1,
+ command=MOCK_COMMAND_STR,
+ priority_weight=100,
+ queued_dttm=queued_dttm - timedelta(seconds=2),
+ )
+ )
+ session.commit()
+
+ body = WorkerQueuesBody(free_concurrency=1, queues=[QUEUE],
team_name=None)
+ fetched = [fetch("worker1", body, session) for _ in range(3)]
+
+ assert [job.task_id for job in fetched if job] ==
[callbacks[1][0], callbacks[0][0], "high"]
+
+ @pytest.mark.parametrize(
+ ("dag_id", "run_id", "try_number", "map_index"),
+ [
+ pytest.param(
+ DAG_ID,
+ build_callback_run_id("low"),
+ CALLBACK_JOB_TRY_NUMBER,
+ CALLBACK_JOB_MAP_INDEX,
+ id="only-dag-id-differs",
+ ),
+ pytest.param(
+ EXECUTE_CALLBACK_TAG,
+ build_callback_run_id("other"),
+ CALLBACK_JOB_TRY_NUMBER,
+ CALLBACK_JOB_MAP_INDEX,
+ id="only-run-id-differs",
+ ),
+ pytest.param(
+ EXECUTE_CALLBACK_TAG,
+ build_callback_run_id("low"),
+ CALLBACK_JOB_TRY_NUMBER + 1,
+ CALLBACK_JOB_MAP_INDEX,
+ id="only-try-number-differs",
+ ),
+ pytest.param(
+ EXECUTE_CALLBACK_TAG,
+ build_callback_run_id("low"),
+ CALLBACK_JOB_TRY_NUMBER,
+ CALLBACK_JOB_MAP_INDEX + 1,
+ id="only-map-index-differs",
+ ),
+ ],
+ )
+ @patch(f"{Stats.__module__}.Stats.incr")
+ def test_fetch_needs_full_callback_identity_to_jump_the_queue(
+ self, mock_stats_incr, dag_id, run_id, try_number, map_index, session:
Session
+ ):
+ queued_dttm = timezone.datetime(2026, 1, 1)
+ with create_session() as session:
+ session.add(
+ EdgeWorkerModel(
+ worker_name="worker1", state=EdgeWorkerState.IDLE,
queues=[QUEUE], team_name=None
+ )
+ )
+ session.add(
+ EdgeJobModel(
+ dag_id=DAG_ID,
+ task_id="high",
+ run_id=RUN_ID,
+ try_number=1,
+ map_index=-1,
+ state=TaskInstanceState.QUEUED,
+ queue=QUEUE,
+ concurrency_slots=1,
+ command=MOCK_COMMAND_STR,
+ priority_weight=100,
+ queued_dttm=queued_dttm + timedelta(seconds=1),
+ )
+ )
+ session.add(
+ EdgeJobModel(
+ dag_id=dag_id,
+ task_id="low",
+ run_id=run_id,
+ try_number=try_number,
+ map_index=map_index,
+ state=TaskInstanceState.QUEUED,
+ queue=QUEUE,
+ concurrency_slots=1,
+ command=MOCK_COMMAND_STR,
+ priority_weight=1,
+ queued_dttm=queued_dttm,
+ )
+ )
+ session.commit()
+
+ body = WorkerQueuesBody(free_concurrency=1, queues=[QUEUE],
team_name=None)
+ fetched = [fetch("worker1", body, session) for _ in range(2)]
+
+ assert [job.task_id for job in fetched if job] == ["high", "low"]
+
@patch(f"{Stats.__module__}.Stats.incr")
def test_fetch_filters_by_worker_team_name(self, mock_stats_incr, session:
Session):
with create_session() as session:
@@ -469,7 +657,14 @@ class TestParseCommand:
workload = self._make_execute_task()
command_json = workload.model_dump_json()
- result = parse_command(command_json, dag_id="test_dag",
run_id="test_run")
+ result = parse_command(
+ command_json,
+ dag_id="test_dag",
+ task_id="test_task",
+ run_id="test_run",
+ try_number=1,
+ map_index=-1,
+ )
assert isinstance(result, ExecuteTask)
assert result.ti.dag_id == "test_dag"
@@ -479,10 +674,16 @@ class TestParseCommand:
workload = self._make_execute_callback()
command_json = workload.model_dump_json()
- dag_id = EXECUTE_CALLBACK_TAG
- run_id = f"{EXECUTE_CALLBACK_TAG}-{workload.callback.key}"
+ callback_id = workload.callback.id
- result = parse_command(command_json, dag_id=dag_id, run_id=run_id)
+ result = parse_command(
+ command_json,
+ dag_id=EXECUTE_CALLBACK_TAG,
+ task_id=callback_id,
+ run_id=build_callback_run_id(callback_id),
+ try_number=CALLBACK_JOB_TRY_NUMBER,
+ map_index=CALLBACK_JOB_MAP_INDEX,
+ )
assert isinstance(result, ExecuteCallback)
assert result.callback.id == "12345678-1234-5678-1234-567812345678"
@@ -493,6 +694,13 @@ class TestParseCommand:
workload = self._make_execute_task()
command_json = workload.model_dump_json()
- result = parse_command(command_json, dag_id="some_dag",
run_id=f"{EXECUTE_CALLBACK_TAG}-something")
+ result = parse_command(
+ command_json,
+ dag_id="some_dag",
+ task_id="something",
+ run_id=build_callback_run_id("something"),
+ try_number=CALLBACK_JOB_TRY_NUMBER,
+ map_index=CALLBACK_JOB_MAP_INDEX,
+ )
assert isinstance(result, ExecuteTask)
