Merge pull request #425 from puneetmatharu/tidy-up-pytorch-build-scripts

puneetmatharu · web-flow · commit c065c1833863 · 2026-02-04T15:41:14.000Z
Tidy up pytorch-aarch64 build scripts
diff --git a/.github/workflows/pytorch.yml b/.github/workflows/pytorch.yml
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright 2025 Arm Limited and affiliates.
+# SPDX-FileCopyrightText: Copyright 2025, 2026 Arm Limited and affiliates.
 #
 # SPDX-License-Identifier: Apache-2.0
 
@@ -55,7 +55,7 @@ jobs:
            ]
     runs-on: ${{ matrix.config.label }}
     env:
-      CCACHE_HOST_DIR: ${{ github.workspace }}/Tool-Solutions/ML-Frameworks/pytorch-aarch64/.ccache
+      CCACHE_LOCAL_DIR: ${{ github.workspace }}/Tool-Solutions/ML-Frameworks/pytorch-aarch64/.ccache
       TORCH_BUILD_CONTAINER_ID_FILE: ${{ github.workspace }}/Tool-Solutions/ML-Frameworks/pytorch-aarch64/.torch_build_container_id
     steps:
       - name: Checkout Tool-Solutions
@@ -71,14 +71,14 @@ jobs:
         run: echo "week=$(date -u +%G-%V)" >> "$GITHUB_OUTPUT"
 
       # Restore cache if available. GitHub automatically evicts cache entries that have not been
-      # accessed for over 7 days. We rotate the cache key weekly; if no cache exists for the 
-      # current week, a cache from a previous week (via the prefix restore key) will be restored 
-      # and then saved under the current week's key at the end of the job. This effectively limits 
+      # accessed for over 7 days. We rotate the cache key weekly; if no cache exists for the
+      # current week, a cache from a previous week (via the prefix restore key) will be restored
+      # and then saved under the current week's key at the end of the job. This effectively limits
       # the cache to at most two weeks of cache data.
       - name: Restore ccache cache
         uses: actions/cache@v5
         with:
-          path: ${{ env.CCACHE_HOST_DIR }}
+          path: ${{ env.CCACHE_LOCAL_DIR }}
           key: ccache-${{ matrix.config.name }}-${{ steps.cache_suffix.outputs.week }}
           restore-keys: |
             ccache-${{ matrix.config.name }}-
@@ -88,11 +88,11 @@ jobs:
         run: ${{ github.workspace }}/Tool-Solutions/ML-Frameworks/pytorch-aarch64/build.sh
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          CCACHE_HOST_DIR: ${{ env.CCACHE_HOST_DIR }}
+          CCACHE_LOCAL_DIR: ${{ env.CCACHE_LOCAL_DIR }}
           CCACHE_MAXSIZE: 2G
 
       - name: Print ccache disk usage
-        run: du -sh "${{ env.CCACHE_HOST_DIR }}" || true
+        run: du -sh "${{ env.CCACHE_LOCAL_DIR }}" || true
 
       - name: Report final ccache build stats
         run: docker exec "$(cat ${{ env.TORCH_BUILD_CONTAINER_ID_FILE }})" ccache -s || true
diff --git a/ML-Frameworks/pytorch-aarch64/build-torch-ao-wheel.sh b/ML-Frameworks/pytorch-aarch64/build-torch-ao-wheel.sh
@@ -9,27 +9,35 @@
 
 set -eux -o pipefail
 
+docker_exec() {
+    docker exec "$TORCH_BUILD_CONTAINER" "$@"
+}
+
 PYTHON_VERSION="3.12"
 
 # Specify DOCKER_IMAGE_MIRROR if you want to use a mirror of hub.docker.com
 IMAGE_NAME="${DOCKER_IMAGE_MIRROR:-}pytorch/manylinux2_28_aarch64-builder:cpu-aarch64-69d4c1f80b5e7da224d4f9c2170ef100e75dfe03"
 TORCH_BUILD_CONTAINER_ID_FILE="${PWD}/.torch_ao_build_container_id"
 
-TEST_VENV=aarch64_env_test_torch_ao
+# Name of the virtual environment inside the container
+TEST_VENV_CONTAINER_DIR=/aarch64_env_test_torch_ao
 
-TORCH_AO_HOST_DIR="${PWD}/ao"
+# Local and container dirs for the torchao source code
+TORCH_AO_LOCAL_DIR="${PWD}/ao"
+TORCH_AO_CONTAINER_DIR=/ao
 
-TORCH_AO_ROOT=/ao
+# Output dir for PyTorch wheel and other artifacts
+OUTPUT_LOCAL_DIR="${OUTPUT_DIR:-"${PWD}/results"}"
 
 if [ -f "$TORCH_BUILD_CONTAINER_ID_FILE" ]; then
-    TORCH_BUILD_CONTAINER=$(cat $TORCH_BUILD_CONTAINER_ID_FILE)
+    TORCH_BUILD_CONTAINER=$(cat "$TORCH_BUILD_CONTAINER_ID_FILE")
     echo "Found an existing torch build container id: $TORCH_BUILD_CONTAINER"
 else
     TORCH_BUILD_CONTAINER=""
     echo "Did not find torch build container id in $(readlink -f $TORCH_BUILD_CONTAINER_ID_FILE), we will create one later"
 fi
 
-if ! docker container inspect $TORCH_BUILD_CONTAINER >/dev/null 2>&1 ; then
+if ! docker container inspect "$TORCH_BUILD_CONTAINER" >/dev/null 2>&1 ; then
 
     TORCH_BUILD_CONTAINER=$(docker run -t -d \
         -e BINARY_ENV_FILE=/tmp/env \
@@ -39,36 +47,42 @@ if ! docker container inspect $TORCH_BUILD_CONTAINER >/dev/null 2>&1 ; then
         -e GITHUB_ACTIONS=0 \
         -e GPU_ARCH_TYPE=cpu-aarch64 \
         -e PACKAGE_TYPE=manywheel \
-        -e TORCH_AO_ROOT="${TORCH_AO_ROOT}" \
+        -e TORCH_AO_ROOT="${TORCH_AO_CONTAINER_DIR}" \
         -e SKIP_ALL_TESTS=1 \
-        -e TEST_VENV="${TEST_VENV}" \
-        -v "${TORCH_AO_HOST_DIR}:${TORCH_AO_ROOT}" \
+        -e TEST_VENV="${TEST_VENV_CONTAINER_DIR}" \
+        -v "${TORCH_AO_LOCAL_DIR}:${TORCH_AO_CONTAINER_DIR}" \
         -w / \
         "${IMAGE_NAME}")
 
     # The Docker image comes with a copy of ACL, but we want to build against whichever
     # version is compatible with the pip-installed torch (which ships its own copy of ACL),
     # so we remove the existing copy here
-    docker exec -t $TORCH_BUILD_CONTAINER bash -c "cd / && rm -rf acl/"
+    docker_exec rm -rf /acl
 
-    docker exec -t $TORCH_BUILD_CONTAINER bash -c "python${PYTHON_VERSION} -m venv $TEST_VENV"
-    docker exec -t $TORCH_BUILD_CONTAINER bash -c "source $TEST_VENV/bin/activate && pip install typing_extensions torch wheel numpy --no-deps"
+    docker_exec python${PYTHON_VERSION} -m venv ${TEST_VENV_CONTAINER_DIR}
+    docker_exec bash -c "source ${TEST_VENV_CONTAINER_DIR}/bin/activate && pip install typing_extensions torch wheel numpy --no-deps"
 
     echo "Storing torch build container id in $TORCH_BUILD_CONTAINER_ID_FILE for reuse: $TORCH_BUILD_CONTAINER"
-    echo $TORCH_BUILD_CONTAINER > "$TORCH_BUILD_CONTAINER_ID_FILE"
+    echo "$TORCH_BUILD_CONTAINER" > "$TORCH_BUILD_CONTAINER_ID_FILE"
 else
-    docker restart $TORCH_BUILD_CONTAINER
+    docker restart "$TORCH_BUILD_CONTAINER"
 fi
 
 # Clean up any old builds
-docker exec -t $TORCH_BUILD_CONTAINER bash -c "rm -rf ${TORCH_AO_ROOT}/build ${TORCH_AO_ROOT}/dist"
+docker_exec rm -rf ${TORCH_AO_CONTAINER_DIR}/build ${TORCH_AO_CONTAINER_DIR}/dist
 
 # Build
-docker exec -t $TORCH_BUILD_CONTAINER bash -c \
-  "source ${TEST_VENV}/bin/activate && \
-   cd ${TORCH_AO_ROOT} && \
-   python${PYTHON_VERSION} -m pip install setuptools build && \
-   python${PYTHON_VERSION} -m build --wheel --no-isolation"
+docker_exec bash -lc "
+  source ${TEST_VENV_CONTAINER_DIR}/bin/activate && \
+  cd ${TORCH_AO_CONTAINER_DIR} && \
+  python${PYTHON_VERSION} -m pip install setuptools build && \
+  python${PYTHON_VERSION} -m build --wheel --no-isolation
+"
 
 # directories generated by the docker container are owned by root, so transfer ownership to user
-docker exec $TORCH_BUILD_CONTAINER chown -R "$(id -u)":"$(id -g)" "${TORCH_AO_ROOT}"
+docker_exec chown -R "$(id -u)":"$(id -g)" "${TORCH_AO_CONTAINER_DIR}"
+
+# Place the torchao wheel next to the torch wheel
+TORCH_AO_WHEEL_PATH=$(find "${TORCH_AO_LOCAL_DIR}/dist" -name '*.whl' -type f)
+mkdir -p "${OUTPUT_LOCAL_DIR}"
+cp "${TORCH_AO_WHEEL_PATH}" "${OUTPUT_LOCAL_DIR}/"
diff --git a/ML-Frameworks/pytorch-aarch64/build-wheel.sh b/ML-Frameworks/pytorch-aarch64/build-wheel.sh
@@ -17,39 +17,42 @@
 
 set -eux -o pipefail
 
+docker_exec() {
+    docker exec "$TORCH_BUILD_CONTAINER" "$@"
+}
+
 PYTHON_VERSION="3.12"
 OPENBLAS_VERSION="v0.3.30"
 ACL_VERSION="v52.8.0"
 
-PYTHON_TAG="cp$(echo "$PYTHON_VERSION" | tr -d .)-cp$(echo "$PYTHON_VERSION" | tr -d .)"
-PYTHON_BIN="/opt/python/${PYTHON_TAG}/bin"
-
 # Specify DOCKER_IMAGE_MIRROR if you want to use a mirror of hub.docker.com
 IMAGE_NAME="${DOCKER_IMAGE_MIRROR:-}pytorch/manylinux2_28_aarch64-builder:cpu-aarch64-69d4c1f80b5e7da224d4f9c2170ef100e75dfe03"
 TORCH_BUILD_CONTAINER_ID_FILE="${PWD}/.torch_build_container_id"
 
-# Output dir for PyTorch wheel and other artifacts
-OUTPUT_DIR=${OUTPUT_DIR:-"${PWD}/results"}
+PYTHON_TAG="cp$(echo "$PYTHON_VERSION" | tr -d .)-cp$(echo "$PYTHON_VERSION" | tr -d .)"
+PYTHON_CONTAINER_BIN="/opt/python/${PYTHON_TAG}/bin"
+
+# Output dir for PyTorch wheel and other artifacts. Rename with the "_LOCAL_DIR"
+# suffix for consistency with other variables
+OUTPUT_LOCAL_DIR="${OUTPUT_DIR:-"${PWD}/results"}"
 
 # Where folders sit locally
-PYTORCH_HOST_DIR="${PWD}/pytorch"
-OPENSSL_HOST_DIR="/opt/openssl"
-PYTORCH_FINAL_PACKAGE_HOST_DIR="${OUTPUT_DIR}"
+PYTORCH_LOCAL_DIR="${PWD}/pytorch"
+PYTORCH_FINAL_PACKAGE_LOCAL_DIR="${OUTPUT_LOCAL_DIR}"
 
-# Where folders sit mounted in the container
-PYTORCH_ROOT=/pytorch
-UTILS=/utils
-COMMON_UTILS=/common_utils
-PYTORCH_FINAL_PACKAGE_DIR=/artifacts
+# Where folders sit in the container
+PYTORCH_CONTAINER_DIR=/pytorch
+PYTORCH_FINAL_PACKAGE_CONTAINER_DIR=/artifacts
+OPENSSL_CONTAINER_DIR=/opt/openssl
 
 # Enable ccache support by default.
 # NOTE: The default behaviour is to have a project-specific cache directory that we cache
 # build artefacts inside and can be easily wiped. These build artefacts are specific to the
 # manylinux builder container (and thus compilers) that we use to build the torch wheel. As
 # such, you may not want to populate the global ccache cache with them. However, if you wish
-# to do so, simply set CCACHE_HOST_DIR to that directory.
-CCACHE_HOST_DIR=${CCACHE_HOST_DIR:-"${PWD}/.ccache"}
-LOCAL_CCACHE_DIR=/.ccache
+# to do so, simply set CCACHE_LOCAL_DIR to that directory.
+CCACHE_LOCAL_DIR="${CCACHE_LOCAL_DIR:-"${PWD}/.ccache"}"
+CCACHE_CONTAINER_DIR=/.ccache
 CCACHE_MAXSIZE=${CCACHE_MAXSIZE:-}
 
 # If the user wants to use ccache for build caching
@@ -58,10 +61,10 @@ if [[ "$*" == *--disable-ccache* ]]; then
     ccache_args+=(-e USE_CCACHE=0)
 else
     ccache_args+=(-e USE_CCACHE=1)
-    mkdir -p "${CCACHE_HOST_DIR}"
+    mkdir -p "${CCACHE_LOCAL_DIR}"
     ccache_args+=(
-        -e CCACHE_DIR="${LOCAL_CCACHE_DIR}"
-        -v "${CCACHE_HOST_DIR}:${LOCAL_CCACHE_DIR}"
+        -e CCACHE_DIR="${CCACHE_CONTAINER_DIR}"
+        -v "${CCACHE_LOCAL_DIR}:${CCACHE_CONTAINER_DIR}"
     )
 fi
 
@@ -73,14 +76,14 @@ GPU_ARCH_TYPE=cpu-aarch64
 MAX_JOBS=${MAX_JOBS:-$(nproc --ignore=2)}
 
 if [ -f "$TORCH_BUILD_CONTAINER_ID_FILE" ]; then
-    TORCH_BUILD_CONTAINER=$(cat $TORCH_BUILD_CONTAINER_ID_FILE)
+    TORCH_BUILD_CONTAINER=$(cat "$TORCH_BUILD_CONTAINER_ID_FILE")
     echo "Found an existing torch build container id: $TORCH_BUILD_CONTAINER"
 else
     TORCH_BUILD_CONTAINER=""
     echo "Did not find torch build container id in $(readlink -f $TORCH_BUILD_CONTAINER_ID_FILE), we will create one later"
 fi
 
-if ! docker container inspect $TORCH_BUILD_CONTAINER >/dev/null 2>&1 ; then
+if ! docker container inspect "$TORCH_BUILD_CONTAINER" >/dev/null 2>&1 ; then
     # Based on environment used in pytorch/.github/workflows/_binary-build-linux.yml
     # and pytorch/.github/workflows/generated-linux-aarch64-binary-manywheel-nightly.yml
     TORCH_BUILD_CONTAINER=$(docker run -t -d \
@@ -94,79 +97,80 @@ if ! docker container inspect $TORCH_BUILD_CONTAINER >/dev/null 2>&1 ; then
         -e GITHUB_ACTIONS=0 \
         -e GPU_ARCH_TYPE=${GPU_ARCH_TYPE} \
         -e PACKAGE_TYPE=manywheel \
-        -e PYTORCH_FINAL_PACKAGE_DIR="${PYTORCH_FINAL_PACKAGE_DIR}" \
-        -e PYTORCH_ROOT="${PYTORCH_ROOT}" \
+        -e PYTORCH_FINAL_PACKAGE_DIR="${PYTORCH_FINAL_PACKAGE_CONTAINER_DIR}" \
+        -e PYTORCH_ROOT="${PYTORCH_CONTAINER_DIR}" \
         -e SKIP_ALL_TESTS=1 \
-        -e OPENSSL_ROOT_DIR="${OPENSSL_HOST_DIR}" \
-        -e CMAKE_INCLUDE_PATH="${OPENSSL_HOST_DIR}/include" \
+        -e OPENSSL_ROOT_DIR="${OPENSSL_CONTAINER_DIR}" \
+        -e CMAKE_INCLUDE_PATH="${OPENSSL_CONTAINER_DIR}/include" \
         "${ccache_args[@]}" \
-        -v "${PYTORCH_HOST_DIR}:${PYTORCH_ROOT}" \
-        -v "${PYTORCH_FINAL_PACKAGE_HOST_DIR}:${PYTORCH_FINAL_PACKAGE_DIR}" \
-        -v "${PWD}/utils:${UTILS}" \
-        -v "${PWD}/../utils:${COMMON_UTILS}" \
+        -v "${PYTORCH_LOCAL_DIR}:${PYTORCH_CONTAINER_DIR}" \
+        -v "${PYTORCH_FINAL_PACKAGE_LOCAL_DIR}:${PYTORCH_FINAL_PACKAGE_CONTAINER_DIR}" \
         -w / \
         "${IMAGE_NAME}")
 
     # Provide ccache support
     if [[ "$*" != *--disable-ccache* ]]; then
-        docker exec "$TORCH_BUILD_CONTAINER" yum install -y ccache || true
+        docker_exec yum install -y ccache || true
         if [ -n "${CCACHE_MAXSIZE}" ]; then
-            docker exec "$TORCH_BUILD_CONTAINER" ccache --max-size="$CCACHE_MAXSIZE" || true
+            docker_exec ccache --max-size="$CCACHE_MAXSIZE" || true
         fi
-        docker exec "$TORCH_BUILD_CONTAINER" ccache -z || true
-        docker exec "$TORCH_BUILD_CONTAINER" ccache -o compression=true || true
-        docker exec "$TORCH_BUILD_CONTAINER" ccache -o compression_level=6 || true
-        docker exec "$TORCH_BUILD_CONTAINER" ccache -s || true
+        docker_exec ccache -z || true
+        docker_exec ccache -o compression=true || true
+        docker_exec ccache -o compression_level=6 || true
+        docker_exec ccache -s || true
     fi
 
     # Currently changes in these scripts will not be applied without a clean
     # build, which is not ideal for dev work. But we have to balance this with
     # extra time/network traffic when rebuilding many times.
-    docker exec $TORCH_BUILD_CONTAINER bash "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+    docker_exec bash "${PYTORCH_CONTAINER_DIR}/.circleci/scripts/binary_populate_env.sh"
 
-    # Install scons for ACL build
-    docker exec $TORCH_BUILD_CONTAINER ${PYTHON_BIN}/python3 -m pip install scons==4.7.0
-    docker exec $TORCH_BUILD_CONTAINER ln -sf ${PYTHON_BIN}/scons /usr/local/bin
+    # Install scons for the Compute Library (ACL) build
+    docker_exec ${PYTHON_CONTAINER_BIN}/python3 -m pip install scons==4.7.0
+    docker_exec ln -sf "${PYTHON_CONTAINER_BIN}/scons" /usr/local/bin
 
-    # The Docker image comes with a pre-built version of ACL, but we want to build our own
-    # so we remove the provided version here
-    docker exec "$TORCH_BUILD_CONTAINER" rm -rf /acl
+    # The Docker image comes with a pre-built version of ACL, but we
+    # want to build our own so we remove the provided version here
+    docker_exec rm -rf /acl
 
     # Affected by ACL_VERSION set as an environment variable above
     echo "Overriding Arm Compute Library version: ${ACL_VERSION}"
-    docker exec "$TORCH_BUILD_CONTAINER" "${PYTORCH_ROOT}/.ci/docker/common/install_acl.sh"
+    docker_exec "${PYTORCH_CONTAINER_DIR}/.ci/docker/common/install_acl.sh"
 
     # Affected by OPENBLAS_VERSION set as an environment variable above
     echo "Installing OpenBLAS version: ${OPENBLAS_VERSION}"
-    docker exec "$TORCH_BUILD_CONTAINER" "${PYTORCH_ROOT}/.ci/docker/common/install_openblas.sh"
+    docker_exec "${PYTORCH_CONTAINER_DIR}/.ci/docker/common/install_openblas.sh"
 
     echo "Storing torch build container ID in ${TORCH_BUILD_CONTAINER_ID_FILE} for reuse: ${TORCH_BUILD_CONTAINER}"
-    echo $TORCH_BUILD_CONTAINER > "${TORCH_BUILD_CONTAINER_ID_FILE}"
+    echo "$TORCH_BUILD_CONTAINER" > "${TORCH_BUILD_CONTAINER_ID_FILE}"
 else
-    docker restart $TORCH_BUILD_CONTAINER
+    docker restart "$TORCH_BUILD_CONTAINER"
 fi
 
 # If there are multiple wheels in the dist directory, an old wheel can be
 # erroneously copied to results, so we clear the directory to be sure
-docker exec $TORCH_BUILD_CONTAINER rm -rf "${PYTORCH_ROOT}/dist"
+docker_exec rm -rf "${PYTORCH_CONTAINER_DIR}/dist"
 
 # We set OVERRIDE_PACKAGE_VERSION to be based on the date of the latest torch
 # commit, this allows us to also install the matching torch* packages, set in
 # the Dockerfile. This is what PyTorch does in its nightly pipeline, see
 # pytorch/.ci/aarch64_linux/aarch64_wheel_ci_build.py for this logic.
-build_date=$(cd $PYTORCH_HOST_DIR && git log --pretty=format:%cs -1 | tr -d '-')
-version=$(cat $PYTORCH_HOST_DIR/version.txt| tr -d "[:space:]" )
+build_date=$(cd "$PYTORCH_LOCAL_DIR" && git log --pretty=format:%cs -1 | tr -d '-')
+version=$(cat "$PYTORCH_LOCAL_DIR/version.txt" | tr -d "[:space:]")
 OVERRIDE_PACKAGE_VERSION="${version%??}.dev${build_date}${TORCH_RELEASE_ID:+"+$TORCH_RELEASE_ID"}"
 
-docker exec $TORCH_BUILD_CONTAINER bash -lc "
+# Build the wheel!
+docker_exec bash -lc "
   source /tmp/env &&
   BUILD_TEST=0 \
   DO_SETUP_PY_CLEAN_BEFORE_BUILD=0 \
   WIPE_RH_CUDA_AFTER_BUILD=0 \
   OVERRIDE_PACKAGE_VERSION=$OVERRIDE_PACKAGE_VERSION \
-  bash ${PYTORCH_ROOT}/.ci/manywheel/build.sh
+  bash ${PYTORCH_CONTAINER_DIR}/.ci/manywheel/build.sh
 "
 
-# directories generated by the docker container are owned by root, so transfer ownership to user
-docker exec "$TORCH_BUILD_CONTAINER" chown -R "$(id -u)":"$(id -g)" \
-    "${PYTORCH_ROOT}" "${PYTORCH_FINAL_PACKAGE_DIR}" "${LOCAL_CCACHE_DIR}"
+# Directories generated by the docker container are owned by root, so transfer ownership to user
+docker_exec chown -R "$(id -u)":"$(id -g)" \
+    "${PYTORCH_CONTAINER_DIR}" \
+    "${PYTORCH_FINAL_PACKAGE_CONTAINER_DIR}" \
+    "${CCACHE_CONTAINER_DIR}"
diff --git a/ML-Frameworks/pytorch-aarch64/build.sh b/ML-Frameworks/pytorch-aarch64/build.sh
@@ -1,14 +1,14 @@
 #!/bin/bash
 
-# SPDX-FileCopyrightText: Copyright 2020-2025 Arm Limited and affiliates.
+# SPDX-FileCopyrightText: Copyright 2020-2026 Arm Limited and affiliates.
 #
 # SPDX-License-Identifier: Apache-2.0
 
 set -eux -o pipefail
 
 # exec redirects all output from now on into a file and stdout
 build_log=build-$(git rev-parse --short=7 HEAD)-$(date '+%Y-%m-%dT%H-%M-%S').log
-exec &> >(tee -a $build_log)
+exec &> >(tee -a "$build_log")
 
 # Bail out if sources are already there
 if [ -f .torch_build_container_id ] || [ -f .torch_ao_build_container_id ] || \
@@ -70,6 +70,10 @@ if ! [[ $* == *--use-existing-sources* ]]; then
     ./get-source.sh
 fi
 
+# Set the output dir for the wheels
+OUTPUT_DIR=${OUTPUT_DIR:-"${PWD}/results"}
+export OUTPUT_DIR="${OUTPUT_DIR}"
+
 # We build the wheel with ccache by default; allow disabling it via the --disable-ccache flag
 build_wheel_args=()
 if [[ "$*" == *--disable-ccache* ]]; then
@@ -80,14 +84,11 @@ fi
 [[ $* == *--wheel-only* ]] && exit 0
 
 # Use the second to last match, otherwise grep finds itself
-torch_wheel_name=$(grep -o "torch-.*.whl" $build_log | head -n -1 | tail -n 1)
+torch_wheel_name=$(grep -o "torch-.*.whl" "$build_log" | head -n -1 | tail -n 1)
 
 ./build-torch-ao-wheel.sh
 
 # Use the second to last match, otherwise grep finds itself
-torch_ao_wheel_name=$(grep -o "torchao-.*.whl" $build_log | head -n -1 | tail -n 1)
-
-# Place the torchao wheel next to the torch wheel
-cp "ao/dist/$torch_ao_wheel_name" "results/$torch_ao_wheel_name"
+torch_ao_wheel_name=$(grep -o "torchao-.*.whl" "$build_log" | head -n -1 | tail -n 1)
 
 ./dockerize.sh "results/$torch_wheel_name" "results/$torch_ao_wheel_name" --build-only
