Coverage for oc_meta / lib / timer.py: 23%
158 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-07-25 10:39 +0000
1#!/usr/bin/env python
3# SPDX-FileCopyrightText: 2025 Arcangelo Massari <arcangelo.massari@unibo.it>
4#
5# SPDX-License-Identifier: ISC
7# -*- coding: utf-8 -*-
8"""
9Timing and metrics collection utilities for OpenCitations Meta processing.
11This module provides reusable timing infrastructure for both production
12processing and benchmarking, with optional activation to avoid overhead.
13"""
15import threading
16import time
17from typing import Any, Callable, Dict, List, Optional
19import psutil
22class _MemorySampler:
23 """Background thread that samples RSS to capture true peak memory."""
25 def __init__(self, process: psutil.Process, interval: float = 0.1):
26 self._process = process
27 self._interval = interval
28 self._stop = threading.Event()
29 self._peak: int = 0
30 self._thread = threading.Thread(target=self._run, daemon=True)
32 def start(self):
33 self._thread.start()
35 def stop(self) -> int:
36 self._stop.set()
37 self._thread.join()
38 return self._peak
40 def _run(self):
41 while not self._stop.is_set():
42 rss = self._process.memory_info().rss
43 if rss > self._peak:
44 self._peak = rss
45 self._stop.wait(self._interval)
48class BenchmarkTimer:
49 """Context manager for timing code blocks and collecting memory metrics."""
51 def __init__(
52 self,
53 name: str,
54 verbose: bool = False,
55 on_exit: Optional[Callable[[], None]] = None,
56 ):
57 self.name = name
58 self.verbose = verbose
59 self.start_time: Optional[float] = None
60 self.end_time: Optional[float] = None
61 self.duration: Optional[float] = None
62 self.start_memory: Optional[int] = None
63 self.end_memory: Optional[int] = None
64 self.peak_memory: Optional[int] = None
65 self._sampler: Optional[_MemorySampler] = None
66 self._on_exit: Optional[Callable[[], None]] = on_exit
68 def __enter__(self):
69 self.start_time = time.time()
70 process = psutil.Process()
71 self.start_memory = process.memory_info().rss
72 self._sampler = _MemorySampler(process)
73 self._sampler.start()
74 if self.verbose:
75 print(f" [{self.name}] Starting...")
76 return self
78 def __exit__(self, exc_type, exc_val, exc_tb):
79 assert (
80 self.start_time is not None
81 and self.start_memory is not None
82 and self._sampler is not None
83 )
84 self.end_time = time.time()
85 self.duration = self.end_time - self.start_time
86 process = psutil.Process()
87 end_memory = process.memory_info().rss
88 self.end_memory = end_memory
89 sampled_peak = self._sampler.stop()
90 self.peak_memory = max(self.start_memory, end_memory, sampled_peak)
91 if self.verbose:
92 print(f" [{self.name}] Completed in {self.duration:.2f}s")
93 if self._on_exit:
94 self._on_exit()
96 def to_dict(self) -> Dict[str, Any]:
97 """Convert timing data to dictionary."""
98 return {
99 "name": self.name,
100 "duration_seconds": round(self.duration, 3) if self.duration else None,
101 "start_memory_mb": round(self.start_memory / 1024 / 1024, 2)
102 if self.start_memory
103 else None,
104 "end_memory_mb": round(self.end_memory / 1024 / 1024, 2)
105 if self.end_memory
106 else None,
107 "peak_memory_mb": round(self.peak_memory / 1024 / 1024, 2)
108 if self.peak_memory
109 else None,
110 }
113class DummyTimer:
114 """No-op timer for when timing is disabled."""
116 def __enter__(self):
117 return self
119 def __exit__(self, *args):
120 pass
123class ProcessTimer:
124 """Optional timing wrapper for MetaProcess operations."""
126 def __init__(
127 self,
128 enabled: bool = False,
129 verbose: bool = False,
130 on_phase_complete: Optional[Callable[["ProcessTimer"], None]] = None,
131 ):
132 self.enabled = enabled
133 self.verbose = verbose
134 self.timers: List[BenchmarkTimer] = []
135 self.metrics: Dict[str, Any] = {}
136 self._on_phase_complete: Optional[Callable[["ProcessTimer"], None]] = (
137 on_phase_complete
138 )
140 def timer(self, name: str):
141 """Create a timer context manager (or no-op if disabled)."""
142 if self.enabled:
143 # Don't show verbose for total_processing and sub-timers
144 show_verbose = self.verbose and name not in [
145 "total_processing",
146 "creator_execution",
147 "provenance_generation",
148 ]
149 timer = BenchmarkTimer(
150 name, verbose=show_verbose, on_exit=self._notify_phase
151 )
152 self.timers.append(timer)
153 return timer
154 else:
155 return DummyTimer()
157 def _notify_phase(self):
158 if self._on_phase_complete:
159 self._on_phase_complete(self)
161 def record_metric(self, key: str, value: Any):
162 """Record a metric."""
163 if self.enabled:
164 self.metrics[key] = value
166 def record_phase(self, name: str, duration: float):
167 """Record a phase with a specific duration (e.g., 0 for unused phases)."""
168 if self.enabled:
169 timer = BenchmarkTimer(name, verbose=False)
170 timer.start_time = 0
171 timer.end_time = duration
172 timer.duration = duration
173 timer.start_memory = 0
174 timer.end_memory = 0
175 timer.peak_memory = 0
176 self.timers.append(timer)
178 def get_report(self) -> Dict[str, Any]:
179 """Generate timing report."""
180 if not self.enabled:
181 return {}
183 total_time = (
184 next(
185 (t.duration for t in self.timers if t.name == "total_processing"), None
186 )
187 or 0.0
188 )
189 input_records = self.metrics.get("input_records", 0)
191 return {
192 "metrics": {
193 **self.metrics,
194 "total_duration_seconds": round(total_time, 3),
195 "throughput_records_per_sec": round(input_records / total_time, 2)
196 if total_time > 0
197 else 0,
198 },
199 "phases": [t.to_dict() for t in self.timers],
200 }
202 def print_summary(self):
203 """Print timing summary to console."""
204 if not self.enabled:
205 return
207 report = self.get_report()
208 metrics = report["metrics"]
210 print(f"\n{'=' * 60}")
211 print("Timing Summary")
212 print(f"{'=' * 60}")
213 print(f"Total Duration: {metrics.get('total_duration_seconds', 0)}s")
214 print(f"Throughput: {metrics.get('throughput_records_per_sec', 0)} records/sec")
215 print(f"Input Records: {metrics.get('input_records', 0)}")
216 print(f"Curated Records: {metrics.get('curated_records', 0)}")
217 print(f"New Entities: {metrics.get('new_entities', 0)}")
218 print(f"Modified Entities: {metrics.get('modified_entities', 0)}")
219 print("\nPhase Breakdown:")
220 for phase in report["phases"]:
221 if phase["name"] not in [
222 "total_processing",
223 "creator_execution",
224 "provenance_generation",
225 ]:
226 print(f" {phase['name']}: {phase['duration_seconds']}s")
227 print(f"{'=' * 60}\n")
229 def print_phase_breakdown(self):
230 """Print detailed phase breakdown for a single file."""
231 if not self.enabled:
232 return
234 report = self.get_report()
235 phases = report["phases"]
237 print("\n Phase Breakdown:")
238 for phase in phases:
239 if phase["name"] == "total_processing":
240 continue
241 name = phase["name"]
242 duration = phase["duration_seconds"]
243 peak = phase["peak_memory_mb"]
244 if peak:
245 delta = phase["end_memory_mb"] - phase["start_memory_mb"]
246 sign = "+" if delta >= 0 else ""
247 print(
248 f" {name:30s} {duration:10.2f}s {peak:10.1f} MB peak {sign}{delta:.1f} MB"
249 )
250 else:
251 print(f" {name:30s} {duration:10.2f}s")
253 def print_file_summary(self, filename: str):
254 """Print complete summary for a single file with metrics and phases."""
255 if not self.enabled:
256 return
258 report = self.get_report()
259 metrics = report["metrics"]
261 total_time = metrics.get("total_duration_seconds", 0)
262 records = metrics.get("input_records", 0)
263 entities = metrics.get("new_entities", 0)
264 throughput = metrics.get("throughput_records_per_sec", 0)
266 total_phase = next(
267 (p for p in report["phases"] if p["name"] == "total_processing"), None
268 )
269 peak_mb = total_phase["peak_memory_mb"] if total_phase else 0
270 delta_mb = (
271 (total_phase["end_memory_mb"] - total_phase["start_memory_mb"])
272 if total_phase
273 else 0
274 )
276 print(f" ✓ Completed in {total_time:.2f}s")
277 self.print_phase_breakdown()
278 print("\n Metrics:")
279 print(f" Records processed: {records}")
280 print(f" New entities: {entities}")
281 print(f" Throughput: {throughput:.2f} rec/s")
282 if peak_mb:
283 sign = "+" if delta_mb >= 0 else ""
284 print(f" Peak memory (RSS): {peak_mb:.1f} MB")
285 print(f" Memory growth: {sign}{delta_mb:.1f} MB")