Fix circuit breaker, nested concurrency, and add test coverage

- Circuit breaker: Use a shared threading.Event to cancel remaining
  audit tasks when the circuit breaker fires. Previously, CircuitBreakerError
  was collected like any other error and all tasks ran to completion.

- Nested concurrency: Pass audit_concurrent_tasks=1 from the scheduler's
  flat pool to the evaluator, preventing max_workers * concurrent_tasks
  threads from hitting the DB simultaneously. Add audit_concurrent_tasks
  parameter to SnapshotEvaluator.audit() for this override.

- Add tests for circuit breaker short-circuiting, blocking audit error
  collection (NodeAuditsErrors), and nested concurrency prevention.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: airhorns

sqlmesh/core/scheduler.py

@@ -1014,14 +1014,14 @@ def _run_audits_concurrently(
 
         errors: t.List[NodeExecutionFailedError[SchedulingUnit]] = []
         errors_lock = threading.Lock()
+        cancelled = threading.Event()
 
         def run_audit_task(node: EvaluateNode) -> None:
-            # The circuit breaker is checked at task start. Tasks already submitted to the
-            # thread pool will run to completion — unlike the DAG executor's level-by-level
-            # cancellation, this is acceptable for audit-only runs because audits are
-            # read-only and have no side effects.
+            if cancelled.is_set():
+                return
             if circuit_breaker and circuit_breaker():
-                raise CircuitBreakerError()
+                cancelled.set()
+                return
 
             snapshot = self.snapshots_by_name[node.snapshot_name]
             node_start, node_end = node.interval
@@ -1035,6 +1035,7 @@ def _do_audit() -> t.List[AuditResult]:
                     start=node_start,
                     end=node_end,
                     execution_time=execution_time,
+                    audit_concurrent_tasks=1,
                 )
 
             self._run_node_with_progress(
@@ -1048,6 +1049,8 @@ def _do_audit() -> t.List[AuditResult]:
         def run_audit_task_collecting_errors(node: EvaluateNode) -> None:
             try:
                 run_audit_task(node)
+            except CircuitBreakerError:
+                cancelled.set()
             except Exception as ex:
                 error: NodeExecutionFailedError[SchedulingUnit] = NodeExecutionFailedError(node)
                 error.__cause__ = ex
@@ -1056,6 +1059,9 @@ def run_audit_task_collecting_errors(node: EvaluateNode) -> None:
 
         concurrent_apply_to_values(audit_tasks, run_audit_task_collecting_errors, self.max_workers)
 
+        if cancelled.is_set():
+            raise CircuitBreakerError()
+
         return errors, []
 
     def _check_ready_intervals(

sqlmesh/core/snapshot/evaluator.py

@@ -555,6 +555,7 @@ def audit(
         execution_time: t.Optional[TimeLike] = None,
         deployability_index: t.Optional[DeployabilityIndex] = None,
         wap_id: t.Optional[str] = None,
+        audit_concurrent_tasks: t.Optional[int] = None,
         **kwargs: t.Any,
     ) -> t.List[AuditResult]:
         """Execute a snapshot's node's audit queries.
@@ -632,10 +633,11 @@ def _run_audit(
                 **kwargs,
             )
 
+        tasks_num = audit_concurrent_tasks if audit_concurrent_tasks is not None else self.concurrent_tasks
         results = concurrent_apply_to_values(
             prepared_audits,
             _run_audit,
-            self.concurrent_tasks,
+            tasks_num,
         )
 
         if wap_id is not None:

tests/core/test_scheduler.py

@@ -1488,3 +1488,183 @@ def fake_audit(snapshot: Snapshot, **kwargs: t.Any) -> t.List[AuditResult]:
     assert all(tid != main_thread_id for tid in audit_call_thread_ids), (
         "Both audits should run on worker threads regardless of DAG dependencies"
     )
+
+
+@pytest.mark.fast
+def test_audit_only_circuit_breaker_stops_remaining_tasks(mocker: MockerFixture, make_snapshot):
+    """When the circuit breaker fires, remaining audit tasks are skipped and CircuitBreakerError is raised."""
+    audit_calls: t.List[str] = []
+    audit_lock = threading.Lock()
+
+    snapshot_a = make_snapshot(SqlModel(name="a", query=parse_one("SELECT 1 as id")))
+    snapshot_b = make_snapshot(SqlModel(name="b", query=parse_one("SELECT 2 as id")))
+    snapshot_c = make_snapshot(SqlModel(name="c", query=parse_one("SELECT 3 as id")))
+    snapshot_a.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_b.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_c.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    def fake_audit(snapshot: Snapshot, **kwargs: t.Any) -> t.List[AuditResult]:
+        with audit_lock:
+            audit_calls.append(snapshot.name)
+        return []
+
+    mock_evaluator = mocker.MagicMock()
+    mock_evaluator.audit.side_effect = fake_audit
+    mock_evaluator.get_snapshots_to_create.return_value = []
+    mock_evaluator.concurrent_context.return_value.__enter__ = mocker.Mock(return_value=None)
+    mock_evaluator.concurrent_context.return_value.__exit__ = mocker.Mock(return_value=False)
+
+    # Circuit breaker fires immediately on the first check
+    scheduler = Scheduler(
+        snapshots=[snapshot_a, snapshot_b, snapshot_c],
+        snapshot_evaluator=mock_evaluator,
+        state_sync=mocker.MagicMock(),
+        default_catalog=None,
+        max_workers=1,  # Sequential so we can reason about ordering
+    )
+
+    interval = (to_timestamp("2023-01-01"), to_timestamp("2023-01-02"))
+    merged_intervals: SnapshotToIntervals = {
+        snapshot_a: [interval],
+        snapshot_b: [interval],
+        snapshot_c: [interval],
+    }
+
+    with pytest.raises(CircuitBreakerError):
+        scheduler.run_merged_intervals(
+            merged_intervals=merged_intervals,
+            deployability_index=DeployabilityIndex.all_deployable(),
+            environment_naming_info=EnvironmentNamingInfo(),
+            audit_only=True,
+            circuit_breaker=lambda: True,
+        )
+
+    # With circuit breaker always-true, no audits should run
+    assert len(audit_calls) == 0
+
+
+@pytest.mark.fast
+def test_audit_only_blocking_audit_error_collected(mocker: MockerFixture, make_snapshot):
+    """When a blocking audit fails (raises NodeAuditsErrors), the error is collected and other audits still run."""
+    audit_calls: t.List[str] = []
+    audit_lock = threading.Lock()
+
+    snapshot_a = make_snapshot(SqlModel(name="a", query=parse_one("SELECT 1 as id")))
+    snapshot_b = make_snapshot(SqlModel(name="b", query=parse_one("SELECT 2 as id")))
+    snapshot_a.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_b.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    def fake_audit(snapshot: Snapshot, **kwargs: t.Any) -> t.List[AuditResult]:
+        with audit_lock:
+            audit_calls.append(snapshot.name)
+        if snapshot.name == '"a"':
+            from sqlmesh.utils.errors import AuditError
+            from sqlglot import exp
+
+            audit_error = AuditError(
+                audit_name="not_null",
+                audit_args={},
+                model=snapshot.model_or_none,
+                count=5,
+                query=exp.select("1"),
+                adapter_dialect="duckdb",
+            )
+            raise NodeAuditsErrors([audit_error])
+        return []
+
+    mock_evaluator = mocker.MagicMock()
+    mock_evaluator.audit.side_effect = fake_audit
+    mock_evaluator.get_snapshots_to_create.return_value = []
+    mock_evaluator.concurrent_context.return_value.__enter__ = mocker.Mock(return_value=None)
+    mock_evaluator.concurrent_context.return_value.__exit__ = mocker.Mock(return_value=False)
+
+    mock_console = mocker.MagicMock()
+
+    scheduler = Scheduler(
+        snapshots=[snapshot_a, snapshot_b],
+        snapshot_evaluator=mock_evaluator,
+        state_sync=mocker.MagicMock(),
+        default_catalog=None,
+        max_workers=2,
+        console=mock_console,
+    )
+
+    interval = (to_timestamp("2023-01-01"), to_timestamp("2023-01-02"))
+    merged_intervals: SnapshotToIntervals = {
+        snapshot_a: [interval],
+        snapshot_b: [interval],
+    }
+
+    errors, skipped = scheduler.run_merged_intervals(
+        merged_intervals=merged_intervals,
+        deployability_index=DeployabilityIndex.all_deployable(),
+        environment_naming_info=EnvironmentNamingInfo(),
+        audit_only=True,
+    )
+
+    # The NodeAuditsErrors should be collected as an error, not re-raised
+    assert len(errors) == 1
+    assert isinstance(errors[0].__cause__, NodeAuditsErrors)
+    assert skipped == []
+    # Both audits should have been attempted despite one failing
+    assert len(audit_calls) == 2
+    assert '"a"' in audit_calls
+    assert '"b"' in audit_calls
+
+
+@pytest.mark.fast
+def test_audit_only_no_nested_concurrency(mocker: MockerFixture, make_snapshot):
+    """With scheduler max_workers > 1, each evaluator audit call uses sequential execution (audit_concurrent_tasks=1).
+
+    This prevents nested thread pool multiplication: max_workers * concurrent_tasks threads hitting
+    the DB at the same time.
+    """
+    import sqlmesh.core.snapshot.evaluator as evaluator_module
+
+    spy = mocker.spy(evaluator_module, "concurrent_apply_to_values")
+
+    snapshot_a = make_snapshot(SqlModel(name="a", query=parse_one("SELECT 1 as id")))
+    snapshot_b = make_snapshot(SqlModel(name="b", query=parse_one("SELECT 2 as id")))
+    snapshot_a.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_b.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    mock_evaluator = mocker.MagicMock()
+    mock_evaluator.audit.return_value = []
+    mock_evaluator.get_snapshots_to_create.return_value = []
+    mock_evaluator.concurrent_context.return_value.__enter__ = mocker.Mock(return_value=None)
+    mock_evaluator.concurrent_context.return_value.__exit__ = mocker.Mock(return_value=False)
+
+    # Use the real SnapshotEvaluator to test the audit_concurrent_tasks parameter flows through
+    real_evaluator = SnapshotEvaluator(adapters=mocker.MagicMock(), concurrent_tasks=4)
+    real_evaluator.audit = mocker.MagicMock(return_value=[])  # type: ignore
+
+    scheduler = Scheduler(
+        snapshots=[snapshot_a, snapshot_b],
+        snapshot_evaluator=real_evaluator,
+        state_sync=mocker.MagicMock(),
+        default_catalog=None,
+        max_workers=2,
+    )
+
+    interval = (to_timestamp("2023-01-01"), to_timestamp("2023-01-02"))
+    merged_intervals: SnapshotToIntervals = {
+        snapshot_a: [interval],
+        snapshot_b: [interval],
+    }
+
+    errors, skipped = scheduler.run_merged_intervals(
+        merged_intervals=merged_intervals,
+        deployability_index=DeployabilityIndex.all_deployable(),
+        environment_naming_info=EnvironmentNamingInfo(),
+        audit_only=True,
+    )
+
+    assert errors == []
+    assert skipped == []
+    assert real_evaluator.audit.call_count == 2
+
+    # Verify that audit_concurrent_tasks=1 was passed to each audit call to prevent nested pools
+    for call in real_evaluator.audit.call_args_list:
+        assert call.kwargs.get("audit_concurrent_tasks") == 1, (
+            "audit_concurrent_tasks=1 must be passed to prevent nested thread pool multiplication"
+        )