Add option to override Celery worker images by sg-c0de · Pull Request #63514 · apache/airflow

sg-c0de · 2026-03-13T12:39:49Z

Summary

Since we have the option to configure multiple Celery worker sets, it is useful to be able to use different images for specific sets (for example, to include CUDA drivers). This adds the option to override the image for all Celery workers globally or on a per-set basis using Values.workers.celery.image and Values.workers.celery.sets[].image.

Changes

chart/templates/_helpers.yaml: Add airflow_worker_image, airflow_worker_image_pull_policy, and airflow_worker_image_for_migrations helpers, and patch the fullOverwrite helper to support image overrides.

chart/templates/workers/worker-deployment.yaml: Update image and imagePullPolicy references to use the new helper templates.

chart/values.yaml: Add new values to support the image overrides.

boring-cyborg · 2026-03-13T12:39:52Z

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide (https://github.com/apache/airflow/blob/main/contributing-docs/README.rst)
Here are some useful points:

Pay attention to the quality of your code (ruff, mypy and type annotations). Our prek-hooks will help you with that.
In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example DAG that shows how users should use it.
Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
Be sure to read the Airflow Coding style.
Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
Apache Airflow is a community-driven project and together we are making it better 🚀.
In case of doubts contact the developers at:
Mailing List: dev@airflow.apache.org
Slack: https://s.apache.org/airflow-slack

jscheffl

Thanks, that option makes sense to me. Can you also please add some tests?

potiuk · 2026-03-16T07:20:47Z

@sg-c0de This PR has a few issues that need to be addressed before it can be reviewed — please see our Pull Request quality criteria.

Issues found:

⚠️ Missing Tests: The PR adds new Helm chart behavior (Celery worker image overrides via new helper templates and new values) but no Helm unit test files are included in the changeset. Per review guidelines, new public behavior requires tests covering success, failure, and edge cases. The Helm tests suite has a dedicated area for exactly this kind of templating logic.
⚠️ Description Does Not Match Code: The PR description states the new values are at Values.workers.celery.image and Values.workers.celery.sets[].image, but the _helpers.yaml patch references .Values.workers.image (not .Values.workers.celery.image). The values.yaml patch adds the image block inside the workers.celery section (after queue: "default"). This discrepancy between the described API surface and the actual template code is a code-correctness concern that warrants scrutiny from maintainers.

Note: Your branch is 73 commits behind main. Some check failures may be caused by changes in the base branch rather than by your PR. Please rebase your branch and push again to get up-to-date CI results.

What to do next:

The comment informs you what you need to do.
Fix each issue, then mark the PR as "Ready for review" in the GitHub UI - but only after making sure that all the issues are fixed.
There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates.
Maintainers will then proceed with a normal review.

There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates. If you have questions, feel free to ask on the Airflow Slack.

…rity

sg-c0de · 2026-03-19T12:17:23Z

@potiuk
Thanks for the review feedback! Here's what I've addressed:

Tests added:

Added test_worker_image.py with 24 tests covering:

Default image behavior (no override, global images.airflow override)
workers.celery.image overrides (full, partial repository/tag, digest precedence)
Per-set image overrides (workers.celery.sets[].image) including full overwrite
imagePullPolicy at celery and per-set levels
wait-for-airflow-migrations init container image behavior with useDefaultImageForMigration
Container consistency (worker, log-groomer, kerberos containers)

Description vs code mismatch clarified:

Updated the airflow_worker_image helper comment in _helpers.yaml to explain the merge indirection: users configure at Values.workers.celery.image, but workersMergeValues in worker-deployment.yaml promotes it to Values.workers.image before the helper runs.

All 24 tests pass locally. Branch has been rebased onto main.

Miretpl

Overall, this is a nice change that aligns with the Workers Sets feature. Let's wait a bit and revisit it in the next round after some discussions on the Helm Chart releases (the images section and design of it are one of the action items) are finished.

Miretpl · 2026-03-21T08:18:07Z

chart/templates/_helpers.yaml

+  in worker-deployment.yaml has already merged celery/set values into Values.workers, so the
+  effective image is at Values.workers.image (post-merge).
+*/}}
+{{- define "airflow_worker_image" -}}


Most of this function copies logic from airflow_image. Maybe there is some nice way to not duplicate code here?

Sure. I created common helpers for images and tags/digests, and reused them across the image helpers.

chart/values.yaml

Miretpl · 2026-03-21T08:21:10Z

helm-tests/tests/helm_tests/airflow_core/test_worker_image.py

+def _all_airflow_images(doc):
+    """Extract all airflow images from a rendered worker deployment/statefulset."""
+    init_images = jmespath.search("spec.template.spec.initContainers[*].image", doc) or []
+    container_images = jmespath.search("spec.template.spec.containers[*].image", doc) or []
+    # Filter out git-sync sidecar images (they don't use airflow_worker_image)
+    all_images = init_images + container_images
+    return [img for img in all_images if "git-sync" not in img]
+
+
+def _all_airflow_pull_policies(doc):
+    """Extract all imagePullPolicy values from airflow containers."""
+    init_policies = jmespath.search("spec.template.spec.initContainers[*].imagePullPolicy", doc) or []
+    container_policies = jmespath.search("spec.template.spec.containers[*].imagePullPolicy", doc) or []
+    return init_policies + container_policies


This added test complexity, which can lead to not-so-straightforward debugging in case of failure. Maybe we could simplify the test cases a little?

Removed this helpers and to make more straightforward tests

Miretpl · 2026-03-21T08:22:03Z

helm-tests/tests/helm_tests/airflow_core/test_worker_image.py

+    return init_policies + container_policies
+
+
+class TestWorkerImageDefault:


Not sure where the images section tests are stored, but I think we should move them there to have all images test in one place.

Image tests are divided on a per-component basis. I removed the separate file for the worker image and moved the tests to the test_worker.py and test_worker_sets.py files.

helm-tests/tests/helm_tests/airflow_core/test_worker_image.py

Miretpl · 2026-03-21T08:24:17Z

helm-tests/tests/helm_tests/airflow_core/test_worker_image.py

+            assert image == "celery-custom/airflow@sha256:abcdef1234567890"
+
+
+class TestWorkerImagePerSetOverride:


All of the tests for sets overwrite of celery, and the workers section is in the workers sets dedicated test file. Could we move these tests there?

potiuk · 2026-04-02T17:28:20Z

Are you still working ont it @sg-c0de ? can you please address @Miretpl comments and solve the issues

sg-c0de · 2026-04-03T17:36:04Z

@potiuk Sorry for delay, I'll address @Miretpl questions in the near feature

…pecific tests

sg-c0de · 2026-04-07T12:21:21Z

@Miretpl, I’ve pushed a new commit with the changes resolving all the issues.

Miretpl · 2026-04-07T18:37:23Z

@sg-c0de thanks for addressing the comments. Could we get back to it when the situation regarding future of Helm Chart (basically how it will looks like) will be more clear (if you want to know more details you can take a look at #64037)?

sg-c0de · 2026-04-08T08:25:11Z

@Miretpl OK, let's circle back once the Helm Chart direction is a bit clearer.

sg-c0de requested review from bugraoz93, hussein-awala, jedcunningham and jscheffl as code owners March 13, 2026 12:39

boring-cyborg bot added the area:helm-chart Airflow Helm Chart label Mar 13, 2026

jscheffl reviewed Mar 13, 2026

View reviewed changes

sg-c0de requested a review from dstandish as a code owner March 19, 2026 12:09

sg-c0de added 2 commits March 19, 2026 14:11

Add option to override Celery worker images

564c501

Add worker image tests and update worker image helper comment for cla…

9105d17

…rity

sg-c0de force-pushed the celery-worker-image-override branch from 366cbfc to 9105d17 Compare March 19, 2026 12:11

jscheffl added this to the Airflow Helm Chart 1.21.0 milestone Mar 19, 2026

Miretpl suggested changes Mar 21, 2026

View reviewed changes

Refactor worker image rendering and move image coverage into worker-s…

1ab16c8

…pecific tests

		return init_policies + container_policies


		class TestWorkerImageDefault:

		assert image == "celery-custom/airflow@sha256:abcdef1234567890"


		class TestWorkerImagePerSetOverride:

Conversation

sg-c0de commented Mar 13, 2026

Summary

Changes

Uh oh!

boring-cyborg bot commented Mar 13, 2026

Uh oh!

jscheffl left a comment

Choose a reason for hiding this comment

Uh oh!

potiuk commented Mar 16, 2026

Uh oh!

sg-c0de commented Mar 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Tests added:

Description vs code mismatch clarified:

Uh oh!

Miretpl left a comment

Choose a reason for hiding this comment

Uh oh!

Miretpl Mar 21, 2026

Choose a reason for hiding this comment

Uh oh!

sg-c0de Apr 7, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Miretpl Mar 21, 2026

Choose a reason for hiding this comment

Uh oh!

sg-c0de Apr 7, 2026

Choose a reason for hiding this comment

Uh oh!

Miretpl Mar 21, 2026

Choose a reason for hiding this comment

Uh oh!

sg-c0de Apr 7, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Miretpl Mar 21, 2026

Choose a reason for hiding this comment

Uh oh!

potiuk commented Apr 2, 2026

Uh oh!

sg-c0de commented Apr 3, 2026

Uh oh!

sg-c0de commented Apr 7, 2026

Uh oh!

Miretpl commented Apr 7, 2026

Uh oh!

sg-c0de commented Apr 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

sg-c0de commented Mar 19, 2026 •

edited

Loading