build/download/docker_downloader.py

   1 #! /usr/bin/env python3
   2 # -*- coding: utf-8 -*-
   3
   4 #   COPYRIGHT NOTICE STARTS HERE
   5
   6 #   Copyright 2022 © Samsung Electronics Co., Ltd.
   7 #
   8 #   Licensed under the Apache License, Version 2.0 (the "License");
   9 #   you may not use this file except in compliance with the License.
  10 #   You may obtain a copy of the License at
  11 #
  12 #       http://www.apache.org/licenses/LICENSE-2.0
  13 #
  14 #   Unless required by applicable law or agreed to in writing, software
  15 #   distributed under the License is distributed on an "AS IS" BASIS,
  16 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  17 #   See the License for the specific language governing permissions and
  18 #   limitations under the License.
  19
  20 #   COPYRIGHT NOTICE ENDS HERE
  21
  22 import argparse
  23 import datetime
  24 import itertools
  25 import logging
  26 import os
  27 import sys
  28 import timeit
  29
  30 import docker
  31 from retrying import retry
  32
  33 from concurrent_downloader import ConcurrentDownloader
  34
  35 log = logging.getLogger(__name__)
  36
  37
  38 class DockerDownloader(ConcurrentDownloader):
  39     def __init__(self, save, *list_args, workers=3):
  40         self._save = save
  41         try:
  42             # big timeout in case of massive images like pnda-mirror-container:5.0.0 (11.4GB)
  43             self._docker_client = docker.from_env(timeout=300)
  44         except docker.errors.DockerException as err:
  45             log.exception(
  46                 'Error creating docker client. Check if docker is installed and running'
  47                 ' or if you have right permissions.')
  48             raise err
  49         self._pulled_images = set(itertools.chain.from_iterable((image.tags for image
  50                                                                  in self._docker_client.images.list())))
  51         list_args = ([*x, None] if len(x) < 2 else x for x in list_args)
  52         super().__init__('docker images', *list_args, workers=workers)
  53
  54     @staticmethod
  55     def image_registry_name(image_name):
  56         """
  57         Get the name as shown in local registry. Since some strings are not part of name
  58         when using default registry e.g. docker.io
  59         :param image_name: name of the image from the list
  60         :return: name of the image as it is shown by docker
  61         """
  62         name = image_name
  63
  64         if name.startswith('docker.io/'):
  65             name = name.replace('docker.io/', '')
  66
  67         if name.startswith('library/'):
  68             name = name.replace('library/', '')
  69
  70         if ':' not in name.rsplit('/')[-1]:
  71             name = '{}:latest'.format(name)
  72
  73         return name
  74
  75     @property
  76     def check_table(self):
  77         """
  78         Table showing information of which images are pulled/saved
  79         """
  80         self.missing()
  81         return self._table(self._data_list)
  82
  83     @property
  84     def fail_table(self):
  85         """
  86         Table showing information about state of download of images
  87         that encountered problems while downloading
  88         """
  89         return self._table(self.missing())
  90
  91     @staticmethod
  92     def _image_filename(image_name):
  93         """
  94         Get a name of a file where image will be saved.
  95         :param image_name: Name of the image from list
  96         :return: Filename of the image
  97         """
  98         return '{}.tar'.format(image_name.replace(':', '_').replace('/', '_'))
  99
 100     def _table(self, images):
 101         """
 102         Get table in format for images
 103         :param images: images to put into table
 104         :return: check table format with specified images
 105         """
 106         header = ['Name', 'Pulled', 'Saved']
 107         data = []
 108         for item in images:
 109             if item not in self._missing:
 110                 data.append((item, True, True if self._save else 'N/A'))
 111             else:
 112                 data.append((item, self._missing[item]['pulled'], self._missing[item]['saved']))
 113         return self._check_table(header, {'Name': 'l'}, data)
 114
 115     def _is_pulled(self, image):
 116         return self.image_registry_name(image) in self._pulled_images
 117
 118     def _is_saved(self, image):
 119         dst = '{}/{}'.format(self._data_list[image], self._image_filename(image))
 120         return os.path.isfile(dst)
 121
 122     def _is_missing(self, item):
 123         """
 124         Missing docker images are checked slightly differently.
 125         """
 126         pass
 127
 128     def missing(self):
 129         """
 130         Get dictionary of images not present locally.
 131         """
 132         missing = dict()
 133         for image, dst in self._data_list.items():
 134             pulled = self._is_pulled(image)
 135             if self._save:
 136                 # if pulling and save is True. Save every pulled image to assure parity
 137                 saved = False if not pulled else self._is_saved(image)
 138             else:
 139                 saved = 'N/A'
 140             if not pulled or not saved:
 141                 missing[image] = {'dst': dst, 'pulled': pulled, 'saved': saved}
 142         self._missing = missing
 143         return self._missing
 144
 145     @retry(stop_max_attempt_number=5, wait_fixed=5000)
 146     def _pull_image(self, image_name):
 147         """
 148         Pull docker image.
 149         :param image_name: name of the image to be pulled
 150         :return: pulled image (image object)
 151         :raises docker.errors.APIError: after unsuccessful retries
 152         """
 153         if ':' not in image_name.rsplit('/')[-1]:
 154             image_name = '{}:latest'.format(image_name)
 155         try:
 156             image = self._docker_client.images.pull(image_name)
 157             log.info('Image {} pulled'.format(image_name))
 158             return image
 159         except docker.errors.APIError as err:
 160             log.warning('Failed: {}: {}. Retrying...'.format(image_name, err))
 161             raise err
 162
 163     def _save_image(self, image_name, image, output_dir):
 164         """
 165         Save image to tar.
 166         :param output_dir: path to destination directory
 167         :param image: image object from pull_image function
 168         :param image_name: name of the image from list
 169         """
 170         dst = '{}/{}'.format(output_dir, self._image_filename(image_name))
 171         os.makedirs(output_dir, exist_ok=True)
 172         try:
 173             with open(dst, 'wb') as f:
 174                 for chunk in image.save(named=self.image_registry_name(image_name)):
 175                     f.write(chunk)
 176             log.info('Image {} saved as {}'.format(image_name, dst))
 177         except Exception as err:
 178             if os.path.isfile(dst):
 179                 os.remove(dst)
 180             raise err
 181
 182     def _download_item(self, image):
 183         """ Pull and save docker image from specified docker registry
 184         :param image: image to be downloaded
 185         """
 186         image_name, image_dict = image
 187         log.info('Downloading image: {}'.format(image_name))
 188         try:
 189             if image_dict['pulled']:
 190                 image_to_save = self._docker_client.images.get(image_name)
 191             else:
 192                 image_to_save = self._pull_image(image_name)
 193             if self._save:
 194                 self._save_image(image_name, image_to_save, image_dict['dst'])
 195         except Exception as err:
 196             log.exception('Error downloading {}: {}'.format(image_name, err))
 197             raise err
 198
 199
 200 def run_cli():
 201     parser = argparse.ArgumentParser(description='Download docker images from list')
 202     parser.add_argument('image_list', metavar='image-list',
 203                         help='File with list of images to download.')
 204     parser.add_argument('--save', '-s', action='store_true', default=False,
 205                         help='Save images (without it only pull is executed)')
 206     parser.add_argument('--output-dir', '-o', default=os.getcwd(),
 207                         help='Download destination')
 208     parser.add_argument('--check', '-c', action='store_true', default=False,
 209                         help='Check what is missing. No download.'
 210                              'Use with combination with -s to check saved images as well.')
 211     parser.add_argument('--debug', action='store_true', default=False,
 212                         help='Turn on debug output')
 213     parser.add_argument('--workers', type=int, default=3,
 214                         help='Set maximum workers for parallel download (default: 3)')
 215
 216     args = parser.parse_args()
 217
 218     if args.debug:
 219         logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
 220     else:
 221         logging.basicConfig(stream=sys.stdout, level=logging.INFO, format='%(message)s')
 222
 223     downloader = DockerDownloader(args.save, [args.image_list, args.output_dir], workers=args.workers)
 224
 225     if args.check:
 226         log.info('Check mode. No download will be executed.')
 227         log.info(downloader.check_table)
 228         sys.exit(0)
 229
 230     timer_start = timeit.default_timer()
 231     try:
 232         downloader.download()
 233     except RuntimeError:
 234         sys.exit(1)
 235     finally:
 236         log.info('Downloading finished in {}'.format(
 237             datetime.timedelta(seconds=timeit.default_timer() - timer_start)))
 238
 239
 240 if __name__ == '__main__':
 241     run_cli()