Repository: mlflow/mlflow

Issue:
Split tests to run them faster
## Motivation

The number of tests is growing. It takes a long time to run them. We should consider splitting them up for a faster iteration cycle.

## Example code:

```diff
diff --git a/.github/workflows/master.yml b/.github/workflows/master.yml
index 22e09d7bb..605f46762 100644
--- a/.github/workflows/master.yml
+++ b/.github/workflows/master.yml
@@ -235,6 +235,11 @@ jobs:
     if: github.event_name != 'pull_request' || github.event.pull_request.draft == false
     runs-on: ubuntu-latest
     timeout-minutes: 120
+    strategy:
+      matrix:
+        group: [1, 2, 3]
+        include:
+          - splits: 3
     steps:
       - uses: actions/checkout@v3
         with:
@@ -252,7 +257,7 @@ jobs:
       - name: Run tests
         run: |
           export MLFLOW_HOME=$(pwd)
-          pytest tests/models
+          pytest --splits=${{ matrix.splits }} --group=${{ matrix.group }} tests/models
 
   # NOTE: numpy is pinned in this suite due to its heavy reliance on shap, which internally uses
   # references to the now fully deprecated (as of 1.24.x) numpy types (i.e., np.bool).
diff --git a/conftest.py b/conftest.py
index 640680cb3..266702d90 100644
--- a/conftest.py
+++ b/conftest.py
@@ -27,6 +27,18 @@ def pytest_addoption(parser):
         default=False,
         help="Ignore tests for model flavors.",
     )
+    parser.addoption(
+        "--splits",
+        default=None,
+        type=int,
+        help="The number of groups to split tests into.",
+    )
+    parser.addoption(
+        "--group",
+        default=None,
+        type=int,
+        help="The group of tests to run.",
+    )
 
 
 def pytest_configure(config):
@@ -36,6 +48,29 @@ def pytest_configure(config):
     config.addinivalue_line("markers", "allow_infer_pip_requirements_fallback")
 
 
+@pytest.hookimpl(tryfirst=True)
+def pytest_cmdline_main(config):
+    group = config.getoption("group")
+    splits = config.getoption("splits")
+
+    if splits is None and group is None:
+        return None
+
+    if splits and group is None:
+        raise pytest.UsageError("`--group` is required")
+
+    if group and splits is None:
+        raise pytest.UsageError("`--splits` is required")
+
+    if splits < 0:
+        raise pytest.UsageError("`--splits` must be >= 1")
+
+    if group < 1 or group > splits:
+        raise pytest.UsageError(f"`--group` must be between 1 and {splits}")
+
+    return None
+
+
 def pytest_sessionstart(session):
     if uri := MLFLOW_TRACKING_URI.get():
         click.echo(
@@ -147,6 +182,7 @@ def pytest_ignore_collect(path, config):
             outcome.force_result(True)
 
 
+@pytest.hookimpl(trylast=True)
 def pytest_collection_modifyitems(session, config, items):  # pylint: disable=unused-argument
     # Executing `tests.server.test_prometheus_exporter` after `tests.server.test_handlers`
     # results in an error because Flask >= 2.2.0 doesn't allow calling setup method such as
@@ -154,6 +190,10 @@ def pytest_collection_modifyitems(session, config, items):  # pylint: disable=un
     # execute `tests.server.test_prometheus_exporter` first by reordering the test items.
     items.sort(key=lambda item: item.module.__name__ != "tests.server.test_prometheus_exporter")
 
+    # Select the tests to run based on the group and splits
+    if (splits := config.getoption("--splits")) and (group := config.getoption("--group")):
+        items[:] = items[(group - 1) :: splits]
+
 
 @pytest.hookimpl(hookwrapper=True)
 def pytest_terminal_summary(

```

## Notes

- This is an experimental attempt. We can first parallelize the `models` job (takes 45 ~ 50 min.). If this succeeds, we can incrementally update other jobs.
 


Discussion:
@TomeHirata Interested in working on this?
Mention me if anyone wants to work on this.
Hi @harupy 
Can I work on this?
@Gekko0114 Thanks, assigned!

Resolving patch:
diff --git a/conftest.py b/conftest.py
--- a/conftest.py
+++ b/conftest.py
@@ -27,6 +27,18 @@ def pytest_addoption(parser):
         default=False,
         help="Ignore tests for model flavors.",
     )
+    parser.addoption(
+        "--splits",
+        default=None,
+        type=int,
+        help="The number of groups to split tests into.",
+    )
+    parser.addoption(
+        "--group",
+        default=None,
+        type=int,
+        help="The group of tests to run.",
+    )
 
 
 def pytest_configure(config):
@@ -36,6 +48,29 @@ def pytest_configure(config):
     config.addinivalue_line("markers", "allow_infer_pip_requirements_fallback")
 
 
+@pytest.hookimpl(tryfirst=True)
+def pytest_cmdline_main(config):
+    group = config.getoption("group")
+    splits = config.getoption("splits")
+
+    if splits is None and group is None:
+        return None
+
+    if splits and group is None:
+        raise pytest.UsageError("`--group` is required")
+
+    if group and splits is None:
+        raise pytest.UsageError("`--splits` is required")
+
+    if splits < 0:
+        raise pytest.UsageError("`--splits` must be >= 1")
+
+    if group < 1 or group > splits:
+        raise pytest.UsageError("`--group` must be between 1 and {splits}")
+
+    return None
+
+
 def pytest_sessionstart(session):
     if uri := MLFLOW_TRACKING_URI.get():
         click.echo(
@@ -147,6 +182,7 @@ def pytest_ignore_collect(path, config):
             outcome.force_result(True)
 
 
+@pytest.hookimpl(trylast=True)
 def pytest_collection_modifyitems(session, config, items):  # pylint: disable=unused-argument
     # Executing `tests.server.test_prometheus_exporter` after `tests.server.test_handlers`
     # results in an error because Flask >= 2.2.0 doesn't allow calling setup method such as
@@ -154,6 +190,10 @@ def pytest_collection_modifyitems(session, config, items):  # pylint: disable=un
     # execute `tests.server.test_prometheus_exporter` first by reordering the test items.
     items.sort(key=lambda item: item.module.__name__ != "tests.server.test_prometheus_exporter")
 
+    # Select the tests to run based on the group and splits
+    if (splits := config.getoption("--splits")) and (group := config.getoption("--group")):
+        items[:] = items[(group - 1) :: splits]
+
 
 @pytest.hookimpl(hookwrapper=True)
 def pytest_terminal_summary(


Test patch: