Use worksteal distribution strategy with pytest-xdist

pyproject.toml

@@ -272,4 +272,3 @@ ignore_missing_imports = true
 [[tool.mypy.overrides]]
 module = "dev_tools.qualtran_dev_tools.bloq_report_card"
 ignore_errors = true
-

qualtran/conftest.py

@@ -247,7 +247,40 @@ def get_available_cpu_count() -> int:
     return cpus
 
 
-def pytest_configure(config):
+def _config_set_xdist_worksteal(config) -> None:
+    """Sets `--dist worksteal` as the default distribution mode if not
+    explicitly overridden by the user."""
+    num_workers = config.getoption("numprocesses")
+    if num_workers in (None, 0, 1, "0", "1"):
+        return
+
+    if (
+        hasattr(config, "option")
+        and hasattr(config.option, "dist")
+        and config.getoption("dist") in (None, "no", "load")
+    ):
+        # Check if the user explicitly provided a distribution option. If they
+        # did, we shouldn't overwrite it. Since dist defaults to "load" when
+        # -n is set, we check if --dist is explicitly passed.
+        args = []
+        if hasattr(config, "invocation_params") and config.invocation_params is not None:
+            args.extend(config.invocation_params.args)
+        try:
+            addopts = config.getini("addopts")
+            if isinstance(addopts, list):
+                args.extend(addopts)
+        except (ValueError, AttributeError):
+            pass
+
+        for arg in args:
+            if arg.startswith("--dist") or arg == "-d":
+                break
+        else:
+            # Checked all args and didn't find --dist or -d.
+            config.option.dist = "worksteal"
+
+
+def _config_set_thread_limits(config) -> None:
     """Limit number of threads to prevent oversubscription with pytest-xdist.
 
     This only influences parallelism in some core numerical libraries used in
@@ -257,16 +290,7 @@ def pytest_configure(config):
     numerical operations in some tests spawn as many parallel threads as CPUs,
     overwhelming host resources when pytest runs the tests in parallel.
     """
-    # Only run in the controlling process, before workers are started.
-    if hasattr(config, "workerinput"):
-        return
-
-    # Get the -n value from the Pytest invocation.
-    try:
-        num_workers = config.getoption("numprocesses")
-    except ValueError:
-        # pytest-xdist is not being used.
-        return
+    num_workers = config.getoption("numprocesses")
     if num_workers is None or not isinstance(num_workers, (int, str)):
         num_workers = 1
 
@@ -292,3 +316,19 @@ def pytest_configure(config):
         ]
         for var in env_vars:
             os.environ[var] = str(limit)
+
+
+def pytest_configure(config):
+    """Configure pytest environment settings, especially for pytest-xdist."""
+
+    # Only run in the controlling process, before workers are started.
+    if hasattr(config, "workerinput"):
+        return
+    try:
+        config.getoption("numprocesses")
+    except ValueError:
+        # pytest-xdist is not being used.
+        return
+
+    _config_set_thread_limits(config)
+    _config_set_xdist_worksteal(config)

qualtran/conftest_test.py

@@ -249,3 +249,45 @@ def test_pytest_configure_cpus_non_positive():
         ]
         for var in env_vars:
             assert var not in os.environ
+
+
+def test_pytest_configure_sets_dist_worksteal():
+    config = mock.MagicMock(spec=pytest.Config)
+    del config.workerinput
+    config.getoption.side_effect = lambda name: '4' if name == 'numprocesses' else 'load'
+    config.option = mock.MagicMock()
+    config.option.dist = 'load'
+    config.invocation_params = mock.MagicMock()
+    config.invocation_params.args = ['-n', '4']
+
+    with mock.patch('qualtran.conftest.get_available_cpu_count', return_value=8):
+        pytest_configure(config)
+        assert config.option.dist == 'worksteal'
+
+
+def test_pytest_configure_preserves_user_dist():
+    config = mock.MagicMock(spec=pytest.Config)
+    del config.workerinput
+    config.getoption.side_effect = lambda name: '4' if name == 'numprocesses' else 'loadscope'
+    config.option = mock.MagicMock()
+    config.option.dist = 'loadscope'
+    config.invocation_params = mock.MagicMock()
+    config.invocation_params.args = ['-n', '4', '--dist', 'loadscope']
+
+    with mock.patch('qualtran.conftest.get_available_cpu_count', return_value=8):
+        pytest_configure(config)
+        assert config.option.dist == 'loadscope'
+
+
+def test_pytest_configure_preserves_user_dist_equals():
+    config = mock.MagicMock(spec=pytest.Config)
+    del config.workerinput
+    config.getoption.side_effect = lambda name: '4' if name == 'numprocesses' else 'each'
+    config.option = mock.MagicMock()
+    config.option.dist = 'each'
+    config.invocation_params = mock.MagicMock()
+    config.invocation_params.args = ['-n', '4', '--dist=each']
+
+    with mock.patch('qualtran.conftest.get_available_cpu_count', return_value=8):
+        pytest_configure(config)
+        assert config.option.dist == 'each'