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"]))