"""Post-mortem performance / debug log for vibe-qc.
The v0.5.1 :mod:`vibeqc.progress` module shows progress *during* a
run; this module shows where the time went *afterwards*. The two
pair: live for "is the SCF stuck?", perf for "why is my LiH /
pob-TZVP run taking 20 minutes — is it J, K, XC quadrature, or
Bloch sums?"
Three ways to enable, all of which feed the same accumulator under
the hood:
* **Environment variable** — common case for one-off jobs::
VIBEQC_PERFLOG=output.perf python my-calc.py
When set, every ``run_job`` call (and every entry point that
threads ``perf_log=``) routes through the tracker.
* **Programmatic context manager** — wrap a region::
import vibeqc as vq
with vq.perf_log("output.perf"):
result = vq.run_rhf(mol, basis, opts)
hess = vq.compute_hessian_rhf_analytic(...)
All work inside the block accumulates into the same tracker;
the report is written when the block exits.
* **``run_job`` kwarg** — one-shot::
vq.run_job(mol, basis="cc-pVDZ", method="rks",
functional="pbe", perf_log="output.perf")
The output file is plain text with section headers, sortable by
wall-time. v0.5.2 instruments the Python-driven periodic SCF
pipeline (the EWALD_3D paths) and the high-level :func:`run_job`
phases. C++ kernel-level scopes (``compute_eri``, ``build_coulomb``,
``xc_eval``) are gated behind a compile-time ``#ifdef
VIBEQC_PERFLOG`` and arrive in a v0.5.2.x patch — release builds pay
zero runtime cost when the feature is off.
Public surface
--------------
.. autoclass:: PerfTracker
.. autoclass:: PerfScope
.. autofunction:: perf_log
.. autofunction:: format_perf_report
.. autofunction:: active_tracker
"""
from __future__ import annotations
import contextvars
import os
import sys
import time
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterator, Optional, Union
__all__ = [
"PerfScope",
"PerfTracker",
"active_tracker",
"format_perf_report",
"perf_log",
]
# ---------------------------------------------------------------------------
# Memory probe — best-effort, never raises.
# ---------------------------------------------------------------------------
def _rss_mb() -> float:
"""Current resident set size in MiB. ``0.0`` when unavailable."""
try:
import resource # POSIX-only
ru = resource.getrusage(resource.RUSAGE_SELF)
# Linux reports ru_maxrss in KiB, macOS in bytes. Detect by magnitude.
if sys.platform == "darwin":
return ru.ru_maxrss / (1024.0 * 1024.0)
return ru.ru_maxrss / 1024.0
except Exception:
try:
import psutil # type: ignore[import-not-found]
return psutil.Process().memory_info().rss / (1024.0 * 1024.0)
except Exception:
return 0.0
# ---------------------------------------------------------------------------
# PerfTracker — accumulator for the duration of a perf_log() block.
# ---------------------------------------------------------------------------
@dataclass
class _PhaseStat:
"""Cumulative timing record for a single named scope."""
name: str
n_calls: int = 0
wall_s: float = 0.0
cpu_s: float = 0.0
def __iadd__(self, other: tuple[float, float]) -> "_PhaseStat":
self.n_calls += 1
self.wall_s += other[0]
self.cpu_s += other[1]
return self
[docs]
@dataclass
class PerfTracker:
"""Accumulator for per-phase wall/CPU timings, per-iteration SCF
rows, and memory snapshots over the course of a calculation.
Constructed inside :func:`perf_log` and passed implicitly to every
:class:`PerfScope` opened in the same async-context (uses
:class:`contextvars.ContextVar`, so perf trackers do NOT leak
across threads or asyncio tasks). Call sites that want to read
the live tracker — to add a custom snapshot, for instance —
use :func:`active_tracker`.
Public attributes:
* ``phases`` — map ``phase_name → _PhaseStat``; ``wall_s`` on
each entry is the cumulative wall time across all calls.
* ``scf_iters`` — list of per-iteration dicts (energy / dE /
grad / DIIS / wall_s); populated by SCF entry points wired
into the progress logger.
* ``memory_snapshots`` — list of ``{label, rss_mb, t_s}`` taken
at SCF transitions.
* ``threads`` — snapshot of the OpenMP thread count at tracker
construction (used to compute parallelism = cpu_s /
(wall_s × threads) per phase).
* ``t_start`` — perf_counter timestamp; phase rows report
wall-time deltas relative to this anchor.
"""
phases: dict[str, _PhaseStat] = field(default_factory=dict)
scf_iters: list[dict[str, Any]] = field(default_factory=list)
memory_snapshots: list[dict[str, Any]] = field(default_factory=list)
threads: int = 1
t_start: float = field(default_factory=time.perf_counter)
[docs]
def add_scope(self, name: str, wall_s: float, cpu_s: float) -> None:
stat = self.phases.setdefault(name, _PhaseStat(name=name))
stat += (wall_s, cpu_s)
[docs]
def add_scf_iter(self, **fields: Any) -> None:
"""Record one SCF iteration. Recognized keys: ``iter``,
``energy``, ``dE``, ``grad``, ``diis``. Extra keys are kept
verbatim in the row."""
self.scf_iters.append(fields)
[docs]
def snapshot_memory(self, label: str) -> None:
"""Take a labeled RSS snapshot. ``label`` is what shows up
in the report (typically ``"start_of_scf"``,
``"after_iter_5"``, ``"end_of_scf"``)."""
self.memory_snapshots.append({
"label": label,
"rss_mb": _rss_mb(),
"t_s": time.perf_counter() - self.t_start,
})
@property
def total_wall_s(self) -> float:
return time.perf_counter() - self.t_start
# ---------------------------------------------------------------------------
# Active-tracker plumbing via ContextVar (thread / task safe).
# ---------------------------------------------------------------------------
_active: contextvars.ContextVar[Optional[PerfTracker]] = contextvars.ContextVar(
"vibeqc_perf_tracker", default=None,
)
[docs]
def active_tracker() -> Optional[PerfTracker]:
"""The currently active :class:`PerfTracker`, or ``None`` if no
:func:`perf_log` block is open. Used by lower-level instrumented
code paths (e.g. periodic SCF iteration loops) to report
per-iter rows without taking an explicit tracker argument."""
return _active.get()
def _detect_threads() -> int:
"""Snapshot the OpenMP thread count vibe-qc would actually use.
Falls back to ``os.cpu_count()`` if the C++ accessor is
unavailable (early-import path / non-built tree)."""
try:
from ._vibeqc_core import get_num_threads
return int(get_num_threads())
except Exception:
return int(os.cpu_count() or 1)
# ---------------------------------------------------------------------------
# PerfScope — RAII timer.
# ---------------------------------------------------------------------------
[docs]
class PerfScope:
"""Context manager that times its body and pushes the wall + CPU
delta into the currently active :class:`PerfTracker`. No-op when
no tracker is active, so call sites can sprinkle ``with
PerfScope("phase_name"): ...`` unconditionally — the cost when
perf logging is off is one ``contextvars.ContextVar.get()`` and
one ``time.perf_counter()`` call.
Re-entrant: nesting two scopes with the same name accumulates
correctly (each scope contributes its own wall/CPU).
"""
__slots__ = ("name", "_t_wall", "_t_cpu")
[docs]
def __init__(self, name: str) -> None:
self.name = name
self._t_wall = 0.0
self._t_cpu = 0.0
def __enter__(self) -> "PerfScope":
self._t_wall = time.perf_counter()
self._t_cpu = time.process_time()
return self
def __exit__(self, *exc: Any) -> None:
wall = time.perf_counter() - self._t_wall
cpu = time.process_time() - self._t_cpu
tracker = _active.get()
if tracker is not None:
tracker.add_scope(self.name, wall, cpu)
# ---------------------------------------------------------------------------
# perf_log() context manager — install/uninstall a tracker.
# ---------------------------------------------------------------------------
[docs]
@contextmanager
def perf_log(
path: Union[str, os.PathLike, None] = None,
) -> Iterator[PerfTracker]:
"""Activate a fresh :class:`PerfTracker` for the duration of the
block; write the formatted report to ``path`` on exit.
Resolution order for ``path``:
1. Explicit argument when not ``None``.
2. ``$VIBEQC_PERFLOG`` env var.
3. No file written (the tracker still accumulates and the user
can read it from the yielded value).
Nesting two ``perf_log()`` blocks is allowed but discouraged —
the inner block's tracker is independent and only inner-block
work routes there. Most callers want a single top-level block.
Example::
import vibeqc as vq
with vq.perf_log("h2o.perf") as tracker:
vq.run_rhf(mol, basis, opts)
# h2o.perf now exists; tracker.phases is also accessible
# post-block for programmatic inspection.
"""
target: Optional[Path] = None
if path is not None:
target = Path(os.fspath(path))
else:
env = os.environ.get("VIBEQC_PERFLOG", "").strip()
if env:
target = Path(env)
tracker = PerfTracker(threads=_detect_threads())
token = _active.set(tracker)
try:
yield tracker
finally:
_active.reset(token)
if target is not None:
try:
target.parent.mkdir(parents=True, exist_ok=True)
target.write_text(format_perf_report(tracker), encoding="utf-8")
except OSError as exc:
# A perf-log write failure must never propagate up
# and crash a finished calculation. Print a warning
# so the user notices, but keep going.
print(
f"vibeqc.perf: warning — could not write {target}: {exc}",
file=sys.stderr,
)
# ---------------------------------------------------------------------------
# Report formatter.
# ---------------------------------------------------------------------------
_BAR = "=" * 72
def _fmt_seconds(s: float) -> str:
if s < 1.0:
return f"{s * 1e3:8.2f} ms"
if s < 60.0:
return f"{s:8.3f} s "
m, sec = divmod(s, 60.0)
return f"{int(m):>4d}m{sec:05.2f}s"