From 7ac6339f3dd112798e447634e2b3896db89c298d Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sat, 3 Oct 2026 22:46:08 +0100 Subject: [PATCH 1/3] perf(rules): split custom rules across host processes and start them first Custom-rule invocations are dealt round-robin by subject across up to one host process per available core, then merged back into single-host order so findings, read sets, untracked and uncacheable rules, observed digests, and the custom-rule cache are identical to one host. Any host failure reruns the whole plan in one host so the reported error is exactly the single-host error. Rule evaluation now launches the custom hosts and native built-in rules before preparing SQL rule inputs, and computes custom rule payloads and implementation fingerprints once instead of up to three times. --- .../rule_engine/_helpers/engine/catalogue.py | 11 +- .../_helpers/engine/custom_rules.py | 185 +++-------- .../rule_engine/_helpers/engine/native.py | 77 ++++- .../_helpers/host/custom_evaluation.py | 4 + .../rule_engine/_helpers/host/custom_host.py | 4 +- .../_helpers/host/custom_host_pool.py | 293 ++++++++++++++++++ .../rule_engine/_helpers/run/rules.py | 55 +++- src/sqlbuild/rule_engine/constants.py | 1 + src/sqlbuild/rule_engine/models.py | 31 +- src/sqlbuild/rule_engine/types.py | 2 + .../sqlbuild/cli/commands/main/_test_types.py | 10 + .../sqlbuild/cli/commands/main/test_rules.py | 80 ++++- .../_helpers/engine/test_custom_rule_reads.py | 4 +- .../rule_engine/_helpers/host/_test_types.py | 35 +++ .../rule_engine/_helpers/host/helpers.py | 100 ++++++ .../_helpers/host/test_custom_host_pool.py | 179 +++++++++++ 16 files changed, 891 insertions(+), 180 deletions(-) create mode 100644 src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py create mode 100644 tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py create mode 100644 tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py create mode 100644 tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py diff --git a/src/sqlbuild/rule_engine/_helpers/engine/catalogue.py b/src/sqlbuild/rule_engine/_helpers/engine/catalogue.py index 626acd3757..1ece17e899 100644 --- a/src/sqlbuild/rule_engine/_helpers/engine/catalogue.py +++ b/src/sqlbuild/rule_engine/_helpers/engine/catalogue.py @@ -174,13 +174,20 @@ def _native_builtin(*, model: object, ctx: object) -> list[Finding]: def select_rules( - *, catalogue: tuple[Rule, ...], config: RulesConfig, project_dir: Path + *, + catalogue: tuple[Rule, ...], + config: RulesConfig, + project_dir: Path, + custom_payloads: list[dict[str, object]] | None = None, ) -> tuple[Rule, ...]: """Resolve the active catalogue through the native Fensu rules owner.""" by_code: dict[str, Rule] = {rule.code: rule for rule in catalogue} codes: tuple[str, ...] = native_selected_codes( - config=config, catalogue=catalogue, project_dir=project_dir + config=config, + catalogue=catalogue, + project_dir=project_dir, + custom_payloads=custom_payloads, ) try: return tuple(by_code[code] for code in codes) diff --git a/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py b/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py index 4b500db707..9833dc1e18 100644 --- a/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py +++ b/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py @@ -4,19 +4,15 @@ import hashlib import os -import pickle -import sys -import tempfile import threading import time from collections.abc import Callable from dataclasses import asdict, dataclass, replace from pathlib import Path -from typing import Any, cast +from typing import Any import orjson -import sqlbuild._native as _native from sqlbuild.compiler.compile.models import CompiledProject from sqlbuild.compiler.fact_cache.main.code_identity import compiled_code_identity from sqlbuild.rule_engine._helpers.engine.custom_rule_evidence import ( @@ -24,21 +20,29 @@ custom_rule_import_closure, ) from sqlbuild.rule_engine._helpers.engine.fact_replay import full_fact_keys, is_tree_fact +from sqlbuild.rule_engine._helpers.host.custom_host_pool import ( + decode_readset, + finding_payload, + run_custom_hosts, +) from sqlbuild.rule_engine.classes.fact_digests import FactDigests from sqlbuild.rule_engine.classes.rule_context import RuleFactViews, build_rule_fact_views from sqlbuild.rule_engine.constants import ( - CUSTOM_HOST_LAUNCH_MODULE, CUSTOM_HOST_RUNTIME_VERSION, + CUSTOM_RULE_PROJECT_SUBJECT, CUSTOM_RULES_CACHE_FILE, CUSTOM_RULES_CACHE_VERSION, ) from sqlbuild.rule_engine.exceptions import FactDigestError, RulesError -from sqlbuild.rule_engine.models import CustomRulesOutcome, Rule, RulesConfig +from sqlbuild.rule_engine.models import ( + CustomHostEvaluation, + CustomHostRun, + CustomRulesOutcome, + Rule, + RulesConfig, +) from sqlbuild.rule_engine.types import FactKey, RuleSubject -_HOST_TIMEOUT_MILLIS: int = 120_000 -_PROJECT_SUBJECT: str = "" - @dataclass(frozen=True) class _Subject: @@ -46,20 +50,6 @@ class _Subject: identity: str -@dataclass(frozen=True) -class _Evaluated: - findings: tuple[dict[str, object], ...] - reads: tuple[FactKey, ...] | None - - -@dataclass(frozen=True) -class _HostRun: - evaluated: dict[tuple[str, str], _Evaluated] - untracked_reads: dict[str, tuple[FactKey, ...]] - uncacheable: frozenset[str] - observed: dict[FactKey, str | None] - - @dataclass(frozen=True) class _CachedEntry: subject_identity: str @@ -83,6 +73,7 @@ def evaluate_custom_rules_cached( rules: tuple[Rule, ...], dialect: str, verify_determinism: bool = False, + implementation_fingerprints: dict[str, str] | None = None, ) -> CustomRulesOutcome: """Reuse results whose recorded fact reads are unchanged and evaluate only the rest.""" @@ -91,11 +82,13 @@ def evaluate_custom_rules_cached( if not custom: return CustomRulesOutcome(findings=(), cache_hits=0, cache_misses=0, custom_ms=0) if not config.cache.enabled: - uncached: _HostRun = _run_host( - project=project, + uncached: CustomHostRun = run_custom_hosts( + project=_host_project(project), config=config, project_dir=project_dir, dialect=dialect, + rules=custom, + model_paths=_model_paths(project), plan={rule.code: None for rule in custom}, track_reads=False, verify_determinism=verify_determinism, @@ -116,7 +109,11 @@ def evaluate_custom_rules_cached( subjects: dict[str, tuple[_Subject, ...]] = _subjects(project=project, rules=custom) identities: dict[str, str] = { rule.code: _rule_identity( - rule=rule, config=config, project_dir=project_dir, dialect=dialect + rule=rule, + config=config, + project_dir=project_dir, + dialect=dialect, + fingerprint=(implementation_fingerprints or {}).get(rule.code), ) for rule in custom } @@ -152,11 +149,13 @@ def evaluate_custom_rules_cached( len(subjects[code]) if planned is None else len(planned) for code, planned in plan.items() ) if plan: - run: _HostRun = _run_host( - project=project, + run: CustomHostRun = run_custom_hosts( + project=_host_project(project), config=config, project_dir=project_dir, dialect=dialect, + rules=custom, + model_paths=_model_paths(project), plan=plan, track_reads=True, verify_determinism=verify_determinism, @@ -200,7 +199,7 @@ def _subjects( ) -> dict[str, tuple[_Subject, ...]]: project_subject: tuple[_Subject, ...] = ( _Subject( - key=_PROJECT_SUBJECT, + key=CUSTOM_RULE_PROJECT_SUBJECT, identity=_digest([len(project.models), project.effective_target_name]), ), ) @@ -217,13 +216,18 @@ def _subjects( } +def _model_paths(project: CompiledProject) -> tuple[str, ...]: + return tuple(sorted(model.relative_path.as_posix() for model in project.models)) + + def _materialization(model: Any) -> str | None: value: object = model.config.values.get("materialized") return value if isinstance(value, str) else None -def _rule_identity(*, rule: Rule, config: RulesConfig, project_dir: Path, dialect: str) -> str: - closure: tuple[Path, ...] = custom_rule_import_closure(rule=rule, project_dir=project_dir) +def _rule_identity( + *, rule: Rule, config: RulesConfig, project_dir: Path, dialect: str, fingerprint: str | None +) -> str: return _digest( [ CUSTOM_RULES_CACHE_VERSION, @@ -235,8 +239,12 @@ def _rule_identity(*, rule: Rule, config: RulesConfig, project_dir: Path, dialec rule.message, rule.remediation, None if rule.subject is None else rule.subject.value, - custom_rule_implementation_fingerprint( - rule=rule, project_dir=project_dir, import_closure=closure + fingerprint + if fingerprint is not None + else custom_rule_implementation_fingerprint( + rule=rule, + project_dir=project_dir, + import_closure=custom_rule_import_closure(rule=rule, project_dir=project_dir), ), asdict(config), dialect, @@ -260,51 +268,6 @@ def _readset_digest( return memo[index] -def _run_host( - *, - project: CompiledProject, - config: RulesConfig, - project_dir: Path, - dialect: str, - plan: dict[str, list[str] | None], - track_reads: bool, - verify_determinism: bool, -) -> _HostRun: - input_dir: Path = project_dir / "target" / "rules-cache" / "host-inputs" - input_dir.mkdir(parents=True, exist_ok=True) - input_path: Path | None = None - try: - with tempfile.NamedTemporaryFile( - dir=input_dir, prefix="project-", suffix=".pickle", delete=False - ) as handle: - input_path = Path(handle.name) - pickle.dump((_host_project(project), config), handle) - response_json: str = _native.run_custom_host_json( - orjson.dumps( - { - "program": sys.executable, - "arguments": ["-m", CUSTOM_HOST_LAUNCH_MODULE], - "timeout_millis": _HOST_TIMEOUT_MILLIS, - "runtime_version": CUSTOM_HOST_RUNTIME_VERSION, - "payload": { - "project_pickle_path": str(input_path.resolve()), - "project_dir": str(project_dir.resolve()), - "dialect": dialect, - "verify_determinism": verify_determinism, - "track_reads": track_reads, - "plan": plan, - }, - } - ).decode() - ) - except (ValueError, TypeError) as error: - raise RulesError(str(error)) from error - finally: - if input_path is not None: - input_path.unlink(missing_ok=True) - return _decode_host_response(orjson.loads(response_json)) - - def _host_project(project: CompiledProject) -> CompiledProject: return replace( project, @@ -318,59 +281,11 @@ def _host_project(project: CompiledProject) -> CompiledProject: ) -def _decode_host_response(response: object) -> _HostRun: - if not isinstance(response, dict): - raise RulesError("custom rule host returned an invalid payload") - payload: dict[str, Any] = cast(dict[str, Any], response) - raw_readsets: Any = payload.get("readsets") - raw_evaluations: Any = payload.get("evaluations") - raw_untracked: Any = payload.get("untracked") - raw_uncacheable: Any = payload.get("uncacheable") - raw_observed: Any = payload.get("observed") - if ( - not isinstance(raw_readsets, list) - or not isinstance(raw_evaluations, list) - or not isinstance(raw_untracked, dict) - or not isinstance(raw_uncacheable, list) - or not isinstance(raw_observed, list) - ): - raise RulesError("custom rule host returned an invalid payload") - readsets: list[tuple[FactKey, ...]] = list(map(_decode_readset, raw_readsets)) - untracked: dict[str, tuple[FactKey, ...]] = { - str(code): _decode_readset(reads) for code, reads in raw_untracked.items() - } - uncacheable: frozenset[str] = frozenset(str(code) for code in raw_uncacheable) - evaluated: dict[tuple[str, str], _Evaluated] = {} - for raw_item in raw_evaluations: - if not isinstance(raw_item, dict): - raise RulesError("custom rule host returned an invalid evaluation") - item: dict[str, Any] = cast(dict[str, Any], raw_item) - code: str = str(item["code"]) - subject: object = item.get("subject") - findings: Any = item.get("findings") - reads: object = item.get("reads") - if not isinstance(findings, list): - raise RulesError("custom rule host returned invalid findings") - evaluated[(code, _PROJECT_SUBJECT if subject is None else str(subject))] = _Evaluated( - findings=tuple(map(_finding_payload, findings)), - reads=(readsets[reads] if isinstance(reads, int) and code not in untracked else None), - ) - return _HostRun( - evaluated=evaluated, - untracked_reads=untracked, - uncacheable=uncacheable, - observed={ - tuple(map(str, key)): None if digest is None else str(digest) - for key, digest in raw_observed - }, - ) - - def _validate_paths( *, project: CompiledProject, rules: tuple[Rule, ...], - evaluated: dict[tuple[str, str], _Evaluated], + evaluated: dict[tuple[str, str], CustomHostEvaluation], ) -> None: codes: frozenset[str] = frozenset(rule.code for rule in rules) model_rules: frozenset[str] = frozenset( @@ -397,7 +312,7 @@ def _store( stored: dict[str, _RuleCache], identities: dict[str, str], subjects: dict[str, tuple[_Subject, ...]], - run: _HostRun, + run: CustomHostRun, digests: FactDigests, ) -> dict[str, _RuleCache]: updated: dict[str, _RuleCache] = {} @@ -495,7 +410,7 @@ def _read_cache(path: Path) -> dict[str, _RuleCache]: try: stored[str(code)] = _RuleCache( identity=str(value["identity"]), - readsets=list(map(_decode_readset, value["readsets"])), + readsets=list(map(decode_readset, value["readsets"])), entries={ str(subject): _decode_entry(entry) for subject, entry in value["entries"].items() @@ -538,23 +453,15 @@ def _write_cache( temporary.replace(path) -def _decode_readset(value: Any) -> tuple[FactKey, ...]: - return tuple(tuple(map(str, key)) for key in value) - - def _decode_entry(value: Any) -> _CachedEntry: return _CachedEntry( subject_identity=str(value[0]), readset=int(value[1]), reads_digest=str(value[2]), - findings=tuple(map(_finding_payload, value[3])), + findings=tuple(map(finding_payload, value[3])), ) -def _finding_payload(value: Any) -> dict[str, object]: - return {str(key): item for key, item in value.items()} - - def _digest(value: object) -> str: return hashlib.sha256(orjson.dumps(value, default=str)).hexdigest() diff --git a/src/sqlbuild/rule_engine/_helpers/engine/native.py b/src/sqlbuild/rule_engine/_helpers/engine/native.py index a609ba1cc8..dcf33baf81 100644 --- a/src/sqlbuild/rule_engine/_helpers/engine/native.py +++ b/src/sqlbuild/rule_engine/_helpers/engine/native.py @@ -4,7 +4,7 @@ import inspect import json -from concurrent.futures import Future, ThreadPoolExecutor +from concurrent.futures import Executor, Future, ThreadPoolExecutor from dataclasses import asdict, replace from decimal import Decimal from pathlib import Path @@ -64,14 +64,15 @@ def evaluate_native( initial_findings: tuple[Finding, ...] = (), defer_suppressions: bool = False, verify_determinism: bool = False, + custom_payloads: list[dict[str, object]] | None = None, + custom_outcome: Future[CustomRulesOutcome] | None = None, ) -> RulesResult: """Evaluate built-in rules natively while custom rules run incrementally beside them.""" - custom_payloads: list[dict[str, object]] = _custom_rule_payloads( - catalogue=catalogue, project_dir=project_dir - ) - custom_rules: tuple[Rule, ...] = _selected_custom_rules( - config=config, catalogue=catalogue, custom_payloads=custom_payloads + payloads: list[dict[str, object]] = ( + custom_rule_payloads(catalogue=catalogue, project_dir=project_dir) + if custom_payloads is None + else custom_payloads ) request: dict[str, object] = { "version": RULES_NATIVE_API_VERSION, @@ -96,21 +97,22 @@ def evaluate_native( "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, + "custom_rules": payloads, } with ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-custom-rules") as pool: custom_future: Future[CustomRulesOutcome] | None = ( - pool.submit( - evaluate_custom_rules_cached, + custom_outcome + if custom_payloads is not None + else start_custom_rules( + executor=pool, project=project, config=config, - project_dir=project_dir.resolve(), - rules=custom_rules, + project_dir=project_dir, + catalogue=catalogue, + custom_payloads=payloads, dialect=dialect, verify_determinism=verify_determinism, ) - if custom_rules - else None ) try: response_json: str @@ -157,6 +159,39 @@ def evaluate_native( ) +def start_custom_rules( + *, + executor: Executor, + project: CompiledProject, + config: RulesConfig, + project_dir: Path, + catalogue: tuple[Rule, ...], + custom_payloads: list[dict[str, object]], + dialect: str, + verify_determinism: bool = False, +) -> Future[CustomRulesOutcome] | None: + """Submit selected custom rules, reusing the implementation fingerprints of their payloads.""" + + custom_rules: tuple[Rule, ...] = _selected_custom_rules( + config=config, catalogue=catalogue, custom_payloads=custom_payloads + ) + if not custom_rules: + return None + return executor.submit( + evaluate_custom_rules_cached, + project=project, + config=config, + project_dir=project_dir.resolve(), + rules=custom_rules, + dialect=dialect, + verify_determinism=verify_determinism, + implementation_fingerprints={ + str(payload["code"]): str(payload["implementation_fingerprint"]) + for payload in custom_payloads + }, + ) + + def _selected_custom_rules( *, config: RulesConfig, @@ -261,13 +296,21 @@ def native_catalogue() -> tuple[dict[str, object], ...]: def native_selected_codes( - *, config: RulesConfig, catalogue: tuple[Rule, ...], project_dir: Path + *, + config: RulesConfig, + catalogue: tuple[Rule, ...], + project_dir: Path, + custom_payloads: list[dict[str, object]] | None = None, ) -> tuple[str, ...]: """Resolve the active ruleset through the native Fensu adapter.""" return _selected_codes( config=config, - custom_payloads=_custom_rule_payloads(catalogue=catalogue, project_dir=project_dir), + custom_payloads=( + custom_rule_payloads(catalogue=catalogue, project_dir=project_dir) + if custom_payloads is None + else custom_payloads + ), ) @@ -572,9 +615,11 @@ def _typed_value_payload(value: SqlValue) -> object: } -def _custom_rule_payloads( +def custom_rule_payloads( *, catalogue: tuple[Rule, ...], project_dir: Path ) -> list[dict[str, object]]: + """Describe every catalogued custom rule for native selection and evaluation.""" + closures: dict[str, tuple[Path, ...]] = {} return [ _custom_rule_payload( diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py b/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py index 17fbe73fe7..8884692f01 100644 --- a/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py @@ -22,6 +22,7 @@ build_rule_fact_views, ) from sqlbuild.rule_engine.classes.runtime_guard import RuntimeGuard +from sqlbuild.rule_engine.constants import CUSTOM_RULE_PROJECT_SUBJECT from sqlbuild.rule_engine.exceptions import ( NonHermeticRuleError, OpaqueModuleStateError, @@ -211,6 +212,9 @@ def context(rule: Rule) -> EvaluationRuleContext: (rule, context(rule), None if plan is None else plan.get(rule.code)) for rule in model_rules ) for rule in project_rules: + planned_project: frozenset[str] | None = None if plan is None else plan.get(rule.code) + if planned_project is not None and CUSTOM_RULE_PROJECT_SUBJECT not in planned_project: + continue ctx: EvaluationRuleContext = context(rule) tracker.start() try: diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_host.py b/src/sqlbuild/rule_engine/_helpers/host/custom_host.py index ce315b2a0b..12701166b9 100644 --- a/src/sqlbuild/rule_engine/_helpers/host/custom_host.py +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_host.py @@ -145,7 +145,8 @@ def _run_payload(run: CustomRuleRun) -> dict[str, object]: "reads": index, } ) - if sum(map(len, readsets)) > CUSTOM_HOST_MAX_TRACKED_READS: + bounded: bool = sum(map(len, readsets)) > CUSTOM_HOST_MAX_TRACKED_READS + if bounded: untracked.update(str(item["code"]) for item in evaluations if item["reads"] is not None) readsets = {} for item in evaluations: @@ -162,6 +163,7 @@ def _run_payload(run: CustomRuleRun) -> dict[str, object]: }, "uncacheable": sorted(run.uncacheable_codes), "observed": sorted([list(key), digest] for key, digest in run.observed), + "bounded": bounded, } diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py b/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py new file mode 100644 index 0000000000..9031eeca15 --- /dev/null +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py @@ -0,0 +1,293 @@ +"""Run planned custom-rule invocations across isolated host processes and merge their results.""" + +from __future__ import annotations + +import os +import pickle +import sys +import tempfile +from concurrent.futures import Future, ThreadPoolExecutor +from pathlib import Path +from typing import Any, cast + +import orjson + +import sqlbuild._native as _native +from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.rule_engine.constants import ( + CUSTOM_HOST_LAUNCH_MODULE, + CUSTOM_HOST_MAX_TRACKED_READS, + CUSTOM_HOST_RUNTIME_VERSION, + CUSTOM_RULE_PROJECT_SUBJECT, +) +from sqlbuild.rule_engine.exceptions import RulesError +from sqlbuild.rule_engine.models import ( + CustomHostEvaluation, + CustomHostRun, + CustomHostSlice, + Rule, + RulesConfig, +) +from sqlbuild.rule_engine.types import CustomHostPlan, FactKey, RuleSubject + +_HOST_TIMEOUT_MILLIS: int = 120_000 +_MIN_INVOCATIONS_PER_HOST: int = 64 + + +def run_custom_hosts( + *, + project: CompiledProject, + config: RulesConfig, + project_dir: Path, + dialect: str, + rules: tuple[Rule, ...], + model_paths: tuple[str, ...], + plan: CustomHostPlan, + track_reads: bool, + verify_determinism: bool, +) -> CustomHostRun: + """Evaluate the plan across host processes and merge their results in single-host order.""" + + invocations: tuple[tuple[str, str], ...] = planned_invocations( + plan=plan, rules=rules, model_paths=model_paths + ) + plans: tuple[CustomHostPlan, ...] = ( + (plan,) + if verify_determinism + else partition_host_plan( + plan=plan, invocations=invocations, hosts=_host_count(len(invocations)) + ) + ) + input_dir: Path = project_dir / "target" / "rules-cache" / "host-inputs" + input_dir.mkdir(parents=True, exist_ok=True) + input_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + dir=input_dir, prefix="project-", suffix=".pickle", delete=False + ) as handle: + input_path = Path(handle.name) + pickle.dump((project, config), handle) + payload: dict[str, object] = { + "project_pickle_path": str(input_path.resolve()), + "project_dir": str(project_dir.resolve()), + "dialect": dialect, + "verify_determinism": verify_determinism, + "track_reads": track_reads, + } + responses: list[CustomHostSlice] = ( + [_run_host(payload={**payload, "plan": plan})] + if len(plans) == 1 + else _run_partitions(payload=payload, plan=plan, plans=plans) + ) + finally: + if input_path is not None: + input_path.unlink(missing_ok=True) + return _merge( + responses=responses, + order={invocation: index for index, invocation in enumerate(invocations)}, + ) + + +def _run_partitions( + *, payload: dict[str, object], plan: CustomHostPlan, plans: tuple[CustomHostPlan, ...] +) -> list[CustomHostSlice]: + """Run every partition; after any failure, one host reruns the whole plan to report it.""" + + with ThreadPoolExecutor( + max_workers=len(plans), thread_name_prefix="sqlbuild-custom-host" + ) as pool: + futures: list[Future[CustomHostSlice]] = [ + pool.submit(_run_host, payload={**payload, "plan": partition}) for partition in plans + ] + failed: bool = any(future.exception() is not None for future in futures) + if failed: + return [_run_host(payload={**payload, "plan": plan})] + return [future.result() for future in futures] + + +def planned_invocations( + *, plan: CustomHostPlan, rules: tuple[Rule, ...], model_paths: tuple[str, ...] +) -> tuple[tuple[str, str], ...]: + """List planned (rule, subject) invocations in the order one host evaluates them.""" + + planned: tuple[Rule, ...] = tuple( + sorted((rule for rule in rules if rule.code in plan), key=lambda rule: rule.code) + ) + invocations: list[tuple[str, str]] = [ + (rule.code, CUSTOM_RULE_PROJECT_SUBJECT) + for rule in planned + if rule.subject is RuleSubject.PROJECT + ] + model_rules: tuple[tuple[str, frozenset[str] | None], ...] = tuple( + (rule.code, None if plan[rule.code] is None else frozenset(plan[rule.code] or ())) + for rule in planned + if rule.subject is RuleSubject.MODEL + ) + for path in model_paths if model_rules else (): + invocations.extend( + (code, path) for code, selected in model_rules if selected is None or path in selected + ) + return tuple(invocations) + + +def partition_host_plan( + *, plan: CustomHostPlan, invocations: tuple[tuple[str, str], ...], hosts: int +) -> tuple[CustomHostPlan, ...]: + """Deal project invocations and whole model subjects round-robin across at most `hosts`.""" + + units: list[list[tuple[str, str]]] = [] + for code, subject in invocations: + if subject == CUSTOM_RULE_PROJECT_SUBJECT or not units or units[-1][0][1] != subject: + units.append([]) + units[-1].append((code, subject)) + count: int = max(1, min(hosts, len(units))) + if count == 1: + return (plan,) + partitions: list[dict[str, list[str]]] = [] + for _ in range(count): + partitions.append({code: [] for code in plan}) + for index, unit in enumerate(units): + for code, subject in unit: + partitions[index % count][code].append(subject) + return tuple(cast(CustomHostPlan, partition) for partition in partitions) + + +def _host_count(invocations: int) -> int: + return max(1, min(_available_cores(), invocations // _MIN_INVOCATIONS_PER_HOST)) + + +def _available_cores() -> int: + affinity: Any = getattr(os, "sched_getaffinity", None) + return len(affinity(0)) if affinity is not None else os.cpu_count() or 1 + + +def _run_host(*, payload: dict[str, object]) -> CustomHostSlice: + try: + response_json: str = _native.run_custom_host_json( + orjson.dumps( + { + "program": sys.executable, + "arguments": ["-m", CUSTOM_HOST_LAUNCH_MODULE], + "timeout_millis": _HOST_TIMEOUT_MILLIS, + "runtime_version": CUSTOM_HOST_RUNTIME_VERSION, + "payload": payload, + } + ).decode() + ) + except (ValueError, TypeError) as error: + raise RulesError(str(error)) from error + return _decode_host_response(orjson.loads(response_json)) + + +def _decode_host_response(response: object) -> CustomHostSlice: + if not isinstance(response, dict): + raise RulesError("custom rule host returned an invalid payload") + payload: dict[str, Any] = cast(dict[str, Any], response) + raw_readsets: Any = payload.get("readsets") + raw_evaluations: Any = payload.get("evaluations") + raw_untracked: Any = payload.get("untracked") + raw_uncacheable: Any = payload.get("uncacheable") + raw_observed: Any = payload.get("observed") + if ( + not isinstance(raw_readsets, list) + or not isinstance(raw_evaluations, list) + or not isinstance(raw_untracked, dict) + or not isinstance(raw_uncacheable, list) + or not isinstance(raw_observed, list) + ): + raise RulesError("custom rule host returned an invalid payload") + readsets: list[tuple[FactKey, ...]] = list(map(decode_readset, raw_readsets)) + untracked: dict[str, tuple[FactKey, ...]] = { + str(code): decode_readset(reads) for code, reads in raw_untracked.items() + } + evaluations: list[tuple[str, str, CustomHostEvaluation]] = [] + for raw_item in raw_evaluations: + if not isinstance(raw_item, dict): + raise RulesError("custom rule host returned an invalid evaluation") + item: dict[str, Any] = cast(dict[str, Any], raw_item) + code: str = str(item["code"]) + subject: object = item.get("subject") + findings: Any = item.get("findings") + reads: object = item.get("reads") + if not isinstance(findings, list): + raise RulesError("custom rule host returned invalid findings") + tracked: bool = isinstance(reads, int) and code not in untracked + evaluations.append( + ( + code, + CUSTOM_RULE_PROJECT_SUBJECT if subject is None else str(subject), + CustomHostEvaluation( + findings=tuple(map(finding_payload, findings)), + reads=readsets[cast(int, reads)] if tracked else None, + ), + ) + ) + return CustomHostSlice( + evaluations=tuple(evaluations), + untracked_reads=untracked, + uncacheable=frozenset(str(code) for code in raw_uncacheable), + observed={ + tuple(map(str, key)): None if digest is None else str(digest) + for key, digest in raw_observed + }, + bounded=payload.get("bounded") is True, + ) + + +def _merge(*, responses: list[CustomHostSlice], order: dict[tuple[str, str], int]) -> CustomHostRun: + """Combine host slices exactly as one host would have reported every invocation.""" + + evaluations: list[tuple[str, str, CustomHostEvaluation]] = [] + untracked: set[str] = set() + for response in responses: + evaluations.extend(response.evaluations) + untracked.update(response.untracked_reads) + if len(responses) > 1: + evaluations.sort(key=lambda item: order.get((item[0], item[1]), len(order))) + tracked_readsets: set[tuple[FactKey, ...]] = { + result.reads + for code, _subject, result in evaluations + if result.reads is not None and code not in untracked + } + if any(response.bounded for response in responses) or ( + sum(map(len, tracked_readsets)) > CUSTOM_HOST_MAX_TRACKED_READS + ): + untracked.update(code for code, _subject, result in evaluations if result.reads is not None) + untracked_reads: dict[str, set[FactKey]] = {code: set() for code in untracked} + for response in responses: + for code, reads in response.untracked_reads.items(): + untracked_reads[code].update(reads) + evaluated: dict[tuple[str, str], CustomHostEvaluation] = {} + for code, subject, result in evaluations: + merged: CustomHostEvaluation = result + if code in untracked and result.reads is not None: + untracked_reads[code].update(result.reads) + merged = CustomHostEvaluation(findings=result.findings, reads=None) + evaluated[(code, subject)] = merged + return CustomHostRun( + evaluated=evaluated, + untracked_reads={code: tuple(sorted(reads)) for code, reads in untracked_reads.items()}, + uncacheable=frozenset().union(*(response.uncacheable for response in responses)), + observed=_merged_observations(responses), + ) + + +def _merged_observations(responses: list[CustomHostSlice]) -> dict[FactKey, str | None]: + """Keep each observed digest, forgetting any that two hosts observed differently.""" + + observed: dict[FactKey, str | None] = {} + for response in responses: + for key, digest in response.observed.items(): + if key in observed and observed[key] != digest: + observed[key] = None + else: + observed.setdefault(key, digest) + return observed + + +def decode_readset(value: Any) -> tuple[FactKey, ...]: + return tuple(tuple(map(str, key)) for key in value) + + +def finding_payload(value: Any) -> dict[str, object]: + return {str(key): item for key, item in value.items()} diff --git a/src/sqlbuild/rule_engine/_helpers/run/rules.py b/src/sqlbuild/rule_engine/_helpers/run/rules.py index a9aa96dd97..a4053439f7 100644 --- a/src/sqlbuild/rule_engine/_helpers/run/rules.py +++ b/src/sqlbuild/rule_engine/_helpers/run/rules.py @@ -51,10 +51,12 @@ from sqlbuild.rule_engine._helpers.engine.config import resolve_rule_ignore_selectors from sqlbuild.rule_engine._helpers.engine.hermeticity import verify_custom_rules from sqlbuild.rule_engine._helpers.engine.native import ( + custom_rule_payloads, decode_rule_finding, evaluate_native, finalize_native_findings, native_catalogue, + start_custom_rules, ) from sqlbuild.rule_engine._helpers.run.findings import group_unevaluated_findings from sqlbuild.rule_engine._helpers.run.literal_duplicates import with_duplicate_literal_hints @@ -62,6 +64,7 @@ from sqlbuild.rule_engine.exceptions import RulesError from sqlbuild.rule_engine.main.load_config import load_rules_config from sqlbuild.rule_engine.models import ( + CustomRulesOutcome, Finding, PreparedSqlLint, PreparedSqlLintResult, @@ -114,8 +117,14 @@ def evaluate_rules( project_dir=resolved_project_dir, include_custom=include_custom, ) + custom_payloads: list[dict[str, object]] = custom_rule_payloads( + catalogue=catalogue, project_dir=resolved_project_dir + ) selected: tuple[Rule, ...] = select_rules( - catalogue=catalogue, config=effective_config, project_dir=resolved_project_dir + catalogue=catalogue, + config=effective_config, + project_dir=resolved_project_dir, + custom_payloads=custom_payloads, ) skipped_type_proof_rules: tuple[str, ...] = () if no_sql_analysis or not graph.project.settings.sql_analysis: @@ -140,20 +149,19 @@ def evaluate_rules( select=tuple(rule.code for rule in selected), ignore=(), ) - lint_config: LintConfig = sql_lint_config( - dialect=dialect, - codes=_sql_rule_codes(native_rules), - thresholds=effective_config.thresholds, - relation_keys=compiled_relation_keys(graph.project), - ) - sql_rule_inputs: _SqlRuleInputs = _prepare_sql_rule_inputs( - rules=native_rules, - project=selected_project, - dialect=_lint_identity(lint_config), - project_dir=resolved_project_dir, - collect_files=prepared_sql is None and model_paths is None, - ) - with ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-rules") as executor: + with ( + ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-custom-rules") as custom, + ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-rules") as executor, + ): + custom_outcome: Future[CustomRulesOutcome] | None = start_custom_rules( + executor=custom, + project=selected_project, + config=selected_config, + project_dir=resolved_project_dir, + catalogue=catalogue, + custom_payloads=custom_payloads, + dialect=dialect, + ) native_result: Future[RulesResult] = executor.submit( evaluate_native, project=selected_project, @@ -162,6 +170,21 @@ def evaluate_rules( catalogue=catalogue, dialect=dialect, defer_suppressions=True, + custom_payloads=custom_payloads, + custom_outcome=custom_outcome, + ) + lint_config: LintConfig = sql_lint_config( + dialect=dialect, + codes=_sql_rule_codes(native_rules), + thresholds=effective_config.thresholds, + relation_keys=compiled_relation_keys(graph.project), + ) + sql_rule_inputs: _SqlRuleInputs = _prepare_sql_rule_inputs( + rules=native_rules, + project=selected_project, + dialect=_lint_identity(lint_config), + project_dir=resolved_project_dir, + collect_files=prepared_sql is None and model_paths is None, ) sql_started: float = time.monotonic() sql_result: _SqlRulesEvaluation = _run_sql_rules( @@ -217,7 +240,7 @@ def _prepare_sql_rule_inputs( project_dir: Path, collect_files: bool, ) -> _SqlRuleInputs: - """Hash SQL and read files first so GIL-releasing work never waits on the native thread.""" + """Hash SQL and read the project files on the caller while built-in rules evaluate natively.""" identities: dict[str, str] = _sql_model_rule_identities( rules=rules, project=project, dialect=dialect diff --git a/src/sqlbuild/rule_engine/constants.py b/src/sqlbuild/rule_engine/constants.py index b9e266f999..4f9230b090 100644 --- a/src/sqlbuild/rule_engine/constants.py +++ b/src/sqlbuild/rule_engine/constants.py @@ -10,6 +10,7 @@ CUSTOM_HOST_PROTOCOL_VERSION: int = 1 CUSTOM_HOST_RUNTIME_VERSION: str = "sqlbuild-rules-custom-v2" CUSTOM_HOST_MAX_TRACKED_READS: int = 100_000 +CUSTOM_RULE_PROJECT_SUBJECT: str = "" CUSTOM_RULES_CACHE_VERSION: str = "custom-rules-reads-v1" CUSTOM_RULES_CACHE_FILE: str = "target/rules-cache/bulk/custom-rules.json" NATIVE_RULES_MEMO_VERSION: str = "native-rules-response-v1" diff --git a/src/sqlbuild/rule_engine/models.py b/src/sqlbuild/rule_engine/models.py index cb745b884b..085353ec41 100644 --- a/src/sqlbuild/rule_engine/models.py +++ b/src/sqlbuild/rule_engine/models.py @@ -10,7 +10,7 @@ from sqlbuild.compiler.compile.models import DeclarationScopeBuild, SqlExpansionContext from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs from sqlbuild.lint.models import LintRunResult -from sqlbuild.rule_engine.types import RuleCheck, RuleOptionValue, RuleSubject +from sqlbuild.rule_engine.types import FactKey, RuleCheck, RuleOptionValue, RuleSubject @dataclass(frozen=True, slots=True) @@ -411,3 +411,32 @@ class CustomRulesOutcome: cache_hits: int cache_misses: int custom_ms: int + + +@dataclass(frozen=True) +class CustomHostEvaluation: + """Decoded findings and recorded reads of one custom-rule invocation.""" + + findings: tuple[dict[str, object], ...] + reads: tuple[FactKey, ...] | None + + +@dataclass(frozen=True) +class CustomHostRun: + """Every planned custom-rule invocation, merged across host processes.""" + + evaluated: dict[tuple[str, str], CustomHostEvaluation] + untracked_reads: dict[str, tuple[FactKey, ...]] + uncacheable: frozenset[str] + observed: dict[FactKey, str | None] + + +@dataclass(frozen=True) +class CustomHostSlice: + """One host process's invocations, in the order that host evaluated them.""" + + evaluations: tuple[tuple[str, str, CustomHostEvaluation], ...] + untracked_reads: dict[str, tuple[FactKey, ...]] + uncacheable: frozenset[str] + observed: dict[FactKey, str | None] + bounded: bool diff --git a/src/sqlbuild/rule_engine/types.py b/src/sqlbuild/rule_engine/types.py index ccaedcbba2..477c7e123c 100644 --- a/src/sqlbuild/rule_engine/types.py +++ b/src/sqlbuild/rule_engine/types.py @@ -28,6 +28,8 @@ type FactKey = tuple[str, ...] type CustomRulePlan = dict[str, frozenset[str] | None] +type CustomHostPlan = dict[str, list[str] | None] + class RuleSubject(StrEnum): """Compiler-owned unit used to invoke and cache one custom rule.""" diff --git a/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py b/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py index 71b10cd215..af6a8c559d 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py @@ -156,6 +156,16 @@ class RulesIntegrationTestCase: expected_code: str +@dataclass(frozen=True) +class CustomHostSplitIntegrationTestCase: + """A compile whose custom rules run on several hosts and must match one host.""" + + description: str + model_count: int + hosts: int + expected_code: str + + @dataclass(frozen=True) class DollarQuotedLiteralBuildTestCase: """One model projection with dollar-quoted literals and its built rows.""" diff --git a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py index 0ec88e2af8..0bef0a630d 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py @@ -11,10 +11,11 @@ import pytest from _pytest.capture import CaptureResult +import sqlbuild._native as native_module import sqlbuild.compiler.compile._helpers.attachment.declaration_scope as declaration_scope_module import sqlbuild.compiler.compile.main._build_compile_inputs as compile_inputs_module import sqlbuild.compiler.compile.main.sql_expansion_context as expansion_context_module -import sqlbuild.rule_engine._helpers.engine.custom_rules as custom_rules_module +import sqlbuild.rule_engine._helpers.host.custom_host_pool as custom_host_pool_module import sqlbuild.rule_engine._helpers.run.rules as rules_module from sqlbuild.cli.commands.main.entrypoint.entry import main from sqlbuild.compiler.compile.models import DeclarationScopeBuild, LoadedMacro @@ -23,6 +24,7 @@ from sqlbuild.lint.classes.stop_context import LintStopContext from sqlbuild.lint.models import LintRunResult from tests.integration.src.sqlbuild.cli.commands.main._test_types import ( + CustomHostSplitIntegrationTestCase, DynamicPivotRulesIntegrationTestCase, ExplicitContractOutputRuleIntegrationTestCase, ImplicitAliasRuleIntegrationTestCase, @@ -1018,13 +1020,13 @@ def final_directory(*, model: Model, ctx: RuleContext) -> list[Finding]: encoding="utf-8", ) host_runs: list[str] = [] - original_host: Callable[[str], str] = custom_rules_module._native.run_custom_host_json + original_host: Callable[[str], str] = native_module.run_custom_host_json def recording_host(spec_json: str) -> str: host_runs.append(spec_json) return original_host(spec_json) - monkeypatch.setattr(custom_rules_module._native, "run_custom_host_json", recording_host) + monkeypatch.setattr(native_module, "run_custom_host_json", recording_host) cold: set[str] = compile_finding_keys(project_dir=tmp_path, capsys=capsys) @@ -1038,6 +1040,78 @@ def recording_host(spec_json: str) -> str: } <= cold +@pytest.mark.parametrize( + "test_case", + [CustomHostSplitIntegrationTestCase("custom rules split across hosts", 9, 3, "XSQBRARCH001")], + ids=lambda case: case.description, +) +def test_given_custom_rules_split_across_hosts_when_compiling_then_output_matches_one_host( + test_case: CustomHostSplitIntegrationTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + (tmp_path / "sqlbuild_project.toml").write_text( + 'name = "orders"\nadapter = "duckdb"\n\n' + f'[rules]\nselect = ["{test_case.expected_code}"]\n\n' + "[rules.thresholds]\nmin_custom_rule_test_cases = 0\n", + encoding="utf-8", + ) + for index in range(test_case.model_count): + folder: str = ("final", "staging", "staging")[index % 3] + model: Path = tmp_path / "models" / folder / f"orders_{index:02d}.sql" + model.parent.mkdir(parents=True, exist_ok=True) + model.write_text( + f'MODEL (description "Orders {index}");\nSELECT {index} AS order_id\n', + encoding="utf-8", + ) + rule_file: Path = tmp_path / "rules" / "architecture.py" + rule_file.parent.mkdir() + rule_file.write_text( + """from sqlbuild.rules import Finding, Model, RuleContext, rule + +@rule( + code="XSQBRARCH001", + message="Final models must use the final directory", + remediation="Move this model beneath models/final/.", +) +def final_directory(*, model: Model, ctx: RuleContext) -> list[Finding]: + return [] if "final" in model.path.parts else [ctx.finding(subject=model)] +""", + encoding="utf-8", + ) + host_runs: list[str] = [] + original_host: Callable[[str], str] = native_module.run_custom_host_json + + def recording_host(spec_json: str) -> str: + host_runs.append(spec_json) + return original_host(spec_json) + + def compile_payload(*, hosts: int) -> dict[str, Any]: + monkeypatch.setattr(custom_host_pool_module, "_MIN_INVOCATIONS_PER_HOST", 1) + monkeypatch.setattr(custom_host_pool_module, "_available_cores", lambda: hosts) + exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) + assert exit_code == 1 + return json.loads(capsys.readouterr().out) + + monkeypatch.setattr(native_module, "run_custom_host_json", recording_host) + single: dict[str, Any] = compile_payload(hosts=1) + single_runs: int = len(host_runs) + (tmp_path / "target" / "rules-cache").rename(tmp_path / "single-rules-cache") + split: dict[str, Any] = compile_payload(hosts=test_case.hosts) + split_runs: int = len(host_runs) - single_runs + warm: dict[str, Any] = compile_payload(hosts=test_case.hosts) + + assert (single_runs, split_runs, len(host_runs)) == (1, test_case.hosts, 1 + test_case.hosts) + assert split["diagnostics"] == single["diagnostics"] + assert len(split["diagnostics"]) == test_case.model_count - test_case.model_count // 3 + assert warm["diagnostics"] == single["diagnostics"] + assert warm["compile_timings"]["rule_cache_hits"] >= test_case.model_count + assert (tmp_path / "single-rules-cache" / "bulk" / "custom-rules.json").read_bytes() == ( + tmp_path / "target" / "rules-cache" / "bulk" / "custom-rules.json" + ).read_bytes() + + @pytest.mark.parametrize( "test_case", [RulesIntegrationTestCase("custom rule is enforced and cached", 1, "XSQBRARCH001")], diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_reads.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_reads.py index a0fe767bcf..617e9cf8b9 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_reads.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_reads.py @@ -245,14 +245,14 @@ def test_given_file_edited_while_rules_run_when_rerunning_then_new_contents_are_ module_import=test_case.module_prelude, ) project: CompiledProject = two_model_project(customers_sql=_CUSTOMERS_SQL, customers_config={}) - run_host: Callable[..., object] = custom_rules._run_host + run_host: Callable[..., object] = custom_rules.run_custom_hosts def run_host_then_edit(**kwargs: Any) -> object: result: object = run_host(**kwargs) owners.write_text(_OWNERS_AFTER, encoding="utf-8") return result - monkeypatch.setattr(custom_rules, "_run_host", run_host_then_edit) + monkeypatch.setattr(custom_rules, "run_custom_hosts", run_host_then_edit) during_edit: CustomRulesOutcome = evaluate_cached_custom_rules( project=project, project_dir=tmp_path, cache_enabled=True ) diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py new file mode 100644 index 0000000000..7596c1c17d --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py @@ -0,0 +1,35 @@ +"""Test case declarations for custom-rule host partitioning.""" + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class HostPartitionTestCase: + """One partitioned custom-rule evaluation compared with a single host.""" + + description: str + model_count: int + hosts: int + expected_host_runs: int + records_module_state: bool = False + + +@dataclass(frozen=True) +class HostFailureTestCase: + """Failing subjects whose first single-host failure must be reported.""" + + description: str + model_count: int + failing_models: tuple[str, ...] + expected_failing_model: str + + +@dataclass(frozen=True) +class HostPlanTestCase: + """One deterministic split of planned invocations.""" + + description: str + plan: dict[str, list[str] | None] + invocations: tuple[tuple[str, str], ...] + hosts: int + expected_plans: tuple[dict[str, list[str] | None], ...] diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py new file mode 100644 index 0000000000..8f7c52cf6d --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py @@ -0,0 +1,100 @@ +"""Builders for custom-rule host partitioning tests.""" + +from collections.abc import Callable +from dataclasses import replace +from operator import attrgetter +from pathlib import Path + +import pytest + +import sqlbuild._native as native_module +from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.rule_engine._helpers.engine.catalogue import build_catalogue +from sqlbuild.rule_engine._helpers.engine.custom_rules import evaluate_custom_rules_cached +from sqlbuild.rule_engine._helpers.host import custom_host_pool +from sqlbuild.rule_engine.models import CustomRulesOutcome, Rule, RulesCacheConfig, RulesConfig +from tests.unit.src.sqlbuild.rule_engine.main.evaluate.helpers import build_project + +_MODEL_RULE: str = "XSQBRT101" +_PROJECT_RULE: str = "XSQBRT102" +_RULES_MODULE: str = f'''from sqlbuild.rules import Finding, Model, Project, RuleContext, rule + + +@rule(code="{_MODEL_RULE}", message="order models need an even suffix", remediation="Rename it.") +def even_suffix(*, model: Model, ctx: RuleContext) -> list[Finding]: + if model.name in FAILING: + raise ValueError(f"cannot inspect {{model.name}}") + if RECORD_SEEN: + SEEN[model.name] = True + source: str = ctx.sql.for_model(model).expanded.source + return [ctx.finding(subject=model)] if "odd" in source else [] + + +@rule(code="{_PROJECT_RULE}", message="projects need few models", remediation="Split it.") +def few_models(*, project: Project, ctx: RuleContext) -> list[Finding]: + del project + first: Model = ctx.project.models[0] + return [ctx.finding(subject=first)] if len(ctx.project.models) > 3 else [] +''' +_PARITIES: tuple[str, str] = ("even", "odd") + + +def write_order_rules( + *, root: Path, failing: tuple[str, ...] = (), records_module_state: bool = False +) -> None: + """Write one model rule and one project rule, optionally failing or keeping module state.""" + + rules: Path = root / "rules" / "orders.py" + rules.parent.mkdir(parents=True, exist_ok=True) + rules.write_text( + f"FAILING = {set(failing)!r}\nRECORD_SEEN = {records_module_state!r}\n" + f"SEEN: dict[str, bool] = {{}}\n{_RULES_MODULE}", + encoding="utf-8", + ) + + +def orders_project(*, model_count: int, edits: dict[str, str] | None = None) -> CompiledProject: + """Return numbered order models whose SQL marks them odd or even, with optional edits.""" + + sql_by_name: dict[str, str] = { + f"orders_{index:03d}": f"SELECT 1 AS order_id -- {_PARITIES[index % 2]}" + for index in range(model_count) + } + sql_by_name.update(edits or {}) + projects: list[CompiledProject] = [ + build_project(name=name, relative_path=f"models/{name}.sql", sql=sql, config_values={}) + for name, sql in sql_by_name.items() + ] + return replace(projects[0], models=tuple(project.models[0] for project in projects)) + + +def evaluate_order_rules( + *, project: CompiledProject, root: Path, cache_enabled: bool = True +) -> CustomRulesOutcome: + """Evaluate both order rules through the incremental custom-rule cache.""" + + config: RulesConfig = RulesConfig( + select=(_MODEL_RULE, _PROJECT_RULE), cache=RulesCacheConfig(enabled=cache_enabled) + ) + rules: tuple[Rule, ...] = tuple( + filter(attrgetter("custom"), build_catalogue(config=config, project_dir=root)) + ) + return evaluate_custom_rules_cached( + project=project, config=config, project_dir=root, rules=rules, dialect="duckdb" + ) + + +def use_hosts(*, monkeypatch: pytest.MonkeyPatch, hosts: int) -> list[str]: + """Allow up to `hosts` host processes for any plan and record every host launch.""" + + launches: list[str] = [] + launch: Callable[[str], str] = native_module.run_custom_host_json + + def recording_launch(spec_json: str) -> str: + launches.append(spec_json) + return launch(spec_json) + + monkeypatch.setattr(custom_host_pool, "_MIN_INVOCATIONS_PER_HOST", 1) + monkeypatch.setattr(custom_host_pool, "_available_cores", lambda: hosts) + monkeypatch.setattr(native_module, "run_custom_host_json", recording_launch) + return launches diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py new file mode 100644 index 0000000000..86c402b735 --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py @@ -0,0 +1,179 @@ +"""Custom rules split across host processes report exactly what one host reports.""" + +from pathlib import Path + +import pytest + +from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.rule_engine._helpers.host import custom_host_pool +from sqlbuild.rule_engine.constants import CUSTOM_RULES_CACHE_FILE +from sqlbuild.rule_engine.exceptions import RulesError +from sqlbuild.rule_engine.models import CustomRulesOutcome +from tests.unit.src.sqlbuild.rule_engine._helpers.host._test_types import ( + HostFailureTestCase, + HostPartitionTestCase, + HostPlanTestCase, +) +from tests.unit.src.sqlbuild.rule_engine._helpers.host.helpers import ( + evaluate_order_rules, + orders_project, + use_hosts, + write_order_rules, +) + + +@pytest.mark.parametrize( + "test_case", + [ + HostPartitionTestCase( + description="three hosts", model_count=7, hosts=3, expected_host_runs=3 + ), + HostPartitionTestCase( + description="more hosts than subjects", model_count=2, hosts=8, expected_host_runs=3 + ), + HostPartitionTestCase( + description="module state makes the rule untracked on every host", + model_count=7, + hosts=3, + expected_host_runs=3, + records_module_state=True, + ), + ], + ids=lambda case: case.description, +) +def test_given_partitioned_hosts_when_evaluating_then_findings_and_cache_match_one_host( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostPartitionTestCase +) -> None: + write_order_rules(root=tmp_path) + project: CompiledProject = orders_project(model_count=test_case.model_count) + cache_path: Path = tmp_path / CUSTOM_RULES_CACHE_FILE + + single_launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=1) + single: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) + single_cache: bytes = cache_path.read_bytes() + cache_path.unlink() + monkeypatch.undo() + split_launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + split: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) + + assert len(single_launches) == 1 + assert len(split_launches) == test_case.expected_host_runs + assert split.findings == single.findings + assert (split.cache_hits, split.cache_misses) == (single.cache_hits, single.cache_misses) + assert cache_path.read_bytes() == single_cache + assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").glob("*.pickle")) + + +@pytest.mark.parametrize( + "test_case", + [ + HostPartitionTestCase( + description="edit after split", model_count=6, hosts=3, expected_host_runs=3 + ) + ], + ids=lambda case: case.description, +) +def test_given_results_cached_by_split_hosts_when_rerunning_then_only_edited_subjects_rerun( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostPartitionTestCase +) -> None: + write_order_rules(root=tmp_path) + project: CompiledProject = orders_project(model_count=test_case.model_count) + edited: CompiledProject = orders_project( + model_count=test_case.model_count, edits={"orders_003": "SELECT 2 AS order_id"} + ) + launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + + cold: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) + cold_launches: int = len(launches) + warm: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) + warm_launches: int = len(launches) - cold_launches + after_edit: CustomRulesOutcome = evaluate_order_rules(project=edited, root=tmp_path) + oracle: CustomRulesOutcome = evaluate_order_rules( + project=edited, root=tmp_path, cache_enabled=False + ) + + assert cold_launches == test_case.expected_host_runs + assert warm_launches == 0 + assert (warm.cache_hits, warm.cache_misses) == (cold.cache_misses, 0) + assert warm.findings == cold.findings + assert (after_edit.cache_hits, after_edit.cache_misses) == (cold.cache_misses - 1, 1) + assert after_edit.findings == oracle.findings + + +@pytest.mark.parametrize( + "test_case", + [ + HostFailureTestCase( + description="earliest failure lives on a later host", + model_count=6, + failing_models=("orders_001", "orders_003"), + expected_failing_model="orders_001", + ), + HostFailureTestCase( + description="one failing subject", + model_count=6, + failing_models=("orders_005",), + expected_failing_model="orders_005", + ), + ], + ids=lambda case: case.description, +) +def test_given_failing_rule_on_split_hosts_when_evaluating_then_single_host_error_is_raised( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostFailureTestCase +) -> None: + write_order_rules(root=tmp_path, failing=test_case.failing_models) + project: CompiledProject = orders_project(model_count=test_case.model_count) + + _ = use_hosts(monkeypatch=monkeypatch, hosts=1) + with pytest.raises(RulesError) as single: + _ = evaluate_order_rules(project=project, root=tmp_path) + monkeypatch.undo() + _ = use_hosts(monkeypatch=monkeypatch, hosts=3) + with pytest.raises(RulesError) as split: + _ = evaluate_order_rules(project=project, root=tmp_path) + + assert str(split.value) == str(single.value) + assert f"cannot inspect {test_case.expected_failing_model}" in str(split.value) + + +@pytest.mark.parametrize( + "test_case", + [ + HostPlanTestCase( + description="project rules and whole subjects are dealt round-robin", + plan={"XSQBRT101": None, "XSQBRT102": None, "XSQBRT103": ["models/b.sql"]}, + invocations=( + ("XSQBRT102", ""), + ("XSQBRT101", "models/a.sql"), + ("XSQBRT101", "models/b.sql"), + ("XSQBRT103", "models/b.sql"), + ("XSQBRT101", "models/c.sql"), + ), + hosts=2, + expected_plans=( + {"XSQBRT101": ["models/b.sql"], "XSQBRT102": [""], "XSQBRT103": ["models/b.sql"]}, + {"XSQBRT101": ["models/a.sql", "models/c.sql"], "XSQBRT102": [], "XSQBRT103": []}, + ), + ), + HostPlanTestCase( + description="one subject keeps the original plan", + plan={"XSQBRT101": None}, + invocations=(("XSQBRT101", "models/a.sql"),), + hosts=4, + expected_plans=({"XSQBRT101": None},), + ), + ], + ids=lambda case: case.description, +) +def test_given_planned_invocations_when_partitioning_then_split_is_deterministic( + test_case: HostPlanTestCase, +) -> None: + plans: tuple[dict[str, list[str] | None], ...] = custom_host_pool.partition_host_plan( + plan=test_case.plan, invocations=test_case.invocations, hosts=test_case.hosts + ) + + assert plans == test_case.expected_plans + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) From 9974af64f6be58bde251def5e9c355a70b1a76e7 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sat, 3 Oct 2026 22:46:11 +0100 Subject: [PATCH 2/3] perf(rules): prepare speculative lint bodies in one native batch The early SQL lint expanded and prepared each body with its own native call, handing the GIL back and forth with graph building thousands of times. Bodies are now expanded first and prepared in one GIL-free native call; a body the batch cannot prepare is prepared again on its own so failures surface exactly as before, in file order relative to expansion failures. --- .../src/bindings/_helpers/functions.rs | 26 ++ src/sqlbuild/_native.pyi | 3 + src/sqlbuild/lint/_helpers/expansion.py | 229 ++++++++++++++++-- src/sqlbuild/lint/main/run_lint.py | 183 +++++++------- src/sqlbuild/lint/models.py | 43 +++- src/sqlbuild/lint/types.py | 9 +- .../src/sqlbuild/lint/_helpers/_test_types.py | 19 ++ .../src/sqlbuild/lint/_helpers/helpers.py | 27 +++ .../sqlbuild/lint/_helpers/test_expansion.py | 116 +++++++++ 9 files changed, 525 insertions(+), 130 deletions(-) create mode 100644 tests/unit/src/sqlbuild/lint/_helpers/test_expansion.py diff --git a/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs b/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs index fdc5c9b548..59b33f5732 100644 --- a/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs +++ b/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs @@ -117,6 +117,31 @@ fn prepare_lint_sql( .map_err(value_error) } +/// Prepare many bodies in one detached call; `false` marks a body to prepare again individually. +#[pyfunction] +fn prepare_lint_sql_batch( + py: Python<'_>, + requests: Vec, +) -> PyResult)>> { + py.compiler_detach(|| { + Ok(requests + .iter() + .map(|request| { + crate::bindings::_helpers::panics::catch_compiler_panic(|| { + crate::sql_lint::main::preparation::prepare( + &request.expanded, + &request.before_expansion, + &request.prior_sites, + &request.dialect, + ) + }) + .map_or((false, None), |prepared| (true, prepared)) + }) + .collect()) + }) + .map_err(value_error) +} + #[pyfunction] fn lint_backtick_identifiers(dialect: &str) -> PyResult { compiler_guard(|| { @@ -431,6 +456,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(run_custom_host_json, module)?)?; module.add_function(wrap_pyfunction!(lint_sql_json, module)?)?; module.add_function(wrap_pyfunction!(prepare_lint_sql, module)?)?; + module.add_function(wrap_pyfunction!(prepare_lint_sql_batch, module)?)?; module.add( "NativeCompilerError", module diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 5c54c88e00..d90c492dac 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -23,6 +23,9 @@ def lint_sql_batch_json(request_json: str) -> str: ... def prepare_lint_sql( request: LintPreparationRequest, ) -> tuple[str, list[tuple[str, int, int, int, int, str]], list[str]] | None: ... +def prepare_lint_sql_batch( + requests: list[LintPreparationRequest], +) -> list[tuple[bool, tuple[str, list[tuple[str, int, int, int, int, str]], list[str]] | None]]: ... def lint_backtick_identifiers(dialect: str) -> bool: ... def load_config_json(project_dir: str) -> str: ... def selected_codes_json(request_json: str) -> str: ... diff --git a/src/sqlbuild/lint/_helpers/expansion.py b/src/sqlbuild/lint/_helpers/expansion.py index 5b641ee845..29ae933e18 100644 --- a/src/sqlbuild/lint/_helpers/expansion.py +++ b/src/sqlbuild/lint/_helpers/expansion.py @@ -3,6 +3,7 @@ from __future__ import annotations import re +from dataclasses import replace from pathlib import Path import sqlbuild._native as _native @@ -23,15 +24,32 @@ from sqlbuild.compiler.discovery.main.discover import discover_project_inputs from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs from sqlbuild.compiler.scopes.types import DeclarationKind +from sqlbuild.lint._helpers.headers import lint_body_ranges, lint_file_role +from sqlbuild.lint._helpers.native import external_identifiers_for_headers +from sqlbuild.lint._helpers.native_sql import check_lint_not_stopped from sqlbuild.lint._helpers.sqlbuild_tokens import ( neutralize_context_interpolation, neutralize_generic_audit_parameters, neutralize_interpolation, sentinel_spans, ) -from sqlbuild.lint.constants import RELATION_IDENTITY_TEMPLATE, TEMPLATE_INTERPOLATION_START +from sqlbuild.lint.constants import ( + HEADER_KIND_SCENARIO, + HEADER_KIND_TEST, + RELATION_IDENTITY_TEMPLATE, + TEMPLATE_INTERPOLATION_START, +) from sqlbuild.lint.exceptions import ProjectCompileError -from sqlbuild.lint.models import InterpolationSite, LintBody, LintFileRole +from sqlbuild.lint.models import ( + ExpandedLintFile, + HeaderSpan, + InterpolationSite, + LintBody, + LintFileRole, + PendingLintBody, + PreparedLintFiles, +) +from sqlbuild.lint.types import NativePreparedSql from sqlbuild.spec.contracts.main.resolve_effective_adapter_name import ( resolve_effective_adapter_name, ) @@ -146,6 +164,38 @@ def prepare_lint_body( ) -> LintBody: """Expand one authored body and neutralize whatever interpolation remains.""" + pending: PendingLintBody = expand_lint_body( + role=role, + file_path=file_path, + contents=contents, + body_range=body_range, + context=context, + dialect=dialect, + external_identifiers=external_identifiers, + allows_ceremonial_select=allows_ceremonial_select, + allows_dynamic_output_star=allows_dynamic_output_star, + compiled_expansion=compiled_expansion, + ) + return finish_lint_body( + pending=pending, prepared=_native.prepare_lint_sql(pending.request), dialect=dialect + ) + + +def expand_lint_body( + *, + role: LintFileRole, + file_path: Path, + contents: str, + body_range: tuple[int, int], + context: SqlExpansionContext, + dialect: str, + external_identifiers: tuple[str, ...] = (), + allows_ceremonial_select: bool = False, + allows_dynamic_output_star: bool = False, + compiled_expansion: CompiledSqlExpansion | None = None, +) -> PendingLintBody: + """Expand one authored body like compile does, leaving its native preparation pending.""" + body_start: int body_end: int body_start, body_end = body_range @@ -184,50 +234,84 @@ def prepare_lint_body( raise ProjectCompileError( f"{file_path} could not be expanded, so its SQL cannot be linted: {error}" ) from error + return PendingLintBody( + file_path=file_path, + body_start=body_start, + body_end=body_end, + expansion_input=expansion_input, + expanded=expanded, + expansion_passes=expansion_passes, + pre_expansion_sites=pre_expansion_sites, + external_identifiers=external_identifiers, + allows_ceremonial_select=allows_ceremonial_select, + allows_dynamic_output_star=allows_dynamic_output_star, + allows_empty_fixture_star=False, + request={ + "expanded": expanded, + "before_expansion": expansion_input, + "prior_sites": [site.neutralized_start for site in pre_expansion_sites], + "dialect": dialect, + }, + ) + + +def prepare_lint_sql_batch( + pending: tuple[PendingLintBody, ...], +) -> tuple[NativePreparedSql | None, ...]: + """Prepare pending bodies in one GIL-free native call; failures re-raise per body, in order.""" + + if not pending: + return () + results: list[tuple[bool, NativePreparedSql | None]] = _native.prepare_lint_sql_batch( + [body.request for body in pending] + ) + return tuple( + prepared if succeeded else _native.prepare_lint_sql(body.request) + for body, (succeeded, prepared) in zip(pending, results, strict=True) + ) + + +def finish_lint_body( + *, pending: PendingLintBody, prepared: NativePreparedSql | None, dialect: str +) -> LintBody: + """Build the lint body from one body's native preparation, or the Python fallback.""" + + expanded: str = pending.expanded neutralized: str sites: tuple[InterpolationSite, ...] - native_prepared: tuple[str, list[tuple[str, int, int, int, int, str]], list[str]] | None = ( - _native.prepare_lint_sql( - { - "expanded": expanded, - "before_expansion": expansion_input, - "prior_sites": [site.neutralized_start for site in pre_expansion_sites], - "dialect": dialect, - } - ) - ) - if native_prepared is None: + if prepared is None: neutralized, sites = neutralize_interpolation(body=expanded, dialect=dialect) externally_referenced_ctes: tuple[str, ...] = _externally_referenced_ctes( expanded=expanded, interpolation_sites=sites, - pre_expansion_body=expansion_input, - pre_expansion_sites=pre_expansion_sites, + pre_expansion_body=pending.expansion_input, + pre_expansion_sites=pending.pre_expansion_sites, ) else: - neutralized = native_prepared[0] - sites = tuple(InterpolationSite(*site) for site in native_prepared[1]) - externally_referenced_ctes = tuple(native_prepared[2]) + neutralized = prepared[0] + sites = tuple(InterpolationSite(*site) for site in prepared[1]) + externally_referenced_ctes = tuple(prepared[2]) dependency_identifiers: tuple[str, ...] = tuple( site.sentinel for site in sites if _DEPENDENCY_INTRINSIC_PATTERN.match(site.original_text) is not None ) return LintBody( - file_path=file_path, - body_start=body_start, - body_end=body_end, + file_path=pending.file_path, + body_start=pending.body_start, + body_end=pending.body_end, lint_text=neutralized, passes=( - sentinel_spans(sites=pre_expansion_sites), - *expansion_passes, + sentinel_spans(sites=pending.pre_expansion_sites), + *pending.expansion_passes, sentinel_spans(sites=sites), ), - external_identifiers=external_identifiers, + external_identifiers=pending.external_identifiers, dependency_identifiers=dependency_identifiers, externally_referenced_ctes=externally_referenced_ctes, - allows_ceremonial_select=allows_ceremonial_select, - allows_dynamic_output_star=allows_dynamic_output_star, + allows_ceremonial_select=pending.allows_ceremonial_select, + allows_dynamic_output_star=pending.allows_dynamic_output_star, + allows_empty_fixture_star=pending.allows_empty_fixture_star, dependency_relations=tuple( ( site.sentinel, @@ -249,3 +333,96 @@ def _is_sql_hook_path(*, role: LintFileRole) -> bool: return role.in_hook_directory or ( role.in_project and role.declaration_kind is DeclarationKind.SQL_HOOK ) + + +def expand_file_bodies( + *, + file_path: Path, + contents: str, + headers: tuple[HeaderSpan, ...], + context: SqlExpansionContext, + dialect: str, + project_dir: Path, + relative_path: Path, + dynamic_output_paths: frozenset[Path], + compiled_expansions: dict[Path, CompiledSqlExpansion] | None, +) -> ExpandedLintFile: + """Expand one file's bodies up to the first body that cannot be expanded.""" + + check_lint_not_stopped() + bodies: list[PendingLintBody] = [] + role: LintFileRole = lint_file_role( + file_path=file_path, project_dir=project_dir, relative_path=relative_path + ) + allows_dynamic_output_star: bool = ( + bool(dynamic_output_paths) and file_path.resolve() in dynamic_output_paths + ) + compiled_expansion: CompiledSqlExpansion | None = (compiled_expansions or {}).get(file_path) + external_identifiers: tuple[str, ...] = external_identifiers_for_headers( + contents=contents, headers=headers + ) + allows_ceremonial_select: bool = any( + header.kind in {HEADER_KIND_TEST, HEADER_KIND_SCENARIO} for header in headers + ) + allows_empty_fixture_star: bool = any(header.kind == HEADER_KIND_TEST for header in headers) + body_start: int + body_end: int + try: + for body_start, body_end in lint_body_ranges( + contents=contents, + headers=headers, + file_path=file_path, + project_dir=project_dir, + role=role, + ): + bodies.append( + replace( + expand_lint_body( + role=role, + file_path=file_path, + contents=contents, + body_range=(body_start, body_end), + context=context, + dialect=dialect, + external_identifiers=external_identifiers, + allows_ceremonial_select=allows_ceremonial_select, + allows_dynamic_output_star=allows_dynamic_output_star, + compiled_expansion=compiled_expansion, + ), + allows_empty_fixture_star=allows_empty_fixture_star, + ) + ) + except ProjectCompileError as error: + return ExpandedLintFile(file_path=file_path, bodies=tuple(bodies), failure=error) + return ExpandedLintFile(file_path=file_path, bodies=tuple(bodies), failure=None) + + +def prepare_expanded_files( + *, expanded_files: tuple[ExpandedLintFile, ...], dialect: str, skip_unexpandable: bool +) -> PreparedLintFiles: + """Prepare every expanded body natively at once, then report failures in file order.""" + + check_lint_not_stopped() + pending: list[PendingLintBody] = [] + for expanded in expanded_files: + pending.extend(expanded.bodies) + prepared: tuple[NativePreparedSql | None, ...] = prepare_lint_sql_batch(tuple(pending)) + bodies: list[LintBody] = [] + unexpandable: dict[Path, str] = {} + position: int = 0 + for expanded in expanded_files: + file_bodies: list[LintBody] = [] + for body in expanded.bodies: + file_bodies.append( + finish_lint_body(pending=body, prepared=prepared[position], dialect=dialect) + ) + position += 1 + if expanded.failure is None: + bodies.extend(file_bodies) + elif skip_unexpandable: + unexpandable[expanded.file_path] = str(expanded.failure).removeprefix( + f"{expanded.file_path} " + ) + else: + raise expanded.failure + return PreparedLintFiles(bodies=tuple(bodies), unexpandable=unexpandable) diff --git a/src/sqlbuild/lint/main/run_lint.py b/src/sqlbuild/lint/main/run_lint.py index f746cd5975..9756337ee1 100644 --- a/src/sqlbuild/lint/main/run_lint.py +++ b/src/sqlbuild/lint/main/run_lint.py @@ -2,7 +2,6 @@ from __future__ import annotations -from dataclasses import replace from pathlib import Path from sqlbuild.compiler.compile.constants import MODEL_DIRECTORY_NAME @@ -14,21 +13,24 @@ INHERITED_DECLARATION_DIRECTORIES, LOCAL_DECLARATION_DIRECTORIES, ) -from sqlbuild.lint._helpers.expansion import build_lint_expansion_context, prepare_lint_body -from sqlbuild.lint._helpers.headers import lint_body_ranges, lint_file_role, scan_headers -from sqlbuild.lint._helpers.native import external_identifiers_for_headers, lint_native_headers -from sqlbuild.lint._helpers.native_sql import check_lint_not_stopped, run_native_sql_lint +from sqlbuild.lint._helpers.expansion import ( + build_lint_expansion_context, + expand_file_bodies, + prepare_expanded_files, +) +from sqlbuild.lint._helpers.headers import scan_headers +from sqlbuild.lint._helpers.native import lint_native_headers +from sqlbuild.lint._helpers.native_sql import run_native_sql_lint from sqlbuild.lint._helpers.project_files import collect_project_files, sort_violations from sqlbuild.lint._helpers.suppressions import apply_suppressions -from sqlbuild.lint.constants import HEADER_KIND_SCENARIO, HEADER_KIND_TEST -from sqlbuild.lint.exceptions import ProjectCompileError from sqlbuild.lint.models import ( + ExpandedLintFile, HeaderSpan, LintBody, LintConfig, - LintFileRole, LintRunResult, LintViolation, + PreparedLintFiles, ) @@ -51,62 +53,51 @@ def run_lint( project_dir=project_dir, selected_paths=selected_paths, source_files=source_files ) violations: list[LintViolation] = [] - bodies: list[LintBody] = [] - unexpandable: dict[Path, str] = {} context: SqlExpansionContext | None = expansion_context or _expansion_context( project_dir=project_dir, native_enabled=config.native_enabled, value_renderer=value_renderer, discovered_inputs=discovered_inputs, ) + expanded_files: list[ExpandedLintFile] = [] + failure: Exception | None = None file_path: Path contents: str for file_path, contents in sorted(files.items()): - relative_path: Path = file_path.relative_to(project_dir) - relative_parts: tuple[str, ...] = relative_path.parts - declaration_directories: frozenset[str] = ( - INHERITED_DECLARATION_DIRECTORIES - | LOCAL_DECLARATION_DIRECTORIES - | {DECLARATION_GROUP_DIRECTORY} - ) - headers: tuple[HeaderSpan, ...] = scan_headers( - contents=contents, - first_only=( - relative_parts[:1] == (MODEL_DIRECTORY_NAME,) - and not declaration_directories.intersection(relative_parts[1:-1]) - ), - ) - if config.header_rules_enabled: - violations.extend( - lint_native_headers( - contents=contents, - file_path=file_path, - headers=headers, - config=config, - ) + try: + header_violations: tuple[LintViolation, ...] + expanded: ExpandedLintFile | None + header_violations, expanded = _scan_file( + file_path=file_path, + contents=contents, + project_dir=project_dir, + config=config, + context=context, + dynamic_output_paths=dynamic_output_paths, + compiled_expansions=compiled_expansions, ) - if context is None: + except Exception as error: + failure = error + break + violations.extend(header_violations) + if expanded is None: continue - prepared: tuple[LintBody, ...] | str = _prepared_bodies( - file_path=file_path, - contents=contents, - headers=headers, - context=context, - dialect=config.dialect, - project_dir=project_dir, - relative_path=relative_path, - dynamic_output_paths=dynamic_output_paths, - compiled_expansions=compiled_expansions, - skip_unexpandable=skip_unexpandable, - ) - if isinstance(prepared, str): - unexpandable[file_path] = prepared - else: - bodies.extend(prepared) + expanded_files.append(expanded) + if expanded.failure is not None and not skip_unexpandable: + break + prepared: PreparedLintFiles = prepare_expanded_files( + expanded_files=tuple(expanded_files), + dialect=config.dialect, + skip_unexpandable=skip_unexpandable, + ) + if failure is not None: + raise failure + bodies: tuple[LintBody, ...] = prepared.bodies + unexpandable: dict[Path, str] = prepared.unexpandable if bodies and config.native_enabled: native_violations: dict[Path, tuple[LintViolation, ...]] = run_native_sql_lint( - bodies=tuple(bodies), + bodies=bodies, contents_by_path=files, config=config, ) @@ -139,66 +130,54 @@ def _expansion_context( ) -def _prepared_bodies( +def _scan_file( *, file_path: Path, contents: str, - headers: tuple[HeaderSpan, ...], - context: SqlExpansionContext, - dialect: str, project_dir: Path, - relative_path: Path, + config: LintConfig, + context: SqlExpansionContext | None, dynamic_output_paths: frozenset[Path], compiled_expansions: dict[Path, CompiledSqlExpansion] | None, - skip_unexpandable: bool, -) -> tuple[LintBody, ...] | str: - """Prepared bodies of one file, or the reason its SQL cannot be expanded when skipping.""" +) -> tuple[tuple[LintViolation, ...], ExpandedLintFile | None]: + """Header violations of one file and, when expanding, its bodies awaiting preparation.""" - check_lint_not_stopped() - bodies: list[LintBody] = [] - role: LintFileRole = lint_file_role( - file_path=file_path, project_dir=project_dir, relative_path=relative_path + relative_path: Path = file_path.relative_to(project_dir) + relative_parts: tuple[str, ...] = relative_path.parts + declaration_directories: frozenset[str] = ( + INHERITED_DECLARATION_DIRECTORIES + | LOCAL_DECLARATION_DIRECTORIES + | {DECLARATION_GROUP_DIRECTORY} ) - allows_dynamic_output_star: bool = ( - bool(dynamic_output_paths) and file_path.resolve() in dynamic_output_paths + headers: tuple[HeaderSpan, ...] = scan_headers( + contents=contents, + first_only=( + relative_parts[:1] == (MODEL_DIRECTORY_NAME,) + and not declaration_directories.intersection(relative_parts[1:-1]) + ), ) - compiled_expansion: CompiledSqlExpansion | None = (compiled_expansions or {}).get(file_path) - external_identifiers: tuple[str, ...] = external_identifiers_for_headers( - contents=contents, headers=headers + header_violations: tuple[LintViolation, ...] = ( + tuple( + lint_native_headers( + contents=contents, + file_path=file_path, + headers=headers, + config=config, + ) + ) + if config.header_rules_enabled + else () ) - allows_ceremonial_select: bool = any( - header.kind in {HEADER_KIND_TEST, HEADER_KIND_SCENARIO} for header in headers + if context is None: + return header_violations, None + return header_violations, expand_file_bodies( + file_path=file_path, + contents=contents, + headers=headers, + context=context, + dialect=config.dialect, + project_dir=project_dir, + relative_path=relative_path, + dynamic_output_paths=dynamic_output_paths, + compiled_expansions=compiled_expansions, ) - allows_empty_fixture_star: bool = any(header.kind == HEADER_KIND_TEST for header in headers) - body_start: int - body_end: int - try: - for body_start, body_end in lint_body_ranges( - contents=contents, - headers=headers, - file_path=file_path, - project_dir=project_dir, - role=role, - ): - bodies.append( - replace( - prepare_lint_body( - role=role, - file_path=file_path, - contents=contents, - body_range=(body_start, body_end), - context=context, - dialect=dialect, - external_identifiers=external_identifiers, - allows_ceremonial_select=allows_ceremonial_select, - allows_dynamic_output_star=allows_dynamic_output_star, - compiled_expansion=compiled_expansion, - ), - allows_empty_fixture_star=allows_empty_fixture_star, - ) - ) - except ProjectCompileError as error: - if not skip_unexpandable: - raise - return str(error).removeprefix(f"{file_path} ") - return tuple(bodies) diff --git a/src/sqlbuild/lint/models.py b/src/sqlbuild/lint/models.py index 296ac9ca06..eb88a403d0 100644 --- a/src/sqlbuild/lint/models.py +++ b/src/sqlbuild/lint/models.py @@ -15,7 +15,13 @@ VIOLATION_SEVERITY_FAULT, VIOLATION_SEVERITY_WARNING, ) -from sqlbuild.lint.types import DiagnosticIdentity, LintSeverity, RuleFixStatus +from sqlbuild.lint.exceptions import ProjectCompileError +from sqlbuild.lint.types import ( + DiagnosticIdentity, + LintSeverity, + NativeLintPreparationRequest, + RuleFixStatus, +) @dataclass(frozen=True) @@ -66,6 +72,41 @@ class LintBody: dependency_relations: tuple[tuple[str, str], ...] = () +@dataclass(frozen=True) +class PendingLintBody: + """One expanded authored body awaiting its native lexical preparation.""" + + file_path: Path + body_start: int + body_end: int + expansion_input: str + expanded: str + expansion_passes: tuple[tuple[ExpansionSpan, ...], ...] + pre_expansion_sites: tuple[InterpolationSite, ...] + external_identifiers: tuple[str, ...] + allows_ceremonial_select: bool + allows_dynamic_output_star: bool + allows_empty_fixture_star: bool + request: NativeLintPreparationRequest + + +@dataclass(frozen=True) +class ExpandedLintFile: + """The expanded bodies of one file, up to the first body that could not be expanded.""" + + file_path: Path + bodies: tuple[PendingLintBody, ...] + failure: ProjectCompileError | None + + +@dataclass(frozen=True) +class PreparedLintFiles: + """Lint bodies of every expandable file, and why the skipped files could not expand.""" + + bodies: tuple[LintBody, ...] + unexpandable: dict[Path, str] + + @dataclass(frozen=True) class FixtureNullCandidateScan: """Cheaply discovered SQL-test fixtures that may contain removable typed nulls.""" diff --git a/src/sqlbuild/lint/types.py b/src/sqlbuild/lint/types.py index 6909aadd68..82135aafcf 100644 --- a/src/sqlbuild/lint/types.py +++ b/src/sqlbuild/lint/types.py @@ -3,7 +3,10 @@ from __future__ import annotations from enum import StrEnum -from typing import Literal +from typing import TYPE_CHECKING, Literal + +if TYPE_CHECKING: + from sqlbuild._native import LintPreparationRequest class LintSeverity(StrEnum): @@ -18,3 +21,7 @@ class LintSeverity(StrEnum): type RelationKeys = dict[str, tuple[tuple[str, ...], ...]] type DiagnosticIdentity = tuple[str, str, str, str] + +type NativeLintPreparationRequest = LintPreparationRequest + +type NativePreparedSql = tuple[str, list[tuple[str, int, int, int, int, str]], list[str]] diff --git a/tests/unit/src/sqlbuild/lint/_helpers/_test_types.py b/tests/unit/src/sqlbuild/lint/_helpers/_test_types.py index 3b4427709d..b49d2a90bb 100644 --- a/tests/unit/src/sqlbuild/lint/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/lint/_helpers/_test_types.py @@ -215,3 +215,22 @@ class ManyViolationLocationTestCase: unused_cte_count: int file_count: int expected_code: str + + +@dataclass(frozen=True) +class BatchedPreparationTestCase: + """A project whose lint is compared between batched and per-body native preparation.""" + + description: str + files: dict[str, str] + expected_codes: tuple[str, ...] + + +@dataclass(frozen=True) +class BatchedPreparationFailureTestCase: + """A native preparation failure and an unexpandable file, in a given file order.""" + + description: str + unexpandable_file: str + skip_unexpandable: bool + expected_error: type[Exception] | None diff --git a/tests/unit/src/sqlbuild/lint/_helpers/helpers.py b/tests/unit/src/sqlbuild/lint/_helpers/helpers.py index 9863d1f4fb..cf3135758f 100644 --- a/tests/unit/src/sqlbuild/lint/_helpers/helpers.py +++ b/tests/unit/src/sqlbuild/lint/_helpers/helpers.py @@ -6,6 +6,7 @@ from sqlbuild.lint._helpers.headers import scan_headers, sql_body_ranges from sqlbuild.lint.models import HeaderSpan, LintBody +from sqlbuild.lint.types import NativeLintPreparationRequest, NativePreparedSql def lint_bodies_for(*, file_path: Path, contents: str) -> tuple[LintBody, ...]: @@ -64,3 +65,29 @@ def authored_positions(*, contents: str, needles: tuple[str, ...]) -> list[tuple (contents.count("\n", 0, offset) + 1, offset - contents.rfind("\n", 0, offset)) ) return positions + + +def write_lint_project(*, root: Path, files: dict[str, str]) -> None: + """Write a DuckDB project containing the given relative files.""" + + (root / "sqlbuild_project.toml").write_text( + 'name = "orders"\nadapter = "duckdb"\n', encoding="utf-8" + ) + for relative_path, contents in files.items(): + path: Path = root / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(contents, encoding="utf-8") + + +def prepare_nothing_in_batch( + requests: list[NativeLintPreparationRequest], +) -> list[tuple[bool, NativePreparedSql | None]]: + """Report every body as unprepared so each one takes the per-body native call.""" + + return [(False, None) for _request in requests] + + +def refuse_preparation(request: NativeLintPreparationRequest) -> NativePreparedSql | None: + """Fail the per-body native preparation of any body.""" + + raise ValueError(f"cannot prepare {request['dialect']} body") diff --git a/tests/unit/src/sqlbuild/lint/_helpers/test_expansion.py b/tests/unit/src/sqlbuild/lint/_helpers/test_expansion.py new file mode 100644 index 0000000000..61b810260e --- /dev/null +++ b/tests/unit/src/sqlbuild/lint/_helpers/test_expansion.py @@ -0,0 +1,116 @@ +"""Lint bodies are prepared natively in one batch with per-body results and failure order.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +import sqlbuild._native as native_module +from sqlbuild.lint.exceptions import ProjectCompileError +from sqlbuild.lint.main.run_lint import run_lint +from sqlbuild.lint.models import LintConfig, LintRunResult +from tests.unit.src.sqlbuild.lint._helpers._test_types import ( + BatchedPreparationFailureTestCase, + BatchedPreparationTestCase, +) +from tests.unit.src.sqlbuild.lint._helpers.helpers import ( + prepare_nothing_in_batch, + refuse_preparation, + write_lint_project, +) + +_ORDERS_MODEL: str = ( + 'MODEL (description "Orders");\n' + 'SELECT o.customer_id FROM __ref("customers") AS o WHERE o.region = NULL\n' +) +_CUSTOMERS_MODEL: str = ( + "MODEL (description \"Customers\");\nSELECT 1 AS customer_id, 'west' AS region\n" +) +_ORDERS_TEST: str = ( + "TEST ();\nWITH\n__ref__customers AS (SELECT 1 AS customer_id, 'west' AS region),\n" + "__expected__orders AS (SELECT 1 AS customer_id)\nSELECT 1\n" +) +_NOT_NULL_AUDIT: str = "AUDIT ();\nSELECT 1 FROM @relation WHERE @column IS NULL\n" + + +@pytest.mark.parametrize( + "test_case", + [ + BatchedPreparationTestCase( + description="models, tests, and generic audits", + files={ + "models/customers.sql": _CUSTOMERS_MODEL, + "models/orders.sql": _ORDERS_MODEL, + "tests/unit/test_orders.sql": _ORDERS_TEST, + "audits/generic/not_null.sql": _NOT_NULL_AUDIT, + }, + expected_codes=("SQBRSQL001",), + ) + ], + ids=lambda case: case.description, +) +def test_given_project_when_linting_with_batched_preparation_then_result_matches_per_body( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: BatchedPreparationTestCase +) -> None: + write_lint_project(root=tmp_path, files=test_case.files) + config: LintConfig = LintConfig(dialect="duckdb") + + batched: LintRunResult = run_lint(project_dir=tmp_path, config=config) + monkeypatch.setattr(native_module, "prepare_lint_sql_batch", prepare_nothing_in_batch) + per_body: LintRunResult = run_lint(project_dir=tmp_path, config=config) + + assert batched == per_body + assert tuple(violation.code for violation in batched.violations) == test_case.expected_codes + + +@pytest.mark.parametrize( + "test_case", + [ + BatchedPreparationFailureTestCase( + description="native failure before a later unexpandable file", + unexpandable_file="models/zz_broken.sql", + skip_unexpandable=False, + expected_error=ValueError, + ), + BatchedPreparationFailureTestCase( + description="unexpandable file before a later native failure", + unexpandable_file="models/aa_broken.sql", + skip_unexpandable=False, + expected_error=ProjectCompileError, + ), + BatchedPreparationFailureTestCase( + description="skipped unexpandable file before a later native failure", + unexpandable_file="models/aa_broken.sql", + skip_unexpandable=True, + expected_error=ValueError, + ), + ], + ids=lambda case: case.description, +) +def test_given_native_and_expansion_failures_when_linting_then_first_in_file_order_is_raised( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: BatchedPreparationFailureTestCase +) -> None: + write_lint_project( + root=tmp_path, + files={ + "models/customers.sql": _CUSTOMERS_MODEL, + "models/orders.sql": _ORDERS_MODEL, + test_case.unexpandable_file: 'MODEL (description "Broken");\nSELECT @unknown()\n', + }, + ) + monkeypatch.setattr(native_module, "prepare_lint_sql_batch", prepare_nothing_in_batch) + monkeypatch.setattr(native_module, "prepare_lint_sql", refuse_preparation) + + with pytest.raises(Exception) as raised: + _ = run_lint( + project_dir=tmp_path, + config=LintConfig(dialect="duckdb"), + skip_unexpandable=test_case.skip_unexpandable, + ) + + assert type(raised.value) is test_case.expected_error + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) From 5bf7c4a73fc61deed3e736cba052ce83d02e2e9d Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 4 Oct 2026 00:09:52 +0100 Subject: [PATCH 3/3] fix(rules): keep split custom-rule hosts exact, bounded, and cancellable Rules that change module-level state on any host are evaluated again in one host over all of their subjects, replacing their split results, so cross-subject state such as a duplicate detector reports exactly what one host reports. The host count now honours cgroup v1 and v2 CPU quotas and is capped at eight. The first failing host writes a per-run cancellation marker. The other hosts check it between invocations, outside any Rule's guarded scope, and exit with a distinct cancelled status, and hosts that start late exit at once; no process is ever signalled. One invocation that is already running still ends or reaches its own timeout. One host then reruns the plan to report the single-host error. --- .../custom-rules/testing-and-determinism.md | 21 +++ .../_helpers/host/custom_evaluation.py | 58 ++++++-- .../rule_engine/_helpers/host/custom_host.py | 38 ++++- .../_helpers/host/custom_host_pool.py | 96 ++++++++++--- .../_helpers/host/host_capacity.py | 103 ++++++++++++++ src/sqlbuild/rule_engine/constants.py | 3 + src/sqlbuild/rule_engine/exceptions.py | 4 + src/sqlbuild/rule_engine/models.py | 10 ++ .../sqlbuild/cli/commands/main/test_rules.py | 2 +- .../rule_engine/_helpers/host/_test_types.py | 25 ++++ .../rule_engine/_helpers/host/helpers.py | 48 ++++++- .../_helpers/host/test_custom_host_pool.py | 127 +++++++++++++++-- .../_helpers/host/test_host_capacity.py | 134 ++++++++++++++++++ .../custom-rules/testing-and-determinism.mdx | 13 ++ 14 files changed, 628 insertions(+), 54 deletions(-) create mode 100644 src/sqlbuild/rule_engine/_helpers/host/host_capacity.py create mode 100644 tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_host_capacity.py diff --git a/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/rules/custom-rules/testing-and-determinism.md b/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/rules/custom-rules/testing-and-determinism.md index c3641a48d4..ef11cb0bf5 100644 --- a/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/rules/custom-rules/testing-and-determinism.md +++ b/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/rules/custom-rules/testing-and-determinism.md @@ -6,6 +6,14 @@ Online: https://sqlbuild.com/docs/concepts/rules/custom-rules/testing-and-determinism/ +## Contents + +- Test Rules +- Parallel evaluation +- Helpers +- Deterministic inputs +- Cache granularity + ## Test Rules Use the public harness to exercise discovery, compilation, and Rules evaluation: @@ -39,6 +47,19 @@ the next. Compute findings from the subject and `ctx` alone. Repository pytest files remain outside the SQLBuild project's SQL `tests/` directory. +## Parallel evaluation + +On large projects, SQLBuild splits custom Rule evaluation across several host processes, at most +one per available CPU and never more than eight. Each process evaluates a share of the subjects, +and SQLBuild merges the findings back into the order a single process reports them. + +A Rule that changes module-level state while it runs, for example by adding each model to a set of +already-seen names, could report different findings when its subjects are split. SQLBuild detects +the state change and evaluates that Rule again in one process over all of its subjects, so its +findings and cache entries match a single-process run. The result stays correct, but the Rule no +longer runs in parallel and costs an extra host process. Keep each subject's findings independent +of the others so Rules can run in parallel. + ## Helpers Custom Rules can import repository-owned helpers under `rules/` and these deterministic standard diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py b/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py index 8884692f01..832e467da7 100644 --- a/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py @@ -3,6 +3,7 @@ from __future__ import annotations from collections import Counter +from collections.abc import Callable from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass from itertools import chain @@ -22,13 +23,18 @@ build_rule_fact_views, ) from sqlbuild.rule_engine.classes.runtime_guard import RuntimeGuard -from sqlbuild.rule_engine.constants import CUSTOM_RULE_PROJECT_SUBJECT +from sqlbuild.rule_engine.constants import ( + CUSTOM_HOST_CANCELLED_MESSAGE, + CUSTOM_RULE_PROJECT_SUBJECT, +) from sqlbuild.rule_engine.exceptions import ( + HostCancelledError, NonHermeticRuleError, OpaqueModuleStateError, RulesError, ) from sqlbuild.rule_engine.models import ( + CustomHostPartition, CustomRuleEvaluation, CustomRuleRun, Finding, @@ -52,6 +58,7 @@ class _PassInputs: facts: RuleFactViews plan: CustomRulePlan | None guard: RuntimeGuard | None + cancelled: Callable[[], bool] | None def evaluate_custom_rules( @@ -65,6 +72,7 @@ def evaluate_custom_rules( track_reads: bool = False, verify_determinism: bool = False, guard: RuntimeGuard | None = None, + partition: CustomHostPartition | None = None, ) -> CustomRuleRun: """Run selected custom rules for every planned subject through the Python authoring API.""" @@ -81,7 +89,10 @@ def evaluate_custom_rules( ) reads: FactReads | None = FactReads() if track_reads else None tracker: _ReadTracker = _ReadTracker( - reads=reads, rules=custom_rules, rules_root=project_dir / "rules" + reads=reads, + rules=custom_rules, + rules_root=project_dir / "rules", + detect_state=partition is not None, ) inputs: _PassInputs = _PassInputs( @@ -93,6 +104,7 @@ def evaluate_custom_rules( facts=facts, plan=plan, guard=guard, + cancelled=None if partition is None else partition.cancelled, ) def run_pass( @@ -119,6 +131,7 @@ def run_pass( untracked_codes=frozenset(tracker.untracked_codes), uncacheable_codes=frozenset(tracker.uncacheable_codes), observed=() if reads is None else tuple(reads.observed.items()), + stateful_codes=frozenset(tracker.stateful_codes), ) @@ -126,15 +139,21 @@ class _ReadTracker: """Attribute reads to one invocation; rules whose module state changes become untracked.""" def __init__( - self, *, reads: FactReads | None, rules: tuple[Rule, ...], rules_root: Path + self, + *, + reads: FactReads | None, + rules: tuple[Rule, ...], + rules_root: Path, + detect_state: bool = False, ) -> None: self.reads: FactReads | None = reads self.untracked_codes: set[str] = set() self.uncacheable_codes: set[str] = set() + self.stateful_codes: set[str] = set() self._rules_root: Path = rules_root self._namespaces: dict[str, tuple[dict[str, object], ...]] = ( {rule.code: rule_namespaces((rule.check,)) for rule in rules} - if reads is not None + if reads is not None or detect_state else {} ) self._states: dict[str, str | None] = { @@ -149,20 +168,22 @@ def start(self) -> None: def finish(self, *, code: str) -> frozenset[tuple[object, ...]] | None: reads: FactReads | None = self.reads - if reads is None: - return None - observed: set[tuple[object, ...]] | None = reads.current - reads.current = None - if reads.untracked: - self.untracked_codes.add(code) - if reads.uncacheable: - self.uncacheable_codes.add(code) + observed: set[tuple[object, ...]] | None = None + if reads is not None: + observed = reads.current + reads.current = None + if reads.untracked: + self.untracked_codes.add(code) + if reads.uncacheable: + self.uncacheable_codes.add(code) for rule_code, namespaces in self._namespaces.items(): - if rule_code in self.untracked_codes: + if rule_code in self.stateful_codes: continue state: str | None = self._state_token(namespaces) if state is None or state != self._states[rule_code]: - self.untracked_codes.add(rule_code) + self.stateful_codes.add(rule_code) + if reads is not None: + self.untracked_codes.add(rule_code) return None if observed is None else frozenset(observed) def _state_token(self, namespaces: tuple[dict[str, object], ...]) -> str | None: @@ -216,6 +237,7 @@ def context(rule: Rule) -> EvaluationRuleContext: if planned_project is not None and CUSTOM_RULE_PROJECT_SUBJECT not in planned_project: continue ctx: EvaluationRuleContext = context(rule) + _stop_if_cancelled(inputs) tracker.start() try: findings: list[Finding] = _invoke( @@ -245,6 +267,7 @@ def context(rule: Rule) -> EvaluationRuleContext: for rule, ctx, planned in model_contexts: if planned is not None and path not in planned: continue + _stop_if_cancelled(inputs) tracker.start() try: findings = _invoke(rule=rule, subject=subject, ctx=ctx, guard=guard) @@ -265,6 +288,13 @@ def context(rule: Rule) -> EvaluationRuleContext: return evaluations +def _stop_if_cancelled(inputs: _PassInputs) -> None: + """Stop between invocations, outside any Rule's guarded scope, once the run is cancelled.""" + + if inputs.cancelled is not None and inputs.cancelled(): + raise HostCancelledError(CUSTOM_HOST_CANCELLED_MESSAGE) + + def _invoke( *, rule: Rule, subject: object, ctx: EvaluationRuleContext, guard: RuntimeGuard | None ) -> list[Finding]: diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_host.py b/src/sqlbuild/rule_engine/_helpers/host/custom_host.py index 12701166b9..f19d2b5347 100644 --- a/src/sqlbuild/rule_engine/_helpers/host/custom_host.py +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_host.py @@ -18,6 +18,8 @@ from sqlbuild.rule_engine._helpers.host.custom_evaluation import evaluate_custom_rules from sqlbuild.rule_engine.classes.runtime_guard import RuntimeGuard from sqlbuild.rule_engine.constants import ( + CUSTOM_HOST_CANCELLED_EXIT_CODE, + CUSTOM_HOST_CANCELLED_MESSAGE, CUSTOM_HOST_HASH_SEED, CUSTOM_HOST_INPUT_TUPLE_SIZE, CUSTOM_HOST_LAUNCH_MODULE, @@ -25,8 +27,14 @@ CUSTOM_HOST_PROTOCOL_VERSION, CUSTOM_HOST_RUNTIME_VERSION, ) -from sqlbuild.rule_engine.exceptions import RulesError -from sqlbuild.rule_engine.models import CustomRuleRun, Finding, Rule, RulesConfig +from sqlbuild.rule_engine.exceptions import HostCancelledError, RulesError +from sqlbuild.rule_engine.models import ( + CustomHostPartition, + CustomRuleRun, + Finding, + Rule, + RulesConfig, +) from sqlbuild.rule_engine.types import CustomRulePlan @@ -50,6 +58,9 @@ def main() -> int: payload: object = request["payload"] if not isinstance(payload, dict): raise RulesError("custom host payload must be an object") + cancel_marker: Path | None = _cancel_marker(payload) + if cancel_marker is not None and cancel_marker.exists(): + raise HostCancelledError(CUSTOM_HOST_CANCELLED_MESSAGE) project, config = _decode_inputs(payload) project_dir: Path = Path(str(payload["project_dir"])).resolve() dialect: str = str(payload.get("dialect", "generic")) @@ -77,7 +88,15 @@ def main() -> int: track_reads=track_reads, verify_determinism=verify_determinism, guard=guard, + partition=( + None + if cancel_marker is None + else CustomHostPartition(cancelled=cancel_marker.exists) + ), ) + except HostCancelledError: + sys.stderr.write(CUSTOM_HOST_CANCELLED_MESSAGE) + return CUSTOM_HOST_CANCELLED_EXIT_CODE except Exception as error: violation: Exception | None = None if guard is None else guard.violation return _write_error(str(violation or error)) @@ -114,6 +133,20 @@ def _decode_inputs(payload: dict[str, Any]) -> tuple[CompiledProject, RulesConfi return decoded +def _cancel_marker(payload: dict[str, Any]) -> Path | None: + """Return the run's cancellation marker, which must live beside the project payload.""" + + value: object = payload.get("cancel_marker") + if value is None: + return None + project_dir: Path = Path(str(payload["project_dir"])).resolve() + input_root: Path = (project_dir / "target" / "rules-cache" / "host-inputs").resolve() + marker: Path = Path(str(value)) + if not marker.is_absolute() or not marker.resolve().is_relative_to(input_root): + raise RulesError("custom host cancellation marker path is invalid") + return marker + + def _decode_plan(value: object) -> CustomRulePlan: if not isinstance(value, dict): raise RulesError("custom host plan must be an object") @@ -164,6 +197,7 @@ def _run_payload(run: CustomRuleRun) -> dict[str, object]: "uncacheable": sorted(run.uncacheable_codes), "observed": sorted([list(key), digest] for key, digest in run.observed), "bounded": bounded, + "stateful": sorted(run.stateful_codes), } diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py b/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py index 9031eeca15..e84492b29a 100644 --- a/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py @@ -2,11 +2,12 @@ from __future__ import annotations -import os import pickle +import shutil import sys import tempfile -from concurrent.futures import Future, ThreadPoolExecutor +from concurrent.futures import FIRST_EXCEPTION, Future, ThreadPoolExecutor, wait +from dataclasses import replace from pathlib import Path from typing import Any, cast @@ -14,7 +15,9 @@ import sqlbuild._native as _native from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.rule_engine._helpers.host.host_capacity import available_cores from sqlbuild.rule_engine.constants import ( + CUSTOM_HOST_CANCELLED_MARKER, CUSTOM_HOST_LAUNCH_MODULE, CUSTOM_HOST_MAX_TRACKED_READS, CUSTOM_HOST_RUNTIME_VERSION, @@ -32,6 +35,7 @@ _HOST_TIMEOUT_MILLIS: int = 120_000 _MIN_INVOCATIONS_PER_HOST: int = 64 +_MAX_HOSTS: int = 8 def run_custom_hosts( @@ -55,7 +59,7 @@ def run_custom_hosts( (plan,) if verify_determinism else partition_host_plan( - plan=plan, invocations=invocations, hosts=_host_count(len(invocations)) + plan=plan, invocations=invocations, hosts=host_count(len(invocations)) ) ) input_dir: Path = project_dir / "target" / "rules-cache" / "host-inputs" @@ -77,7 +81,7 @@ def run_custom_hosts( responses: list[CustomHostSlice] = ( [_run_host(payload={**payload, "plan": plan})] if len(plans) == 1 - else _run_partitions(payload=payload, plan=plan, plans=plans) + else _run_partitions(payload=payload, plan=plan, plans=plans, input_dir=input_dir) ) finally: if input_path is not None: @@ -89,20 +93,72 @@ def run_custom_hosts( def _run_partitions( - *, payload: dict[str, object], plan: CustomHostPlan, plans: tuple[CustomHostPlan, ...] + *, + payload: dict[str, object], + plan: CustomHostPlan, + plans: tuple[CustomHostPlan, ...], + input_dir: Path, ) -> list[CustomHostSlice]: - """Run every partition; after any failure, one host reruns the whole plan to report it.""" - - with ThreadPoolExecutor( - max_workers=len(plans), thread_name_prefix="sqlbuild-custom-host" - ) as pool: - futures: list[Future[CustomHostSlice]] = [ - pool.submit(_run_host, payload={**payload, "plan": partition}) for partition in plans - ] - failed: bool = any(future.exception() is not None for future in futures) + """Run partitions; failures and rules with module state fall back to one host.""" + + cancel_dir: Path = Path(tempfile.mkdtemp(dir=input_dir, prefix="hosts-")) + marker: Path = cancel_dir / CUSTOM_HOST_CANCELLED_MARKER + futures: list[Future[CustomHostSlice]] = [] + failed: bool = False + try: + with ThreadPoolExecutor( + max_workers=len(plans), thread_name_prefix="sqlbuild-custom-host" + ) as pool: + futures = [ + pool.submit( + _run_host, + payload={ + **payload, + "plan": partition, + "cancel_marker": str(marker.resolve()), + }, + ) + for partition in plans + ] + done: set[Future[CustomHostSlice]] = wait(futures, return_when=FIRST_EXCEPTION).done + failed = any(future.exception() is not None for future in done) + if failed: + _cancel_hosts(marker) + finally: + shutil.rmtree(cancel_dir, ignore_errors=True) if failed: return [_run_host(payload={**payload, "plan": plan})] - return [future.result() for future in futures] + slices: list[CustomHostSlice] = [future.result() for future in futures] + stateful: frozenset[str] = frozenset().union(*(item.stateful for item in slices)) + if not stateful: + return slices + single: CustomHostSlice = _run_host( + payload={ + **payload, + "plan": {code: plan[code] if code in stateful else [] for code in plan}, + }, + ) + return [*(_without_codes(response=item, codes=stateful) for item in slices), single] + + +def _cancel_hosts(marker: Path) -> None: + """Stop hosts at their next invocation; a running invocation still ends or times out.""" + + marker.touch() + + +def _without_codes(*, response: CustomHostSlice, codes: frozenset[str]) -> CustomHostSlice: + """Drop every result a host reported for the given rules.""" + + return replace( + response, + evaluations=tuple(item for item in response.evaluations if item[0] not in codes), + untracked_reads={ + code: reads for code, reads in response.untracked_reads.items() if code not in codes + }, + uncacheable=response.uncacheable - codes, + stateful=response.stateful - codes, + ) def planned_invocations( @@ -152,13 +208,10 @@ def partition_host_plan( return tuple(cast(CustomHostPlan, partition) for partition in partitions) -def _host_count(invocations: int) -> int: - return max(1, min(_available_cores(), invocations // _MIN_INVOCATIONS_PER_HOST)) - +def host_count(invocations: int) -> int: + """Hosts worth starting: one per 64 invocations, bounded by usable CPUs and eight.""" -def _available_cores() -> int: - affinity: Any = getattr(os, "sched_getaffinity", None) - return len(affinity(0)) if affinity is not None else os.cpu_count() or 1 + return max(1, min(_MAX_HOSTS, available_cores(), invocations // _MIN_INVOCATIONS_PER_HOST)) def _run_host(*, payload: dict[str, object]) -> CustomHostSlice: @@ -231,6 +284,7 @@ def _decode_host_response(response: object) -> CustomHostSlice: for key, digest in raw_observed }, bounded=payload.get("bounded") is True, + stateful=frozenset(str(code) for code in payload.get("stateful") or ()), ) diff --git a/src/sqlbuild/rule_engine/_helpers/host/host_capacity.py b/src/sqlbuild/rule_engine/_helpers/host/host_capacity.py new file mode 100644 index 0000000000..0f6d4e5840 --- /dev/null +++ b/src/sqlbuild/rule_engine/_helpers/host/host_capacity.py @@ -0,0 +1,103 @@ +"""CPU capacity available to this process, including container CPU quotas.""" + +from __future__ import annotations + +import math +import os +from collections.abc import Callable, Iterator +from pathlib import Path + +_CGROUP_ROOT: Path = Path("/sys/fs/cgroup") +_PROC_CGROUP: Path = Path("/proc/self/cgroup") +_CGROUP_V1_CPU_CONTROLLER: str = "cpu" +_CGROUP_V1_CPU_DIRECTORIES: tuple[str, ...] = ("cpu", "cpu,cpuacct", "cpuacct,cpu") +_CPU_MAX_FIELD_COUNT: int = 2 + + +def available_cores(*, cgroup_root: Path = _CGROUP_ROOT, proc_cgroup: Path = _PROC_CGROUP) -> int: + """Return usable CPUs: the process CPU count capped by any cgroup v1 or v2 CPU quota.""" + + counter: Callable[[], int | None] | None = getattr(os, "process_cpu_count", None) + count: int = (counter() if counter is not None else _affinity_count()) or 1 + quotas: list[int] = list(_cgroup_quotas(cgroup_root=cgroup_root, proc_cgroup=proc_cgroup)) + return max(1, min((count, *quotas))) + + +def _affinity_count() -> int | None: + affinity: Callable[[int], set[int]] | None = getattr(os, "sched_getaffinity", None) + return len(affinity(0)) if affinity is not None else os.cpu_count() + + +def _cgroup_quotas(*, cgroup_root: Path, proc_cgroup: Path) -> Iterator[int]: + try: + entries: list[str] = proc_cgroup.read_text(encoding="utf-8").splitlines() + except (OSError, ValueError): + entries = [] + for entry in entries or ["0::/"]: + _hierarchy, _separator, remainder = entry.partition(":") + controllers, _separator, path = remainder.partition(":") + if not controllers: + yield from _quotas_in( + paths=tuple( + directory / "cpu.max" + for directory in _directories(base=cgroup_root, relative=path) + ), + parse=_cpu_max_quota, + ) + elif _CGROUP_V1_CPU_CONTROLLER in controllers.split(","): + for name in (controllers, *_CGROUP_V1_CPU_DIRECTORIES): + yield from _quotas_in( + paths=tuple(_directories(base=cgroup_root / name, relative=path)), + parse=_cfs_quota, + ) + + +def _directories(*, base: Path, relative: str) -> Iterator[Path]: + """Yield the cgroup directory and each ancestor down from the hierarchy root.""" + + directory: Path = base + yield directory + for part in Path(relative.strip("/")).parts: + directory = directory / part + yield directory + + +def _quotas_in(*, paths: tuple[Path, ...], parse: Callable[[Path], int | None]) -> Iterator[int]: + for path in paths: + quota: int | None = parse(path) + if quota is not None: + yield quota + + +def _cpu_max_quota(path: Path) -> int | None: + """Parse cgroup v2 `cpu.max`: `max ` means unlimited, else ` `.""" + + try: + fields: list[str] = path.read_text(encoding="utf-8").split() + except (OSError, ValueError): + return None + return ( + _quota(quota=fields[0], period=fields[1]) if len(fields) == _CPU_MAX_FIELD_COUNT else None + ) + + +def _cfs_quota(directory: Path) -> int | None: + """Parse cgroup v1 `cpu.cfs_quota_us` / `cpu.cfs_period_us`; a negative quota is unlimited.""" + + try: + quota: str = (directory / "cpu.cfs_quota_us").read_text(encoding="utf-8").strip() + period: str = (directory / "cpu.cfs_period_us").read_text(encoding="utf-8").strip() + except (OSError, ValueError): + return None + return _quota(quota=quota, period=period) + + +def _quota(*, quota: str, period: str) -> int | None: + try: + quota_us: int = int(quota) + period_us: int = int(period) + except ValueError: + return None + if quota_us <= 0 or period_us <= 0: + return None + return max(1, math.ceil(quota_us / period_us)) diff --git a/src/sqlbuild/rule_engine/constants.py b/src/sqlbuild/rule_engine/constants.py index 4f9230b090..98078ffac1 100644 --- a/src/sqlbuild/rule_engine/constants.py +++ b/src/sqlbuild/rule_engine/constants.py @@ -17,6 +17,9 @@ NATIVE_RULES_MEMO_FILE: str = "target/rules-cache/bulk/native-response.json" CUSTOM_HOST_INPUT_TUPLE_SIZE: int = 2 CUSTOM_HOST_MODULE: str = "sqlbuild.rule_engine._helpers.host.custom_host" +CUSTOM_HOST_CANCELLED_MARKER: str = "cancelled" +CUSTOM_HOST_CANCELLED_EXIT_CODE: int = 75 +CUSTOM_HOST_CANCELLED_MESSAGE: str = "custom host cancelled after another host failed" CUSTOM_HOST_LAUNCH_MODULE: str = "sqlbuild.rule_engine._helpers.host.custom_host_launch" CUSTOM_HOST_HASH_SEED: str = "0" CUSTOM_HOST_EXEC_OS_NAME: str = "posix" diff --git a/src/sqlbuild/rule_engine/exceptions.py b/src/sqlbuild/rule_engine/exceptions.py index 2d164a99a2..9786c0c1f5 100644 --- a/src/sqlbuild/rule_engine/exceptions.py +++ b/src/sqlbuild/rule_engine/exceptions.py @@ -23,3 +23,7 @@ class FactDigestError(RulesError): class OpaqueModuleStateError(RulesError): """Raised when custom-rule module state holds a value whose mutations cannot be observed.""" + + +class HostCancelledError(RulesError): + """Raised in a custom-rule host that stops because another host already failed.""" diff --git a/src/sqlbuild/rule_engine/models.py b/src/sqlbuild/rule_engine/models.py index 085353ec41..99055f59b4 100644 --- a/src/sqlbuild/rule_engine/models.py +++ b/src/sqlbuild/rule_engine/models.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Callable from concurrent.futures import Future from dataclasses import dataclass, field from pathlib import Path @@ -401,6 +402,7 @@ class CustomRuleRun: untracked_codes: frozenset[str] uncacheable_codes: frozenset[str] = frozenset() observed: tuple[tuple[tuple[str, ...], str | None], ...] = () + stateful_codes: frozenset[str] = frozenset() @dataclass(frozen=True) @@ -440,3 +442,11 @@ class CustomHostSlice: uncacheable: frozenset[str] observed: dict[FactKey, str | None] bounded: bool + stateful: frozenset[str] + + +@dataclass(frozen=True) +class CustomHostPartition: + """How one host of a split run detects module state and learns the run was cancelled.""" + + cancelled: Callable[[], bool] diff --git a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py index 0bef0a630d..c8c4cb3bc5 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py @@ -1089,7 +1089,7 @@ def recording_host(spec_json: str) -> str: def compile_payload(*, hosts: int) -> dict[str, Any]: monkeypatch.setattr(custom_host_pool_module, "_MIN_INVOCATIONS_PER_HOST", 1) - monkeypatch.setattr(custom_host_pool_module, "_available_cores", lambda: hosts) + monkeypatch.setattr(custom_host_pool_module, "available_cores", lambda: hosts) exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) assert exit_code == 1 return json.loads(capsys.readouterr().out) diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py index 7596c1c17d..6ae0c71543 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py @@ -12,6 +12,8 @@ class HostPartitionTestCase: hosts: int expected_host_runs: int records_module_state: bool = False + detects_duplicates: bool = False + cache_enabled: bool = True @dataclass(frozen=True) @@ -33,3 +35,26 @@ class HostPlanTestCase: invocations: tuple[tuple[str, str], ...] hosts: int expected_plans: tuple[dict[str, list[str] | None], ...] + + +@dataclass(frozen=True) +class HostCancellationTestCase: + """A fast failure on one host while the other hosts are still busy.""" + + description: str + model_count: int + hosts: int + timeout_millis: int + delay_seconds: float + expected_error_fragment: str + expected_max_seconds: float + + +@dataclass(frozen=True) +class HostCapacityTestCase: + """Fake cgroup files and the CPU count they allow.""" + + description: str + proc_cgroup: str + files: dict[str, str] + expected_cores: int diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py index 8f7c52cf6d..cf8d7e815b 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py @@ -1,5 +1,6 @@ """Builders for custom-rule host partitioning tests.""" +import os from collections.abc import Callable from dataclasses import replace from operator import attrgetter @@ -24,10 +25,16 @@ def even_suffix(*, model: Model, ctx: RuleContext) -> list[Finding]: if model.name in FAILING: raise ValueError(f"cannot inspect {{model.name}}") + time.sleep(DELAY_SECONDS) if RECORD_SEEN: SEEN[model.name] = True source: str = ctx.sql.for_model(model).expanded.source - return [ctx.finding(subject=model)] if "odd" in source else [] + parity: str = "odd" if "odd" in source else "even" + if DETECT_DUPLICATES: + duplicate: bool = parity in SEEN + SEEN[parity] = True + return [ctx.finding(subject=model)] if duplicate else [] + return [ctx.finding(subject=model)] if parity == "odd" else [] @rule(code="{_PROJECT_RULE}", message="projects need few models", remediation="Split it.") @@ -40,14 +47,20 @@ def few_models(*, project: Project, ctx: RuleContext) -> list[Finding]: def write_order_rules( - *, root: Path, failing: tuple[str, ...] = (), records_module_state: bool = False + *, + root: Path, + failing: tuple[str, ...] = (), + records_module_state: bool = False, + detects_duplicates: bool = False, + delay_seconds: float = 0.0, ) -> None: - """Write one model rule and one project rule, optionally failing or keeping module state.""" + """Write one model rule and one project rule with optional failures, state, or delays.""" rules: Path = root / "rules" / "orders.py" rules.parent.mkdir(parents=True, exist_ok=True) rules.write_text( - f"FAILING = {set(failing)!r}\nRECORD_SEEN = {records_module_state!r}\n" + f"import time\n\nFAILING = {set(failing)!r}\nRECORD_SEEN = {records_module_state!r}\n" + f"DETECT_DUPLICATES = {detects_duplicates!r}\nDELAY_SECONDS = {delay_seconds!r}\n" f"SEEN: dict[str, bool] = {{}}\n{_RULES_MODULE}", encoding="utf-8", ) @@ -95,6 +108,31 @@ def recording_launch(spec_json: str) -> str: return launch(spec_json) monkeypatch.setattr(custom_host_pool, "_MIN_INVOCATIONS_PER_HOST", 1) - monkeypatch.setattr(custom_host_pool, "_available_cores", lambda: hosts) + monkeypatch.setattr(custom_host_pool, "available_cores", lambda: hosts) monkeypatch.setattr(native_module, "run_custom_host_json", recording_launch) return launches + + +def write_cgroup_files(*, root: Path, proc_cgroup: str, files: dict[str, str]) -> Path: + """Write a fake `/proc/self/cgroup` and cgroup hierarchy beneath `root`.""" + + for relative_path, contents in files.items(): + path: Path = root / "cgroup" / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(contents, encoding="utf-8", errors="surrogateescape") + proc: Path = root / "proc-self-cgroup" + proc.write_text(proc_cgroup, encoding="utf-8", errors="surrogateescape") + return proc + + +def record_signals(*, monkeypatch: pytest.MonkeyPatch) -> list[tuple[int, int]]: + """Replace process signalling with a recorder so no process can be signalled.""" + + signals: list[tuple[int, int]] = [] + + def record(pid: int, signal_number: int) -> None: + signals.append((pid, signal_number)) + + monkeypatch.setattr(os, "kill", record) + monkeypatch.setattr(os, "killpg", record, raising=False) + return signals diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py index 86c402b735..7f1eccf405 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py @@ -1,5 +1,6 @@ """Custom rules split across host processes report exactly what one host reports.""" +import time from pathlib import Path import pytest @@ -10,6 +11,7 @@ from sqlbuild.rule_engine.exceptions import RulesError from sqlbuild.rule_engine.models import CustomRulesOutcome from tests.unit.src.sqlbuild.rule_engine._helpers.host._test_types import ( + HostCancellationTestCase, HostFailureTestCase, HostPartitionTestCase, HostPlanTestCase, @@ -17,6 +19,7 @@ from tests.unit.src.sqlbuild.rule_engine._helpers.host.helpers import ( evaluate_order_rules, orders_project, + record_signals, use_hosts, write_order_rules, ) @@ -32,36 +35,75 @@ description="more hosts than subjects", model_count=2, hosts=8, expected_host_runs=3 ), HostPartitionTestCase( - description="module state makes the rule untracked on every host", + description="module state reruns the rule in one host", model_count=7, hosts=3, - expected_host_runs=3, + expected_host_runs=4, records_module_state=True, ), + HostPartitionTestCase( + description="cross-model duplicate detector reruns in one host", + model_count=7, + hosts=3, + expected_host_runs=4, + detects_duplicates=True, + ), + HostPartitionTestCase( + description="uncached cross-model duplicate detector reruns in one host", + model_count=7, + hosts=3, + expected_host_runs=4, + detects_duplicates=True, + cache_enabled=False, + ), ], ids=lambda case: case.description, ) def test_given_partitioned_hosts_when_evaluating_then_findings_and_cache_match_one_host( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostPartitionTestCase ) -> None: - write_order_rules(root=tmp_path) + write_order_rules( + root=tmp_path, + records_module_state=test_case.records_module_state, + detects_duplicates=test_case.detects_duplicates, + ) project: CompiledProject = orders_project(model_count=test_case.model_count) cache_path: Path = tmp_path / CUSTOM_RULES_CACHE_FILE + cache_glob: str = CUSTOM_RULES_CACHE_FILE + cache_enabled: bool = test_case.cache_enabled single_launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=1) - single: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) - single_cache: bytes = cache_path.read_bytes() - cache_path.unlink() + single: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=cache_enabled + ) + single_warm: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=cache_enabled + ) + single_cache: bytes = b"".join(path.read_bytes() for path in tmp_path.glob(cache_glob)) + cache_path.unlink(missing_ok=True) monkeypatch.undo() split_launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) - split: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) + split: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=cache_enabled + ) + split_cold_launches: int = len(split_launches) + split_cache: bytes = b"".join(path.read_bytes() for path in tmp_path.glob(cache_glob)) + split_warm: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=cache_enabled + ) - assert len(single_launches) == 1 - assert len(split_launches) == test_case.expected_host_runs assert split.findings == single.findings + assert split_warm.findings == single_warm.findings == single.findings + assert len(single_launches) == 1 + int(not cache_enabled) + assert split_cold_launches == test_case.expected_host_runs assert (split.cache_hits, split.cache_misses) == (single.cache_hits, single.cache_misses) - assert cache_path.read_bytes() == single_cache - assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").glob("*.pickle")) + assert (split_warm.cache_hits, split_warm.cache_misses) == ( + single_warm.cache_hits, + single_warm.cache_misses, + ) + assert split_cache == single_cache + assert bool(single_cache) is cache_enabled + assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) @pytest.mark.parametrize( @@ -175,5 +217,68 @@ def test_given_planned_invocations_when_partitioning_then_split_is_deterministic assert plans == test_case.expected_plans +@pytest.mark.parametrize( + "test_case", + [ + HostCancellationTestCase( + description="busy hosts stop at their next invocation", + model_count=60, + hosts=3, + timeout_millis=120_000, + delay_seconds=2.0, + expected_error_fragment="cannot inspect orders_000", + expected_max_seconds=20.0, + ), + ], + ids=lambda case: case.description, +) +def test_given_failing_host_when_others_are_busy_then_they_stop_before_later_invocations( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostCancellationTestCase +) -> None: + write_order_rules(root=tmp_path, failing=("orders_000",), delay_seconds=test_case.delay_seconds) + project: CompiledProject = orders_project(model_count=test_case.model_count) + launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + monkeypatch.setattr(custom_host_pool, "_HOST_TIMEOUT_MILLIS", test_case.timeout_millis) + started: float = time.monotonic() + + with pytest.raises(RulesError) as raised: + _ = evaluate_order_rules(project=project, root=tmp_path) + + assert time.monotonic() - started < test_case.expected_max_seconds + assert test_case.expected_error_fragment in str(raised.value) + assert len(launches) == test_case.hosts + 1 + assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) + + +@pytest.mark.parametrize( + "test_case", + [ + HostCancellationTestCase( + description="busy hosts stop at their next invocation", + model_count=60, + hosts=3, + timeout_millis=120_000, + delay_seconds=2.0, + expected_error_fragment="cannot inspect orders_000", + expected_max_seconds=20.0, + ), + ], + ids=lambda case: case.description, +) +def test_given_cancelled_hosts_when_failing_then_no_process_is_signalled( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostCancellationTestCase +) -> None: + write_order_rules(root=tmp_path, failing=("orders_000",), delay_seconds=test_case.delay_seconds) + project: CompiledProject = orders_project(model_count=test_case.model_count) + _ = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + signals: list[tuple[int, int]] = record_signals(monkeypatch=monkeypatch) + + with pytest.raises(RulesError) as raised: + _ = evaluate_order_rules(project=project, root=tmp_path) + + assert signals == [] + assert test_case.expected_error_fragment in str(raised.value) + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_host_capacity.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_host_capacity.py new file mode 100644 index 0000000000..025e8f4757 --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_host_capacity.py @@ -0,0 +1,134 @@ +"""Host capacity honours container CPU quotas from cgroup v1 and v2.""" + +import os +from pathlib import Path + +import pytest + +from sqlbuild.rule_engine._helpers.host import custom_host_pool +from sqlbuild.rule_engine._helpers.host.host_capacity import available_cores +from tests.unit.src.sqlbuild.rule_engine._helpers.host._test_types import HostCapacityTestCase +from tests.unit.src.sqlbuild.rule_engine._helpers.host.helpers import write_cgroup_files + +_NODE_CPUS: int = 64 + + +@pytest.mark.parametrize( + "test_case", + [ + HostCapacityTestCase( + description="cgroup v2 quota on the process cgroup", + proc_cgroup="0::/kubepods/orders\n", + files={"kubepods/orders/cpu.max": "400000 100000\n"}, + expected_cores=4, + ), + HostCapacityTestCase( + description="cgroup v2 fractional quota rounds up", + proc_cgroup="0::/kubepods/orders\n", + files={"kubepods/orders/cpu.max": "150000 100000\n"}, + expected_cores=2, + ), + HostCapacityTestCase( + description="cgroup v2 ancestor quota is tighter", + proc_cgroup="0::/kubepods/orders\n", + files={ + "kubepods/cpu.max": "200000 100000\n", + "kubepods/orders/cpu.max": "800000 100000\n", + }, + expected_cores=2, + ), + HostCapacityTestCase( + description="cgroup v2 without a quota", + proc_cgroup="0::/\n", + files={"cpu.max": "max 100000\n"}, + expected_cores=_NODE_CPUS, + ), + HostCapacityTestCase( + description="cgroup v1 CFS quota", + proc_cgroup="4:cpu,cpuacct:/docker/orders\n", + files={ + "cpu,cpuacct/docker/orders/cpu.cfs_quota_us": "300000\n", + "cpu,cpuacct/docker/orders/cpu.cfs_period_us": "100000\n", + }, + expected_cores=3, + ), + HostCapacityTestCase( + description="cgroup v1 unlimited CFS quota", + proc_cgroup="4:cpu,cpuacct:/\n", + files={ + "cpu,cpuacct/cpu.cfs_quota_us": "-1\n", + "cpu,cpuacct/cpu.cfs_period_us": "100000\n", + }, + expected_cores=_NODE_CPUS, + ), + HostCapacityTestCase( + description="no cgroup files", proc_cgroup="", files={}, expected_cores=_NODE_CPUS + ), + HostCapacityTestCase( + description="non-UTF-8 cgroup v2 quota is ignored", + proc_cgroup="0::/kubepods/orders\n", + files={"kubepods/orders/cpu.max": "\udcff\udcfe 100000\n"}, + expected_cores=_NODE_CPUS, + ), + HostCapacityTestCase( + description="non-UTF-8 process cgroup is ignored", + proc_cgroup="0::/kubepods/\udcff\n", + files={"kubepods/cpu.max": "200000 100000\n"}, + expected_cores=_NODE_CPUS, + ), + HostCapacityTestCase( + description="non-UTF-8 cgroup v1 quota is ignored", + proc_cgroup="4:cpu,cpuacct:/\n", + files={ + "cpu,cpuacct/cpu.cfs_quota_us": "\udcff\n", + "cpu,cpuacct/cpu.cfs_period_us": "100000\n", + }, + expected_cores=_NODE_CPUS, + ), + ], + ids=lambda case: case.description, +) +def test_given_cgroup_quota_when_counting_cores_then_quota_caps_process_cpus( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostCapacityTestCase +) -> None: + proc_cgroup: Path = write_cgroup_files( + root=tmp_path, proc_cgroup=test_case.proc_cgroup, files=test_case.files + ) + monkeypatch.setattr(os, "process_cpu_count", lambda: _NODE_CPUS, raising=False) + + cores: int = available_cores(cgroup_root=tmp_path / "cgroup", proc_cgroup=proc_cgroup) + + assert cores == test_case.expected_cores + + +@pytest.mark.parametrize( + "test_case", + [ + HostCapacityTestCase( + description="hosts are capped on a large node", + proc_cgroup="", + files={}, + expected_cores=8, + ) + ], + ids=lambda case: case.description, +) +def test_given_many_cores_when_partitioning_then_at_most_eight_hosts_start( + monkeypatch: pytest.MonkeyPatch, test_case: HostCapacityTestCase +) -> None: + monkeypatch.setattr(custom_host_pool, "available_cores", lambda: _NODE_CPUS) + invocations: tuple[tuple[str, str], ...] = tuple( + ("XSQBRT101", f"models/orders_{index:04d}.sql") for index in range(10_000) + ) + + plans: tuple[dict[str, list[str] | None], ...] = custom_host_pool.partition_host_plan( + plan={"XSQBRT101": None}, + invocations=invocations, + hosts=custom_host_pool.host_count(len(invocations)), + ) + + assert len(plans) == test_case.expected_cores + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/website/src/content/docs/docs/concepts/rules/custom-rules/testing-and-determinism.mdx b/website/src/content/docs/docs/concepts/rules/custom-rules/testing-and-determinism.mdx index aa6facc605..e0c011bbdc 100644 --- a/website/src/content/docs/docs/concepts/rules/custom-rules/testing-and-determinism.mdx +++ b/website/src/content/docs/docs/concepts/rules/custom-rules/testing-and-determinism.mdx @@ -37,6 +37,19 @@ the next. Compute findings from the subject and `ctx` alone. Repository pytest files remain outside the SQLBuild project's SQL `tests/` directory. +## Parallel evaluation + +On large projects, SQLBuild splits custom Rule evaluation across several host processes, at most +one per available CPU and never more than eight. Each process evaluates a share of the subjects, +and SQLBuild merges the findings back into the order a single process reports them. + +A Rule that changes module-level state while it runs, for example by adding each model to a set of +already-seen names, could report different findings when its subjects are split. SQLBuild detects +the state change and evaluates that Rule again in one process over all of its subjects, so its +findings and cache entries match a single-process run. The result stays correct, but the Rule no +longer runs in parallel and costs an extra host process. Keep each subject's findings independent +of the others so Rules can run in parallel. + ## Helpers Custom Rules can import repository-owned helpers under `rules/` and these deterministic standard