Chore: Report duration for each finished snapshot batch during evaluation (#1571)

Author: izeigerman

sqlmesh/core/console.py

@@ -71,7 +71,9 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
         """Starts the snapshot evaluation progress."""
 
     @abc.abstractmethod
-    def update_snapshot_evaluation_progress(self, snapshot: Snapshot, num_batches: int) -> None:
+    def update_snapshot_evaluation_progress(
+        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
+    ) -> None:
         """Updates the snapshot evaluation progress."""
 
     @abc.abstractmethod
@@ -263,19 +265,29 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
                 total=self.evaluation_model_batches[snapshot],
             )
 
-    def update_snapshot_evaluation_progress(self, snapshot: Snapshot, num_batches: int) -> None:
+    def update_snapshot_evaluation_progress(
+        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
+    ) -> None:
         """Update the snapshot evaluation progress."""
-        if self.evaluation_total_progress and self.evaluation_model_progress:
+        if (
+            self.evaluation_total_progress
+            and self.evaluation_model_progress
+            and self.evaluation_progress_live
+        ):
+            total_batches = self.evaluation_model_batches[snapshot]
+
+            if duration_ms:
+                self.evaluation_progress_live.console.print(
+                    f"[{batch_idx + 1}/{total_batches}] {snapshot.name} [green]finished[/green] in {(duration_ms / 1000.0):.2f}s"
+                )
+
             self.evaluation_total_progress.update(
-                self.evaluation_total_task or TaskID(0), refresh=True, advance=num_batches
+                self.evaluation_total_task or TaskID(0), refresh=True, advance=1
             )
 
             model_task_id = self.evaluation_model_tasks[snapshot.name]
-            self.evaluation_model_progress.update(model_task_id, refresh=True, advance=num_batches)
-            if (
-                self.evaluation_model_progress._tasks[model_task_id].completed
-                >= self.evaluation_model_batches[snapshot]
-            ):
+            self.evaluation_model_progress.update(model_task_id, refresh=True, advance=1)
+            if self.evaluation_model_progress._tasks[model_task_id].completed >= total_batches:
                 self.evaluation_model_progress.remove_task(model_task_id)
 
     def stop_evaluation_progress(self, success: bool = True) -> None:
@@ -1318,11 +1330,13 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
             self.evaluation_batch_progress[snapshot.name] = (view_name, 0)
             print(f"Starting '{view_name}', Total batches: {self.evaluation_batches[snapshot]}")
 
-    def update_snapshot_evaluation_progress(self, snapshot: Snapshot, num_batches: int) -> None:
+    def update_snapshot_evaluation_progress(
+        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
+    ) -> None:
         view_name, loaded_batches = self.evaluation_batch_progress[snapshot.name]
         total_batches = self.evaluation_batches[snapshot]
 
-        loaded_batches += num_batches
+        loaded_batches += 1
         self.evaluation_batch_progress[snapshot.name] = (view_name, loaded_batches)
 
         finished_loading = loaded_batches == total_batches

sqlmesh/core/plan/evaluator.py

@@ -263,6 +263,8 @@ def evaluate(self, plan: Plan) -> None:
             if not plan_application_succeeded:
                 raise SQLMeshError("Plan application failed.")
 
+            self.console.log_success("The plan has been applied successfully")
+
     @property
     def client(self) -> BaseAirflowClient:
         raise NotImplementedError

sqlmesh/core/scheduler.py

@@ -23,14 +23,20 @@
 from sqlmesh.utils import format_exception
 from sqlmesh.utils.concurrency import concurrent_apply_to_dag
 from sqlmesh.utils.dag import DAG
-from sqlmesh.utils.date import TimeLike, now, to_datetime, validate_date_range
+from sqlmesh.utils.date import (
+    TimeLike,
+    now,
+    now_timestamp,
+    to_datetime,
+    validate_date_range,
+)
 from sqlmesh.utils.errors import AuditError, SQLMeshError
 
 logger = logging.getLogger(__name__)
 Interval = t.Tuple[datetime, datetime]
 Batch = t.List[Interval]
 SnapshotToBatches = t.Dict[Snapshot, Batch]
-SchedulingUnit = t.Tuple[Snapshot, Interval]
+SchedulingUnit = t.Tuple[Snapshot, t.Tuple[Interval, int]]
 
 
 class Scheduler:
@@ -246,12 +252,20 @@ def run(
 
         def evaluate_node(node: SchedulingUnit) -> None:
             assert execution_time
-            snapshot, (start, end) = node
+
+            snapshot, ((start, end), batch_idx) = node
             self.console.start_snapshot_evaluation_progress(snapshot)
+
+            execution_start_ts = now_timestamp()
+            evaluation_duration_ms: t.Optional[int] = None
+
             try:
                 self.evaluate(snapshot, start, end, execution_time, is_dev=is_dev)
+                evaluation_duration_ms = now_timestamp() - execution_start_ts
             finally:
-                self.console.update_snapshot_evaluation_progress(snapshot, 1)
+                self.console.update_snapshot_evaluation_progress(
+                    snapshot, batch_idx, evaluation_duration_ms
+                )
 
         try:
             with self.snapshot_evaluator.concurrent_context():
@@ -297,23 +311,25 @@ def _dag(self, batches: SnapshotToBatches) -> DAG[SchedulingUnit]:
             if not intervals:
                 continue
             upstream_dependencies = [
-                (self.snapshots[p_sid], interval)
+                (self.snapshots[p_sid], (interval, i))
                 for p_sid in snapshot.parents
                 if p_sid in self.snapshots
-                for interval in intervals_per_snapshot_version.get(
-                    (
-                        self.snapshots[p_sid].name,
-                        self.snapshots[p_sid].version_get_or_generate(),
-                    ),
-                    [],
+                for i, interval in enumerate(
+                    intervals_per_snapshot_version.get(
+                        (
+                            self.snapshots[p_sid].name,
+                            self.snapshots[p_sid].version_get_or_generate(),
+                        ),
+                        [],
+                    )
                 )
             ]
             for i, interval in enumerate(intervals):
-                dag.add((snapshot, interval), upstream_dependencies)
+                dag.add((snapshot, (interval, i)), upstream_dependencies)
                 if snapshot.depends_on_past:
                     dag.add(
-                        (snapshot, interval),
-                        [(snapshot, _interval) for _interval in intervals[:i]],
+                        (snapshot, (interval, i)),
+                        [(snapshot, (_interval, _i)) for _i, _interval in enumerate(intervals[:i])],
                     )
         return dag
 

tests/core/test_scheduler.py

@@ -124,7 +124,10 @@ def test_incremental_by_unique_key_kind_dag(mocker: MockerFixture, make_snapshot
     dag = scheduler._dag(batches)
     assert dag.graph == {
         # Depends on no one
-        (unique_by_key_snapshot, (to_datetime("2023-01-01"), to_datetime("2023-01-07"))): set(),
+        (
+            unique_by_key_snapshot,
+            ((to_datetime("2023-01-01"), to_datetime("2023-01-07")), 0),
+        ): set(),
     }
     mock_state_sync.refresh_snapshot_intervals.assert_called_once()
 
@@ -161,17 +164,35 @@ def test_incremental_time_self_reference_dag(mocker: MockerFixture, make_snapsho
 
     assert dag.graph == {
         # Only run one day at a time and each day relies on the previous days
-        (incremental_self_snapshot, (to_datetime("2023-01-01"), to_datetime("2023-01-02"))): set(),
-        (incremental_self_snapshot, (to_datetime("2023-01-03"), to_datetime("2023-01-04"))): {
-            (incremental_self_snapshot, (to_datetime("2023-01-01"), to_datetime("2023-01-02")))
+        (
+            incremental_self_snapshot,
+            ((to_datetime("2023-01-01"), to_datetime("2023-01-02")), 0),
+        ): set(),
+        (incremental_self_snapshot, ((to_datetime("2023-01-03"), to_datetime("2023-01-04")), 1)): {
+            (incremental_self_snapshot, ((to_datetime("2023-01-01"), to_datetime("2023-01-02")), 0))
         },
-        (incremental_self_snapshot, (to_datetime("2023-01-04"), to_datetime("2023-01-05"))): {
-            (incremental_self_snapshot, (to_datetime("2023-01-01"), to_datetime("2023-01-02"))),
-            (incremental_self_snapshot, (to_datetime("2023-01-03"), to_datetime("2023-01-04"))),
+        (incremental_self_snapshot, ((to_datetime("2023-01-04"), to_datetime("2023-01-05")), 2)): {
+            (
+                incremental_self_snapshot,
+                ((to_datetime("2023-01-01"), to_datetime("2023-01-02")), 0),
+            ),
+            (
+                incremental_self_snapshot,
+                ((to_datetime("2023-01-03"), to_datetime("2023-01-04")), 1),
+            ),
         },
-        (incremental_self_snapshot, (to_datetime("2023-01-06"), to_datetime("2023-01-07"))): {
-            (incremental_self_snapshot, (to_datetime("2023-01-01"), to_datetime("2023-01-02"))),
-            (incremental_self_snapshot, (to_datetime("2023-01-03"), to_datetime("2023-01-04"))),
-            (incremental_self_snapshot, (to_datetime("2023-01-04"), to_datetime("2023-01-05"))),
+        (incremental_self_snapshot, ((to_datetime("2023-01-06"), to_datetime("2023-01-07")), 3)): {
+            (
+                incremental_self_snapshot,
+                ((to_datetime("2023-01-01"), to_datetime("2023-01-02")), 0),
+            ),
+            (
+                incremental_self_snapshot,
+                ((to_datetime("2023-01-03"), to_datetime("2023-01-04")), 1),
+            ),
+            (
+                incremental_self_snapshot,
+                ((to_datetime("2023-01-04"), to_datetime("2023-01-05")), 2),
+            ),
         },
     }

web/server/console.py

@@ -55,10 +55,12 @@ def start_evaluation_progress(
     def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
         pass
 
-    def update_snapshot_evaluation_progress(self, snapshot: Snapshot, num_batches: int) -> None:
+    def update_snapshot_evaluation_progress(
+        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
+    ) -> None:
         """Update snapshot evaluation progress."""
         if self.current_task_status:
-            self.current_task_status[snapshot.name]["completed"] += num_batches
+            self.current_task_status[snapshot.name]["completed"] += 1
             if (
                 self.current_task_status[snapshot.name]["completed"]
                 >= self.current_task_status[snapshot.name]["total"]