You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2023/04/05 11:38:04 UTC
[airflow] branch main updated: Allow to set limits for XCOM container (#28125)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new dc4dd9178c Allow to set limits for XCOM container (#28125)
dc4dd9178c is described below
commit dc4dd9178cfab46640c02561be63afd1da55fd52
Author: Maxim Martynov <ma...@mail.ru>
AuthorDate: Wed Apr 5 14:37:54 2023 +0300
Allow to set limits for XCOM container (#28125)
---
.../providers/cncf/kubernetes/hooks/kubernetes.py | 11 +++++
airflow/providers/cncf/kubernetes/operators/pod.py | 4 +-
.../cncf/kubernetes/utils/xcom_sidecar.py | 12 +++++-
kubernetes_tests/test_kubernetes_pod_operator.py | 11 +++--
.../cncf/kubernetes/decorators/test_kubernetes.py | 5 +++
.../cncf/kubernetes/hooks/test_kubernetes_pod.py | 33 ++++++++++++++
.../cncf/kubernetes/operators/test_pod.py | 50 ++++++++++++++++++++++
7 files changed, 119 insertions(+), 7 deletions(-)
diff --git a/airflow/providers/cncf/kubernetes/hooks/kubernetes.py b/airflow/providers/cncf/kubernetes/hooks/kubernetes.py
index f2866e5d71..91c00f483c 100644
--- a/airflow/providers/cncf/kubernetes/hooks/kubernetes.py
+++ b/airflow/providers/cncf/kubernetes/hooks/kubernetes.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import contextlib
+import json
import tempfile
import warnings
from typing import TYPE_CHECKING, Any, Generator
@@ -101,6 +102,9 @@ class KubernetesHook(BaseHook):
"xcom_sidecar_container_image": StringField(
lazy_gettext("XCom sidecar image"), widget=BS3TextFieldWidget()
),
+ "xcom_sidecar_container_resources": StringField(
+ lazy_gettext("XCom sidecar resources (JSON format)"), widget=BS3TextFieldWidget()
+ ),
}
@staticmethod
@@ -366,6 +370,13 @@ class KubernetesHook(BaseHook):
"""Returns the xcom sidecar image that defined in the connection"""
return self._get_field("xcom_sidecar_container_image")
+ def get_xcom_sidecar_container_resources(self):
+ """Returns the xcom sidecar resources that defined in the connection"""
+ field = self._get_field("xcom_sidecar_container_resources")
+ if not field:
+ return None
+ return json.loads(field)
+
def get_pod_log_stream(
self,
pod_name: str,
diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py
index 5267d1aeb6..a1ab51878f 100644
--- a/airflow/providers/cncf/kubernetes/operators/pod.py
+++ b/airflow/providers/cncf/kubernetes/operators/pod.py
@@ -810,7 +810,9 @@ class KubernetesPodOperator(BaseOperator):
if self.do_xcom_push:
self.log.debug("Adding xcom sidecar to task %s", self.task_id)
pod = xcom_sidecar.add_xcom_sidecar(
- pod, sidecar_container_image=self.hook.get_xcom_sidecar_container_image()
+ pod,
+ sidecar_container_image=self.hook.get_xcom_sidecar_container_image(),
+ sidecar_container_resources=self.hook.get_xcom_sidecar_container_resources(),
)
labels = self._get_ti_pod_labels(context)
diff --git a/airflow/providers/cncf/kubernetes/utils/xcom_sidecar.py b/airflow/providers/cncf/kubernetes/utils/xcom_sidecar.py
index 81b3047993..462e6870b6 100644
--- a/airflow/providers/cncf/kubernetes/utils/xcom_sidecar.py
+++ b/airflow/providers/cncf/kubernetes/utils/xcom_sidecar.py
@@ -42,12 +42,18 @@ class PodDefaults:
resources=k8s.V1ResourceRequirements(
requests={
"cpu": "1m",
- }
+ "memory": "10Mi",
+ },
),
)
-def add_xcom_sidecar(pod: k8s.V1Pod, *, sidecar_container_image=None) -> k8s.V1Pod:
+def add_xcom_sidecar(
+ pod: k8s.V1Pod,
+ *,
+ sidecar_container_image: str | None = None,
+ sidecar_container_resources: k8s.V1ResourceRequirements | dict | None = None,
+) -> k8s.V1Pod:
"""Adds sidecar"""
pod_cp = copy.deepcopy(pod)
pod_cp.spec.volumes = pod.spec.volumes or []
@@ -56,6 +62,8 @@ def add_xcom_sidecar(pod: k8s.V1Pod, *, sidecar_container_image=None) -> k8s.V1P
pod_cp.spec.containers[0].volume_mounts.insert(0, PodDefaults.VOLUME_MOUNT)
sidecar = copy.deepcopy(PodDefaults.SIDECAR_CONTAINER)
sidecar.image = sidecar_container_image or PodDefaults.SIDECAR_CONTAINER.image
+ if sidecar_container_resources:
+ sidecar.resources = sidecar_container_resources
pod_cp.spec.containers.append(sidecar)
return pod_cp
diff --git a/kubernetes_tests/test_kubernetes_pod_operator.py b/kubernetes_tests/test_kubernetes_pod_operator.py
index 63443bb153..bc7e356975 100644
--- a/kubernetes_tests/test_kubernetes_pod_operator.py
+++ b/kubernetes_tests/test_kubernetes_pod_operator.py
@@ -699,18 +699,18 @@ class TestKubernetesPodOperatorSystem:
assert result == {"hello": "world"}
@pytest.mark.parametrize(
- "input",
+ "env_vars",
[
param([k8s.V1EnvVar(name="env_name", value="value")], id="current"),
param({"env_name": "value"}, id="backcompat"), # todo: remove?
],
)
- def test_pod_template_file_with_overrides_system(self, input, test_label):
+ def test_pod_template_file_with_overrides_system(self, env_vars, test_label):
fixture = sys.path[0] + "/tests/kubernetes/basic_pod.yaml"
k = KubernetesPodOperator(
task_id=str(uuid4()),
labels=self.labels,
- env_vars=[k8s.V1EnvVar(name="env_name", value="value")],
+ env_vars=env_vars,
in_cluster=False,
pod_template_file=fixture,
do_xcom_push=True,
@@ -890,6 +890,7 @@ class TestKubernetesPodOperatorSystem:
await_xcom_sidecar_container_start_mock.return_value = None
hook_mock.return_value.is_in_cluster = False
hook_mock.return_value.get_xcom_sidecar_container_image.return_value = None
+ hook_mock.return_value.get_xcom_sidecar_container_resources.return_value = None
extract_xcom_mock.return_value = "{}"
path = sys.path[0] + "/tests/kubernetes/pod.yaml"
k = KubernetesPodOperator(
@@ -956,7 +957,9 @@ class TestKubernetesPodOperatorSystem:
"command": ["sh", "-c", 'trap "exit 0" INT; while true; do sleep 1; done;'],
"image": "alpine",
"name": "airflow-xcom-sidecar",
- "resources": {"requests": {"cpu": "1m"}},
+ "resources": {
+ "requests": {"cpu": "1m", "memory": "10Mi"},
+ },
"volumeMounts": [{"mountPath": "/airflow/xcom", "name": "xcom"}],
},
],
diff --git a/tests/providers/cncf/kubernetes/decorators/test_kubernetes.py b/tests/providers/cncf/kubernetes/decorators/test_kubernetes.py
index dc1b9b2bc8..4e04e96154 100644
--- a/tests/providers/cncf/kubernetes/decorators/test_kubernetes.py
+++ b/tests/providers/cncf/kubernetes/decorators/test_kubernetes.py
@@ -128,6 +128,10 @@ def test_kubernetes_with_input_output(
(ti,) = dr.task_instances
mock_hook.return_value.get_xcom_sidecar_container_image.return_value = XCOM_IMAGE
+ mock_hook.return_value.get_xcom_sidecar_container_resources.return_value = {
+ "requests": {"cpu": "1m", "memory": "10Mi"},
+ "limits": {"cpu": "1m", "memory": "50Mi"},
+ }
dag.get_task("my_task_id").execute(context=ti.get_template_context(session=session))
@@ -139,6 +143,7 @@ def test_kubernetes_with_input_output(
)
assert mock_create_pod.call_count == 1
assert mock_hook.return_value.get_xcom_sidecar_container_image.call_count == 1
+ assert mock_hook.return_value.get_xcom_sidecar_container_resources.call_count == 1
containers = mock_create_pod.call_args[1]["pod"].spec.containers
diff --git a/tests/providers/cncf/kubernetes/hooks/test_kubernetes_pod.py b/tests/providers/cncf/kubernetes/hooks/test_kubernetes_pod.py
index 2040604497..9ccca08abc 100644
--- a/tests/providers/cncf/kubernetes/hooks/test_kubernetes_pod.py
+++ b/tests/providers/cncf/kubernetes/hooks/test_kubernetes_pod.py
@@ -76,6 +76,18 @@ class TestKubernetesHook:
("disable_tcp_keepalive_empty", {"disable_tcp_keepalive": ""}),
("sidecar_container_image", {"xcom_sidecar_container_image": "private.repo.com/alpine:3.16"}),
("sidecar_container_image_empty", {"xcom_sidecar_container_image": ""}),
+ (
+ "sidecar_container_resources",
+ {
+ "xcom_sidecar_container_resources": json.dumps(
+ {
+ "requests": {"cpu": "1m", "memory": "10Mi"},
+ "limits": {"cpu": "1m", "memory": "50Mi"},
+ }
+ ),
+ },
+ ),
+ ("sidecar_container_resources_empty", {"xcom_sidecar_container_resources": ""}),
]:
db.merge_conn(Connection(conn_type="kubernetes", conn_id=conn_id, extra=json.dumps(extra)))
@@ -341,6 +353,27 @@ class TestKubernetesHook:
hook = KubernetesHook(conn_id=conn_id)
assert hook.get_xcom_sidecar_container_image() == expected
+ @pytest.mark.parametrize(
+ "conn_id, expected",
+ (
+ pytest.param(
+ "sidecar_container_resources",
+ {
+ "requests": {"cpu": "1m", "memory": "10Mi"},
+ "limits": {
+ "cpu": "1m",
+ "memory": "50Mi",
+ },
+ },
+ id="sidecar-with-resources",
+ ),
+ pytest.param("sidecar_container_resources_empty", None, id="sidecar-without-resources"),
+ ),
+ )
+ def test_get_xcom_sidecar_container_resources(self, conn_id, expected):
+ hook = KubernetesHook(conn_id=conn_id)
+ assert hook.get_xcom_sidecar_container_resources() == expected
+
@patch("kubernetes.config.kube_config.KubeConfigLoader")
@patch("kubernetes.config.kube_config.KubeConfigMerger")
def test_client_types(self, mock_kube_config_merger, mock_kube_config_loader):
diff --git a/tests/providers/cncf/kubernetes/operators/test_pod.py b/tests/providers/cncf/kubernetes/operators/test_pod.py
index 53457ae374..c542ccbdb1 100644
--- a/tests/providers/cncf/kubernetes/operators/test_pod.py
+++ b/tests/providers/cncf/kubernetes/operators/test_pod.py
@@ -513,6 +513,39 @@ class TestKubernetesPodOperator:
pod = k.build_pod_request_obj(create_context(k))
assert pod.spec.containers[1].image == "private.repo/alpine:3.13"
+ @patch(HOOK_CLASS)
+ def test_xcom_sidecar_container_resources_default(self, hook_mock):
+ hook_mock.return_value.get_xcom_sidecar_container_resources.return_value = None
+ k = KubernetesPodOperator(
+ name="test",
+ task_id="task",
+ do_xcom_push=True,
+ )
+ pod = k.build_pod_request_obj(create_context(k))
+ assert pod.spec.containers[1].resources == k8s.V1ResourceRequirements(
+ requests={
+ "cpu": "1m",
+ "memory": "10Mi",
+ },
+ )
+
+ @patch(HOOK_CLASS)
+ def test_xcom_sidecar_container_resources_custom(self, hook_mock):
+ hook_mock.return_value.get_xcom_sidecar_container_resources.return_value = {
+ "requests": {"cpu": "1m", "memory": "10Mi"},
+ "limits": {"cpu": "10m", "memory": "50Mi"},
+ }
+ k = KubernetesPodOperator(
+ name="test",
+ task_id="task",
+ do_xcom_push=True,
+ )
+ pod = k.build_pod_request_obj(create_context(k))
+ assert pod.spec.containers[1].resources == {
+ "requests": {"cpu": "1m", "memory": "10Mi"},
+ "limits": {"cpu": "10m", "memory": "50Mi"},
+ }
+
def test_image_pull_policy_correctly_set(self):
k = KubernetesPodOperator(
task_id="task",
@@ -1264,6 +1297,23 @@ class TestKubernetesPodOperatorAsync:
pod = k.build_pod_request_obj(create_context(k))
assert pod.spec.containers[1].image == "alpine"
+ @patch(HOOK_CLASS)
+ def test_async_xcom_sidecar_container_resources_default_should_execute_successfully(self, hook_mock):
+ hook_mock.return_value.get_xcom_sidecar_container_resources.return_value = None
+ k = KubernetesPodOperator(
+ name=TEST_NAME,
+ task_id="task",
+ do_xcom_push=True,
+ deferrable=True,
+ )
+ pod = k.build_pod_request_obj(create_context(k))
+ assert pod.spec.containers[1].resources == k8s.V1ResourceRequirements(
+ requests={
+ "cpu": "1m",
+ "memory": "10Mi",
+ },
+ )
+
@pytest.mark.parametrize("do_xcom_push", [True, False])
@patch(KUB_OP_PATH.format("post_complete_action"))
@patch(KUB_OP_PATH.format("extract_xcom"))