def task_cleanup(e):
"""
Since files written by docker containers are owned by root, we can't
clean them up in the worker process since that typically doesn't run
as root. So, we run a lightweight container to make the temp dir cleanable.
"""
if e.info['task']['mode'] == 'docker' and '_tempdir' in e.info['kwargs']:
tmpdir = e.info['kwargs']['_tempdir']
chmod_writable(tmpdir)
def __call__(self, *args, **kwargs):
default_temp_volume = _RequestDefaultTemporaryVolume()
self.request._default_temp_volume = default_temp_volume
volumes = kwargs.setdefault('volumes', {})
# If we have a list of volumes, the user provide a list of Volume objects,
# we need to transform them.
temp_volumes = []
if isinstance(volumes, list):
# See if we have been passed any TemporaryVolume instances.
for v in volumes:
if isinstance(v, TemporaryVolume):
temp_volumes.append(v)
# First call the transform method, this we replace default temp volumes
# with the instance associated with this task create above. That is any
# reference to TemporaryVolume.default
_walk_obj(volumes, self._maybe_transform_argument)
# Now convert them to JSON
def _json(volume):
return volume._repr_json_()
volumes = _walk_obj(volumes, _json)
# We then need to merge them into a single dict and it will be ready
# for docker-py.
volumes = {k: v for volume in volumes for k, v in volume.items()}
kwargs['volumes'] = volumes
volumes.update(default_temp_volume._repr_json_())
super(DockerTask, self).__call__(*args, **kwargs)
# Set the permission to allow cleanup of temp directories
temp_volumes = [v for v in temp_volumes if os.path.exists(v.host_path)]
to_chmod = temp_volumes[:]
# If our default_temp_volume instance has been transformed then we
# know it has been used and we have to clean it up.
if default_temp_volume._transformed:
to_chmod.append(default_temp_volume)
temp_volumes.append(default_temp_volume)
if len(to_chmod) > 0:
utils.chmod_writable([v.host_path for v in to_chmod])
for v in temp_volumes:
shutil.rmtree(v.host_path)