Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
from sqlalchemy.sql import select
from structlog.contextvars import bind_contextvars

from airflow._shared.observability.metrics import stats
from airflow._shared.observability.traces import override_ids
from airflow._shared.state import TaskScope
from airflow._shared.timezones import timezone
Expand Down Expand Up @@ -159,6 +160,8 @@ def ti_run(
TI.try_number,
TI.max_tries,
TI.start_date,
TI.queue,
TI.queued_dttm,
TI.next_method,
TI.hostname,
TI.unixname,
Expand Down Expand Up @@ -199,6 +202,10 @@ def ti_run(
query = update(TI).where(TI.id == task_instance_id).values(data)

previous_state = ti.state
# Set on a genuine QUEUED -> RUNNING transition so task.queued_duration is emitted once the
# DagRun is loaded below (its stats_tags supply the tags). A duplicate start request falls
# through the branch below without raising, so the flag keeps it from emitting twice.
emit_queued_duration = False

# If we are already running, but this is a duplicate request from the same client return the same OK
# -- it's possible there was a network glitch and they never got the response
Expand Down Expand Up @@ -242,6 +249,17 @@ def ti_run(
extra=json.dumps({"host_name": ti_run_payload.hostname}) if ti_run_payload.hostname else None,
)
)
# Aims at one sample per try: the scheduler refreshes queued_dttm on every queueing, so a
# retry measures its own wait, while a resume from deferral is the same try continuing and
# is skipped via next_method (set on the DEFERRED transition, left in place until resume).
# Two known gaps in that reading: an operator with start_from_trigger=True is deferred
# straight from SCHEDULED without ever being queued, so the resume leg skipped here is its
# only real wait and it goes unmeasured; and this disagrees with task.scheduled_duration on
# retries, which emit_state_change_metric skips because the previous attempt's end_date is
# still on the row when the scheduler queues the TI.
# queued_dttm is None only in rare races and test setups.
emit_queued_duration = ti.queued_dttm is not None and ti.next_method is None

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This guard came in at @henry3260's request and it does what it says, but it leaves task.queued_duration and its sibling task.scheduled_duration disagreeing about which transitions count, in opposite directions. Taking the three ways a TI reaches RUNNING:

  • First run: both metrics emit.
  • Retry: only queued_duration. schedule_tis sets SCHEDULED, scheduled_dttm and try_number without clearing end_date, so the previous attempt's end_date is still on the row when the scheduler queues the TI, and emit_state_change_metric returns early.
  • Deferral resume: only scheduled_duration. The DEFERRED update never touches end_date and ti_run already set it to None, so that guard passes, and the trigger refreshes scheduled_dttm, so the sample is a real measurement.

The resume's queue wait is equally real: queued_dttm is refreshed on every queueing and the critical section selects SCHEDULED with no next_method filter, so a deferrable task sits in QUEUED again waiting for a worker slot before execute_complete runs. Skipping it means that wait is never measured, and for sensor-heavy deployments the resume leg is where most of the queue time lives.

I would drop and ti.next_method is None and emit per queue wait, which is also where @ashb's end_date change pointed: one sample per real wait rather than one per try. The per-try reading is defensible too, but then the retry case should not emit either, and the comment should name the axis so the asymmetry reads as deliberate.

Either answer works for me, and this is the only thing I would like settled before merge. Adding these samples after release shifts percentiles for anyone alerting on the timer, which is cheap to decide now and awkward to change later.

On why this did not come up in my last pass: I was looking at the tag set then, and only walked the resume path through the scheduler this time.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

start_from_trigger=True operators are deferred straight from SCHEDULED, never queued, so
the resume this guard skips is their only queue wait — they emit nothing at all today.
Test added.

Also fixed my comment: legacy skipped retries and emitted on resume (it runs before
end_date is cleared), the inverse of what I claimed.

Per queue wait is the only axis that measures those, so that's my vote — your call.


# Ensure there is no end date set and clear retry policy overrides from the previous attempt.
query = query.values(
end_date=None,
Expand Down Expand Up @@ -297,7 +315,21 @@ def ti_run(
or 0
)

dr.team_name = get_team_name_for_ti(task_instance_id, session)
team_name = get_team_name_for_ti(task_instance_id, session)
dr.team_name = team_name

if emit_queued_duration:
# Tag via dr.stats_tags so this stays sliceable the same way as its sibling
# task.scheduled_duration, which emit_state_change_metric sends as
# {**ti.stats_tags, "queue": ti.queue} -- that is dag_run.stats_tags plus task_id.
# Team lives on the Bundle rather than the DagRun schema, so stats_tags cannot resolve
# it here; add the value looked up above instead. Falsy values are pruned from
# stats_tags, so only set it when there is a team. The registry-derived legacy name
# dag.<dag_id>.<task_id>.queued_duration is emitted by stats.timing automatically.
tags = {**dr.stats_tags, "task_id": ti.task_id, "queue": ti.queue}

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When metrics.dag_tags_in_metrics is on, dr.stats_tags costs two lazy loads per task start here: dag_tags_for_stats touches self.dag_model and then dag_model.tags, and the dr select above only eager-loads consumed_asset_events. The earlier query does join DagModel, but it selects DagModel.owners as a column rather than the entity, so nothing lands in the identity map and the lazy load still fires.

get_running_dag_runs_to_examine eager-loads exactly this to keep it out of the scheduler loop, and the comment on dag_tags_for_stats calls the remaining paths low frequency. ti_run is once per task start on the execution API, so it belongs with the first group. Adding the same eager load to the dr select, conditional on the config so the join is not paid when the feature is off, would cover it.

Non-blocking, and the config is off by default.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed — the earlier select takes DagModel.owners as a column, not the entity. Left out
to keep the diff on the blocking question; happy to add it here or in a follow-up.

if team_name:
tags["team_name"] = team_name

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On the _team_name point from your last reply: it is the convention the scheduler already uses for exactly this, ti.dag_run._team_name = team plus two more sites, and stats_tags reads it through getattr. Setting dr._team_name = team_name alongside the dr.team_name = team_name above would make these three lines and the comment redundant, and keeps the tag set derived in one place if stats_tags grows again. Optional, since what you have produces the same tags today.

On the test: yes please, pin the multi-team leg with conf_vars. Your read of the gap is right, and get_team_name_for_ti returns None unless core.multi_team is on, so today that key agrees for the wrong reason on both sides.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Test added with conf_vars, and it fails if the team tag is dropped.

_team_name: agreed — left out since it touches the same lines as the axis decision.

stats.timing("task.queued_duration", timezone.utcnow() - ti.queued_dttm, tags=tags)

context = TIRunContext(
dag_run=dr,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1072,6 +1072,271 @@ def test_ti_run_creates_audit_log(self, client, session, create_task_instance, t
assert logs[0].owner == ti.task.owner
assert logs[0].extra == '{"host_name": "random-hostname"}'

@pytest.mark.parametrize(
"scenario",
["first_run", "retry"],
)
def test_ti_run_emits_queued_duration_metric(
self, client, session, create_task_instance, time_machine, scenario
):
"""task.queued_duration is emitted on a real QUEUED -> RUNNING transition.

The scheduler refreshes queued_dttm every time it queues a task, so a retry has a
fresh queue wait just like a first run and must emit too. A retry still carries the
previous attempt's end_date on the row when ti_run is reached, so this asserts the
emit does not depend on end_date being unset.
"""
queued_at = timezone.parse("2024-09-30T12:00:00Z")
run_at = queued_at.add(seconds=42)

ti = create_task_instance(
task_id=f"test_ti_run_emits_queued_duration_metric_{scenario}",
state=State.QUEUED,
dagrun_state=DagRunState.RUNNING,
session=session,
start_date=queued_at,
dag_id=str(uuid4()),
)
ti.queued_dttm = queued_at
ti.queue = "default"
if scenario == "retry":
# A retried TI still has the previous attempt's end_date set on the row until
# ti_run clears it; the metric must fire regardless.
ti.end_date = queued_at.add(seconds=10)
session.commit()

# The metric has to stay sliceable the same way as its sibling task.scheduled_duration,
# which emit_state_change_metric sends as {**ti.stats_tags, "queue": ti.queue}. Deriving
# the expectation from that same expression makes the two drift apart only if this fails.
expected_tags = {**ti.stats_tags, "queue": ti.queue}
assert "run_type" in expected_tags

time_machine.move_to(run_at, tick=False)

with mock.patch("airflow.api_fastapi.execution_api.routes.task_instances.stats") as mock_stats:
response = client.patch(
f"/execution/task-instances/{ti.id}/run",
json={
"state": "running",
"hostname": "random-hostname",
"unixname": "random-unixname",
"pid": 100,
"start_date": run_at.isoformat(),
},
)

assert response.status_code == 200
mock_stats.timing.assert_called_once_with(
"task.queued_duration",
run_at - queued_at,
tags=expected_tags,
)

@pytest.mark.parametrize(
"skip_reason",
["deferral_resume", "queued_dttm_missing"],
)
def test_ti_run_skips_queued_duration_metric(
self, client, session, create_task_instance, time_machine, skip_reason
):
"""task.queued_duration is skipped on a resume from deferral (next_method set, so
the queue wait belongs to the same try already counted) and when queued_dttm was
not recorded (rare race / test setups)."""
queued_at = timezone.parse("2024-09-30T12:00:00Z")
run_at = queued_at.add(seconds=42)
time_machine.move_to(run_at, tick=False)

ti = create_task_instance(
task_id=f"test_ti_run_skips_queued_duration_metric_{skip_reason}",
state=State.QUEUED,
dagrun_state=DagRunState.RUNNING,
session=session,
start_date=queued_at,
dag_id=str(uuid4()),
)
if skip_reason == "deferral_resume":
ti.queued_dttm = queued_at
ti.next_method = "execute_complete"
else:
ti.queued_dttm = None
session.commit()

with mock.patch("airflow.api_fastapi.execution_api.routes.task_instances.stats") as mock_stats:
response = client.patch(
f"/execution/task-instances/{ti.id}/run",
json={
"state": "running",
"hostname": "random-hostname",
"unixname": "random-unixname",
"pid": 100,
"start_date": run_at.isoformat(),
},
)

assert response.status_code == 200
mock_stats.timing.assert_not_called()

def test_ti_run_skips_queued_duration_metric_on_duplicate_start(
self, client, session, create_task_instance, time_machine
):
"""A start request replayed by the same worker after a network glitch returns the same
200 without re-running the transition, so it must not double-count the queue wait."""
queued_at = timezone.parse("2024-09-30T12:00:00Z")
run_at = queued_at.add(seconds=42)
time_machine.move_to(run_at, tick=False)

ti = create_task_instance(
task_id="test_ti_run_skips_queued_duration_metric_on_duplicate_start",
state=State.RUNNING,
dagrun_state=DagRunState.RUNNING,
session=session,
start_date=queued_at,
dag_id=str(uuid4()),
)
ti.queued_dttm = queued_at
ti.hostname = "random-hostname"
ti.unixname = "random-unixname"
ti.pid = 100
session.commit()

with mock.patch("airflow.api_fastapi.execution_api.routes.task_instances.stats") as mock_stats:
response = client.patch(
f"/execution/task-instances/{ti.id}/run",
json={
"state": "running",
"hostname": "random-hostname",
"unixname": "random-unixname",
"pid": 100,
"start_date": run_at.isoformat(),
},
)

assert response.status_code == 200
mock_stats.timing.assert_not_called()

def test_ti_run_queued_duration_metric_is_tagged_with_team_name(
self, client, session, dag_maker, time_machine
):
"""The team tag has to be pinned with multi-team on: with it off both the metric and any
expectation derived from stats_tags omit the key, so they would agree for the wrong reason.
"""
from airflow.models.dagbundle import DagBundleModel
from airflow.models.team import Team

queued_at = timezone.parse("2024-09-30T12:00:00Z")
run_at = queued_at.add(seconds=42)
time_machine.move_to(queued_at, tick=False)

dag_id = str(uuid4())
with dag_maker(dag_id=dag_id, session=session):
EmptyOperator(task_id="task")
dr = dag_maker.create_dagrun(
run_id="test", logical_date=queued_at, state=DagRunState.RUNNING, start_date=queued_at
)
ti = dr.get_task_instance(task_id="task")
session.execute(
update(TaskInstance)
.where(TaskInstance.id == ti.id)
.values(state=State.QUEUED, queued_dttm=queued_at, queue="default")
)

bundle_name = f"bundle-{dag_id}"
team_name = f"team-{dag_id[:8]}"
bundle = DagBundleModel(name=bundle_name)
bundle.teams.append(Team(name=team_name))
session.add(bundle)
session.flush()
session.execute(update(DagModel).where(DagModel.dag_id == dag_id).values(bundle_name=bundle_name))
session.commit()

time_machine.move_to(run_at, tick=False)

with conf_vars({("core", "multi_team"): "True"}):
with mock.patch("airflow.api_fastapi.execution_api.routes.task_instances.stats") as mock_stats:
response = client.patch(
f"/execution/task-instances/{ti.id}/run",
json={
"state": "running",
"hostname": "random-hostname",
"unixname": "random-unixname",
"pid": 100,
"start_date": run_at.isoformat(),
},
)

assert response.status_code == 200
mock_stats.timing.assert_called_once()
assert mock_stats.timing.call_args.kwargs["tags"]["team_name"] == team_name

def test_ti_run_emits_no_queued_duration_for_start_from_trigger(
self, client, session, dag_maker, time_machine
):
"""A start_from_trigger operator is deferred straight from SCHEDULED, so the resume that the
next_method guard skips is its only trip through the queue and nothing is ever measured.
"""
from airflow.sdk import BaseOperator
from airflow.triggers.base import StartTriggerArgs

class StartFromTriggerOperator(BaseOperator):
start_trigger_args = StartTriggerArgs(
trigger_cls="airflow.triggers.testing.SuccessTrigger",
trigger_kwargs={},
next_method="execute_complete",
timeout=None,
)
start_from_trigger = True

def execute_complete(self):
pass

queued_at = timezone.parse("2024-09-30T12:00:00Z")
run_at = queued_at.add(seconds=42)
time_machine.move_to(queued_at, tick=False)

with dag_maker(dag_id=str(uuid4()), session=session):
StartFromTriggerOperator(task_id="task")
dr = dag_maker.create_dagrun(run_id="test", logical_date=queued_at, state=DagRunState.RUNNING)

ti = dr.get_task_instance(task_id="task")
ti.task = dr.dag.get_task("task")
assert dr.schedule_tis((ti,), session=session) == 0
session.merge(ti)
session.commit()

# Read the row back rather than the in-memory TI: what ti_run sees is the persisted state,
# and the point of this test is that the deferral lands without the TI ever being queued.
state, next_method, queued_dttm = session.execute(
select(TaskInstance.state, TaskInstance.next_method, TaskInstance.queued_dttm).where(
TaskInstance.id == ti.id
)
).one()
assert (state, next_method, queued_dttm) == (TaskInstanceState.DEFERRED, "execute_complete", None)

# The trigger fires and the scheduler queues it -- the first and only real queue wait.
session.execute(
update(TaskInstance)
.where(TaskInstance.id == ti.id)
.values(state=State.QUEUED, queued_dttm=queued_at, queue="default")
)
session.commit()

time_machine.move_to(run_at, tick=False)

with mock.patch("airflow.api_fastapi.execution_api.routes.task_instances.stats") as mock_stats:
response = client.patch(
f"/execution/task-instances/{ti.id}/run",
json={
"state": "running",
"hostname": "random-hostname",
"unixname": "random-unixname",
"pid": 100,
"start_date": run_at.isoformat(),
},
)

assert response.status_code == 200
assert mock_stats.timing.call_args_list == []


class TestTIUpdateState:
def setup_method(self):
Expand Down Expand Up @@ -1537,6 +1802,8 @@ def test_ti_run_database_error(self, client, session, create_task_instance):
try_number=1,
max_tries=0,
start_date=None,
queue="default",
queued_dttm=timezone.utcnow(),
next_method=None,
hostname=None,
unixname=None,
Expand Down