Skip to content

Avoid a supervisor round trip before poke for rescheduled sensors - #68012

Merged
kaxil merged 5 commits into
apache:mainfrom
hkc-8010:fix/reschedule-start-date-context
Oct 12, 2026
Merged

kaxil merged 5 commits into
apache:mainfrom
hkc-8010:fix/reschedule-start-date-context

Conversation

@hkc-8010

@hkc-8010 hkc-8010 commented Jun 4, 2026 •

Copy link
Copy Markdown
Contributor

Rescheduled tasks currently fetch their first task reschedule start date from the supervisor when RuntimeTaskInstance.get_first_reschedule_date() is called. For sensors, this can happen before poke() starts, which means a worker can block on the supervisor/API round trip before emitting the sensor's usual poke log line.

This PR includes the first task reschedule start date in the task instance run context returned by the Execution API. The Task SDK uses that value directly when present, while keeping the existing supervisor request as a compatibility fallback for older API responses.

closes: #68010

Tests:

  • breeze testing task-sdk-tests --python 3.10 -- task-sdk/tests/task_sdk/execution_time/test_task_runner.py -k get_first_reschedule_date -q
  • breeze testing core-tests --python 3.10 --db-reset -- airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_06_30/test_task_instances.py -q
  • breeze testing core-tests --backend postgres --python 3.10 --db-reset -- airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_06_30/test_task_instances.py -q
  • breeze testing task-sdk-tests --python 3.10
  • prek run --files <changed-files>
  • prek run --files <changed-files> --stage manual

Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

Generated-by: OpenAI Codex following the guidelines


  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

Important

🛠️ Maintainer triage note for @hkc-8010 · by @potiuk · 2026-08-13 12:55 UTC

Helpful heads-up from the maintainers — please address before this PR can be reviewed:

  • ❌ Provider tests. See docs.

Full list of what we check: Pull Request quality criteria.

The ball is in your court — you've been assigned to this PR. Fix the above, then mark it Ready for review.

Automated triage — may be imperfect; a maintainer takes the next look.

@boring-cyborg boring-cyborg Bot added area:API Airflow's REST/HTTP API area:task-sdk labels Jun 4, 2026
@eladkal eladkal added this to the Airflow 3.3.0 milestone Jun 7, 2026
@eladkal eladkal added the type:bug-fix Changelog: Bug Fixes label Jun 7, 2026
@potiuk

potiuk commented Jun 9, 2026

Copy link
Copy Markdown
Member

@hkc-8010 A few things need addressing before review — see our Pull Request quality criteria.

  • ❌ Merge conflicts with main. See docs.

No rush.

Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you.

@hkc-8010
hkc-8010 force-pushed the fix/reschedule-start-date-context branch from 08823ae to f3bd4aa Compare June 12, 2026 08:56
@hkc-8010

Copy link
Copy Markdown
Contributor Author

@potiuk thanks for the heads up. I rebased the branch on current main, resolved the conflicts, and force-pushed the updated commit (f3bd4aae0da6ad23a2b06500b6c9354dec140a77).

The local pre-push hooks are passing now, including the supervisor schema snapshot/version checks that caught one missing generated snapshot update during the rebase. GitHub checks have restarted on the updated commit.

Could you please take another look when you get a chance?

@potiuk potiuk added the ready for maintainer review Set by triage when all criteria pass, or by the fork CI reconciler when fork CI is green. label Jun 17, 2026
@ephraimbuddy
ephraimbuddy requested a review from Copilot June 29, 2026 14:15

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

This PR aims to prevent rescheduled sensors from blocking before poke() by including the first TaskReschedule.start_date in the Execution API’s TIRunContext, and teaching the Task SDK to use that value when present (with a supervisor/API fallback for older server responses).

Changes:

  • Add first_task_reschedule_start_date to TIRunContext (core + Task SDK generated models/schema) and expose it from the /execution/task-instances/{id}/run response.
  • Update RuntimeTaskInstance.get_first_reschedule_date() to use the server-provided context value before falling back to GetTaskRescheduleStartDate.
  • Add/adjust backward-compat and Task SDK tests covering presence/absence of the new field.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
task-sdk/tests/task_sdk/execution_time/test_task_runner.py Adds a test ensuring the server-provided first reschedule date avoids a supervisor request.
task-sdk/tests/conftest.py Extends test context factory helpers to support first_task_reschedule_start_date.
task-sdk/src/airflow/sdk/execution_time/task_runner.py Uses first_task_reschedule_start_date from server context when available.
task-sdk/src/airflow/sdk/execution_time/schema/versions/init.py Adjusts schema version bundle/lazy import plumbing (currently broken as written).
task-sdk/src/airflow/sdk/execution_time/schema/schema.json Adds first_task_reschedule_start_date to the Task SDK schema.
task-sdk/src/airflow/sdk/api/datamodels/_generated.py Regenerates Task SDK datamodels to include first_task_reschedule_start_date.
airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_06_30/test_task_instances.py Adds a backward-compat test for stripping the new field for older API versions (currently has duplicated fixture/import block).
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py Adds a head-version test asserting the run context includes the first reschedule start date.
airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py Adds a Cadwyn version change for the new response field.
airflow-core/src/airflow/api_fastapi/execution_api/versions/init.py Registers the new version change in the 2026-06-30 bundle.
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py Populates first_task_reschedule_start_date from TaskReschedule rows during ti_run.
airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py Adds the new field to the server-side TIRunContext datamodel.
Comments suppressed due to low confidence (1)

airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_06_30/test_task_instances.py:43

  • This file now contains a duplicate old_ver_client fixture plus a second block of repeated imports/pytestmark. The first fixture has no body beyond a docstring, and the redefinition will be flagged by linters (and is easy to miss in review). Remove the duplicate block so there is only one old_ver_client fixture and a single import section.
@pytest.fixture
def old_ver_client(client):
    """Last released execution API before first_task_reschedule_start_date was added."""
import pytest

from airflow._shared.timezones import timezone
from airflow.utils.state import DagRunState, State

from tests_common.test_utils.db import clear_db_runs

pytestmark = pytest.mark.db_test

Comment thread task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py Outdated
hkc-8010 added a commit to hkc-8010/my-airflow-repository that referenced this pull request Jul 3, 2026
Fix the CI failures and review comments from PR apache#68012:
- Remove the dead module-level VersionBundle construction in
  schema/versions/__init__.py that referenced names only imported
  under TYPE_CHECKING, breaking the lazy cadwyn import on any access
  to the module and causing NameError in ~30 task-sdk tests.
- Remove a duplicated old_ver_client fixture and import block that
  had been pasted into test_task_instances.py, which triggered
  mypy no-redef and ruff F811/E402 failures in Static checks.
- Use func.min(start_date) instead of ordering by TaskReschedule.id
  when computing the first reschedule start date, since insertion
  order does not guarantee chronological order.
hkc-8010 added a commit to hkc-8010/my-airflow-repository that referenced this pull request Jul 3, 2026
Fix the CI failures and review comments from PR apache#68012:
- Remove the dead module-level VersionBundle construction in
  schema/versions/__init__.py that referenced names only imported
  under TYPE_CHECKING, breaking the lazy cadwyn import on any access
  to the module and causing NameError in ~30 task-sdk tests.
- Remove a duplicated old_ver_client fixture and import block that
  had been pasted into test_task_instances.py, which triggered
  mypy no-redef and ruff F811/E402 failures in Static checks.
- Use func.min(start_date) instead of ordering by TaskReschedule.id
  when computing the first reschedule start date, since insertion
  order does not guarantee chronological order.
@hkc-8010
hkc-8010 force-pushed the fix/reschedule-start-date-context branch from ccbc4a5 to fea9229 Compare July 3, 2026 07:56

@aaron-y-chen aaron-y-chen left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Brilliant! I have a small suggestion that might be helpful :)

Comment thread airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py Outdated
@eladkal
eladkal force-pushed the fix/reschedule-start-date-context branch from 1b60429 to 171524f Compare August 5, 2026 07:09
eladkal pushed a commit to hkc-8010/my-airflow-repository that referenced this pull request Aug 5, 2026
Fix the CI failures and review comments from PR apache#68012:
- Remove the dead module-level VersionBundle construction in
  schema/versions/__init__.py that referenced names only imported
  under TYPE_CHECKING, breaking the lazy cadwyn import on any access
  to the module and causing NameError in ~30 task-sdk tests.
- Remove a duplicated old_ver_client fixture and import block that
  had been pasted into test_task_instances.py, which triggered
  mypy no-redef and ruff F811/E402 failures in Static checks.
- Use func.min(start_date) instead of ordering by TaskReschedule.id
  when computing the first reschedule start date, since insertion
  order does not guarantee chronological order.
@potiuk potiuk removed the ready for maintainer review Set by triage when all criteria pass, or by the fork CI reconciler when fork CI is green. label Aug 13, 2026
@potiuk
potiuk marked this pull request as draft August 13, 2026 12:56

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 13 out of 14 changed files in this pull request and generated no new comments.

@hkc-8010
hkc-8010 force-pushed the fix/reschedule-start-date-context branch from 3f471bc to 7209bb3 Compare September 10, 2026 07:20

@jason810496 jason810496 left a comment

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.

Thanks for the improvement, LGTM overall.

Additionally, would it be better to rephrase the title as improvement instead of the fix? Since w're improving the sensor startup overload by avoiding the additional round trip to retrieve the reschedule date instead of fixing a logical bug.

Comment thread task-sdk/src/airflow/sdk/execution_time/task_runner.py
@hkc-8010 hkc-8010 changed the title Fix rescheduled sensors hanging before poke Avoid a supervisor round trip before poke for rescheduled sensors Sep 13, 2026
@hkc-8010

Copy link
Copy Markdown
Contributor Author

@jason810496 retitled to "Avoid a supervisor round trip before poke for rescheduled sensors", your framing is more accurate. The change removes the pre-poke round trip rather than fixing a logic bug.

I replied on the three inline nits. I'd rather keep all three as they are and gave the reason on each thread. Happy to change any of them if you disagree.

@hkc-8010
hkc-8010 force-pushed the fix/reschedule-start-date-context branch from 7209bb3 to fbbc067 Compare September 24, 2026 14:38

@potiuk potiuk left a comment

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.

Thanks, this looks good to me. I checked that nothing depended on the removed round trip. The value is snapshotted at ti_run, and TaskReschedule rows for a given ti_id are only written when that same attempt reschedules (which ends it) or are deleted together with a new ti.id on clear/retry. So the snapshot always equals what the supervisor request would have returned before poke(), and the sensor timeout / backoff computation is unchanged. The versioning is also right: 2026-10-30 is still unreleased, there are matching Cadwyn changes on both the Execution API and the supervisor schema, and the SDK falls back for older servers.

It conflicts with main again: new entries were added to both 2026-10-30 version bundles since your last rebase. Could you rebase and regenerate the schema/TS artifacts? I'm happy to approve once CI is green.

One optional nit: GET /task-reschedules/{ti_id}/start_date (routes/task_reschedules.py) still picks the first row by order_by(TaskReschedule.id), while ti_run now uses min(start_date). They are equivalent in practice, but using min(start_date) in both places would keep the two definitions from drifting. Fine as a follow-up.


Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting

@hkc-8010
hkc-8010 force-pushed the fix/reschedule-start-date-context branch from fbbc067 to 4811f41 Compare October 4, 2026 13:09
@hkc-8010

hkc-8010 commented Oct 4, 2026

Copy link
Copy Markdown
Contributor Author

@potiuk Rebased onto main and regenerated the schema/TS artifacts. They came out unchanged after the merge, and both 2026-10-30 bundles now carry main's new entries alongside this one. I also took the nit: get_start_date now uses min(start_date) like ti_run, with a test that inserts reschedules out of order so the old order_by(id) query fails it. It's on the async session version main switched to in #73403.

@potiuk potiuk left a comment

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.

Thanks — that's everything. The rebase picked up main's new 2026-10-30 entries cleanly, and with min(start_date) in both places the two definitions can't drift; the out-of-order test pins it. Approving; merging once CI is green.


Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting

@potiuk

potiuk commented Oct 4, 2026

Copy link
Copy Markdown
Member

@jason810496 could you take another look and confirm you're OK with this as it stands? hkc-8010 rebased and addressed the follow-up from my review, and replied on your three remaining threads (TIRunContext construction, the uuid7 import, and the older-server comment) explaining why they'd prefer to keep them as they are. I'll hold off merging until you've had a chance to re-review — if you're happy with the replies, please resolve the threads.


Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting

@potiuk
potiuk force-pushed the fix/reschedule-start-date-context branch from 4811f41 to e9700ae Compare October 5, 2026 21:20
Comment thread task-sdk/src/airflow/sdk/execution_time/schema/schema.json
@potiuk

potiuk commented Oct 5, 2026

Copy link
Copy Markdown
Member

I pushed a fixup (fab4d19) to get static checks green. The PR adds first_task_reschedule_start_date to the Task SDK supervisor schema, and the Go SDK keeps a vendored copy of that schema plus models generated from it, so sync-go-sdk-schemas and check-go-sdk-generated-drift failed. The commit refreshes go-sdk/schema/supervisor-schema.json and regenerates go-sdk/pkg/execution/genmodels/models.gen.go. Next time the schema changes, running those two prek hooks (they need the Go toolchain) does the same. Please pull before pushing again.


Drafted-by: Claude Code (Opus 5.5); reviewed by @potiuk before posting

@hkc-8010
hkc-8010 requested a review from kaxil October 6, 2026 05:24

@amoghrajesh amoghrajesh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Nits, otherwise LGTM

Comment thread task-sdk/src/airflow/sdk/execution_time/task_runner.py Outdated
Comment thread task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py Outdated
@hkc-8010
hkc-8010 force-pushed the fix/reschedule-start-date-context branch from 98dfb4e to 2447fd1 Compare October 10, 2026 18:24
hkc-8010 and others added 5 commits October 11, 2026 10:18
Rescheduled tasks fetched their first task reschedule start date from the
supervisor when RuntimeTaskInstance.get_first_reschedule_date() was called.
For sensors this happens before poke() starts, so a worker blocked on the
supervisor/API round trip before emitting the sensor's usual poke log line.

The ti_run endpoint already queries TaskReschedule for the reschedule count;
widen that to a single count + min(start_date) query and return the first
reschedule start date in the run context. The Task SDK uses the value when
present and keeps the supervisor request as a fallback for older servers.

closes: apache#68010
ti_run now returns the earliest reschedule start date via min(start_date).
Use the same definition in GET /task-reschedules/{ti_id}/start_date instead
of picking the first row by id, so the two cannot drift apart.

Co-Authored-By: Claude <noreply@anthropic.com>
… field

The Task SDK supervisor schema gained first_task_reschedule_start_date, and the Go SDK keeps a vendored copy of that schema plus models generated from it, so the static checks failed until both were refreshed.

Generated-by: Claude Opus 5
AddFirstTaskRescheduleStartDateToSupervisorTIRunContext is didnt_exist-only,
so it is only visible through SchemaVersionMigrator.downgrade. Pin it the same
way the arg_bindings and retry_reason changes are pinned.

Co-Authored-By: Claude <noreply@anthropic.com>
Drop the API date from the older-server comment, move the migrator test
imports to module level, and add first_task_reschedule_start_date to the
TS SDK's vendored supervisor schema that apache#74415 introduced.

Co-Authored-By: Claude <noreply@anthropic.com>
@hkc-8010
hkc-8010 force-pushed the fix/reschedule-start-date-context branch from 2447fd1 to 4d2464c Compare October 12, 2026 02:52

@kaxil kaxil left a comment

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.

Thanks, LGTM

@kaxil
kaxil merged commit 3af85e2 into apache:main Oct 12, 2026
110 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:API Airflow's REST/HTTP API area:task-sdk type:bug-fix Changelog: Bug Fixes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Rescheduled sensors can hang before poke while fetching first reschedule start date

9 participants