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