diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d6aab5f3..7c36ca0d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -81,9 +81,11 @@ jobs: - name: Test UK calibration source additions run: uv run pytest -q --durations=30 tests/test_chronicle_uk_national_2025.py - # GNU time reports the run's peak memory ("Maximum resident set size"): the - # merged bundle peaked near 15.6 GB in a local build, against the runner's - # 16 GB. + # GNU time reports the run's peak memory ("Maximum resident set size", in + # KiB). build_bundle writes and counts each package's rows as it reads them + # (chronicle#414), so this step peaked at 3.4 GiB for 1.74M facts on Nov 5, + # against 15.0 GiB, close to the runner's memory, when every row was held + # to the end. - name: Test merged Chronicle bundles run: | timer="" diff --git a/changelog.d/bundle-streams-consumer-rows.fixed.md b/changelog.d/bundle-streams-consumer-rows.fixed.md new file mode 100644 index 00000000..674ce96d --- /dev/null +++ b/changelog.d/bundle-streams-consumer-rows.fixed.md @@ -0,0 +1,7 @@ +`build-bundle` writes and counts each source package's consumer rows as it reads +them, and frees them before it builds the next package. It used to hold every +package's rows until the end. CI's merged default bundle (1.74M rows) peaked at +15.0 GiB that way and now peaks at 3.4 GiB; the UK suite's local build went from +12.8 to 4.0 GiB. `consumer_facts.jsonl`, `coverage.json`, `source_packages.json` +and the build report are unchanged (byte-identical on the UK suite); the facts +file takes its name only once every package is in it. diff --git a/chronicle/bundle.py b/chronicle/bundle.py index 3f2d16ed..cbd8d5eb 100644 --- a/chronicle/bundle.py +++ b/chronicle/bundle.py @@ -7,7 +7,7 @@ import shutil from dataclasses import asdict, dataclass from pathlib import Path -from typing import Any, Callable, Sequence +from typing import Any, Callable, Iterable, Sequence from chronicle.dimension_labels import ( cross_package_value_label_remedy, @@ -813,102 +813,113 @@ def build_bundle( errors: list[BuildBundleIssue] = [] warnings: list[BuildBundleIssue] = [] source_reports: list[BundleSourceReport] = [] - consumer_rows: list[dict[str, Any]] = [] - shared_metadata: dict[tuple[str, str], dict[str, Any]] = {} dimension_labels: dict[str, dict[str, list[str]]] = {} geography_names: dict[tuple[str, str], dict[str, list[str]]] = {} value_labels: dict[tuple[str, str], dict[str, list[str]]] = {} + # Each package's rows are written and counted as soon as they are read, then + # freed. Holding every package's rows until the end peaked at 15.0 GiB on CI + # for the default bundle's 1.74M rows, close to the runner's memory. The file + # takes its final name once every package is in it; a failed build removes it. + coverage_counts = _BundleCoverage() + consumer_facts_path = output_path / "consumer_facts.jsonl" + partial_facts_path = output_path / "consumer_facts.jsonl.partial" + partial_facts_path.write_text("", encoding="utf-8") + + try: + for source in build_sources: + suite_dir = sources_path / _safe_source_dir_name(source) + try: + suite_report = build_source_suite( + source, + suite_dir, + year=year, + axiom_command=axiom_command, + axiom_roots=axiom_roots, + require_axiom_validation=require_axiom_validation, + replace=True, + ) + except Exception as exc: + errors.append( + BuildBundleIssue( + code="source_suite_build_failed", + message=str(exc), + source=source, + ) + ) + continue - for source in build_sources: - suite_dir = sources_path / _safe_source_dir_name(source) - try: - suite_report = build_source_suite( - source, - suite_dir, - year=year, - axiom_command=axiom_command, - axiom_roots=axiom_roots, - require_axiom_validation=require_axiom_validation, - replace=True, + if not suite_report.valid: + errors.append( + BuildBundleIssue( + code="source_suite_invalid", + message="Nested source-package suite did not pass validation.", + source=source, + ) + ) + # The metadata pool is the package's own, so it is freed with its rows. + rows = _load_jsonl( + Path(suite_report.outputs["consumer_facts"]), + shared_metadata={}, ) - except Exception as exc: + _append_jsonl(partial_facts_path, rows) + coverage_counts.add(rows) + source_reports.append(_bundle_source_report(source, suite_report)) + label_errors, label_warnings = _dimension_label_reports(source, rows) + errors.extend(label_errors) + warnings.extend(label_warnings) + errors.extend(_geography_name_errors(source, rows)) + for area, names in _geography_names(rows).items(): + for name in names: + geography_names.setdefault(area, {}).setdefault(name, []).append( + source + ) + for pair, labels in _value_labels(rows).items(): + for label in labels: + value_labels.setdefault(pair, {}).setdefault(label, []).append( + source + ) + for dimension_id, labels in dimension_labels_by_id(rows).items(): + for label in labels: + dimension_labels.setdefault(dimension_id, {}).setdefault( + label, [] + ).append(source) + # Free this package's rows and suite reports before the next one builds. + del rows, suite_report + + errors.extend(_cross_package_dimension_label_errors(dimension_labels)) + errors.extend(_cross_package_value_label_errors(value_labels)) + warnings.extend(_cross_package_geography_name_warnings(geography_names)) + aggregate_duplicates = coverage_counts.duplicates("aggregate_fact_key") + semantic_duplicates = coverage_counts.duplicates("semantic_fact_key") + for duplicate in aggregate_duplicates: errors.append( BuildBundleIssue( - code="source_suite_build_failed", - message=str(exc), - source=source, + code="duplicate_aggregate_fact_key", + message="Aggregate fact keys must be unique within a bundle.", + key=duplicate["key"], ) ) - continue - - if not suite_report.valid: - errors.append( + if semantic_duplicates: + warnings.append( BuildBundleIssue( - code="source_suite_invalid", - message="Nested source-package suite did not pass validation.", - source=source, + code="duplicate_semantic_fact_key", + message=( + "One or more semantic facts appear in multiple rows; " + "downstream consumers should reconcile or select sources." + ), ) ) - rows = _load_jsonl( - Path(suite_report.outputs["consumer_facts"]), - shared_metadata=shared_metadata, - ) - consumer_rows.extend(rows) - source_reports.append(_bundle_source_report(source, suite_report)) - label_errors, label_warnings = _dimension_label_reports(source, rows) - errors.extend(label_errors) - warnings.extend(label_warnings) - errors.extend(_geography_name_errors(source, rows)) - for area, names in _geography_names(rows).items(): - for name in names: - geography_names.setdefault(area, {}).setdefault(name, []).append(source) - for pair, labels in _value_labels(rows).items(): - for label in labels: - value_labels.setdefault(pair, {}).setdefault(label, []).append(source) - for dimension_id, labels in dimension_labels_by_id(rows).items(): - for label in labels: - dimension_labels.setdefault(dimension_id, {}).setdefault( - label, [] - ).append(source) - - errors.extend(_cross_package_dimension_label_errors(dimension_labels)) - errors.extend(_cross_package_value_label_errors(value_labels)) - warnings.extend(_cross_package_geography_name_warnings(geography_names)) - aggregate_duplicates = _duplicate_key_reports( - consumer_rows, - "aggregate_fact_key", - ) - semantic_duplicates = _duplicate_key_reports( - consumer_rows, - "semantic_fact_key", - ) - for duplicate in aggregate_duplicates: - errors.append( - BuildBundleIssue( - code="duplicate_aggregate_fact_key", - message="Aggregate fact keys must be unique within a bundle.", - key=duplicate["key"], - ) - ) - if semantic_duplicates: - warnings.append( - BuildBundleIssue( - code="duplicate_semantic_fact_key", - message=( - "One or more semantic facts appear in multiple rows; " - "downstream consumers should reconcile or select sources." - ), - ) - ) - consumer_facts_path = output_path / "consumer_facts.jsonl" + partial_facts_path.replace(consumer_facts_path) + except BaseException: + partial_facts_path.unlink(missing_ok=True) + raise + source_packages_path = output_path / "source_packages.json" coverage_path = output_path / "coverage.json" report_path = reports_path / "build_bundle.json" - _write_jsonl(consumer_facts_path, consumer_rows) - coverage = build_bundle_coverage( - consumer_rows, + coverage = coverage_counts.report( aggregate_duplicates=aggregate_duplicates, semantic_duplicates=semantic_duplicates, ) @@ -1196,46 +1207,12 @@ def build_bundle_coverage( semantic_duplicates: list[dict[str, Any]] | None = None, ) -> dict[str, Any]: """Build source/concept/geography coverage diagnostics for consumer rows.""" - aggregate_duplicates = ( - aggregate_duplicates - if aggregate_duplicates is not None - else _duplicate_key_reports(rows, "aggregate_fact_key") - ) - semantic_duplicates = ( - semantic_duplicates - if semantic_duplicates is not None - else _duplicate_key_reports(rows, "semantic_fact_key") + coverage = _BundleCoverage() + coverage.add(rows) + return coverage.report( + aggregate_duplicates=aggregate_duplicates, + semantic_duplicates=semantic_duplicates, ) - return { - "schema_version": BUNDLE_COVERAGE_SCHEMA_VERSION, - "fact_count": len(rows), - "counts": { - "by_source": _counts_by(rows, _source_name), - "by_source_table": _counts_by(rows, _source_table_key), - "by_period": _counts_by(rows, _period_key), - "by_geography": _counts_by(rows, _geography_key), - "by_entity": _counts_by(rows, _entity_key), - "by_observed_measure": _counts_by(rows, _observed_measure_key), - "by_observed_concept": _counts_by(rows, _observed_concept_key), - "by_canonical_concept": _counts_by(rows, _canonical_concept_key), - }, - "unique_counts": { - "aggregate_fact_key": _unique_count(rows, "aggregate_fact_key"), - "semantic_fact_key": _unique_count(rows, "semantic_fact_key"), - "source_release_key": _unique_count(rows, "source_release_key"), - "source_series_key": _unique_count(rows, "source_series_key"), - "observed_measure_key": _unique_count(rows, "observed_measure_key"), - "dimension_set_key": _unique_count(rows, "dimension_set_key"), - "universe_constraint_set_key": _unique_count( - rows, - "universe_constraint_set_key", - ), - }, - "duplicates": { - "aggregate_fact_keys": aggregate_duplicates, - "semantic_fact_keys": semantic_duplicates, - }, - } def _resolve_bundle_sources( @@ -1303,50 +1280,158 @@ def _bundle_source_report( ) -def _duplicate_key_reports( - rows: list[dict[str, Any]], - key: str, -) -> list[dict[str, Any]]: - grouped: dict[Any, list[dict[str, Any]]] = {} - for row in rows: - grouped.setdefault(_canonical_key(row[key], key), []).append(row) - return [ - { - "key": key_value, - "count": len(key_rows), - "sources": sorted({_source_table_key(row) for row in key_rows}), - "legacy_fact_keys": sorted( - { - _canonical_key(legacy_key, "legacy_fact_key") - for row in key_rows - if (legacy_key := row.get("legacy_fact_key")) - }, - key=_identity_sort_key, - ), +class _BundleCoverage: + """Coverage diagnostics added up one package's rows at a time. + + It keeps counts, each identity's distinct values and, for each fact key, + where its first row came from, never the rows themselves, so build_bundle + can free a package's rows once they are added. + """ + + def __init__(self) -> None: + self.fact_count = 0 + self.count_keys: dict[str, Callable[[dict[str, Any]], str | None]] = { + "by_source": _source_name, + "by_source_table": _source_table_key, + "by_period": _period_key, + "by_geography": _geography_key, + "by_entity": _entity_key, + "by_observed_measure": _observed_measure_key, + "by_observed_concept": _observed_concept_key, + "by_canonical_concept": _canonical_concept_key, + } + self.counts: dict[str, dict[str, int]] = {name: {} for name in self.count_keys} + self.fact_keys = { + key: _FactKeyIndex(key) + for key in ("aggregate_fact_key", "semantic_fact_key") + } + self.identities: dict[str, set[Any]] = { + key: set() + for key in ( + "source_release_key", + "source_series_key", + "observed_measure_key", + "dimension_set_key", + "universe_constraint_set_key", + ) + } + # One string per source table, however many rows name it. + self.source_tables: dict[str | None, str | None] = {} + + def add(self, rows: Iterable[dict[str, Any]]) -> None: + for row in rows: + self.fact_count += 1 + for name, key_fn in self.count_keys.items(): + key = key_fn(row) + if key is not None: + counts = self.counts[name] + counts[key] = counts.get(key, 0) + 1 + for key, values in self.identities.items(): + if key in row: + values.add(_canonical_key(row[key], key)) + source_table = _source_table_key(row) + origin = ( + self.source_tables.setdefault(source_table, source_table), + row.get("legacy_fact_key") or None, + ) + for index in self.fact_keys.values(): + index.add(row, origin) + + def duplicates(self, key: str) -> list[dict[str, Any]]: + return self.fact_keys[key].duplicates() + + def report( + self, + *, + aggregate_duplicates: list[dict[str, Any]] | None = None, + semantic_duplicates: list[dict[str, Any]] | None = None, + ) -> dict[str, Any]: + if aggregate_duplicates is None: + aggregate_duplicates = self.duplicates("aggregate_fact_key") + if semantic_duplicates is None: + semantic_duplicates = self.duplicates("semantic_fact_key") + return { + "schema_version": BUNDLE_COVERAGE_SCHEMA_VERSION, + "fact_count": self.fact_count, + "counts": { + name: dict(sorted(counts.items())) + for name, counts in self.counts.items() + }, + "unique_counts": { + **{key: len(index.first) for key, index in self.fact_keys.items()}, + **{key: len(values) for key, values in self.identities.items()}, + }, + "duplicates": { + "aggregate_fact_keys": aggregate_duplicates, + "semantic_fact_keys": semantic_duplicates, + }, } - for key_value, key_rows in sorted( - grouped.items(), - key=lambda item: _identity_sort_key(item[0]), - ) - if len(key_rows) > 1 - ] -def _counts_by( - rows: list[dict[str, Any]], - key_fn: Callable[[dict[str, Any]], str | None], -) -> dict[str, int]: - counts: dict[str, int] = {} - for row in rows: - key = key_fn(row) - if key is None: - continue - counts[key] = counts.get(key, 0) + 1 - return dict(sorted(counts.items())) +@dataclass(slots=True) +class _RepeatedFactKey: + count: int + source_tables: set[str | None] + legacy_fact_keys: set[Any] -def _unique_count(rows: list[dict[str, Any]], key: str) -> int: - return len({_canonical_key(row[key], key) for row in rows if key in row}) +class _FactKeyIndex: + """Each distinct value of one fact key, and where the repeated ones come from. + + A key keeps its first row's source table and legacy fact key; only keys + seen again collect the rest, so duplicates are reported without any row. + """ + + def __init__(self, key: str) -> None: + self.key = key + self.first: dict[Any, tuple[str | None, Any]] = {} + self.repeated: dict[Any, _RepeatedFactKey] = {} + self.missing = False + + def add(self, row: dict[str, Any], origin: tuple[str | None, Any]) -> None: + if self.key not in row: + self.missing = True + return + value = _canonical_key(row[self.key], self.key) + first = self.first.get(value) + if first is None: + self.first[value] = origin + return + repeated = self.repeated.get(value) + if repeated is None: + repeated = self.repeated[value] = _RepeatedFactKey( + count=1, + source_tables={first[0]}, + legacy_fact_keys=set() if first[1] is None else {first[1]}, + ) + repeated.count += 1 + repeated.source_tables.add(origin[0]) + if origin[1] is not None: + repeated.legacy_fact_keys.add(origin[1]) + + def duplicates(self) -> list[dict[str, Any]]: + if self.missing: + # A row without the key fails here, as it did when every row was + # grouped by it. + raise KeyError(self.key) + return [ + { + "key": value, + "count": repeated.count, + "sources": sorted(repeated.source_tables), + "legacy_fact_keys": sorted( + { + _canonical_key(legacy_key, "legacy_fact_key") + for legacy_key in repeated.legacy_fact_keys + }, + key=_identity_sort_key, + ), + } + for value, repeated in sorted( + self.repeated.items(), + key=lambda item: _identity_sort_key(item[0]), + ) + ] def _canonical_key(value: Any, field_name: str) -> Any: @@ -1477,9 +1562,8 @@ def _load_jsonl( return rows -def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - with path.open("w", encoding="utf-8") as file: +def _append_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> None: + with path.open("a", encoding="utf-8") as file: for row in rows: file.write(json.dumps(row, sort_keys=True)) file.write("\n") diff --git a/tests/test_chronicle_bundle.py b/tests/test_chronicle_bundle.py index fd03791b..77a9f8ce 100644 --- a/tests/test_chronicle_bundle.py +++ b/tests/test_chronicle_bundle.py @@ -3,9 +3,11 @@ from __future__ import annotations import json +import weakref from collections import Counter from datetime import date, timedelta from pathlib import Path +from types import SimpleNamespace import pytest @@ -3190,6 +3192,48 @@ def test_bundle_coverage_canonicalizes_cross_epoch_identities(): ] +def test_bundle_coverage_duplicates_keep_each_keys_first_row(): + # Coverage keeps only a key's first row's source table and legacy key until + # the key repeats, so the first row's must still reach the report. Main's + # build, which grouped every row by key, reports the same. + fixture_rows = _fixture_consumer_rows() + first, second, third = (json.loads(json.dumps(fixture_rows[0])) for _ in range(3)) + first["source"]["source_table"] = "Publication 1304 Table 1.2" + first["legacy_fact_key"] = fixture_rows[1]["legacy_fact_key"] + first["semantic_fact_key"] = fixture_rows[1]["semantic_fact_key"] + del second["legacy_fact_key"] + + coverage = build_bundle_coverage([first, second, third]) + + assert coverage["unique_counts"]["aggregate_fact_key"] == 1 + assert coverage["unique_counts"]["semantic_fact_key"] == 2 + assert coverage["duplicates"] == { + "aggregate_fact_keys": [ + { + "key": fixture_rows[0]["aggregate_fact_key"], + "count": 3, + "sources": [ + "irs_soi:Publication 1304 Table 1.1", + "irs_soi:Publication 1304 Table 1.2", + ], + "legacy_fact_keys": [ + fixture_rows[0]["legacy_fact_key"], + fixture_rows[1]["legacy_fact_key"], + ], + } + ], + # The semantic key's first row is the second, which has no legacy key. + "semantic_fact_keys": [ + { + "key": fixture_rows[0]["semantic_fact_key"], + "count": 2, + "sources": ["irs_soi:Publication 1304 Table 1.1"], + "legacy_fact_keys": [fixture_rows[0]["legacy_fact_key"]], + } + ], + } + + def test_bundle_coverage_preserves_non_string_identity_scalars(tmp_path): identity_fields = ( "aggregate_fact_key", @@ -3254,3 +3298,103 @@ def test_bundle_metadata_pool_preserves_serialized_rows_and_distinct_facts(tmp_p assert loaded[0]["geography"] is loaded[1]["geography"] assert loaded[0]["entity"] is loaded[1]["entity"] assert loaded[0]["period"] is loaded[1]["period"] + + +def _stub_source_suites(monkeypatch, packages, on_build=lambda: None): + """Have build_bundle build each named package as a suite of the given rows.""" + + def build_source_suite(source, suite_dir, **kwargs): + on_build() + suite_dir.mkdir(parents=True) + facts = suite_dir / "consumer_facts.jsonl" + facts.write_text( + "".join(json.dumps(row, sort_keys=True) + "\n" for row in packages[source]), + encoding="utf-8", + ) + return SimpleNamespace( + source=source, + output_dir=str(suite_dir), + valid=True, + outputs={"consumer_facts": str(facts)}, + to_dict=lambda: {"counts": {}}, + ) + + monkeypatch.setattr(chronicle.bundle, "build_source_suite", build_source_suite) + monkeypatch.setattr( + chronicle.bundle, + "validate_source_package", + lambda source, year: SimpleNamespace(valid=True, errors=()), + ) + + +def test_build_bundle_frees_each_packages_rows_before_the_next_builds( + tmp_path, monkeypatch +): + """Holding every package's rows until the end peaked at 15.0 GiB on CI for + the default bundle's 1.74M rows. build_bundle now writes and counts a + package's rows, and frees them before the next package builds.""" + fixture_rows = _fixture_consumer_rows() + # The second package repeats the first package's first fact. + packages = { + "first": fixture_rows[:2], + "second": [fixture_rows[0], fixture_rows[2]], + } + loaded = [] + held_when_built = [] + + class Row(dict): + """A row that can be weakly referenced, to see when it is freed.""" + + load_jsonl = chronicle.bundle._load_jsonl + + def load_weakly_referenced_rows(path, **kwargs): + rows = [Row(row) for row in load_jsonl(path, **kwargs)] + loaded.extend(weakref.ref(row) for row in rows) + return rows + + def note_rows_held(): + # build_bundle records a failed package build as a bundle error, so + # this notes what is held rather than asserting here. + held_when_built.append(sum(row() is not None for row in loaded)) + + monkeypatch.setattr(chronicle.bundle, "_load_jsonl", load_weakly_referenced_rows) + _stub_source_suites(monkeypatch, packages, on_build=note_rows_held) + + report = build_bundle(tmp_path / "bundle", sources=["first", "second"]) + + rows = [*packages["first"], *packages["second"]] + assert held_when_built == [0, 0] + assert _load_jsonl(Path(report.outputs["consumer_facts"])) == rows + assert not (tmp_path / "bundle" / "consumer_facts.jsonl.partial").exists() + assert report.coverage == build_bundle_coverage(rows) + assert report.coverage["fact_count"] == 4 + assert report.coverage["duplicates"]["aggregate_fact_keys"] == [ + { + "key": fixture_rows[0]["aggregate_fact_key"], + "count": 2, + "sources": ["irs_soi:Publication 1304 Table 1.1"], + "legacy_fact_keys": [fixture_rows[0]["legacy_fact_key"]], + } + ] + assert "duplicate_aggregate_fact_key" in [issue.code for issue in report.errors] + + +def test_build_bundle_removes_its_partial_facts_file_when_it_fails( + tmp_path, monkeypatch +): + """The first package's rows are written before the second fails to load. + The old build wrote no facts file before the end, and nor does a failed + build now: the UK suite's partial file reaches about 8 GB.""" + fixture_rows = _fixture_consumer_rows() + unloadable = json.loads(json.dumps(fixture_rows[1])) + digest = unloadable["aggregate_fact_key"].partition(":")[2] + unloadable["aggregate_fact_key"] = f"future.aggregate_fact.v4:{digest}" + _stub_source_suites( + monkeypatch, {"first": fixture_rows[:1], "second": [unloadable]} + ) + + with pytest.raises(ValueError, match="chronicle.aggregate_fact.v3"): + build_bundle(tmp_path / "bundle", sources=["first", "second"]) + + assert not (tmp_path / "bundle" / "consumer_facts.jsonl.partial").exists() + assert not (tmp_path / "bundle" / "consumer_facts.jsonl").exists()