chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Trace execution information in dump comparator (#19682)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-02 18:46:27 +08:00
committed by GitHub
parent 3ebd85bf1c
commit 5bf3deb4bc
17 changed files with 1180 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/debug_utils/comparator/__init__.py
View File

@@ -1,3 +1,6 @@
from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import ( # noqa: F401
TracedAlignerPlan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( # noqa: F401
AlignerPlan,
)

104
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py
View File

@@ -1,13 +1,19 @@
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Optional
from typing import NamedTuple, Optional
import torch
from sglang.srt.debug_utils.comparator.aligner.axis_aligner import (
execute_axis_aligner_plan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
TracedAlignerPlan,
TracedSidePlan,
TracedStepPlan,
TracedSubPlan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
AlignerPerStepPlan,
AlignerPerStepSubPlan,
@@ -28,15 +34,31 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import (
execute_unsharder_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
from sglang.srt.debug_utils.comparator.output_types import (
ReplicatedCheckResult,
ShapeSnapshot,
)
from sglang.srt.debug_utils.comparator.utils import Pair
class StepPlansResult(NamedTuple):
tensors: dict[int, torch.Tensor]
checks: list[ReplicatedCheckResult]
traced_side: TracedSidePlan
class SubPlansResult(NamedTuple):
tensor: Optional[torch.Tensor]
checks: list[ReplicatedCheckResult]
snapshots: list[ShapeSnapshot]
@dataclass(frozen=True)
class AlignerResult:
tensors: Optional[Pair[torch.Tensor]]
failed_side_xy: Optional[str] # "x" or "y"; None if success
replicated_checks: list[ReplicatedCheckResult] = field(default_factory=list)
traced_plan: Optional[TracedAlignerPlan] = None
def execute_aligner_plan(
@@ -48,26 +70,34 @@ def execute_aligner_plan(
all_checks: list[ReplicatedCheckResult] = []
# Per-side: unshard + reorder -> dict[step, tensor]
step_tensors_x, checks_x = _execute_step_plans(
result_x: StepPlansResult = _execute_step_plans(
tensors=tensors_pair.x, step_plans=plan.per_step_plans.x
)
all_checks.extend(checks_x)
all_checks.extend(result_x.checks)
step_tensors_y, checks_y = _execute_step_plans(
result_y: StepPlansResult = _execute_step_plans(
tensors=tensors_pair.y, step_plans=plan.per_step_plans.y
)
all_checks.extend(checks_y)
all_checks.extend(result_y.checks)
if not step_tensors_x or not step_tensors_y:
failed_side_xy: str = "x" if not step_tensors_x else "y"
traced_plan: TracedAlignerPlan = TracedAlignerPlan(
plan=plan,
per_side=Pair(x=result_x.traced_side, y=result_y.traced_side),
)
if not result_x.tensors or not result_y.tensors:
failed_side_xy: str = "x" if not result_x.tensors else "y"
return AlignerResult(
tensors=None,
failed_side_xy=failed_side_xy,
replicated_checks=all_checks,
traced_plan=traced_plan,
)
# Cross-side: token alignment (or direct extraction for single-step)
step_pair: Pair[dict[int, torch.Tensor]] = Pair(x=step_tensors_x, y=step_tensors_y)
step_pair: Pair[dict[int, torch.Tensor]] = Pair(
x=result_x.tensors, y=result_y.tensors
)
combined: Pair[torch.Tensor]
if plan.token_aligner_mode == "concat_steps":
combined = execute_token_aligner_concat_steps(tensor_of_step_pair=step_pair)
@@ -78,10 +108,10 @@ def execute_aligner_plan(
tensor_of_step_pair=step_pair,
)
else:
assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1
assert len(result_x.tensors) == 1 and len(result_y.tensors) == 1
combined = Pair(
x=list(step_tensors_x.values())[0],
y=list(step_tensors_y.values())[0],
x=list(result_x.tensors.values())[0],
y=list(result_y.tensors.values())[0],
)
# Cross-side: axis alignment (squeeze singletons + rearrange dim order)
@@ -95,50 +125,78 @@ def execute_aligner_plan(
tensors=combined,
failed_side_xy=None,
replicated_checks=all_checks,
traced_plan=traced_plan,
)
def _execute_step_plans(
tensors: list[torch.Tensor],
step_plans: list[AlignerPerStepPlan],
) -> tuple[dict[int, torch.Tensor], list[ReplicatedCheckResult]]:
) -> StepPlansResult:
result: dict[int, torch.Tensor] = {}
all_checks: list[ReplicatedCheckResult] = []
traced_steps: list[TracedStepPlan] = []
for step_plan in step_plans:
step_tensors: list[torch.Tensor] = [
tensors[i] for i in step_plan.input_object_indices
]
tensor, checks = execute_sub_plans(
sub_result: SubPlansResult = execute_sub_plans(
tensors=step_tensors, plans=step_plan.sub_plans
)
all_checks.extend(checks)
if tensor is not None:
result[step_plan.step] = tensor
all_checks.extend(sub_result.checks)
return result, all_checks
traced_subs: list[TracedSubPlan] = [
TracedSubPlan(plan=sub_plan, snapshot=snapshot)
for sub_plan, snapshot in zip(step_plan.sub_plans, sub_result.snapshots)
]
traced_steps.append(
TracedStepPlan(
step=step_plan.step,
input_object_indices=step_plan.input_object_indices,
sub_plans=traced_subs,
)
)
if sub_result.tensor is not None:
result[step_plan.step] = sub_result.tensor
return StepPlansResult(
tensors=result,
checks=all_checks,
traced_side=TracedSidePlan(step_plans=traced_steps),
)
def execute_sub_plans(
tensors: list[torch.Tensor],
plans: list[AlignerPerStepSubPlan],
) -> tuple[Optional[torch.Tensor], list[ReplicatedCheckResult]]:
) -> SubPlansResult:
if not tensors:
return None, []
return SubPlansResult(tensor=None, checks=[], snapshots=[])
if not plans:
if len(tensors) != 1:
return None, []
return tensors[0], []
return SubPlansResult(tensor=None, checks=[], snapshots=[])
return SubPlansResult(tensor=tensors[0], checks=[], snapshots=[])
current: list[torch.Tensor] = tensors
all_checks: list[ReplicatedCheckResult] = []
all_snapshots: list[ShapeSnapshot] = []
for plan in plans:
input_shapes: list[list[int]] = [list(t.shape) for t in current]
current, checks = execute_sub_plan(tensors=current, plan=plan)
output_shapes: list[list[int]] = [list(t.shape) for t in current]
all_checks.extend(checks)
all_snapshots.append(
ShapeSnapshot(
input_shapes=input_shapes,
output_shapes=output_shapes,
)
)
assert len(current) == 1
return current[0], all_checks
return SubPlansResult(tensor=current[0], checks=all_checks, snapshots=all_snapshots)
def execute_sub_plan(

37
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/traced_types.py Normal file
View File

@@ -0,0 +1,37 @@
"""Traced wrapper types that embed execution traces (ShapeSnapshots) into plan nodes.
These types are created *after* execution, pairing each sub-plan with its
observed shape snapshot so that downstream formatters never need to manually
zip plan + trace by index.
"""
from __future__ import annotations
from typing import Optional
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
AlignerPerStepSubPlan,
AlignerPlan,
)
from sglang.srt.debug_utils.comparator.output_types import ShapeSnapshot
from sglang.srt.debug_utils.comparator.utils import Pair, _StrictBase
class TracedSubPlan(_StrictBase):
plan: AlignerPerStepSubPlan
snapshot: Optional[ShapeSnapshot] = None
class TracedStepPlan(_StrictBase):
step: int
input_object_indices: list[int]
sub_plans: list[TracedSubPlan]
class TracedSidePlan(_StrictBase):
step_plans: list[TracedStepPlan]
class TracedAlignerPlan(_StrictBase):
plan: AlignerPlan
per_side: Pair[TracedSidePlan]

8
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_loader.py
View File

@@ -240,9 +240,11 @@ def _load_and_align_aux_tensor(
dim_names: list[str] = resolve_dim_names(dims_str)
tensors = [apply_dim_names(t, dim_names) for t in tensors]
result, _replicated_checks = execute_sub_plans(tensors=tensors, plans=sub_plans)
assert result is not None
return result.rename(None) # strip named dims before returning to plugin
sub_result = execute_sub_plans(tensors=tensors, plans=sub_plans)
assert sub_result.tensor is not None
return sub_result.tensor.rename(
None
) # strip named dims before returning to plugin
log_sink.add(
InfoLog(

42
python/sglang/srt/debug_utils/comparator/bundle_comparator.py
View File

@@ -29,6 +29,8 @@ from sglang.srt.debug_utils.comparator.dp_utils import filter_to_non_empty_dp_ra
from sglang.srt.debug_utils.comparator.log_sink import log_sink
from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider
from sglang.srt.debug_utils.comparator.output_types import (
BundleFileInfo,
BundleSideInfo,
ErrorLog,
NonTensorComparisonRecord,
SkipComparisonRecord,
@@ -44,6 +46,37 @@ from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta
_FAILED_SIDE_MAP: dict[str, str] = {"x": "baseline", "y": "target"}
def _collect_bundle_side_info(
items: list[ValueWithMeta],
metas: list[dict[str, Any]],
) -> BundleSideInfo:
from sglang.srt.debug_utils.comparator.display import (
PARALLEL_INFO_KEYS,
extract_parallel_info,
)
files: list[BundleFileInfo] = []
for item, meta in zip(items, metas):
assert isinstance(item.value, torch.Tensor)
tensor: torch.Tensor = item.value
parallel_info: dict[str, str] = {}
for key in PARALLEL_INFO_KEYS:
extract_parallel_info(row_data=parallel_info, info=meta.get(key, {}))
files.append(
BundleFileInfo(
shape=list(tensor.shape),
dtype=str(tensor.dtype),
rank=meta.get("rank"),
parallel_info=parallel_info if parallel_info else None,
)
)
dims: Optional[str] = metas[0].get("dims") if metas else None
return BundleSideInfo(num_files=len(files), files=files, dims=dims)
def compare_bundle_pair(
*,
name: str,
@@ -186,6 +219,12 @@ def _compare_bundle_pair_tensor_type(
thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
)
# Collect raw bundle info before alignment
raw_bundle_info: Pair[BundleSideInfo] = Pair(
x=_collect_bundle_side_info(items=valid_pair.x, metas=metas_pair.x),
y=_collect_bundle_side_info(items=valid_pair.y, metas=metas_pair.y),
)
# Apply dim names to tensors, then execute
tensors_pair: Pair[list[torch.Tensor]] = Pair(
x=_apply_dim_names_from_meta(
@@ -226,8 +265,9 @@ def _compare_bundle_pair_tensor_type(
)
record = TensorComparisonRecord(
**info.model_dump(),
aligner_plan=plan,
traced_plan=aligner_result.traced_plan,
replicated_checks=replicated_checks,
raw_bundle_info=raw_bundle_info,
)
if viz_output_dir is not None:

10
python/sglang/srt/debug_utils/comparator/display.py
View File

@@ -10,11 +10,11 @@ import polars as pl
from sglang.srt.debug_utils.comparator.output_types import (
InputIdsRecord,
RankInfoRecord,
report_sink,
)
from sglang.srt.debug_utils.comparator.report_sink import report_sink
from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta
_PARALLEL_INFO_KEYS: list[str] = ["sglang_parallel_info", "megatron_parallel_info"]
PARALLEL_INFO_KEYS: list[str] = ["sglang_parallel_info", "megatron_parallel_info"]
def emit_display_records(
@@ -68,8 +68,8 @@ def _collect_rank_info(
meta: dict[str, Any] = ValueWithMeta.load(dump_dir / row["filename"]).meta
row_data: dict[str, Any] = {"rank": row["rank"]}
for key in _PARALLEL_INFO_KEYS:
_extract_parallel_info(row_data=row_data, info=meta.get(key, {}))
for key in PARALLEL_INFO_KEYS:
extract_parallel_info(row_data=row_data, info=meta.get(key, {}))
table_rows.append(row_data)
return table_rows or None
@@ -119,7 +119,7 @@ def _collect_input_ids_and_positions(
return table_rows or None
def _extract_parallel_info(row_data: dict[str, Any], info: dict[str, Any]) -> None:
def extract_parallel_info(row_data: dict[str, Any], info: dict[str, Any]) -> None:
if not info or info.get("error"):
return

26
python/sglang/srt/debug_utils/comparator/entrypoint.py
View File

@@ -31,12 +31,12 @@ from sglang.srt.debug_utils.comparator.output_types import (
SkipComparisonRecord,
SummaryRecord,
TensorComparisonRecord,
report_sink,
)
from sglang.srt.debug_utils.comparator.per_token_visualizer import (
generate_per_token_heatmap,
)
from sglang.srt.debug_utils.comparator.preset import PRESETS, expand_preset
from sglang.srt.debug_utils.comparator.report_sink import report_sink
from sglang.srt.debug_utils.comparator.utils import (
Pair,
auto_descend_dir,
@@ -53,7 +53,11 @@ def main() -> None:
def run(args: argparse.Namespace) -> int:
report_sink.configure(output_format=args.output_format, report_path=None)
report_sink.configure(
output_format=args.output_format,
report_path=None,
verbosity=args.verbosity,
)
dir_pair: Pair[Path] = Pair(
x=auto_descend_dir(Path(args.baseline_path), label="baseline_path"),
@@ -69,6 +73,16 @@ def run(args: argparse.Namespace) -> int:
Path(args.override_config) if args.override_config else None
)
report_path: Optional[Path] = _resolve_report_path(
target_path=dir_pair.y,
report_path_arg=args.report_path,
)
report_sink.configure(
output_format=args.output_format,
report_path=report_path,
verbosity=args.verbosity,
)
report_path: Optional[Path] = _resolve_report_path(
target_path=dir_pair.y,
report_path_arg=args.report_path,
@@ -294,6 +308,14 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
default="text",
help="Output format: text (default) or json (JSONL, one JSON object per line)",
)
parser.add_argument(
"--verbosity",
type=str,
choices=["minimal", "normal", "verbose"],
default="normal",
help="Output verbosity: minimal (1 line per tensor), normal (compact lifecycle), "
"verbose (full detail). Default: normal",
)
parser.add_argument(
"--preset",
type=str,

2
python/sglang/srt/debug_utils/comparator/log_sink.py
View File

@@ -27,8 +27,8 @@ class LogSink:
from sglang.srt.debug_utils.comparator.output_types import (
LogRecord,
_split_logs,
report_sink,
)
from sglang.srt.debug_utils.comparator.report_sink import report_sink
errors, infos = _split_logs([log])
report_sink.add(LogRecord(errors=errors, infos=infos))

158
python/sglang/srt/debug_utils/comparator/output_types.py
View File

@@ -1,12 +1,12 @@
from __future__ import annotations
import sys
from abc import abstractmethod
from pathlib import Path
from typing import IO, TYPE_CHECKING, Annotated, Any, Literal, Optional, Union
from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional, Union
import polars as pl
from pydantic import ConfigDict, Discriminator, Field, TypeAdapter, model_validator
from rich.console import RenderableType
from rich.markup import escape
from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
format_comparison,
@@ -16,12 +16,15 @@ from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorComparisonInfo,
)
from sglang.srt.debug_utils.comparator.utils import _StrictBase
from sglang.srt.debug_utils.comparator.utils import Pair, _StrictBase
if TYPE_CHECKING:
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
AlignerPlan,
from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
TracedAlignerPlan,
TracedSubPlan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
from sglang.srt.debug_utils.comparator.report_sink import Verbosity
class BaseLog(_StrictBase):
@@ -59,6 +62,26 @@ class ReplicatedCheckResult(_StrictBase):
diff: Optional[DiffInfo] = None
class BundleFileInfo(_StrictBase):
"""Per-file info within a bundle (one rank's raw tensor)."""
shape: list[int]
dtype: str
rank: Optional[int] = None
parallel_info: Optional[dict[str, str]] = None # e.g. {"tp": "0/4", "ep": "1/2"}
class BundleSideInfo(_StrictBase):
num_files: int
files: list[BundleFileInfo]
dims: Optional[str] = None # e.g. "b s h(tp) d"
class ShapeSnapshot(_StrictBase):
input_shapes: list[list[int]]
output_shapes: list[list[int]]
class _OutputRecord(_StrictBase):
errors: list[ErrorLog] = Field(default_factory=list)
infos: list[InfoLog] = Field(default_factory=list)
@@ -66,6 +89,12 @@ class _OutputRecord(_StrictBase):
@abstractmethod
def _format_body(self) -> str: ...
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return self._format_body()
def to_rich(self, verbosity: Verbosity = "normal") -> RenderableType:
return self._format_body()
def to_text(self) -> str:
body = self._format_body()
if self.errors:
@@ -87,6 +116,11 @@ class _BaseComparisonRecord(_OutputRecord):
return f"[step={self.location.step}] "
return ""
def _format_location_prefix_rich(self) -> str:
if self.location.step is not None:
return escape(f"[step={self.location.step}]") + " "
return ""
def _format_location_suffix(self) -> str:
if self.location.step is not None:
return f" (step={self.location.step})"
@@ -149,8 +183,9 @@ class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord):
model_config = ConfigDict(extra="forbid", defer_build=True)
type: Literal["comparison"] = "comparison"
aligner_plan: Optional[AlignerPlan] = None
traced_plan: Optional[TracedAlignerPlan] = None
replicated_checks: list[ReplicatedCheckResult] = Field(default_factory=list)
raw_bundle_info: Optional[Pair[BundleSideInfo]] = None
@property
def category(self) -> str:
@@ -164,8 +199,8 @@ class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord):
body: str = self._format_location_prefix() + format_comparison(self)
if self.replicated_checks:
body += "\n" + format_replicated_checks(self.replicated_checks)
if self.aligner_plan is not None:
body += "\n" + _format_aligner_plan(self.aligner_plan)
if self.traced_plan is not None:
body += "\n" + _format_aligner_plan(self.traced_plan)
return body
@@ -225,26 +260,48 @@ class LogRecord(_OutputRecord):
return ""
def _format_aligner_plan(plan: AlignerPlan) -> str:
def _format_aligner_plan(traced_plan: TracedAlignerPlan) -> str:
lines: list[str] = ["Aligner Plan:"]
for side_label, side_plans in [
("baseline", plan.per_step_plans.x),
("target", plan.per_step_plans.y),
for side_label, traced_side in [
("baseline", traced_plan.per_side.x),
("target", traced_plan.per_side.y),
]:
if not side_plans:
if not traced_side.step_plans:
lines.append(f" {side_label}: (no steps)")
continue
step_summaries: list[str] = []
for step_plan in side_plans:
sub_strs: list[str] = []
for sub in step_plan.sub_plans:
sub_strs.append(f"{sub.type}")
for traced_step in traced_side.step_plans:
sub_strs: list[str] = [
_format_sub_plan_text(traced_sub)
for traced_sub in traced_step.sub_plans
]
summary: str = ", ".join(sub_strs) if sub_strs else "passthrough"
step_summaries.append(f"step={step_plan.step}: {summary}")
step_summaries.append(f"step={traced_step.step}: {summary}")
lines.append(f" {side_label}: [{'; '.join(step_summaries)}]")
lines.extend(_format_cross_side_plan_text(traced_plan.plan))
return "\n".join(lines)
def _format_sub_plan_text(traced_sub: TracedSubPlan) -> str:
sub_desc: str = f"{traced_sub.plan.type}"
if traced_sub.snapshot is not None:
snap = traced_sub.snapshot
in_count: int = len(snap.input_shapes)
out_count: int = len(snap.output_shapes)
in_shape: str = str(snap.input_shapes[0]) if snap.input_shapes else "?"
out_shape: str = str(snap.output_shapes[0]) if snap.output_shapes else "?"
sub_desc += f" {in_count}x{in_shape} -> {out_count}x{out_shape}"
return sub_desc
def _format_cross_side_plan_text(plan: AlignerPlan) -> list[str]:
lines: list[str] = []
if plan.token_aligner_plan is not None:
num_tokens: int = len(plan.token_aligner_plan.locators.x.steps)
lines.append(f" token_aligner: {num_tokens} tokens aligned")
@@ -257,7 +314,7 @@ def _format_aligner_plan(plan: AlignerPlan) -> str:
parts.append(f"y: {plan.axis_aligner_plan.pattern.y}")
lines.append(f" axis_aligner: {', '.join(parts)}")
return "\n".join(lines)
return lines
AnyRecord = Annotated[
@@ -281,64 +338,3 @@ def _get_any_record_adapter() -> TypeAdapter:
def parse_record_json(json_str: str | bytes) -> AnyRecord:
return _get_any_record_adapter().validate_json(json_str)
def _print_to_stdout(record: _OutputRecord, *, output_format: str) -> None:
if output_format == "json":
print(record.model_dump_json())
else:
print(record.to_text())
class ReportSink:
"""Unified entry point for all record output."""
def __init__(self) -> None:
self._output_format: str = "text"
self._report_file: Optional[IO[str]] = None
self._report_path: Optional[Path] = None
def configure(
self,
*,
output_format: str = "text",
report_path: Optional[Path] = None,
) -> None:
self._output_format = output_format
if report_path is not None:
try:
report_path.parent.mkdir(parents=True, exist_ok=True)
self._report_file = open(report_path, "w", encoding="utf-8")
self._report_path = report_path
except OSError as exc:
print(
f"Warning: cannot open report file {report_path}: {exc}",
file=sys.stderr,
)
def add(self, record: _OutputRecord) -> None:
_print_to_stdout(record, output_format=self._output_format)
if self._report_file is not None:
self._report_file.write(record.model_dump_json())
self._report_file.write("\n")
self._report_file.flush()
def close(self) -> None:
if self._report_file is not None:
self._report_file.close()
self._report_file = None
@property
def report_path(self) -> Optional[Path]:
return self._report_path
def _reset(self) -> None:
"""Reset state for test isolation."""
self.close()
self._output_format = "text"
self._report_path = None
report_sink = ReportSink()

87
python/sglang/srt/debug_utils/comparator/report_sink.py Normal file
View File

@@ -0,0 +1,87 @@
from __future__ import annotations
import sys
from pathlib import Path
from typing import IO, Literal, Optional
from rich.console import Console
from sglang.srt.debug_utils.comparator.output_types import _OutputRecord
Verbosity = Literal["minimal", "normal", "verbose"]
class ReportSink:
"""Unified entry point for all record output."""
def __init__(self) -> None:
self._output_format: str = "text"
self._verbosity: Verbosity = "normal"
self._report_file: Optional[IO[str]] = None
self._report_path: Optional[Path] = None
self._console: Optional[Console] = None
@property
def verbosity(self) -> Verbosity:
return self._verbosity
def configure(
self,
*,
output_format: str = "text",
report_path: Optional[Path] = None,
verbosity: Verbosity = "normal",
) -> None:
self._output_format = output_format
self._verbosity = verbosity
if report_path is not None:
try:
report_path.parent.mkdir(parents=True, exist_ok=True)
self._report_file = open(report_path, "w", encoding="utf-8")
self._report_path = report_path
except OSError as exc:
print(
f"Warning: cannot open report file {report_path}: {exc}",
file=sys.stderr,
)
def add(self, record: _OutputRecord) -> None:
self._print_to_stdout(record)
if self._report_file is not None:
self._report_file.write(record.model_dump_json())
self._report_file.write("\n")
self._report_file.flush()
def close(self) -> None:
if self._report_file is not None:
self._report_file.close()
self._report_file = None
@property
def report_path(self) -> Optional[Path]:
return self._report_path
def _reset(self) -> None:
self.close()
self._output_format = "text"
self._verbosity = "normal"
self._report_path = None
self._console = None
def _get_console(self) -> Console:
if self._console is None:
self._console = Console()
return self._console
def _print_to_stdout(self, record: _OutputRecord) -> None:
if self._output_format == "json":
print(record.model_dump_json())
else:
console: Console = self._get_console()
console.print(record.to_rich(verbosity=self._verbosity))
console.print() # blank line between records
report_sink = ReportSink()