We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.
 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `conftest.py`
Content:
```
1 import json
2 import os
3 import posixpath
4 import shutil
5 import subprocess
6 
7 import click
8 import pytest
9 
10 from mlflow.environment_variables import _MLFLOW_TESTING, MLFLOW_TRACKING_URI
11 
12 
13 def pytest_addoption(parser):
14     parser.addoption(
15         "--requires-ssh",
16         action="store_true",
17         dest="requires_ssh",
18         default=False,
19         help="Run tests decorated with 'requires_ssh' annotation. "
20         "These tests require keys to be configured locally "
21         "for SSH authentication.",
22     )
23     parser.addoption(
24         "--ignore-flavors",
25         action="store_true",
26         dest="ignore_flavors",
27         default=False,
28         help="Ignore tests for model flavors.",
29     )
30 
31 
32 def pytest_configure(config):
33     # Register markers to suppress `PytestUnknownMarkWarning`
34     config.addinivalue_line("markers", "requires_ssh")
35     config.addinivalue_line("markers", "notrackingurimock")
36     config.addinivalue_line("markers", "allow_infer_pip_requirements_fallback")
37 
38 
39 def pytest_sessionstart(session):
40     if uri := MLFLOW_TRACKING_URI.get():
41         click.echo(
42             click.style(
43                 (
44                     f"Environment variable {MLFLOW_TRACKING_URI} is set to {uri!r}, "
45                     "which may interfere with tests."
46                 ),
47                 fg="red",
48             )
49         )
50 
51 
52 def pytest_runtest_setup(item):
53     markers = [mark.name for mark in item.iter_markers()]
54     if "requires_ssh" in markers and not item.config.getoption("--requires-ssh"):
55         pytest.skip("use `--requires-ssh` to run this test")
56 
57 
58 @pytest.hookimpl(hookwrapper=True)
59 def pytest_report_teststatus(report, config):
60     outcome = yield
61     if report.when == "call":
62         try:
63             import psutil
64         except ImportError:
65             return
66 
67         (*rest, result) = outcome.get_result()
68         mem = psutil.virtual_memory()
69         mem_used = mem.used / 1024**3
70         mem_total = mem.total / 1024**3
71 
72         disk = psutil.disk_usage("/")
73         disk_used = disk.used / 1024**3
74         disk_total = disk.total / 1024**3
75         outcome.force_result(
76             (
77                 *rest,
78                 (
79                     f"{result} | "
80                     f"MEM {mem_used:.1f}/{mem_total:.1f} GB | "
81                     f"DISK {disk_used:.1f}/{disk_total:.1f} GB"
82                 ),
83             )
84         )
85 
86 
87 @pytest.hookimpl(hookwrapper=True)
88 def pytest_ignore_collect(path, config):
89     outcome = yield
90     if not outcome.get_result() and config.getoption("ignore_flavors"):
91         # If not ignored by the default hook and `--ignore-flavors` specified
92 
93         # Ignored files and directories must be included in dev/run-python-flavor-tests.sh
94         model_flavors = [
95             # Tests of flavor modules.
96             "tests/azureml",
97             "tests/catboost",
98             "tests/diviner",
99             "tests/fastai",
100             "tests/gluon",
101             "tests/h2o",
102             "tests/johnsnowlabs",
103             "tests/keras",
104             "tests/keras_core",
105             "tests/langchain",
106             "tests/lightgbm",
107             "tests/mleap",
108             "tests/models",
109             "tests/onnx",
110             "tests/openai",
111             "tests/paddle",
112             "tests/pmdarima",
113             "tests/prophet",
114             "tests/pyfunc",
115             "tests/pytorch",
116             "tests/sagemaker",
117             "tests/sentence_transformers",
118             "tests/shap",
119             "tests/sklearn",
120             "tests/spacy",
121             "tests/spark",
122             "tests/statsmodels",
123             "tests/tensorflow",
124             "tests/transformers",
125             "tests/xgboost",
126             # Lazy loading test.
127             "tests/test_mlflow_lazily_imports_ml_packages.py",
128             # Tests of utils.
129             "tests/utils/test_model_utils.py",
130             # This test is included here because it imports many big libraries like tf, keras, etc.
131             "tests/tracking/fluent/test_fluent_autolog.py",
132             # Cross flavor autologging related tests.
133             "tests/autologging/test_autologging_safety_unit.py",
134             "tests/autologging/test_autologging_behaviors_unit.py",
135             "tests/autologging/test_autologging_behaviors_integration.py",
136             "tests/autologging/test_autologging_utils.py",
137             "tests/autologging/test_training_session.py",
138             # Opt in authentication feature.
139             "tests/server/auth",
140             "tests/gateway",
141         ]
142 
143         relpath = os.path.relpath(str(path))
144         relpath = relpath.replace(os.sep, posixpath.sep)  # for Windows
145 
146         if relpath in model_flavors:
147             outcome.force_result(True)
148 
149 
150 def pytest_collection_modifyitems(session, config, items):  # pylint: disable=unused-argument
151     # Executing `tests.server.test_prometheus_exporter` after `tests.server.test_handlers`
152     # results in an error because Flask >= 2.2.0 doesn't allow calling setup method such as
153     # `before_request` on the application after the first request. To avoid this issue,
154     # execute `tests.server.test_prometheus_exporter` first by reordering the test items.
155     items.sort(key=lambda item: item.module.__name__ != "tests.server.test_prometheus_exporter")
156 
157 
158 @pytest.hookimpl(hookwrapper=True)
159 def pytest_terminal_summary(
160     terminalreporter, exitstatus, config
161 ):  # pylint: disable=unused-argument
162     yield
163     failed_test_reports = terminalreporter.stats.get("failed", [])
164     if failed_test_reports:
165         if len(failed_test_reports) <= 30:
166             terminalreporter.section("command to run failed test cases")
167             ids = [repr(report.nodeid) for report in failed_test_reports]
168         else:
169             terminalreporter.section("command to run failed test suites")
170             # Use dict.fromkeys to preserve the order
171             ids = list(dict.fromkeys(report.fspath for report in failed_test_reports))
172         terminalreporter.write(" ".join(["pytest"] + ids))
173         terminalreporter.write("\n" * 2)
174 
175 
176 @pytest.fixture(scope="module", autouse=True)
177 def clean_up_envs():
178     yield
179 
180     if "GITHUB_ACTIONS" in os.environ:
181         from mlflow.utils.virtualenv import _get_mlflow_virtualenv_root
182 
183         shutil.rmtree(_get_mlflow_virtualenv_root(), ignore_errors=True)
184         if os.name != "nt":
185             conda_info = json.loads(subprocess.check_output(["conda", "info", "--json"], text=True))
186             root_prefix = conda_info["root_prefix"]
187             for env in conda_info["envs"]:
188                 if env != root_prefix:
189                     shutil.rmtree(env, ignore_errors=True)
190 
191 
192 @pytest.fixture(scope="session", autouse=True)
193 def enable_mlflow_testing():
194     with pytest.MonkeyPatch.context() as mp:
195         mp.setenv(_MLFLOW_TESTING.name, "TRUE")
196         yield
197
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/conftest.py b/conftest.py
--- a/conftest.py
+++ b/conftest.py
@@ -27,6 +27,18 @@
         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 @@
     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 @@
             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 @@
     # 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(

{"golden_diff": "diff --git a/conftest.py b/conftest.py\n--- a/conftest.py\n+++ b/conftest.py\n@@ -27,6 +27,18 @@\n         default=False,\n         help=\"Ignore tests for model flavors.\",\n     )\n+    parser.addoption(\n+        \"--splits\",\n+        default=None,\n+        type=int,\n+        help=\"The number of groups to split tests into.\",\n+    )\n+    parser.addoption(\n+        \"--group\",\n+        default=None,\n+        type=int,\n+        help=\"The group of tests to run.\",\n+    )\n \n \n def pytest_configure(config):\n@@ -36,6 +48,29 @@\n     config.addinivalue_line(\"markers\", \"allow_infer_pip_requirements_fallback\")\n \n \n+@pytest.hookimpl(tryfirst=True)\n+def pytest_cmdline_main(config):\n+    group = config.getoption(\"group\")\n+    splits = config.getoption(\"splits\")\n+\n+    if splits is None and group is None:\n+        return None\n+\n+    if splits and group is None:\n+        raise pytest.UsageError(\"`--group` is required\")\n+\n+    if group and splits is None:\n+        raise pytest.UsageError(\"`--splits` is required\")\n+\n+    if splits < 0:\n+        raise pytest.UsageError(\"`--splits` must be >= 1\")\n+\n+    if group < 1 or group > splits:\n+        raise pytest.UsageError(\"`--group` must be between 1 and {splits}\")\n+\n+    return None\n+\n+\n def pytest_sessionstart(session):\n     if uri := MLFLOW_TRACKING_URI.get():\n         click.echo(\n@@ -147,6 +182,7 @@\n             outcome.force_result(True)\n \n \n+@pytest.hookimpl(trylast=True)\n def pytest_collection_modifyitems(session, config, items):  # pylint: disable=unused-argument\n     # Executing `tests.server.test_prometheus_exporter` after `tests.server.test_handlers`\n     # results in an error because Flask >= 2.2.0 doesn't allow calling setup method such as\n@@ -154,6 +190,10 @@\n     # execute `tests.server.test_prometheus_exporter` first by reordering the test items.\n     items.sort(key=lambda item: item.module.__name__ != \"tests.server.test_prometheus_exporter\")\n \n+    # Select the tests to run based on the group and splits\n+    if (splits := config.getoption(\"--splits\")) and (group := config.getoption(\"--group\")):\n+        items[:] = items[(group - 1) :: splits]\n+\n \n @pytest.hookimpl(hookwrapper=True)\n def pytest_terminal_summary(\n", "issue": "Split tests to run them faster\n## Motivation\r\n\r\nThe number of tests is growing. It takes a long time to run them. We should consider splitting them up for a faster iteration cycle.\r\n\r\n## Example code:\r\n\r\n```diff\r\ndiff --git a/.github/workflows/master.yml b/.github/workflows/master.yml\r\nindex 22e09d7bb..605f46762 100644\r\n--- a/.github/workflows/master.yml\r\n+++ b/.github/workflows/master.yml\r\n@@ -235,6 +235,11 @@ jobs:\r\n     if: github.event_name != 'pull_request' || github.event.pull_request.draft == false\r\n     runs-on: ubuntu-latest\r\n     timeout-minutes: 120\r\n+    strategy:\r\n+      matrix:\r\n+        group: [1, 2, 3]\r\n+        include:\r\n+          - splits: 3\r\n     steps:\r\n       - uses: actions/checkout@v3\r\n         with:\r\n@@ -252,7 +257,7 @@ jobs:\r\n       - name: Run tests\r\n         run: |\r\n           export MLFLOW_HOME=$(pwd)\r\n-          pytest tests/models\r\n+          pytest --splits=${{ matrix.splits }} --group=${{ matrix.group }} tests/models\r\n \r\n   # NOTE: numpy is pinned in this suite due to its heavy reliance on shap, which internally uses\r\n   # references to the now fully deprecated (as of 1.24.x) numpy types (i.e., np.bool).\r\ndiff --git a/conftest.py b/conftest.py\r\nindex 640680cb3..266702d90 100644\r\n--- a/conftest.py\r\n+++ b/conftest.py\r\n@@ -27,6 +27,18 @@ def pytest_addoption(parser):\r\n         default=False,\r\n         help=\"Ignore tests for model flavors.\",\r\n     )\r\n+    parser.addoption(\r\n+        \"--splits\",\r\n+        default=None,\r\n+        type=int,\r\n+        help=\"The number of groups to split tests into.\",\r\n+    )\r\n+    parser.addoption(\r\n+        \"--group\",\r\n+        default=None,\r\n+        type=int,\r\n+        help=\"The group of tests to run.\",\r\n+    )\r\n \r\n \r\n def pytest_configure(config):\r\n@@ -36,6 +48,29 @@ def pytest_configure(config):\r\n     config.addinivalue_line(\"markers\", \"allow_infer_pip_requirements_fallback\")\r\n \r\n \r\n+@pytest.hookimpl(tryfirst=True)\r\n+def pytest_cmdline_main(config):\r\n+    group = config.getoption(\"group\")\r\n+    splits = config.getoption(\"splits\")\r\n+\r\n+    if splits is None and group is None:\r\n+        return None\r\n+\r\n+    if splits and group is None:\r\n+        raise pytest.UsageError(\"`--group` is required\")\r\n+\r\n+    if group and splits is None:\r\n+        raise pytest.UsageError(\"`--splits` is required\")\r\n+\r\n+    if splits < 0:\r\n+        raise pytest.UsageError(\"`--splits` must be >= 1\")\r\n+\r\n+    if group < 1 or group > splits:\r\n+        raise pytest.UsageError(f\"`--group` must be between 1 and {splits}\")\r\n+\r\n+    return None\r\n+\r\n+\r\n def pytest_sessionstart(session):\r\n     if uri := MLFLOW_TRACKING_URI.get():\r\n         click.echo(\r\n@@ -147,6 +182,7 @@ def pytest_ignore_collect(path, config):\r\n             outcome.force_result(True)\r\n \r\n \r\n+@pytest.hookimpl(trylast=True)\r\n def pytest_collection_modifyitems(session, config, items):  # pylint: disable=unused-argument\r\n     # Executing `tests.server.test_prometheus_exporter` after `tests.server.test_handlers`\r\n     # results in an error because Flask >= 2.2.0 doesn't allow calling setup method such as\r\n@@ -154,6 +190,10 @@ def pytest_collection_modifyitems(session, config, items):  # pylint: disable=un\r\n     # execute `tests.server.test_prometheus_exporter` first by reordering the test items.\r\n     items.sort(key=lambda item: item.module.__name__ != \"tests.server.test_prometheus_exporter\")\r\n \r\n+    # Select the tests to run based on the group and splits\r\n+    if (splits := config.getoption(\"--splits\")) and (group := config.getoption(\"--group\")):\r\n+        items[:] = items[(group - 1) :: splits]\r\n+\r\n \r\n @pytest.hookimpl(hookwrapper=True)\r\n def pytest_terminal_summary(\r\n\r\n```\r\n\r\n## Notes\r\n\r\n- 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.\r\n \n", "before_files": [{"content": "import json\nimport os\nimport posixpath\nimport shutil\nimport subprocess\n\nimport click\nimport pytest\n\nfrom mlflow.environment_variables import _MLFLOW_TESTING, MLFLOW_TRACKING_URI\n\n\ndef pytest_addoption(parser):\n    parser.addoption(\n        \"--requires-ssh\",\n        action=\"store_true\",\n        dest=\"requires_ssh\",\n        default=False,\n        help=\"Run tests decorated with 'requires_ssh' annotation. \"\n        \"These tests require keys to be configured locally \"\n        \"for SSH authentication.\",\n    )\n    parser.addoption(\n        \"--ignore-flavors\",\n        action=\"store_true\",\n        dest=\"ignore_flavors\",\n        default=False,\n        help=\"Ignore tests for model flavors.\",\n    )\n\n\ndef pytest_configure(config):\n    # Register markers to suppress `PytestUnknownMarkWarning`\n    config.addinivalue_line(\"markers\", \"requires_ssh\")\n    config.addinivalue_line(\"markers\", \"notrackingurimock\")\n    config.addinivalue_line(\"markers\", \"allow_infer_pip_requirements_fallback\")\n\n\ndef pytest_sessionstart(session):\n    if uri := MLFLOW_TRACKING_URI.get():\n        click.echo(\n            click.style(\n                (\n                    f\"Environment variable {MLFLOW_TRACKING_URI} is set to {uri!r}, \"\n                    \"which may interfere with tests.\"\n                ),\n                fg=\"red\",\n            )\n        )\n\n\ndef pytest_runtest_setup(item):\n    markers = [mark.name for mark in item.iter_markers()]\n    if \"requires_ssh\" in markers and not item.config.getoption(\"--requires-ssh\"):\n        pytest.skip(\"use `--requires-ssh` to run this test\")\n\n\n@pytest.hookimpl(hookwrapper=True)\ndef pytest_report_teststatus(report, config):\n    outcome = yield\n    if report.when == \"call\":\n        try:\n            import psutil\n        except ImportError:\n            return\n\n        (*rest, result) = outcome.get_result()\n        mem = psutil.virtual_memory()\n        mem_used = mem.used / 1024**3\n        mem_total = mem.total / 1024**3\n\n        disk = psutil.disk_usage(\"/\")\n        disk_used = disk.used / 1024**3\n        disk_total = disk.total / 1024**3\n        outcome.force_result(\n            (\n                *rest,\n                (\n                    f\"{result} | \"\n                    f\"MEM {mem_used:.1f}/{mem_total:.1f} GB | \"\n                    f\"DISK {disk_used:.1f}/{disk_total:.1f} GB\"\n                ),\n            )\n        )\n\n\n@pytest.hookimpl(hookwrapper=True)\ndef pytest_ignore_collect(path, config):\n    outcome = yield\n    if not outcome.get_result() and config.getoption(\"ignore_flavors\"):\n        # If not ignored by the default hook and `--ignore-flavors` specified\n\n        # Ignored files and directories must be included in dev/run-python-flavor-tests.sh\n        model_flavors = [\n            # Tests of flavor modules.\n            \"tests/azureml\",\n            \"tests/catboost\",\n            \"tests/diviner\",\n            \"tests/fastai\",\n            \"tests/gluon\",\n            \"tests/h2o\",\n            \"tests/johnsnowlabs\",\n            \"tests/keras\",\n            \"tests/keras_core\",\n            \"tests/langchain\",\n            \"tests/lightgbm\",\n            \"tests/mleap\",\n            \"tests/models\",\n            \"tests/onnx\",\n            \"tests/openai\",\n            \"tests/paddle\",\n            \"tests/pmdarima\",\n            \"tests/prophet\",\n            \"tests/pyfunc\",\n            \"tests/pytorch\",\n            \"tests/sagemaker\",\n            \"tests/sentence_transformers\",\n            \"tests/shap\",\n            \"tests/sklearn\",\n            \"tests/spacy\",\n            \"tests/spark\",\n            \"tests/statsmodels\",\n            \"tests/tensorflow\",\n            \"tests/transformers\",\n            \"tests/xgboost\",\n            # Lazy loading test.\n            \"tests/test_mlflow_lazily_imports_ml_packages.py\",\n            # Tests of utils.\n            \"tests/utils/test_model_utils.py\",\n            # This test is included here because it imports many big libraries like tf, keras, etc.\n            \"tests/tracking/fluent/test_fluent_autolog.py\",\n            # Cross flavor autologging related tests.\n            \"tests/autologging/test_autologging_safety_unit.py\",\n            \"tests/autologging/test_autologging_behaviors_unit.py\",\n            \"tests/autologging/test_autologging_behaviors_integration.py\",\n            \"tests/autologging/test_autologging_utils.py\",\n            \"tests/autologging/test_training_session.py\",\n            # Opt in authentication feature.\n            \"tests/server/auth\",\n            \"tests/gateway\",\n        ]\n\n        relpath = os.path.relpath(str(path))\n        relpath = relpath.replace(os.sep, posixpath.sep)  # for Windows\n\n        if relpath in model_flavors:\n            outcome.force_result(True)\n\n\ndef pytest_collection_modifyitems(session, config, items):  # pylint: disable=unused-argument\n    # Executing `tests.server.test_prometheus_exporter` after `tests.server.test_handlers`\n    # results in an error because Flask >= 2.2.0 doesn't allow calling setup method such as\n    # `before_request` on the application after the first request. To avoid this issue,\n    # execute `tests.server.test_prometheus_exporter` first by reordering the test items.\n    items.sort(key=lambda item: item.module.__name__ != \"tests.server.test_prometheus_exporter\")\n\n\n@pytest.hookimpl(hookwrapper=True)\ndef pytest_terminal_summary(\n    terminalreporter, exitstatus, config\n):  # pylint: disable=unused-argument\n    yield\n    failed_test_reports = terminalreporter.stats.get(\"failed\", [])\n    if failed_test_reports:\n        if len(failed_test_reports) <= 30:\n            terminalreporter.section(\"command to run failed test cases\")\n            ids = [repr(report.nodeid) for report in failed_test_reports]\n        else:\n            terminalreporter.section(\"command to run failed test suites\")\n            # Use dict.fromkeys to preserve the order\n            ids = list(dict.fromkeys(report.fspath for report in failed_test_reports))\n        terminalreporter.write(\" \".join([\"pytest\"] + ids))\n        terminalreporter.write(\"\\n\" * 2)\n\n\n@pytest.fixture(scope=\"module\", autouse=True)\ndef clean_up_envs():\n    yield\n\n    if \"GITHUB_ACTIONS\" in os.environ:\n        from mlflow.utils.virtualenv import _get_mlflow_virtualenv_root\n\n        shutil.rmtree(_get_mlflow_virtualenv_root(), ignore_errors=True)\n        if os.name != \"nt\":\n            conda_info = json.loads(subprocess.check_output([\"conda\", \"info\", \"--json\"], text=True))\n            root_prefix = conda_info[\"root_prefix\"]\n            for env in conda_info[\"envs\"]:\n                if env != root_prefix:\n                    shutil.rmtree(env, ignore_errors=True)\n\n\n@pytest.fixture(scope=\"session\", autouse=True)\ndef enable_mlflow_testing():\n    with pytest.MonkeyPatch.context() as mp:\n        mp.setenv(_MLFLOW_TESTING.name, \"TRUE\")\n        yield\n", "path": "conftest.py"}], "after_files": [{"content": "import json\nimport os\nimport posixpath\nimport shutil\nimport subprocess\n\nimport click\nimport pytest\n\nfrom mlflow.environment_variables import _MLFLOW_TESTING, MLFLOW_TRACKING_URI\n\n\ndef pytest_addoption(parser):\n    parser.addoption(\n        \"--requires-ssh\",\n        action=\"store_true\",\n        dest=\"requires_ssh\",\n        default=False,\n        help=\"Run tests decorated with 'requires_ssh' annotation. \"\n        \"These tests require keys to be configured locally \"\n        \"for SSH authentication.\",\n    )\n    parser.addoption(\n        \"--ignore-flavors\",\n        action=\"store_true\",\n        dest=\"ignore_flavors\",\n        default=False,\n        help=\"Ignore tests for model flavors.\",\n    )\n    parser.addoption(\n        \"--splits\",\n        default=None,\n        type=int,\n        help=\"The number of groups to split tests into.\",\n    )\n    parser.addoption(\n        \"--group\",\n        default=None,\n        type=int,\n        help=\"The group of tests to run.\",\n    )\n\n\ndef pytest_configure(config):\n    # Register markers to suppress `PytestUnknownMarkWarning`\n    config.addinivalue_line(\"markers\", \"requires_ssh\")\n    config.addinivalue_line(\"markers\", \"notrackingurimock\")\n    config.addinivalue_line(\"markers\", \"allow_infer_pip_requirements_fallback\")\n\n\n@pytest.hookimpl(tryfirst=True)\ndef pytest_cmdline_main(config):\n    group = config.getoption(\"group\")\n    splits = config.getoption(\"splits\")\n\n    if splits is None and group is None:\n        return None\n\n    if splits and group is None:\n        raise pytest.UsageError(\"`--group` is required\")\n\n    if group and splits is None:\n        raise pytest.UsageError(\"`--splits` is required\")\n\n    if splits < 0:\n        raise pytest.UsageError(\"`--splits` must be >= 1\")\n\n    if group < 1 or group > splits:\n        raise pytest.UsageError(\"`--group` must be between 1 and {splits}\")\n\n    return None\n\n\ndef pytest_sessionstart(session):\n    if uri := MLFLOW_TRACKING_URI.get():\n        click.echo(\n            click.style(\n                (\n                    f\"Environment variable {MLFLOW_TRACKING_URI} is set to {uri!r}, \"\n                    \"which may interfere with tests.\"\n                ),\n                fg=\"red\",\n            )\n        )\n\n\ndef pytest_runtest_setup(item):\n    markers = [mark.name for mark in item.iter_markers()]\n    if \"requires_ssh\" in markers and not item.config.getoption(\"--requires-ssh\"):\n        pytest.skip(\"use `--requires-ssh` to run this test\")\n\n\n@pytest.hookimpl(hookwrapper=True)\ndef pytest_report_teststatus(report, config):\n    outcome = yield\n    if report.when == \"call\":\n        try:\n            import psutil\n        except ImportError:\n            return\n\n        (*rest, result) = outcome.get_result()\n        mem = psutil.virtual_memory()\n        mem_used = mem.used / 1024**3\n        mem_total = mem.total / 1024**3\n\n        disk = psutil.disk_usage(\"/\")\n        disk_used = disk.used / 1024**3\n        disk_total = disk.total / 1024**3\n        outcome.force_result(\n            (\n                *rest,\n                (\n                    f\"{result} | \"\n                    f\"MEM {mem_used:.1f}/{mem_total:.1f} GB | \"\n                    f\"DISK {disk_used:.1f}/{disk_total:.1f} GB\"\n                ),\n            )\n        )\n\n\n@pytest.hookimpl(hookwrapper=True)\ndef pytest_ignore_collect(path, config):\n    outcome = yield\n    if not outcome.get_result() and config.getoption(\"ignore_flavors\"):\n        # If not ignored by the default hook and `--ignore-flavors` specified\n\n        # Ignored files and directories must be included in dev/run-python-flavor-tests.sh\n        model_flavors = [\n            # Tests of flavor modules.\n            \"tests/azureml\",\n            \"tests/catboost\",\n            \"tests/diviner\",\n            \"tests/fastai\",\n            \"tests/gluon\",\n            \"tests/h2o\",\n            \"tests/johnsnowlabs\",\n            \"tests/keras\",\n            \"tests/keras_core\",\n            \"tests/langchain\",\n            \"tests/lightgbm\",\n            \"tests/mleap\",\n            \"tests/models\",\n            \"tests/onnx\",\n            \"tests/openai\",\n            \"tests/paddle\",\n            \"tests/pmdarima\",\n            \"tests/prophet\",\n            \"tests/pyfunc\",\n            \"tests/pytorch\",\n            \"tests/sagemaker\",\n            \"tests/sentence_transformers\",\n            \"tests/shap\",\n            \"tests/sklearn\",\n            \"tests/spacy\",\n            \"tests/spark\",\n            \"tests/statsmodels\",\n            \"tests/tensorflow\",\n            \"tests/transformers\",\n            \"tests/xgboost\",\n            # Lazy loading test.\n            \"tests/test_mlflow_lazily_imports_ml_packages.py\",\n            # Tests of utils.\n            \"tests/utils/test_model_utils.py\",\n            # This test is included here because it imports many big libraries like tf, keras, etc.\n            \"tests/tracking/fluent/test_fluent_autolog.py\",\n            # Cross flavor autologging related tests.\n            \"tests/autologging/test_autologging_safety_unit.py\",\n            \"tests/autologging/test_autologging_behaviors_unit.py\",\n            \"tests/autologging/test_autologging_behaviors_integration.py\",\n            \"tests/autologging/test_autologging_utils.py\",\n            \"tests/autologging/test_training_session.py\",\n            # Opt in authentication feature.\n            \"tests/server/auth\",\n            \"tests/gateway\",\n        ]\n\n        relpath = os.path.relpath(str(path))\n        relpath = relpath.replace(os.sep, posixpath.sep)  # for Windows\n\n        if relpath in model_flavors:\n            outcome.force_result(True)\n\n\n@pytest.hookimpl(trylast=True)\ndef pytest_collection_modifyitems(session, config, items):  # pylint: disable=unused-argument\n    # Executing `tests.server.test_prometheus_exporter` after `tests.server.test_handlers`\n    # results in an error because Flask >= 2.2.0 doesn't allow calling setup method such as\n    # `before_request` on the application after the first request. To avoid this issue,\n    # execute `tests.server.test_prometheus_exporter` first by reordering the test items.\n    items.sort(key=lambda item: item.module.__name__ != \"tests.server.test_prometheus_exporter\")\n\n    # Select the tests to run based on the group and splits\n    if (splits := config.getoption(\"--splits\")) and (group := config.getoption(\"--group\")):\n        items[:] = items[(group - 1) :: splits]\n\n\n@pytest.hookimpl(hookwrapper=True)\ndef pytest_terminal_summary(\n    terminalreporter, exitstatus, config\n):  # pylint: disable=unused-argument\n    yield\n    failed_test_reports = terminalreporter.stats.get(\"failed\", [])\n    if failed_test_reports:\n        if len(failed_test_reports) <= 30:\n            terminalreporter.section(\"command to run failed test cases\")\n            ids = [repr(report.nodeid) for report in failed_test_reports]\n        else:\n            terminalreporter.section(\"command to run failed test suites\")\n            # Use dict.fromkeys to preserve the order\n            ids = list(dict.fromkeys(report.fspath for report in failed_test_reports))\n        terminalreporter.write(\" \".join([\"pytest\"] + ids))\n        terminalreporter.write(\"\\n\" * 2)\n\n\n@pytest.fixture(scope=\"module\", autouse=True)\ndef clean_up_envs():\n    yield\n\n    if \"GITHUB_ACTIONS\" in os.environ:\n        from mlflow.utils.virtualenv import _get_mlflow_virtualenv_root\n\n        shutil.rmtree(_get_mlflow_virtualenv_root(), ignore_errors=True)\n        if os.name != \"nt\":\n            conda_info = json.loads(subprocess.check_output([\"conda\", \"info\", \"--json\"], text=True))\n            root_prefix = conda_info[\"root_prefix\"]\n            for env in conda_info[\"envs\"]:\n                if env != root_prefix:\n                    shutil.rmtree(env, ignore_errors=True)\n\n\n@pytest.fixture(scope=\"session\", autouse=True)\ndef enable_mlflow_testing():\n    with pytest.MonkeyPatch.context() as mp:\n        mp.setenv(_MLFLOW_TESTING.name, \"TRUE\")\n        yield\n", "path": "conftest.py"}]}