Fix Macos arm64 build.

PiperOrigin-RevId: 708484831
google · Dec 21, 2024 · b54c76c · b54c76c
1 parent 77c56dc
commit b54c76c
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 15 deletions.
diff --git a/grain/oss/Dockerfile b/grain/oss/Dockerfile
@@ -13,8 +13,10 @@ ENV DEBIAN_FRONTEND=noninteractive
 
 RUN ulimit -n 1024 && yum install -y rsync
 
-ENV PYTHON_BIN=/opt/python/cp${PYTHON_VERSION}-cp${PYTHON_VERSION}/bin
-ENV PATH="${PYTHON_BIN}:${PATH}"
+ENV PYTHON_BIN_PATH=/opt/python/cp${PYTHON_VERSION}-cp${PYTHON_VERSION}/bin
+ENV PATH="${PYTHON_BIN_PATH}:${PATH}"
+
+ENV PYTHON_BIN=${PYTHON_BIN_PATH}/python
 
 # Download the correct bazel version and make sure it's on path.
 RUN BAZEL_ARCH_SUFFIX="$(uname -m | sed s/aarch64/arm64/)" \
@@ -28,12 +30,12 @@ RUN BAZEL_ARCH_SUFFIX="$(uname -m | sed s/aarch64/arm64/)" \
 # exists.
 COPY /array_recor[d] /tmp/grain/array_record
 RUN --mount=type=cache,target=/root/.cache \
-  ${PYTHON_BIN}/python -m pip install -U \
+  $PYTHON_BIN -m pip install -U \
     --find-links=/tmp/grain/array_record array_record;
 
 # Install dependencies needed for grain.
 RUN --mount=type=cache,target=/root/.cache \
-  ${PYTHON_BIN}/python -m pip install -U \
+  $PYTHON_BIN -m pip install -U \
     absl-py \
     build \
     cloudpickle \
@@ -45,7 +47,7 @@ RUN --mount=type=cache,target=/root/.cache \
 
 # Install dependencies needed for grain tests
 RUN --mount=type=cache,target=/root/.cache \
-  ${PYTHON_BIN}/python -m pip install -U \
+  $PYTHON_BIN -m pip install -U \
     attrs \
     auditwheel \
     dill \

diff --git a/grain/oss/build_whl.sh b/grain/oss/build_whl.sh
@@ -3,13 +3,10 @@
 
 set -e -x
 
-CP_VERSION="cp${PYTHON_MAJOR_VERSION}${PYTHON_MINOR_VERSION}"
-PYTHON_BIN_PATH="/opt/python/${CP_VERSION}-${CP_VERSION}/bin/python"
-
 function main() {
   bazel clean
-  bazel build ... --action_env PYTHON_BIN_PATH="${PYTHON_BIN_PATH}" --action_env MACOSX_DEPLOYMENT_TARGET=11.0
-  bazel test --verbose_failures --test_output=errors ... --action_env PYTHON_BIN_PATH="${PYTHON_BIN_PATH}"
+  bazel build ... --action_env PYTHON_BIN_PATH="${PYTHON_BIN}" --action_env MACOSX_DEPLOYMENT_TARGET=11.0
+  bazel test --verbose_failures --test_output=errors ... --action_env PYTHON_BIN_PATH="${PYTHON_BIN}"
 
   DEST="/tmp/grain/all_dist"
   mkdir -p "${DEST}"
@@ -38,11 +35,13 @@ function main() {
     plat_name="--plat-name macosx_11_0_$(uname -m)"
   fi
 
-  "python${PYTHON_VERSION}" setup.py bdist_wheel --python-tag py3${PYTHON_MINOR_VERSION} $plat_name
+  $PYTHON_BIN setup.py bdist_wheel --python-tag py3${PYTHON_MINOR_VERSION} $plat_name
   cp dist/*.whl "${DEST}"
 
-  echo $(date) : "=== Auditing wheel"
-  auditwheel repair --plat ${AUDITWHEEL_PLATFORM} -w dist dist/*.whl
+  if [ -n "${AUDITWHEEL_PLATFORM}" ]; then
+    echo $(date) : "=== Auditing wheel"
+    auditwheel repair --plat ${AUDITWHEEL_PLATFORM} -w dist dist/*.whl
+  fi
 
   echo $(date) : "=== Listing wheel"
   ls -lrt dist/*.whl

diff --git a/grain/oss/runner_common.sh b/grain/oss/runner_common.sh
@@ -80,7 +80,7 @@ function update_bazel_macos {
 }
 
 function install_grain_deps {
-  AR_DIR="./grain/oss/array_record"
+  AR_DIR="$SOURCE_DIR/grain/oss/array_record"
   $PYTHON_BIN -m pip install -U --find-links=$AR_DIR array_record --no-cache-dir;
   $PYTHON_BIN -m pip install -U \
     absl-py \
@@ -121,6 +121,6 @@ function build_and_test_grain_macos() {
     install_and_init_pyenv
     install_grain_deps
 
-    bash grain/oss/build_whl.sh
+    bash "${SOURCE_DIR}/grain/oss/build_whl.sh"
   done
 }