Improve gather_test_models (#1038)

Summary:
* In export_models_for_ci(), use itertools.product to simplify the code
* Add a helper model_should_run_on_event() to check whether we are doing pull.yml (pull_request) or trunk.yml (push). Only run high priority models for pull_request.

To test, the envvar set by the script should be same, except for the models filtered by model_should_run_on_event() when event is pull_request

Pull Request resolved: #1038

Reviewed By: huydhn

Differential Revision: D50485044

Pulled By: kirklandsign

fbshipit-source-id: 366cc51dfc9b2aef734e1b8acbef0fe56c9d6644

Author: kirklandsign

.ci/scripts/gather_test_models.py

@@ -5,6 +5,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import itertools
 import json
 import os
 from typing import Any
@@ -47,6 +48,15 @@ def parse_args() -> Any:
         default="linux",
         help="the target OS",
     )
+    parser.add_argument(
+        "-e",
+        "--event",
+        type=str,
+        choices=["pull_request", "push"],
+        required=True,
+        help="GitHub CI Event. See https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions#on",
+    )
+
     return parser.parse_args()
 
 
@@ -63,49 +73,64 @@ def set_output(name: str, val: Any) -> None:
         print(f"::set-output name={name}::{val}")
 
 
-def export_models_for_ci() -> None:
+def model_should_run_on_event(model: str, event: str) -> bool:
+    """
+    A helper function to decide whether a model should be tested on an event (pull_request/push)
+    We put higher priority and fast models to pull request and rest to push.
+    """
+    if event == "pull_request":
+        return model in ["add", "ic3", "mv2", "mv3", "resnet18", "vit"]
+    return True
+
+
+def export_models_for_ci() -> dict[str, dict]:
     """
     This gathers all the example models that we want to test on GitHub OSS CI
     """
     args = parse_args()
     target_os = args.target_os
+    event = args.event
 
     # 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": []}
-    for name in MODEL_NAME_TO_MODEL.keys():
-        quantization_configs = {
-            False,
-            name in MODEL_NAME_TO_OPTIONS and MODEL_NAME_TO_OPTIONS[name].quantization,
-        }
-        delegation_configs = {
-            False,
-            name in MODEL_NAME_TO_OPTIONS and MODEL_NAME_TO_OPTIONS[name].delegation,
+    for (name, build_tool, q_config, d_config) in itertools.product(
+        MODEL_NAME_TO_MODEL.keys(), BUILD_TOOLS.keys(), [False, True], [False, True]
+    ):
+        if not model_should_run_on_event(name, event):
+            continue
+
+        if q_config and (
+            (name not in MODEL_NAME_TO_OPTIONS)
+            or (not MODEL_NAME_TO_OPTIONS[name].quantization)
+        ):
+            continue
+
+        if d_config and (
+            (name not in MODEL_NAME_TO_OPTIONS)
+            or (not MODEL_NAME_TO_OPTIONS[name].delegation)
+        ):
+            continue
+
+        if target_os not in BUILD_TOOLS[build_tool]:
+            continue
+
+        record = {
+            "build-tool": build_tool,
+            "model": name,
+            "xnnpack_quantization": q_config,
+            "xnnpack_delegation": d_config,
+            "runner": DEFAULT_RUNNERS.get(target_os, "linux.2xlarge"),
+            # demo_backend_delegation test only supports add_mul model
+            "demo_backend_delegation": name == "add_mul",
         }
-        for build_tool in BUILD_TOOLS.keys():
-            if target_os not in BUILD_TOOLS[build_tool]:
-                continue
-
-            for q_config in quantization_configs:
-                for d_config in delegation_configs:
-                    record = {
-                        "build-tool": build_tool,
-                        "model": name,
-                        "xnnpack_quantization": q_config,
-                        "xnnpack_delegation": d_config,
-                        "runner": DEFAULT_RUNNERS.get(target_os, "linux.2xlarge"),
-                        # demo_backend_delegation test only supports add_mul model
-                        "demo_backend_delegation": name == "add_mul",
-                    }
-
-                    # NB: Some model requires much bigger Linux runner to avoid
-                    # running OOM. The team is investigating the root cause
-                    if target_os in CUSTOM_RUNNERS and name in CUSTOM_RUNNERS.get(
-                        target_os, {}
-                    ):
-                        record["runner"] = CUSTOM_RUNNERS[target_os][name]
-
-                    models["include"].append(record)
+
+        # NB: Some model requires much bigger Linux runner to avoid
+        # running OOM. The team is investigating the root cause
+        if target_os in CUSTOM_RUNNERS and name in CUSTOM_RUNNERS.get(target_os, {}):
+            record["runner"] = CUSTOM_RUNNERS[target_os][name]
+
+        models["include"].append(record)
 
     set_output("models", json.dumps(models))
 

.github/workflows/pull.yml

@@ -36,7 +36,7 @@ jobs:
           install_pip_dependencies
           install_executorch
 
-          PYTHONPATH="${PWD}" python .ci/scripts/gather_test_models.py
+          PYTHONPATH="${PWD}" python .ci/scripts/gather_test_models.py --event "${GITHUB_EVENT_NAME}"
 
   test-models-linux:
     name: test-models-linux

.github/workflows/trunk.yml

@@ -35,7 +35,7 @@ jobs:
           install_pip_dependencies
           install_executorch
 
-          PYTHONPATH="${PWD}" python .ci/scripts/gather_test_models.py --target-os macos
+          PYTHONPATH="${PWD}" python .ci/scripts/gather_test_models.py --target-os macos --event "${GITHUB_EVENT_NAME}"
 
   test-models-macos:
     name: test-models-macos