diff --git a/python/tvm/testing/plugin.py b/python/tvm/testing/plugin.py index 1f4f983b7210..2d845b70ff11 100644 --- a/python/tvm/testing/plugin.py +++ b/python/tvm/testing/plugin.py @@ -37,6 +37,13 @@ import tvm from tvm.testing import utils +try: + from xdist.scheduler.loadscope import LoadScopeScheduling + + HAVE_XDIST = True +except ImportError: + HAVE_XDIST = False + MARKERS = { "gpu": "mark a test as requiring a gpu", @@ -319,3 +326,38 @@ def _parametrize_correlated_parameters(metafunc): names = ",".join(name for name, values in params) value_sets = zip(*[values for name, values in params]) metafunc.parametrize(names, value_sets, indirect=True, ids=ids) + + +# pytest-xdist isn't required but is used in CI, so guard on its presence +if HAVE_XDIST: + + def pytest_xdist_make_scheduler(config, log): + """ + Serialize certain tests for pytest-xdist that have inter-test + dependencies + """ + + class TvmTestScheduler(LoadScopeScheduling): + """ + Scheduler to serializer tests + """ + + def _split_scope(self, nodeid): + """ + Returns a specific string for classes of nodeids + """ + # NOTE: these tests contain inter-test dependencies and must be + # serialized + items = { + "test_tvm_testing_features": "functional-tests", + "tests/python/unittest/test_crt": "crt-tests", + "tests/python/driver/tvmc": "tvmc-tests", + } + + for nodeid_pattern, suite_name in items.items(): + if nodeid_pattern in nodeid: + return suite_name + + return nodeid + + return TvmTestScheduler(config, log)