alphafold3: tolerate large dependency downloads

images/alphafold3/build.py

@@ -19,6 +19,7 @@ def main() -> None:
         image_dir=Path(__file__).parent,
         image="registry.sjanglab.org/sjanglab/alphafold3",
         dockerfile="docker/Dockerfile",
+        build_arguments={"UV_HTTP_TIMEOUT": "300"},
         push=args.push,
     )
     print(result.digest)

scripts/builder/build.py

@@ -4,6 +4,7 @@
 
 import json
 import os
+import re
 import subprocess
 import tarfile
 import tempfile
@@ -86,6 +87,34 @@ def _require_unpublished(image: str) -> None:
         )
 
 
+def _derive_dockerfile(
+    source: Path, destination: Path, build_arguments: list[str]
+) -> None:
+    """Declare non-secret build arguments in each stage of a derived Dockerfile."""
+    for name in build_arguments:
+        if re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name) is None:
+            raise ValueError(f"invalid build argument name: {name}")
+
+    lines: list[str] = []
+    stages = 0
+    for line in source.read_text().splitlines(keepends=True):
+        lines.append(line)
+        if re.match(r"^\s*FROM(?:\s|$)", line, re.IGNORECASE):
+            stages += 1
+            lines.extend(f"ARG {name}\n" for name in build_arguments)
+    if stages == 0:
+        raise ValueError(f"Dockerfile has no FROM instruction: {source}")
+    destination.write_text("".join(lines))
+
+
+def _build_argument_arguments(build_arguments: dict[str, str]) -> list[str]:
+    return [
+        argument
+        for name, value in build_arguments.items()
+        for argument in ("--build-arg", f"{name}={value}")
+    ]
+
+
 def _attestation_arguments() -> list[str]:
     return ["--provenance=mode=max,version=v1", "--sbom=true"]
 
@@ -139,10 +168,12 @@ def build(
     image_dir: Path,
     image: str,
     dockerfile: str = "Dockerfile",
+    build_arguments: dict[str, str] | None = None,
     push: bool = False,
 ) -> BuildResult:
     """Build source recorded by committed submodule gitlink."""
     image_dir = image_dir.resolve()
+    build_arguments = build_arguments or {}
     _root, source, revision = _snapshot_revision(image_dir)
     version = _version_for_revision(source, revision)
     tag = f"{image}:{version}"
@@ -154,17 +185,25 @@ def build(
         context = temporary / "src"
         metadata = temporary / "metadata.json"
         _export_snapshot(source, revision, context)
+        source_dockerfile = context / dockerfile
+        build_dockerfile = source_dockerfile
+        if build_arguments:
+            build_dockerfile = temporary / "Dockerfile.derived"
+            _derive_dockerfile(
+                source_dockerfile, build_dockerfile, list(build_arguments)
+            )
         subprocess.run(
             [
                 "docker",
                 "buildx",
                 "build",
                 "--file",
-                str(context / dockerfile),
+                str(build_dockerfile),
                 "--platform",
                 "linux/amd64",
                 "--tag",
                 tag,
+                *_build_argument_arguments(build_arguments),
                 *_attestation_arguments(),
                 "--metadata-file",
                 str(metadata),

scripts/builder/dockerfile_test.py

@@ -0,0 +1,49 @@
+import tempfile
+import unittest
+from pathlib import Path
+
+from scripts.builder.build import _derive_dockerfile
+
+
+class DerivedDockerfileTest(unittest.TestCase):
+    def test_declares_build_arguments_without_changing_source(self) -> None:
+        source_text = "# syntax=docker/dockerfile:1\nFROM example AS build\nRUN tool\n"
+        with tempfile.TemporaryDirectory() as directory:
+            source = Path(directory) / "Dockerfile"
+            destination = Path(directory) / "Derivedfile"
+            source.write_text(source_text)
+
+            _derive_dockerfile(source, destination, ["UV_HTTP_TIMEOUT"])
+
+            self.assertEqual(source.read_text(), source_text)
+            self.assertEqual(
+                destination.read_text(),
+                "# syntax=docker/dockerfile:1\n"
+                "FROM example AS build\n"
+                "ARG UV_HTTP_TIMEOUT\n"
+                "RUN tool\n",
+            )
+
+    def test_declares_arguments_in_each_stage(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            source = Path(directory) / "Dockerfile"
+            destination = Path(directory) / "Derivedfile"
+            source.write_text("FROM one\nFROM two\n")
+
+            _derive_dockerfile(source, destination, ["NETWORK_TIMEOUT"])
+
+            self.assertEqual(
+                destination.read_text(),
+                "FROM one\nARG NETWORK_TIMEOUT\nFROM two\nARG NETWORK_TIMEOUT\n",
+            )
+
+    def test_rejects_invalid_argument_name(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            source = Path(directory) / "Dockerfile"
+            source.write_text("FROM example\n")
+            with self.assertRaisesRegex(ValueError, "invalid build argument"):
+                _derive_dockerfile(source, Path(directory) / "out", ["BAD-NAME"])
+
+
+if __name__ == "__main__":
+    unittest.main()