feat(fill): Improve mem usage by: fixture collector flush + lazy loading pre-alloc groups (#2032)

* feat(FixtureCollector): Add flush interval to the fixture collector

* fix: typing

* fix: typing

* fix: unit test

Author: marioevz

src/cli/show_pre_alloc_group_stats.py

@@ -102,7 +102,7 @@ def calculate_size_distribution(
 
 def analyze_pre_alloc_folder(folder: Path, verbose: int = 0) -> Dict:
     """Analyze pre-allocation folder and return statistics."""
-    pre_alloc_groups = PreAllocGroups.from_folder(folder)
+    pre_alloc_groups = PreAllocGroups.from_folder(folder, lazy_load=False)
 
     # Basic stats
     total_groups = len(pre_alloc_groups)
@@ -147,7 +147,7 @@ def analyze_pre_alloc_folder(folder: Path, verbose: int = 0) -> Dict:
 
     # Calculate frequency distribution of group sizes
     group_distribution, test_distribution = calculate_size_distribution(
-        [g["tests"] for g in group_details]
+        [g["tests"] for g in group_details]  # type: ignore
     )
 
     # Analyze test functions split across multiple size-1 groups

src/ethereum_test_fixtures/collector.py

@@ -113,6 +113,7 @@ class FixtureCollector:
     single_fixture_per_file: bool
     filler_path: Path
     base_dump_dir: Optional[Path] = None
+    flush_interval: int = 1000
 
     # Internal state
     all_fixtures: Dict[Path, Fixtures] = field(default_factory=dict)
@@ -151,6 +152,9 @@ def add_fixture(self, info: TestInfo, fixture: BaseFixture) -> Path:
 
         self.all_fixtures[fixture_path][info.get_id()] = fixture
 
+        if self.flush_interval > 0 and len(self.all_fixtures) >= self.flush_interval:
+            self.dump_fixtures()
+
         return fixture_path
 
     def dump_fixtures(self) -> None:
@@ -168,6 +172,8 @@ def dump_fixtures(self) -> None:
                 raise TypeError("All fixtures in a single file must have the same format.")
             fixtures.collect_into_file(fixture_path)
 
+        self.all_fixtures.clear()
+
     def verify_fixture_files(self, evm_fixture_verification: FixtureConsumer) -> None:
         """Run `evm [state|block]test` on each fixture."""
         for fixture_path, name_fixture_dict in self.all_fixtures.items():

src/ethereum_test_fixtures/pre_alloc_groups.py

@@ -2,10 +2,10 @@
 
 import json
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Dict, Iterator, List, Tuple
 
 from filelock import FileLock
-from pydantic import Field, computed_field
+from pydantic import Field, PrivateAttr, computed_field
 
 from ethereum_test_base_types import CamelModel, EthereumTestRootModel
 from ethereum_test_forks import Fork
@@ -84,35 +84,62 @@ def to_file(self, file: Path) -> None:
 
 
 class PreAllocGroups(EthereumTestRootModel):
-    """Root model mapping pre-allocation group hashes to test groups."""
+    """
+    Root model mapping pre-allocation group hashes to test groups.
+
+    If lazy_load is True, the groups are not loaded from the folder until they are accessed.
+
+    Iterating will fail if lazy_load is True.
+    """
 
-    root: Dict[str, PreAllocGroup]
+    root: Dict[str, PreAllocGroup | None]
+
+    _folder_source: Path | None = PrivateAttr(None)
 
     def __setitem__(self, key: str, value: Any):
         """Set item in root dict."""
+        assert self._folder_source is None, (
+            "Cannot set item in root dict after folder source is set"
+        )
         self.root[key] = value
 
     @classmethod
-    def from_folder(cls, folder: Path) -> "PreAllocGroups":
+    def from_folder(cls, folder: Path, *, lazy_load: bool = False) -> "PreAllocGroups":
         """Create PreAllocGroups from a folder of pre-allocation files."""
         # First check for collision failures
         for fail_file in folder.glob("*.fail"):
             with open(fail_file) as f:
                 raise Alloc.CollisionError.from_json(json.loads(f.read()))
-        data = {}
+
+        data: Dict[str, PreAllocGroup | None] = {}
         for file in folder.glob("*.json"):
-            with open(file) as f:
-                data[file.stem] = PreAllocGroup.model_validate_json(f.read())
-        return cls(root=data)
+            if lazy_load:
+                data[file.stem] = None
+            else:
+                with open(file) as f:
+                    data[file.stem] = PreAllocGroup.model_validate_json(f.read())
+        instance = cls(root=data)
+        if lazy_load:
+            instance._folder_source = folder
+        return instance
 
     def to_folder(self, folder: Path) -> None:
         """Save PreAllocGroups to a folder of pre-allocation files."""
         for key, value in self.root.items():
+            assert value is not None, f"Value for key {key} is None"
             value.to_file(folder / f"{key}.json")
 
     def __getitem__(self, item):
         """Get item from root dict."""
-        return self.root[item]
+        if self._folder_source is None:
+            item = self.root[item]
+            assert item is not None, f"Item {item} is None"
+            return item
+        else:
+            if self.root[item] is None:
+                with open(self._folder_source / f"{item}.json") as f:
+                    self.root[item] = PreAllocGroup.model_validate_json(f.read())
+            return self.root[item]
 
     def __iter__(self):
         """Iterate over root dict."""
@@ -130,10 +157,14 @@ def keys(self):
         """Get keys from root dict."""
         return self.root.keys()
 
-    def values(self):
+    def values(self) -> Iterator[PreAllocGroup]:
         """Get values from root dict."""
-        return self.root.values()
+        for value in self.root.values():
+            assert value is not None, "Value is None"
+            yield value
 
-    def items(self):
+    def items(self) -> Iterator[Tuple[str, PreAllocGroup]]:
         """Get items from root dict."""
-        return self.root.items()
+        for key, value in self.root.items():
+            assert value is not None, f"Value for key {key} is None"
+            yield key, value

src/pytest_plugins/filler/filler.py

@@ -226,7 +226,7 @@ def _load_pre_alloc_groups_from_folder(self) -> None:
         """Load pre-allocation groups from the output folder."""
         pre_alloc_folder = self.fixture_output.pre_alloc_groups_folder_path
         if pre_alloc_folder.exists():
-            self.pre_alloc_groups = PreAllocGroups.from_folder(pre_alloc_folder)
+            self.pre_alloc_groups = PreAllocGroups.from_folder(pre_alloc_folder, lazy_load=True)
         else:
             raise FileNotFoundError(
                 f"Pre-allocation groups folder not found: {pre_alloc_folder}. "
@@ -315,7 +315,7 @@ def aggregate_pre_alloc_groups(self, worker_groups: PreAllocGroups) -> None:
         if self.pre_alloc_groups is None:
             self.pre_alloc_groups = PreAllocGroups(root={})
 
-        for hash_key, group in worker_groups.root.items():
+        for hash_key, group in worker_groups.items():
             if hash_key in self.pre_alloc_groups:
                 # Merge if exists (should not happen in practice)
                 existing = self.pre_alloc_groups[hash_key]
@@ -699,16 +699,15 @@ def pytest_terminal_summary(
             if config.pluginmanager.hasplugin("xdist"):
                 # Load pre-allocation groups from disk
                 pre_alloc_groups = PreAllocGroups.from_folder(
-                    config.fixture_output.pre_alloc_groups_folder_path  # type: ignore[attr-defined]
+                    config.fixture_output.pre_alloc_groups_folder_path,  # type: ignore[attr-defined]
+                    lazy_load=False,
                 )
             else:
                 assert session_instance.pre_alloc_groups is not None
                 pre_alloc_groups = session_instance.pre_alloc_groups
 
             total_groups = len(pre_alloc_groups.root)
-            total_accounts = sum(
-                group.pre_account_count for group in pre_alloc_groups.root.values()
-            )
+            total_accounts = sum(group.pre_account_count for group in pre_alloc_groups.values())
 
             terminalreporter.write_sep(
                 "=",

src/pytest_plugins/filler/tests/test_prealloc_group.py

@@ -436,7 +436,7 @@ def test_pre_alloc_grouping_by_test_type(
         Path(default_output_directory()).absolute() / "blockchain_tests_engine_x" / "pre_alloc"
     )
     assert output_dir.exists()
-    groups = PreAllocGroups.from_folder(output_dir)
+    groups = PreAllocGroups.from_folder(output_dir, lazy_load=False)
     if (
         len([f for f in output_dir.iterdir() if f.name.endswith(".json")])
         != expected_different_pre_alloc_groups