Update

[ghstack-poisoned]

Author: swolchok

.github/workflows/pull.yml

@@ -632,32 +632,33 @@ jobs:
         # run eval_llama wikitext task
         PYTHON_EXECUTABLE=python bash .ci/scripts/test_eval_llama_wikitext.sh
 
-  test-eval_llama-mmlu-linux:
-    name: test-eval_llama-mmlu-linux
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
-    permissions:
-      id-token: write
-      contents: read
-    strategy:
-      fail-fast: false
-    with:
-      runner: linux.24xlarge
-      docker-image: ci-image:executorch-ubuntu-22.04-clang12
-      submodules: 'recursive'
-      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-      timeout: 90
-      script: |
-        # The generic Linux job chooses to use base env, not the one setup by the image
-        CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
-        conda activate "${CONDA_ENV}"
-
-        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh --build-tool "cmake"
-
-        # install llama requirements
-        bash examples/models/llama/install_requirements.sh
-
-        # run eval_llama mmlu task
-        PYTHON_EXECUTABLE=python bash .ci/scripts/test_eval_llama_mmlu.sh
+  # TODO(larryliu0820): Fix this issue before reenabling it: https://gist.github.com/larryliu0820/7377ecd0d79dbc06076cec8d9f2b85d2
+  # test-eval_llama-mmlu-linux:
+  #   name: test-eval_llama-mmlu-linux
+  #   uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+  #   permissions:
+  #     id-token: write
+  #     contents: read
+  #   strategy:
+  #     fail-fast: false
+  #   with:
+  #     runner: linux.24xlarge
+  #     docker-image: ci-image:executorch-ubuntu-22.04-clang12
+  #     submodules: 'recursive'
+  #     ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+  #     timeout: 90
+  #     script: |
+  #       # The generic Linux job chooses to use base env, not the one setup by the image
+  #       CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
+  #       conda activate "${CONDA_ENV}"
+
+  #       PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh --build-tool "cmake"
+
+  #       # install llama requirements
+  #       bash examples/models/llama/install_requirements.sh
+
+  #       # run eval_llama mmlu task
+  #       PYTHON_EXECUTABLE=python bash .ci/scripts/test_eval_llama_mmlu.sh
 
   test-llama_runner_eager-linux:
     name: test-llama_runner_eager-linux

backends/apple/coreml/partition/coreml_partitioner.py

@@ -23,25 +23,27 @@
 from torch.fx.passes.operator_support import OperatorSupportBase
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.WARNING)
+logger.setLevel(logging.INFO)
 
 
-class OperatorsSupportedForCoreMLBackend(OperatorSupportBase):
+class _OperatorsSupportedForCoreMLBackend(OperatorSupportBase):
     def __init__(
         self,
         skip_ops_for_coreml_delegation: Optional[List[str]] = None,
         lower_full_graph: bool = False,
+        log: bool = False,
     ) -> None:
         if skip_ops_for_coreml_delegation is None:
             skip_ops_for_coreml_delegation = []
         super().__init__()
         self.skip_ops_for_coreml_delegation = skip_ops_for_coreml_delegation
         self.lower_full_graph = lower_full_graph
         self._logged_msgs = set()
+        self._log = log
 
     def log_once(self, msg: str) -> None:
-        if msg not in self._logged_msgs:
-            logging.info(msg)
+        if self._log and msg not in self._logged_msgs:
+            logger.info(msg)
             self._logged_msgs.add(msg)
 
     def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
@@ -154,8 +156,10 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
 
         capability_partitioner = CapabilityBasedPartitioner(
             exported_program.graph_module,
-            OperatorsSupportedForCoreMLBackend(
-                self.skip_ops_for_coreml_delegation, self.lower_full_graph
+            _OperatorsSupportedForCoreMLBackend(
+                self.skip_ops_for_coreml_delegation,
+                self.lower_full_graph,
+                log=True,
             ),
             allows_single_node_partition=True,
         )
@@ -191,8 +195,10 @@ def ops_to_not_decompose(
         self, ep: ExportedProgram
     ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]:
         do_not_decompose = []
-        op_support = OperatorsSupportedForCoreMLBackend(
-            self.skip_ops_for_coreml_delegation, self.lower_full_graph
+        op_support = _OperatorsSupportedForCoreMLBackend(
+            self.skip_ops_for_coreml_delegation,
+            self.lower_full_graph,
+            log=False,
         )
 
         # CoreML prevents certain ops (like triu) from lowering to CoreML when put in the ExecuTorch op namespace

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -16,7 +16,6 @@
 from executorch.backends.apple.coreml.compiler import CoreMLBackend
 from executorch.backends.apple.coreml.partition import CoreMLPartitioner
 from executorch.exir.backend.utils import format_delegated_graph
-from executorch.runtime import Runtime
 
 
 @torch.library.custom_op("unsupported::linear", mutates_args=())
@@ -37,7 +36,13 @@ def _(
     return torch.ops.aten.linear.default(x, w, b)
 
 
-_TEST_RUNTIME = sys.platform == "darwin"
+def is_fbcode():
+    return not hasattr(torch.version, "git_version")
+
+
+_TEST_RUNTIME = (sys.platform == "darwin") and not is_fbcode()
+if _TEST_RUNTIME:
+    from executorch.runtime import Runtime
 
 
 class TestCoreMLPartitioner(unittest.TestCase):

backends/apple/coreml/test/test_torch_ops.py

@@ -14,12 +14,20 @@
 
 from executorch.backends.apple.coreml.compiler import CoreMLBackend
 from executorch.backends.apple.coreml.partition import CoreMLPartitioner
-from executorch.runtime import Runtime
 from torchao.quantization import IntxWeightOnlyConfig, PerAxis, PerGroup, quantize_
 
-_TEST_RUNTIME = sys.platform == "darwin" and tuple(
-    map(int, platform.mac_ver()[0].split("."))
-) >= (15, 0)
+
+def is_fbcode():
+    return not hasattr(torch.version, "git_version")
+
+
+_TEST_RUNTIME = (
+    (sys.platform == "darwin")
+    and not is_fbcode()
+    and tuple(map(int, platform.mac_ver()[0].split("."))) >= (15, 0)
+)
+if _TEST_RUNTIME:
+    from executorch.runtime import Runtime
 
 
 class TestTorchOps(unittest.TestCase):

backends/test/suite/__init__.py

@@ -12,11 +12,13 @@
 import unittest
 
 from enum import Enum
-from typing import Any, Callable, Tuple
+from typing import Callable
+
+import executorch.backends.test.suite.flow
 
 import torch
-from executorch.backends.test.harness import Tester
 from executorch.backends.test.suite.context import get_active_test_context, TestContext
+from executorch.backends.test.suite.flow import TestFlow
 from executorch.backends.test.suite.reporting import log_test_summary
 from executorch.backends.test.suite.runner import run_test, runner_main
 
@@ -44,22 +46,20 @@ def is_backend_enabled(backend):
         return backend in _ENABLED_BACKENDS
 
 
-ALL_TEST_FLOWS = []
+_ALL_TEST_FLOWS: dict[str, TestFlow] = {}
 
-if is_backend_enabled("xnnpack"):
-    from executorch.backends.xnnpack.test.tester import Tester as XnnpackTester
 
-    XNNPACK_TEST_FLOW = ("xnnpack", XnnpackTester)
-    ALL_TEST_FLOWS.append(XNNPACK_TEST_FLOW)
+def get_test_flows() -> dict[str, TestFlow]:
+    global _ALL_TEST_FLOWS
 
-if is_backend_enabled("coreml"):
-    try:
-        from executorch.backends.apple.coreml.test.tester import CoreMLTester
+    if not _ALL_TEST_FLOWS:
+        _ALL_TEST_FLOWS = {
+            name: f
+            for name, f in executorch.backends.test.suite.flow.all_flows().items()
+            if is_backend_enabled(f.backend)
+        }
 
-        COREML_TEST_FLOW = ("coreml", CoreMLTester)
-        ALL_TEST_FLOWS.append(COREML_TEST_FLOW)
-    except Exception:
-        print("Core ML AOT is not available.")
+    return _ALL_TEST_FLOWS
 
 
 DTYPES = [
@@ -115,53 +115,51 @@ def _create_tests(cls):
 # Expand a test into variants for each registered flow.
 def _expand_test(cls, test_name: str):
     test_func = getattr(cls, test_name)
-    for flow_name, tester_factory in ALL_TEST_FLOWS:
-        _create_test_for_backend(cls, test_func, flow_name, tester_factory)
+    for flow in get_test_flows().values():
+        _create_test_for_backend(cls, test_func, flow)
     delattr(cls, test_name)
 
 
 def _make_wrapped_test(
     test_func: Callable,
     test_name: str,
-    test_flow: str,
-    tester_factory: Callable,
+    flow: TestFlow,
     params: dict | None = None,
 ):
     def wrapped_test(self):
-        with TestContext(test_name, test_flow, params):
+        with TestContext(test_name, flow.name, params):
             test_kwargs = params or {}
-            test_kwargs["tester_factory"] = tester_factory
+            test_kwargs["tester_factory"] = flow.tester_factory
 
             test_func(self, **test_kwargs)
 
+    wrapped_test._name = test_name
+    wrapped_test._flow = flow
+
     return wrapped_test
 
 
 def _create_test_for_backend(
     cls,
     test_func: Callable,
-    flow_name: str,
-    tester_factory: Callable[[torch.nn.Module, Tuple[Any]], Tester],
+    flow: TestFlow,
 ):
     test_type = getattr(test_func, "test_type", TestType.STANDARD)
 
     if test_type == TestType.STANDARD:
-        wrapped_test = _make_wrapped_test(
-            test_func, test_func.__name__, flow_name, tester_factory
-        )
-        test_name = f"{test_func.__name__}_{flow_name}"
+        wrapped_test = _make_wrapped_test(test_func, test_func.__name__, flow)
+        test_name = f"{test_func.__name__}_{flow.name}"
         setattr(cls, test_name, wrapped_test)
     elif test_type == TestType.DTYPE:
         for dtype in DTYPES:
             wrapped_test = _make_wrapped_test(
                 test_func,
                 test_func.__name__,
-                flow_name,
-                tester_factory,
+                flow,
                 {"dtype": dtype},
             )
             dtype_name = str(dtype)[6:]  # strip "torch."
-            test_name = f"{test_func.__name__}_{dtype_name}_{flow_name}"
+            test_name = f"{test_func.__name__}_{dtype_name}_{flow.name}"
             setattr(cls, test_name, wrapped_test)
     else:
         raise NotImplementedError(f"Unknown test type {test_type}.")

backends/test/suite/discovery.py

@@ -0,0 +1,63 @@
+# 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.
+
+# pyre-unsafe
+
+import os
+import unittest
+
+from types import ModuleType
+
+from executorch.backends.test.suite.flow import TestFlow
+
+#
+# This file contains logic related to test discovery and filtering.
+#
+
+
+def discover_tests(
+    root_module: ModuleType, backends: set[str] | None
+) -> unittest.TestSuite:
+    # Collect all tests using the unittest discovery mechanism then filter down.
+
+    # Find the file system path corresponding to the root module.
+    module_file = root_module.__file__
+    if module_file is None:
+        raise RuntimeError(f"Module {root_module} has no __file__ attribute")
+
+    loader = unittest.TestLoader()
+    module_dir = os.path.dirname(module_file)
+    suite = loader.discover(module_dir)
+
+    return _filter_tests(suite, backends)
+
+
+def _filter_tests(
+    suite: unittest.TestSuite, backends: set[str] | None
+) -> unittest.TestSuite:
+    # Recursively traverse the test suite and add them to the filtered set.
+    filtered_suite = unittest.TestSuite()
+
+    for child in suite:
+        if isinstance(child, unittest.TestSuite):
+            filtered_suite.addTest(_filter_tests(child, backends))
+        elif isinstance(child, unittest.TestCase):
+            if _is_test_enabled(child, backends):
+                filtered_suite.addTest(child)
+        else:
+            raise RuntimeError(f"Unexpected test type: {type(child)}")
+
+    return filtered_suite
+
+
+def _is_test_enabled(test_case: unittest.TestCase, backends: set[str] | None) -> bool:
+    test_method = getattr(test_case, test_case._testMethodName)
+
+    if backends is not None:
+        flow: TestFlow = test_method._flow
+        return flow.backend in backends
+    else:
+        return True