diff --git a/src/fromager/bootstrapper/_bootstrapper.py b/src/fromager/bootstrapper/_bootstrapper.py index dfec81e2..027e28d1 100644 --- a/src/fromager/bootstrapper/_bootstrapper.py +++ b/src/fromager/bootstrapper/_bootstrapper.py @@ -3,12 +3,14 @@ import concurrent.futures import contextlib import datetime +import io import json import logging import operator import pathlib import shutil import tempfile +import time import typing from packaging.requirements import Requirement @@ -104,6 +106,16 @@ def __init__( self._stack_filename = self.ctx.work_dir / "bootstrap-stack.json" logger.info("recording bootstrap stack state to %s", self._stack_filename) + # Single-threaded pool for background file writes (serialized, never interleave) + self._write_pool: concurrent.futures.ThreadPoolExecutor | None = ( + concurrent.futures.ThreadPoolExecutor( + max_workers=1, thread_name_prefix="fromager-writes" + ) + ) + # Throttle stack-state writes to at most once per interval + self._last_stack_write: float = 0.0 + self._STACK_WRITE_INTERVAL: float = 5.0 + # Track failed packages in test mode (list of typed dicts for JSON export) self.failed_packages: list[FailureRecord] = [] @@ -850,7 +862,7 @@ def add_to_graph( pre_built=pbi.pre_built, constraint=self.ctx.constraints.get_constraint(req.name), ) - self.ctx.write_to_graph_to_file() + self._write_graph_async() def _sort_requirements( self, @@ -954,21 +966,31 @@ def add_to_build_order( if req.url: info["source_url"] = req.url self._build_stack.append(info) - with open(self._build_order_filename, "w") as f: - # Set default=str because the why value includes - # Requirement and Version instances that can't be - # converted to JSON without help. - json.dump(self._build_stack, f, indent=2, default=str) + + def _schedule_write(self, path: pathlib.Path, content: str) -> None: + """Submit a pre-serialized write to the background write pool.""" + if self._write_pool is not None: + self._write_pool.submit(path.write_text, content, "utf-8") + + def _write_graph_async(self) -> None: + """Serialize the dependency graph on the main thread, write it in background.""" + buf = io.StringIO() + self.ctx.dependency_graph.serialize(buf) + self._schedule_write(self.ctx.graph_file, buf.getvalue()) def _record_stack_state(self, stack: list[Phase]) -> None: """Write the current bootstrap stack to `self._stack_filename`. Index 0 in the output corresponds to `stack[-1]`, the next item to be - processed. Overwrites the file on each call. + processed. Throttled to at most once every ``_STACK_WRITE_INTERVAL`` seconds. """ + now = time.monotonic() + if now - self._last_stack_write < self._STACK_WRITE_INTERVAL: + return records = [item.as_json() for item in reversed(stack)] with open(self._stack_filename, "w") as f: json.dump(records, f, indent=2, default=str) + self._last_stack_write = now # ---- Iterative bootstrap: phase handlers and helpers ---- @@ -1198,7 +1220,7 @@ def _handle_phase_error( self._seen_requirements.discard( self._resolved_key(wi.req, wi.resolved_version, "wheel") ) - self.ctx.write_to_graph_to_file() + self._write_graph_async() return [] # Normal mode: fail-fast @@ -1246,6 +1268,21 @@ def finalize(self) -> int: self._bg_pool.shutdown(wait=True, cancel_futures=True) self._bg_pool = None + # Write build-order once at the end (data was buffered in _build_stack) + with open(self._build_order_filename, "w") as f: + # Set default=str because values may include Requirement and Version + # instances that can't be converted to JSON without help. + json.dump(self._build_stack, f, indent=2, default=str) + + # Final graph write and stack flush, then drain the write pool + self._write_graph_async() + records: list[typing.Any] = [] + with open(self._stack_filename, "w") as f: + json.dump(records, f, indent=2) + if self._write_pool is not None: + self._write_pool.shutdown(wait=True, cancel_futures=False) + self._write_pool = None + if self.multiple_versions and self._failed_versions: self._log_failed_versions_table() @@ -1284,3 +1321,6 @@ def __exit__( if self._bg_pool is not None: self._bg_pool.shutdown(wait=False, cancel_futures=True) self._bg_pool = None + if self._write_pool is not None: + self._write_pool.shutdown(wait=False, cancel_futures=True) + self._write_pool = None diff --git a/src/fromager/context.py b/src/fromager/context.py index 58866aed..40d12639 100644 --- a/src/fromager/context.py +++ b/src/fromager/context.py @@ -156,7 +156,8 @@ def uv_clean_cache(self, *reqs: Requirement) -> None: cmd.extend(req_list) external_commands.run(cmd, extra_environ=extra_environ) - def write_to_graph_to_file(self) -> None: + def write_graph_file(self) -> None: + """Write the dependency graph to ``self.graph_file``.""" with self.graph_file.open("w", encoding="utf-8") as f: self.dependency_graph.serialize(f) diff --git a/tests/test_bootstrapper.py b/tests/test_bootstrapper.py index 797121f5..409588e3 100644 --- a/tests/test_bootstrapper.py +++ b/tests/test_bootstrapper.py @@ -129,6 +129,7 @@ def test_build_order(tmp_context: WorkContext) -> None: source_url="url", source_type=SourceType.SDIST, ) + bt.finalize() contents_str = bt._build_order_filename.read_text() contents = json.loads(contents_str) expected = [ @@ -174,6 +175,7 @@ def test_build_order_repeats(tmp_context: WorkContext) -> None: "url", SourceType.SDIST, ) + bt.finalize() contents_str = bt._build_order_filename.read_text() contents = json.loads(contents_str) expected = [ @@ -204,6 +206,7 @@ def test_build_order_name_canonicalization(tmp_context: WorkContext) -> None: "url", SourceType.SDIST, ) + bt.finalize() contents_str = bt._build_order_filename.read_text() contents = json.loads(contents_str) expected = [ @@ -764,12 +767,14 @@ def test_record_stack_state_ordering(tmp_context: WorkContext) -> None: def test_record_stack_state_overwrites_each_call(tmp_context: WorkContext) -> None: - """Second call replaces first call's content.""" + """Second call replaces first call's content when throttle interval has elapsed.""" bt = bootstrapper.Bootstrapper(tmp_context) bt._record_stack_state([_make_resolve_item("pkga"), _make_resolve_item("pkgb")]) first_content = bt._stack_filename.read_text() + # Reset throttle so the second call is not suppressed + bt._last_stack_write = 0.0 bt._record_stack_state([_make_resolve_item("pkgc")]) second_content = bt._stack_filename.read_text() @@ -779,6 +784,46 @@ def test_record_stack_state_overwrites_each_call(tmp_context: WorkContext) -> No assert contents[0]["req"] == "pkgc" +def test_record_stack_state_throttled_when_called_rapidly( + tmp_context: WorkContext, +) -> None: + """Rapid successive calls do not overwrite the file (throttle active).""" + bt = bootstrapper.Bootstrapper(tmp_context) + stack: list[Phase] = [_make_resolve_item("pkga"), _make_resolve_item("pkgb")] + + # First call writes (interval has elapsed from epoch) + bt._record_stack_state(stack) + first_mtime = bt._stack_filename.stat().st_mtime + + # Second call is throttled — file should not change + bt._record_stack_state([_make_resolve_item("pkgc")]) + second_mtime = bt._stack_filename.stat().st_mtime + + assert first_mtime == second_mtime + + +def test_finalize_writes_build_order_and_graph(tmp_context: WorkContext) -> None: + """finalize() writes build-order.json and graph.json, and drains the write pool.""" + bt = bootstrapper.Bootstrapper(tmp_context) + bt.add_to_build_order( + req=Requirement("mypkg==1.0"), + version=Version("1.0"), + source_url="https://pypi.test/mypkg-1.0.tar.gz", + source_type=SourceType.SDIST, + ) + + assert not bt._build_order_filename.exists() + + bt.finalize() + + assert bt._build_order_filename.exists() + assert bt._write_pool is None # pool was drained and closed + + contents = json.loads(bt._build_order_filename.read_text()) + assert len(contents) == 1 + assert contents[0]["dist"] == "mypkg" + + def test_bootstrap_calls_record_stack_state(tmp_context: WorkContext) -> None: """`_record_stack_state` is called at least once during `bootstrap()`.""" bt = bootstrapper.Bootstrapper(tmp_context) diff --git a/tests/test_context.py b/tests/test_context.py index 948ca2a3..1ebe0b9b 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -155,12 +155,12 @@ def test_wheels_build_parallel(tmp_context: context.WorkContext) -> None: assert result.is_dir() -def test_write_to_graph_to_file(tmp_path: pathlib.Path) -> None: +def test_write_graph_file(tmp_path: pathlib.Path) -> None: ctx = _make_context(tmp_path) ctx.setup() with patch.object(ctx.dependency_graph, "serialize") as mock_serialize: - ctx.write_to_graph_to_file() + ctx.write_graph_file() mock_serialize.assert_called_once() assert ctx.graph_file.exists()