feat(consume): create pre_hash subgroups w/max size for better xdist balancing

danceratopz · danceratopz · commit 7f72f591669f · 2025-07-15T12:55:27.000+02:00
diff --git a/src/cli/pytest_commands/consume.py b/src/cli/pytest_commands/consume.py
@@ -124,7 +124,7 @@ def engine() -> None:
 def enginex(enginex_fcu_frequency: int, pytest_args: List[str], **_kwargs) -> None:
     """Client consumes Engine X Fixtures via the Engine API."""
     command_name = "enginex"
-    command_paths = get_command_paths(command_name, is_hive=True)
+    static_test_paths = get_static_test_paths(command_name, is_hive=True)
 
     # Validate the frequency parameter
     if enginex_fcu_frequency < 0:
@@ -133,7 +133,9 @@ def enginex(enginex_fcu_frequency: int, pytest_args: List[str], **_kwargs) -> No
     # Add the FCU frequency to pytest args as a custom config option
     pytest_args_with_fcu = [f"--enginex-fcu-frequency={enginex_fcu_frequency}"] + list(pytest_args)
 
-    consume_cmd = ConsumeCommand(command_paths, is_hive=True, command_name=command_name)
+    consume_cmd = create_consume_command(
+        static_test_paths=static_test_paths, is_hive=True, command_name=command_name
+    )
     consume_cmd.execute(pytest_args_with_fcu)
 
 
diff --git a/src/cli/pytest_commands/processors.py b/src/cli/pytest_commands/processors.py
@@ -110,7 +110,7 @@ def process_args(self, args: List[str]) -> List[str]:
         if self.command_name == "engine":
             modified_args.extend(["-p", "pytest_plugins.consume.simulators.engine.conftest"])
         elif self.command_name == "enginex":
-            modified_args.extend(["-p", "pytest_plugins.consume.simulators.enginex.conftest"])
+            modified_args.extend(["-p", "pytest_plugins.consume.simulators.engine_reorg.conftest"])
         elif self.command_name == "rlp":
             modified_args.extend(["-p", "pytest_plugins.consume.simulators.rlp.conftest"])
         else:
diff --git a/src/pytest_plugins/consume/consume.py b/src/pytest_plugins/consume/consume.py
@@ -1,12 +1,20 @@
-"""A pytest plugin providing common functionality for consuming test fixtures."""
+"""
+A pytest plugin providing common functionality for consuming test fixtures.
 
+Features:
+- Downloads and caches test fixtures from various sources (local, URL, release).
+- Manages test case generation from fixture files.
+- Provides xdist load balancing for large pre-allocation groups (enginex simulator).
+"""
+
+import logging
 import re
 import sys
 import tarfile
 from dataclasses import dataclass
 from io import BytesIO
 from pathlib import Path
-from typing import List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 from urllib.parse import urlparse
 
 import platformdirs
@@ -15,18 +23,125 @@
 import rich
 
 from cli.gen_index import generate_fixtures_index
-from ethereum_test_fixtures import BaseFixture
+from ethereum_test_fixtures import BaseFixture, BlockchainEngineXFixture
 from ethereum_test_fixtures.consume import IndexFile, TestCases
 from ethereum_test_forks import get_forks, get_relative_fork_markers, get_transition_forks
 from ethereum_test_tools.utility.versioning import get_current_commit_hash_or_tag
 
 from .releases import ReleaseTag, get_release_page_url, get_release_url, is_release_url, is_url
 
+logger = logging.getLogger(__name__)
+
 CACHED_DOWNLOADS_DIRECTORY = (
     Path(platformdirs.user_cache_dir("ethereum-execution-spec-tests")) / "cached_downloads"
 )
 
 
+class XDistGroupMapper:
+    """
+    Maps test cases to xdist groups, splitting large pre-allocation groups into sub-groups.
+
+    This class helps improve load balancing when using pytest-xdist with --dist=loadgroup
+    by breaking up large pre-allocation groups (e.g., 1000+ tests) into smaller virtual
+    sub-groups while maintaining the constraint that tests from the same pre-allocation
+    group must run on the same worker.
+    """
+
+    def __init__(self, max_group_size: int = 400):
+        """Initialize the mapper with a maximum group size."""
+        self.max_group_size = max_group_size
+        self.group_sizes: Dict[str, int] = {}
+        self.test_to_subgroup: Dict[str, int] = {}
+        self._built = False
+
+    def build_mapping(self, test_cases: TestCases) -> None:
+        """
+        Build the mapping of test cases to sub-groups.
+
+        This analyzes all test cases and determines which pre-allocation groups
+        need to be split into sub-groups based on the max_group_size.
+        """
+        if self._built:
+            return
+
+        # Count tests per pre-allocation group
+        for test_case in test_cases:
+            if hasattr(test_case, "pre_hash") and test_case.pre_hash:
+                pre_hash = test_case.pre_hash
+                self.group_sizes[pre_hash] = self.group_sizes.get(pre_hash, 0) + 1
+
+        # Assign sub-groups for large groups
+        group_counters: Dict[str, int] = {}
+        for test_case in test_cases:
+            if hasattr(test_case, "pre_hash") and test_case.pre_hash:
+                pre_hash = test_case.pre_hash
+                group_size = self.group_sizes[pre_hash]
+
+                if group_size <= self.max_group_size:
+                    # Small group, no sub-group needed
+                    self.test_to_subgroup[test_case.id] = 0
+                else:
+                    # Large group, assign to sub-group using round-robin
+                    counter = group_counters.get(pre_hash, 0)
+                    sub_group = counter // self.max_group_size
+                    self.test_to_subgroup[test_case.id] = sub_group
+                    group_counters[pre_hash] = counter + 1
+
+        self._built = True
+
+        # Log summary of large groups
+        large_groups = [
+            (pre_hash, size)
+            for pre_hash, size in self.group_sizes.items()
+            if size > self.max_group_size
+        ]
+        if large_groups:
+            logger.info(
+                f"Found {len(large_groups)} pre-allocation groups larger than "
+                f"{self.max_group_size} tests that will be split for better load balancing"
+            )
+
+    def get_xdist_group_name(self, test_case) -> str:
+        """
+        Get the xdist group name for a test case.
+
+        For small groups, returns the pre_hash as-is.
+        For large groups, returns "{pre_hash}:{sub_group_index}".
+        """
+        if not hasattr(test_case, "pre_hash") or not test_case.pre_hash:
+            # No pre_hash, use test ID as fallback
+            return test_case.id
+
+        pre_hash = test_case.pre_hash
+        group_size = self.group_sizes.get(pre_hash, 0)
+
+        if group_size <= self.max_group_size:
+            # Small group, use pre_hash as-is
+            return pre_hash
+
+        # Large group, include sub-group index
+        sub_group = self.test_to_subgroup.get(test_case.id, 0)
+        return f"{pre_hash}:{sub_group}"
+
+    def get_split_statistics(self) -> Dict[str, Dict[str, int]]:
+        """
+        Get statistics about how groups were split.
+
+        Returns a dict with information about each pre-allocation group
+        and how many sub-groups it was split into.
+        """
+        stats = {}
+        for pre_hash, size in self.group_sizes.items():
+            if size > self.max_group_size:
+                num_subgroups = (size + self.max_group_size - 1) // self.max_group_size
+                stats[pre_hash] = {
+                    "total_tests": size,
+                    "num_subgroups": num_subgroups,
+                    "tests_per_subgroup": size // num_subgroups,
+                }
+        return stats
+
+
 def default_input() -> str:
     """
     Directory (default) to consume generated test fixtures from. Defined as a
@@ -357,6 +472,29 @@ def pytest_configure(config):  # noqa: D103
     index = IndexFile.model_validate_json(index_file.read_text())
     config.test_cases = index.test_cases
 
+    # Create XDistGroupMapper for enginex simulator if needed
+    if (
+        hasattr(config, "_supported_fixture_formats")
+        and BlockchainEngineXFixture.format_name in config._supported_fixture_formats
+    ):
+        max_group_size = getattr(config, "enginex_max_group_size", 400)
+        config.xdist_group_mapper = XDistGroupMapper(max_group_size)
+        config.xdist_group_mapper.build_mapping(config.test_cases)
+
+        # Log statistics about group splitting
+        split_stats = config.xdist_group_mapper.get_split_statistics()
+        if split_stats:
+            rich.print("[bold yellow]Pre-allocation group splitting for load balancing:[/]")
+            for pre_hash, stats in split_stats.items():
+                rich.print(
+                    f"  Group {pre_hash[:8]}: {stats['total_tests']} tests → "
+                    f"{stats['num_subgroups']} sub-groups "
+                    f"(~{stats['tests_per_subgroup']} tests each)"
+                )
+            rich.print(f"  Max group size: {max_group_size}")
+    else:
+        config.xdist_group_mapper = None
+
     for fixture_format in BaseFixture.formats.values():
         config.addinivalue_line(
             "markers",
@@ -428,6 +566,7 @@ def pytest_generate_tests(metafunc):
         return
 
     test_cases = metafunc.config.test_cases
+    xdist_group_mapper = getattr(metafunc.config, "xdist_group_mapper", None)
     param_list = []
     for test_case in test_cases:
         if test_case.format.format_name not in metafunc.config._supported_fixture_formats:
@@ -439,12 +578,23 @@ def pytest_generate_tests(metafunc):
         if hasattr(test_case, "pre_hash") and test_case.pre_hash:
             test_id = f"{test_case.id}[{test_case.pre_hash[:8]}]"
 
+        # Determine xdist group name
+        if xdist_group_mapper and hasattr(test_case, "pre_hash") and test_case.pre_hash:
+            # Use the mapper to get potentially split group name
+            xdist_group_name = xdist_group_mapper.get_xdist_group_name(test_case)
+        elif hasattr(test_case, "pre_hash") and test_case.pre_hash:
+            # No mapper or not enginex, use pre_hash directly
+            xdist_group_name = test_case.pre_hash
+        else:
+            # No pre_hash, use test ID
+            xdist_group_name = test_case.id
+
         param = pytest.param(
             test_case,
             id=test_id,
             marks=[getattr(pytest.mark, m) for m in fork_markers]
             + [getattr(pytest.mark, test_case.format.format_name)]
-            + [pytest.mark.xdist_group(name=test_case.pre_hash)],
+            + [pytest.mark.xdist_group(name=xdist_group_name)],
         )
         param_list.append(param)
 
diff --git a/src/pytest_plugins/consume/simulators/enginex/conftest.py b/src/pytest_plugins/consume/simulators/enginex/conftest.py
@@ -39,6 +39,16 @@ def pytest_addoption(parser):
             "pre-allocation group."
         ),
     )
+    enginex_group.addoption(
+        "--enginex-max-group-size",
+        action="store",
+        type=int,
+        default=400,
+        help=(
+            "Maximum number of tests per xdist group. Large pre-allocation groups will be "
+            "split into virtual sub-groups to improve load balancing. Default: 400."
+        ),
+    )
 
 
 def pytest_configure(config):
@@ -48,6 +58,9 @@ def pytest_configure(config):
     # Store FCU frequency on config for access by fixtures
     config.enginex_fcu_frequency = config.getoption("--enginex-fcu-frequency", 1)
 
+    # Store max group size on config for access during test generation
+    config.enginex_max_group_size = config.getoption("--enginex-max-group-size", 400)
+
 
 @pytest.fixture(scope="module")
 def test_suite_name() -> str:
