Fix rescheduled sensors hanging before poke - #68012
Conversation
|
@hkc-8010 A few things need addressing before review — see our Pull Request quality criteria.
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. |
08823ae to
f3bd4aa
Compare
|
@potiuk thanks for the heads up. I rebased the branch on current 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? |
There was a problem hiding this comment.
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_datetoTIRunContext(core + Task SDK generated models/schema) and expose it from the/execution/task-instances/{id}/runresponse. - Update
RuntimeTaskInstance.get_first_reschedule_date()to use the server-provided context value before falling back toGetTaskRescheduleStartDate. - 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_clientfixture 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 oneold_ver_clientfixture 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
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.
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.
ccbc4a5 to
fea9229
Compare
aaron-y-chen
left a comment
There was a problem hiding this comment.
Brilliant! I have a small suggestion that might be helpful :)
1b60429 to
171524f
Compare
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.
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
171524f to
3f471bc
Compare
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 beforepoke()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 -qbreeze 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 -qbreeze 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 -qbreeze testing task-sdk-tests --python 3.10prek run --files <changed-files>prek run --files <changed-files> --stage manualWas generative AI tooling used to co-author this PR?
Generated-by: OpenAI Codex following the guidelines
{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 UTCHelpful heads-up from the maintainers — please address before this PR can be reviewed:
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.