Repository: incubator-airflow Updated Branches: refs/heads/master d9f22e44c -> e50eea629
[AIRFLOW-1633] docker_operator needs a way to set shm_size Closes #3199 from feng-tao/airflow_1633 Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/e50eea62 Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/e50eea62 Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/e50eea62 Branch: refs/heads/master Commit: e50eea629506bcafc5f2778eb3d0dfe0faf2670c Parents: d9f22e4 Author: Tao feng <tf...@lyft.com> Authored: Tue Apr 10 09:17:15 2018 +0200 Committer: Fokko Driesprong <fokkodriespr...@godatadriven.com> Committed: Tue Apr 10 09:17:15 2018 +0200 ---------------------------------------------------------------------- airflow/operators/docker_operator.py | 21 ++++++++++++--------- tests/operators/docker_operator.py | 5 +++-- 2 files changed, 15 insertions(+), 11 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/e50eea62/airflow/operators/docker_operator.py ---------------------------------------------------------------------- diff --git a/airflow/operators/docker_operator.py b/airflow/operators/docker_operator.py index 6801da4..77e6f1c 100644 --- a/airflow/operators/docker_operator.py +++ b/airflow/operators/docker_operator.py @@ -140,6 +140,7 @@ class DockerOperator(BaseOperator): self.xcom_push_flag = xcom_push self.xcom_all = xcom_all self.docker_conn_id = docker_conn_id + self.shm_size = kwargs.get('shm_size') self.cli = None self.container = None @@ -184,15 +185,17 @@ class DockerOperator(BaseOperator): self.volumes.append('{0}:{1}'.format(host_tmp_dir, self.tmp_dir)) self.container = self.cli.create_container( - command=self.get_command(), - cpu_shares=cpu_shares, - environment=self.environment, - host_config=self.cli.create_host_config(binds=self.volumes, - network_mode=self.network_mode), - image=image, - mem_limit=self.mem_limit, - user=self.user, - working_dir=self.working_dir + command=self.get_command(), + cpu_shares=cpu_shares, + environment=self.environment, + host_config=self.cli.create_host_config( + binds=self.volumes, + network_mode=self.network_mode, + shm_size=self.shm_size), + image=image, + mem_limit=self.mem_limit, + user=self.user, + working_dir=self.working_dir ) self.cli.start(self.container['Id']) http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/e50eea62/tests/operators/docker_operator.py ---------------------------------------------------------------------- diff --git a/tests/operators/docker_operator.py b/tests/operators/docker_operator.py index 9ff60d2..5123880 100644 --- a/tests/operators/docker_operator.py +++ b/tests/operators/docker_operator.py @@ -53,7 +53,7 @@ class DockerOperatorTestCase(unittest.TestCase): operator = DockerOperator(api_version='1.19', command='env', environment={'UNIT': 'TEST'}, image='ubuntu:latest', network_mode='bridge', owner='unittest', task_id='unittest', volumes=['/host/path:/container/path'], - working_dir='/container/path') + working_dir='/container/path', shm_size=1000) operator.execute(None) client_class_mock.assert_called_with(base_url='unix://var/run/docker.sock', tls=None, @@ -71,7 +71,8 @@ class DockerOperatorTestCase(unittest.TestCase): ) client_mock.create_host_config.assert_called_with(binds=['/host/path:/container/path', '/mkdtemp:/tmp/airflow'], - network_mode='bridge') + network_mode='bridge', + shm_size=1000) client_mock.images.assert_called_with(name='ubuntu:latest') client_mock.logs.assert_called_with(container='some_id', stream=True) client_mock.pull.assert_called_with('ubuntu:latest', stream=True)
