From a723aef66da49dde9842395a431c8be9bbdd4072 Mon Sep 17 00:00:00 2001 From: Nischith Shetty Date: Sat, 24 Sep 2022 06:49:35 -0700 Subject: [PATCH 1/3] #24921 - add logging options to docker operator --- airflow/providers/docker/operators/docker.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/airflow/providers/docker/operators/docker.py b/airflow/providers/docker/operators/docker.py index 6456348188beb..5801d86a0802f 100644 --- a/airflow/providers/docker/operators/docker.py +++ b/airflow/providers/docker/operators/docker.py @@ -29,7 +29,7 @@ from docker import APIClient, tls # type: ignore[attr-defined] from docker.constants import DEFAULT_TIMEOUT_SECONDS # type: ignore[attr-defined] from docker.errors import APIError # type: ignore[attr-defined] -from docker.types import DeviceRequest, Mount # type: ignore[attr-defined] +from docker.types import DeviceRequest, Mount, LogConfig # type: ignore[attr-defined] from airflow.exceptions import AirflowException from airflow.models import BaseOperator @@ -140,6 +140,12 @@ class DockerOperator(BaseOperator): output that is not posted to logs :param retrieve_output_path: path for output file that will be retrieved and passed to xcom :param device_requests: Expose host resources such as GPUs to the container. + :param log_opts_max_size: The maximum size of the log before it is rolled. + A positive integer plus a modifier representing the unit of measure (k, m, or g). + Eg: 10m or 1g Defaults to -1 (unlimited). + :param log_opts_max_file: The maximum number of log files that can be present. + If rolling the logs creates excess files, the oldest file is removed. + Only effective when max-size is also set. A positive integer. Defaults to 1. """ template_fields: Sequence[str] = ('image', 'command', 'environment', 'container_name') @@ -188,6 +194,8 @@ def __init__( retrieve_output_path: str | None = None, timeout: int = DEFAULT_TIMEOUT_SECONDS, device_requests: list[DeviceRequest] | None = None, + log_opts_max_size: str | None = None, + log_opts_max_file: str | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -244,6 +252,8 @@ def __init__( self.retrieve_output_path = retrieve_output_path self.timeout = timeout self.device_requests = device_requests + self.log_opts_max_size = log_opts_max_size + self.log_opts_max_file = log_opts_max_file def get_hook(self) -> DockerHook: """ @@ -289,6 +299,11 @@ def _run_image_with_mounts(self, target_mounts, add_tmp_variable: bool) -> list[ self.environment.pop('AIRFLOW_TMP_DIR', None) if not self.cli: raise Exception("The 'cli' should be initialized before!") + docker_log_config = {} + if self.log_opts_max_size is not None: + docker_log_config['max-size'] = self.log_opts_max_size + if self.log_opts_max_file is not None: + docker_log_config['max-file'] = self.log_opts_max_file self.container = self.cli.create_container( command=self.format_command(self.command), name=self.container_name, @@ -306,6 +321,7 @@ def _run_image_with_mounts(self, target_mounts, add_tmp_variable: bool) -> list[ extra_hosts=self.extra_hosts, privileged=self.privileged, device_requests=self.device_requests, + log_config=LogConfig(config=docker_log_config), ), image=self.image, user=self.user, From f95e14a0184f1b1efdbf6a89fc40604516beb4ba Mon Sep 17 00:00:00 2001 From: Nischith Shetty Date: Sat, 24 Sep 2022 08:38:08 -0700 Subject: [PATCH 2/3] fix import based on isort suggestion --- airflow/providers/docker/operators/docker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/docker/operators/docker.py b/airflow/providers/docker/operators/docker.py index 5801d86a0802f..cd97b99f4e393 100644 --- a/airflow/providers/docker/operators/docker.py +++ b/airflow/providers/docker/operators/docker.py @@ -29,7 +29,7 @@ from docker import APIClient, tls # type: ignore[attr-defined] from docker.constants import DEFAULT_TIMEOUT_SECONDS # type: ignore[attr-defined] from docker.errors import APIError # type: ignore[attr-defined] -from docker.types import DeviceRequest, Mount, LogConfig # type: ignore[attr-defined] +from docker.types import DeviceRequest, LogConfig, Mount # type: ignore[attr-defined] from airflow.exceptions import AirflowException from airflow.models import BaseOperator From 0c030f8a3f05417a68527a92ed68283791702564 Mon Sep 17 00:00:00 2001 From: Nischith Shetty Date: Sat, 24 Sep 2022 10:06:42 -0700 Subject: [PATCH 3/3] modify test cases to include the new log opts parameters --- tests/providers/docker/operators/test_docker.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/providers/docker/operators/test_docker.py b/tests/providers/docker/operators/test_docker.py index 04e9676412bf1..72b544e321756 100644 --- a/tests/providers/docker/operators/test_docker.py +++ b/tests/providers/docker/operators/test_docker.py @@ -30,7 +30,7 @@ try: from docker import APIClient - from docker.types import DeviceRequest, Mount + from docker.types import DeviceRequest, LogConfig, Mount from airflow.providers.docker.hooks.docker import DockerHook from airflow.providers.docker.operators.docker import DockerOperator @@ -91,6 +91,8 @@ def test_execute(self): container_name='test_container', tty=True, device_requests=[DeviceRequest(count=-1, capabilities=[['gpu']])], + log_opts_max_file='5', + log_opts_max_size='10m', ) operator.execute(None) @@ -125,6 +127,7 @@ def test_execute(self): extra_hosts=None, privileged=False, device_requests=[DeviceRequest(count=-1, capabilities=[['gpu']])], + log_config=LogConfig(config={'max-size': '10m', 'max-file': '5'}), ) self.tempdir_mock.assert_called_once_with(dir='/host/airflow', prefix='airflowtmp') self.client_mock.images.assert_called_once_with(name='ubuntu:latest') @@ -188,6 +191,7 @@ def test_execute_no_temp_dir(self): extra_hosts=None, privileged=False, device_requests=None, + log_config=LogConfig(config={}), ) self.tempdir_mock.assert_not_called() self.client_mock.images.assert_called_once_with(name='ubuntu:latest') @@ -276,6 +280,7 @@ def test_execute_fallback_temp_dir(self): extra_hosts=None, privileged=False, device_requests=None, + log_config=LogConfig(config={}), ), call( mounts=[ @@ -292,6 +297,7 @@ def test_execute_fallback_temp_dir(self): extra_hosts=None, privileged=False, device_requests=None, + log_config=LogConfig(config={}), ), ] )