update smoke test

Author: jathu

.ci/scripts/gather_test_models.py

@@ -113,13 +113,10 @@ def model_should_run_on_target_os(model: str, target_os: str) -> bool:
     return model not in ["llava"]
 
 
-def export_models_for_ci() -> dict[str, dict]:
+def export_models_for_ci(target_os: str, event: str) -> dict[str, Any]:
     """
     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
@@ -175,9 +172,14 @@ def export_models_for_ci() -> dict[str, dict]:
             record["runner"] = CUSTOM_RUNNERS[target_os][name]
 
         models["include"].append(record)
-
-    set_output("models", json.dumps(models))
+    
+    return models
 
 
 if __name__ == "__main__":
-    export_models_for_ci()
+    args = parse_args()
+    models = export_models_for_ci(
+        target_os = args.target_os,
+        event = args.event
+    )
+    set_output("models", json.dumps(models))

.github/workflows/build-wheels-linux.yml

@@ -41,7 +41,7 @@ jobs:
           - repository: pytorch/executorch
             pre-script: build/packaging/pre_build_script.sh
             post-script: build/packaging/post_build_script.sh
-            smoke-test-script: build/packaging/smoke_test.py
+            smoke-test-script: build/util/packaging/wheel_test_linux.py
             package-name: executorch
     name: ${{ matrix.repository }}
     uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main

.github/workflows/build-wheels-macos.yml

@@ -41,7 +41,7 @@ jobs:
           - repository: pytorch/executorch
             pre-script: build/packaging/pre_build_script.sh
             post-script: build/packaging/post_build_script.sh
-            smoke-test-script: build/packaging/smoke_test.py
+            smoke-test-script: build/util/packaging/wheel_test_macos.py
             package-name: executorch
     name: ${{ matrix.repository }}
     uses: pytorch/test-infra/.github/workflows/build_wheels_macos.yml@main

build/packaging/smoke_test.py

@@ -0,0 +1,44 @@
+import sys
+import os
+import argparse
+import subprocess
+from typing import List, Any
+
+_PYTHON_PATH = os.getenv("PYTHONPATH")
+if _PYTHON_PATH is None:
+  print("PYTHONPATH is not set")
+  exit(1)
+
+# Since the .ci folder starts with a period, it is not possible to import it
+# directly, so we must use the importlib module to load it.
+sys.path.append(os.path.join(_PYTHON_PATH, ".ci"))
+from scripts import gather_test_models
+
+
+def list_models(target_os: str) -> List[Any]:
+  if target_os not in gather_test_models.DEFAULT_RUNNERS:
+    print(f"unknown os '{target_os}' provided")
+    exit(1)
+
+  return gather_test_models.export_models_for_ci(
+    target_os = target_os,
+    # Event refers to the type of models that will be downloaded. "pull_request"
+    # uses higher priority and fast models.
+    event="pull_request",
+  ).get("include", [])
+
+def run_tests(models: List[Any]) -> None:
+  if len(models) == 0:
+    print("No models found")
+    exit(1)
+
+  for model in models:
+    subprocess.run(
+        [
+          os.path.join(_PYTHON_PATH, ".ci/scripts/test_model.sh"), 
+          model["model"], 
+          model["build-tool"], 
+          model["backend"],
+        ],
+        check=True,
+    )

build/util/packaging/__init__.py

@@ -0,0 +1,6 @@
+import wheel_test_base
+
+if __name__ == "__main__":
+    wheel_test_base.run_tests(
+      models=wheel_test_base.list_models("linux")
+    )

build/util/packaging/wheel_test_base.py

@@ -0,0 +1,6 @@
+import wheel_test_base
+
+if __name__ == "__main__":
+    wheel_test_base.run_tests(
+      models=wheel_test_base.list_models("macos")
+    )