Do not precompile model evaluation info

Author: treysp

sqlmesh/core/console.py

@@ -95,7 +95,7 @@ def stop_plan_evaluation(self) -> None:
     @abc.abstractmethod
     def start_evaluation_progress(
         self,
-        batches: t.Dict[Snapshot, Intervals],
+        batch_sizes: t.Dict[Snapshot, int],
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
     ) -> None:
@@ -109,6 +109,7 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
     def update_snapshot_evaluation_progress(
         self,
         snapshot: Snapshot,
+        interval: Interval,
         batch_idx: int,
         duration_ms: t.Optional[int],
         num_audits_passed: int,
@@ -350,7 +351,7 @@ def stop_plan_evaluation(self) -> None:
 
     def start_evaluation_progress(
         self,
-        batches: t.Dict[Snapshot, Intervals],
+        batch_sizes: t.Dict[Snapshot, int],
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
     ) -> None:
@@ -362,6 +363,7 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
     def update_snapshot_evaluation_progress(
         self,
         snapshot: Snapshot,
+        interval: Interval,
         batch_idx: int,
         duration_ms: t.Optional[int],
         num_audits_passed: int,
@@ -531,6 +533,13 @@ class TerminalConsole(Console):
 
     TABLE_DIFF_SOURCE_BLUE = "#0248ff"
 
+    EVAL_PROGRESS_BAR_COLUMN_WIDTHS: t.Dict[str, int] = {
+        "batch": 9,
+        "name": 50,
+        "annotation": 50,
+        "duration": 8,
+    }
+
     def __init__(
         self,
         console: t.Optional[RichConsole] = None,
@@ -546,9 +555,6 @@ def __init__(
         self.evaluation_total_task: t.Optional[TaskID] = None
         self.evaluation_model_progress: t.Optional[Progress] = None
         self.evaluation_model_tasks: t.Dict[str, TaskID] = {}
-        self.evaluation_model_batch_sizes: t.Dict[Snapshot, int] = {}
-        self.evaluation_model_info: t.Dict[Snapshot, t.Dict[str, t.Any]] = {}
-        self.evaluation_model_column_widths: t.Dict[str, int] = {}
 
         # Put in temporary values that are replaced when evaluating
         self.environment_naming_info = EnvironmentNamingInfo()
@@ -589,28 +595,12 @@ def stop_plan_evaluation(self) -> None:
 
     def start_evaluation_progress(
         self,
-        batched_intervals: t.Dict[Snapshot, Intervals],
+        batch_sizes: t.Dict[Snapshot, int],
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
     ) -> None:
         """Indicates that a new snapshot evaluation progress has begun."""
         if not self.evaluation_progress_live:
-            self.evaluation_model_batch_sizes = {
-                snapshot: len(intervals) for snapshot, intervals in batched_intervals.items()
-            }
-            self.environment_naming_info = environment_naming_info
-            self.default_catalog = default_catalog
-
-            self.evaluation_model_info, self.evaluation_model_column_widths = (
-                _create_evaluation_model_info(
-                    batched_intervals,
-                    self.evaluation_model_batch_sizes,
-                    environment_naming_info,
-                    default_catalog,
-                    self.dialect,
-                )
-            )
-
             self.evaluation_total_progress = make_progress_bar(
                 "Evaluating model batches", self.console
             )
@@ -629,9 +619,13 @@ def start_evaluation_progress(
             self.evaluation_progress_live.start()
 
             self.evaluation_total_task = self.evaluation_total_progress.add_task(
-                "Evaluating models...", total=sum(self.evaluation_model_batch_sizes.values())
+                "Evaluating models...", total=sum(batch_sizes.values())
             )
 
+            self.evaluation_model_batch_sizes = batch_sizes
+            self.environment_naming_info = environment_naming_info
+            self.default_catalog = default_catalog
+
     def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
         if self.evaluation_model_progress and snapshot.name not in self.evaluation_model_tasks:
             display_name = snapshot.display_name(
@@ -648,6 +642,7 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
     def update_snapshot_evaluation_progress(
         self,
         snapshot: Snapshot,
+        interval: Interval,
         batch_idx: int,
         duration_ms: t.Optional[int],
         num_audits_passed: int,
@@ -661,26 +656,32 @@ def update_snapshot_evaluation_progress(
         ):
             total_batches = self.evaluation_model_batch_sizes[snapshot]
             batch_num = str(batch_idx + 1).rjust(len(str(total_batches)))
-            batch = f"[{batch_num}/{total_batches}] "
+            batch = f"[{batch_num}/{total_batches}]".ljust(
+                self.EVAL_PROGRESS_BAR_COLUMN_WIDTHS["batch"]
+            )
 
             if duration_ms:
-                display_name = self.evaluation_model_info[snapshot]["display_name"].ljust(
-                    self.evaluation_model_column_widths["display_name"]
+                display_name = snapshot.display_name(
+                    self.environment_naming_info,
+                    self.default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None,
+                    dialect=self.dialect,
+                ).ljust(self.EVAL_PROGRESS_BAR_COLUMN_WIDTHS["name"])
+
+                annotation = _create_evaluation_model_annotation(
+                    snapshot, _format_evaluation_model_interval(snapshot, interval)
                 )
 
-                annotation = self.evaluation_model_info[snapshot]["annotation"][batch_idx]
                 if num_audits_passed:
                     annotation += f", {num_audits_passed} audits pass"
                 if num_audits_failed:
                     annotation += f", {num_audits_failed} audits fail {RED_X_MARK}"
                 annotation = (annotation + "]").ljust(
-                    self.evaluation_model_column_widths["annotation"]
+                    self.EVAL_PROGRESS_BAR_COLUMN_WIDTHS["annotation"]
                 )
 
-                # 8 characters for duration
-                # if the failed audit red X is present, the console adds an extra space
-                duration_width = 7 if num_audits_failed else 8
-                duration = f"{(duration_ms / 1000.0):.2f}s".rjust(duration_width)
+                duration = f"{(duration_ms / 1000.0):.2f}s".rjust(
+                    self.EVAL_PROGRESS_BAR_COLUMN_WIDTHS["duration"]
+                )
 
                 self.evaluation_progress_live.console.print(
                     f"{GREEN_CHECK_MARK} {batch}{display_name}{annotation} {duration}"
@@ -708,8 +709,6 @@ def stop_evaluation_progress(self, success: bool = True) -> None:
         self.evaluation_model_progress = None
         self.evaluation_model_tasks = {}
         self.evaluation_model_batch_sizes = {}
-        self.evaluation_model_info = {}
-        self.evaluation_model_column_widths = {}
         self.environment_naming_info = EnvironmentNamingInfo()
         self.default_catalog = None
 
@@ -2303,13 +2302,11 @@ def _confirm(self, message: str, **kwargs: t.Any) -> bool:
 
     def start_evaluation_progress(
         self,
-        batched_intervals: t.Dict[Snapshot, Intervals],
+        batch_sizes: t.Dict[Snapshot, int],
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
     ) -> None:
-        self.evaluation_model_batch_sizes = {
-            snapshot: len(intervals) for snapshot, intervals in batched_intervals.items()
-        }
+        self.evaluation_model_batch_sizes = batch_sizes
         self.evaluation_environment_naming_info = environment_naming_info
         self.default_catalog = default_catalog
 
@@ -2328,6 +2325,7 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
     def update_snapshot_evaluation_progress(
         self,
         snapshot: Snapshot,
+        interval: Interval,
         batch_idx: int,
         duration_ms: t.Optional[int],
         num_audits_passed: int,
@@ -2472,18 +2470,19 @@ def stop_plan_evaluation(self) -> None:
 
     def start_evaluation_progress(
         self,
-        batched_intervals: t.Dict[Snapshot, Intervals],
+        batch_sizes: t.Dict[Snapshot, int],
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
     ) -> None:
-        self._write(f"Starting evaluation for {len(batched_intervals)} snapshots")
+        self._write(f"Starting evaluation for {sum(batch_sizes.values())} snapshots")
 
     def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
         self._write(f"Evaluating {snapshot.name}")
 
     def update_snapshot_evaluation_progress(
         self,
         snapshot: Snapshot,
+        interval: Interval,
         batch_idx: int,
         duration_ms: t.Optional[int],
         num_audits_passed: int,

sqlmesh/core/scheduler.py

@@ -436,7 +436,7 @@ def run_merged_intervals(
         batched_intervals = self.batch_intervals(merged_intervals)
 
         self.console.start_evaluation_progress(
-            batched_intervals,
+            {snapshot: len(intervals) for snapshot, intervals in batched_intervals.items()},
             environment_naming_info,
             self.default_catalog,
         )
@@ -495,6 +495,7 @@ def evaluate_node(node: SchedulingUnit) -> None:
                 num_audits_failed = sum(1 for result in audit_results if result.count)
                 self.console.update_snapshot_evaluation_progress(
                     snapshot,
+                    batched_intervals[snapshot][batch_idx],
                     batch_idx,
                     evaluation_duration_ms,
                     num_audits - num_audits_failed,

tests/cli/test_cli.py

@@ -220,7 +220,7 @@ def test_plan_restate_model(runner, tmp_path):
     assert result.exit_code == 0
     assert_duckdb_test(result)
     assert "No changes to plan: project files match the `prod` environment" in result.output
-    assert "sqlmesh_example.full_model [full refresh" in result.output
+    assert "sqlmesh_example.full_model                         [full refresh" in result.output
     assert_model_batches_evaluated(result)
     assert_env_views_updated(result)
 
@@ -530,7 +530,7 @@ def test_plan_nonbreaking(runner, tmp_path):
     assert "+  'a' AS new_col" in result.output
     assert "Directly Modified: sqlmesh_example.incremental_model (Non-breaking)" in result.output
     assert "sqlmesh_example.full_model (Indirect Non-breaking)" in result.output
-    assert "sqlmesh_example.incremental_model [insert" in result.output
+    assert "sqlmesh_example.incremental_model                  [insert" in result.output
     assert "sqlmesh_example.full_model evaluated [full refresh" not in result.output
     assert_backfill_success(result)
 
@@ -588,8 +588,8 @@ def test_plan_breaking(runner, tmp_path):
     assert result.exit_code == 0
     assert "+  item_id + 1 AS item_id," in result.output
     assert "Directly Modified: sqlmesh_example.full_model (Breaking)" in result.output
-    assert "sqlmesh_example.full_model [full refresh" in result.output
-    assert "sqlmesh_example.incremental_model [insert" not in result.output
+    assert "sqlmesh_example.full_model                         [full refresh" in result.output
+    assert "sqlmesh_example.incremental_model                   [insert" not in result.output
     assert_backfill_success(result)
 
 
@@ -627,8 +627,8 @@ def test_plan_dev_select(runner, tmp_path):
     assert "+  item_id + 1 AS item_id," not in result.output
     assert "Directly Modified: sqlmesh_example__dev.full_model (Breaking)" not in result.output
     # only incremental_model backfilled
-    assert "sqlmesh_example__dev.incremental_model [insert" in result.output
-    assert "sqlmesh_example__dev.full_model [full refresh" not in result.output
+    assert "sqlmesh_example__dev.incremental_model             [insert" in result.output
+    assert "sqlmesh_example__dev.full_model                   [full refresh" not in result.output
     assert_backfill_success(result)
 
 
@@ -666,8 +666,8 @@ def test_plan_dev_backfill(runner, tmp_path):
         "Directly Modified: sqlmesh_example__dev.incremental_model (Non-breaking)" in result.output
     )
     # only incremental_model backfilled
-    assert "sqlmesh_example__dev.incremental_model [insert" in result.output
-    assert "sqlmesh_example__dev.full_model [full refresh" not in result.output
+    assert "sqlmesh_example__dev.incremental_model             [insert" in result.output
+    assert "sqlmesh_example__dev.full_model                    [full refresh" not in result.output
     assert_backfill_success(result)
 
 

web/server/console.py

@@ -7,12 +7,11 @@
 
 from fastapi.encoders import jsonable_encoder
 from sse_starlette.sse import ServerSentEvent
-
+from sqlmesh.core.snapshot.definition import Interval
 from sqlmesh.core.console import TerminalConsole
 from sqlmesh.core.environment import EnvironmentNamingInfo
 from sqlmesh.core.plan.definition import EvaluatablePlan
 from sqlmesh.core.snapshot import Snapshot, SnapshotInfoLike
-from sqlmesh.core.snapshot.definition import Intervals
 from sqlmesh.core.test import ModelTest
 from sqlmesh.utils.date import now_timestamp
 from web.server import models
@@ -92,7 +91,7 @@ def stop_restate_progress(self, success: bool) -> None:
 
     def start_evaluation_progress(
         self,
-        batched_intervals: t.Dict[Snapshot, Intervals],
+        batched_intervals: t.Dict[Snapshot, int],
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
     ) -> None:
@@ -126,6 +125,7 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
     def update_snapshot_evaluation_progress(
         self,
         snapshot: Snapshot,
+        interval: Interval,
         batch_idx: int,
         duration_ms: t.Optional[int],
         audits_passed: int,