Coverage for .github / scripts / preload_kind_images.py: 100.00%
257 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1#!/usr/bin/env python3
2"""Preload the images a rendered Helm chart runs into a kind node.
4``helm install --wait`` in the kind jobs leaves every image pull to the
5kubelet, which retries a failed pull on its own exponential backoff (10 s,
620 s, ... up to 5 min between attempts), so an eight-minute wait gets about six
7tries. That is not enough for Amazon ECR Public. Anonymous clients get one
8image-pull request per second (the documented "Rate of unauthenticated image
9pulls" quota, which is not adjustable), and ECR Public answers the next request
10with 429 Too Many Requests. A pull by tag sends its manifest requests
11back-to-back (the tag, the image index, then the platform manifest), so the
12second is often refused even from an idle client. GitHub-hosted runners also
13share their egress addresses with other jobs. Argo CD's chart runs Redis from
14ECR Public, and ``integration:kind:platform-addons`` kept failing on exactly
15that 429 (``argocd-redis`` stuck in ImagePullBackOff until the wait expired).
17For every image the rendered manifests run, this pulls the image on the runner
18with retry and then ``kind load``s it, so the kubelet (the charts pull
19``IfNotPresent``) finds it on the node and never contacts a registry. Helm test
20hooks are skipped because ``helm install`` never starts them. A Docker Official
21Image served from ECR Public (``<ecr-public>/docker/library/<name>``) is first
22resolved to its digest with paced requests to ECR Public. It is then pulled by
23that digest from Docker Hub, which publishes the same index, and tagged back to
24the chart's own reference. The bytes are the ones ECR Public serves, and nothing
25the chart renders changes. If Docker Hub does not serve that digest, the image is
26pulled from ECR Public by digest, with more attempts. If ECR Public never answers
27the digest lookup (every request throttled), Docker Hub's copy of the same tag is
28used, with a warning; a definitive answer, such as a tag ECR Public does not
29have, still fails the step. Any other image is pulled as written. An image pinned
30by digest cannot be addressed by ``kind load``, so it is left to the kubelet,
31with a warning.
33``--pinned IMAGE=LOCAL`` covers a pod the job writes itself that must run
34digest-pinned bytes. Anonymous ECR Public pulls also run into its data limit
35(429 "Data limit exceeded") from shared runner addresses, which kept failing
36``integration:kind:examples-smoke`` on the MLflow TLS sidecar's pinned Python
37image. The image is pulled by its digest (Docker Hub's copy first for a Docker
38Official Image on ECR Public, then ECR Public), tagged ``LOCAL`` and loaded. The
39pod names ``LOCAL`` with ``imagePullPolicy: Never``, so it never reaches a
40registry.
42``--image NAME`` loads an image a previous step already put in the runner's
43Docker (a service image built with ``load: true``, or an example's image pulled
44with retry) without pulling anything.
46Every load goes through one path. ``kind load docker-image`` runs ``docker
47save`` for the whole image and imports the archive with ``ctr images import
48--all-platforms``. Docker 29 (the Ubuntu 26.04 runner image) stores images in
49containerd, where a pulled or built image is recorded under its multi-platform
50index while only the daemon's own platform has content, so that import fails
51with ``content digest sha256:...: not found`` (kubernetes-sigs/kind#4224). This
52script instead saves each image with ``docker save --platform`` for the
53daemon's platform (``docker version`` reports it), so the archive carries one
54manifest and the blobs behind it, and loads that archive with ``kind load
55image-archive``. One archive at a time keeps the runner's disk peak at one
56image. The archive lives under ``TMPDIR``, which the workflows point at the
57runner's work disk.
59Usage::
61 python3 .github/scripts/preload_kind_images.py --cluster NAME RENDERED.yaml [...]
62 python3 .github/scripts/preload_kind_images.py --cluster NAME --pinned IMAGE@DIGEST=LOCAL:TAG
63 python3 .github/scripts/preload_kind_images.py --cluster NAME --image NAME:TAG [--image ...]
64"""
66from __future__ import annotations
68import argparse
69import json
70import re
71import subprocess
72import sys
73import tempfile
74import time
75import urllib.error
76import urllib.request
77from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
78from dataclasses import dataclass
79from email.message import Message
80from pathlib import Path
81from typing import Any
83import yaml
85#: Amazon ECR Public's registry hosts: the gallery name and the newer endpoint
86#: charts reference (Argo CD's runs ``ecr-public.aws.com/docker/library/redis``).
87ECR_PUBLIC_HOSTS = frozenset({"public.ecr.aws", "ecr-public.aws.com"})
88#: ECR Public's service name in its token realm, the same for both hosts.
89ECR_PUBLIC_SERVICE = "public.ecr.aws"
90#: ECR Public's mirror of the Docker Official Images, published by Docker.
91DOCKER_LIBRARY = "docker/library/"
92DOCKER_HUB_LIBRARY = "docker.io/library/"
94MANIFEST_ACCEPT = ", ".join(
95 (
96 "application/vnd.oci.image.index.v1+json",
97 "application/vnd.docker.distribution.manifest.list.v2+json",
98 "application/vnd.oci.image.manifest.v1+json",
99 "application/vnd.docker.distribution.manifest.v2+json",
100 )
101)
102DIGEST_RE = re.compile(r"^sha256:[0-9a-f]{64}$")
103#: ``os/arch`` as ``docker version`` prints it (``linux/amd64``, ``linux/arm64``).
104PLATFORM_RE = re.compile(r"^[a-z0-9]+/[a-z0-9]+$")
105#: The daemon's own platform, the one an image it built or pulled has content for.
106DAEMON_PLATFORM_ARGV = ("docker", "version", "--format", "{{.Server.Os}}/{{.Server.Arch}}")
108#: Pod-spec keys whose entries carry an ``image``.
109CONTAINER_KEYS = ("initContainers", "containers", "ephemeralContainers")
110#: ``helm.sh/hook`` values ``helm install`` never runs (``helm test`` does).
111TEST_HOOKS = frozenset({"test", "test-success", "test-failure"})
113#: ECR Public admits one anonymous request per second: space them further apart.
114REQUEST_SPACING_SECONDS = 1.5
115REGISTRY_ATTEMPTS = 8
116REGISTRY_BACKOFF_CAP_SECONDS = 20.0
117REQUEST_TIMEOUT_SECONDS = 20.0
118PULL_ATTEMPTS = 4
119PULL_DELAY_SECONDS = 15.0
120#: The ECR Public fallback pull trips the same throttle, so it gets more tries.
121ECR_PUBLIC_PULL_ATTEMPTS = 12
122ECR_PUBLIC_PULL_DELAY_SECONDS = 5.0
124Runner = Callable[[Sequence[str]], int]
125#: Runs a command and returns its stdout; an empty string when it fails.
126Capturer = Callable[[Sequence[str]], str]
127Opener = Callable[[urllib.request.Request, float], Any]
128Sleeper = Callable[[float], None]
131class PreloadError(RuntimeError):
132 """An image could not be resolved, pulled, tagged or loaded."""
135class RegistryUnreachable(PreloadError):
136 """Every attempt was throttled (429) or lost in transport: no answer about the image."""
139@dataclass(frozen=True)
140class ImageReference:
141 """One image reference, normalized the way Docker and containerd read it."""
143 registry: str
144 repository: str
145 tag: str | None
146 digest: str | None
148 @classmethod
149 def parse(cls, reference: str) -> ImageReference:
150 """Split ``[registry/]repository[:tag][@digest]``; Docker Hub is the default registry."""
151 name, _, digest = reference.strip().partition("@")
152 first, slash, rest = name.partition("/")
153 if slash and ("." in first or ":" in first or first == "localhost"):
154 registry, path = first, rest
155 else:
156 registry, path = "docker.io", name
157 if "/" not in path:
158 path = f"library/{path}"
159 tag: str | None = None
160 if ":" in path.rsplit("/", 1)[-1]:
161 path, tag = path.rsplit(":", 1)
162 if not all(path.split("/")) or tag == "" or (digest and not DIGEST_RE.match(digest)):
163 raise PreloadError(f"cannot parse the image reference {reference!r}")
164 if not digest and tag is None:
165 tag = "latest"
166 return cls(registry=registry, repository=path, tag=tag, digest=digest or None)
168 def docker_hub_twin(self) -> str | None:
169 """Docker Hub's name for a Docker Official Image mirrored on ECR Public, else None."""
170 if self.registry in ECR_PUBLIC_HOSTS and self.repository.startswith(DOCKER_LIBRARY):
171 return DOCKER_HUB_LIBRARY + self.repository.removeprefix(DOCKER_LIBRARY)
172 return None
175def _is_test_hook(document: Mapping[str, Any]) -> bool:
176 metadata = document.get("metadata")
177 annotations = metadata.get("annotations") if isinstance(metadata, Mapping) else None
178 hooks = annotations.get("helm.sh/hook", "") if isinstance(annotations, Mapping) else ""
179 return any(hook.strip() in TEST_HOOKS for hook in str(hooks).split(","))
182def _container_images(node: Any) -> Iterator[str]:
183 """Every container image under ``node``, wherever a pod template sits."""
184 if isinstance(node, Mapping):
185 for key, value in node.items():
186 if key in CONTAINER_KEYS and isinstance(value, list):
187 for container in value:
188 image = container.get("image") if isinstance(container, Mapping) else None
189 if isinstance(image, str) and image.strip():
190 yield image.strip()
191 else:
192 yield from _container_images(value)
193 elif isinstance(node, list):
194 for item in node:
195 yield from _container_images(item)
198def collect_images(texts: Iterable[str]) -> list[str]:
199 """The images the rendered manifests run, in first-seen order, test hooks excluded."""
200 images: dict[str, None] = {}
201 for text in texts:
202 for document in yaml.safe_load_all(text):
203 if isinstance(document, Mapping) and not _is_test_hook(document):
204 for image in _container_images(document):
205 images.setdefault(image, None)
206 return list(images)
209def _open(request: urllib.request.Request, timeout: float) -> Any:
210 return urllib.request.urlopen( # nosec B310 # nosemgrep: dynamic-urllib-use-detected - every URL is https:// on a host from the fixed ECR_PUBLIC_HOSTS set; a reference on any other registry never reaches this function
211 request, timeout=timeout
212 )
215def _run(argv: Sequence[str]) -> int:
216 return subprocess.run( # nosemgrep: dangerous-subprocess-use-audit - fixed docker/kind argv, no shell=True
217 list(argv), check=False
218 ).returncode
221def _capture(argv: Sequence[str]) -> str:
222 completed = subprocess.run( # nosemgrep: dangerous-subprocess-use-audit - fixed docker argv, no shell=True
223 list(argv), check=False, capture_output=True, text=True
224 )
225 return completed.stdout if completed.returncode == 0 else ""
228def _registry_request(
229 url: str,
230 *,
231 method: str = "GET",
232 headers: Mapping[str, str] | None = None,
233 opener: Opener = _open,
234 sleep: Sleeper = time.sleep,
235 attempts: int = REGISTRY_ATTEMPTS,
236 spacing: float = REQUEST_SPACING_SECONDS,
237) -> tuple[Message, bytes]:
238 """One ECR Public request, retried on 429 and transport errors with growing waits.
240 Any other HTTP status is an answer (a missing tag is a 404) and raises
241 :class:`PreloadError` at once; running out of attempts raises
242 :class:`RegistryUnreachable`.
243 """
244 last_error = ""
245 for attempt in range(1, attempts + 1):
246 retry_after = ""
247 request = urllib.request.Request(url, method=method, headers=dict(headers or {}))
248 try:
249 with opener(request, REQUEST_TIMEOUT_SECONDS) as response:
250 return response.headers, response.read()
251 except urllib.error.HTTPError as exc:
252 if exc.code != 429:
253 raise PreloadError(f"{method} {url} answered HTTP {exc.code}") from exc
254 last_error = "HTTP 429 Too Many Requests"
255 retry_after = str((exc.headers.get("Retry-After") if exc.headers else None) or "")
256 except (urllib.error.URLError, TimeoutError, OSError) as exc:
257 last_error = f"{type(exc).__name__}: {exc}"
258 if attempt < attempts:
259 wait = min(spacing * 2 ** (attempt - 1), REGISTRY_BACKOFF_CAP_SECONDS)
260 if retry_after.isdigit():
261 wait = max(wait, float(retry_after))
262 print(f"{method} {url}: {last_error}; retrying in {wait:.1f}s", flush=True)
263 sleep(wait)
264 raise RegistryUnreachable(f"{method} {url} failed after {attempts} attempts ({last_error})")
267def ecr_public_digest(
268 reference: ImageReference,
269 *,
270 opener: Opener = _open,
271 sleep: Sleeper = time.sleep,
272) -> str:
273 """Resolve an ECR Public tag to the digest it serves: one token and one HEAD request."""
274 if reference.digest is not None:
275 return reference.digest
276 scope = f"repository:{reference.repository}:pull"
277 token_url = f"https://{reference.registry}/token/?service={ECR_PUBLIC_SERVICE}&scope={scope}"
278 _headers, body = _registry_request(token_url, opener=opener, sleep=sleep)
279 try:
280 token = json.loads(body)["token"]
281 except (ValueError, KeyError, TypeError) as exc:
282 raise PreloadError(f"{token_url} returned no anonymous token") from exc
283 sleep(REQUEST_SPACING_SECONDS)
284 manifest_url = (
285 f"https://{reference.registry}/v2/{reference.repository}/manifests/{reference.tag}"
286 )
287 headers, _body = _registry_request(
288 manifest_url,
289 method="HEAD",
290 headers={"Authorization": f"Bearer {token}", "Accept": MANIFEST_ACCEPT},
291 opener=opener,
292 sleep=sleep,
293 )
294 digest = str(headers.get("Docker-Content-Digest") or "").strip()
295 if not DIGEST_RE.match(digest):
296 raise PreloadError(f"{manifest_url} returned no usable Docker-Content-Digest ({digest!r})")
297 return digest
300def pull_with_retry(
301 image: str,
302 *,
303 run: Runner = _run,
304 sleep: Sleeper = time.sleep,
305 attempts: int = PULL_ATTEMPTS,
306 delay: float = PULL_DELAY_SECONDS,
307) -> bool:
308 """``docker pull`` with a fixed wait between attempts; False once they run out."""
309 for attempt in range(1, attempts + 1):
310 print(f"::group::docker pull {image} (attempt {attempt}/{attempts})", flush=True)
311 code = run(["docker", "pull", image])
312 print("::endgroup::", flush=True)
313 if code == 0:
314 return True
315 if attempt < attempts:
316 print(
317 f"::warning::docker pull {image} failed (attempt {attempt}/{attempts}); "
318 f"retrying in {delay:.0f}s",
319 flush=True,
320 )
321 sleep(delay)
322 return False
325def _tag(source: str, target: str, run: Runner) -> None:
326 if run(["docker", "tag", source, target]) != 0:
327 raise PreloadError(f"docker tag {source} {target} failed")
330def stage_image(
331 image: str,
332 *,
333 run: Runner = _run,
334 sleep: Sleeper = time.sleep,
335 resolve: Callable[[ImageReference], str] = ecr_public_digest,
336) -> bool:
337 """Put ``image`` in the runner's Docker under the chart's own name; False when skipped."""
338 reference = ImageReference.parse(image)
339 if reference.digest is not None:
340 print(
341 f"::warning::{image} is pinned by digest, which kind load cannot address; "
342 "the kubelet pulls it",
343 flush=True,
344 )
345 return False
346 twin = reference.docker_hub_twin()
347 if twin is None:
348 if not pull_with_retry(image, run=run, sleep=sleep):
349 raise PreloadError(f"could not pull {image} after {PULL_ATTEMPTS} attempts")
350 return True
352 try:
353 digest = resolve(reference)
354 except RegistryUnreachable as exc:
355 # ECR Public never answered, so its digest cannot be compared. The
356 # tag is the Docker Official Image either way. A definitive answer
357 # (a tag ECR Public does not have) is a PreloadError and is not caught.
358 source = f"{twin}:{reference.tag}"
359 print(f"::warning::{exc}; pulling {source} without the digest check", flush=True)
360 if not pull_with_retry(source, run=run, sleep=sleep):
361 raise PreloadError(f"could not pull {source} after {PULL_ATTEMPTS} attempts") from exc
362 _tag(source, image, run)
363 return True
364 source = f"{twin}@{digest}"
365 if pull_with_retry(source, run=run, sleep=sleep):
366 _tag(source, image, run)
367 print(f"{image}: pulled {source}, Docker Hub's copy of the same index", flush=True)
368 return True
369 print(f"::warning::Docker Hub does not serve {source}; pulling it from ECR Public", flush=True)
370 fallback = f"{reference.registry}/{reference.repository}@{digest}"
371 if not pull_with_retry(
372 fallback,
373 run=run,
374 sleep=sleep,
375 attempts=ECR_PUBLIC_PULL_ATTEMPTS,
376 delay=ECR_PUBLIC_PULL_DELAY_SECONDS,
377 ):
378 raise PreloadError(
379 f"could not pull {image} ({digest}) from Docker Hub or, after "
380 f"{ECR_PUBLIC_PULL_ATTEMPTS} attempts, from ECR Public"
381 )
382 _tag(fallback, image, run)
383 return True
386def stage_pinned_image(
387 image: str,
388 local: str,
389 *,
390 run: Runner = _run,
391 sleep: Sleeper = time.sleep,
392) -> None:
393 """Put the digest-pinned ``image`` in the runner's Docker under the local tag ``local``.
395 ``kind load`` cannot address a digest reference, so a pod that must run
396 pinned bytes without a registry names a local tag and pulls ``Never``.
397 Docker checks a pull by digest against that digest, so the bytes are the
398 pinned ones whichever registry serves them. A Docker Official Image on ECR
399 Public comes from Docker Hub first, which publishes the same index, then
400 from ECR Public with more attempts. Any other image is pulled by digest
401 from its own registry.
402 """
403 reference = ImageReference.parse(image)
404 if reference.digest is None:
405 raise PreloadError(f"{image} is not pinned by digest")
406 if ImageReference.parse(local).digest is not None:
407 raise PreloadError(f"the local name {local} must be a tag, not a digest")
408 own = f"{reference.registry}/{reference.repository}@{reference.digest}"
409 twin = reference.docker_hub_twin()
410 sources = (
411 [
412 (f"{twin}@{reference.digest}", PULL_ATTEMPTS, PULL_DELAY_SECONDS),
413 (own, ECR_PUBLIC_PULL_ATTEMPTS, ECR_PUBLIC_PULL_DELAY_SECONDS),
414 ]
415 if twin is not None
416 else [(own, PULL_ATTEMPTS, PULL_DELAY_SECONDS)]
417 )
418 for source, attempts, delay in sources:
419 if pull_with_retry(source, run=run, sleep=sleep, attempts=attempts, delay=delay):
420 _tag(source, local, run)
421 print(f"{image}: pulled {source} and tagged it {local}", flush=True)
422 return
423 print(f"::warning::could not pull {source}", flush=True)
424 raise PreloadError(
425 f"could not pull {image} from " + " or ".join(source for source, _, _ in sources)
426 )
429def _pinned_argument(value: str) -> tuple[str, str]:
430 """Parse ``--pinned IMAGE=LOCAL``."""
431 image, separator, local = value.partition("=")
432 if not separator or not image.strip() or not local.strip():
433 raise argparse.ArgumentTypeError(f"expected IMAGE=LOCAL, got {value!r}")
434 return image.strip(), local.strip()
437def daemon_platform(*, capture: Capturer = _capture) -> str:
438 """The Docker daemon's ``os/arch``: the one platform its images have content for."""
439 platform = capture(DAEMON_PLATFORM_ARGV).strip()
440 if not PLATFORM_RE.match(platform):
441 raise PreloadError(f"docker version reported no usable daemon platform ({platform!r})")
442 return platform
445def load_into_kind(
446 images: Sequence[str],
447 cluster: str,
448 *,
449 run: Runner = _run,
450 capture: Capturer = _capture,
451) -> None:
452 """Load every staged image into the kind node, one single-platform archive at a time.
454 ``docker save --platform`` writes only the daemon's platform, so the
455 archive is complete and ``kind load image-archive`` (``ctr images import
456 --all-platforms``) accepts it from Docker's containerd image store, where a
457 plain ``kind load docker-image`` fails on the other platforms' missing
458 content. Each archive is removed before the next is written.
459 """
460 if not images:
461 return
462 platform = daemon_platform(capture=capture)
463 for image in images:
464 with tempfile.TemporaryDirectory(prefix="kind-image-") as directory:
465 archive = str(Path(directory) / "image.tar")
466 if run(["docker", "save", "--platform", platform, "--output", archive, image]) != 0:
467 raise PreloadError(f"docker save --platform {platform} {image} failed")
468 if run(["kind", "load", "image-archive", "--name", cluster, archive]) != 0:
469 raise PreloadError(f"kind load image-archive of {image} into {cluster} failed")
470 print(f"loaded {image} ({platform}) into kind cluster {cluster}", flush=True)
473def _image_argument(value: str) -> str:
474 """Parse ``--image NAME``: a tag reference already in the runner's Docker."""
475 image = value.strip()
476 try:
477 reference = ImageReference.parse(image)
478 except PreloadError as exc:
479 raise argparse.ArgumentTypeError(str(exc)) from exc
480 if reference.digest is not None:
481 raise argparse.ArgumentTypeError(
482 f"{image} is pinned by digest, which kind cannot address; give the tag it carries"
483 )
484 return image
487def main(
488 argv: Sequence[str] | None = None,
489 *,
490 run: Runner = _run,
491 capture: Capturer = _capture,
492 sleep: Sleeper = time.sleep,
493 resolve: Callable[[ImageReference], str] = ecr_public_digest,
494) -> int:
495 parser = argparse.ArgumentParser(
496 description="Pull the images rendered Helm manifests run and load them into kind."
497 )
498 parser.add_argument("--cluster", required=True, help="kind cluster name")
499 parser.add_argument(
500 "--pinned",
501 action="append",
502 default=[],
503 type=_pinned_argument,
504 metavar="IMAGE=LOCAL",
505 help=(
506 "also stage the digest-pinned IMAGE under the local tag LOCAL, for a pod "
507 "that runs it with imagePullPolicy: Never (repeatable)"
508 ),
509 )
510 parser.add_argument(
511 "--image",
512 action="append",
513 default=[],
514 type=_image_argument,
515 metavar="NAME",
516 help=(
517 "also load NAME, an image a previous step built or pulled into the runner's "
518 "Docker, without pulling it (repeatable)"
519 ),
520 )
521 parser.add_argument("manifests", nargs="*", type=Path, help="rendered manifest files")
522 args = parser.parse_args(argv)
523 if not args.manifests and not args.pinned and not args.image:
524 parser.error("give rendered manifests, --pinned IMAGE=LOCAL, --image NAME, or a mix")
526 images = collect_images(path.read_text(encoding="utf-8") for path in args.manifests)
527 if args.manifests and not images:
528 print("::error::the rendered manifests run no container images", flush=True)
529 return 1
530 try:
531 staged = [
532 image for image in images if stage_image(image, run=run, sleep=sleep, resolve=resolve)
533 ]
534 for image, local in args.pinned:
535 stage_pinned_image(image, local, run=run, sleep=sleep)
536 load_into_kind(
537 [*staged, *(local for _, local in args.pinned), *args.image],
538 args.cluster,
539 run=run,
540 capture=capture,
541 )
542 except PreloadError as exc:
543 print(f"::error::{exc}", flush=True)
544 return 1
545 if args.manifests:
546 print(f"preloaded {len(staged)} of {len(images)} image(s) into kind cluster {args.cluster}")
547 if args.pinned:
548 print(
549 f"loaded {len(args.pinned)} pinned image(s) under local tags into kind cluster "
550 f"{args.cluster}"
551 )
552 if args.image:
553 print(f"loaded {len(args.image)} local image(s) into kind cluster {args.cluster}")
554 return 0
557if __name__ == "__main__":
558 sys.exit(main())