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.optimize.strategies import IndexSearch
  13 from taskgraph.parameters import Parameters
  14 from taskgraph.util.taskcluster import get_artifact_url, get_session
  15
  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             tarin = tarfile.open(
 138                 mode="r|",
 139                 fileobj=ifh,
 140                 bufsize=zstd.DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE,
 141             )
 142
 143             # Stream through each member of the downloaded tar file individually.
 144             for member in tarin:
 145                 # Non-file members only need a tar header. Emit one.
 146                 if not member.isfile():
 147                     yield member.tobuf(tarfile.GNU_FORMAT)
 148                     continue
 149
 150                 # Open stream reader for the member
 151                 reader = tarin.extractfile(member)
 152
 153                 # If member is `repositories`, we parse and possibly rewrite the
 154                 # image tags.
 155                 if member.name == "repositories":
 156                     # Read and parse repositories
 157                     repos = json.loads(reader.read())
 158                     reader.close()
 159
 160                     # If there is more than one image or tag, we can't handle it
 161                     # here.
 162                     if len(repos.keys()) > 1:
 163                         raise Exception("file contains more than one image")
 164                     info["image"] = image = list(repos.keys())[0]
 165                     if len(repos[image].keys()) > 1:
 166                         raise Exception("file contains more than one tag")
 167                     info["tag"] = tag = list(repos[image].keys())[0]
 168                     info["layer"] = layer = repos[image][tag]
 169
 170                     # Rewrite the repositories file
 171                     data = json.dumps({imageName or image: {imageTag or tag: layer}})
 172                     reader = BytesIO(data.encode("utf-8"))
 173                     member.size = len(data)
 174
 175                 # Emit the tar header for this member.
 176                 yield member.tobuf(tarfile.GNU_FORMAT)
 177                 # Then emit its content.
 178                 remaining = member.size
 179                 while remaining:
 180                     length = min(remaining, zstd.DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE)
 181                     buf = reader.read(length)
 182                     remaining -= len(buf)
 183                     yield buf
 184                 # Pad to fill a 512 bytes block, per tar format.
 185                 remainder = member.size % 512
 186                 if remainder:
 187                     yield ("\0" * (512 - remainder)).encode("utf-8")
 188
 189                 reader.close()
 190
 191     docker.post_to_docker(download_and_modify_image(), "/images/load", quiet=0)
 192
 193     # Check that we found a repositories file
 194     if not info.get("image") or not info.get("tag") or not info.get("layer"):
 195         raise Exception("No repositories file found!")
 196
 197     return info