Skip to content

Commit b156151

Browse files
XD-DENGFokko
authored andcommitted
[AIRFLOW-3203] Fix DockerOperator & some operator test (#4049)
- For argument `image`, no need to explicitly add "latest" if tag is omitted.   "latest" will be used by default if no tag provided. This is handled by `docker` package itself. - Intermediate variable `cpu_shares` is not needed. - Fix wrong usage of `cpu_shares` and `cpu_shares`. Based on https://docker-py.readthedocs.io/en/stable/api.html#docker.api.container.ContainerApiMixin.create_host_config, They should be an arguments of self.cli.create_host_config() rather than APIClient.create_container(). - Change name of the corresponding test script, to ensure it can be discovered. - Fix the test itself. - Some other test scripts are not named properly, which result in failure of test discovery.
1 parent a5b0a3d commit b156151

8 files changed

+13
-17
lines changed

airflow/operators/docker_operator.py

+9-15
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ class DockerOperator(BaseOperator):
4343
be provided with the parameter ``docker_conn_id``.
4444
4545
:param image: Docker image from which to create the container.
46+
If image tag is omitted, "latest" will be used.
4647
:type image: str
4748
:param api_version: Remote API version. Set to ``auto`` to automatically
4849
detect the server's version.
@@ -62,7 +63,7 @@ class DockerOperator(BaseOperator):
6263
:type docker_url: str
6364
:param environment: Environment variables to set in the container. (templated)
6465
:type environment: dict
65-
:param force_pull: Pull the docker image on every run. Default is false.
66+
:param force_pull: Pull the docker image on every run. Default is False.
6667
:type force_pull: bool
6768
:param mem_limit: Maximum amount of memory the container can use.
6869
Either a float value, which represents the limit in bytes,
@@ -187,35 +188,28 @@ def execute(self, context):
187188
tls=tls_config
188189
)
189190

190-
if ':' not in self.image:
191-
image = self.image + ':latest'
192-
else:
193-
image = self.image
194-
195-
if self.force_pull or len(self.cli.images(name=image)) == 0:
196-
self.log.info('Pulling docker image %s', image)
197-
for l in self.cli.pull(image, stream=True):
191+
if self.force_pull or len(self.cli.images(name=self.image)) == 0:
192+
self.log.info('Pulling docker image %s', self.image)
193+
for l in self.cli.pull(self.image, stream=True):
198194
output = json.loads(l.decode('utf-8'))
199195
self.log.info("%s", output['status'])
200196

201-
cpu_shares = int(round(self.cpus * 1024))
202-
203197
with TemporaryDirectory(prefix='airflowtmp') as host_tmp_dir:
204198
self.environment['AIRFLOW_TMP_DIR'] = self.tmp_dir
205199
self.volumes.append('{0}:{1}'.format(host_tmp_dir, self.tmp_dir))
206200

207201
self.container = self.cli.create_container(
208202
command=self.get_command(),
209-
cpu_shares=cpu_shares,
210203
environment=self.environment,
211204
host_config=self.cli.create_host_config(
212205
binds=self.volumes,
213206
network_mode=self.network_mode,
214207
shm_size=self.shm_size,
215208
dns=self.dns,
216-
dns_search=self.dns_search),
217-
image=image,
218-
mem_limit=self.mem_limit,
209+
dns_search=self.dns_search,
210+
cpu_shares=int(round(self.cpus * 1024)),
211+
mem_limit=self.mem_limit),
212+
image=self.image,
219213
user=self.user,
220214
working_dir=self.working_dir
221215
)

tests/operators/docker_operator.py tests/operators/test_docker_operator.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -64,20 +64,22 @@ def test_execute(self, client_class_mock, mkdtemp_mock):
6464
client_class_mock.assert_called_with(base_url='unix://var/run/docker.sock', tls=None,
6565
version='1.19')
6666

67-
client_mock.create_container.assert_called_with(command='env', cpu_shares=1024,
67+
client_mock.create_container.assert_called_with(command='env',
6868
environment={
6969
'AIRFLOW_TMP_DIR': '/tmp/airflow',
7070
'UNIT': 'TEST'
7171
},
7272
host_config=host_config,
7373
image='ubuntu:latest',
74-
mem_limit=None, user=None,
74+
user=None,
7575
working_dir='/container/path'
7676
)
7777
client_mock.create_host_config.assert_called_with(binds=['/host/path:/container/path',
7878
'/mkdtemp:/tmp/airflow'],
7979
network_mode='bridge',
8080
shm_size=1000,
81+
cpu_shares=1024,
82+
mem_limit=None,
8183
dns=None,
8284
dns_search=None)
8385
client_mock.images.assert_called_with(name='ubuntu:latest')
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.

0 commit comments

Comments
 (0)