taskcluster/gecko_taskgraph/docker.py

   1 # This Source Code Form is subject to the terms of the Mozilla Public
   2 # License, v. 2.0. If a copy of the MPL was not distributed with this
   3 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
   4
   5
   6 import json
   7 import os
   8 import tarfile
   9 from io import BytesIO
  10
  11 from taskgraph.generator import load_tasks_for_kind
  12 from taskgraph.parameters import Parameters
  13 from taskgraph.util.taskcluster import get_artifact_url, get_session
  14
  15 from gecko_taskgraph.optimize.strategies import IndexSearch
  16 from gecko_taskgraph.util import docker
  17
  18 from . import GECKO
  19
  20
  21 def get_image_digest(image_name):
  22     params = Parameters(
  23         level=os.environ.get("MOZ_SCM_LEVEL", "3"),
  24         strict=False,
  25     )
  26     tasks = load_tasks_for_kind(params, "docker-image")
  27     task = tasks[f"docker-image-{image_name}"]
  28     return task.attributes["cached_task"]["digest"]
  29
  30
  31 def load_image_by_name(image_name, tag=None):
  32     params = {"level": os.environ.get("MOZ_SCM_LEVEL", "3")}
  33     tasks = load_tasks_for_kind(params, "docker-image")
  34     task = tasks[f"docker-image-{image_name}"]
  35     deadline = None
  36     task_id = IndexSearch().should_replace_task(
  37         task, {}, deadline, task.optimization.get("index-search", [])
  38     )
  39
  40     if task_id in (True, False):
  41         print(
  42             "Could not find artifacts for a docker image "
  43             "named `{image_name}`. Local commits and other changes "
  44             "in your checkout may cause this error. Try "
  45             "updating to a fresh checkout of mozilla-central "
  46             "to download image.".format(image_name=image_name)
  47         )
  48         return False
  49
  50     return load_image_by_task_id(task_id, tag)
  51
  52
  53 def load_image_by_task_id(task_id, tag=None):
  54     artifact_url = get_artifact_url(task_id, "public/image.tar.zst")
  55     result = load_image(artifact_url, tag)
  56     print("Found docker image: {}:{}".format(result["image"], result["tag"]))
  57     if tag:
  58         print(f"Re-tagged as: {tag}")
  59     else:
  60         tag = "{}:{}".format(result["image"], result["tag"])
  61     print(f"Try: docker run -ti --rm {tag} bash")
  62     return True
  63
  64
  65 def build_context(name, outputFile, args=None):
  66     """Build a context.tar for image with specified name."""
  67     if not name:
  68         raise ValueError("must provide a Docker image name")
  69     if not outputFile:
  70         raise ValueError("must provide a outputFile")
  71
  72     image_dir = docker.image_path(name)
  73     if not os.path.isdir(image_dir):
  74         raise Exception("image directory does not exist: %s" % image_dir)
  75
  76     docker.create_context_tar(GECKO, image_dir, outputFile, image_name=name, args=args)
  77
  78
  79 def build_image(name, tag, args=None):
  80     """Build a Docker image of specified name.
  81
  82     Output from image building process will be printed to stdout.
  83     """
  84     if not name:
  85         raise ValueError("must provide a Docker image name")
  86
  87     image_dir = docker.image_path(name)
  88     if not os.path.isdir(image_dir):
  89         raise Exception("image directory does not exist: %s" % image_dir)
  90
  91     tag = tag or docker.docker_image(name, by_tag=True)
  92
  93     buf = BytesIO()
  94     docker.stream_context_tar(GECKO, image_dir, buf, name, args)
  95     docker.post_to_docker(buf.getvalue(), "/build", nocache=1, t=tag)
  96
  97     print(f"Successfully built {name} and tagged with {tag}")
  98
  99     if tag.endswith(":latest"):
 100         print("*" * 50)
 101         print("WARNING: no VERSION file found in image directory.")
 102         print("Image is not suitable for deploying/pushing.")
 103         print("Create an image suitable for deploying/pushing by creating")
 104         print("a VERSION file in the image directory.")
 105         print("*" * 50)
 106
 107
 108 def load_image(url, imageName=None, imageTag=None):
 109     """
 110     Load docker image from URL as imageName:tag, if no imageName or tag is given
 111     it will use whatever is inside the zstd compressed tarball.
 112
 113     Returns an object with properties 'image', 'tag' and 'layer'.
 114     """
 115     import zstandard as zstd
 116
 117     # If imageName is given and we don't have an imageTag
 118     # we parse out the imageTag from imageName, or default it to 'latest'
 119     # if no imageName and no imageTag is given, 'repositories' won't be rewritten
 120     if imageName and not imageTag:
 121         if ":" in imageName:
 122             imageName, imageTag = imageName.split(":", 1)
 123         else:
 124             imageTag = "latest"
 125
 126     info = {}
 127
 128     def download_and_modify_image():
 129         # This function downloads and edits the downloaded tar file on the fly.
 130         # It emits chunked buffers of the editted tar file, as a generator.
 131         print(f"Downloading from {url}")
 132         # get_session() gets us a requests.Session set to retry several times.
 133         req = get_session().get(url, stream=True)
 134         req.raise_for_status()
 135
 136         with zstd.ZstdDecompressor().stream_reader(req.raw) as ifh:
 137
 138             tarin = tarfile.open(
 139                 mode="r|",
 140                 fileobj=ifh,
 141                 bufsize=zstd.DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE,
 142             )
 143
 144             # Stream through each member of the downloaded tar file individually.
 145             for member in tarin:
 146                 # Non-file members only need a tar header. Emit one.
 147                 if not member.isfile():
 148                     yield member.tobuf(tarfile.GNU_FORMAT)
 149                     continue
 150
 151                 # Open stream reader for the member
 152                 reader = tarin.extractfile(member)
 153
 154                 # If member is `repositories`, we parse and possibly rewrite the
 155                 # image tags.
 156                 if member.name == "repositories":
 157                     # Read and parse repositories
 158                     repos = json.loads(reader.read())
 159                     reader.close()
 160
 161                     # If there is more than one image or tag, we can't handle it
 162                     # here.
 163                     if len(repos.keys()) > 1:
 164                         raise Exception("file contains more than one image")
 165                     info["image"] = image = list(repos.keys())[0]
 166                     if len(repos[image].keys()) > 1:
 167                         raise Exception("file contains more than one tag")
 168                     info["tag"] = tag = list(repos[image].keys())[0]
 169                     info["layer"] = layer = repos[image][tag]
 170
 171                     # Rewrite the repositories file
 172                     data = json.dumps({imageName or image: {imageTag or tag: layer}})
 173                     reader = BytesIO(data.encode("utf-8"))
 174                     member.size = len(data)
 175
 176                 # Emit the tar header for this member.
 177                 yield member.tobuf(tarfile.GNU_FORMAT)
 178                 # Then emit its content.
 179                 remaining = member.size
 180                 while remaining:
 181                     length = min(remaining, zstd.DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE)
 182                     buf = reader.read(length)
 183                     remaining -= len(buf)
 184                     yield buf
 185                 # Pad to fill a 512 bytes block, per tar format.
 186                 remainder = member.size % 512
 187                 if remainder:
 188                     yield ("\0" * (512 - remainder)).encode("utf-8")
 189
 190                 reader.close()
 191
 192     docker.post_to_docker(download_and_modify_image(), "/images/load", quiet=0)
 193
 194     # Check that we found a repositories file
 195     if not info.get("image") or not info.get("tag") or not info.get("layer"):
 196         raise Exception("No repositories file found!")
 197
 198     return info