Skip to content

fix(scheduler): load dag_run eagering for execution API validation#56919

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
eladkal merged 1 commit into from
Oct 21, 2025
Merged

fix(scheduler): load dag_run eagering for execution API validation #56919

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 31 additions & 3 deletions airflow-core/src/airflow/jobs/scheduler_job_runner.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -93,7 +93,7 @@
from types import FrameType

from pendulum.datetime import DateTime
from sqlalchemy.orm import Query, Session
from sqlalchemy.orm import Load, Query, Session

from airflow._shared.logging.types import Logger
from airflow.executors.base_executor import BaseExecutor
Expand All@@ -110,6 +110,31 @@
""":meta private:"""


def _eager_load_dag_run_for_validation() -> tuple[Load, Load]:
"""
Eager-load DagRun relations required for execution API datamodel validation.

When building TaskCallbackRequest with DRDataModel.model_validate(ti.dag_run),
the consumed_asset_events collection and nested asset/source_aliases must be
preloaded to avoid DetachedInstanceError after the session closes.

Returns a tuple of two load options:
- Asset loader: TI.dag_run → consumed_asset_events → asset
- Alias loader: TI.dag_run → consumed_asset_events → source_aliases

Example usage::

asset_loader, alias_loader = _eager_load_dag_run_for_validation()
query = select(TI).options(asset_loader).options(alias_loader)
"""
# Traverse TI → dag_run → consumed_asset_events once, then branch to asset/aliases
base = selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events)
return (
base.selectinload(AssetEvent.asset),
base.selectinload(AssetEvent.source_aliases),
)


def _get_current_dag(dag_id: str, session: Session) -> SerializedDAG | None:
serdag = SerializedDagModel.get(dag_id=dag_id, session=session) # grabs the latest version
if not serdag:
Expand DownExpand Up@@ -806,11 +831,12 @@ def process_executor_events(

# Check state of finished tasks
filter_for_tis = TI.filter_for_tis(tis_with_right_state)
asset_loader, alias_loader = _eager_load_dag_run_for_validation()
query = (
select(TI)
.where(filter_for_tis)
.options(selectinload(TI.dag_model))
.options(joinedload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
.options(asset_loader)
.options(joinedload(TI.dag_run).selectinload(DagRun.created_dag_version))
.options(joinedload(TI.dag_version))
)
Expand DownExpand Up@@ -2375,10 +2401,12 @@ def _find_and_purge_task_instances_without_heartbeats(self) -> None:
def _find_task_instances_without_heartbeats(self, *, session: Session) -> list[TI]:
self.log.debug("Finding 'running' jobs without a recent heartbeat")
limit_dttm = timezone.utcnow() - timedelta(seconds=self._task_instance_heartbeat_timeout_secs)
asset_loader, alias_loader = _eager_load_dag_run_for_validation()
task_instances_without_heartbeats = session.scalars(
select(TI)
.options(selectinload(TI.dag_model))
.options(selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
.options(asset_loader)
.options(alias_loader)
.options(selectinload(TI.dag_version))
.with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
.join(DM, TI.dag_id == DM.dag_id)
Expand Down
150 changes: 150 additions & 0 deletions airflow-core/tests/unit/jobs/test_scheduler_job.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -608,6 +608,66 @@ def test_process_executor_events_ti_requeued(self, mock_stats_incr, mock_task_ca
scheduler_job.executor.callback_sink.send.assert_not_called()
mock_stats_incr.assert_not_called()

@pytest.mark.usefixtures("testing_dag_bundle")
@mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr")
def test_process_executor_events_with_asset_events(self, mock_stats_incr, session, dag_maker):
"""
Test that _process_executor_events handles asset events without DetachedInstanceError.

Regression test for scheduler crashes when task callbacks are built with
consumed_asset_events that weren't eager-loaded.
"""
asset1 = Asset(uri="test://asset1", name="test_asset_executor", group="test_group")
asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
session.add(asset_model)
session.flush()

with dag_maker(dag_id="test_executor_events_with_assets", schedule=[asset1], fileloc="/test_path1/"):
EmptyOperator(task_id="dummy_task", on_failure_callback=lambda ctx: None)

dag = dag_maker.dag
sync_dag_to_db(dag)
DagVersion.get_latest_version(dag.dag_id)

dr = dag_maker.create_dagrun()

# Create asset event and attach to dag run
asset_event = AssetEvent(
asset_id=asset_model.id,
source_task_id="upstream_task",
source_dag_id="upstream_dag",
source_run_id="upstream_run",
source_map_index=-1,
)
session.add(asset_event)
session.flush()
dr.consumed_asset_events.append(asset_event)
session.flush()

executor = MockExecutor(do_update=False)
scheduler_job = Job(executor=executor)
self.job_runner = SchedulerJobRunner(scheduler_job)

ti1 = dr.get_task_instance("dummy_task")
ti1.state = State.QUEUED
session.merge(ti1)
session.commit()

executor.event_buffer[ti1.key] = State.FAILED, None

# This should not raise DetachedInstanceError
self.job_runner._process_executor_events(executor=executor, session=session)

ti1.refresh_from_db(session=session)
assert ti1.state == State.FAILED

# Verify callback was created with asset event data
scheduler_job.executor.callback_sink.send.assert_called_once()
callback_request = scheduler_job.executor.callback_sink.send.call_args.args[0]
assert callback_request.context_from_server is not None
assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
assert callback_request.context_from_server.dag_run.consumed_asset_events[0].asset.uri == asset1.uri

def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker):
dag_id = "SchedulerJobTest.test_execute_task_instances_is_paused_wont_execute"
task_id_1 = "dummy_task"
Expand All@@ -628,6 +688,96 @@ def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker)
assert ti1.state == State.SCHEDULED
session.rollback()

@pytest.mark.usefixtures("testing_dag_bundle")
def test_find_and_purge_task_instances_without_heartbeats_with_asset_events(
self, session, dag_maker, create_dagrun
):
"""
Test that heartbeat purge succeeds when DagRun has consumed_asset_events.

Regression test for DetachedInstanceError when building TaskCallbackRequest
with asset event data after session expunge.
"""
asset1 = Asset(uri="test://asset1", name="test_asset", group="test_group")
asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
session.add(asset_model)
session.flush()

with dag_maker(dag_id="test_heartbeat_with_assets", schedule=[asset1]):
EmptyOperator(task_id="dummy_task")

dag = dag_maker.dag
scheduler_dag = sync_dag_to_db(dag)
dag_v = DagVersion.get_latest_version(dag.dag_id)

data_interval = infer_automated_data_interval(scheduler_dag.timetable, DEFAULT_LOGICAL_DATE)
dag_run = create_dagrun(
scheduler_dag,
logical_date=DEFAULT_DATE,
run_type=DagRunType.SCHEDULED,
data_interval=data_interval,
)

# Create asset alias and event with full relationships
asset_alias = AssetAliasModel(name="test_alias", group="test_group")
session.add(asset_alias)
session.flush()

asset_event = AssetEvent(
asset_id=asset_model.id,
source_task_id="upstream_task",
source_dag_id="upstream_dag",
source_run_id="upstream_run",
source_map_index=-1,
)
session.add(asset_event)
session.flush()

# Attach alias to event and event to dag run
asset_event.source_aliases.append(asset_alias)
dag_run.consumed_asset_events.append(asset_event)
session.flush()

executor = MockExecutor()
scheduler_job = Job(executor=executor)
with mock.patch("airflow.executors.executor_loader.ExecutorLoader.load_executor") as loader_mock:
loader_mock.return_value = executor
self.job_runner = SchedulerJobRunner(job=scheduler_job)

ti = dag_run.get_task_instance("dummy_task")
assert ti is not None # sanity check: dag_maker.create_dagrun created the TI

ti.state = State.RUNNING
ti.last_heartbeat_at = timezone.utcnow() - timedelta(minutes=6)
ti.start_date = timezone.utcnow() - timedelta(minutes=10)
ti.queued_by_job_id = scheduler_job.id
ti.dag_version = dag_v
session.merge(ti)
session.flush()

executor.running.add(ti.key)

tis_without_heartbeats = self.job_runner._find_task_instances_without_heartbeats(session=session)
assert len(tis_without_heartbeats) == 1
ti_from_query = tis_without_heartbeats[0]
ti_key = ti_from_query.key

# Detach all ORM objects to mirror scheduler behaviour after session closes
session.expunge_all()

# This should not raise DetachedInstanceError now that eager loads are in place
self.job_runner._purge_task_instances_without_heartbeats(tis_without_heartbeats, session=session)
assert ti_key not in executor.running

executor.callback_sink.send.assert_called_once()
callback_request = executor.callback_sink.send.call_args.args[0]
assert callback_request.context_from_server is not None
assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
consumed_event = callback_request.context_from_server.dag_run.consumed_asset_events[0]
assert consumed_event.asset.uri == asset1.uri
assert len(consumed_event.source_aliases) == 1
assert consumed_event.source_aliases[0].name == "test_alias"

# @pytest.mark.usefixtures("mock_executor")
def test_execute_task_instances_backfill_tasks_will_execute(self, dag_maker):
"""
Expand Down
, 'i'); if (__m === '*' || __re.test(location.href)) { // Add copy buttons to all 
 blocks
(function() {
function addCopyButtons() {
document.querySelectorAll('pre code').forEach(function(codeBlock) {
if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;
codeBlock.parentElement.setAttribute('data-copy-added', 'true');
var btn = document.createElement('button');
btn.textContent = 'Copy';
btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';
btn.onmouseover = function() { this.style.opacity = '1'; };
btn.onmouseout = function() { this.style.opacity = '0.7'; };
btn.onclick = function() {
navigator.clipboard.writeText(codeBlock.textContent).then(function() {
btn.textContent = 'Copied!';
setTimeout(function() { btn.textContent = 'Copy'; }, 1500);
});
};
codeBlock.parentElement.style.position = 'relative';
codeBlock.parentElement.appendChild(btn);
});
}
addCopyButtons();
// Re-run on dynamic content
var observer = new MutationObserver(addCopyButtons);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
fix(scheduler): load dag_run eagering for execution API validation by Lee-W · Pull Request #56919 · apache/airflow · GitHub
Skip to content
fix(scheduler): load dag_run eagering for execution API validation#56919

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

eladkal
merged 1 commit into

from

Oct 21, 2025
 Merged


fix(scheduler): load dag_run eagering for execution API validation
#56919
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 31 additions & 3 deletions
airflow-core/src/airflow/jobs/scheduler_job_runner.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -93,7 +93,7 @@
 from types import FrameType

 from pendulum.datetime import DateTime
 from sqlalchemy.orm import Query, Session
 from sqlalchemy.orm import Load, Query, Session

 from airflow._shared.logging.types import Logger
 from airflow.executors.base_executor import BaseExecutor
Expand All@@ -110,6 +110,31 @@
""":meta private:"""


def _eager_load_dag_run_for_validation() -> tuple[Load, Load]:
 """
 Eager-load DagRun relations required for execution API datamodel validation.

 When building TaskCallbackRequest with DRDataModel.model_validate(ti.dag_run),
 the consumed_asset_events collection and nested asset/source_aliases must be
 preloaded to avoid DetachedInstanceError after the session closes.

 Returns a tuple of two load options:
 - Asset loader: TI.dag_run → consumed_asset_events → asset
 - Alias loader: TI.dag_run → consumed_asset_events → source_aliases

 Example usage::

 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = select(TI).options(asset_loader).options(alias_loader)
 """
 # Traverse TI → dag_run → consumed_asset_events once, then branch to asset/aliases
 base = selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events)
 return (
 base.selectinload(AssetEvent.asset),
 base.selectinload(AssetEvent.source_aliases),
 )


def _get_current_dag(dag_id: str, session: Session) -> SerializedDAG | None:
 serdag = SerializedDagModel.get(dag_id=dag_id, session=session) # grabs the latest version
 if not serdag:
Expand DownExpand Up@@ -806,11 +831,12 @@ def process_executor_events(

 # Check state of finished tasks
 filter_for_tis = TI.filter_for_tis(tis_with_right_state)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = (
 select(TI)
 .where(filter_for_tis)
 .options(selectinload(TI.dag_model))
 .options(joinedload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(joinedload(TI.dag_run).selectinload(DagRun.created_dag_version))
 .options(joinedload(TI.dag_version))
 )
Expand DownExpand Up@@ -2375,10 +2401,12 @@ def _find_and_purge_task_instances_without_heartbeats(self) -> None:
 def _find_task_instances_without_heartbeats(self, *, session: Session) -> list[TI]:
 self.log.debug("Finding 'running' jobs without a recent heartbeat")
 limit_dttm = timezone.utcnow() - timedelta(seconds=self._task_instance_heartbeat_timeout_secs)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 task_instances_without_heartbeats = session.scalars(
 select(TI)
 .options(selectinload(TI.dag_model))
 .options(selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(alias_loader)
 .options(selectinload(TI.dag_version))
 .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
 .join(DM, TI.dag_id == DM.dag_id)
Expand Down

150 changes: 150 additions & 0 deletions
airflow-core/tests/unit/jobs/test_scheduler_job.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -608,6 +608,66 @@ def test_process_executor_events_ti_requeued(self, mock_stats_incr, mock_task_ca
 scheduler_job.executor.callback_sink.send.assert_not_called()
 mock_stats_incr.assert_not_called()

 @pytest.mark.usefixtures("testing_dag_bundle")
 @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr")
 def test_process_executor_events_with_asset_events(self, mock_stats_incr, session, dag_maker):
 """
 Test that _process_executor_events handles asset events without DetachedInstanceError.

 Regression test for scheduler crashes when task callbacks are built with
 consumed_asset_events that weren't eager-loaded.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset_executor", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_executor_events_with_assets", schedule=[asset1], fileloc="/test_path1/"):
 EmptyOperator(task_id="dummy_task", on_failure_callback=lambda ctx: None)

 dag = dag_maker.dag
 sync_dag_to_db(dag)
 DagVersion.get_latest_version(dag.dag_id)

 dr = dag_maker.create_dagrun()

 # Create asset event and attach to dag run
 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()
 dr.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor(do_update=False)
 scheduler_job = Job(executor=executor)
 self.job_runner = SchedulerJobRunner(scheduler_job)

 ti1 = dr.get_task_instance("dummy_task")
 ti1.state = State.QUEUED
 session.merge(ti1)
 session.commit()

 executor.event_buffer[ti1.key] = State.FAILED, None

 # This should not raise DetachedInstanceError
 self.job_runner._process_executor_events(executor=executor, session=session)

 ti1.refresh_from_db(session=session)
 assert ti1.state == State.FAILED

 # Verify callback was created with asset event data
 scheduler_job.executor.callback_sink.send.assert_called_once()
 callback_request = scheduler_job.executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 assert callback_request.context_from_server.dag_run.consumed_asset_events[0].asset.uri == asset1.uri

 def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker):
 dag_id = "SchedulerJobTest.test_execute_task_instances_is_paused_wont_execute"
 task_id_1 = "dummy_task"
Expand All@@ -628,6 +688,96 @@ def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker)
 assert ti1.state == State.SCHEDULED
 session.rollback()

 @pytest.mark.usefixtures("testing_dag_bundle")
 def test_find_and_purge_task_instances_without_heartbeats_with_asset_events(
 self, session, dag_maker, create_dagrun
 ):
 """
 Test that heartbeat purge succeeds when DagRun has consumed_asset_events.

 Regression test for DetachedInstanceError when building TaskCallbackRequest
 with asset event data after session expunge.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_heartbeat_with_assets", schedule=[asset1]):
 EmptyOperator(task_id="dummy_task")

 dag = dag_maker.dag
 scheduler_dag = sync_dag_to_db(dag)
 dag_v = DagVersion.get_latest_version(dag.dag_id)

 data_interval = infer_automated_data_interval(scheduler_dag.timetable, DEFAULT_LOGICAL_DATE)
 dag_run = create_dagrun(
 scheduler_dag,
 logical_date=DEFAULT_DATE,
 run_type=DagRunType.SCHEDULED,
 data_interval=data_interval,
 )

 # Create asset alias and event with full relationships
 asset_alias = AssetAliasModel(name="test_alias", group="test_group")
 session.add(asset_alias)
 session.flush()

 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()

 # Attach alias to event and event to dag run
 asset_event.source_aliases.append(asset_alias)
 dag_run.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor()
 scheduler_job = Job(executor=executor)
 with mock.patch("airflow.executors.executor_loader.ExecutorLoader.load_executor") as loader_mock:
 loader_mock.return_value = executor
 self.job_runner = SchedulerJobRunner(job=scheduler_job)

 ti = dag_run.get_task_instance("dummy_task")
 assert ti is not None # sanity check: dag_maker.create_dagrun created the TI

 ti.state = State.RUNNING
 ti.last_heartbeat_at = timezone.utcnow() - timedelta(minutes=6)
 ti.start_date = timezone.utcnow() - timedelta(minutes=10)
 ti.queued_by_job_id = scheduler_job.id
 ti.dag_version = dag_v
 session.merge(ti)
 session.flush()

 executor.running.add(ti.key)

 tis_without_heartbeats = self.job_runner._find_task_instances_without_heartbeats(session=session)
 assert len(tis_without_heartbeats) == 1
 ti_from_query = tis_without_heartbeats[0]
 ti_key = ti_from_query.key

 # Detach all ORM objects to mirror scheduler behaviour after session closes
 session.expunge_all()

 # This should not raise DetachedInstanceError now that eager loads are in place
 self.job_runner._purge_task_instances_without_heartbeats(tis_without_heartbeats, session=session)
 assert ti_key not in executor.running

 executor.callback_sink.send.assert_called_once()
 callback_request = executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 consumed_event = callback_request.context_from_server.dag_run.consumed_asset_events[0]
 assert consumed_event.asset.uri == asset1.uri
 assert len(consumed_event.source_aliases) == 1
 assert consumed_event.source_aliases[0].name == "test_alias"

 # @pytest.mark.usefixtures("mock_executor")
 def test_execute_task_instances_backfill_tasks_will_execute(self, dag_maker):
 """
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Force GitHub README to respect dark mode
(function() {
var style = document.createElement('style');
style.textContent = '
.markdown-body {
color-scheme: dark light;
}
.markdown-body pre { background: #161b22 !important; }
.markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }
.markdown-body table th, .markdown-body table td { border-color: #30363d !important; }
.markdown-body img { background: #0d1117; }
.markdown-body blockquote { border-left-color: #8b949e; }
.markdown-body hr { border-color: #30363d; }
';
document.head.appendChild(style);
})();
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix(scheduler): load dag_run eagering for execution API validation by Lee-W · Pull Request #56919 · apache/airflow · GitHub
Skip to content
fix(scheduler): load dag_run eagering for execution API validation#56919

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

eladkal
merged 1 commit into

from

Oct 21, 2025
 Merged


fix(scheduler): load dag_run eagering for execution API validation
#56919
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 31 additions & 3 deletions
airflow-core/src/airflow/jobs/scheduler_job_runner.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -93,7 +93,7 @@
 from types import FrameType

 from pendulum.datetime import DateTime
 from sqlalchemy.orm import Query, Session
 from sqlalchemy.orm import Load, Query, Session

 from airflow._shared.logging.types import Logger
 from airflow.executors.base_executor import BaseExecutor
Expand All@@ -110,6 +110,31 @@
""":meta private:"""


def _eager_load_dag_run_for_validation() -> tuple[Load, Load]:
 """
 Eager-load DagRun relations required for execution API datamodel validation.

 When building TaskCallbackRequest with DRDataModel.model_validate(ti.dag_run),
 the consumed_asset_events collection and nested asset/source_aliases must be
 preloaded to avoid DetachedInstanceError after the session closes.

 Returns a tuple of two load options:
 - Asset loader: TI.dag_run → consumed_asset_events → asset
 - Alias loader: TI.dag_run → consumed_asset_events → source_aliases

 Example usage::

 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = select(TI).options(asset_loader).options(alias_loader)
 """
 # Traverse TI → dag_run → consumed_asset_events once, then branch to asset/aliases
 base = selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events)
 return (
 base.selectinload(AssetEvent.asset),
 base.selectinload(AssetEvent.source_aliases),
 )


def _get_current_dag(dag_id: str, session: Session) -> SerializedDAG | None:
 serdag = SerializedDagModel.get(dag_id=dag_id, session=session) # grabs the latest version
 if not serdag:
Expand DownExpand Up@@ -806,11 +831,12 @@ def process_executor_events(

 # Check state of finished tasks
 filter_for_tis = TI.filter_for_tis(tis_with_right_state)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = (
 select(TI)
 .where(filter_for_tis)
 .options(selectinload(TI.dag_model))
 .options(joinedload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(joinedload(TI.dag_run).selectinload(DagRun.created_dag_version))
 .options(joinedload(TI.dag_version))
 )
Expand DownExpand Up@@ -2375,10 +2401,12 @@ def _find_and_purge_task_instances_without_heartbeats(self) -> None:
 def _find_task_instances_without_heartbeats(self, *, session: Session) -> list[TI]:
 self.log.debug("Finding 'running' jobs without a recent heartbeat")
 limit_dttm = timezone.utcnow() - timedelta(seconds=self._task_instance_heartbeat_timeout_secs)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 task_instances_without_heartbeats = session.scalars(
 select(TI)
 .options(selectinload(TI.dag_model))
 .options(selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(alias_loader)
 .options(selectinload(TI.dag_version))
 .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
 .join(DM, TI.dag_id == DM.dag_id)
Expand Down

150 changes: 150 additions & 0 deletions
airflow-core/tests/unit/jobs/test_scheduler_job.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -608,6 +608,66 @@ def test_process_executor_events_ti_requeued(self, mock_stats_incr, mock_task_ca
 scheduler_job.executor.callback_sink.send.assert_not_called()
 mock_stats_incr.assert_not_called()

 @pytest.mark.usefixtures("testing_dag_bundle")
 @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr")
 def test_process_executor_events_with_asset_events(self, mock_stats_incr, session, dag_maker):
 """
 Test that _process_executor_events handles asset events without DetachedInstanceError.

 Regression test for scheduler crashes when task callbacks are built with
 consumed_asset_events that weren't eager-loaded.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset_executor", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_executor_events_with_assets", schedule=[asset1], fileloc="/test_path1/"):
 EmptyOperator(task_id="dummy_task", on_failure_callback=lambda ctx: None)

 dag = dag_maker.dag
 sync_dag_to_db(dag)
 DagVersion.get_latest_version(dag.dag_id)

 dr = dag_maker.create_dagrun()

 # Create asset event and attach to dag run
 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()
 dr.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor(do_update=False)
 scheduler_job = Job(executor=executor)
 self.job_runner = SchedulerJobRunner(scheduler_job)

 ti1 = dr.get_task_instance("dummy_task")
 ti1.state = State.QUEUED
 session.merge(ti1)
 session.commit()

 executor.event_buffer[ti1.key] = State.FAILED, None

 # This should not raise DetachedInstanceError
 self.job_runner._process_executor_events(executor=executor, session=session)

 ti1.refresh_from_db(session=session)
 assert ti1.state == State.FAILED

 # Verify callback was created with asset event data
 scheduler_job.executor.callback_sink.send.assert_called_once()
 callback_request = scheduler_job.executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 assert callback_request.context_from_server.dag_run.consumed_asset_events[0].asset.uri == asset1.uri

 def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker):
 dag_id = "SchedulerJobTest.test_execute_task_instances_is_paused_wont_execute"
 task_id_1 = "dummy_task"
Expand All@@ -628,6 +688,96 @@ def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker)
 assert ti1.state == State.SCHEDULED
 session.rollback()

 @pytest.mark.usefixtures("testing_dag_bundle")
 def test_find_and_purge_task_instances_without_heartbeats_with_asset_events(
 self, session, dag_maker, create_dagrun
 ):
 """
 Test that heartbeat purge succeeds when DagRun has consumed_asset_events.

 Regression test for DetachedInstanceError when building TaskCallbackRequest
 with asset event data after session expunge.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_heartbeat_with_assets", schedule=[asset1]):
 EmptyOperator(task_id="dummy_task")

 dag = dag_maker.dag
 scheduler_dag = sync_dag_to_db(dag)
 dag_v = DagVersion.get_latest_version(dag.dag_id)

 data_interval = infer_automated_data_interval(scheduler_dag.timetable, DEFAULT_LOGICAL_DATE)
 dag_run = create_dagrun(
 scheduler_dag,
 logical_date=DEFAULT_DATE,
 run_type=DagRunType.SCHEDULED,
 data_interval=data_interval,
 )

 # Create asset alias and event with full relationships
 asset_alias = AssetAliasModel(name="test_alias", group="test_group")
 session.add(asset_alias)
 session.flush()

 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()

 # Attach alias to event and event to dag run
 asset_event.source_aliases.append(asset_alias)
 dag_run.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor()
 scheduler_job = Job(executor=executor)
 with mock.patch("airflow.executors.executor_loader.ExecutorLoader.load_executor") as loader_mock:
 loader_mock.return_value = executor
 self.job_runner = SchedulerJobRunner(job=scheduler_job)

 ti = dag_run.get_task_instance("dummy_task")
 assert ti is not None # sanity check: dag_maker.create_dagrun created the TI

 ti.state = State.RUNNING
 ti.last_heartbeat_at = timezone.utcnow() - timedelta(minutes=6)
 ti.start_date = timezone.utcnow() - timedelta(minutes=10)
 ti.queued_by_job_id = scheduler_job.id
 ti.dag_version = dag_v
 session.merge(ti)
 session.flush()

 executor.running.add(ti.key)

 tis_without_heartbeats = self.job_runner._find_task_instances_without_heartbeats(session=session)
 assert len(tis_without_heartbeats) == 1
 ti_from_query = tis_without_heartbeats[0]
 ti_key = ti_from_query.key

 # Detach all ORM objects to mirror scheduler behaviour after session closes
 session.expunge_all()

 # This should not raise DetachedInstanceError now that eager loads are in place
 self.job_runner._purge_task_instances_without_heartbeats(tis_without_heartbeats, session=session)
 assert ti_key not in executor.running

 executor.callback_sink.send.assert_called_once()
 callback_request = executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 consumed_event = callback_request.context_from_server.dag_run.consumed_asset_events[0]
 assert consumed_event.asset.uri == asset1.uri
 assert len(consumed_event.source_aliases) == 1
 assert consumed_event.source_aliases[0].name == "test_alias"

 # @pytest.mark.usefixtures("mock_executor")
 def test_execute_task_instances_backfill_tasks_will_execute(self, dag_maker):
 """
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Highlight search terms from Google/DuckDuckGo/Bing referrer
(function() {
var ref = document.referrer;
var terms = [];
if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {
var url = new URL(ref);
var q = url.searchParams.get('q') || url.searchParams.get('p');
if (q) {
terms = q.split(/\s+/).filter(function(t) { return t.length > 2; });
}
}
if (terms.length === 0) return;
var style = document.createElement('style');
style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';
document.head.appendChild(style);
function highlight(node) {
if (node.nodeType === 3) { // text node
var text = node.textContent;
var found = false;
terms.forEach(function(term) {
var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\]\\]/g, '\\') + ')', 'gi');
if (regex.test(text)) {
found = true;
var frag = document.createDocumentFragment();
var parts = text.split(regex);
parts.forEach(function(part, i) {
if (i % 2 === 0) {
frag.appendChild(document.createTextNode(part));
} else {
var span = document.createElement('span');
span.className = 'userscript-highlight';
span.textContent = part;
frag.appendChild(span);
}
});
node.parentNode.replaceChild(frag, node);
}
});
} else if (node.nodeType === 1 && node.childNodes) { // element
var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];
if (!skipTags.includes(node.tagName)) {
Array.from(node.childNodes).forEach(highlight);
}
}
}
highlight(document.body);
// Re-highlight on dynamic content
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1 || node.nodeType === 3) highlight(node);
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix(scheduler): load dag_run eagering for execution API validation by Lee-W · Pull Request #56919 · apache/airflow · GitHub
Skip to content
fix(scheduler): load dag_run eagering for execution API validation#56919

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

eladkal
merged 1 commit into

from

Oct 21, 2025
 Merged


fix(scheduler): load dag_run eagering for execution API validation
#56919
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 31 additions & 3 deletions
airflow-core/src/airflow/jobs/scheduler_job_runner.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -93,7 +93,7 @@
 from types import FrameType

 from pendulum.datetime import DateTime
 from sqlalchemy.orm import Query, Session
 from sqlalchemy.orm import Load, Query, Session

 from airflow._shared.logging.types import Logger
 from airflow.executors.base_executor import BaseExecutor
Expand All@@ -110,6 +110,31 @@
""":meta private:"""


def _eager_load_dag_run_for_validation() -> tuple[Load, Load]:
 """
 Eager-load DagRun relations required for execution API datamodel validation.

 When building TaskCallbackRequest with DRDataModel.model_validate(ti.dag_run),
 the consumed_asset_events collection and nested asset/source_aliases must be
 preloaded to avoid DetachedInstanceError after the session closes.

 Returns a tuple of two load options:
 - Asset loader: TI.dag_run → consumed_asset_events → asset
 - Alias loader: TI.dag_run → consumed_asset_events → source_aliases

 Example usage::

 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = select(TI).options(asset_loader).options(alias_loader)
 """
 # Traverse TI → dag_run → consumed_asset_events once, then branch to asset/aliases
 base = selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events)
 return (
 base.selectinload(AssetEvent.asset),
 base.selectinload(AssetEvent.source_aliases),
 )


def _get_current_dag(dag_id: str, session: Session) -> SerializedDAG | None:
 serdag = SerializedDagModel.get(dag_id=dag_id, session=session) # grabs the latest version
 if not serdag:
Expand DownExpand Up@@ -806,11 +831,12 @@ def process_executor_events(

 # Check state of finished tasks
 filter_for_tis = TI.filter_for_tis(tis_with_right_state)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = (
 select(TI)
 .where(filter_for_tis)
 .options(selectinload(TI.dag_model))
 .options(joinedload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(joinedload(TI.dag_run).selectinload(DagRun.created_dag_version))
 .options(joinedload(TI.dag_version))
 )
Expand DownExpand Up@@ -2375,10 +2401,12 @@ def _find_and_purge_task_instances_without_heartbeats(self) -> None:
 def _find_task_instances_without_heartbeats(self, *, session: Session) -> list[TI]:
 self.log.debug("Finding 'running' jobs without a recent heartbeat")
 limit_dttm = timezone.utcnow() - timedelta(seconds=self._task_instance_heartbeat_timeout_secs)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 task_instances_without_heartbeats = session.scalars(
 select(TI)
 .options(selectinload(TI.dag_model))
 .options(selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(alias_loader)
 .options(selectinload(TI.dag_version))
 .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
 .join(DM, TI.dag_id == DM.dag_id)
Expand Down

150 changes: 150 additions & 0 deletions
airflow-core/tests/unit/jobs/test_scheduler_job.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -608,6 +608,66 @@ def test_process_executor_events_ti_requeued(self, mock_stats_incr, mock_task_ca
 scheduler_job.executor.callback_sink.send.assert_not_called()
 mock_stats_incr.assert_not_called()

 @pytest.mark.usefixtures("testing_dag_bundle")
 @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr")
 def test_process_executor_events_with_asset_events(self, mock_stats_incr, session, dag_maker):
 """
 Test that _process_executor_events handles asset events without DetachedInstanceError.

 Regression test for scheduler crashes when task callbacks are built with
 consumed_asset_events that weren't eager-loaded.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset_executor", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_executor_events_with_assets", schedule=[asset1], fileloc="/test_path1/"):
 EmptyOperator(task_id="dummy_task", on_failure_callback=lambda ctx: None)

 dag = dag_maker.dag
 sync_dag_to_db(dag)
 DagVersion.get_latest_version(dag.dag_id)

 dr = dag_maker.create_dagrun()

 # Create asset event and attach to dag run
 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()
 dr.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor(do_update=False)
 scheduler_job = Job(executor=executor)
 self.job_runner = SchedulerJobRunner(scheduler_job)

 ti1 = dr.get_task_instance("dummy_task")
 ti1.state = State.QUEUED
 session.merge(ti1)
 session.commit()

 executor.event_buffer[ti1.key] = State.FAILED, None

 # This should not raise DetachedInstanceError
 self.job_runner._process_executor_events(executor=executor, session=session)

 ti1.refresh_from_db(session=session)
 assert ti1.state == State.FAILED

 # Verify callback was created with asset event data
 scheduler_job.executor.callback_sink.send.assert_called_once()
 callback_request = scheduler_job.executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 assert callback_request.context_from_server.dag_run.consumed_asset_events[0].asset.uri == asset1.uri

 def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker):
 dag_id = "SchedulerJobTest.test_execute_task_instances_is_paused_wont_execute"
 task_id_1 = "dummy_task"
Expand All@@ -628,6 +688,96 @@ def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker)
 assert ti1.state == State.SCHEDULED
 session.rollback()

 @pytest.mark.usefixtures("testing_dag_bundle")
 def test_find_and_purge_task_instances_without_heartbeats_with_asset_events(
 self, session, dag_maker, create_dagrun
 ):
 """
 Test that heartbeat purge succeeds when DagRun has consumed_asset_events.

 Regression test for DetachedInstanceError when building TaskCallbackRequest
 with asset event data after session expunge.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_heartbeat_with_assets", schedule=[asset1]):
 EmptyOperator(task_id="dummy_task")

 dag = dag_maker.dag
 scheduler_dag = sync_dag_to_db(dag)
 dag_v = DagVersion.get_latest_version(dag.dag_id)

 data_interval = infer_automated_data_interval(scheduler_dag.timetable, DEFAULT_LOGICAL_DATE)
 dag_run = create_dagrun(
 scheduler_dag,
 logical_date=DEFAULT_DATE,
 run_type=DagRunType.SCHEDULED,
 data_interval=data_interval,
 )

 # Create asset alias and event with full relationships
 asset_alias = AssetAliasModel(name="test_alias", group="test_group")
 session.add(asset_alias)
 session.flush()

 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()

 # Attach alias to event and event to dag run
 asset_event.source_aliases.append(asset_alias)
 dag_run.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor()
 scheduler_job = Job(executor=executor)
 with mock.patch("airflow.executors.executor_loader.ExecutorLoader.load_executor") as loader_mock:
 loader_mock.return_value = executor
 self.job_runner = SchedulerJobRunner(job=scheduler_job)

 ti = dag_run.get_task_instance("dummy_task")
 assert ti is not None # sanity check: dag_maker.create_dagrun created the TI

 ti.state = State.RUNNING
 ti.last_heartbeat_at = timezone.utcnow() - timedelta(minutes=6)
 ti.start_date = timezone.utcnow() - timedelta(minutes=10)
 ti.queued_by_job_id = scheduler_job.id
 ti.dag_version = dag_v
 session.merge(ti)
 session.flush()

 executor.running.add(ti.key)

 tis_without_heartbeats = self.job_runner._find_task_instances_without_heartbeats(session=session)
 assert len(tis_without_heartbeats) == 1
 ti_from_query = tis_without_heartbeats[0]
 ti_key = ti_from_query.key

 # Detach all ORM objects to mirror scheduler behaviour after session closes
 session.expunge_all()

 # This should not raise DetachedInstanceError now that eager loads are in place
 self.job_runner._purge_task_instances_without_heartbeats(tis_without_heartbeats, session=session)
 assert ti_key not in executor.running

 executor.callback_sink.send.assert_called_once()
 callback_request = executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 consumed_event = callback_request.context_from_server.dag_run.consumed_asset_events[0]
 assert consumed_event.asset.uri == asset1.uri
 assert len(consumed_event.source_aliases) == 1
 assert consumed_event.source_aliases[0].name == "test_alias"

 # @pytest.mark.usefixtures("mock_executor")
 def test_execute_task_instances_backfill_tasks_will_execute(self, dag_maker):
 """
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Strip utm_, fbclid, gclid, etc. from all links on page
(function() {
var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',
'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',
'ref', 'ref_src', 'source', 'medium', 'campaign'];
function cleanUrl(url) {
try {
var u = new URL(url, window.location.origin);
var changed = false;
trackingParams.forEach(function(p) {
if (u.searchParams.has(p)) {
u.searchParams.delete(p);
changed = true;
}
});
return changed ? u.toString() : url;
} catch (e) {
return url;
}
}
function cleanLinks() {
document.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
cleanLinks();
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1) {
if (node.tagName === 'A') cleanLinks();
node.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '
fix(scheduler): load dag_run eagering for execution API validation by Lee-W · Pull Request #56919 · apache/airflow · GitHub
Skip to content
fix(scheduler): load dag_run eagering for execution API validation#56919

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

eladkal
merged 1 commit into

from

Oct 21, 2025
 Merged


fix(scheduler): load dag_run eagering for execution API validation
#56919
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 31 additions & 3 deletions
airflow-core/src/airflow/jobs/scheduler_job_runner.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -93,7 +93,7 @@
 from types import FrameType

 from pendulum.datetime import DateTime
 from sqlalchemy.orm import Query, Session
 from sqlalchemy.orm import Load, Query, Session

 from airflow._shared.logging.types import Logger
 from airflow.executors.base_executor import BaseExecutor
Expand All@@ -110,6 +110,31 @@
""":meta private:"""


def _eager_load_dag_run_for_validation() -> tuple[Load, Load]:
 """
 Eager-load DagRun relations required for execution API datamodel validation.

 When building TaskCallbackRequest with DRDataModel.model_validate(ti.dag_run),
 the consumed_asset_events collection and nested asset/source_aliases must be
 preloaded to avoid DetachedInstanceError after the session closes.

 Returns a tuple of two load options:
 - Asset loader: TI.dag_run → consumed_asset_events → asset
 - Alias loader: TI.dag_run → consumed_asset_events → source_aliases

 Example usage::

 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = select(TI).options(asset_loader).options(alias_loader)
 """
 # Traverse TI → dag_run → consumed_asset_events once, then branch to asset/aliases
 base = selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events)
 return (
 base.selectinload(AssetEvent.asset),
 base.selectinload(AssetEvent.source_aliases),
 )


def _get_current_dag(dag_id: str, session: Session) -> SerializedDAG | None:
 serdag = SerializedDagModel.get(dag_id=dag_id, session=session) # grabs the latest version
 if not serdag:
Expand DownExpand Up@@ -806,11 +831,12 @@ def process_executor_events(

 # Check state of finished tasks
 filter_for_tis = TI.filter_for_tis(tis_with_right_state)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = (
 select(TI)
 .where(filter_for_tis)
 .options(selectinload(TI.dag_model))
 .options(joinedload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(joinedload(TI.dag_run).selectinload(DagRun.created_dag_version))
 .options(joinedload(TI.dag_version))
 )
Expand DownExpand Up@@ -2375,10 +2401,12 @@ def _find_and_purge_task_instances_without_heartbeats(self) -> None:
 def _find_task_instances_without_heartbeats(self, *, session: Session) -> list[TI]:
 self.log.debug("Finding 'running' jobs without a recent heartbeat")
 limit_dttm = timezone.utcnow() - timedelta(seconds=self._task_instance_heartbeat_timeout_secs)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 task_instances_without_heartbeats = session.scalars(
 select(TI)
 .options(selectinload(TI.dag_model))
 .options(selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(alias_loader)
 .options(selectinload(TI.dag_version))
 .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
 .join(DM, TI.dag_id == DM.dag_id)
Expand Down

150 changes: 150 additions & 0 deletions
airflow-core/tests/unit/jobs/test_scheduler_job.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -608,6 +608,66 @@ def test_process_executor_events_ti_requeued(self, mock_stats_incr, mock_task_ca
 scheduler_job.executor.callback_sink.send.assert_not_called()
 mock_stats_incr.assert_not_called()

 @pytest.mark.usefixtures("testing_dag_bundle")
 @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr")
 def test_process_executor_events_with_asset_events(self, mock_stats_incr, session, dag_maker):
 """
 Test that _process_executor_events handles asset events without DetachedInstanceError.

 Regression test for scheduler crashes when task callbacks are built with
 consumed_asset_events that weren't eager-loaded.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset_executor", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_executor_events_with_assets", schedule=[asset1], fileloc="/test_path1/"):
 EmptyOperator(task_id="dummy_task", on_failure_callback=lambda ctx: None)

 dag = dag_maker.dag
 sync_dag_to_db(dag)
 DagVersion.get_latest_version(dag.dag_id)

 dr = dag_maker.create_dagrun()

 # Create asset event and attach to dag run
 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()
 dr.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor(do_update=False)
 scheduler_job = Job(executor=executor)
 self.job_runner = SchedulerJobRunner(scheduler_job)

 ti1 = dr.get_task_instance("dummy_task")
 ti1.state = State.QUEUED
 session.merge(ti1)
 session.commit()

 executor.event_buffer[ti1.key] = State.FAILED, None

 # This should not raise DetachedInstanceError
 self.job_runner._process_executor_events(executor=executor, session=session)

 ti1.refresh_from_db(session=session)
 assert ti1.state == State.FAILED

 # Verify callback was created with asset event data
 scheduler_job.executor.callback_sink.send.assert_called_once()
 callback_request = scheduler_job.executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 assert callback_request.context_from_server.dag_run.consumed_asset_events[0].asset.uri == asset1.uri

 def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker):
 dag_id = "SchedulerJobTest.test_execute_task_instances_is_paused_wont_execute"
 task_id_1 = "dummy_task"
Expand All@@ -628,6 +688,96 @@ def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker)
 assert ti1.state == State.SCHEDULED
 session.rollback()

 @pytest.mark.usefixtures("testing_dag_bundle")
 def test_find_and_purge_task_instances_without_heartbeats_with_asset_events(
 self, session, dag_maker, create_dagrun
 ):
 """
 Test that heartbeat purge succeeds when DagRun has consumed_asset_events.

 Regression test for DetachedInstanceError when building TaskCallbackRequest
 with asset event data after session expunge.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_heartbeat_with_assets", schedule=[asset1]):
 EmptyOperator(task_id="dummy_task")

 dag = dag_maker.dag
 scheduler_dag = sync_dag_to_db(dag)
 dag_v = DagVersion.get_latest_version(dag.dag_id)

 data_interval = infer_automated_data_interval(scheduler_dag.timetable, DEFAULT_LOGICAL_DATE)
 dag_run = create_dagrun(
 scheduler_dag,
 logical_date=DEFAULT_DATE,
 run_type=DagRunType.SCHEDULED,
 data_interval=data_interval,
 )

 # Create asset alias and event with full relationships
 asset_alias = AssetAliasModel(name="test_alias", group="test_group")
 session.add(asset_alias)
 session.flush()

 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()

 # Attach alias to event and event to dag run
 asset_event.source_aliases.append(asset_alias)
 dag_run.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor()
 scheduler_job = Job(executor=executor)
 with mock.patch("airflow.executors.executor_loader.ExecutorLoader.load_executor") as loader_mock:
 loader_mock.return_value = executor
 self.job_runner = SchedulerJobRunner(job=scheduler_job)

 ti = dag_run.get_task_instance("dummy_task")
 assert ti is not None # sanity check: dag_maker.create_dagrun created the TI

 ti.state = State.RUNNING
 ti.last_heartbeat_at = timezone.utcnow() - timedelta(minutes=6)
 ti.start_date = timezone.utcnow() - timedelta(minutes=10)
 ti.queued_by_job_id = scheduler_job.id
 ti.dag_version = dag_v
 session.merge(ti)
 session.flush()

 executor.running.add(ti.key)

 tis_without_heartbeats = self.job_runner._find_task_instances_without_heartbeats(session=session)
 assert len(tis_without_heartbeats) == 1
 ti_from_query = tis_without_heartbeats[0]
 ti_key = ti_from_query.key

 # Detach all ORM objects to mirror scheduler behaviour after session closes
 session.expunge_all()

 # This should not raise DetachedInstanceError now that eager loads are in place
 self.job_runner._purge_task_instances_without_heartbeats(tis_without_heartbeats, session=session)
 assert ti_key not in executor.running

 executor.callback_sink.send.assert_called_once()
 callback_request = executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 consumed_event = callback_request.context_from_server.dag_run.consumed_asset_events[0]
 assert consumed_event.asset.uri == asset1.uri
 assert len(consumed_event.source_aliases) == 1
 assert consumed_event.source_aliases[0].name == "test_alias"

 # @pytest.mark.usefixtures("mock_executor")
 def test_execute_task_instances_backfill_tasks_will_execute(self, dag_maker):
 """
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Auto-enable theater mode on YouTube
(function() {
function tryTheater() {
var btn = document.querySelector('button[aria-label="Theater mode"], ytd-player #player button[title="Theater mode"]');
if (btn && !btn.classList.contains('activated')) {
btn.click();
}
}
// Try immediately
tryTheater();
// Try after navigation (SPA)
var lastUrl = location.href;
setInterval(function() {
if (location.href !== lastUrl) {
lastUrl = location.href;
setTimeout(tryTheater, 500);
}
}, 1000);
// Also try on player load
var observer = new MutationObserver(tryTheater);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix(scheduler): load dag_run eagering for execution API validation by Lee-W · Pull Request #56919 · apache/airflow · GitHub
Skip to content
fix(scheduler): load dag_run eagering for execution API validation#56919

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

eladkal
merged 1 commit into

from

Oct 21, 2025
 Merged


fix(scheduler): load dag_run eagering for execution API validation
#56919
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 31 additions & 3 deletions
airflow-core/src/airflow/jobs/scheduler_job_runner.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -93,7 +93,7 @@
 from types import FrameType

 from pendulum.datetime import DateTime
 from sqlalchemy.orm import Query, Session
 from sqlalchemy.orm import Load, Query, Session

 from airflow._shared.logging.types import Logger
 from airflow.executors.base_executor import BaseExecutor
Expand All@@ -110,6 +110,31 @@
""":meta private:"""


def _eager_load_dag_run_for_validation() -> tuple[Load, Load]:
 """
 Eager-load DagRun relations required for execution API datamodel validation.

 When building TaskCallbackRequest with DRDataModel.model_validate(ti.dag_run),
 the consumed_asset_events collection and nested asset/source_aliases must be
 preloaded to avoid DetachedInstanceError after the session closes.

 Returns a tuple of two load options:
 - Asset loader: TI.dag_run → consumed_asset_events → asset
 - Alias loader: TI.dag_run → consumed_asset_events → source_aliases

 Example usage::

 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = select(TI).options(asset_loader).options(alias_loader)
 """
 # Traverse TI → dag_run → consumed_asset_events once, then branch to asset/aliases
 base = selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events)
 return (
 base.selectinload(AssetEvent.asset),
 base.selectinload(AssetEvent.source_aliases),
 )


def _get_current_dag(dag_id: str, session: Session) -> SerializedDAG | None:
 serdag = SerializedDagModel.get(dag_id=dag_id, session=session) # grabs the latest version
 if not serdag:
Expand DownExpand Up@@ -806,11 +831,12 @@ def process_executor_events(

 # Check state of finished tasks
 filter_for_tis = TI.filter_for_tis(tis_with_right_state)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = (
 select(TI)
 .where(filter_for_tis)
 .options(selectinload(TI.dag_model))
 .options(joinedload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(joinedload(TI.dag_run).selectinload(DagRun.created_dag_version))
 .options(joinedload(TI.dag_version))
 )
Expand DownExpand Up@@ -2375,10 +2401,12 @@ def _find_and_purge_task_instances_without_heartbeats(self) -> None:
 def _find_task_instances_without_heartbeats(self, *, session: Session) -> list[TI]:
 self.log.debug("Finding 'running' jobs without a recent heartbeat")
 limit_dttm = timezone.utcnow() - timedelta(seconds=self._task_instance_heartbeat_timeout_secs)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 task_instances_without_heartbeats = session.scalars(
 select(TI)
 .options(selectinload(TI.dag_model))
 .options(selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(alias_loader)
 .options(selectinload(TI.dag_version))
 .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
 .join(DM, TI.dag_id == DM.dag_id)
Expand Down

150 changes: 150 additions & 0 deletions
airflow-core/tests/unit/jobs/test_scheduler_job.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -608,6 +608,66 @@ def test_process_executor_events_ti_requeued(self, mock_stats_incr, mock_task_ca
 scheduler_job.executor.callback_sink.send.assert_not_called()
 mock_stats_incr.assert_not_called()

 @pytest.mark.usefixtures("testing_dag_bundle")
 @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr")
 def test_process_executor_events_with_asset_events(self, mock_stats_incr, session, dag_maker):
 """
 Test that _process_executor_events handles asset events without DetachedInstanceError.

 Regression test for scheduler crashes when task callbacks are built with
 consumed_asset_events that weren't eager-loaded.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset_executor", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_executor_events_with_assets", schedule=[asset1], fileloc="/test_path1/"):
 EmptyOperator(task_id="dummy_task", on_failure_callback=lambda ctx: None)

 dag = dag_maker.dag
 sync_dag_to_db(dag)
 DagVersion.get_latest_version(dag.dag_id)

 dr = dag_maker.create_dagrun()

 # Create asset event and attach to dag run
 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()
 dr.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor(do_update=False)
 scheduler_job = Job(executor=executor)
 self.job_runner = SchedulerJobRunner(scheduler_job)

 ti1 = dr.get_task_instance("dummy_task")
 ti1.state = State.QUEUED
 session.merge(ti1)
 session.commit()

 executor.event_buffer[ti1.key] = State.FAILED, None

 # This should not raise DetachedInstanceError
 self.job_runner._process_executor_events(executor=executor, session=session)

 ti1.refresh_from_db(session=session)
 assert ti1.state == State.FAILED

 # Verify callback was created with asset event data
 scheduler_job.executor.callback_sink.send.assert_called_once()
 callback_request = scheduler_job.executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 assert callback_request.context_from_server.dag_run.consumed_asset_events[0].asset.uri == asset1.uri

 def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker):
 dag_id = "SchedulerJobTest.test_execute_task_instances_is_paused_wont_execute"
 task_id_1 = "dummy_task"
Expand All@@ -628,6 +688,96 @@ def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker)
 assert ti1.state == State.SCHEDULED
 session.rollback()

 @pytest.mark.usefixtures("testing_dag_bundle")
 def test_find_and_purge_task_instances_without_heartbeats_with_asset_events(
 self, session, dag_maker, create_dagrun
 ):
 """
 Test that heartbeat purge succeeds when DagRun has consumed_asset_events.

 Regression test for DetachedInstanceError when building TaskCallbackRequest
 with asset event data after session expunge.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_heartbeat_with_assets", schedule=[asset1]):
 EmptyOperator(task_id="dummy_task")

 dag = dag_maker.dag
 scheduler_dag = sync_dag_to_db(dag)
 dag_v = DagVersion.get_latest_version(dag.dag_id)

 data_interval = infer_automated_data_interval(scheduler_dag.timetable, DEFAULT_LOGICAL_DATE)
 dag_run = create_dagrun(
 scheduler_dag,
 logical_date=DEFAULT_DATE,
 run_type=DagRunType.SCHEDULED,
 data_interval=data_interval,
 )

 # Create asset alias and event with full relationships
 asset_alias = AssetAliasModel(name="test_alias", group="test_group")
 session.add(asset_alias)
 session.flush()

 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()

 # Attach alias to event and event to dag run
 asset_event.source_aliases.append(asset_alias)
 dag_run.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor()
 scheduler_job = Job(executor=executor)
 with mock.patch("airflow.executors.executor_loader.ExecutorLoader.load_executor") as loader_mock:
 loader_mock.return_value = executor
 self.job_runner = SchedulerJobRunner(job=scheduler_job)

 ti = dag_run.get_task_instance("dummy_task")
 assert ti is not None # sanity check: dag_maker.create_dagrun created the TI

 ti.state = State.RUNNING
 ti.last_heartbeat_at = timezone.utcnow() - timedelta(minutes=6)
 ti.start_date = timezone.utcnow() - timedelta(minutes=10)
 ti.queued_by_job_id = scheduler_job.id
 ti.dag_version = dag_v
 session.merge(ti)
 session.flush()

 executor.running.add(ti.key)

 tis_without_heartbeats = self.job_runner._find_task_instances_without_heartbeats(session=session)
 assert len(tis_without_heartbeats) == 1
 ti_from_query = tis_without_heartbeats[0]
 ti_key = ti_from_query.key

 # Detach all ORM objects to mirror scheduler behaviour after session closes
 session.expunge_all()

 # This should not raise DetachedInstanceError now that eager loads are in place
 self.job_runner._purge_task_instances_without_heartbeats(tis_without_heartbeats, session=session)
 assert ti_key not in executor.running

 executor.callback_sink.send.assert_called_once()
 callback_request = executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 consumed_event = callback_request.context_from_server.dag_run.consumed_asset_events[0]
 assert consumed_event.asset.uri == asset1.uri
 assert len(consumed_event.source_aliases) == 1
 assert consumed_event.source_aliases[0].name == "test_alias"

 # @pytest.mark.usefixtures("mock_executor")
 def test_execute_task_instances_backfill_tasks_will_execute(self, dag_maker):
 """
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Remove or un-stick sticky/fixed headers that block content
(function() {
function unstick() {
document.querySelectorAll('header, nav, [role="banner"], .header, .navbar, .sticky, .fixed-top, [style*="position: fixed"], [style*="position:sticky"]').forEach(function(el) {
if (el.style.position === 'fixed' || el.style.position === 'sticky' || getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {
el.style.position = 'static';
el.style.top = 'auto';
el.style.zIndex = 'auto';
}
});
}
unstick();
var observer = new MutationObserver(unstick);
observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });
})();
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix(scheduler): load dag_run eagering for execution API validation by Lee-W · Pull Request #56919 · apache/airflow · GitHub
Skip to content
fix(scheduler): load dag_run eagering for execution API validation#56919

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

eladkal
merged 1 commit into

from

Oct 21, 2025
 Merged


fix(scheduler): load dag_run eagering for execution API validation
#56919
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 31 additions & 3 deletions
airflow-core/src/airflow/jobs/scheduler_job_runner.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -93,7 +93,7 @@
 from types import FrameType

 from pendulum.datetime import DateTime
 from sqlalchemy.orm import Query, Session
 from sqlalchemy.orm import Load, Query, Session

 from airflow._shared.logging.types import Logger
 from airflow.executors.base_executor import BaseExecutor
Expand All@@ -110,6 +110,31 @@
""":meta private:"""


def _eager_load_dag_run_for_validation() -> tuple[Load, Load]:
 """
 Eager-load DagRun relations required for execution API datamodel validation.

 When building TaskCallbackRequest with DRDataModel.model_validate(ti.dag_run),
 the consumed_asset_events collection and nested asset/source_aliases must be
 preloaded to avoid DetachedInstanceError after the session closes.

 Returns a tuple of two load options:
 - Asset loader: TI.dag_run → consumed_asset_events → asset
 - Alias loader: TI.dag_run → consumed_asset_events → source_aliases

 Example usage::

 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = select(TI).options(asset_loader).options(alias_loader)
 """
 # Traverse TI → dag_run → consumed_asset_events once, then branch to asset/aliases
 base = selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events)
 return (
 base.selectinload(AssetEvent.asset),
 base.selectinload(AssetEvent.source_aliases),
 )


def _get_current_dag(dag_id: str, session: Session) -> SerializedDAG | None:
 serdag = SerializedDagModel.get(dag_id=dag_id, session=session) # grabs the latest version
 if not serdag:
Expand DownExpand Up@@ -806,11 +831,12 @@ def process_executor_events(

 # Check state of finished tasks
 filter_for_tis = TI.filter_for_tis(tis_with_right_state)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = (
 select(TI)
 .where(filter_for_tis)
 .options(selectinload(TI.dag_model))
 .options(joinedload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(joinedload(TI.dag_run).selectinload(DagRun.created_dag_version))
 .options(joinedload(TI.dag_version))
 )
Expand DownExpand Up@@ -2375,10 +2401,12 @@ def _find_and_purge_task_instances_without_heartbeats(self) -> None:
 def _find_task_instances_without_heartbeats(self, *, session: Session) -> list[TI]:
 self.log.debug("Finding 'running' jobs without a recent heartbeat")
 limit_dttm = timezone.utcnow() - timedelta(seconds=self._task_instance_heartbeat_timeout_secs)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 task_instances_without_heartbeats = session.scalars(
 select(TI)
 .options(selectinload(TI.dag_model))
 .options(selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(alias_loader)
 .options(selectinload(TI.dag_version))
 .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
 .join(DM, TI.dag_id == DM.dag_id)
Expand Down

150 changes: 150 additions & 0 deletions
airflow-core/tests/unit/jobs/test_scheduler_job.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -608,6 +608,66 @@ def test_process_executor_events_ti_requeued(self, mock_stats_incr, mock_task_ca
 scheduler_job.executor.callback_sink.send.assert_not_called()
 mock_stats_incr.assert_not_called()

 @pytest.mark.usefixtures("testing_dag_bundle")
 @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr")
 def test_process_executor_events_with_asset_events(self, mock_stats_incr, session, dag_maker):
 """
 Test that _process_executor_events handles asset events without DetachedInstanceError.

 Regression test for scheduler crashes when task callbacks are built with
 consumed_asset_events that weren't eager-loaded.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset_executor", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_executor_events_with_assets", schedule=[asset1], fileloc="/test_path1/"):
 EmptyOperator(task_id="dummy_task", on_failure_callback=lambda ctx: None)

 dag = dag_maker.dag
 sync_dag_to_db(dag)
 DagVersion.get_latest_version(dag.dag_id)

 dr = dag_maker.create_dagrun()

 # Create asset event and attach to dag run
 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()
 dr.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor(do_update=False)
 scheduler_job = Job(executor=executor)
 self.job_runner = SchedulerJobRunner(scheduler_job)

 ti1 = dr.get_task_instance("dummy_task")
 ti1.state = State.QUEUED
 session.merge(ti1)
 session.commit()

 executor.event_buffer[ti1.key] = State.FAILED, None

 # This should not raise DetachedInstanceError
 self.job_runner._process_executor_events(executor=executor, session=session)

 ti1.refresh_from_db(session=session)
 assert ti1.state == State.FAILED

 # Verify callback was created with asset event data
 scheduler_job.executor.callback_sink.send.assert_called_once()
 callback_request = scheduler_job.executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 assert callback_request.context_from_server.dag_run.consumed_asset_events[0].asset.uri == asset1.uri

 def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker):
 dag_id = "SchedulerJobTest.test_execute_task_instances_is_paused_wont_execute"
 task_id_1 = "dummy_task"
Expand All@@ -628,6 +688,96 @@ def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker)
 assert ti1.state == State.SCHEDULED
 session.rollback()

 @pytest.mark.usefixtures("testing_dag_bundle")
 def test_find_and_purge_task_instances_without_heartbeats_with_asset_events(
 self, session, dag_maker, create_dagrun
 ):
 """
 Test that heartbeat purge succeeds when DagRun has consumed_asset_events.

 Regression test for DetachedInstanceError when building TaskCallbackRequest
 with asset event data after session expunge.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_heartbeat_with_assets", schedule=[asset1]):
 EmptyOperator(task_id="dummy_task")

 dag = dag_maker.dag
 scheduler_dag = sync_dag_to_db(dag)
 dag_v = DagVersion.get_latest_version(dag.dag_id)

 data_interval = infer_automated_data_interval(scheduler_dag.timetable, DEFAULT_LOGICAL_DATE)
 dag_run = create_dagrun(
 scheduler_dag,
 logical_date=DEFAULT_DATE,
 run_type=DagRunType.SCHEDULED,
 data_interval=data_interval,
 )

 # Create asset alias and event with full relationships
 asset_alias = AssetAliasModel(name="test_alias", group="test_group")
 session.add(asset_alias)
 session.flush()

 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()

 # Attach alias to event and event to dag run
 asset_event.source_aliases.append(asset_alias)
 dag_run.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor()
 scheduler_job = Job(executor=executor)
 with mock.patch("airflow.executors.executor_loader.ExecutorLoader.load_executor") as loader_mock:
 loader_mock.return_value = executor
 self.job_runner = SchedulerJobRunner(job=scheduler_job)

 ti = dag_run.get_task_instance("dummy_task")
 assert ti is not None # sanity check: dag_maker.create_dagrun created the TI

 ti.state = State.RUNNING
 ti.last_heartbeat_at = timezone.utcnow() - timedelta(minutes=6)
 ti.start_date = timezone.utcnow() - timedelta(minutes=10)
 ti.queued_by_job_id = scheduler_job.id
 ti.dag_version = dag_v
 session.merge(ti)
 session.flush()

 executor.running.add(ti.key)

 tis_without_heartbeats = self.job_runner._find_task_instances_without_heartbeats(session=session)
 assert len(tis_without_heartbeats) == 1
 ti_from_query = tis_without_heartbeats[0]
 ti_key = ti_from_query.key

 # Detach all ORM objects to mirror scheduler behaviour after session closes
 session.expunge_all()

 # This should not raise DetachedInstanceError now that eager loads are in place
 self.job_runner._purge_task_instances_without_heartbeats(tis_without_heartbeats, session=session)
 assert ti_key not in executor.running

 executor.callback_sink.send.assert_called_once()
 callback_request = executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 consumed_event = callback_request.context_from_server.dag_run.consumed_asset_events[0]
 assert consumed_event.asset.uri == asset1.uri
 assert len(consumed_event.source_aliases) == 1
 assert consumed_event.source_aliases[0].name == "test_alias"

 # @pytest.mark.usefixtures("mock_executor")
 def test_execute_task_instances_backfill_tasks_will_execute(self, dag_maker):
 """
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Universal Dark Mode - works on any site
(function() {
var enabled = true;
function applyDarkMode() {
if (!enabled) return;
// Create style element if it doesn't exist
var style = document.getElementById('universal-dark-mode-style');
if (!style) {
style = document.createElement('style');
style.id = 'universal-dark-mode-style';
document.head.appendChild(style);
}
// Dark mode CSS - inverts colors but preserves images/video
style.textContent = '
/* Invert everything except media */
html {
filter: invert(1) hue-rotate(180deg) !important;
background: #1a1a2e !important;
}
/* Restore images, videos, iframes, canvas */
img, video, iframe, canvas, svg, picture, [style*="background-image"] {
filter: invert(1) hue-rotate(180deg) !important;
}
/* Preserve specific elements that should not be inverted */
.no-dark-mode, .no-dark-mode *,
[data-theme="light"], [data-theme="light"],
.ace_editor, .ace_editor *,
.CodeMirror, .CodeMirror *,
.monaco-editor, .monaco-editor *,
.markdown-body pre, .markdown-body pre *,
.highlight, .highlight *,
pre code, pre code * {
filter: none !important;
}
/* Fix common UI elements */
.modal, .popup, .dropdown-menu, .tooltip, .popover {
filter: invert(1) hue-rotate(180deg) !important;
background: #2d2d44 !important;
border-color: #444 !important;
}
/* Scrollbars */
::-webkit-scrollbar { background: #1a1a2e !important; }
::-webkit-scrollbar-thumb { background: #444 !important; }
::-webkit-scrollbar-thumb:hover { background: #555 !important; }
/* Selection */
::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
';
}
function removeDarkMode() {
var style = document.getElementById('universal-dark-mode-style');
if (style) style.remove();
}
// Toggle with Alt+Shift+D
document.addEventListener('keydown', function(e) {
if (e.altKey && e.shiftKey && e.key === 'D') {
e.preventDefault();
enabled = !enabled;
if (enabled) {
applyDarkMode();
console.log('[Universal Dark Mode] Enabled');
} else {
removeDarkMode();
console.log('[Universal Dark Mode] Disabled');
}
}
});
// Apply on load
applyDarkMode();
// Re-apply on dynamic content
var observer = new MutationObserver(function(mutations) {
if (enabled && !document.getElementById('universal-dark-mode-style')) {
applyDarkMode();
}
});
observer.observe(document.head, { childList: true });
console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');
})();
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();
fix(scheduler): load dag_run eagering for execution API validation by Lee-W · Pull Request #56919 · apache/airflow · GitHub
Skip to content
fix(scheduler): load dag_run eagering for execution API validation#56919

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

eladkal
merged 1 commit into

from

Oct 21, 2025
 Merged


fix(scheduler): load dag_run eagering for execution API validation
#56919
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 31 additions & 3 deletions
airflow-core/src/airflow/jobs/scheduler_job_runner.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -93,7 +93,7 @@
 from types import FrameType

 from pendulum.datetime import DateTime
 from sqlalchemy.orm import Query, Session
 from sqlalchemy.orm import Load, Query, Session

 from airflow._shared.logging.types import Logger
 from airflow.executors.base_executor import BaseExecutor
Expand All@@ -110,6 +110,31 @@
""":meta private:"""


def _eager_load_dag_run_for_validation() -> tuple[Load, Load]:
 """
 Eager-load DagRun relations required for execution API datamodel validation.

 When building TaskCallbackRequest with DRDataModel.model_validate(ti.dag_run),
 the consumed_asset_events collection and nested asset/source_aliases must be
 preloaded to avoid DetachedInstanceError after the session closes.

 Returns a tuple of two load options:
 - Asset loader: TI.dag_run → consumed_asset_events → asset
 - Alias loader: TI.dag_run → consumed_asset_events → source_aliases

 Example usage::

 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = select(TI).options(asset_loader).options(alias_loader)
 """
 # Traverse TI → dag_run → consumed_asset_events once, then branch to asset/aliases
 base = selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events)
 return (
 base.selectinload(AssetEvent.asset),
 base.selectinload(AssetEvent.source_aliases),
 )


def _get_current_dag(dag_id: str, session: Session) -> SerializedDAG | None:
 serdag = SerializedDagModel.get(dag_id=dag_id, session=session) # grabs the latest version
 if not serdag:
Expand DownExpand Up@@ -806,11 +831,12 @@ def process_executor_events(

 # Check state of finished tasks
 filter_for_tis = TI.filter_for_tis(tis_with_right_state)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 query = (
 select(TI)
 .where(filter_for_tis)
 .options(selectinload(TI.dag_model))
 .options(joinedload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(joinedload(TI.dag_run).selectinload(DagRun.created_dag_version))
 .options(joinedload(TI.dag_version))
 )
Expand DownExpand Up@@ -2375,10 +2401,12 @@ def _find_and_purge_task_instances_without_heartbeats(self) -> None:
 def _find_task_instances_without_heartbeats(self, *, session: Session) -> list[TI]:
 self.log.debug("Finding 'running' jobs without a recent heartbeat")
 limit_dttm = timezone.utcnow() - timedelta(seconds=self._task_instance_heartbeat_timeout_secs)
 asset_loader, alias_loader = _eager_load_dag_run_for_validation()
 task_instances_without_heartbeats = session.scalars(
 select(TI)
 .options(selectinload(TI.dag_model))
 .options(selectinload(TI.dag_run).selectinload(DagRun.consumed_asset_events))
 .options(asset_loader)
 .options(alias_loader)
 .options(selectinload(TI.dag_version))
 .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
 .join(DM, TI.dag_id == DM.dag_id)
Expand Down

150 changes: 150 additions & 0 deletions
airflow-core/tests/unit/jobs/test_scheduler_job.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -608,6 +608,66 @@ def test_process_executor_events_ti_requeued(self, mock_stats_incr, mock_task_ca
 scheduler_job.executor.callback_sink.send.assert_not_called()
 mock_stats_incr.assert_not_called()

 @pytest.mark.usefixtures("testing_dag_bundle")
 @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr")
 def test_process_executor_events_with_asset_events(self, mock_stats_incr, session, dag_maker):
 """
 Test that _process_executor_events handles asset events without DetachedInstanceError.

 Regression test for scheduler crashes when task callbacks are built with
 consumed_asset_events that weren't eager-loaded.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset_executor", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_executor_events_with_assets", schedule=[asset1], fileloc="/test_path1/"):
 EmptyOperator(task_id="dummy_task", on_failure_callback=lambda ctx: None)

 dag = dag_maker.dag
 sync_dag_to_db(dag)
 DagVersion.get_latest_version(dag.dag_id)

 dr = dag_maker.create_dagrun()

 # Create asset event and attach to dag run
 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()
 dr.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor(do_update=False)
 scheduler_job = Job(executor=executor)
 self.job_runner = SchedulerJobRunner(scheduler_job)

 ti1 = dr.get_task_instance("dummy_task")
 ti1.state = State.QUEUED
 session.merge(ti1)
 session.commit()

 executor.event_buffer[ti1.key] = State.FAILED, None

 # This should not raise DetachedInstanceError
 self.job_runner._process_executor_events(executor=executor, session=session)

 ti1.refresh_from_db(session=session)
 assert ti1.state == State.FAILED

 # Verify callback was created with asset event data
 scheduler_job.executor.callback_sink.send.assert_called_once()
 callback_request = scheduler_job.executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 assert callback_request.context_from_server.dag_run.consumed_asset_events[0].asset.uri == asset1.uri

 def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker):
 dag_id = "SchedulerJobTest.test_execute_task_instances_is_paused_wont_execute"
 task_id_1 = "dummy_task"
Expand All@@ -628,6 +688,96 @@ def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker)
 assert ti1.state == State.SCHEDULED
 session.rollback()

 @pytest.mark.usefixtures("testing_dag_bundle")
 def test_find_and_purge_task_instances_without_heartbeats_with_asset_events(
 self, session, dag_maker, create_dagrun
 ):
 """
 Test that heartbeat purge succeeds when DagRun has consumed_asset_events.

 Regression test for DetachedInstanceError when building TaskCallbackRequest
 with asset event data after session expunge.
 """
 asset1 = Asset(uri="test://asset1", name="test_asset", group="test_group")
 asset_model = AssetModel(name=asset1.name, uri=asset1.uri, group=asset1.group)
 session.add(asset_model)
 session.flush()

 with dag_maker(dag_id="test_heartbeat_with_assets", schedule=[asset1]):
 EmptyOperator(task_id="dummy_task")

 dag = dag_maker.dag
 scheduler_dag = sync_dag_to_db(dag)
 dag_v = DagVersion.get_latest_version(dag.dag_id)

 data_interval = infer_automated_data_interval(scheduler_dag.timetable, DEFAULT_LOGICAL_DATE)
 dag_run = create_dagrun(
 scheduler_dag,
 logical_date=DEFAULT_DATE,
 run_type=DagRunType.SCHEDULED,
 data_interval=data_interval,
 )

 # Create asset alias and event with full relationships
 asset_alias = AssetAliasModel(name="test_alias", group="test_group")
 session.add(asset_alias)
 session.flush()

 asset_event = AssetEvent(
 asset_id=asset_model.id,
 source_task_id="upstream_task",
 source_dag_id="upstream_dag",
 source_run_id="upstream_run",
 source_map_index=-1,
 )
 session.add(asset_event)
 session.flush()

 # Attach alias to event and event to dag run
 asset_event.source_aliases.append(asset_alias)
 dag_run.consumed_asset_events.append(asset_event)
 session.flush()

 executor = MockExecutor()
 scheduler_job = Job(executor=executor)
 with mock.patch("airflow.executors.executor_loader.ExecutorLoader.load_executor") as loader_mock:
 loader_mock.return_value = executor
 self.job_runner = SchedulerJobRunner(job=scheduler_job)

 ti = dag_run.get_task_instance("dummy_task")
 assert ti is not None # sanity check: dag_maker.create_dagrun created the TI

 ti.state = State.RUNNING
 ti.last_heartbeat_at = timezone.utcnow() - timedelta(minutes=6)
 ti.start_date = timezone.utcnow() - timedelta(minutes=10)
 ti.queued_by_job_id = scheduler_job.id
 ti.dag_version = dag_v
 session.merge(ti)
 session.flush()

 executor.running.add(ti.key)

 tis_without_heartbeats = self.job_runner._find_task_instances_without_heartbeats(session=session)
 assert len(tis_without_heartbeats) == 1
 ti_from_query = tis_without_heartbeats[0]
 ti_key = ti_from_query.key

 # Detach all ORM objects to mirror scheduler behaviour after session closes
 session.expunge_all()

 # This should not raise DetachedInstanceError now that eager loads are in place
 self.job_runner._purge_task_instances_without_heartbeats(tis_without_heartbeats, session=session)
 assert ti_key not in executor.running

 executor.callback_sink.send.assert_called_once()
 callback_request = executor.callback_sink.send.call_args.args[0]
 assert callback_request.context_from_server is not None
 assert len(callback_request.context_from_server.dag_run.consumed_asset_events) == 1
 consumed_event = callback_request.context_from_server.dag_run.consumed_asset_events[0]
 assert consumed_event.asset.uri == asset1.uri
 assert len(consumed_event.source_aliases) == 1
 assert consumed_event.source_aliases[0].name == "test_alias"

 # @pytest.mark.usefixtures("mock_executor")
 def test_execute_task_instances_backfill_tasks_will_execute(self, dag_maker):
 """
Expand Down