[ci] arm64 wheels (ray-project#41056)

Build arm64 wheels in civ2 Signed-off-by: can <can@anyscale.com>
ujjawal-khare-27 · Nov 29, 2023 · ff16a00 · ff16a00
1 parent 89f0640
commit ff16a00
Show file tree

Hide file tree

Showing 13 changed files with 105 additions and 38 deletions.
diff --git a/.buildkite/_forge.rayci.yml b/.buildkite/_forge.rayci.yml
@@ -3,9 +3,17 @@ steps:
   - name: forge
     wanda: ci/docker/forge.wanda.yaml
 
+  - name: forge-aarch64
+    wanda: ci/docker/forge.aarch64.wanda.yaml
+    instance_type: builder-arm64
+
   - name: manylinux
     wanda: ci/docker/manylinux.wanda.yaml
 
+  - name: manylinux-aarch64
+    wanda: ci/docker/manylinux.aarch64.wanda.yaml
+    instance_type: builder-arm64
+
   - name: raycudabase
     label: "wanda: ray.py{{matrix.python}}.cu{{matrix.cuda}}.base"
     tags:

diff --git a/.buildkite/build.rayci.yml b/.buildkite/build.rayci.yml
@@ -1,12 +1,12 @@
 group: build
 steps:
-  - label: ":tapioca: build: wheel {{matrix}}"
+  - label: ":tapioca: build: wheel {{matrix}} (x86_64)"
     tags:
       - linux_wheels
       - oss
     instance_type: medium
     commands:
-      - bazel run //ci/ray_ci:build_in_docker -- wheel --python-version {{matrix}}
+      - bazel run //ci/ray_ci:build_in_docker -- wheel --python-version {{matrix}} --architecture x86_64
     matrix:
       - "3.8"
       - "3.9"
@@ -17,6 +17,23 @@ steps:
       - forge
     job_env: forge
 
+  - label: ":tapioca: build: wheel {{matrix}} (aarch64)"
+    tags:
+      - linux_wheels
+      - oss
+    instance_type: medium-arm64
+    commands:
+      - bazel run //ci/ray_ci:build_in_docker -- wheel --python-version {{matrix}} --architecture aarch64
+    matrix:
+      - "3.8"
+      - "3.9"
+      - "3.10"
+      - "3.11"
+    depends_on:
+      - manylinux-aarch64
+      - forge-aarch64
+    job_env: forge-aarch64
+
   - label: ":tapioca: build: debug wheel"
     tags:
       - linux_wheels

diff --git a/.buildkite/pipeline.arm64.yml b/.buildkite/pipeline.arm64.yml
@@ -1,21 +1,5 @@
 #ci:group=:mechanical_arm: :ferris_wheel: ARM64 wheels and containers
 
-- label: ":mechanical_arm: :ferris_wheel: ARM64 Linux wheels"
-  conditions: ["RAY_CI_LINUX_WHEELS_AFFECTED"]
-  instance_size: arm64-medium
-  commands:
-    # Build the wheels
-    - UPLOAD_WHEELS_AS_ARTIFACTS=1 LINUX_WHEELS=1 ./ci/ci.sh build
-    # Upload the wheels
-    # We don't want to push on PRs, in fact, the copy_files will fail because unauthenticated.
-    - if [ "$BUILDKITE_PULL_REQUEST" != "false" ]; then exit 0; fi
-    - pip install -q docker aws_requests_auth boto3
-    - ./ci/env/env_info.sh
-    # Upload to branch directory.
-    - python .buildkite/copy_files.py --destination branch_wheels --path ./.whl
-    # Upload to latest directory.
-    - if [ "$BUILDKITE_BRANCH" == "master" ]; then python .buildkite/copy_files.py --destination wheels --path ./.whl; fi
-
 - label: ":mechanical_arm: :ferris_wheel: ARM64 Post-wheels tests"
   conditions: ["RAY_CI_LINUX_WHEELS_AFFECTED"]
   instance_size: arm64-medium

diff --git a/ci/docker/forge.aarch64.wanda.yaml b/ci/docker/forge.aarch64.wanda.yaml
@@ -0,0 +1,6 @@
+name: "forge-aarch64"
+froms:
+  - ubuntu:20.04
+build_args:
+  - BUILDKITE_BAZEL_CACHE_URL
+dockerfile: ci/docker/forge.Dockerfile
diff --git a/ci/docker/manylinux.Dockerfile b/ci/docker/manylinux.Dockerfile
@@ -1,6 +1,7 @@
 # syntax=docker/dockerfile:1.3-labs
 
-FROM quay.io/pypa/manylinux2014_x86_64:2022-12-20-b4884d9
+ARG HOSTTYPE
+FROM quay.io/pypa/manylinux2014_${HOSTTYPE}:2022-12-20-b4884d9
 
 ARG BUILDKITE_BAZEL_CACHE_URL
 

diff --git a/ci/docker/manylinux.aarch64.wanda.yaml b/ci/docker/manylinux.aarch64.wanda.yaml
@@ -0,0 +1,9 @@
+name: "manylinux-aarch64"
+froms:
+  - quay.io/pypa/manylinux2014_aarch64:2022-12-20-b4884d9
+srcs:
+  - ci/build/build-manylinux-forge.sh
+build_args:
+  - BUILDKITE_BAZEL_CACHE_URL
+  - HOSTTYPE=aarch64
+dockerfile: ci/docker/manylinux.Dockerfile
diff --git a/ci/docker/manylinux.wanda.yaml b/ci/docker/manylinux.wanda.yaml
@@ -5,4 +5,5 @@ srcs:
   - ci/build/build-manylinux-forge.sh
 build_args:
   - BUILDKITE_BAZEL_CACHE_URL
+  - HOSTTYPE=x86_64
 dockerfile: ci/docker/manylinux.Dockerfile
diff --git a/ci/pipeline/determine_tests_to_run.py b/ci/pipeline/determine_tests_to_run.py
@@ -250,7 +250,6 @@ def get_commit_range():
                         break
             elif (
                 changed_file == ".buildkite/core.rayci.yml"
-                or changed_file == ".buildkite/_forge.rayci.yml"
                 or changed_file == "ci/docker/min.build.Dockerfile"
                 or changed_file == "ci/docker/min.build.wanda.yaml"
                 or changed_file == ".buildkite/serverless.rayci.yml"
@@ -310,6 +309,9 @@ def get_commit_range():
             elif (
                 changed_file.startswith("ci/pipeline")
                 or changed_file.startswith("ci/ray_ci")
+                or changed_file == ".buildkite/_forge.rayci.yml"
+                or changed_file == "ci/docker/forge.wanda.yaml"
+                or changed_file == "ci/docker/forge.aarch64.wanda.yaml"
                 or changed_file == ".buildkite/pipeline.build.yml"
                 or changed_file == ".buildkite/pipeline.ml.yml"
             ):
@@ -318,6 +320,10 @@ def get_commit_range():
             elif (
                 changed_file.endswith("build-docker-images.py")
                 or changed_file == ".buildkite/build.rayci.yml"
+                or changed_file == ".buildkite/pipeline.arm64.yml"
+                or changed_file == "ci/docker/manylinux.Dockerfile"
+                or changed_file == "ci/docker/manylinux.wanda.yaml"
+                or changed_file == "ci/docker/manylinux.aarch64.wanda.yaml"
             ):
                 RAY_CI_DOCKER_AFFECTED = 1
                 RAY_CI_LINUX_WHEELS_AFFECTED = 1

diff --git a/ci/ray_ci/builder.py b/ci/ray_ci/builder.py
@@ -2,7 +2,12 @@
 
 import click
 
-from ci.ray_ci.builder_container import PYTHON_VERSIONS, BUILD_TYPES, BuilderContainer
+from ci.ray_ci.builder_container import (
+    PYTHON_VERSIONS,
+    BUILD_TYPES,
+    ARCHITECTURE,
+    BuilderContainer,
+)
 from ci.ray_ci.doc_builder_container import DocBuilderContainer
 from ci.ray_ci.forge_container import ForgeContainer
 from ci.ray_ci.docker_container import PLATFORM
@@ -40,32 +45,39 @@
     type=click.Choice(list(PLATFORM)),
     help=("Platform to build the docker with"),
 )
+@click.option(
+    "--architecture",
+    default="x86_64",
+    type=click.Choice(list(ARCHITECTURE)),
+    help=("Platform to build the docker with"),
+)
 def main(
     artifact_type: str,
     image_type: str,
     build_type: str,
     python_version: str,
     platform: List[str],
+    architecture: str,
 ) -> None:
     """
     Build a wheel or jar artifact
     """
     docker_login(_DOCKER_ECR_REPO.split("/")[0])
     if artifact_type == "wheel":
         logger.info(f"Building wheel for {python_version}")
-        build_wheel(python_version, build_type)
+        build_wheel(python_version, build_type, architecture)
         return
 
     if artifact_type == "docker":
         logger.info(f"Building {image_type} docker for {python_version} on {platform}")
-        build_docker(image_type, python_version, build_type, platform)
+        build_docker(image_type, python_version, build_type, platform, architecture)
         return
 
     if artifact_type == "anyscale":
         logger.info(
             f"Building {image_type} anyscale for {python_version} on {platform}"
         )
-        build_anyscale(image_type, python_version, build_type, platform)
+        build_anyscale(image_type, python_version, build_type, platform, architecture)
         return
 
     if artifact_type == "doc":
@@ -76,32 +88,40 @@ def main(
     raise ValueError(f"Invalid artifact type {artifact_type}")
 
 
-def build_wheel(python_version: str, build_type: str) -> None:
+def build_wheel(python_version: str, build_type: str, architecture: str) -> None:
     """
     Build a wheel artifact.
     """
-    BuilderContainer(python_version, build_type).run()
-    ForgeContainer().upload_wheel()
+    BuilderContainer(python_version, build_type, architecture).run()
+    ForgeContainer(architecture).upload_wheel()
 
 
 def build_docker(
-    image_type: str, python_version: str, build_type: str, platform: List[str]
+    image_type: str,
+    python_version: str,
+    build_type: str,
+    platform: List[str],
+    architecture: str,
 ) -> None:
     """
     Build a container artifact.
     """
-    BuilderContainer(python_version, build_type).run()
+    BuilderContainer(python_version, build_type, architecture).run()
     for p in platform:
         RayDockerContainer(python_version, p, image_type).run()
 
 
 def build_anyscale(
-    image_type: str, python_version: str, build_type: str, platform: List[str]
+    image_type: str,
+    python_version: str,
+    build_type: str,
+    platform: List[str],
+    architecture: str,
 ) -> None:
     """
     Build an anyscale container artifact.
     """
-    BuilderContainer(python_version, build_type).run()
+    BuilderContainer(python_version, build_type, architecture).run()
     for p in platform:
         RayDockerContainer(python_version, p, image_type).run()
         AnyscaleDockerContainer(python_version, p, image_type).run()

diff --git a/ci/ray_ci/builder_container.py b/ci/ray_ci/builder_container.py
@@ -13,6 +13,10 @@ class PythonVersionInfo(TypedDict):
     "optimized",
     "debug",
 ]
+ARCHITECTURE = [
+    "x86_64",
+    "aarch64",
+]
 PYTHON_VERSIONS = {
     "3.8": PythonVersionInfo(bin_path="cp38-cp38", numpy_version="1.19.3"),
     "3.9": PythonVersionInfo(bin_path="cp39-cp39", numpy_version="1.19.3"),
@@ -21,12 +25,13 @@ class PythonVersionInfo(TypedDict):
 }
 DEFAULT_PYTHON_VERSION = "3.8"
 DEFAULT_BUILD_TYPE = "optimized"
+DEFAULT_ARCHITECTURE = "x86_64"
 
 
 class BuilderContainer(Container):
-    def __init__(self, python_version: str, build_type: str) -> None:
+    def __init__(self, python_version: str, build_type: str, architecture: str) -> None:
         super().__init__(
-            "manylinux",
+            "manylinux" if architecture == "x86_64" else f"manylinux-{architecture}",
             volumes=[f"{os.environ.get('RAYCI_CHECKOUT_DIR')}:/rayci"],
         )
         python_version_info = PYTHON_VERSIONS.get(python_version)

diff --git a/ci/ray_ci/forge_container.py b/ci/ray_ci/forge_container.py
@@ -4,9 +4,9 @@
 
 
 class ForgeContainer(Container):
-    def __init__(self) -> None:
+    def __init__(self, architecture: str) -> None:
         super().__init__(
-            "forge",
+            "forge" if architecture == "x86_64" else "forge-aarch64",
             volumes=[f"{os.environ.get('RAYCI_CHECKOUT_DIR')}:/rayci"],
         )
 

diff --git a/ci/ray_ci/test_builder_container.py b/ci/ray_ci/test_builder_container.py
@@ -6,6 +6,13 @@
 from ci.ray_ci.builder_container import BuilderContainer
 
 
+def test_init() -> None:
+    builder = BuilderContainer("3.10", "optimized", "aarch64")
+    assert builder.docker_tag == "manylinux-aarch64"
+    builder = BuilderContainer("3.10", "optimized", "x86_64")
+    assert builder.docker_tag == "manylinux"
+
+
 def test_run() -> None:
     cmds = []
 
@@ -17,15 +24,15 @@ def _mock_run_script(input: List[str]) -> None:
         side_effect=_mock_run_script,
     ):
         # test optimized build
-        BuilderContainer("3.10", "optimized").run()
+        BuilderContainer("3.10", "optimized", "x86_64").run()
         assert cmds[-1] == [
             "./ci/build/build-manylinux-ray.sh",
             "./ci/build/build-manylinux-wheel.sh cp310-cp310 1.22.0",
             "chown -R 2000:100 /artifact-mount",
         ]
 
         # test debug build
-        BuilderContainer("3.9", "debug").run()
+        BuilderContainer("3.9", "debug", "x86_64").run()
         assert cmds[-1] == [
             "export RAY_DEBUG_BUILD=debug",
             "./ci/build/build-manylinux-ray.sh",

diff --git a/ci/ray_ci/tester.py b/ci/ray_ci/tester.py
@@ -10,6 +10,7 @@
     BuilderContainer,
     DEFAULT_BUILD_TYPE,
     DEFAULT_PYTHON_VERSION,
+    DEFAULT_ARCHITECTURE,
 )
 from ci.ray_ci.tester_container import TesterContainer
 from ci.ray_ci.utils import docker_login
@@ -156,7 +157,9 @@ def main(
 
     if build_type == "wheel":
         # for wheel testing, we first build the wheel and then use it for running tests
-        BuilderContainer(DEFAULT_PYTHON_VERSION, DEFAULT_BUILD_TYPE).run()
+        BuilderContainer(
+            DEFAULT_PYTHON_VERSION, DEFAULT_BUILD_TYPE, DEFAULT_ARCHITECTURE
+        ).run()
     container = _get_container(
         team,
         workers,