From b17735630883204c05ee9c347db1d273b286ad00 Mon Sep 17 00:00:00 2001
From: anakinxc <zhengyancheng.zyc@antgroup.com>
Date: Mon, 11 Dec 2023 17:51:30 +0800
Subject: [PATCH] Tweak script to handle aarch64 only dockerfiles

---
 dockerfiles/release-ci-aarch64.DockerFile |  1 +
 dockerfiles/release-ci.DockerFile         |  2 +-
 publish_docker_image.py                   | 21 ++++++++++++++++-----
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/dockerfiles/release-ci-aarch64.DockerFile b/dockerfiles/release-ci-aarch64.DockerFile
index 3ce5c6b..9d40ac0 100644
--- a/dockerfiles/release-ci-aarch64.DockerFile
+++ b/dockerfiles/release-ci-aarch64.DockerFile
@@ -1,3 +1,4 @@
+# Conda does not support CentOS 7 on Linux aarch64, so use CentOS 8
 FROM centos:centos8
 
 RUN cd /etc/yum.repos.d/ \
diff --git a/dockerfiles/release-ci.DockerFile b/dockerfiles/release-ci.DockerFile
index 11716ff..2977b04 100644
--- a/dockerfiles/release-ci.DockerFile
+++ b/dockerfiles/release-ci.DockerFile
@@ -1,4 +1,4 @@
-FROM centos:centos7
+FROM centos:centos7 
 
 RUN yum install -y centos-release-scl epel-release \
     && yum update -y \
diff --git a/publish_docker_image.py b/publish_docker_image.py
index 0574d7b..ee2b5de 100755
--- a/publish_docker_image.py
+++ b/publish_docker_image.py
@@ -38,6 +38,11 @@ def _check_is_multiarch_dockerfile(file):
         else:
             return False
 
+def _check_is_aarch64_dockerfile(file):
+    if 'aarch64' in file:
+        return True
+    return False
+
 
 COLOR_GREEN = "\033[92m"
 COLOR_END = "\033[0m"
@@ -72,9 +77,15 @@ def main():
 
     print(f"{COLOR_GREEN}Build docker image {name}{COLOR_END}")
 
-    if _check_is_multiarch_dockerfile(dockerfile):
+    is_multi_arch = _check_is_multiarch_dockerfile(dockerfile)
+    is_aarch64 = _check_is_aarch64_dockerfile(dockerfile)
+
+    if is_multi_arch or is_aarch64:
+        if is_multi_arch:
+            arch_str = "linux/amd64,linux/arm64"
+        elif is_aarch64:
+            arch_str = "linux/arm64"
         print(f"{COLOR_GREEN}Creating buildx")
-        is_multi_arch = True
         _run_shell_command_with_live_output(
             [
                 "docker",
@@ -83,7 +94,7 @@ def main():
                 "--name",
                 "sf-image-builder",
                 "--platform",
-                "linux/amd64,linux/arm64",
+                arch_str,
                 "--use",
             ],
             ".",
@@ -96,7 +107,7 @@ def main():
                 "buildx",
                 "build",
                 "--platform",
-                "linux/amd64,linux/arm64",
+                arch_str,
                 "--no-cache",
                 ".",
                 "-f",
@@ -114,7 +125,7 @@ def main():
                 "buildx",
                 "build",
                 "--platform",
-                "linux/amd64,linux/arm64",
+                arch_str,
                 ".",
                 "-f",
                 dockerfile,