Repository: pantsbuild/pants

Issue:
Allow for automatic use of pytest-xdist via `Process.concurrency_available`
With the [`Process.concurrency_available`](https://github.com/pantsbuild/pants/blob/511691657d00827e5e44ab5118688b99c221d744/src/rust/engine/process_execution/src/lib.rs#L378-L387) facility, a process can hint at how much concurrency Pants should attempt to allocate for it. This would allow for safe automatic usage of `pytest-xdist` (behind a flag, and probably optionally disabled via a per-test `Field`), by passing `pytest -n {pants_concurrency}`.


Discussion:
This relates to #14886 as well: assuming that you had enough tests to saturate your cores initially, you might end up with a large single test file still running after the others had finished. That issue discusses dynamically scaling up for a straggler process.
I would be interested in this as well. Parallelizing by test rather than test file currently gives us much better performance.  Yes, we could do the manual effort of splitting out long running files, but it'd be less work to have the option of setting pants to work like this.
+1, also interested. 

While it _is_ a best practice to separate modularize tests such that any test file in itself is not too large, different pieces of code with varying degrees of complexity will have tests that vary in performance.

It feels that the Pants tests execution layer should be robust enough to handle this without requiring the developer to modify the way they write tests.
I think that to do this, you would:
1. Add a boolean option to the `pytest` subsystem to control whether it is enabled, [similar to this one](https://github.com/pantsbuild/pants/blob/258ab2602a05b1450540d8d3d89f51407aa416fa/src/python/pants/backend/python/subsystems/pytest.py#L137-L148). It will likely need to be disabled by default, since concurrency can change test behavior.
2. Add `pytest-xdist` to the [`default_extra_requirements` of `pytest`](https://github.com/pantsbuild/pants/blob/258ab2602a05b1450540d8d3d89f51407aa416fa/src/python/pants/backend/python/subsystems/pytest.py#L75).
3. When the flag is enabled, pass an (initially hardcoded... i.e. to `4`) `VenvPexProcess(.., concurrency_available=4)` value for [the pytest process](https://github.com/pantsbuild/pants/blob/fca8babc1543444067abec9efe753d72a1ee27d2/src/python/pants/backend/python/goals/pytest_runner.py#L328-L350), and add an `-n {concurrency_available}` argument for `pytest` (similar to [what `black` does here](https://github.com/pantsbuild/pants/blob/511691657d00827e5e44ab5118688b99c221d744/src/python/pants/backend/python/lint/black/rules.py#L86-L87)).
4. Once that's working, add a `@_rule_helper` in the same file that will use a heuristic to choose how many threads to use for a particular file (i.e., what to pass to `concurrency_available`). This could probably be as simple as a regex that matches `def test` and a few other common test-method-name patterns.
    * This will need to take _only_ the test sources, which will be similar to https://github.com/pantsbuild/pants/blob/fca8babc1543444067abec9efe753d72a1ee27d2/src/python/pants/backend/python/goals/pytest_runner.py#L202-L204, but for only the `TestSetupRequest` target which the method is running for.
    * It will then need to use `await Get(DigestContents, Digest, source_files.snapshot.digest)` to load the content of the test file.
    * Finally, match a regex/etc against the content of the source files in the `DigestContents`.
5. Add an `IntField` to the `python_tests` target ([similar to this field](https://github.com/pantsbuild/pants/blob/fca8babc1543444067abec9efe753d72a1ee27d2/src/python/pants/backend/python/target_types.py#L805)), which will allow for overriding the concurrency detected for a particular test.

To test this, you'll probably want to [add a test here](https://github.com/pantsbuild/pants/blob/fca8babc1543444067abec9efe753d72a1ee27d2/src/python/pants/backend/python/goals/pytest_runner_integration_test.py#L251-L270), which will really just be a "coverage" test: you won't be able to test that concurrency is actually used, but you can test that nothing fails with the option enabled.

Resolving patch:
diff --git a/src/python/pants/backend/python/goals/pytest_runner.py b/src/python/pants/backend/python/goals/pytest_runner.py
--- a/src/python/pants/backend/python/goals/pytest_runner.py
+++ b/src/python/pants/backend/python/goals/pytest_runner.py
@@ -2,6 +2,7 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 import logging
+import re
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Optional, Tuple
@@ -44,6 +45,7 @@
     EMPTY_DIGEST,
     CreateDigest,
     Digest,
+    DigestContents,
     DigestSubset,
     Directory,
     MergeDigests,
@@ -164,6 +166,13 @@ class TestSetup:
     __test__ = False
 
 
+_TEST_PATTERN = re.compile(b"def\\s+test_")
+
+
+def _count_pytest_tests(contents: DigestContents) -> int:
+    return sum(len(_TEST_PATTERN.findall(file.content)) for file in contents)
+
+
 @rule(level=LogLevel.DEBUG)
 async def setup_pytest_for_target(
     request: TestSetupRequest,
@@ -325,12 +334,22 @@ async def setup_pytest_for_target(
     cache_scope = (
         ProcessCacheScope.PER_SESSION if test_subsystem.force else ProcessCacheScope.SUCCESSFUL
     )
+
+    xdist_concurrency = 0
+    if pytest.xdist_enabled and not request.is_debug:
+        concurrency = request.field_set.xdist_concurrency.value
+        if concurrency is None:
+            contents = await Get(DigestContents, Digest, field_set_source_files.snapshot.digest)
+            concurrency = _count_pytest_tests(contents)
+        xdist_concurrency = concurrency
+
     process = await Get(
         Process,
         VenvPexProcess(
             pytest_runner_pex,
             argv=(
                 *(("-c", pytest.config) if pytest.config else ()),
+                *(("-n", "{pants_concurrency}") if xdist_concurrency else ()),
                 *request.prepend_argv,
                 *pytest.args,
                 *coverage_args,
@@ -344,6 +363,7 @@ async def setup_pytest_for_target(
                 test_subsystem, pytest
             ),
             execution_slot_variable=pytest.execution_slot_var,
+            concurrency_available=xdist_concurrency,
             description=f"Run Pytest for {request.field_set.address}",
             level=LogLevel.DEBUG,
             cache_scope=cache_scope,
diff --git a/src/python/pants/backend/python/subsystems/pytest.py b/src/python/pants/backend/python/subsystems/pytest.py
--- a/src/python/pants/backend/python/subsystems/pytest.py
+++ b/src/python/pants/backend/python/subsystems/pytest.py
@@ -25,6 +25,7 @@
     PythonTestsExtraEnvVarsField,
     PythonTestSourceField,
     PythonTestsTimeoutField,
+    PythonTestsXdistConcurrencyField,
     SkipPythonTestsField,
 )
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
@@ -56,6 +57,7 @@ class PythonTestFieldSet(TestFieldSet):
     timeout: PythonTestsTimeoutField
     runtime_package_dependencies: RuntimePackageDependenciesField
     extra_env_vars: PythonTestsExtraEnvVarsField
+    xdist_concurrency: PythonTestsXdistConcurrencyField
 
     @classmethod
     def opt_out(cls, tgt: Target) -> bool:
@@ -75,7 +77,7 @@ class PyTest(PythonToolBase):
     # TODO: Once this issue is fixed, loosen this to allow the version to float above the bad ones.
     #  E.g., as default_version = "pytest>=7,<8,!=7.1.0,!=7.1.1"
     default_version = "pytest==7.0.1"
-    default_extra_requirements = ["pytest-cov>=2.12,!=2.12.1,<3.1"]
+    default_extra_requirements = ["pytest-cov>=2.12,!=2.12.1,<3.1", "pytest-xdist>=2.5,<3"]
 
     default_main = ConsoleScript("pytest")
 
@@ -164,6 +166,20 @@ class PyTest(PythonToolBase):
             """
         ),
     )
+    xdist_enabled = BoolOption(
+        default=False,
+        advanced=False,
+        help=softwrap(
+            """
+            If true, Pants will use `pytest-xdist` (https://pytest-xdist.readthedocs.io/en/latest/)
+            to parallelize tests within each `python_test` target.
+
+            NOTE: Enabling `pytest-xdist` can cause high-level scoped fixtures (for example `session`)
+            to execute more than once. See the `pytest-xdist` docs for more info:
+            https://pypi.org/project/pytest-xdist/#making-session-scoped-fixtures-execute-only-once
+            """
+        ),
+    )
 
     export = ExportToolOption()
 
diff --git a/src/python/pants/backend/python/target_types.py b/src/python/pants/backend/python/target_types.py
--- a/src/python/pants/backend/python/target_types.py
+++ b/src/python/pants/backend/python/target_types.py
@@ -858,6 +858,25 @@ class PythonTestsExtraEnvVarsField(TestExtraEnvVarsField):
     pass
 
 
+class PythonTestsXdistConcurrencyField(IntField):
+    alias = "xdist_concurrency"
+    help = softwrap(
+        """
+        Maximum number of CPUs to allocate to run each test file belonging to this target.
+
+        Tests are spread across multiple CPUs using `pytest-xdist`
+        (https://pytest-xdist.readthedocs.io/en/latest/index.html).
+        Use of `pytest-xdist` must be enabled using the `[pytest].xdist_enabled` option for
+        this field to have an effect.
+
+        If `pytest-xdist` is enabled and this field is unset, Pants will attempt to derive
+        the concurrency for test sources by counting the number of tests in each file.
+
+        Set this field to `0` to explicitly disable use of `pytest-xdist` for a target.
+        """
+    )
+
+
 class SkipPythonTestsField(BoolField):
     alias = "skip_tests"
     default = False
@@ -869,6 +888,7 @@ class SkipPythonTestsField(BoolField):
     PythonResolveField,
     PythonRunGoalUseSandboxField,
     PythonTestsTimeoutField,
+    PythonTestsXdistConcurrencyField,
     RuntimePackageDependenciesField,
     PythonTestsExtraEnvVarsField,
     InterpreterConstraintsField,


Test patch:
diff --git a/src/python/pants/backend/python/goals/pytest_runner_integration_test.py b/src/python/pants/backend/python/goals/pytest_runner_integration_test.py
--- a/src/python/pants/backend/python/goals/pytest_runner_integration_test.py
+++ b/src/python/pants/backend/python/goals/pytest_runner_integration_test.py
@@ -39,6 +39,7 @@
 from pants.core.util_rules import config_files, distdir
 from pants.engine.addresses import Address
 from pants.engine.fs import CreateDigest, Digest, DigestContents, FileContent
+from pants.engine.process import InteractiveProcessResult
 from pants.engine.rules import Get, rule
 from pants.engine.target import Target
 from pants.engine.unions import UnionRule
@@ -94,19 +95,28 @@ def test():
 )
 
 
-def run_pytest(
+def _configure_pytest_runner(
     rule_runner: RuleRunner,
-    test_target: Target,
     *,
     extra_args: list[str] | None = None,
     env: dict[str, str] | None = None,
-) -> TestResult:
+) -> None:
     args = [
         "--backend-packages=pants.backend.python",
         f"--source-root-patterns={SOURCE_ROOT}",
         *(extra_args or ()),
     ]
     rule_runner.set_options(args, env=env, env_inherit={"PATH", "PYENV_ROOT", "HOME"})
+
+
+def run_pytest(
+    rule_runner: RuleRunner,
+    test_target: Target,
+    *,
+    extra_args: list[str] | None = None,
+    env: dict[str, str] | None = None,
+) -> TestResult:
+    _configure_pytest_runner(rule_runner, extra_args=extra_args, env=env)
     inputs = [PythonTestFieldSet.create(test_target)]
     test_result = rule_runner.request(TestResult, inputs)
     debug_request = rule_runner.request(TestDebugRequest, inputs)
@@ -117,6 +127,32 @@ def run_pytest(
     return test_result
 
 
+def run_pytest_noninteractive(
+    rule_runner: RuleRunner,
+    test_target: Target,
+    *,
+    extra_args: list[str] | None = None,
+    env: dict[str, str] | None = None,
+) -> TestResult:
+    _configure_pytest_runner(rule_runner, extra_args=extra_args, env=env)
+    inputs = [PythonTestFieldSet.create(test_target)]
+    return rule_runner.request(TestResult, inputs)
+
+
+def run_pytest_interactive(
+    rule_runner: RuleRunner,
+    test_target: Target,
+    *,
+    extra_args: list[str] | None = None,
+    env: dict[str, str] | None = None,
+) -> InteractiveProcessResult:
+    _configure_pytest_runner(rule_runner, extra_args=extra_args, env=env)
+    inputs = [PythonTestFieldSet.create(test_target)]
+    debug_request = rule_runner.request(TestDebugRequest, inputs)
+    with mock_console(rule_runner.options_bootstrapper):
+        return rule_runner.run_interactive_process(debug_request.process)
+
+
 @pytest.mark.platform_specific_behavior
 @pytest.mark.parametrize(
     "major_minor_interpreter",
@@ -231,7 +267,11 @@ def test() -> None:
             ),
         }
     )
-    extra_args = ["--pytest-version=pytest>=4.6.6,<4.7", "--pytest-lockfile=<none>"]
+    extra_args = [
+        "--pytest-version=pytest>=4.6.6,<4.7",
+        "--pytest-extra-requirements=[]",
+        "--pytest-lockfile=<none>",
+    ]
 
     py2_tgt = rule_runner.get_target(
         Address(PACKAGE, target_name="py2", relative_file_path="tests.py")
@@ -270,6 +310,72 @@ def test_ignore_me():
     assert "collected 2 items / 1 deselected / 1 selected" in result.stdout
 
 
+def test_xdist_enabled_noninteractive(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            f"{PACKAGE}/tests.py": dedent(
+                """\
+                import os
+
+                def test_worker_id_set():
+                  assert "PYTEST_XDIST_WORKER" in os.environ
+
+                def test_worker_count_set():
+                  assert "PYTEST_XDIST_WORKER_COUNT" in os.environ
+                """
+            ),
+            f"{PACKAGE}/BUILD": "python_tests(xdist_concurrency=2)",
+        }
+    )
+    tgt = rule_runner.get_target(Address(PACKAGE, relative_file_path="tests.py"))
+    result = run_pytest_noninteractive(rule_runner, tgt, extra_args=["--pytest-xdist-enabled"])
+    assert result.exit_code == 0
+
+
+def test_xdist_enabled_but_disabled_for_target(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            f"{PACKAGE}/tests.py": dedent(
+                """\
+                import os
+
+                def test_worker_id_not_set():
+                  assert "PYTEST_XDIST_WORKER" not in os.environ
+
+                def test_worker_count_not_set():
+                  assert "PYTEST_XDIST_WORKER_COUNT" not in os.environ
+                """
+            ),
+            f"{PACKAGE}/BUILD": "python_tests(xdist_concurrency=0)",
+        }
+    )
+    tgt = rule_runner.get_target(Address(PACKAGE, relative_file_path="tests.py"))
+    result = run_pytest_noninteractive(rule_runner, tgt, extra_args=["--pytest-xdist-enabled"])
+    assert result.exit_code == 0
+
+
+def test_xdist_enabled_interactive(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            f"{PACKAGE}/tests.py": dedent(
+                """\
+                import os
+
+                def test_worker_id_not_set():
+                  assert "PYTEST_XDIST_WORKER" not in os.environ
+
+                def test_worker_count_not_set():
+                  assert "PYTEST_XDIST_WORKER_COUNT" not in os.environ
+                """
+            ),
+            f"{PACKAGE}/BUILD": "python_tests(xdist_concurrency=2)",
+        }
+    )
+    tgt = rule_runner.get_target(Address(PACKAGE, relative_file_path="tests.py"))
+    result = run_pytest_interactive(rule_runner, tgt, extra_args=["--pytest-xdist-enabled"])
+    assert result.exit_code == 0
+
+
 @pytest.mark.parametrize(
     "config_path,extra_args",
     (["pytest.ini", []], ["custom_config.ini", ["--pytest-config=custom_config.ini"]]),
diff --git a/src/python/pants/backend/python/goals/pytest_runner_test.py b/src/python/pants/backend/python/goals/pytest_runner_test.py
new file mode 100644
--- /dev/null
+++ b/src/python/pants/backend/python/goals/pytest_runner_test.py
@@ -0,0 +1,56 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from pants.backend.python.goals.pytest_runner import _count_pytest_tests
+from pants.engine.fs import DigestContents, FileContent
+
+EXAMPLE_TEST1 = b"""
+def test_foo():
+    pass
+
+def test_bar():
+    pass
+"""
+
+EXAMPLE_TEST2 = b"""
+class TestStuff(TestCase):
+    def test_baz():
+        pass
+
+    def testHelper():
+        pass
+"""
+
+
+def test_count_pytest_tests_empty() -> None:
+    digest_contents = DigestContents([FileContent(path="tests/test_empty.py", content=b"")])
+    test_count = _count_pytest_tests(digest_contents)
+    assert test_count == 0
+
+
+def test_count_pytest_tests_methods() -> None:
+    digest_contents = DigestContents(
+        [FileContent(path="tests/test_example1.py", content=EXAMPLE_TEST1)]
+    )
+    test_count = _count_pytest_tests(digest_contents)
+    assert test_count == 2
+
+
+def test_count_pytest_tests_in_class() -> None:
+    digest_contents = DigestContents(
+        [FileContent(path="tests/test_example1.py", content=EXAMPLE_TEST2)]
+    )
+    test_count = _count_pytest_tests(digest_contents)
+    assert test_count == 1
+
+
+def test_count_pytest_tests_multiple() -> None:
+    digest_contents = DigestContents(
+        [
+            FileContent(path="tests/test_empty.py", content=b""),
+            FileContent(path="tests/test_example1.py", content=EXAMPLE_TEST1),
+            FileContent(path="tests/test_example2.py", content=EXAMPLE_TEST2),
+        ]
+    )
+    test_count = _count_pytest_tests(digest_contents)
+    assert test_count == 3