Make platform configurable instead of only auto-detecting

manics · manics · commit 8f3f57ada738 · 2023-01-02T19:51:51.000Z
diff --git a/repo2docker/app.py b/repo2docker/app.py
@@ -15,6 +15,7 @@
 import sys
 import tempfile
 import time
+import warnings
 from urllib.parse import urlparse
 
 import entrypoints
@@ -36,7 +37,7 @@
     RBuildPack,
 )
 from .engine import BuildError, ContainerEngineException, ImageLoadError
-from .utils import ByteSpecification, R2dState, chdir
+from .utils import ByteSpecification, R2dState, chdir, get_platform
 
 
 class Repo2Docker(Application):
@@ -250,6 +251,27 @@ def _user_name_default(self):
         config=True,
     )
 
+    platform = Unicode(
+        config=True,
+        help="""
+        Platform to build for, linux/amd64 (recommended) or linux/arm64 (experimental).
+        """,
+    )
+
+    @default("platform")
+    def _platform_default(self):
+        """
+        Default platform
+        """
+        p = get_platform()
+        if p == "linux/arm64":
+            warnings.warn(
+                "Building for linux/arm64 is experimental. "
+                "To use the recommended platform set --Repo2Docker.platform=linux/amd64. "
+                "To silence this warning set --Repo2Docker.platform=linux/arm64."
+            )
+        return p
+
     extra_build_args = Dict(
         {},
         help="""
@@ -778,6 +800,7 @@ def build(self):
                 else:
                     picked_buildpack = self.default_buildpack()
 
+                picked_buildpack.platform = self.platform
                 picked_buildpack.appendix = self.appendix
                 # Add metadata labels
                 picked_buildpack.labels["repo2docker.version"] = self.version
@@ -819,6 +842,7 @@ def build(self):
                         build_args,
                         self.cache_from,
                         self.extra_build_kwargs,
+                        platform=self.platform,
                     ):
                         if docker_client.string_output:
                             self.log.info(l, extra=dict(phase=R2dState.BUILDING))
diff --git a/repo2docker/buildpacks/_r_base.py b/repo2docker/buildpacks/_r_base.py
@@ -4,15 +4,11 @@
 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,8 +11,6 @@
 import escapism
 import jinja2
 
-from ..utils import get_platform
-
 # Only use syntax features supported by Docker 17.09
 TEMPLATE = r"""
 FROM buildpack-deps:bionic
@@ -231,7 +229,7 @@ def __init__(self):
                 "Windows environment detected. Note that Windows "
                 "support is experimental in repo2docker."
             )
-        self.platform = get_platform()
+        self.platform = ""
 
     def get_packages(self):
         """
@@ -562,6 +560,7 @@ def build(
         build_args,
         cache_from,
         extra_build_kwargs,
+        platform=None,
     ):
         tarf = io.BytesIO()
         tar = tarfile.open(fileobj=tarf, mode="w")
@@ -616,6 +615,7 @@ def _filter_tar(tar):
             buildargs=build_args,
             container_limits=limits,
             cache_from=cache_from,
+            platform=platform,
         )
 
         build_kwargs.update(extra_build_kwargs)
diff --git a/repo2docker/buildpacks/conda/__init__.py b/repo2docker/buildpacks/conda/__init__.py
@@ -35,6 +35,14 @@ class CondaBuildPack(BaseImage):
     # extra pip requirements.txt for the notebook env
     _nb_requirements_file = ""
 
+    def _conda_platform(self):
+        """Return the conda platform name for the current platform"""
+        if self.platform == "linux/amd64":
+            return "linux-64"
+        if self.platform == "linux/arm64":
+            return "linux-aarch64"
+        raise ValueError(f"Unknown platform {self.platform}")
+
     def get_build_env(self):
         """Return environment variables to be set.
 
@@ -59,7 +67,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),
+            ("CONDA_PLATFORM", self._conda_platform()),
         ]
         if self._nb_requirements_file:
             env.append(("NB_REQUIREMENTS_FILE", self._nb_requirements_file))
@@ -161,7 +169,7 @@ 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 = f"environment-{self.platform}.lock"
+        frozen_name = f"environment-{self._conda_platform()}.lock"
         pip_frozen_name = "requirements.txt"
         if py_version:
             if self.python_version == "2.7":
@@ -177,13 +185,15 @@ def get_build_script_files(self):
                         self._kernel_requirements_file
                     ) = "/tmp/env/kernel-requirements.txt"
             else:
-                py_frozen_name = f"environment.py-{py_version}-{self.platform}.lock"
+                py_frozen_name = (
+                    f"environment.py-{py_version}-{self._conda_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} {self.platform} is not supported!"
+                        f"Python version {py_version} {self._conda_platform()} is not supported!"
                     )
         files[
             "conda/" + frozen_name
@@ -369,6 +379,10 @@ def get_env_scripts(self):
                 """,
                 )
             )
+            if self.platform != "linux/amd64":
+                raise RuntimeError(
+                    f"RStudio is only available for linux/amd64 ({self.platform})"
+                )
             scripts += rstudio_base_scripts(self.r_version)
             scripts += [
                 (
diff --git a/repo2docker/buildpacks/docker.py b/repo2docker/buildpacks/docker.py
@@ -30,6 +30,7 @@ def build(
         build_args,
         cache_from,
         extra_build_kwargs,
+        platform=None,
     ):
         """Build a Docker image based on the Dockerfile in the source repo."""
         # If you work on this bit of code check the corresponding code in
@@ -55,6 +56,7 @@ def build(
             container_limits=limits,
             cache_from=cache_from,
             labels=self.get_labels(),
+            platform=platform,
         )
 
         build_kwargs.update(extra_build_kwargs)
diff --git a/repo2docker/buildpacks/julia/julia_project.py b/repo2docker/buildpacks/julia/julia_project.py
@@ -87,7 +87,7 @@ def get_build_env(self):
             For example, a tuple may be `('JULIA_VERSION', '0.6.0')`.
 
         """
-        if self.platform == "linux-aarch64":
+        if self.platform == "linux/arm64":
             julia_arch = julia_arch_short = "aarch64"
         else:
             julia_arch = "x86_64"
diff --git a/repo2docker/buildpacks/julia/julia_require.py b/repo2docker/buildpacks/julia/julia_require.py
@@ -77,7 +77,7 @@ def get_build_env(self):
             For example, a tuple may be `('JULIA_VERSION', '0.6.0')`.
 
         """
-        if self.platform == "linux-aarch64":
+        if self.platform == "linux/arm64":
             julia_arch = julia_arch_short = "aarch64"
         else:
             julia_arch = "x86_64"
diff --git a/repo2docker/buildpacks/nix/__init__.py b/repo2docker/buildpacks/nix/__init__.py
@@ -31,7 +31,7 @@ def get_build_scripts(self):
          - install nix package manager for user
 
         """
-        if self.platform == "linux-aarch64":
+        if self.platform == "linux/arm64":
             nix_arch = "aarch64"
         else:
             nix_arch = "x86_64"
diff --git a/repo2docker/buildpacks/r.py b/repo2docker/buildpacks/r.py
@@ -5,7 +5,6 @@
 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
 
@@ -278,8 +277,10 @@ 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")
+        if self.platform != "linux/amd64":
+            raise RuntimeError(
+                f"RStudio is only available for linux/amd64 ({self.platform})"
+            )
         scripts = [
             (
                 "root",
diff --git a/repo2docker/docker.py b/repo2docker/docker.py
@@ -8,7 +8,6 @@
 import docker
 
 from .engine import Container, ContainerEngine, ContainerEngineException, Image
-from .utils import get_platform
 
 
 class DockerContainer(Container):
@@ -92,14 +91,9 @@ def build(
         fileobj=None,
         path="",
         labels=None,
+        platform=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,
@@ -113,7 +107,7 @@ def build(
             fileobj=fileobj,
             path=path,
             labels=labels,
-            platform=docker_platform,
+            platform=platform,
             **kwargs,
         )
 
diff --git a/repo2docker/engine.py b/repo2docker/engine.py
@@ -177,6 +177,7 @@ def build(
         fileobj=None,
         path="",
         labels=None,
+        platform=None,
         **kwargs,
     ):
         """
@@ -207,6 +208,8 @@ def build(
             path to the Dockerfile
         labels : dict
             Dictionary of labels to set on the image
+        platform: str
+            Platform to build for
 
         Returns
         -------
diff --git a/repo2docker/utils.py b/repo2docker/utils.py
@@ -529,15 +529,15 @@ def is_local_pip_requirement(line):
 
 
 def get_platform():
-    """Return the platform of the image
+    """Return the target platform of the container image
 
-    Returns either `linux-64` or `linux-aarch64`
+    Returns either `linux/amd64` or `linux/arm64`
     """
     m = platform.machine()
     if m == "x86_64":
-        return "linux-64"
+        return "linux/amd64"
     elif m == "aarch64":
-        return "linux-aarch64"
+        return "linux/arm64"
     else:
-        warnings.warn(f"Unexpected platform '{m}', defaulting to linux-64")
-        return "linux-64"
+        warnings.warn(f"Unexpected platform '{m}', defaulting to linux/amd64")
+        return "linux/amd64"
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
@@ -2,6 +2,7 @@
 Tests for repo2docker/utils.py
 """
 import os
+import platform
 import subprocess
 import tempfile
 
@@ -160,3 +161,16 @@ def test_open_guess_encoding():
 )
 def test_local_pip_requirement(req, is_local):
     assert utils.is_local_pip_requirement(req) == is_local
+
+
+@pytest.mark.parametrize(
+    "machine_name,expected",
+    [
+        ("x86_64", "linux/amd64"),
+        ("aarch64", "linux/arm64"),
+        ("other", "linux/amd64"),
+    ],
+)
+def test_get_platform(monkeypatch, machine_name, expected):
+    monkeypatch.setattr(platform, "machine", lambda: machine_name)
+    assert utils.get_platform() == expected
