BuildPacks: support arm64 where possible (i.e. not R)

manics · manics · commit 064d91ae2248 · 2023-01-02T19:51:51.000Z
diff --git a/repo2docker/buildpacks/_r_base.py b/repo2docker/buildpacks/_r_base.py
@@ -4,11 +4,15 @@
 Keeping this in r.py would lead to cyclic imports.
 """
 from ..semver import parse_version as V
+from ..utils import get_platform
 
 
 def rstudio_base_scripts(r_version):
     """Base steps to install RStudio and shiny-server."""
 
+    if get_platform() != "linux-64":
+        raise RuntimeError("RStudio is only available for linux-64")
+
     # Shiny server (not the package!) seems to be the same version for all R versions
     shiny_server_url = "https://download3.rstudio.org/ubuntu-14.04/x86_64/shiny-server-1.5.17.973-amd64.deb"
     shiny_proxy_version = "1.1"
diff --git a/repo2docker/buildpacks/base.py b/repo2docker/buildpacks/base.py
@@ -11,6 +11,8 @@
 import escapism
 import jinja2
 
+from ..utils import get_platform
+
 # Only use syntax features supported by Docker 17.09
 TEMPLATE = r"""
 FROM buildpack-deps:bionic
@@ -229,6 +231,7 @@ def __init__(self):
                 "Windows environment detected. Note that Windows "
                 "support is experimental in repo2docker."
             )
+        self.platform = get_platform()
 
     def get_packages(self):
         """
diff --git a/repo2docker/buildpacks/conda/__init__.py b/repo2docker/buildpacks/conda/__init__.py
@@ -59,6 +59,7 @@ def get_build_env(self):
             # this exe should be used for installs after bootstrap with micromamba
             # switch this to /usr/local/bin/micromamba to use it for all installs
             ("MAMBA_EXE", "${CONDA_DIR}/bin/mamba"),
+            ("CONDA_PLATFORM", self.platform),
         ]
         if self._nb_requirements_file:
             env.append(("NB_REQUIREMENTS_FILE", self._nb_requirements_file))
@@ -160,13 +161,13 @@ def get_build_script_files(self):
         # major Python versions during upgrade.
         # If no version is specified or no matching X.Y version is found,
         # the default base environment is used.
-        frozen_name = "environment.lock"
+        frozen_name = f"environment-{self.platform}.lock"
         pip_frozen_name = "requirements.txt"
         if py_version:
             if self.python_version == "2.7":
                 # python 2 goes in a different env
                 files[
-                    "conda/environment.py-2.7.lock"
+                    "conda/environment.py-2.7-linux-64.lock"
                 ] = self._kernel_environment_file = "/tmp/env/kernel-environment.lock"
                 # additional pip requirements for kernel env
                 if os.path.exists(os.path.join(HERE, "requirements.py-2.7.txt")):
@@ -176,12 +177,14 @@ def get_build_script_files(self):
                         self._kernel_requirements_file
                     ) = "/tmp/env/kernel-requirements.txt"
             else:
-                py_frozen_name = f"environment.py-{py_version}.lock"
+                py_frozen_name = f"environment.py-{py_version}-{self.platform}.lock"
                 if os.path.exists(os.path.join(HERE, py_frozen_name)):
                     frozen_name = py_frozen_name
                     pip_frozen_name = f"requirements.py-{py_version}.pip"
                 else:
-                    raise ValueError(f"Python version {py_version} is not supported!")
+                    raise ValueError(
+                        f"Python version {py_version} {self.platform} is not supported!"
+                    )
         files[
             "conda/" + frozen_name
         ] = self._nb_environment_file = "/tmp/env/environment.lock"
diff --git a/repo2docker/buildpacks/conda/freeze.py b/repo2docker/buildpacks/conda/freeze.py
@@ -131,7 +131,9 @@ def set_python(py_env_file, py):
         for platform in args.platform:
             env_file = pathlib.Path(str(ENV_FILE_T).format(py=py))
             set_python(env_file, py)
-            frozen_file = pathlib.Path(os.path.splitext(env_file)[0] + f"-{platform}.lock")
+            frozen_file = pathlib.Path(
+                os.path.splitext(env_file)[0] + f"-{platform}.lock"
+            )
             freeze(env_file, frozen_file, platform)
             if py == default_py:
                 shutil.copy(frozen_file, FROZEN_FILE_T.format(platform=platform))
diff --git a/repo2docker/buildpacks/conda/install-base-env.bash b/repo2docker/buildpacks/conda/install-base-env.bash
@@ -8,7 +8,7 @@ cd $(dirname $0)
 export MAMBA_VERSION=1.0.0
 export CONDA_VERSION=4.13.0
 
-URL="https://anaconda.org/conda-forge/micromamba/${MAMBA_VERSION}/download/linux-64/micromamba-${MAMBA_VERSION}-0.tar.bz2"
+URL="https://anaconda.org/conda-forge/micromamba/${MAMBA_VERSION}/download/${CONDA_PLATFORM}/micromamba-${MAMBA_VERSION}-0.tar.bz2"
 
 # make sure we don't do anything funky with user's $HOME
 # since this is run as root
diff --git a/repo2docker/buildpacks/julia/julia_project.py b/repo2docker/buildpacks/julia/julia_project.py
@@ -79,16 +79,25 @@ def get_build_env(self):
                 will be installed
             - `JULIA_DEPOT_PATH`: path where Julia libraries are installed.
             - `JULIA_VERSION`: default version of julia to be installed
+            - `JULIA_ARCH`: machine architecture used in Julia download URLs
+            - `JULIA_ARCH_SHORT`: machine architecture used in Julia download URLs
             - `JUPYTER`: environment variable required by IJulia to point to
                 the `jupyter` executable
 
             For example, a tuple may be `('JULIA_VERSION', '0.6.0')`.
 
         """
+        if self.platform == "linux-aarch64":
+            julia_arch = julia_arch_short = "aarch64"
+        else:
+            julia_arch = "x86_64"
+            julia_arch_short = "x64"
         return super().get_build_env() + [
             ("JULIA_PATH", "${APP_BASE}/julia"),
             ("JULIA_DEPOT_PATH", "${JULIA_PATH}/pkg"),
             ("JULIA_VERSION", self.julia_version),
+            ("JULIA_ARCH", julia_arch),
+            ("JULIA_ARCH_SHORT", julia_arch_short),
             ("JUPYTER", "${NB_PYTHON_PREFIX}/bin/jupyter"),
             ("JUPYTER_DATA_DIR", "${NB_PYTHON_PREFIX}/share/jupyter"),
         ]
@@ -129,7 +138,7 @@ def get_build_scripts(self):
                 "root",
                 r"""
                 mkdir -p ${JULIA_PATH} && \
-                curl -sSL "https://julialang-s3.julialang.org/bin/linux/x64/${JULIA_VERSION%[.-]*}/julia-${JULIA_VERSION}-linux-x86_64.tar.gz" | tar -xz -C ${JULIA_PATH} --strip-components 1
+                curl -sSL "https://julialang-s3.julialang.org/bin/linux/${JULIA_ARCH_SHORT}/${JULIA_VERSION%[.-]*}/julia-${JULIA_VERSION}-linux-${JULIA_ARCH}.tar.gz" | tar -xz -C ${JULIA_PATH} --strip-components 1
                 """,
             ),
             (
diff --git a/repo2docker/buildpacks/julia/julia_require.py b/repo2docker/buildpacks/julia/julia_require.py
@@ -69,19 +69,28 @@ def get_build_env(self):
             - `JULIA_DEPOT_PATH`: path where Julia libraries are installed.
                                   Similar to JULIA_PKGDIR, used in 1.x.
             - `JULIA_VERSION`: default version of julia to be installed
+            - `JULIA_ARCH`: machine architecture used in Julia download URLs
+            - `JULIA_ARCH_SHORT`: machine architecture used in Julia download URLs
             - `JUPYTER`: environment variable required by IJulia to point to
                 the `jupyter` executable
 
             For example, a tuple may be `('JULIA_VERSION', '0.6.0')`.
 
         """
+        if self.platform == "linux-aarch64":
+            julia_arch = julia_arch_short = "aarch64"
+        else:
+            julia_arch = "x86_64"
+            julia_arch_short = "x64"
         return super().get_build_env() + [
             ("JULIA_PATH", "${APP_BASE}/julia"),
             ("JULIA_HOME", "${JULIA_PATH}/bin"),  # julia <= 0.6
             ("JULIA_BINDIR", "${JULIA_HOME}"),  # julia >= 0.7
             ("JULIA_PKGDIR", "${JULIA_PATH}/pkg"),
             ("JULIA_DEPOT_PATH", "${JULIA_PKGDIR}"),  # julia >= 0.7
             ("JULIA_VERSION", self.julia_version),
+            ("JULIA_ARCH", julia_arch),
+            ("JULIA_ARCH_SHORT", julia_arch_short),
             ("JUPYTER", "${NB_PYTHON_PREFIX}/bin/jupyter"),
         ]
 
@@ -111,7 +120,7 @@ def get_build_scripts(self):
                 "root",
                 r"""
                 mkdir -p ${JULIA_PATH} && \
-                curl -sSL "https://julialang-s3.julialang.org/bin/linux/x64/${JULIA_VERSION%[.-]*}/julia-${JULIA_VERSION}-linux-x86_64.tar.gz" | tar -xz -C ${JULIA_PATH} --strip-components 1
+                curl -sSL "https://julialang-s3.julialang.org/bin/linux/${JULIA_ARCH_SHORT}/${JULIA_VERSION%[.-]*}/julia-${JULIA_VERSION}-linux-${JULIA_ARCH}.tar.gz" | tar -xz -C ${JULIA_PATH} --strip-components 1
                 """,
             ),
             (
diff --git a/repo2docker/buildpacks/nix/__init__.py b/repo2docker/buildpacks/nix/__init__.py
@@ -13,6 +13,7 @@ def get_path(self):
 
     def get_env(self):
         """Ordered list of environment variables to be set for this image"""
+
         return super().get_env() + [
             ("NIX_PATH", "nixpkgs=/home/${NB_USER}/.nix-defexpr/channels/nixpkgs"),
             ("NIX_SSL_CERT_FILE", "/etc/ssl/certs/ca-certificates.crt"),
@@ -30,6 +31,10 @@ def get_build_scripts(self):
          - install nix package manager for user
 
         """
+        if self.platform == "linux-aarch64":
+            nix_arch = "aarch64"
+        else:
+            nix_arch = "x86_64"
         return super().get_build_scripts() + [
             (
                 "root",
@@ -43,9 +48,9 @@ def get_build_scripts(self):
             ),
             (
                 "${NB_USER}",
-                """
-            bash /home/${NB_USER}/.local/bin/install-nix.bash && \
-            rm /home/${NB_USER}/.local/bin/install-nix.bash
+                f"""
+            NIX_ARCH={nix_arch} bash /home/${{NB_USER}}/.local/bin/install-nix.bash && \
+            rm /home/${{NB_USER}}/.local/bin/install-nix.bash
             """,
             ),
         ]
diff --git a/repo2docker/buildpacks/nix/install-nix.bash b/repo2docker/buildpacks/nix/install-nix.bash
@@ -3,12 +3,16 @@
 set -ex
 
 NIX_VERSION="2.3.9"
-NIX_SHA256="49763fd7fa06bcb712ced2f3f11afd275e3a4d7bc5ff0d6fd1d50a4c3ce7bbf4"
+if [ "$NIX_ARCH" = "aarch64" ]; then
+  NIX_SHA256="733a26911193fdd44d5d68342075af5924d8c0701aae877e51a38d74ee9f4ff8"
+else
+  NIX_SHA256="49763fd7fa06bcb712ced2f3f11afd275e3a4d7bc5ff0d6fd1d50a4c3ce7bbf4"
+fi
 
 # Do all our operations in /tmp, since we can't rely on current directory being writeable yet.
 cd /tmp
-wget --quiet https://nixos.org/releases/nix/nix-$NIX_VERSION/nix-$NIX_VERSION-x86_64-linux.tar.xz
-echo "$NIX_SHA256  nix-$NIX_VERSION-x86_64-linux.tar.xz" | sha256sum -c
-tar xJf nix-*-x86_64-linux.tar.xz
-sh nix-*-x86_64-linux/install
-rm -r nix-*-x86_64-linux*
+wget --quiet https://nixos.org/releases/nix/nix-$NIX_VERSION/nix-$NIX_VERSION-$NIX_ARCH-linux.tar.xz
+echo "$NIX_SHA256  nix-$NIX_VERSION-$NIX_ARCH-linux.tar.xz" | sha256sum -c
+tar xJf nix-*-$NIX_ARCH-linux.tar.xz
+sh nix-*-$NIX_ARCH-linux/install
+rm -r nix-*-$NIX_ARCH-linux*
diff --git a/repo2docker/buildpacks/r.py b/repo2docker/buildpacks/r.py
@@ -5,6 +5,7 @@
 import requests
 
 from ..semver import parse_version as V
+from ..utils import get_platform
 from ._r_base import rstudio_base_scripts
 from .python import PythonBuildPack
 
@@ -277,6 +278,8 @@ def get_build_scripts(self):
 
         cran_mirror_url = self.get_cran_mirror_url(self.checkpoint_date)
 
+        if get_platform() != "linux-64":
+            raise RuntimeError("RStudio is only available for linux-64")
         scripts = [
             (
                 "root",
diff --git a/repo2docker/docker.py b/repo2docker/docker.py
@@ -8,6 +8,7 @@
 import docker
 
 from .engine import Container, ContainerEngine, ContainerEngineException, Image
+from .utils import get_platform
 
 
 class DockerContainer(Container):
@@ -93,6 +94,12 @@ def build(
         labels=None,
         **kwargs,
     ):
+        platform = get_platform()
+        if platform == "linux-aarch64":
+            docker_platform = "linux/arm64"
+        else:
+            docker_platform = "linux/amd64"
+
         return self._apiclient.build(
             buildargs=buildargs,
             cache_from=cache_from,
@@ -106,7 +113,7 @@ def build(
             fileobj=fileobj,
             path=path,
             labels=labels,
-            platform="linux/amd64",
+            platform=docker_platform,
             **kwargs,
         )
 
diff --git a/repo2docker/utils.py b/repo2docker/utils.py
@@ -1,6 +1,8 @@
 import os
+import platform
 import re
 import subprocess
+import warnings
 from contextlib import contextmanager
 from enum import Enum
 from functools import partial
@@ -524,3 +526,18 @@ def is_local_pip_requirement(line):
         return True
 
     return False
+
+
+def get_platform():
+    """Return the platform of the image
+
+    Returns either `linux-64` or `linux-aarch64`
+    """
+    m = platform.machine()
+    if m == "x86_64":
+        return "linux-64"
+    elif m == "aarch64":
+        return "linux-aarch64"
+    else:
+        warnings.warn(f"Unexpected platform '{m}', defaulting to linux-64")
+        return "linux-64"
