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 = []
  40         if project_uuid:
  41             args.append("--project-uuid="+project_uuid)
  42         args.append(image_name)
  43         if image_tag:
  44             args.append(image_tag)
  45         logger.info("Uploading Docker image %s", ":".join(args[1:]))
  46         try:
  47             arvados.commands.keepdocker.main(args, stdout=sys.stderr)
  48         except SystemExit as e:
  49             if e.code:
  50                 raise WorkflowException("keepdocker exited with code %s" % e.code)
  51
  52         images = arvados.commands.keepdocker.list_images_in_arv(api_client, 3,
  53                                                                 image_name=image_name,
  54                                                                 image_tag=image_tag)
  55
  56     if not images:
  57         raise WorkflowException("Could not find Docker image %s:%s" % (image_name, image_tag))
  58
  59     pdh = api_client.collections().get(uuid=images[0][0]).execute()["portable_data_hash"]
  60
  61     with cached_lookups_lock:
  62         cached_lookups[dockerRequirement["dockerImageId"]] = pdh
  63
  64     return pdh
  65
  66 def arv_docker_clear_cache():
  67     global cached_lookups
  68     global cached_lookups_lock
  69     with cached_lookups_lock:
  70         cached_lookups = {}