feat: Multi-arch Python base image (#614)

.changeset/bunq2ynab-ninety-planets-yell.md

@@ -0,0 +1,5 @@
+---
+"bunq2ynab": minor
+---
+
+feat: Release a multi-architecture docker image, both in amd64 and arm64.

.gitignore

@@ -50,4 +50,5 @@ node_modules
 .pdm-python
 __pypackages__
 keys
-nixos.qcow2
+nixos.qcow2
+result

BUILD.bazel

@@ -68,7 +68,7 @@ platform(
 )
 
 platform(
-    name = "linux_x86_64",
+    name = "linux_amd64",
     constraint_values = _linux_amd64,
 )
 
@@ -77,6 +77,20 @@ platform(
     constraint_values = _linux_arm64,
 )
 
+platform(
+    name = "python_container_linux_amd64",
+    constraint_values = _linux_amd64 + [
+        "//tools/python:python_run_in_container",
+    ],
+)
+
+platform(
+    name = "python_container_linux_arm64",
+    constraint_values = _linux_arm64 + [
+        "//tools/python:python_run_in_container",
+    ],
+)
+
 pycross_target_environment(
     name = "python_darwin_arm64",
     abis = ["cp310"],
@@ -238,3 +252,13 @@ release_manager(
         "@rules_task//:release",
     ],
 )
+
+alias(
+    name = "regctl",
+    actual = "//tools/regctl",
+)
+
+alias(
+    name = "tsh",
+    actual = "//tools/teleport:tsh",
+)

MODULE.bazel

@@ -48,7 +48,7 @@ python.toolchain(
 use_repo(python, "python_versions")
 
 # ------------------------------------ rules_oci ------------------------------------ #
-bazel_dep(name = "rules_oci", version = "1.4.0")
+bazel_dep(name = "rules_oci", version = "1.4.3")
 
 # ------------------------------------ rules_release ------------------------------------ #
 bazel_dep(name = "rules_release", version = "0.0.0")

WORKSPACE.bzlmod

@@ -318,11 +318,31 @@ nixpkgs_git_repository(
 )
 
 nixpkgs_package(
-    name = "python310_base_image",
+    name = "python_base_image_amd64",
     build_file_content = """
 package(default_visibility = [ "//visibility:public" ])
 exports_files(["image.tar.gz"])
     """,
-    nix_file = "//:python310_base_image.nix",
+    nix_file = "//tools/python:python_base_image.nix",
+    nixopts = [
+        "--argstr",
+        "targetArch",
+        "x86_64",
+    ],
+    repository = "@nixpkgs//:default.nix",
+)
+
+nixpkgs_package(
+    name = "python_base_image_arm64",
+    build_file_content = """
+package(default_visibility = [ "//visibility:public" ])
+exports_files(["image.tar.gz"])
+    """,
+    nix_file = "//tools/python:python_base_image.nix",
+    nixopts = [
+        "--argstr",
+        "targetArch",
+        "aarch64",
+    ],
     repository = "@nixpkgs//:default.nix",
 )

tools/bunq2ynab/BUILD.bazel

@@ -119,25 +119,28 @@ task(
 
 py_image(
     name = "bunq2ynab_image",
-    base = "@python310_base_image//:image.tar.gz",
+    base = "//tools/python:python_base_image_file",
     binary = ":bunq2ynab",
-    host_container_platform = "//:host_container_platform",
+    platforms = [
+        "//:python_container_linux_amd64",
+        "//:python_container_linux_arm64",
+    ],
     prefix = "opt/",
 )
 
 task(
     name = "bunq2ynab_image_run",
     cmds = [
         cmd.executable("bunq2ynab_image.load"),
-        "docker run --rm --env OP_SERVICE_ACCOUNT_TOKEN=$ONEPASSWORD_SERVICE_ACCOUNT_TOKEN_DEV -it --entrypoint='' bunq2ynab:latest $CLI_ARGS",
+        "docker run --rm --env OP_SERVICE_ACCOUNT_TOKEN=$ONEPASSWORD_SERVICE_ACCOUNT_TOKEN_DEV -it --entrypoint='' localhost/bunq2ynab:latest $CLI_ARGS",
     ],
 )
 
 task_test(
     name = "bunq2ynab_image_test",
     cmds = [
         cmd.executable("bunq2ynab_image.load"),
-        "docker run --rm --env OP_SERVICE_ACCOUNT_TOKEN=$ONEPASSWORD_SERVICE_ACCOUNT_TOKEN_DEV bunq2ynab:latest",
+        "docker run --rm --env OP_SERVICE_ACCOUNT_TOKEN=$ONEPASSWORD_SERVICE_ACCOUNT_TOKEN_DEV localhost/bunq2ynab:latest",
     ],
     exec_properties = {
         "workload-isolation-type": "firecracker",

tools/python/BUILD.bazel

@@ -1,12 +1,17 @@
 load("@bazel_tools//tools/python:toolchain.bzl", "py_runtime_pair")
-load(":defs.bzl", "host_python_container_platform")
 
 package(default_visibility = ["//visibility:public"])
 
 constraint_setting(name = "python_containerized")
 
-host_python_container_platform(
-    name = "host_python_container",
+python_base_image_file = select({
+    "//:is_linux_amd64": ["@python_base_image_amd64//:image.tar.gz"],
+    "//:is_linux_arm64": ["@python_base_image_arm64//:image.tar.gz"],
+})
+
+filegroup(
+    name = "python_base_image_file",
+    srcs = python_base_image_file,
 )
 
 constraint_value(

tools/python/defs.bzl

@@ -1,25 +1,9 @@
 load("@aspect_bazel_lib//lib:tar.bzl", "mtree_spec", "tar")
-load("@rules_oci//oci:defs.bzl", "oci_image", "oci_tarball")
+load("@rules_oci//oci:defs.bzl", "oci_image", "oci_image_index", "oci_tarball")
 load("@aspect_bazel_lib//lib:transitions.bzl", "platform_transition_filegroup")
 load("@rules_task//task:defs.bzl", "cmd", "task")
 load("@local_config_platform//:constraints.bzl", "HOST_CONSTRAINTS")
 
-# This sets up a platform for the Python toolchain to run in a container.
-# This is way to prevent the Python hermetic interpreter to be copied into the container
-# And instead we rely on a shebang and a Python interpreter in the container base image.
-# Currently this only works for the host cpu, but can be extended for multi-arch images
-def host_python_container_platform(name):
-    host_cpu, _host_os = HOST_CONSTRAINTS
-
-    native.platform(
-        name = name,
-        constraint_values = [
-            host_cpu,
-            "@platforms//os:linux",
-            "//tools/python:python_run_in_container",
-        ],
-    )
-
 def py_image_layer(name, binary, prefix = "", **kwargs):
     mtree_spec_name = "{}_mtree".format(name)
     prefixed_mtree_spec_name = "{}_prefixed".format(mtree_spec_name)
@@ -44,16 +28,16 @@ def py_image_layer(name, binary, prefix = "", **kwargs):
         mtree = prefixed_mtree_spec_name,
     )
 
-def py_image(name, base, binary, host_container_platform, prefix = ""):
+def py_image(name, base, binary, platforms, prefix = ""):
     binary_name = Label(binary).name
     package_name = native.package_name()
     entrypoint = ["/{}{}/{}".format(prefix, package_name, binary_name)]
 
-    image_name = name
-    transitioned_image = "{}_transitioned".format(name)
-    image_load_name = "{}.load".format(name)
-    image_python_layer_name = "{}_python_layer".format(name)
-    tarball_name = "{}.tarball".format(transitioned_image)
+    image_index_name = name
+    image_name = "{}.image".format(image_index_name)
+    image_load_name = "{}.load".format(image_index_name)
+    image_python_layer_name = "{}_python_layer".format(image_index_name)
+    tarball_name = "{}.tarball".format(image_index_name)
 
     repo_tags = [
         "{}:{}".format(binary_name, "latest"),
@@ -74,27 +58,49 @@ def py_image(name, base, binary, host_container_platform, prefix = ""):
         ],
     )
 
-    # This can be extended to multi-arch images. For example see:
-    # https://github.com/macourteau/aspect-rules_oci/blob/master/container.bzl#L85
-    platform_transition_filegroup(
-        name = transitioned_image,
-        srcs = [image_name],
-        target_platform = "//tools/python:host_python_container",
+    transitioned_images = []
+
+    for platform in platforms:
+        platform_name = Label(platform).name
+        transitioned_image = "{}_{}".format(image_name, platform_name)
+
+        platform_transition_filegroup(
+            name = transitioned_image,
+            srcs = [image_name],
+            target_platform = platform,
+        )
+
+        transitioned_images.append(transitioned_image)
+
+    oci_image_index(
+        name = image_index_name,
+        images = transitioned_images,
     )
 
     oci_tarball(
         name = tarball_name,
-        image = transitioned_image,
+        image = image_index_name,
         repo_tags = repo_tags,
+        format = "oci",
     )
 
+    # From https://stackoverflow.com/questions/72945407/how-do-i-import-and-run-a-multi-platform-oci-image-in-docker-for-macos
+    # We need to load the multi-arch image using regctl
+    # Export the platform specific digest into a tar
+    # And load that tar into the daemon
     task(
         name = image_load_name,
         cmds = [
-            "docker load < $TARBALL",
+            "$REGCTL image import ocidir://{} $TARBALL".format(binary_name),
+            "digest=$($REGCTL image digest --platform local ocidir://{})".format(binary_name),
+            "export LOCAL_TARBALL=$(pwd)/{}.tar".format(binary_name),
+            "$REGCTL image export ocidir://{}@$digest $LOCAL_TARBALL".format(binary_name),
+            {"defer": "rm -f $LOCAL_TARBALL"},
+            "docker load < $LOCAL_TARBALL",
         ],
         env = {
             "TARBALL": cmd.file(tarball_name),
+            "REGCTL": cmd.executable("//tools/regctl:regctl"),
         },
         exec_properties = {
             "workload-isolation-type": "firecracker",

python310_base_image.nix → tools/python/python_base_image.nix

@@ -1,35 +1,29 @@
-let
-  currentSystem = builtins.currentSystem;
-  targetSystem = {
-    "x86_64-linux" = "x86_64-linux";
-    "aarch64-darwin" = "aarch64-linux";
-    "aarch64-linux" = "aarch64-linux";
-  }."${currentSystem}";
-in
-with import <nixpkgs>
-{
-  system = targetSystem;
-};
+{ targetArch }:
 
 let
-  dockerEtc = runCommand "docker-etc" { } ''
+  localPkgs = import <nixpkgs> { };
+  targetPkgs = import <nixpkgs> { system = targetArch + "-linux"; };
+
+  dockerEtc = localPkgs.runCommand "docker-etc" { } ''
     mkdir -p $out/etc/pam.d
 
     echo "root:x:0:0::/root:/bin/bash" > $out/etc/passwd
     echo "root:!x:::::::" > $out/etc/shadow
     echo "root:x:0:" > $out/etc/group
   '';
 
-  pythonBase = dockerTools.buildLayeredImage {
-    name = "python310-base-image-unwrapped";
+  pythonBaseImage = localPkgs.dockerTools.buildLayeredImage {
+    name = "python_base_image";
+    tag = "latest";
     created = "now";
+    architecture = targetArch;
     maxLayers = 2;
     contents = [
-      busybox
-      bashInteractive
-      python310
-      stdenv.cc.cc.lib
-      cacert
+      targetPkgs.busybox
+      targetPkgs.bashInteractive
+      targetPkgs.python310
+      targetPkgs.stdenv.cc.cc.lib
+      targetPkgs.cacert
       dockerEtc
     ];
     extraCommands = ''
@@ -45,9 +39,9 @@ let
             ln -s /usr/bin/python3 usr/bin/python
     '';
   };
-
 in
-runCommand "python310-base-image" { } ''
+localPkgs.runCommand "pythonBaseImage" { } ''
   mkdir -p $out
-  gunzip -c ${pythonBase} > $out/image.tar.gz
+  gunzip -c ${pythonBaseImage} > $out/image.tar.gz
 ''
+