Spaces:
Sleeping
Sleeping
File size: 5,288 Bytes
399944f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | from __future__ import annotations
"""
Run-level timing utilities for the KBDebugger pipeline.
Goal
----
We want to measure and record how long each pipeline stage takes, without
polluting `pipeline/run.py` with timing logic.
This module provides:
1) RunTimer
- collects per-stage timings
- can dump a single JSON file for the whole run
2) timed_stage(...)
- a context manager that:
- shows a Rich spinner (nice UX for blocking stages)
- measures elapsed time using time.perf_counter()
- registers stage timing into a RunTimer
Design constraints
------------------
- Lightweight (Rich only; no external tracing frameworks)
- Stable output structure (so we can diff across runs)
- Works for both:
- blocking single-call stages (Docling, Neo4j retrieval, LLM calls)
- loops (we can still use rich.track inside loops)
"""
from dataclasses import dataclass, field
from contextlib import contextmanager
from time import perf_counter
from typing import Any, Dict, Iterator, Optional
from rich.console import Console
from rich.status import Status
from kbdebugger.utils.json import write_json
from kbdebugger.utils.time import now_utc_compact
from .time import _format_seconds_human, now_utc_iso
@dataclass
class StageTiming:
"""
Timing record for a single stage.
"""
name: str
elapsed_seconds: float
elapsed_human: str
started_at_utc: str
finished_at_utc: str
@dataclass
class RunTimer:
"""
Collect timings for an entire pipeline run.
Typical usage
-------------
>>> timer = RunTimer(run_name="kbdebugger_pipeline")
>>> with timer.stage("KG retrieval"):
... kg = retrieve_keyword_subgraph(...)
>>> timer.save_json()
"""
run_name: str = "kbdebugger_pipeline"
created_at_utc: str = field(default_factory=lambda: now_utc_iso())
stages: Dict[str, StageTiming] = field(default_factory=dict)
def record(
self,
*,
stage_name: str,
started_at_utc: str,
finished_at_utc: str,
elapsed_seconds: float,
) -> None:
"""
Record a completed stage timing.
Notes
-----
- If a stage name is reused, the last timing wins.
(This is usually what we want when rerunning a stage in one run.)
"""
self.stages[stage_name] = StageTiming(
name=stage_name,
elapsed_seconds=float(elapsed_seconds),
elapsed_human=_format_seconds_human(elapsed_seconds),
started_at_utc=started_at_utc,
finished_at_utc=finished_at_utc,
)
def as_json_dict(self) -> Dict[str, Any]:
"""
Produce a stable JSON-serializable structure for saving/logging.
"""
# Keep insertion order in Python 3.7+ dicts; but we also sort by name for stability.
stages_sorted = {k: vars(self.stages[k]) for k in sorted(self.stages.keys())}
total_seconds = sum(t.elapsed_seconds for t in self.stages.values())
return {
"run_name": self.run_name,
"created_at_utc": self.created_at_utc,
"total_elapsed_seconds": float(total_seconds),
"total_elapsed_human": _format_seconds_human(total_seconds),
"stages": stages_sorted,
}
def save_json(
self,
*,
path: Optional[str] = None,
prefix: str = "logs/00_pipeline_timing",
) -> str:
"""
Save a single JSON file for the entire run.
Parameters
----------
path:
Optional explicit file path. If not provided, we generate one.
prefix:
Default prefix used when autogenerating a filename.
Returns
-------
str
The path written.
"""
if path is None:
created_at = now_utc_compact()
path = f"{prefix}_{created_at}.json"
write_json(path, self.as_json_dict())
return path
@contextmanager
def stage(
self,
title: str,
*,
console: Optional[Console] = None,
show_spinner: bool = False,
spinner: str = "dots",
print_done: bool = True,
) -> Iterator[None]:
"""
Context manager: show a spinner + measure + record timing.
This is the "one-liner" you wrap around each pipeline stage.
Example
-------
>>> with timer.stage("🦆 Docling: PDF → paragraphs"):
... paragraphs = extract_paragraphs_from_pdf(...)
"""
c = console or Console()
started_at = now_utc_iso() # e.g., "2026-02-11T15:26:55"
t0 = perf_counter()
if show_spinner:
with Status(title, console=c, spinner=spinner):
yield
else:
# No status UI — timing only.
yield
elapsed = perf_counter() - t0
finished_at = now_utc_iso()
self.record(
stage_name=title,
started_at_utc=started_at,
finished_at_utc=finished_at,
elapsed_seconds=elapsed,
)
if print_done:
self.console.print(
f"[green]✅ {title}[/green] [dim](took {_format_seconds_human(elapsed)})[/dim]"
)
|