Skip to content

feat(console): show running model batches #6102

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Vamsi-klu wants to merge 1 commit into
base: main
Choose a base branch
from
+593 −52
Open

feat(console): show running model batches #6102

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
131 changes: 93 additions & 38 deletions sqlmesh/core/console.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -431,6 +431,15 @@ def start_snapshot_evaluation_progress(
) -> None:
"""Starts the snapshot evaluation progress."""

def start_snapshot_evaluation_batch(
self, snapshot: Snapshot, interval: Interval, batch_idx: int
) -> None:
"""Starts progress for a snapshot evaluation batch.

Consoles without per-batch output can inherit this default no-op implementation.
"""
pass

@abc.abstractmethod
def update_snapshot_evaluation_progress(
self,
Expand Down Expand Up @@ -1083,6 +1092,63 @@ def start_evaluation_progress(
self.environment_naming_info = environment_naming_info
self.default_catalog = default_catalog

def _print_evaluation_row(
self,
snapshot: Snapshot,
interval: Interval,
batch_idx: int,
status: str,
audit_only: bool = False,
num_audits_passed: int = 0,
num_audits_failed: int = 0,
execution_stats: t.Optional[QueryExecutionStats] = None,
) -> None:
"""Print one evaluation row (Running or completion) to the live console.

Both start and completion rows share this formatter so their columns align.
"""
if not (
self.evaluation_total_progress
and self.evaluation_model_progress
and self.evaluation_progress_live
):
return

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}]".ljust(self.evaluation_column_widths["batch"])

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.evaluation_column_widths["name"])

annotation = _create_evaluation_model_annotation(
snapshot,
_format_evaluation_model_interval(snapshot, interval),
execution_stats,
)
audits_str = ""
if num_audits_passed:
audits_str += f" {self.AUDIT_PASS_MARK}{num_audits_passed}"
if num_audits_failed:
audits_str += f" {self.AUDIT_FAIL_MARK}{num_audits_failed}"
audits_str = f", audits{audits_str}" if audits_str else ""
annotation_len = self.evaluation_column_widths["annotation"]
# don't adjust the annotation_len if we're using AUDIT_PADDING
annotation = f"\\[{annotation + audits_str}]".ljust(
annotation_len - 1 if num_audits_failed and self.AUDIT_PADDING == 0 else annotation_len
)

status_col = status.ljust(self.evaluation_column_widths["duration"])

msg = (
f"{f'{batch} ' if not audit_only else ''}{display_name} {annotation} {status_col}"
).replace(self.AUDIT_PASS_MARK, self.GREEN_AUDIT_PASS_MARK)

self.evaluation_progress_live.console.print(msg, soft_wrap=True)

def start_snapshot_evaluation_progress(
self, snapshot: Snapshot, audit_only: bool = False
) -> None:
Expand All @@ -1098,6 +1164,12 @@ def start_snapshot_evaluation_progress(
total=self.evaluation_model_batch_sizes[snapshot],
)

def start_snapshot_evaluation_batch(
self, snapshot: Snapshot, interval: Interval, batch_idx: int
) -> None:
"""Print a durable running row so the active batch appears in captured output."""
self._print_evaluation_row(snapshot, interval, batch_idx, "Running")

def update_snapshot_evaluation_progress(
self,
snapshot: Snapshot,
Expand All @@ -1117,53 +1189,29 @@ def update_snapshot_evaluation_progress(
and self.evaluation_progress_live
):
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}]".ljust(self.evaluation_column_widths["batch"])

if duration_ms:
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.evaluation_column_widths["name"])

annotation = _create_evaluation_model_annotation(
snapshot, _format_evaluation_model_interval(snapshot, interval), execution_stats
)
audits_str = ""
if num_audits_passed:
audits_str += f" {self.AUDIT_PASS_MARK}{num_audits_passed}"
if num_audits_failed:
audits_str += f" {self.AUDIT_FAIL_MARK}{num_audits_failed}"
audits_str = f", audits{audits_str}" if audits_str else ""
annotation_len = self.evaluation_column_widths["annotation"]
# don't adjust the annotation_len if we're using AUDIT_PADDING
annotation = f"\\[{annotation + audits_str}]".ljust(
annotation_len - 1
if num_audits_failed and self.AUDIT_PADDING == 0
else annotation_len
)

duration = f"{(duration_ms / 1000.0):.2f}s".ljust(
self.evaluation_column_widths["duration"]
)

msg = f"{f'{batch} ' if not audit_only else ''}{display_name} {annotation} {duration}".replace(
self.AUDIT_PASS_MARK, self.GREEN_AUDIT_PASS_MARK
if duration_ms is not None:
duration = f"{(duration_ms / 1000.0):.2f}s"
self._print_evaluation_row(
snapshot,
interval,
batch_idx,
duration,
audit_only=audit_only,
num_audits_passed=num_audits_passed,
num_audits_failed=num_audits_failed,
execution_stats=execution_stats,
)

self.evaluation_progress_live.console.print(msg)
elif not audit_only:
self._print_evaluation_row(snapshot, interval, batch_idx, "Failed")

self.evaluation_total_progress.update(
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=1)
if (
self.evaluation_model_progress._tasks[model_task_id].completed >= total_batches
or audit_only
):
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:
Expand Down Expand Up @@ -3925,6 +3973,13 @@ def start_snapshot_evaluation_progress(
) -> None:
self._write(f"{'Evaluating' if not audit_only else 'Auditing'} {snapshot.name}")

def start_snapshot_evaluation_batch(
self, snapshot: Snapshot, interval: Interval, batch_idx: int
) -> None:
# The debugger console already logs the model start above and intentionally has no
# TerminalConsole progress state because it doesn't call TerminalConsole.__init__.
pass

def update_snapshot_evaluation_progress(
self,
snapshot: Snapshot,
Expand Down
48 changes: 35 additions & 13 deletions sqlmesh/core/scheduler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -526,7 +526,14 @@ def run_node(node: SchedulingUnit) -> None:
snapshot = self.snapshots_by_name[node.snapshot_name]

if isinstance(node, EvaluateNode):
self.console.start_snapshot_evaluation_progress(snapshot)
if audit_only:
self.console.start_snapshot_evaluation_progress(snapshot, audit_only=True)
else:
self.console.start_snapshot_evaluation_progress(snapshot)
if not audit_only:
self.console.start_snapshot_evaluation_batch(
snapshot, node.interval, node.batch_index
)
execution_start_ts = now_timestamp()
evaluation_duration_ms: t.Optional[int] = None
start, end = node.interval
Expand Down Expand Up @@ -574,18 +581,33 @@ def run_node(node: SchedulingUnit) -> None:
SnapshotIdBatch(snapshot_id=snapshot.snapshot_id, batch_id=node.batch_index)
)

self.console.update_snapshot_evaluation_progress(
snapshot,
batched_intervals[snapshot][node.batch_index],
node.batch_index,
evaluation_duration_ms,
num_audits - num_audits_failed,
num_audits_failed,
execution_stats=execution_stats,
auto_restatement_triggers=auto_restatement_triggers.get(
snapshot.snapshot_id
),
)
if audit_only:
self.console.update_snapshot_evaluation_progress(
snapshot,
batched_intervals[snapshot][node.batch_index],
node.batch_index,
evaluation_duration_ms,
num_audits - num_audits_failed,
num_audits_failed,
audit_only=True,
execution_stats=execution_stats,
auto_restatement_triggers=auto_restatement_triggers.get(
snapshot.snapshot_id
),
)
else:
self.console.update_snapshot_evaluation_progress(
snapshot,
batched_intervals[snapshot][node.batch_index],
node.batch_index,
evaluation_duration_ms,
num_audits - num_audits_failed,
num_audits_failed,
execution_stats=execution_stats,
auto_restatement_triggers=auto_restatement_triggers.get(
snapshot.snapshot_id
),
)
elif isinstance(node, CreateNode):
self.snapshot_evaluator.create_snapshot(
snapshot=snapshot,
Expand Down
4 changes: 4 additions & 0 deletions tests/cli/test_cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -942,6 +942,10 @@ def test_run_dev(runner, tmp_path, flag):
result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "run", "dev"])
assert result.exit_code == 0
assert_model_batches_executed(result)
running_lines = [
line for line in result.output.splitlines() if line.rstrip().endswith("Running")
]
assert any("[1/1]" in line and "[insert " in line for line in running_lines)


@time_machine.travel(FREEZE_TIME)
Expand Down
Loading