sdk/python/arvados/commands/keepdocker.py

   1 # Copyright (C) The Arvados Authors. All rights reserved.
   2 #
   3 # SPDX-License-Identifier: Apache-2.0
   4
   5 import argparse
   6 import collections
   7 import datetime
   8 import errno
   9 import fcntl
  10 import json
  11 import logging
  12 import os
  13 import re
  14 import subprocess
  15 import sys
  16 import tarfile
  17 import tempfile
  18
  19 import ciso8601
  20 from operator import itemgetter
  21 from stat import *
  22
  23 import arvados
  24 import arvados.config
  25 import arvados.util
  26 import arvados.commands._util as arv_cmd
  27 import arvados.commands.put as arv_put
  28 from arvados._version import __version__
  29
  30 logger = logging.getLogger('arvados.keepdocker')
  31 logger.setLevel(logging.DEBUG if arvados.config.get('ARVADOS_DEBUG')
  32                 else logging.INFO)
  33
  34 EARLIEST_DATETIME = datetime.datetime(datetime.MINYEAR, 1, 1, 0, 0, 0)
  35 STAT_CACHE_ERRORS = (IOError, OSError, ValueError)
  36
  37 DockerImage = collections.namedtuple(
  38     'DockerImage', ['repo', 'tag', 'hash', 'created', 'vsize'])
  39
  40 keepdocker_parser = argparse.ArgumentParser(add_help=False)
  41 keepdocker_parser.add_argument(
  42     '--version', action='version', version="%s %s" % (sys.argv[0], __version__),
  43     help='Print version and exit.')
  44 keepdocker_parser.add_argument(
  45     '-f', '--force', action='store_true', default=False,
  46     help="Re-upload the image even if it already exists on the server")
  47 keepdocker_parser.add_argument(
  48     '--force-image-format', action='store_true', default=False,
  49     help="Proceed even if the image format is not supported by the server")
  50
  51 _group = keepdocker_parser.add_mutually_exclusive_group()
  52 _group.add_argument(
  53     '--pull', action='store_true', default=False,
  54     help="Try to pull the latest image from Docker registry")
  55 _group.add_argument(
  56     '--no-pull', action='store_false', dest='pull',
  57     help="Use locally installed image only, don't pull image from Docker registry (default)")
  58
  59 # Combine keepdocker options listed above with run_opts options of arv-put.
  60 # The options inherited from arv-put include --name, --project-uuid,
  61 # --progress/--no-progress/--batch-progress and --resume/--no-resume.
  62 arg_parser = argparse.ArgumentParser(
  63         description="Upload or list Docker images in Arvados",
  64         parents=[keepdocker_parser, arv_put.run_opts, arv_cmd.retry_opt])
  65
  66 arg_parser.add_argument(
  67     'image', nargs='?',
  68     help="Docker image to upload: repo, repo:tag, or hash")
  69 arg_parser.add_argument(
  70     'tag', nargs='?',
  71     help="Tag of the Docker image to upload (default 'latest'), if image is given as an untagged repo name")
  72
  73 class DockerError(Exception):
  74     pass
  75
  76
  77 def popen_docker(cmd, *args, **kwargs):
  78     manage_stdin = ('stdin' not in kwargs)
  79     kwargs.setdefault('stdin', subprocess.PIPE)
  80     kwargs.setdefault('stdout', subprocess.PIPE)
  81     kwargs.setdefault('stderr', subprocess.PIPE)
  82     try:
  83         docker_proc = subprocess.Popen(['docker'] + cmd, *args, **kwargs)
  84     except OSError:  # No docker in $PATH, try docker.io
  85         docker_proc = subprocess.Popen(['docker.io'] + cmd, *args, **kwargs)
  86     if manage_stdin:
  87         docker_proc.stdin.close()
  88     return docker_proc
  89
  90 def check_docker(proc, description):
  91     proc.wait()
  92     if proc.returncode != 0:
  93         raise DockerError("docker {} returned status code {}".
  94                           format(description, proc.returncode))
  95
  96 def docker_image_format(image_hash):
  97     """Return the registry format ('v1' or 'v2') of the given image."""
  98     cmd = popen_docker(['inspect', '--format={{.Id}}', image_hash],
  99                         stdout=subprocess.PIPE)
 100     try:
 101         image_id = next(cmd.stdout).decode('utf-8').strip()
 102         if image_id.startswith('sha256:'):
 103             return 'v2'
 104         elif ':' not in image_id:
 105             return 'v1'
 106         else:
 107             return 'unknown'
 108     finally:
 109         check_docker(cmd, "inspect")
 110
 111 def docker_image_compatible(api, image_hash):
 112     supported = api._rootDesc.get('dockerImageFormats', [])
 113     if not supported:
 114         logger.warning("server does not specify supported image formats (see docker_image_formats in server config).")
 115         return False
 116
 117     fmt = docker_image_format(image_hash)
 118     if fmt in supported:
 119         return True
 120     else:
 121         logger.error("image format is {!r} " \
 122             "but server supports only {!r}".format(fmt, supported))
 123         return False
 124
 125 def docker_images():
 126     # Yield a DockerImage tuple for each installed image.
 127     list_proc = popen_docker(['images', '--no-trunc'], stdout=subprocess.PIPE)
 128     list_output = iter(list_proc.stdout)
 129     next(list_output)  # Ignore the header line
 130     for line in list_output:
 131         words = line.split()
 132         words = [word.decode('utf-8') for word in words]
 133         size_index = len(words) - 2
 134         repo, tag, imageid = words[:3]
 135         ctime = ' '.join(words[3:size_index])
 136         vsize = ' '.join(words[size_index:])
 137         yield DockerImage(repo, tag, imageid, ctime, vsize)
 138     list_proc.stdout.close()
 139     check_docker(list_proc, "images")
 140
 141 def find_image_hashes(image_search, image_tag=None):
 142     # Query for a Docker images with the repository and tag and return
 143     # the image ids in a list.  Returns empty list if no match is
 144     # found.
 145
 146     list_proc = popen_docker(['inspect', "%s%s" % (image_search, ":"+image_tag if image_tag else "")], stdout=subprocess.PIPE)
 147
 148     inspect = list_proc.stdout.read()
 149     list_proc.stdout.close()
 150
 151     imageinfo = json.loads(inspect)
 152
 153     return [i["Id"] for i in imageinfo]
 154
 155 def find_one_image_hash(image_search, image_tag=None):
 156     hashes = find_image_hashes(image_search, image_tag)
 157     hash_count = len(hashes)
 158     if hash_count == 1:
 159         return hashes.pop()
 160     elif hash_count == 0:
 161         raise DockerError("no matching image found")
 162     else:
 163         raise DockerError("{} images match {}".format(hash_count, image_search))
 164
 165 def stat_cache_name(image_file):
 166     return getattr(image_file, 'name', image_file) + '.stat'
 167
 168 def pull_image(image_name, image_tag):
 169     check_docker(popen_docker(['pull', '{}:{}'.format(image_name, image_tag)]),
 170                  "pull")
 171
 172 def save_image(image_hash, image_file):
 173     # Save the specified Docker image to image_file, then try to save its
 174     # stats so we can try to resume after interruption.
 175     check_docker(popen_docker(['save', image_hash], stdout=image_file),
 176                  "save")
 177     image_file.flush()
 178     try:
 179         with open(stat_cache_name(image_file), 'w') as statfile:
 180             json.dump(tuple(os.fstat(image_file.fileno())), statfile)
 181     except STAT_CACHE_ERRORS:
 182         pass  # We won't resume from this cache.  No big deal.
 183
 184 def get_cache_dir():
 185     return arv_cmd.make_home_conf_dir(
 186         os.path.join('.cache', 'arvados', 'docker'), 0o700)
 187
 188 def prep_image_file(filename):
 189     # Return a file object ready to save a Docker image,
 190     # and a boolean indicating whether or not we need to actually save the
 191     # image (False if a cached save is available).
 192     cache_dir = get_cache_dir()
 193     if cache_dir is None:
 194         image_file = tempfile.NamedTemporaryFile(suffix='.tar')
 195         need_save = True
 196     else:
 197         file_path = os.path.join(cache_dir, filename)
 198         try:
 199             with open(stat_cache_name(file_path)) as statfile:
 200                 prev_stat = json.load(statfile)
 201             now_stat = os.stat(file_path)
 202             need_save = any(prev_stat[field] != now_stat[field]
 203                             for field in [ST_MTIME, ST_SIZE])
 204         except STAT_CACHE_ERRORS + (AttributeError, IndexError):
 205             need_save = True  # We couldn't compare against old stats
 206         image_file = open(file_path, 'w+b' if need_save else 'rb')
 207     return image_file, need_save
 208
 209 def make_link(api_client, num_retries, link_class, link_name, **link_attrs):
 210     link_attrs.update({'link_class': link_class, 'name': link_name})
 211     return api_client.links().create(body=link_attrs).execute(
 212         num_retries=num_retries)
 213
 214 def docker_link_sort_key(link):
 215     """Build a sort key to find the latest available Docker image.
 216
 217     To find one source collection for a Docker image referenced by
 218     name or image id, the API server looks for a link with the most
 219     recent `image_timestamp` property; then the most recent
 220     `created_at` timestamp.  This method generates a sort key for
 221     Docker metadata links to sort them from least to most preferred.
 222     """
 223     try:
 224         image_timestamp = ciso8601.parse_datetime_as_naive(
 225             link['properties']['image_timestamp'])
 226     except (KeyError, ValueError):
 227         image_timestamp = EARLIEST_DATETIME
 228     try:
 229         created_timestamp = ciso8601.parse_datetime_as_naive(link['created_at'])
 230     except ValueError:
 231         created_timestamp = None
 232     return (image_timestamp, created_timestamp)
 233
 234 def _get_docker_links(api_client, num_retries, **kwargs):
 235     links = list(arvados.util.keyset_list_all(
 236         api_client.links().list, num_retries=num_retries, **kwargs,
 237     ))
 238     for link in links:
 239         link['_sort_key'] = docker_link_sort_key(link)
 240     links.sort(key=itemgetter('_sort_key'), reverse=True)
 241     return links
 242
 243 def _new_image_listing(link, dockerhash, repo='<none>', tag='<none>'):
 244     timestamp_index = 1 if (link['_sort_key'][0] is EARLIEST_DATETIME) else 0
 245     return {
 246         '_sort_key': link['_sort_key'],
 247         'timestamp': link['_sort_key'][timestamp_index],
 248         'collection': link['head_uuid'],
 249         'dockerhash': dockerhash,
 250         'repo': repo,
 251         'tag': tag,
 252         }
 253
 254 def list_images_in_arv(api_client, num_retries, image_name=None, image_tag=None, project_uuid=None):
 255     """List all Docker images known to the api_client with image_name and
 256     image_tag.  If no image_name is given, defaults to listing all
 257     Docker images.
 258
 259     Returns a list of tuples representing matching Docker images,
 260     sorted in preference order (i.e. the first collection in the list
 261     is the one that the API server would use). Each tuple is a
 262     (collection_uuid, collection_info) pair, where collection_info is
 263     a dict with fields "dockerhash", "repo", "tag", and "timestamp".
 264
 265     """
 266     search_filters = []
 267     repo_links = None
 268     hash_links = None
 269
 270     project_filter = []
 271     if project_uuid is not None:
 272         project_filter = [["owner_uuid", "=", project_uuid]]
 273
 274     if image_name:
 275         # Find images with the name the user specified.
 276         search_links = _get_docker_links(
 277             api_client, num_retries,
 278             filters=[['link_class', '=', 'docker_image_repo+tag'],
 279                      ['name', '=',
 280                       '{}:{}'.format(image_name, image_tag or 'latest')]]+project_filter)
 281         if search_links:
 282             repo_links = search_links
 283         else:
 284             # Fall back to finding images with the specified image hash.
 285             search_links = _get_docker_links(
 286                 api_client, num_retries,
 287                 filters=[['link_class', '=', 'docker_image_hash'],
 288                          ['name', 'ilike', image_name + '%']]+project_filter)
 289             hash_links = search_links
 290         # Only list information about images that were found in the search.
 291         search_filters.append(['head_uuid', 'in',
 292                                [link['head_uuid'] for link in search_links]])
 293
 294     # It should be reasonable to expect that each collection only has one
 295     # image hash (though there may be many links specifying this).  Find
 296     # the API server's most preferred image hash link for each collection.
 297     if hash_links is None:
 298         hash_links = _get_docker_links(
 299             api_client, num_retries,
 300             filters=search_filters + [['link_class', '=', 'docker_image_hash']]+project_filter)
 301     hash_link_map = {link['head_uuid']: link for link in reversed(hash_links)}
 302
 303     # Each collection may have more than one name (though again, one name
 304     # may be specified more than once).  Build an image listing from name
 305     # tags, sorted by API server preference.
 306     if repo_links is None:
 307         repo_links = _get_docker_links(
 308             api_client, num_retries,
 309             filters=search_filters + [['link_class', '=',
 310                                        'docker_image_repo+tag']]+project_filter)
 311     seen_image_names = collections.defaultdict(set)
 312     images = []
 313     for link in repo_links:
 314         collection_uuid = link['head_uuid']
 315         if link['name'] in seen_image_names[collection_uuid]:
 316             continue
 317         seen_image_names[collection_uuid].add(link['name'])
 318         try:
 319             dockerhash = hash_link_map[collection_uuid]['name']
 320         except KeyError:
 321             dockerhash = '<unknown>'
 322         name_parts = link['name'].rsplit(':', 1)
 323         images.append(_new_image_listing(link, dockerhash, *name_parts))
 324
 325     # Find any image hash links that did not have a corresponding name link,
 326     # and add image listings for them, retaining the API server preference
 327     # sorting.
 328     images_start_size = len(images)
 329     for collection_uuid, link in hash_link_map.items():
 330         if not seen_image_names[collection_uuid]:
 331             images.append(_new_image_listing(link, link['name']))
 332     if len(images) > images_start_size:
 333         images.sort(key=itemgetter('_sort_key'), reverse=True)
 334
 335     # Remove any image listings that refer to unknown collections.
 336     existing_coll_uuids = {coll['uuid'] for coll in arvados.util.keyset_list_all(
 337         api_client.collections().list,
 338         num_retries=num_retries,
 339         filters=[['uuid', 'in', [im['collection'] for im in images]]]+project_filter,
 340         select=['uuid'],
 341     )}
 342     return [(image['collection'], image) for image in images
 343             if image['collection'] in existing_coll_uuids]
 344
 345 def items_owned_by(owner_uuid, arv_items):
 346     return (item for item in arv_items if item['owner_uuid'] == owner_uuid)
 347
 348 def _uuid2pdh(api, uuid):
 349     return api.collections().list(
 350         filters=[['uuid', '=', uuid]],
 351         select=['portable_data_hash'],
 352     ).execute()['items'][0]['portable_data_hash']
 353
 354 def load_image_metadata(image_file):
 355     """Load an image manifest and config from an archive
 356
 357     Given an image archive as an open binary file object, this function loads
 358     the image manifest and configuration, deserializing each from JSON and
 359     returning them in a 2-tuple of dicts.
 360     """
 361     image_file.seek(0)
 362     with tarfile.open(fileobj=image_file) as image_tar:
 363         with image_tar.extractfile('manifest.json') as manifest_file:
 364             image_manifest_list = json.load(manifest_file)
 365         # Because arv-keepdocker only saves one image, there should only be
 366         # one manifest.  This extracts that from the list and raises
 367         # ValueError if there's not exactly one.
 368         image_manifest, = image_manifest_list
 369         with image_tar.extractfile(image_manifest['Config']) as config_file:
 370             image_config = json.load(config_file)
 371     return image_manifest, image_config
 372
 373 def main(arguments=None, stdout=sys.stdout, install_sig_handlers=True, api=None):
 374     args = arg_parser.parse_args(arguments)
 375     if api is None:
 376         api = arvados.api('v1', num_retries=args.retries)
 377
 378     if args.image is None or args.image == 'images':
 379         fmt = "{:30}  {:10}  {:12}  {:29}  {:20}\n"
 380         stdout.write(fmt.format("REPOSITORY", "TAG", "IMAGE ID", "COLLECTION", "CREATED"))
 381         try:
 382             for i, j in list_images_in_arv(api, args.retries):
 383                 stdout.write(fmt.format(j["repo"], j["tag"], j["dockerhash"][0:12], i, j["timestamp"].strftime("%c")))
 384         except IOError as e:
 385             if e.errno == errno.EPIPE:
 386                 pass
 387             else:
 388                 raise
 389         sys.exit(0)
 390
 391     if re.search(r':\w[-.\w]{0,127}$', args.image):
 392         # image ends with :valid-tag
 393         if args.tag is not None:
 394             logger.error(
 395                 "image %r already includes a tag, cannot add tag argument %r",
 396                 args.image, args.tag)
 397             sys.exit(1)
 398         # rsplit() accommodates "myrepo.example:8888/repo/image:tag"
 399         args.image, args.tag = args.image.rsplit(':', 1)
 400     elif args.tag is None:
 401         args.tag = 'latest'
 402
 403     if '/' in args.image:
 404         hostport, path = args.image.split('/', 1)
 405         if hostport.endswith(':443'):
 406             # "docker pull host:443/asdf" transparently removes the
 407             # :443 (which is redundant because https is implied) and
 408             # after it succeeds "docker images" will list "host/asdf",
 409             # not "host:443/asdf".  If we strip the :443 then the name
 410             # doesn't change underneath us.
 411             args.image = '/'.join([hostport[:-4], path])
 412
 413     # Pull the image if requested, unless the image is specified as a hash
 414     # that we already have.
 415     if args.pull and not find_image_hashes(args.image):
 416         pull_image(args.image, args.tag)
 417
 418     images_in_arv = list_images_in_arv(api, args.retries, args.image, args.tag)
 419
 420     image_hash = None
 421     try:
 422         image_hash = find_one_image_hash(args.image, args.tag)
 423         if not docker_image_compatible(api, image_hash):
 424             if args.force_image_format:
 425                 logger.warning("forcing incompatible image")
 426             else:
 427                 logger.error("refusing to store " \
 428                     "incompatible format (use --force-image-format to override)")
 429                 sys.exit(1)
 430     except DockerError as error:
 431         if images_in_arv:
 432             # We don't have Docker / we don't have the image locally,
 433             # use image that's already uploaded to Arvados
 434             image_hash = images_in_arv[0][1]['dockerhash']
 435         else:
 436             logger.error(str(error))
 437             sys.exit(1)
 438
 439     image_repo_tag = '{}:{}'.format(args.image, args.tag) if not image_hash.startswith(args.image.lower()) else None
 440
 441     if args.name is None:
 442         if image_repo_tag:
 443             collection_name = 'Docker image {} {}'.format(image_repo_tag.replace("/", " "), image_hash[0:12])
 444         else:
 445             collection_name = 'Docker image {}'.format(image_hash[0:12])
 446     else:
 447         collection_name = args.name
 448
 449     # Acquire a lock so that only one arv-keepdocker process will
 450     # dump/upload a particular docker image at a time.  Do this before
 451     # checking if the image already exists in Arvados so that if there
 452     # is an upload already underway, when that upload completes and
 453     # this process gets a turn, it will discover the Docker image is
 454     # already available and exit quickly.
 455     outfile_name = '{}.tar'.format(image_hash)
 456     lockfile_name = '{}.lock'.format(outfile_name)
 457     lockfile = None
 458     cache_dir = get_cache_dir()
 459     if cache_dir:
 460         lockfile = open(os.path.join(cache_dir, lockfile_name), 'w+')
 461         fcntl.flock(lockfile, fcntl.LOCK_EX)
 462
 463     try:
 464         if not args.force:
 465             # Check if this image is already in Arvados.
 466
 467             # Project where everything should be owned
 468             parent_project_uuid = args.project_uuid or api.users().current().execute(
 469                 num_retries=args.retries)['uuid']
 470
 471             # Find image hash tags
 472             existing_links = _get_docker_links(
 473                 api, args.retries,
 474                 filters=[['link_class', '=', 'docker_image_hash'],
 475                          ['name', '=', image_hash]])
 476             if existing_links:
 477                 # get readable collections
 478                 collections = api.collections().list(
 479                     filters=[['uuid', 'in', [link['head_uuid'] for link in existing_links]]],
 480                     select=["uuid", "owner_uuid", "name", "manifest_text"]
 481                     ).execute(num_retries=args.retries)['items']
 482
 483                 if collections:
 484                     # check for repo+tag links on these collections
 485                     if image_repo_tag:
 486                         existing_repo_tag = _get_docker_links(
 487                             api, args.retries,
 488                             filters=[['link_class', '=', 'docker_image_repo+tag'],
 489                                      ['name', '=', image_repo_tag],
 490                                      ['head_uuid', 'in', [c["uuid"] for c in collections]]])
 491                     else:
 492                         existing_repo_tag = []
 493
 494                     try:
 495                         coll_uuid = next(items_owned_by(parent_project_uuid, collections))['uuid']
 496                     except StopIteration:
 497                         # create new collection owned by the project
 498                         coll_uuid = api.collections().create(
 499                             body={"manifest_text": collections[0]['manifest_text'],
 500                                   "name": collection_name,
 501                                   "owner_uuid": parent_project_uuid,
 502                                   "properties": {"docker-image-repo-tag": image_repo_tag}},
 503                             ensure_unique_name=True
 504                             ).execute(num_retries=args.retries)['uuid']
 505
 506                     link_base = {'owner_uuid': parent_project_uuid,
 507                                  'head_uuid':  coll_uuid,
 508                                  'properties': existing_links[0]['properties']}
 509
 510                     if not any(items_owned_by(parent_project_uuid, existing_links)):
 511                         # create image link owned by the project
 512                         make_link(api, args.retries,
 513                                   'docker_image_hash', image_hash, **link_base)
 514
 515                     if image_repo_tag and not any(items_owned_by(parent_project_uuid, existing_repo_tag)):
 516                         # create repo+tag link owned by the project
 517                         make_link(api, args.retries, 'docker_image_repo+tag',
 518                                   image_repo_tag, **link_base)
 519
 520                     stdout.write(coll_uuid + "\n")
 521
 522                     sys.exit(0)
 523
 524         # Open a file for the saved image, and write it if needed.
 525         image_file, need_save = prep_image_file(outfile_name)
 526         if need_save:
 527             save_image(image_hash, image_file)
 528
 529         # Call arv-put with switches we inherited from it
 530         # (a.k.a., switches that aren't our own).
 531         if arguments is None:
 532             arguments = sys.argv[1:]
 533         arguments = [i for i in arguments if i not in (args.image, args.tag, image_repo_tag)]
 534         put_args = keepdocker_parser.parse_known_args(arguments)[1]
 535
 536         # Don't fail when cached manifest is invalid, just ignore the cache.
 537         put_args += ['--batch']
 538
 539         if args.name is None:
 540             put_args += ['--name', collection_name]
 541
 542         coll_uuid = arv_put.main(
 543             put_args + ['--filename', outfile_name, image_file.name], stdout=stdout,
 544             install_sig_handlers=install_sig_handlers).strip()
 545
 546         # Managed properties could be already set
 547         coll_properties = api.collections().get(uuid=coll_uuid).execute(num_retries=args.retries).get('properties', {})
 548         coll_properties.update({"docker-image-repo-tag": image_repo_tag})
 549         api.collections().update(uuid=coll_uuid, body={"properties": coll_properties}).execute(num_retries=args.retries)
 550
 551         _, image_metadata = load_image_metadata(image_file)
 552         link_base = {'head_uuid': coll_uuid, 'properties': {}}
 553         if 'created' in image_metadata:
 554             link_base['properties']['image_timestamp'] = image_metadata['created']
 555         if args.project_uuid is not None:
 556             link_base['owner_uuid'] = args.project_uuid
 557
 558         make_link(api, args.retries, 'docker_image_hash', image_hash, **link_base)
 559         if image_repo_tag:
 560             make_link(api, args.retries,
 561                       'docker_image_repo+tag', image_repo_tag, **link_base)
 562
 563         # Clean up.
 564         image_file.close()
 565         for filename in [stat_cache_name(image_file), image_file.name]:
 566             try:
 567                 os.unlink(filename)
 568             except OSError as error:
 569                 if error.errno != errno.ENOENT:
 570                     raise
 571     finally:
 572         if lockfile is not None:
 573             # Closing the lockfile unlocks it.
 574             lockfile.close()
 575
 576 if __name__ == '__main__':
 577     main()