pytorch: fix dockerize and deduplicate build logic

Add missing USERNAME variable from dockerize.sh, and use dockerize.sh
in build.sh instead of duplictaing the build invocation. This needed the
addition of a `--build-only` flag for dockerize.sh, but it means the
existing behaviour doesn't need to change.

Author: jondea

ML-Frameworks/pytorch-aarch64/build.sh

@@ -50,9 +50,4 @@ torch_wheel_name=$(grep -o "torch-.*.whl" $build_log | head -n -1 | tail -n 1)
 # 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)
 
-docker build -t toolsolutions-pytorch:latest \
-    --build-arg TORCH_WHEEL=results/$torch_wheel_name \
-    --build-arg DOCKER_IMAGE_MIRROR \
-    --build-arg TORCH_AO_WHEEL=ao/dist/$torch_ao_wheel_name \
-    --build-arg USERNAME=ubuntu \
-    .
+./dockerize.sh "results/$torch_wheel_name" "ao/dist/$torch_ao_wheel_name" --build-only

ML-Frameworks/pytorch-aarch64/dockerize.sh

@@ -17,15 +17,18 @@
 # limitations under the License.
 # *******************************************************************************
 
-help_str="dockerize.sh takes a PyTorch wheel as the first argument and an ao wheel as second argument. It \
-installs the wheel inside a Docker container with examples and requirements."
-if [ "$#" -ne 2 ]; then
+help_str="dockerize.sh takes a PyTorch wheel as the first argument and an ao wheel
+as the second argument. It installs the wheel inside a Docker container with examples
+and requirements. The docker image will then be run unless you pass in the optional
+--build-only argument"
+
+if [ "$#" -lt 2 ]; then
     echo $help_str
     exit 1
 fi
 
 if ! [ -e "$1" ] || ! [ -e "$2" ]; then
-    echo "I couldn't find a wheel at $1 or $2"
+    echo "I couldn't find wheels at $1 and $2"
     echo $help_str
     exit 1
 fi
@@ -34,5 +37,8 @@ docker build -t toolsolutions-pytorch:latest  \
     --build-arg TORCH_WHEEL=$1 \
     --build-arg DOCKER_IMAGE_MIRROR \
     --build-arg TORCH_AO_WHEEL=$2 \
+    --build-arg USERNAME=ubuntu \
     .
+
+[[ $* == *--build-only* ]] && exit 0
 docker run --rm -it toolsolutions-pytorch:latest