sdk/cwl/arvados_cwl/arvdocker.py

   1 import logging
   2 import sys
   3 import threading
   4
   5 import cwltool.docker
   6 from cwltool.errors import WorkflowException
   7 import arvados.commands.keepdocker
   8
   9 logger = logging.getLogger('arvados.cwl-runner')
  10
  11 cached_lookups = {}
  12 cached_lookups_lock = threading.Lock()
  13
  14 def arv_docker_get_image(api_client, dockerRequirement, pull_image, project_uuid):
  15     """Check if a Docker image is available in Keep, if not, upload it using arv-keepdocker."""
  16
  17     if "dockerImageId" not in dockerRequirement and "dockerPull" in dockerRequirement:
  18         dockerRequirement["dockerImageId"] = dockerRequirement["dockerPull"]
  19
  20     global cached_lookups
  21     global cached_lookups_lock
  22     with cached_lookups_lock:
  23         if dockerRequirement["dockerImageId"] in cached_lookups:
  24             return cached_lookups[dockerRequirement["dockerImageId"]]
  25
  26     sp = dockerRequirement["dockerImageId"].split(":")
  27     image_name = sp[0]
  28     image_tag = sp[1] if len(sp) > 1 else None
  29
  30     images = arvados.commands.keepdocker.list_images_in_arv(api_client, 3,
  31                                                             image_name=image_name,
  32                                                             image_tag=image_tag)
  33
  34     if not images:
  35         # Fetch Docker image if necessary.
  36         cwltool.docker.get_image(dockerRequirement, pull_image)
  37
  38         # Upload image to Arvados
  39         args = ["--project-uuid="+project_uuid, image_name]
  40         if image_tag:
  41             args.append(image_tag)
  42         logger.info("Uploading Docker image %s", ":".join(args[1:]))
  43         try:
  44             arvados.commands.keepdocker.main(args, stdout=sys.stderr)
  45         except SystemExit as e:
  46             if e.code:
  47                 raise WorkflowException("keepdocker exited with code %s" % e.code)
  48
  49         images = arvados.commands.keepdocker.list_images_in_arv(api_client, 3,
  50                                                                 image_name=image_name,
  51                                                                 image_tag=image_tag)
  52
  53     if not images:
  54         raise WorkflowException("Could not find Docker image %s:%s" % (image_name, image_tag))
  55
  56     pdh = api_client.collections().get(uuid=images[0][0]).execute()["portable_data_hash"]
  57
  58     with cached_lookups_lock:
  59         cached_lookups[dockerRequirement["dockerImageId"]] = pdh
  60
  61     return pdh
  62
  63 def arv_docker_clear_cache():
  64     global cached_lookups
  65     global cached_lookups_lock
  66     with cached_lookups_lock:
  67         cached_lookups = {}