From 84da39b8aa43f859f69c5d381401705987ee7f8c Mon Sep 17 00:00:00 2001 From: Sam Wheating Date: Wed, 30 Jul 2025 21:04:01 -0700 Subject: [PATCH] Allow overriding schedulerName on worker/tasks pods --- .../pod-template-file.kubernetes-helm-yaml | 5 ++-- .../templates/workers/worker-deployment.yaml | 5 ++-- chart/values.schema.json | 9 ++++++ .../airflow_aux/test_pod_template_file.py | 28 +++++++++++++------ .../helm_tests/airflow_core/test_worker.py | 28 +++++++++++++------ 5 files changed, 53 insertions(+), 22 deletions(-) diff --git a/chart/files/pod-template-file.kubernetes-helm-yaml b/chart/files/pod-template-file.kubernetes-helm-yaml index aa42c1b2542c3..0b1f2158d6d51 100644 --- a/chart/files/pod-template-file.kubernetes-helm-yaml +++ b/chart/files/pod-template-file.kubernetes-helm-yaml @@ -28,6 +28,7 @@ {{- $containerLifecycleHooks := or .Values.workers.containerLifecycleHooks .Values.containerLifecycleHooks }} {{- $safeToEvict := dict "cluster-autoscaler.kubernetes.io/safe-to-evict" (.Values.workers.safeToEvict | toString) }} {{- $podAnnotations := mergeOverwrite (deepCopy .Values.airflowPodAnnotations) $safeToEvict .Values.workers.podAnnotations }} +{{- $schedulerName := or .Values.workers.schedulerName .Values.schedulerName }} apiVersion: v1 kind: Pod metadata: @@ -205,8 +206,8 @@ spec: securityContext: {{ $securityContext | nindent 4 }} nodeSelector: {{- toYaml $nodeSelector | nindent 4 }} affinity: {{- toYaml $affinity | nindent 4 }} - {{- if .Values.schedulerName }} - schedulerName: {{ .Values.schedulerName }} + {{- if $schedulerName }} + schedulerName: {{ $schedulerName }} {{- end }} terminationGracePeriodSeconds: {{ .Values.workers.terminationGracePeriodSeconds }} tolerations: {{- toYaml $tolerations | nindent 4 }} diff --git a/chart/templates/workers/worker-deployment.yaml b/chart/templates/workers/worker-deployment.yaml index 9078560988a35..39a8fd84c792b 100644 --- a/chart/templates/workers/worker-deployment.yaml +++ b/chart/templates/workers/worker-deployment.yaml @@ -40,6 +40,7 @@ {{- $containerLifecycleHooksKerberosSidecar := or .Values.workers.kerberosSidecar.containerLifecycleHooks .Values.containerLifecycleHooks }} {{- $safeToEvict := dict "cluster-autoscaler.kubernetes.io/safe-to-evict" (.Values.workers.safeToEvict | toString) }} {{- $podAnnotations := mergeOverwrite (deepCopy .Values.airflowPodAnnotations) $safeToEvict .Values.workers.podAnnotations }} +{{- $schedulerName := or .Values.workers.schedulerName .Values.schedulerName }} apiVersion: apps/v1 kind: {{ if $persistence }}StatefulSet{{ else }}Deployment{{ end }} metadata: @@ -111,8 +112,8 @@ spec: {{- if .Values.workers.priorityClassName }} priorityClassName: {{ .Values.workers.priorityClassName }} {{- end }} - {{- if .Values.schedulerName }} - schedulerName: {{ .Values.schedulerName }} + {{- if $schedulerName }} + schedulerName: {{ $schedulerName }} {{- end }} nodeSelector: {{- toYaml $nodeSelector | nindent 8 }} affinity: diff --git a/chart/values.schema.json b/chart/values.schema.json index 4c1d8c0346146..1fde4243eb76a 100644 --- a/chart/values.schema.json +++ b/chart/values.schema.json @@ -2297,6 +2297,15 @@ "type": "string" } }, + "schedulerName": { + "description": "Specify kube scheduler name for Airflow Celery workers objects and pods created with pod-template-file.", + "type": [ + "string", + "null" + ], + "default": null, + "x-docsSection": "Common" + }, "labels": { "description": "Labels to add to the Airflow Celery workers objects and pods created with pod-template-file.", "type": "object", diff --git a/helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py b/helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py index 727540fbd51d6..dff77a2e03f2c 100644 --- a/helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py +++ b/helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py @@ -575,20 +575,30 @@ def test_affinity_tolerations_topology_spread_constraints_and_node_selector_prec "spec.topologySpreadConstraints[0]", docs[0] ) - def test_scheduler_name(self): + @pytest.mark.parametrize( + "base_scheduler_name, worker_scheduler_name, expected", + [ + ("default-scheduler", "most-allocated", "most-allocated"), + ("default-scheduler", None, "default-scheduler"), + (None, None, None), + ], + ) + def test_scheduler_name(self, base_scheduler_name, worker_scheduler_name, expected): docs = render_chart( - values={"schedulerName": "airflow-scheduler"}, + values={ + "schedulerName": base_scheduler_name, + "workers": {"schedulerName": worker_scheduler_name}, + }, show_only=["templates/pod-template-file.yaml"], chart_dir=self.temp_chart_dir, ) - assert ( - jmespath.search( - "spec.schedulerName", - docs[0], - ) - == "airflow-scheduler" - ) + scheduler_name = jmespath.search("spec.schedulerName", docs[0]) + + if expected is not None: + assert scheduler_name == expected + else: + assert scheduler_name is None def test_should_not_create_default_affinity(self): docs = render_chart(show_only=["templates/pod-template-file.yaml"], chart_dir=self.temp_chart_dir) diff --git a/helm-tests/tests/helm_tests/airflow_core/test_worker.py b/helm-tests/tests/helm_tests/airflow_core/test_worker.py index 570c0ccdc8abf..0a9fbab30cf14 100644 --- a/helm-tests/tests/helm_tests/airflow_core/test_worker.py +++ b/helm-tests/tests/helm_tests/airflow_core/test_worker.py @@ -446,19 +446,29 @@ def test_affinity_tolerations_topology_spread_constraints_and_node_selector_prec "spec.template.spec.topologySpreadConstraints[0]", docs[0] ) - def test_scheduler_name(self): + @pytest.mark.parametrize( + "base_scheduler_name, worker_scheduler_name, expected", + [ + ("default-scheduler", "most-allocated", "most-allocated"), + ("default-scheduler", None, "default-scheduler"), + (None, None, None), + ], + ) + def test_scheduler_name(self, base_scheduler_name, worker_scheduler_name, expected): docs = render_chart( - values={"schedulerName": "airflow-scheduler"}, + values={ + "schedulerName": base_scheduler_name, + "workers": {"schedulerName": worker_scheduler_name}, + }, show_only=["templates/workers/worker-deployment.yaml"], ) - assert ( - jmespath.search( - "spec.template.spec.schedulerName", - docs[0], - ) - == "airflow-scheduler" - ) + scheduler_name = jmespath.search("spec.template.spec.schedulerName", docs[0]) + + if expected is not None: + assert scheduler_name == expected + else: + assert scheduler_name is None def test_should_create_default_affinity(self): docs = render_chart(show_only=["templates/workers/worker-deployment.yaml"])