fix(release): try an infra image's build, push or tag again while a registry rate-limits it

.github/scripts/infra_build.py

@@ -6,7 +6,7 @@
 ships, with its Dockerfile, context, repository and a source tag. ``exists REF`` exits 0 when the registry has REF, 1
 when it says it doesn't, and 2 when it can't tell. ``smoke NAME IMAGE`` checks that IMAGE, a local build of NAME, runs.
 ``verify-record PATH`` fails unless PATH records each image agent-env ships once, with the digest the registry serves
-under its tag now.
+under its tag now. ``retry COMMAND...`` runs COMMAND, again while a registry rate-limits it.
 """
 
 from __future__ import annotations
@@ -18,6 +18,7 @@
 import subprocess
 import sys
 import time
+from collections.abc import Callable
 from pathlib import Path
 
 REGISTRY = "ghcr.io/scaleapi"
@@ -35,6 +36,11 @@
 _HEALTH_POLL_SECONDS = 2
 _LOG_TAIL_LINES = 30  # of a container that never turned healthy, in the error
 _PRINTED_CHARS = 200  # of what docker printed in place of a manifest, in the error
+# What a registry answers when it rate-limits a request: ECR Public's 429, and the toomanyrequests error code GHCR and
+# Docker Hub send.
+_RATE_LIMITED = re.compile(r"429 Too Many Requests|toomanyrequests", re.IGNORECASE)
+# How long to wait before each try again, so a burst of builds starting together spreads out.
+_RATE_LIMIT_WAITS_SECONDS = (30, 60, 120)
 
 # Each image's check, beyond its build: the gateway imports its server, Postgres turns healthy, and the sidecars and the
 # browser answer their own commands.
@@ -57,7 +63,10 @@ def main() -> int:
     smoke.add_argument("name", choices=sorted(IMAGES))
     smoke.add_argument("image")
     commands.add_parser("verify-record").add_argument("path", type=Path)
+    commands.add_parser("retry").add_argument("argv", nargs=argparse.REMAINDER, metavar="COMMAND")
     args = parser.parse_args()
+    if args.command == "retry":
+        return _until_not_rate_limited(lambda: _streamed(args.argv)).returncode
     if args.command == "exists":
         return _exists(args.ref)
     if args.command == "smoke":
@@ -87,11 +96,34 @@ def _check(entry: str) -> list[str]:
     return problems
 
 
+def _until_not_rate_limited(attempt: Callable[[], subprocess.CompletedProcess]) -> subprocess.CompletedProcess:
+    """``attempt()``, and again after each of ``_RATE_LIMIT_WAITS_SECONDS`` while what it printed says a registry
+    rate-limited it: public.ecr.aws allows one anonymous pull a second, which builds starting together pass."""
+    for wait in _RATE_LIMIT_WAITS_SECONDS:
+        result = attempt()
+        if result.returncode == 0 or not _RATE_LIMITED.search(f"{result.stdout}{result.stderr}"):
+            return result
+        print(f"::warning::a registry rate-limited `{' '.join(result.args)[:_PRINTED_CHARS]}`; trying again in {wait}s")
+        time.sleep(wait)
+    return attempt()
+
+
+def _streamed(argv: list[str]) -> subprocess.CompletedProcess:
+    """Run ``argv``, printing what it prints as it prints it, which is kept too, to look for a rate limit in."""
+    process = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+    printed = []
+    for line in process.stdout:
+        sys.stdout.write(line)
+        printed.append(line)
+    return subprocess.CompletedProcess(argv, process.wait(), "".join(printed), "")
+
+
 def _served(ref: str) -> tuple[str | None, str]:
     """The digest the registry serves under ``ref``, or None, and what docker said when it serves none. The manifest as
     JSON, since buildx before 0.35 doesn't know the ``.Manifest.Digest`` template and prints the whole listing."""
-    result = subprocess.run(["docker", "buildx", "imagetools", "inspect", ref, "--format", "{{json .Manifest}}"],
-                            capture_output=True, text=True)
+    result = _until_not_rate_limited(lambda: subprocess.run(
+        ["docker", "buildx", "imagetools", "inspect", ref, "--format", "{{json .Manifest}}"],
+        capture_output=True, text=True))
     if result.returncode != 0:
         return None, result.stderr.strip()
     try:

.github/workflows/publish-pypi.yml

@@ -88,6 +88,8 @@ jobs:
     if: needs.infra-plan.outputs.missing != '[]'
     strategy:
       fail-fast: true
+      # A few at a time: each pulls its base images anonymously, and public.ecr.aws allows one such pull a second.
+      max-parallel: 4
       matrix:
         image: ${{ fromJSON(needs.infra-plan.outputs.missing) }}
         arch: [amd64, arm64]
@@ -127,7 +129,8 @@ jobs:
         if: steps.wanted.outputs.build == 'true'
         run: |
           # No provenance attestation, which would show in the multi-platform image as an unknown platform.
-          docker build --platform "$PLATFORM" --provenance=false -f "$DOCKERFILE" -t "$BUILT" \
+          python3 .github/scripts/infra_build.py retry \
+            docker build --platform "$PLATFORM" --provenance=false -f "$DOCKERFILE" -t "$BUILT" \
             --label "org.opencontainers.image.source=https://github.com/${GITHUB_REPOSITORY}" \
             --label "org.opencontainers.image.revision=${GITHUB_SHA}" \
             --label "io.agentenv.source-tag=${TAG}" "$CONTEXT"
@@ -136,7 +139,7 @@ jobs:
         run: python3 .github/scripts/infra_build.py smoke "$NAME" "$BUILT"
       - name: Push it under its tag and platform
         if: steps.wanted.outputs.build == 'true'
-        run: docker push "$BUILT"
+        run: python3 .github/scripts/infra_build.py retry docker push "$BUILT"
 
   # Tags each new image's two platform builds as one multi-platform image under its source tag, which installs look up.
   # A tag that exists is never moved, so an image a wheel or an install pinned stays the one its tag names.
@@ -168,7 +171,7 @@ jobs:
               continue
             fi
             [ "$code" = 1 ] || exit 1
-            docker buildx imagetools create -t "$ref" \
+            python3 .github/scripts/infra_build.py retry docker buildx imagetools create -t "$ref" \
               --annotation "index:org.opencontainers.image.source=https://github.com/${GITHUB_REPOSITORY}" \
               --annotation "index:org.opencontainers.image.revision=${GITHUB_SHA}" \
               "${ref}-amd64" "${ref}-arm64"

tst/unit/test_infra_build_script.py

@@ -218,3 +218,75 @@ def test_exists_cant_tell_when_docker_prints_no_manifest(script, monkeypatch, ca
     assert _run(script, monkeypatch, "exists", ref) == 2
 
     assert "docker printed no manifest: Name: a listing" in capsys.readouterr().out
+
+
+_LIMITED = ("#2 ERROR: failed to copy: httpReadSeeker: failed open: unexpected status code https://public.ecr.aws/v2/"
+            "docker/library/python/manifests/sha256:c208: 429 Too Many Requests\n")
+
+
+class _Process:
+    """What ``subprocess.Popen`` returns, printing ``printed`` and exiting ``code``."""
+
+    def __init__(self, printed: str, code: int):
+        self.stdout, self._code = iter(printed.splitlines(keepends=True)), code
+
+    def wait(self) -> int:
+        return self._code
+
+
+def _attempts(script, monkeypatch, answers: list[tuple[str, int]]) -> tuple[list[list[str]], list[float]]:
+    """Each run of the command prints and exits as the next of ``answers``; returns the commands run and the waits."""
+    runs, waits = [], []
+
+    def popen(argv, **kw):
+        runs.append(argv)
+        return _Process(*answers[len(runs) - 1])
+
+    monkeypatch.setattr(script.subprocess, "Popen", popen)
+    monkeypatch.setattr(script.time, "sleep", waits.append)
+    return runs, waits
+
+
+def test_retry_runs_a_command_once_when_it_works_and_prints_what_it_printed(script, monkeypatch, capsys):
+    runs, waits = _attempts(script, monkeypatch, [("#1 DONE 0.0s\n", 0)])
+
+    assert _run(script, monkeypatch, "retry", "docker", "build", "--platform", "linux/amd64", ".") == 0
+
+    assert runs == [["docker", "build", "--platform", "linux/amd64", "."]] and waits == []
+    assert capsys.readouterr().out == "#1 DONE 0.0s\n"
+
+
+def test_retry_tries_again_while_a_registry_rate_limits_it(script, monkeypatch, capsys):
+    runs, waits = _attempts(script, monkeypatch, [(_LIMITED, 1), ("toomanyrequests: retry-after: 1s\n", 1),
+                                                  ("pushed\n", 0)])
+
+    assert _run(script, monkeypatch, "retry", "docker", "push", "img") == 0
+
+    assert len(runs) == 3 and waits == list(script._RATE_LIMIT_WAITS_SECONDS[:2])
+    assert capsys.readouterr().out.count("::warning::a registry rate-limited `docker push img`; trying again in") == 2
+
+
+def test_retry_gives_up_after_its_waits(script, monkeypatch, capsys):
+    runs, waits = _attempts(script, monkeypatch, [(_LIMITED, 1)] * (len(script._RATE_LIMIT_WAITS_SECONDS) + 1))
+
+    assert _run(script, monkeypatch, "retry", "docker", "build", ".") == 1
+
+    assert len(runs) == len(script._RATE_LIMIT_WAITS_SECONDS) + 1 and waits == list(script._RATE_LIMIT_WAITS_SECONDS)
+
+
+def test_retry_doesnt_try_again_a_command_that_fails_otherwise(script, monkeypatch, capsys):
+    runs, waits = _attempts(script, monkeypatch, [("ERROR: failed to solve: process \"/bin/sh -c pip install\"\n", 1)])
+
+    assert _run(script, monkeypatch, "retry", "docker", "build", ".") == 1
+
+    assert len(runs) == 1 and waits == []
+
+
+def test_a_rate_limited_registry_lookup_is_tried_again(script, monkeypatch, capsys):
+    ref = f"{REGISTRY}/agentenv-gateway:{_TAG}"
+    answers = [subprocess.CompletedProcess([], 1, "", "ERROR: toomanyrequests: retry-after: 1s\n"),
+               subprocess.CompletedProcess([], 0, json.dumps({"digest": _DIGEST}), "")]
+    monkeypatch.setattr(script.subprocess, "run", lambda command, **kw: answers.pop(0))
+    monkeypatch.setattr(script.time, "sleep", lambda seconds: None)
+
+    assert _run(script, monkeypatch, "exists", ref) == 0

tst/unit/test_release_infra_images.py

@@ -108,3 +108,14 @@ def test_nothing_reaches_pypi_from_a_run_by_hand():
 def test_a_run_by_hand_goes_ahead_only_on_main_or_a_release_tag(job):
     """main-ci waits for CI runs on main, which a branch has none of: off main, every job is skipped instead."""
     assert "github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/v')" in JOBS[job]["if"]
+
+
+def test_what_a_registry_can_rate_limit_is_tried_again():
+    """A 429 on a base image's pull would fail the build, and with it the release's PyPI jobs."""
+    build = "\n".join(_runs("infra-build"))
+    publish = "\n".join(_runs("infra-publish"))
+
+    assert re.search(r"infra_build\.py retry \\\s+docker build", build)
+    assert 'infra_build.py retry docker push "$BUILT"' in build
+    assert "infra_build.py retry docker buildx imagetools create" in publish
+    assert JOBS["infra-build"]["strategy"]["max-parallel"] <= 4