Add more sample models to OSS CI

.ci/docker/build.sh

@@ -17,17 +17,23 @@ OS_VERSION=22.04
 CLANG_VERSION=12
 PYTHON_VERSION=3.10
 MINICONDA_VERSION=23.5.1-0
-TORCH_VERSION=$(cat ci_commit_pins/pytorch.txt)
 BUCK2_VERSION=$(cat ci_commit_pins/buck2.txt)
 
+NIGHTLY=$(cat ci_commit_pins/nightly.txt)
+TORCH_VERSION=$(cat ci_commit_pins/pytorch.txt)
+TORCHAUDIO_VERSION=$(cat ci_commit_pins/audio.txt)
+TORCHVISION_VERSION=$(cat ci_commit_pins/vision.txt)
+
 docker build \
   --no-cache \
   --progress=plain \
   --build-arg "OS_VERSION=${OS_VERSION}" \
   --build-arg "CLANG_VERSION=${CLANG_VERSION}" \
   --build-arg "PYTHON_VERSION=${PYTHON_VERSION}" \
   --build-arg "MINICONDA_VERSION=${MINICONDA_VERSION}" \
-  --build-arg "TORCH_VERSION=${TORCH_VERSION}" \
+  --build-arg "TORCH_VERSION=${TORCH_VERSION}.${NIGHTLY}" \
+  --build-arg "TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION}.${NIGHTLY}" \
+  --build-arg "TORCHVISION_VERSION=${TORCHVISION_VERSION}.${NIGHTLY}" \
   --build-arg "BUCK2_VERSION=${BUCK2_VERSION}" \
   -f "${OS}"/Dockerfile \
   "$@" \

.ci/docker/ci_commit_pins/audio.txt

@@ -0,0 +1 @@
+2.1.0

.ci/docker/ci_commit_pins/nightly.txt

@@ -0,0 +1 @@
+dev20230813

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-2.1.0.dev20230813
+2.1.0

.ci/docker/ci_commit_pins/vision.txt

@@ -0,0 +1 @@
+0.16.0

.ci/docker/common/install_conda.sh

@@ -40,7 +40,11 @@ install_pip_dependencies() {
   pushd /opt/conda
   # Install all Python dependencies, including PyTorch
   pip_install -r /opt/conda/requirements-ci.txt
-  pip_install --pre torch=="${TORCH_VERSION}" --index-url https://download.pytorch.org/whl/nightly/cpu
+  pip_install --pre \
+    torch=="${TORCH_VERSION}" \
+    torchaudio=="${TORCHAUDIO_VERSION}" \
+    torchvision=="${TORCHVISION_VERSION}" \
+    --index-url https://download.pytorch.org/whl/nightly/cpu
   popd
 }
 

.ci/docker/ubuntu/Dockerfile

@@ -28,6 +28,8 @@ RUN bash ./install_user.sh && rm install_user.sh
 ARG MINICONDA_VERSION
 ARG PYTHON_VERSION
 ARG TORCH_VERSION
+ARG TORCHAUDIO_VERSION
+ARG TORCHVISION_VERSION
 ENV PYTHON_VERSION=$PYTHON_VERSION
 ENV PATH /opt/conda/envs/py_$PYTHON_VERSION/bin:/opt/conda/bin:$PATH
 COPY requirements-ci.txt /opt/conda/

.ci/scripts/gather_test_models.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+import os
+from typing import Any
+
+from examples.models import MODEL_NAME_TO_MODEL
+
+
+def set_output(name: str, val: Any) -> None:
+    """
+    Set the GitHb output so that it can be accessed by other jobs
+    """
+    print(f"Setting {val} to GitHub output")
+
+    if os.getenv("GITHUB_OUTPUT"):
+        with open(str(os.getenv("GITHUB_OUTPUT")), "a") as env:
+            print(f"{name}={val}", file=env)
+    else:
+        print(f"::set-output name={name}::{val}")
+
+
+def export_models_for_ci() -> None:
+    """
+    This gathers all the example models that we want to test on GitHub OSS CI
+    """
+    # This is the JSON syntax for configuration matrix used by GitHub
+    # https://docs.github.com/en/actions/using-jobs/using-a-matrix-for-your-jobs
+    models = {"include": [{"model": name} for name in MODEL_NAME_TO_MODEL.keys()]}
+    set_output("models", json.dumps(models))
+
+
+if __name__ == "__main__":
+    export_models_for_ci()

.ci/scripts/setup-macos.sh

@@ -7,6 +7,9 @@
 
 set -exu
 
+# shellcheck source=/dev/null
+source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
+
 install_buck() {
   if ! command -v zstd &> /dev/null; then
     brew install zstd
@@ -33,23 +36,6 @@ install_buck() {
   fi
 }
 
-install_conda() {
-  pushd .ci/docker
-  # Install conda dependencies like flatbuffer
-  conda install --file conda-env-ci.txt
-  popd
-}
-
-install_pip_dependencies() {
-  pushd .ci/docker
-  # Install all Python dependencies, including PyTorch
-  pip install --progress-bar off -r requirements-ci.txt
-
-  TORCH_VERSION=$(cat ci_commit_pins/pytorch.txt)
-  pip install --progress-bar off --pre torch=="${TORCH_VERSION}" --index-url https://download.pytorch.org/whl/nightly/cpu
-  popd
-}
-
 install_buck
 install_conda
 install_pip_dependencies

.ci/scripts/test-cmake.sh

@@ -10,8 +10,15 @@ set -exu
 # shellcheck source=/dev/null
 source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
 
+MODEL_NAME=$1
+if [[ -z "${MODEL_NAME:-}" ]]; then
+  echo "Missing model name, exiting..."
+  exit 1
+else
+  echo "Testing ${MODEL_NAME} ..."
+fi
+
 test_model() {
-  MODEL_NAME=$1
   python -m examples.export.export_example --model_name="${MODEL_NAME}"
 
   # Run test model
@@ -24,7 +31,7 @@ build_and_test_executorch() {
   rm -rf "${CMAKE_OUTPUT_DIR}" && mkdir "${CMAKE_OUTPUT_DIR}"
 
   pushd "${CMAKE_OUTPUT_DIR}"
-  cmake -DBUCK2=buck2 ..
+  cmake -DBUCK2=buck2 -DPYTHON_EXECUTABLE="${PYTHON_EXECUTABLE}" ..
   popd
 
   if [ "$(uname)" == "Darwin" ]; then
@@ -35,9 +42,13 @@ build_and_test_executorch() {
   cmake --build "${CMAKE_OUTPUT_DIR}" -j "${CMAKE_JOBS}"
 
   which python
-  # Test the example linear model
-  test_model "linear"
+  # Test the select model
+  test_model
 }
 
+if [[ -z "${PYTHON_EXECUTABLE:-}" ]]; then
+  PYTHON_EXECUTABLE=python3
+fi
+
 install_executorch
 build_and_test_executorch

.ci/scripts/test.sh

@@ -10,8 +10,15 @@ set -exu
 # shellcheck source=/dev/null
 source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
 
+MODEL_NAME=$1
+if [[ -z "${MODEL_NAME:-}" ]]; then
+  echo "Missing model name, exiting..."
+  exit 1
+else
+  echo "Testing ${MODEL_NAME} ..."
+fi
+
 test_model() {
-  MODEL_NAME=$1
   python -m examples.export.export_example --model_name="${MODEL_NAME}"
 
   # Run test model
@@ -23,8 +30,8 @@ build_and_test_executorch() {
   buck2 build //examples/executor_runner:executor_runner
 
   which python
-  # Test the example linear model
-  test_model "linear"
+  # Test the select model
+  test_model
 }
 
 install_executorch

.ci/scripts/utils.sh

@@ -13,3 +13,28 @@ install_executorch() {
   # Just print out the list of packages for debugging
   pip list
 }
+
+install_conda() {
+  pushd .ci/docker || return
+  # Install conda dependencies like flatbuffer
+  conda install --file conda-env-ci.txt
+  popd || return
+}
+
+install_pip_dependencies() {
+  pushd .ci/docker || return
+  # Install all Python dependencies, including PyTorch
+  pip install --progress-bar off -r requirements-ci.txt
+
+  NIGHTLY=$(cat ci_commit_pins/nightly.txt)
+  TORCH_VERSION=$(cat ci_commit_pins/pytorch.txt)
+  TORCHAUDIO_VERSION=$(cat ci_commit_pins/audio.txt)
+  TORCHVISION_VERSION=$(cat ci_commit_pins/vision.txt)
+
+  pip install --progress-bar off --pre \
+    torch=="${TORCH_VERSION}.${NIGHTLY}" \
+    torchaudio=="${TORCHAUDIO_VERSION}.${NIGHTLY}" \
+    torchvision=="${TORCHVISION_VERSION}.${NIGHTLY}" \
+    --index-url https://download.pytorch.org/whl/nightly/cpu
+  popd || return
+}

.github/workflows/pull.yml

@@ -12,13 +12,39 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  gather-models:
+    runs-on: ubuntu-22.04
+    outputs:
+      models: ${{ steps.gather-models.outputs.models }}
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: pip
+      - name: Extract the list of models to test
+        id: gather-models
+        run: |
+          set -eux
+
+          source .ci/scripts/utils.sh
+          # This is a simple Python script but as it tries to import executorch.examples.models,
+          # it requires a whole bunch of Executorch dependencies on the Docker image
+          install_pip_dependencies
+          install_executorch
+
+          PYTHONPATH="${PWD}" python .ci/scripts/gather_test_models.py
+
   buck-build-test-linux:
     name: buck-build-test-linux
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    needs: gather-models
+    strategy:
+      matrix: ${{ fromJSON(needs.gather-models.outputs.models) }}
+      fail-fast: false
     with:
       runner: linux.2xlarge
       docker-image: executorch-ubuntu-22.04-clang12
-      fetch-depth: 0
       submodules: 'true'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       script: |
@@ -27,15 +53,20 @@ jobs:
         # here, as it's there in the container
         export PATH="/opt/conda/envs/py_${PYTHON_VERSION}/bin:${PATH}"
 
+        # The name of model we are going to test
+        MODEL_NAME=${{ matrix.model }}
         # Build and test Executorch
-        bash .ci/scripts/test.sh
-
+        bash .ci/scripts/test.sh "${MODEL_NAME}"
         # Test custom ops
         bash examples/custom_ops/test_custom_ops.sh buck2
 
   buck-build-test-macos:
     name: buck-build-test-macos
     uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
+    needs: gather-models
+    strategy:
+      matrix: ${{ fromJSON(needs.gather-models.outputs.models) }}
+      fail-fast: false
     with:
       runner: macos-m1-12
       submodules: 'true'
@@ -47,20 +78,24 @@ jobs:
         # Setup MacOS dependencies as there is no Docker support on MacOS atm
         bash .ci/scripts/setup-macos.sh
 
+        # The name of model we are going to test
+        MODEL_NAME=${{ matrix.model }}
         # Build and test Executorch
-        bash .ci/scripts/test.sh
-
+        PYTHON_EXECUTABLE=python bash .ci/scripts/test.sh "${MODEL_NAME}"
         # Test custom ops
         PYTHON_EXECUTABLE=python bash examples/custom_ops/test_custom_ops.sh buck2
         popd
 
   cmake-build-test-linux:
     name: cmake-build-test-linux
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    needs: gather-models
+    strategy:
+      matrix: ${{ fromJSON(needs.gather-models.outputs.models) }}
+      fail-fast: false
     with:
       runner: linux.2xlarge
       docker-image: executorch-ubuntu-22.04-clang12
-      fetch-depth: 0
       submodules: 'true'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       script: |
@@ -69,15 +104,20 @@ jobs:
         # here, as it's there in the container
         export PATH="/opt/conda/envs/py_${PYTHON_VERSION}/bin:${PATH}"
 
+        # The name of model we are going to test
+        MODEL_NAME=${{ matrix.model }}
         # Build and test Executorch
-        bash .ci/scripts/test-cmake.sh
-
+        bash .ci/scripts/test-cmake.sh "${MODEL_NAME}"
         # Build and test custom ops
         bash examples/custom_ops/test_custom_ops.sh cmake
 
   cmake-build-test-macos:
     name: cmake-build-test-macos
     uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
+    needs: gather-models
+    strategy:
+      matrix: ${{ fromJSON(needs.gather-models.outputs.models) }}
+      fail-fast: false
     with:
       runner: macos-m1-12
       submodules: 'true'
@@ -89,6 +129,10 @@ jobs:
         # Setup MacOS dependencies as there is no Docker support on MacOS atm
         bash .ci/scripts/setup-macos.sh
 
+        # The name of model we are going to test
+        MODEL_NAME=${{ matrix.model }}
+        # Build and test Executorch
+        PYTHON_EXECUTABLE=python bash .ci/scripts/test-cmake.sh "${MODEL_NAME}"
         # Build and test custom ops
         PYTHON_EXECUTABLE=python bash examples/custom_ops/test_custom_ops.sh cmake
         popd