From 3906f8e8e00b5923be05f5b9a07442baa0ada2fa Mon Sep 17 00:00:00 2001
From: Yuki Iwai <yuki.iwai.tz@gmail.com>
Date: Tue, 30 Aug 2022 03:16:53 +0900
Subject: [PATCH] Support for arm64 in mxnet-mnist image (#1940)

* support for arm64 in mxnet-mnist image

* fix build script
---
 .../trial-images/mxnet-mnist/Dockerfile       | 21 ++++++++---
 .../install-arm-performance-libraries.sh      | 35 +++++++++++++++++++
 scripts/v1beta1/build.sh                      |  6 ++--
 3 files changed, 54 insertions(+), 8 deletions(-)
 create mode 100755 examples/v1beta1/trial-images/mxnet-mnist/install-arm-performance-libraries.sh

diff --git a/examples/v1beta1/trial-images/mxnet-mnist/Dockerfile b/examples/v1beta1/trial-images/mxnet-mnist/Dockerfile
index 4c46894aa1a..d19c1aac175 100644
--- a/examples/v1beta1/trial-images/mxnet-mnist/Dockerfile
+++ b/examples/v1beta1/trial-images/mxnet-mnist/Dockerfile
@@ -1,15 +1,26 @@
 FROM python:3.9-slim
 
-ADD examples/v1beta1/trial-images/mxnet-mnist /opt/mxnet-mnist
+ARG TARGETARCH
+ENV LD_LIBRARY_PATH "${LD_LIBRARY_PATH}:/opt/arm/armpl_22.0.2_gcc-11.2/lib"
+
+COPY examples/v1beta1/trial-images/mxnet-mnist /opt/mxnet-mnist
 
 WORKDIR /opt/mxnet-mnist
 
 RUN apt-get -y update \
-  && apt-get -y install libgomp1 libquadmath0 \
-  && apt-get clean \
-  && rm -rf /var/lib/apt/lists/*
+ && apt-get -y install libgomp1 \
+ && if [ "${TARGETARCH}" = "arm64" ]; then \
+      apt-get -y install wget; \
+    elif [ "${TARGETARCH}" = "amd64" ]; then \
+      apt-get -y install libquadmath0; \
+    fi \
+ && apt-get clean \
+ && rm -rf /var/lib/apt/lists/*
 
-RUN pip install mxnet==1.9.0
+RUN if [ "${TARGETARCH}" = "arm64" ]; then \
+      /opt/mxnet-mnist/install-arm-performance-libraries.sh; \
+    fi
+RUN pip install mxnet==1.9.1
 RUN chgrp -R 0 /opt/mxnet-mnist \
   && chmod -R g+rwX /opt/mxnet-mnist
 
diff --git a/examples/v1beta1/trial-images/mxnet-mnist/install-arm-performance-libraries.sh b/examples/v1beta1/trial-images/mxnet-mnist/install-arm-performance-libraries.sh
new file mode 100755
index 00000000000..aa372cfb6af
--- /dev/null
+++ b/examples/v1beta1/trial-images/mxnet-mnist/install-arm-performance-libraries.sh
@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+
+# Copyright 2022 The Kubeflow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -o errexit
+set -o nounset
+set -o pipefail
+cd "$(dirname "$0")"
+
+# Download Arm Performance Libraries for Ubuntu 20.04
+# Ref: https://developer.arm.com/downloads/-/arm-performance-libraries
+echo "Downloading Arm Performance Libraries for Ubuntu 20.04..."
+wget -qO - \
+  "https://developer.arm.com/-/media/Files/downloads/hpc/arm-performance-libraries/22-0-2/Ubuntu20.04/arm-performance-libraries_22.0.2_Ubuntu-20.04_gcc-11.2.tar?rev=577d3dbcff7847b9af57399b2978f9a6&revision=577d3dbc-ff78-47b9-af57-399b2978f9a6" \
+  | tar -xf -
+
+# Install Arm Performance Libraries
+echo "Installing Arm Performance Libraries for Ubuntu 20.04..."
+./arm-performance-libraries_22.0.2_Ubuntu-20.04/arm-performance-libraries_22.0.2_Ubuntu-20.04.sh -a
+
+# Clean up
+echo "Removing installer..."
+rm -rf ./arm-performance-libraries_22.0.2_Ubuntu-20.04
diff --git a/scripts/v1beta1/build.sh b/scripts/v1beta1/build.sh
index 93f74430748..0acc611e04a 100755
--- a/scripts/v1beta1/build.sh
+++ b/scripts/v1beta1/build.sh
@@ -120,9 +120,6 @@ echo -e "\nBuilding training container images..."
 if [ ! "$ARCH" = "amd64" ]; then
   echo -e "\nSome training container images are supported only amd64."
 else
-  echo -e "\nBuilding mxnet mnist training container example...\n"
-  docker buildx build --platform linux/amd64 -t "${REGISTRY}/mxnet-mnist:${TAG}" -f examples/${VERSION}/trial-images/mxnet-mnist/Dockerfile .
-
   echo -e "\nBuilding PyTorch mnist training container example with CPU support...\n"
   docker buildx build --platform linux/amd64 -t "${REGISTRY}/pytorch-mnist-cpu:${TAG}" -f examples/${VERSION}/trial-images/pytorch-mnist/Dockerfile.cpu .
 
@@ -142,6 +139,9 @@ else
   docker buildx build --platform linux/amd64 -t "${REGISTRY}/simple-pbt:${TAG}" -f examples/${VERSION}/trial-images/simple-pbt/Dockerfile .
 fi
 
+echo -e "\nBuilding mxnet mnist training container example...\n"
+docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/mxnet-mnist:${TAG}" -f examples/${VERSION}/trial-images/mxnet-mnist/Dockerfile .
+
 echo -e "\nBuilding Tensorflow with summaries mnist training container example...\n"
 docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/tf-mnist-with-summaries:${TAG}" -f examples/${VERSION}/trial-images/tf-mnist-with-summaries/Dockerfile .