[pytorch][ci] add custom selective build flow for android build (pytorch#40199)

Summary:
Pull Request resolved: pytorch#40199

Mobile custom selective build has already been covered by `test/mobile/custom_build/build.sh`.
It builds a CLI binary with host-toolchain and runs on host machine to
check correctness of the result.

But that custom build test doesn't cover the android/gradle build part.
And we cannot use it to measure and track the in-APK size of custom
build library.

So this PR adds the selective build test coverage for android NDK build.
Also integrate with the CI to upload the custom build size to scuba.

TODO:
Ideally it should build android/test_app and measure the in-APK size.
But the test_app hasn't been covered by any CI yet and is currently
broken, so build & measure AAR instead (which can be inaccurate as we
plan to pack C++ header files into AAR soon).

Sample result: https://fburl.com/scuba/pytorch_binary_size/skxwb1gh
```

+---------------------+-------------+-------------------+-----------+----------+
|     build_mode      |    arch     |        lib        | Build Num |   Size   |
+---------------------+-------------+-------------------+-----------+----------+
| custom-build-single | armeabi-v7a | libpytorch_jni.so |   5901579 | 3.68 MiB |
| prebuild            | armeabi-v7a | libpytorch_jni.so |   5901014 | 6.23 MiB |
| prebuild            | x86_64      | libpytorch_jni.so |   5901014 | 7.67 MiB |
+---------------------+-------------+-------------------+-----------+----------+
```

Test Plan: Imported from OSS

Differential Revision: D22111115

Pulled By: ljk53

fbshipit-source-id: 11d24efbc49a85f851ecd0e481d14123f405b3a9

Author: ljk53

.circleci/cimodel/data/simple/android_definitions.py

@@ -1,4 +1,4 @@
-import cimodel.data.simple.util.branch_filters
+import cimodel.data.simple.util.branch_filters as branch_filters
 from cimodel.data.simple.util.docker_constants import DOCKER_IMAGE_NDK
 
 
@@ -37,7 +37,7 @@ def gen_tree(self):
         }
 
         if self.is_master_only:
-            props_dict["filters"] = cimodel.data.simple.util.branch_filters.gen_filter_dict()
+            props_dict["filters"] = branch_filters.gen_filter_dict(branch_filters.NON_PR_BRANCH_LIST)
 
         return [{self.template_name: props_dict}]
 
@@ -47,12 +47,14 @@ def __init__(self,
                  job_name,
                  template_name,
                  dependencies,
-                 is_master_only=True):
+                 is_master_only=True,
+                 is_pr_only=False):
 
         self.job_name = job_name
         self.template_name = template_name
         self.dependencies = dependencies
         self.is_master_only = is_master_only
+        self.is_pr_only = is_pr_only
 
     def gen_tree(self):
 
@@ -62,7 +64,9 @@ def gen_tree(self):
         }
 
         if self.is_master_only:
-            props_dict["filters"] = cimodel.data.simple.util.branch_filters.gen_filter_dict()
+            props_dict["filters"] = branch_filters.gen_filter_dict(branch_filters.NON_PR_BRANCH_LIST)
+        elif self.is_pr_only:
+            props_dict["filters"] = branch_filters.gen_filter_dict(branch_filters.PR_BRANCH_LIST)
 
         return [{self.template_name: props_dict}]
 
@@ -77,7 +81,14 @@ def gen_tree(self):
         "pytorch-linux-xenial-py3-clang5-android-ndk-r19c-gradle-build-x86_32",
         "pytorch_android_gradle_build-x86_32",
         ["pytorch_linux_xenial_py3_clang5_android_ndk_r19c_x86_32_build"],
-        is_master_only=False),
+        is_master_only=False,
+        is_pr_only=True),
+    AndroidGradleJob(
+        "pytorch-linux-xenial-py3-clang5-android-ndk-r19c-gradle-custom-build-single",
+        "pytorch_android_gradle_custom_build_single",
+        [],
+        is_master_only=False,
+        is_pr_only=True),
     AndroidGradleJob(
         "pytorch-linux-xenial-py3-clang5-android-ndk-r19c-gradle-build",
         "pytorch_android_gradle_build",

.circleci/cimodel/data/simple/util/branch_filters.py

@@ -4,6 +4,11 @@
     r"/release\/.*/",
 ]
 
+PR_BRANCH_LIST = [
+    r"/gh\/.*\/head/",
+    r"/pull\/.*/",
+]
+
 RC_PATTERN = r"/v[0-9]+(\.[0-9]+)*-rc[0-9]+/"
 
 def gen_filter_dict(

.circleci/config.yml

@@ -1624,14 +1624,6 @@ jobs:
       image: ubuntu-1604:201903-01
     steps:
     - checkout
-    - run:
-        name: filter out not PR runs
-        no_output_timeout: "5m"
-        command: |
-          echo "CIRCLE_PULL_REQUEST: ${CIRCLE_PULL_REQUEST:-}"
-          if [ -z "${CIRCLE_PULL_REQUEST:-}" ]; then
-            circleci step halt
-          fi
     - setup_linux_system_environment
     - checkout
     - setup_ci_environment
@@ -1663,6 +1655,46 @@ jobs:
         path: ~/workspace/build_android_x86_32_artifacts/artifacts.tgz
         destination: artifacts.tgz
 
+  pytorch_android_gradle_custom_build_single:
+    environment:
+      BUILD_ENVIRONMENT: pytorch-linux-xenial-py3-clang5-android-ndk-r19c-gradle-custom-build-single
+      DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-py3-clang5-android-ndk-r19c:209062ef-ab58-422a-b295-36c4eed6e906"
+      PYTHON_VERSION: "3.6"
+    resource_class: large
+    machine:
+      image: ubuntu-1604:201903-01
+    steps:
+    - checkout
+    - setup_linux_system_environment
+    - checkout
+    - setup_ci_environment
+    - run:
+        name: pytorch android gradle custom build single architecture (for PR)
+        no_output_timeout: "1h"
+        command: |
+          set -e
+          # Unlike other gradle jobs, it's not worth building libtorch in a separate CI job and share via docker, because:
+          # 1) Not shareable: it's custom selective build, which is different from default libtorch mobile build;
+          # 2) Not parallelizable by architecture: it only builds libtorch for one architecture;
+
+          echo "DOCKER_IMAGE: ${DOCKER_IMAGE}"
+          time docker pull ${DOCKER_IMAGE} >/dev/null
+          export id=$(docker run --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -t -d -w /var/lib/jenkins ${DOCKER_IMAGE})
+
+          git submodule sync && git submodule update -q --init --recursive
+
+          VOLUME_MOUNTS="-v /home/circleci/project/:/var/lib/jenkins/workspace"
+          export id=$(docker run ${VOLUME_MOUNTS} --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -t -d -w /var/lib/jenkins ${DOCKER_IMAGE})
+
+          export COMMAND='((echo "source ./workspace/env" && echo "export BUILD_ENVIRONMENT=${BUILD_ENVIRONMENT}" && echo "export GRADLE_OFFLINE=1" && echo "sudo chown -R jenkins workspace && cd workspace && ./.circleci/scripts/build_android_gradle.sh") | docker exec -u jenkins -i "$id" bash) 2>&1'
+          echo ${COMMAND} > ./command.sh && unbuffer bash ./command.sh | ts
+
+          # Skip docker push as this job is purely for size analysis purpose.
+          # Result binaries are already in `/home/circleci/project/` as it's mounted instead of copied.
+
+    - upload_binary_size_for_android_build:
+        build_type: custom-build-single
+
   pytorch_ios_build:
     <<: *pytorch_ios_params
     macos:
@@ -7554,9 +7586,21 @@ workflows:
           docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-py3-clang5-android-ndk-r19c:fff7795428560442086f7b2bb6004b65245dc11a"
           name: pytorch_linux_xenial_py3_clang5_android_ndk_r19c_vulkan_x86_32_build
       - pytorch_android_gradle_build-x86_32:
+          filters:
+            branches:
+              only:
+                - /gh\/.*\/head/
+                - /pull\/.*/
           name: pytorch-linux-xenial-py3-clang5-android-ndk-r19c-gradle-build-x86_32
           requires:
             - pytorch_linux_xenial_py3_clang5_android_ndk_r19c_x86_32_build
+      - pytorch_android_gradle_custom_build_single:
+          filters:
+            branches:
+              only:
+                - /gh\/.*\/head/
+                - /pull\/.*/
+          name: pytorch-linux-xenial-py3-clang5-android-ndk-r19c-gradle-custom-build-single
       - pytorch_android_gradle_build:
           filters:
             branches:

.circleci/scripts/build_android_gradle.sh

@@ -1,7 +1,11 @@
 #!/usr/bin/env bash
 set -eux -o pipefail
 
+env
+echo "BUILD_ENVIRONMENT:$BUILD_ENVIRONMENT"
+
 export ANDROID_NDK_HOME=/opt/ndk
+export ANDROID_NDK=/opt/ndk
 export ANDROID_HOME=/opt/android/sdk
 
 # Must be in sync with GRADLE_VERSION in docker image for android
@@ -10,6 +14,31 @@ export GRADLE_VERSION=4.10.3
 export GRADLE_HOME=/opt/gradle/gradle-$GRADLE_VERSION
 export GRADLE_PATH=$GRADLE_HOME/bin/gradle
 
+# touch gradle cache files to prevent expiration
+while IFS= read -r -d '' file
+do
+  touch "$file" || true
+done < <(find /var/lib/jenkins/.gradle -type f -print0)
+
+export GRADLE_LOCAL_PROPERTIES=~/workspace/android/local.properties
+rm -f $GRADLE_LOCAL_PROPERTIES
+echo "sdk.dir=/opt/android/sdk" >> $GRADLE_LOCAL_PROPERTIES
+echo "ndk.dir=/opt/ndk" >> $GRADLE_LOCAL_PROPERTIES
+echo "cmake.dir=/usr/local" >> $GRADLE_LOCAL_PROPERTIES
+
+retry () {
+  $* || (sleep 1 && $*) || (sleep 2 && $*) || (sleep 4 && $*) || (sleep 8 && $*)
+}
+
+# Run custom build script
+if [[ "${BUILD_ENVIRONMENT}" == *-gradle-custom-build* ]]; then
+  # Install torch & torchvision - used to download & dump used ops from test model.
+  retry pip install torch torchvision --progress-bar off
+
+  exec "$(dirname "${BASH_SOURCE[0]}")/../../android/build_test_app_custom.sh" armeabi-v7a
+fi
+
+# Run default build
 BUILD_ANDROID_INCLUDE_DIR_x86=~/workspace/build_android/install/include
 BUILD_ANDROID_LIB_DIR_x86=~/workspace/build_android/install/lib
 
@@ -44,9 +73,6 @@ ln -s ${BUILD_ANDROID_INCLUDE_DIR_arm_v8a} ${JNI_INCLUDE_DIR}/arm64-v8a
 ln -s ${BUILD_ANDROID_LIB_DIR_arm_v8a} ${JNI_LIBS_DIR}/arm64-v8a
 fi
 
-env
-echo "BUILD_ENVIRONMENT:$BUILD_ENVIRONMENT"
-
 GRADLE_PARAMS="-p android assembleRelease --debug --stacktrace"
 if [[ "${BUILD_ENVIRONMENT}" == *-gradle-build-only-x86_32* ]]; then
     GRADLE_PARAMS+=" -PABI_FILTERS=x86"
@@ -56,20 +82,6 @@ if [ -n "{GRADLE_OFFLINE:-}" ]; then
     GRADLE_PARAMS+=" --offline"
 fi
 
-# touch gradle cache files to prevent expiration
-while IFS= read -r -d '' file
-do
-  touch "$file" || true
-done < <(find /var/lib/jenkins/.gradle -type f -print0)
-
-env
-
-export GRADLE_LOCAL_PROPERTIES=~/workspace/android/local.properties
-rm -f $GRADLE_LOCAL_PROPERTIES
-echo "sdk.dir=/opt/android/sdk" >> $GRADLE_LOCAL_PROPERTIES
-echo "ndk.dir=/opt/ndk" >> $GRADLE_LOCAL_PROPERTIES
-echo "cmake.dir=/usr/local" >> $GRADLE_LOCAL_PROPERTIES
-
 $GRADLE_PATH $GRADLE_PARAMS
 
 find . -type f -name "*.a" -exec ls -lh {} \;

.circleci/verbatim-sources/job-specs/job-specs-custom.yml

@@ -283,14 +283,6 @@
       image: ubuntu-1604:201903-01
     steps:
     - checkout
-    - run:
-        name: filter out not PR runs
-        no_output_timeout: "5m"
-        command: |
-          echo "CIRCLE_PULL_REQUEST: ${CIRCLE_PULL_REQUEST:-}"
-          if [ -z "${CIRCLE_PULL_REQUEST:-}" ]; then
-            circleci step halt
-          fi
     - setup_linux_system_environment
     - checkout
     - setup_ci_environment
@@ -322,6 +314,46 @@
         path: ~/workspace/build_android_x86_32_artifacts/artifacts.tgz
         destination: artifacts.tgz
 
+  pytorch_android_gradle_custom_build_single:
+    environment:
+      BUILD_ENVIRONMENT: pytorch-linux-xenial-py3-clang5-android-ndk-r19c-gradle-custom-build-single
+      DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-py3-clang5-android-ndk-r19c:209062ef-ab58-422a-b295-36c4eed6e906"
+      PYTHON_VERSION: "3.6"
+    resource_class: large
+    machine:
+      image: ubuntu-1604:201903-01
+    steps:
+    - checkout
+    - setup_linux_system_environment
+    - checkout
+    - setup_ci_environment
+    - run:
+        name: pytorch android gradle custom build single architecture (for PR)
+        no_output_timeout: "1h"
+        command: |
+          set -e
+          # Unlike other gradle jobs, it's not worth building libtorch in a separate CI job and share via docker, because:
+          # 1) Not shareable: it's custom selective build, which is different from default libtorch mobile build;
+          # 2) Not parallelizable by architecture: it only builds libtorch for one architecture;
+
+          echo "DOCKER_IMAGE: ${DOCKER_IMAGE}"
+          time docker pull ${DOCKER_IMAGE} >/dev/null
+          export id=$(docker run --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -t -d -w /var/lib/jenkins ${DOCKER_IMAGE})
+
+          git submodule sync && git submodule update -q --init --recursive
+
+          VOLUME_MOUNTS="-v /home/circleci/project/:/var/lib/jenkins/workspace"
+          export id=$(docker run ${VOLUME_MOUNTS} --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -t -d -w /var/lib/jenkins ${DOCKER_IMAGE})
+
+          export COMMAND='((echo "source ./workspace/env" && echo "export BUILD_ENVIRONMENT=${BUILD_ENVIRONMENT}" && echo "export GRADLE_OFFLINE=1" && echo "sudo chown -R jenkins workspace && cd workspace && ./.circleci/scripts/build_android_gradle.sh") | docker exec -u jenkins -i "$id" bash) 2>&1'
+          echo ${COMMAND} > ./command.sh && unbuffer bash ./command.sh | ts
+
+          # Skip docker push as this job is purely for size analysis purpose.
+          # Result binaries are already in `/home/circleci/project/` as it's mounted instead of copied.
+
+    - upload_binary_size_for_android_build:
+        build_type: custom-build-single
+
   pytorch_ios_build:
     <<: *pytorch_ios_params
     macos:

android/build_test_app_custom.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+###############################################################################
+# This script tests the custom selective build flow for PyTorch Android, which
+# optimizes library size by only including ops used by a specific model.
+###############################################################################
+
+set -eux
+
+PYTORCH_DIR="$(cd $(dirname $0)/..; pwd -P)"
+PYTORCH_ANDROID_DIR="${PYTORCH_DIR}/android"
+BUILD_ROOT="${PYTORCH_DIR}/build_pytorch_android_custom"
+
+source "${PYTORCH_ANDROID_DIR}/common.sh"
+
+prepare_model_and_dump_root_ops() {
+  cd "${BUILD_ROOT}"
+  MODEL="${BUILD_ROOT}/MobileNetV2.pt"
+  ROOT_OPS="${BUILD_ROOT}/MobileNetV2.yaml"
+  python "${PYTORCH_ANDROID_DIR}/test_app/make_assets_custom.py"
+  cp "${MODEL}" "${PYTORCH_ANDROID_DIR}/test_app/app/src/main/assets/mobilenet2.pt"
+}
+
+# Start building
+mkdir -p "${BUILD_ROOT}"
+check_android_sdk
+check_gradle
+parse_abis_list "$@"
+prepare_model_and_dump_root_ops
+SELECTED_OP_LIST="${ROOT_OPS}" build_android
+
+# TODO: change this to build test_app instead
+$GRADLE_PATH -PABI_FILTERS=$ABIS_LIST -p $PYTORCH_ANDROID_DIR clean assembleRelease

android/test_app/make_assets_custom.py

@@ -0,0 +1,25 @@
+"""
+This is a script for PyTorch Android custom selective build test. It prepares
+MobileNetV2 TorchScript model, and dumps root ops used by the model for custom
+build script to create a tailored build which only contains these used ops.
+"""
+
+import torch
+import torchvision
+import yaml
+
+# Download and trace the model.
+model = torchvision.models.mobilenet_v2(pretrained=True)
+model.eval()
+example = torch.rand(1, 3, 224, 224)
+# TODO: create script model with `torch.jit.script`
+traced_script_module = torch.jit.trace(model, example)
+
+# Save traced TorchScript model.
+traced_script_module.save("MobileNetV2.pt")
+
+# Dump root ops used by the model (for custom build optimization).
+ops = torch.jit.export_opnames(traced_script_module)
+
+with open('MobileNetV2.yaml', 'w') as output:
+    yaml.dump(ops, output)