arvados.commands.keepdocker

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