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,16 @@ def export_models_for_ci() -> dict[str, dict]:
             record["runner"] = CUSTOM_RUNNERS[target_os][name]
 
         models["include"].append(record)
+    
+    return models
 
-    set_output("models", json.dumps(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.py --target-os linux
             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.py --target-os macos
             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,59 @@
+import sys
+import os
+import argparse
+import subprocess
+# 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.
+_PYTHON_PATH = os.getenv("PYTHONPATH")
+if _PYTHON_PATH is None:
+  print("PYTHONPATH is not set")
+  exit(1)
+
+sys.path.append(os.path.join(_PYTHON_PATH, ".ci"))
+from scripts import gather_test_models
+
+
+def _create_arg_parser() -> argparse.ArgumentParser:
+  parser = argparse.ArgumentParser()
+
+  parser.add_argument(
+      "--target-os",
+      type=str,
+      required=True,
+      choices=gather_test_models.DEFAULT_RUNNERS.keys(),
+      help="the target OS",
+  )
+
+  return parser
+
+def _run_test(model_name: str, build_tool: str, backend: str) -> None:
+  subprocess.run(
+      [
+        os.path.join(_PYTHON_PATH, ".ci/scripts/test_model.sh"), 
+        model_name, 
+        build_tool, 
+        backend,
+      ],
+      env={**os.environ, "PYTHON_EXECUTABLE": "python"},
+      check=True,
+  )
+
+if __name__ == "__main__":
+  args = _create_arg_parser().parse_args()
+  models = gather_test_models.export_models_for_ci(
+    target_os = args.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", [])
+
+  if len(models) == 0:
+    print("No models found")
+    exit(1)
+
+  for model in models:
+    _run_test(
+      model_name=model["model"],
+      build_tool=model["build-tool"],
+      backend=model["backend"],
+    )