From f561daab2b39686e8d9e30c9e1d39f857d13e9d0 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 4 Oct 2026 05:22:59 +0100 Subject: [PATCH] perf(compile): keep native analysis parallel at 10k models Drop the size heuristic that ran compact native analysis on one worker for projects with at least 10,000 models and 48 MiB of SQL. Peak memory is set by the rules phase, not analysis, so four workers do not raise it. Build and serialize the native rules request before custom rules start and outside the native error boundary, then release the request object tree and its JSON bytes before native evaluation, removing two full-project copies from the peak. Builder programming errors keep propagating unchanged. Pin the worker decision and the builder error semantics with unit tests and tighten the dense 10,000-model peak RSS budget from 4 GiB to 3.5 GiB. --- .../compile/_helpers/analysis/compact.py | 12 +- .../rule_engine/_helpers/engine/native.py | 140 ++++++++++++------ .../compile/test_dense_compile_performance.py | 2 +- .../compile/_helpers/analysis/_test_types.py | 8 + .../_helpers/analysis/test_compact_workers.py | 58 ++++++++ .../_helpers/engine/_test_types.py | 17 +++ .../_helpers/engine/test_native_request.py | 109 ++++++++++++++ 7 files changed, 293 insertions(+), 53 deletions(-) create mode 100644 tests/unit/src/sqlbuild/compiler/compile/_helpers/analysis/test_compact_workers.py create mode 100644 tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_native_request.py diff --git a/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py b/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py index 9dc9c4d0c6..cdd4dc7b7c 100644 --- a/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py +++ b/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py @@ -180,8 +180,7 @@ from sqlbuild.diagnostics.main.log_debug_event import log_debug_event _DEBUG_LOGGER: logging.Logger = logging.getLogger("sqlbuild.compile") -_LARGE_COMPACT_PROJECT_MODELS: int = 10_000 -_LARGE_COMPACT_SQL_BYTES: int = 48 * 1024 * 1024 +_COMPACT_ANALYSIS_WORKERS: int = 4 _NATIVE_LEGACY_FALLBACK: str = "native project type recovery requires legacy fallback" @@ -886,14 +885,9 @@ def _prepare_compact_analysis_batch( def _run_compact_analysis_batch(*, preparation: CompactBatchPreparation) -> object: - """Bound resident analysis heaps before the isolated custom-rule host starts.""" + """Analyse every project size with the full worker pool; the rules phase sets peak memory.""" - large_project: bool = ( - len(preparation.projections) >= _LARGE_COMPACT_PROJECT_MODELS - and sum(len(sql.encode("utf-8")) for sql in preparation.cleaned_sql) - >= _LARGE_COMPACT_SQL_BYTES - ) - workers: int = 1 if large_project else 4 + workers: int = _COMPACT_ANALYSIS_WORKERS if preparation.binding_catalog is not None: return orjson.loads( preparation.binding_catalog.native.analyze_compact( diff --git a/src/sqlbuild/rule_engine/_helpers/engine/native.py b/src/sqlbuild/rule_engine/_helpers/engine/native.py index dcf33baf81..9d31604b6c 100644 --- a/src/sqlbuild/rule_engine/_helpers/engine/native.py +++ b/src/sqlbuild/rule_engine/_helpers/engine/native.py @@ -5,7 +5,7 @@ import inspect import json from concurrent.futures import Executor, Future, ThreadPoolExecutor -from dataclasses import asdict, replace +from dataclasses import asdict, dataclass, replace from decimal import Decimal from pathlib import Path from typing import Any, cast @@ -54,6 +54,13 @@ from sqlbuild.sql_values.types import SqlValueKind +@dataclass(frozen=True) +class _PreparedNativeRequest: + identity: str | None + reused: str | None + request_text: str | None + + def evaluate_native( *, project: CompiledProject, @@ -74,31 +81,20 @@ def evaluate_native( if custom_payloads is None else custom_payloads ) - request: dict[str, object] = { - "version": RULES_NATIVE_API_VERSION, - "project_dir": str(project_dir.resolve()), - "dialect": dialect, - "config": _config_payload(config), - "models": _model_payloads( - project=project, - dialect=dialect, - include_type_proof=any( - _rule_selected(config=config, code=code) for code in TYPE_PROOF_RULE_CODES - ), + prepared: _PreparedNativeRequest = _prepare_native_request( + request_json=_serialise_native_request( + _native_request( + project=project, + config=config, + project_dir=project_dir, + dialect=dialect, + initial_findings=initial_findings, + custom_payloads=payloads, + ) ), - "sql_tests": _sql_test_payloads(project), - "sql_scenarios": _sql_scenario_payloads(project), - "public_enums": [ - _enum_payload(declaration) for declaration in project.public_enums.values() - ], - "public_constants": [ - _constant_payload(declaration) for declaration in project.public_constants.values() - ], - "scope_index": scope_metadata_projection(index=project.scope_index), - "initial_findings": [_finding_payload(finding) for finding in initial_findings], - "defer_suppressions": True, - "custom_rules": payloads, - } + project_dir=project_dir, + cache_enabled=config.cache.enabled, + ) with ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-custom-rules") as pool: custom_future: Future[CustomRulesOutcome] | None = ( custom_outcome @@ -114,15 +110,14 @@ def evaluate_native( verify_determinism=verify_determinism, ) ) + reused: bool = prepared.reused is not None try: - response_json: str - reused: bool - response_json, reused = _evaluate_request( - request=request, project_dir=project_dir, cache_enabled=config.cache.enabled + response: object = orjson.loads( + _evaluate_request(prepared=prepared, project_dir=project_dir) ) - response: object = orjson.loads(response_json) except (ValueError, TypeError) as error: raise RulesError(str(error)) from error + del prepared custom: CustomRulesOutcome = ( CustomRulesOutcome(findings=(), cache_hits=0, cache_misses=0, custom_ms=0) if custom_future is None @@ -159,6 +154,42 @@ def evaluate_native( ) +def _native_request( + *, + project: CompiledProject, + config: RulesConfig, + project_dir: Path, + dialect: str, + initial_findings: tuple[Finding, ...], + custom_payloads: list[dict[str, object]], +) -> dict[str, object]: + return { + "version": RULES_NATIVE_API_VERSION, + "project_dir": str(project_dir.resolve()), + "dialect": dialect, + "config": _config_payload(config), + "models": _model_payloads( + project=project, + dialect=dialect, + include_type_proof=any( + _rule_selected(config=config, code=code) for code in TYPE_PROOF_RULE_CODES + ), + ), + "sql_tests": _sql_test_payloads(project), + "sql_scenarios": _sql_scenario_payloads(project), + "public_enums": [ + _enum_payload(declaration) for declaration in project.public_enums.values() + ], + "public_constants": [ + _constant_payload(declaration) for declaration in project.public_constants.values() + ], + "scope_index": scope_metadata_projection(index=project.scope_index), + "initial_findings": [_finding_payload(finding) for finding in initial_findings], + "defer_suppressions": True, + "custom_rules": custom_payloads, + } + + def start_custom_rules( *, executor: Executor, @@ -206,19 +237,42 @@ def _selected_custom_rules( return tuple(rule for rule in catalogue if rule.custom and rule.code in selected) -def _evaluate_request( - *, request: dict[str, object], project_dir: Path, cache_enabled: bool -) -> tuple[str, bool]: - request_json: bytes = orjson.dumps(request, option=orjson.OPT_SORT_KEYS, default=str) - if not cache_enabled: - return _native.evaluate_json(request_json.decode()), False - identity: str = native_request_identity(request_json) - reused: str | None = read_native_response(project_dir=project_dir, identity=identity) - if reused is not None: - return reused, True - response: str = _native.evaluate_json(request_json.decode()) - write_native_response(project_dir=project_dir, identity=identity, response=response) - return response, False +def _serialise_native_request(request: dict[str, object]) -> bytes: + """Encode the native request, reporting values the encoder rejects as rules errors.""" + + try: + return orjson.dumps(request, option=orjson.OPT_SORT_KEYS, default=str) + except orjson.JSONEncodeError as error: + raise RulesError(str(error)) from error + + +def _prepare_native_request( + *, request_json: bytes, project_dir: Path, cache_enabled: bool +) -> _PreparedNativeRequest: + """Resolve the memo and decode once; the caller's bytes die when this returns.""" + + identity: str | None = native_request_identity(request_json) if cache_enabled else None + reused: str | None = ( + read_native_response(project_dir=project_dir, identity=identity) + if identity is not None + else None + ) + return _PreparedNativeRequest( + identity=identity, + reused=reused, + request_text=request_json.decode() if reused is None else None, + ) + + +def _evaluate_request(*, prepared: _PreparedNativeRequest, project_dir: Path) -> str: + if prepared.reused is not None: + return prepared.reused + response: str = _native.evaluate_json(cast(str, prepared.request_text)) + if prepared.identity is not None: + write_native_response( + project_dir=project_dir, identity=prepared.identity, response=response + ) + return response def finalize_native_findings( diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_compile_performance.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_compile_performance.py index 27cb152395..dc07122711 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_compile_performance.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_compile_performance.py @@ -62,7 +62,7 @@ "dense_models_10000_all_rules", 10000, 145.0, - 4 * _GIB, + 7 * _GIB // 2, "913a08c2277304efacf75e293728767a3fa416f6b32f5488dca6ab5baea80668", ), ), diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/analysis/_test_types.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/analysis/_test_types.py index f422157cba..ae3d136001 100644 --- a/tests/unit/src/sqlbuild/compiler/compile/_helpers/analysis/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/analysis/_test_types.py @@ -14,3 +14,11 @@ class DynamicPivotAnalysisTestCase: expected_family_types: tuple[str | None, ...] = ("DECIMAL(12,2)",) expected_input_relations: tuple[str, ...] = () aggregate: str = "MAX" + + +@dataclass(frozen=True) +class CompactAnalysisWorkersTestCase: + description: str + model_count: int + sql_bytes_per_model: int + expected_workers: int diff --git a/tests/unit/src/sqlbuild/compiler/compile/_helpers/analysis/test_compact_workers.py b/tests/unit/src/sqlbuild/compiler/compile/_helpers/analysis/test_compact_workers.py new file mode 100644 index 0000000000..2c70492d97 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/compile/_helpers/analysis/test_compact_workers.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import orjson +import pytest + +from sqlbuild.compiler.compile._helpers.analysis import compact +from sqlbuild.compiler.compile.models import CompactBatchPreparation +from tests.unit.src.sqlbuild.compiler.compile._helpers.analysis._test_types import ( + CompactAnalysisWorkersTestCase, +) + +_MIB: int = 1024 * 1024 + + +@pytest.mark.parametrize( + "test_case", + ( + CompactAnalysisWorkersTestCase( + description="small project", + model_count=3, + sql_bytes_per_model=64, + expected_workers=4, + ), + CompactAnalysisWorkersTestCase( + description="dense ten thousand model project", + model_count=10_000, + sql_bytes_per_model=(64 * _MIB) // 10_000, + expected_workers=4, + ), + ), + ids=lambda case: case.description, +) +def test_given_project_size_when_running_compact_analysis_then_native_uses_all_workers( + test_case: CompactAnalysisWorkersTestCase, + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: list[dict[str, object]] = [] + + def analyze(request_json: str) -> str: + captured.append(orjson.loads(request_json)) + return "{}" + + monkeypatch.setattr(compact._native, "analyze_project_queries_compact_json", analyze) + sql: str = "SELECT 1 AS order_id" + " " * (test_case.sql_bytes_per_model - 20) + preparation: CompactBatchPreparation = CompactBatchPreparation( + cleaned_sql=(sql,) * test_case.model_count, + queries=({"sql": "SELECT 1 AS order_id", "dialect": "duckdb"},), + templates=({"queryIndex": 0},), + projections=({"templateIndex": 0},) * test_case.model_count, + ) + + _ = compact._run_compact_analysis_batch(preparation=preparation) + + assert [request["workers"] for request in captured] == [test_case.expected_workers] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/_test_types.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/_test_types.py index 10d5f74f42..79617b04b7 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/_test_types.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/_test_types.py @@ -202,3 +202,20 @@ class NativeBuildIdentityTestCase: description: str expected_rebuilt_evaluations: int expected_warm_evaluations: int + + +@dataclass(frozen=True) +class NativeRequestBuildErrorTestCase: + """One programming error raised while building the native rules request.""" + + description: str + expected_error: Exception + + +@dataclass(frozen=True) +class NativeRequestEncodeErrorTestCase: + """One request value the JSON encoder rejects.""" + + description: str + rejected_value: object + expected_message: str diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_native_request.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_native_request.py new file mode 100644 index 0000000000..589e4707e2 --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_native_request.py @@ -0,0 +1,109 @@ +from pathlib import Path +from typing import NoReturn + +import pytest + +from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.rule_engine._helpers.engine import native +from sqlbuild.rule_engine.models import RulesCacheConfig, RulesConfig +from tests.unit.src.sqlbuild.rule_engine._helpers.engine._test_types import ( + NativeRequestBuildErrorTestCase, + NativeRequestEncodeErrorTestCase, +) +from tests.unit.src.sqlbuild.rule_engine.main.evaluate.helpers import build_project + + +@pytest.mark.parametrize( + "test_case", + [ + NativeRequestBuildErrorTestCase( + description="type error", expected_error=TypeError("orders payload is not a mapping") + ), + NativeRequestBuildErrorTestCase( + description="value error", expected_error=ValueError("orders payload has no name") + ), + ], + ids=lambda case: case.description, +) +def test_given_failing_request_builder_when_evaluating_native_then_error_propagates_unchanged( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + test_case: NativeRequestBuildErrorTestCase, +) -> None: + project: CompiledProject = build_project( + name="orders", relative_path="models/orders.sql", sql="SELECT 1", config_values={} + ) + started: list[str] = [] + + def failing_model_payloads(**_: object) -> NoReturn: + raise test_case.expected_error + + def record_custom_start(**_: object) -> None: + started.append("custom") + + def record_native_call(_: str) -> str: + started.append("native") + return "{}" + + monkeypatch.setattr(native, "_model_payloads", failing_model_payloads) + monkeypatch.setattr(native, "start_custom_rules", record_custom_start) + monkeypatch.setattr(native._native, "evaluate_json", record_native_call) + + with pytest.raises(type(test_case.expected_error)) as raised: + native.evaluate_native( + project=project, + config=RulesConfig(cache=RulesCacheConfig(enabled=False)), + project_dir=tmp_path, + catalogue=(), + ) + + assert raised.value is test_case.expected_error + assert started == [] + + +@pytest.mark.parametrize( + "test_case", + [ + NativeRequestEncodeErrorTestCase( + description="integer beyond 64 bits", rejected_value=2**70, expected_message="64-bit" + ), + ], + ids=lambda case: case.description, +) +def test_given_value_the_encoder_rejects_when_evaluating_native_then_rules_error_is_raised( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + test_case: NativeRequestEncodeErrorTestCase, +) -> None: + project: CompiledProject = build_project( + name="orders", relative_path="models/orders.sql", sql="SELECT 1", config_values={} + ) + started: list[str] = [] + + def oversized_model_payloads(**_: object) -> list[dict[str, object]]: + return [{"name": "orders", "priority": test_case.rejected_value}] + + def record_custom_start(**_: object) -> None: + started.append("custom") + + def record_native_call(_: str) -> str: + started.append("native") + return "{}" + + monkeypatch.setattr(native, "_model_payloads", oversized_model_payloads) + monkeypatch.setattr(native, "start_custom_rules", record_custom_start) + monkeypatch.setattr(native._native, "evaluate_json", record_native_call) + + with pytest.raises(native.RulesError, match=test_case.expected_message): + native.evaluate_native( + project=project, + config=RulesConfig(cache=RulesCacheConfig(enabled=False)), + project_dir=tmp_path, + catalogue=(), + ) + + assert started == [] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"]))