Thread version_data through BundleInfo to worker-side bundle initialization by o-nikolas · Pull Request #67217 · apache/airflow

o-nikolas · 2026-05-19T22:29:30Z

This is PR 2 of the S3 Dag Bundle versioning series. PR 1 (#66491) added the BundleVersion dataclass, Alembic migration, and persistence path. This PR completes the worker-side plumbing so that version data reaches the bundle instance at task execution time.

Adds version_data to BundleInfo and threads it through the worker-side bundle initialization path so that structured version metadata (e.g., S3 manifests) reaches the bundle at task execution time.

Changes:

BundleInfo gains version_data: dict | None = None field
ExecuteTask.make() reads version_data from DagVersion (via eagerly-loaded relationship)
Scheduler query adds selectinload(TI.dag_version) to avoid N+1 queries
BaseDagBundle.__init__ accepts and stores version_data
DagBundlesManager.get_bundle() passes version_data to the bundle constructor
task_runner.parse() and callback_supervisor pass version_data through
Task SDK _generated.py updated with the new field

related: #66491

Was generative AI tooling used to co-author this PR?

Yes — Claude Code (Opus 4)

Generated-by: Claude Code (Opus 4) following the guidelines

ferruzzi

Do we want any unit tests? It's just plumbing a value through, so maybe not necessary. Feels pretty trivial on this PR so I'll approve but felt odd not calling it out.

o-nikolas · 2026-05-22T18:12:33Z

@ashb and @amoghrajesh You might be interested in this one.

github-actions · 2026-05-27T12:11:38Z

uv.lock on main just moved via commit 76bed66 and this PR currently conflicts.

Quickest fix:

git fetch upstream main && git rebase upstream/main
rm uv.lock && uv lock
git add uv.lock && git rebase --continue
git push --force-with-lease

Automated nudge — ignore if you're not ready to rebase. This comment is updated in place on future uv.lock bumps.

o-nikolas · 2026-06-05T23:58:44Z

@ashb and @kaxil this one is ready for another pass when you get a minute 😃

…zation Add version_data to the push path so structured bundle metadata (e.g., S3 manifests) reaches workers at task execution time. Changes: - Add version_data field to BundleInfo (workloads/base.py) - Populate version_data from DagVersion in ExecuteTask.make() - Add selectinload(TI.dag_version) to scheduler enqueue query to avoid N+1 queries when reading version_data - Add version_data parameter to BaseDagBundle.__init__ (stored as self.version_data) and DagBundlesManager.get_bundle() - Pass version_data through task_runner.py and callback_supervisor.py - Regenerate task-sdk datamodels to include version_data in BundleInfo Existing bundles ignore version_data (defaults to None). The S3 bundle will use self.version_data in initialize() to fetch specific object versions (follow-up PR).

…ta field

Address review feedback: - Use dict[str, Any] | None instead of bare dict | None for version_data in both BaseDagBundle.__init__ and BundleInfo - Add minimal tests verifying version_data plumbing through the bundle constructor

…pat with 3.2.x The test_process_workloads_routes_execute_callback test uses CallbackKey(id=...) which requires the dataclass form introduced in 3.3. In Airflow 3.2.x, CallbackKey is a str type alias and does not accept keyword arguments. Change the skipif guard from AIRFLOW_V_3_2_PLUS to AIRFLOW_V_3_3_PLUS.

…rialization - Only populate version_data when dag_run.bundle_version is not None, mirroring the pinning rule for bundle_version (kaxil feedback) - Add model_serializer to BundleInfo so version_data is absent (not null) on the wire when None (ashb feedback) - Update edge OpenAPI spec: version_data is type:object, not anyOf with null - Add :param version_data: to BaseDagBundle docstring (kaxil feedback) - Remove unrelated changes from bad rebase (types.py, celery test) that were fixed separately in apache#66973 (ashb feedback)

…generate artifacts The model_serializer(mode='wrap') on BundleInfo caused Pydantic to lose JSON schema information, making the OpenAPI generator produce a generic 'additionalProperties: true, type: object' instead of the full BundleInfo schema with name/version/version_data fields. Removing the custom serializer restores correct schema generation. The version_data field is Optional so receivers already handle null. Also regenerates: - edge OpenAPI spec (v2-edge-generated.yaml) - supervisor schema snapshot (schema.json) - uv.lock (reflects upstream dependency changes after rebase)

…e() tests

Scope the scheduler's dag_version eager-load with load_only(version_data) so the batched SELECT reads two columns instead of the full row. Drop the empty supervisor schema version file, which carried no VersionChange and only existed to satisfy the version-bump hook. Also declare version_data on the _BundleInfoLike Protocol and fix the BundleInfo.version_data docstring.

PR apache#67217 wired version_data into the task execution path but not callbacks, so callbacks for a pinned run initialized their bundle without the manifest tasks used. Populate version_data on the callback producer paths (ExecuteCallback.make and the CallbackRequest creation sites) under the same pin guard as tasks, carry it on BaseCallbackRequest, and forward it through prepare_callback_bundle to get_bundle.

…zation (apache#67217) Add version_data to the push path so structured bundle metadata (e.g., S3 manifests) reaches workers at task execution time. Changes: - Add version_data field to BundleInfo (workloads/base.py) - Populate version_data from DagVersion in ExecuteTask.make() - Add selectinload(TI.dag_version) to scheduler enqueue query to avoid N+1 queries when reading version_data - Add version_data parameter to BaseDagBundle.__init__ (stored as self.version_data) and DagBundlesManager.get_bundle() - Pass version_data through task_runner.py and callback_supervisor.py - Regenerate task-sdk datamodels to include version_data in BundleInfo Existing bundles ignore version_data (defaults to None). The S3 bundle will use self.version_data in initialize() to fetch specific object versions (follow-up PR).

PR apache#67217 wired version_data into the task execution path but not callbacks, so callbacks for a pinned run initialized their bundle without the manifest tasks used. Populate version_data on the callback producer paths (ExecuteCallback.make and the CallbackRequest creation sites) under the same pin guard as tasks, carry it on BaseCallbackRequest, and forward it through prepare_callback_bundle to get_bundle.

PR #67217 wired version_data into the task execution path but not callbacks, so callbacks for a pinned run initialized their bundle without the manifest tasks used. Populate version_data on the callback producer paths (ExecuteCallback.make and the CallbackRequest creation sites) under the same pin guard as tasks, carry it on BaseCallbackRequest, and forward it through prepare_callback_bundle to get_bundle.

o-nikolas requested review from XD-DENG, amoghrajesh, ashb, dheerajturaga, ephraimbuddy, hussein-awala, jedcunningham, kaxil and pierrejeambrun as code owners May 19, 2026 22:29

boring-cyborgBot added area:DAG-processing area:Executors-core LocalExecutor & SequentialExecutor area:Scheduler including HA (high availability) scheduler area:task-sdk labels May 19, 2026

o-nikolas requested a review from jscheffl as a code owner May 20, 2026 00:14

ferruzzi approved these changes May 20, 2026
View reviewed changes

Comment threadairflow-core/src/airflow/dag_processing/bundles/base.py Outdated

o-nikolas force-pushed the onikolas/pr2/s3-bundle-version-worker-flow branch from 97dc507 to 48a9b1bCompare May 20, 2026 08:23

vincbeck approved these changes May 20, 2026
View reviewed changes

o-nikolas force-pushed the onikolas/pr2/s3-bundle-version-worker-flow branch from e5e65a1 to a5bc3b2Compare May 21, 2026 18:35

ramitkataria approved these changes May 22, 2026
View reviewed changes

ashb reviewed May 22, 2026
View reviewed changes

kaxil reviewed May 22, 2026
View reviewed changes

Comment threadairflow-core/src/airflow/dag_processing/bundles/base.py

o-nikolas force-pushed the onikolas/pr2/s3-bundle-version-worker-flow branch from 3fb016b to 87dd9a8Compare May 26, 2026 00:21

kaxil reviewed May 27, 2026
View reviewed changes

Comment threadairflow-core/src/airflow/dag_processing/bundles/manager.py Outdated
Comment threadairflow-core/src/airflow/jobs/scheduler_job_runner.py Outdated
Comment threadairflow-core/src/airflow/executors/workloads/base.py
Comment threadairflow-core/tests/unit/dag_processing/bundles/test_base.py

o-nikolas force-pushed the onikolas/pr2/s3-bundle-version-worker-flow branch from 87dd9a8 to 82399f9Compare June 5, 2026 16:00

o-nikolas force-pushed the onikolas/pr2/s3-bundle-version-worker-flow branch from 3d59905 to e00f681Compare June 9, 2026 17:07

o-nikolas requested a review from shahar1 as a code owner June 9, 2026 19:24

o-nikolas force-pushed the onikolas/pr2/s3-bundle-version-worker-flow branch from e400d62 to 1ffd2dcCompare June 9, 2026 21:35

o-nikolas force-pushed the onikolas/pr2/s3-bundle-version-worker-flow branch from 1ffd2dc to 99a3398Compare June 12, 2026 16:19

o-nikolas force-pushed the onikolas/pr2/s3-bundle-version-worker-flow branch from 3c82b47 to 35b29bdCompare June 20, 2026 00:33

vincbeck approved these changes Jun 22, 2026
View reviewed changes

kaxil reviewed Jun 22, 2026
View reviewed changes

Comment threadtask-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_06_16.py

kaxil reviewed Jun 22, 2026
View reviewed changes

Comment threadtask-sdk/src/airflow/sdk/execution_time/callback_supervisor.py

kaxil reviewed Jun 22, 2026
View reviewed changes

Comment threadairflow-core/src/airflow/executors/workloads/base.py Outdated

kaxil reviewed Jun 22, 2026
View reviewed changes

Comment threadairflow-core/tests/unit/executors/test_workloads.py Outdated

kaxil approved these changes Jun 23, 2026
View reviewed changes

o-nikolas added 10 commits June 24, 2026 15:17

Fix: add dag_version mock and update edge OpenAPI spec for version_da…
6e6ae90
…ta field

Address kaxil review: type consistency, docs, inline comment, and mak…
6c1d650
…e() tests

Simplify version_data docstring: raise soft cap to 256KB
4bfd304

Fix: update callback_supervisor test to expect version_data parameter
d63c1e6

o-nikolas added this to the Airflow 3.4.0 milestone Jun 24, 2026

o-nikolas added the full tests needed We need to run full set of tests for this PR to merge label Jun 24, 2026

Set version_data on mock TI in batch executor adopt test
703510d

o-nikolas force-pushed the onikolas/pr2/s3-bundle-version-worker-flow branch from efc148e to 703510dCompare June 24, 2026 22:36

o-nikolas merged commit 8e9a13b into apache:mainJun 25, 2026
145 checks passed

o-nikolas deleted the onikolas/pr2/s3-bundle-version-worker-flow branch June 25, 2026 21:37

o-nikolas mentioned this pull request Jun 30, 2026
Thread version_data to callbacks #69185
Merged
1 task

vatsrahul1001 mentioned this pull request Jul 30, 2026
Fix duplicate heartbeat-timeout task callbacks #68008
Merged

Navigation Menu

Uh oh!

Thread version_data through BundleInfo to worker-side bundle initialization - #67217

Thread version_data through BundleInfo to worker-side bundle initialization#67217
o-nikolas merged 11 commits into
apache:mainfrom
aws-mwaa:onikolas/pr2/s3-bundle-version-worker-flow

o-nikolas commented May 19, 2026•
edited
Loading

Uh oh!

ferruzzi left a comment

Uh oh!

Uh oh!

o-nikolas commented May 22, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actionsBot commented May 27, 2026•
edited
Loading

Uh oh!

o-nikolas commented Jun 5, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

o-nikolas commented May 19, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Was generative AI tooling used to co-author this PR?

Uh oh!

ferruzzi left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

o-nikolas commented May 22, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actionsBot commented May 27, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

o-nikolas commented Jun 5, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

o-nikolas commented May 19, 2026•
edited
Loading

github-actionsBot commented May 27, 2026•
edited
Loading