Skip to content

Commit ad29d02

Browse files
authored
perf(compile): keep native analysis parallel at 10k models (#934)
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.
1 parent e51b2d6 commit ad29d02

7 files changed

Lines changed: 293 additions & 53 deletions

File tree

‎src/sqlbuild/compiler/compile/_helpers/analysis/compact.py‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -180,8 +180,7 @@
180180
from sqlbuild.diagnostics.main.log_debug_event import log_debug_event
181181

182182
_DEBUG_LOGGER: logging.Logger = logging.getLogger("sqlbuild.compile")
183-
_LARGE_COMPACT_PROJECT_MODELS: int = 10_000
184-
_LARGE_COMPACT_SQL_BYTES: int = 48 * 1024 * 1024
183+
_COMPACT_ANALYSIS_WORKERS: int = 4
185184
_NATIVE_LEGACY_FALLBACK: str = "native project type recovery requires legacy fallback"
186185

187186

@@ -886,14 +885,9 @@ def _prepare_compact_analysis_batch(
886885

887886

888887
def _run_compact_analysis_batch(*, preparation: CompactBatchPreparation) -> object:
889-
"""Bound resident analysis heaps before the isolated custom-rule host starts."""
888+
"""Analyse every project size with the full worker pool; the rules phase sets peak memory."""
890889

891-
large_project: bool = (
892-
len(preparation.projections) >= _LARGE_COMPACT_PROJECT_MODELS
893-
and sum(len(sql.encode("utf-8")) for sql in preparation.cleaned_sql)
894-
>= _LARGE_COMPACT_SQL_BYTES
895-
)
896-
workers: int = 1 if large_project else 4
890+
workers: int = _COMPACT_ANALYSIS_WORKERS
897891
if preparation.binding_catalog is not None:
898892
return orjson.loads(
899893
preparation.binding_catalog.native.analyze_compact(

‎src/sqlbuild/rule_engine/_helpers/engine/native.py‎

Lines changed: 97 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import inspect
66
import json
77
from concurrent.futures import Executor, Future, ThreadPoolExecutor
8-
from dataclasses import asdict, replace
8+
from dataclasses import asdict, dataclass, replace
99
from decimal import Decimal
1010
from pathlib import Path
1111
from typing import Any, cast
@@ -54,6 +54,13 @@
5454
from sqlbuild.sql_values.types import SqlValueKind
5555

5656

57+
@dataclass(frozen=True)
58+
class _PreparedNativeRequest:
59+
identity: str | None
60+
reused: str | None
61+
request_text: str | None
62+
63+
5764
def evaluate_native(
5865
*,
5966
project: CompiledProject,
@@ -74,31 +81,20 @@ def evaluate_native(
7481
if custom_payloads is None
7582
else custom_payloads
7683
)
77-
request: dict[str, object] = {
78-
"version": RULES_NATIVE_API_VERSION,
79-
"project_dir": str(project_dir.resolve()),
80-
"dialect": dialect,
81-
"config": _config_payload(config),
82-
"models": _model_payloads(
83-
project=project,
84-
dialect=dialect,
85-
include_type_proof=any(
86-
_rule_selected(config=config, code=code) for code in TYPE_PROOF_RULE_CODES
87-
),
84+
prepared: _PreparedNativeRequest = _prepare_native_request(
85+
request_json=_serialise_native_request(
86+
_native_request(
87+
project=project,
88+
config=config,
89+
project_dir=project_dir,
90+
dialect=dialect,
91+
initial_findings=initial_findings,
92+
custom_payloads=payloads,
93+
)
8894
),
89-
"sql_tests": _sql_test_payloads(project),
90-
"sql_scenarios": _sql_scenario_payloads(project),
91-
"public_enums": [
92-
_enum_payload(declaration) for declaration in project.public_enums.values()
93-
],
94-
"public_constants": [
95-
_constant_payload(declaration) for declaration in project.public_constants.values()
96-
],
97-
"scope_index": scope_metadata_projection(index=project.scope_index),
98-
"initial_findings": [_finding_payload(finding) for finding in initial_findings],
99-
"defer_suppressions": True,
100-
"custom_rules": payloads,
101-
}
95+
project_dir=project_dir,
96+
cache_enabled=config.cache.enabled,
97+
)
10298
with ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-custom-rules") as pool:
10399
custom_future: Future[CustomRulesOutcome] | None = (
104100
custom_outcome
@@ -114,15 +110,14 @@ def evaluate_native(
114110
verify_determinism=verify_determinism,
115111
)
116112
)
113+
reused: bool = prepared.reused is not None
117114
try:
118-
response_json: str
119-
reused: bool
120-
response_json, reused = _evaluate_request(
121-
request=request, project_dir=project_dir, cache_enabled=config.cache.enabled
115+
response: object = orjson.loads(
116+
_evaluate_request(prepared=prepared, project_dir=project_dir)
122117
)
123-
response: object = orjson.loads(response_json)
124118
except (ValueError, TypeError) as error:
125119
raise RulesError(str(error)) from error
120+
del prepared
126121
custom: CustomRulesOutcome = (
127122
CustomRulesOutcome(findings=(), cache_hits=0, cache_misses=0, custom_ms=0)
128123
if custom_future is None
@@ -159,6 +154,42 @@ def evaluate_native(
159154
)
160155

161156

157+
def _native_request(
158+
*,
159+
project: CompiledProject,
160+
config: RulesConfig,
161+
project_dir: Path,
162+
dialect: str,
163+
initial_findings: tuple[Finding, ...],
164+
custom_payloads: list[dict[str, object]],
165+
) -> dict[str, object]:
166+
return {
167+
"version": RULES_NATIVE_API_VERSION,
168+
"project_dir": str(project_dir.resolve()),
169+
"dialect": dialect,
170+
"config": _config_payload(config),
171+
"models": _model_payloads(
172+
project=project,
173+
dialect=dialect,
174+
include_type_proof=any(
175+
_rule_selected(config=config, code=code) for code in TYPE_PROOF_RULE_CODES
176+
),
177+
),
178+
"sql_tests": _sql_test_payloads(project),
179+
"sql_scenarios": _sql_scenario_payloads(project),
180+
"public_enums": [
181+
_enum_payload(declaration) for declaration in project.public_enums.values()
182+
],
183+
"public_constants": [
184+
_constant_payload(declaration) for declaration in project.public_constants.values()
185+
],
186+
"scope_index": scope_metadata_projection(index=project.scope_index),
187+
"initial_findings": [_finding_payload(finding) for finding in initial_findings],
188+
"defer_suppressions": True,
189+
"custom_rules": custom_payloads,
190+
}
191+
192+
162193
def start_custom_rules(
163194
*,
164195
executor: Executor,
@@ -206,19 +237,42 @@ def _selected_custom_rules(
206237
return tuple(rule for rule in catalogue if rule.custom and rule.code in selected)
207238

208239

209-
def _evaluate_request(
210-
*, request: dict[str, object], project_dir: Path, cache_enabled: bool
211-
) -> tuple[str, bool]:
212-
request_json: bytes = orjson.dumps(request, option=orjson.OPT_SORT_KEYS, default=str)
213-
if not cache_enabled:
214-
return _native.evaluate_json(request_json.decode()), False
215-
identity: str = native_request_identity(request_json)
216-
reused: str | None = read_native_response(project_dir=project_dir, identity=identity)
217-
if reused is not None:
218-
return reused, True
219-
response: str = _native.evaluate_json(request_json.decode())
220-
write_native_response(project_dir=project_dir, identity=identity, response=response)
221-
return response, False
240+
def _serialise_native_request(request: dict[str, object]) -> bytes:
241+
"""Encode the native request, reporting values the encoder rejects as rules errors."""
242+
243+
try:
244+
return orjson.dumps(request, option=orjson.OPT_SORT_KEYS, default=str)
245+
except orjson.JSONEncodeError as error:
246+
raise RulesError(str(error)) from error
247+
248+
249+
def _prepare_native_request(
250+
*, request_json: bytes, project_dir: Path, cache_enabled: bool
251+
) -> _PreparedNativeRequest:
252+
"""Resolve the memo and decode once; the caller's bytes die when this returns."""
253+
254+
identity: str | None = native_request_identity(request_json) if cache_enabled else None
255+
reused: str | None = (
256+
read_native_response(project_dir=project_dir, identity=identity)
257+
if identity is not None
258+
else None
259+
)
260+
return _PreparedNativeRequest(
261+
identity=identity,
262+
reused=reused,
263+
request_text=request_json.decode() if reused is None else None,
264+
)
265+
266+
267+
def _evaluate_request(*, prepared: _PreparedNativeRequest, project_dir: Path) -> str:
268+
if prepared.reused is not None:
269+
return prepared.reused
270+
response: str = _native.evaluate_json(cast(str, prepared.request_text))
271+
if prepared.identity is not None:
272+
write_native_response(
273+
project_dir=project_dir, identity=prepared.identity, response=response
274+
)
275+
return response
222276

223277

224278
def finalize_native_findings(

‎tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_compile_performance.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@
6262
"dense_models_10000_all_rules",
6363
10000,
6464
145.0,
65-
4 * _GIB,
65+
7 * _GIB // 2,
6666
"913a08c2277304efacf75e293728767a3fa416f6b32f5488dca6ab5baea80668",
6767
),
6868
),

‎tests/unit/src/sqlbuild/compiler/compile/_helpers/analysis/_test_types.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,3 +14,11 @@ class DynamicPivotAnalysisTestCase:
1414
expected_family_types: tuple[str | None, ...] = ("DECIMAL(12,2)",)
1515
expected_input_relations: tuple[str, ...] = ()
1616
aggregate: str = "MAX"
17+
18+
19+
@dataclass(frozen=True)
20+
class CompactAnalysisWorkersTestCase:
21+
description: str
22+
model_count: int
23+
sql_bytes_per_model: int
24+
expected_workers: int
Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
from __future__ import annotations
2+
3+
import orjson
4+
import pytest
5+
6+
from sqlbuild.compiler.compile._helpers.analysis import compact
7+
from sqlbuild.compiler.compile.models import CompactBatchPreparation
8+
from tests.unit.src.sqlbuild.compiler.compile._helpers.analysis._test_types import (
9+
CompactAnalysisWorkersTestCase,
10+
)
11+
12+
_MIB: int = 1024 * 1024
13+
14+
15+
@pytest.mark.parametrize(
16+
"test_case",
17+
(
18+
CompactAnalysisWorkersTestCase(
19+
description="small project",
20+
model_count=3,
21+
sql_bytes_per_model=64,
22+
expected_workers=4,
23+
),
24+
CompactAnalysisWorkersTestCase(
25+
description="dense ten thousand model project",
26+
model_count=10_000,
27+
sql_bytes_per_model=(64 * _MIB) // 10_000,
28+
expected_workers=4,
29+
),
30+
),
31+
ids=lambda case: case.description,
32+
)
33+
def test_given_project_size_when_running_compact_analysis_then_native_uses_all_workers(
34+
test_case: CompactAnalysisWorkersTestCase,
35+
monkeypatch: pytest.MonkeyPatch,
36+
) -> None:
37+
captured: list[dict[str, object]] = []
38+
39+
def analyze(request_json: str) -> str:
40+
captured.append(orjson.loads(request_json))
41+
return "{}"
42+
43+
monkeypatch.setattr(compact._native, "analyze_project_queries_compact_json", analyze)
44+
sql: str = "SELECT 1 AS order_id" + " " * (test_case.sql_bytes_per_model - 20)
45+
preparation: CompactBatchPreparation = CompactBatchPreparation(
46+
cleaned_sql=(sql,) * test_case.model_count,
47+
queries=({"sql": "SELECT 1 AS order_id", "dialect": "duckdb"},),
48+
templates=({"queryIndex": 0},),
49+
projections=({"templateIndex": 0},) * test_case.model_count,
50+
)
51+
52+
_ = compact._run_compact_analysis_batch(preparation=preparation)
53+
54+
assert [request["workers"] for request in captured] == [test_case.expected_workers]
55+
56+
57+
if __name__ == "__main__":
58+
raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"]))

‎tests/unit/src/sqlbuild/rule_engine/_helpers/engine/_test_types.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,3 +202,20 @@ class NativeBuildIdentityTestCase:
202202
description: str
203203
expected_rebuilt_evaluations: int
204204
expected_warm_evaluations: int
205+
206+
207+
@dataclass(frozen=True)
208+
class NativeRequestBuildErrorTestCase:
209+
"""One programming error raised while building the native rules request."""
210+
211+
description: str
212+
expected_error: Exception
213+
214+
215+
@dataclass(frozen=True)
216+
class NativeRequestEncodeErrorTestCase:
217+
"""One request value the JSON encoder rejects."""
218+
219+
description: str
220+
rejected_value: object
221+
expected_message: str

0 commit comments

Comments
 (0)