666 lines
20 KiB
Python
666 lines
20 KiB
Python
"""Low-overhead, run-local performance instrumentation for simulations.
|
|
|
|
The profiling mode is intentionally read once when this module is imported.
|
|
``standard`` records low-frequency pipeline stages, while ``audit`` also wraps
|
|
hot RHS/property operations and computes exact-input reuse metrics. With
|
|
profiling disabled, decorators return the original callable while classes are
|
|
being defined, so ordinary simulation calls do not pass through a wrapper.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Generator, Mapping
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar, Token
|
|
from dataclasses import dataclass, field
|
|
from functools import wraps
|
|
import inspect
|
|
import os
|
|
import struct
|
|
from time import perf_counter_ns
|
|
from typing import Any, Literal, ParamSpec, TypeVar, cast
|
|
|
|
|
|
ProfileMode = Literal["off", "standard", "audit"]
|
|
|
|
_P = ParamSpec("_P")
|
|
_R = TypeVar("_R")
|
|
_MODE_RANK: Mapping[ProfileMode, int] = {"off": 0, "standard": 1, "audit": 2}
|
|
|
|
|
|
def _read_startup_mode() -> ProfileMode:
|
|
raw_mode = os.getenv("SIMULATIONAPP_PROFILE", "off").strip().lower()
|
|
aliases: dict[str, ProfileMode] = {
|
|
"": "off",
|
|
"0": "off",
|
|
"false": "off",
|
|
"no": "off",
|
|
"off": "off",
|
|
"1": "standard",
|
|
"true": "standard",
|
|
"yes": "standard",
|
|
"on": "standard",
|
|
"standard": "standard",
|
|
"audit": "audit",
|
|
}
|
|
try:
|
|
return aliases[raw_mode]
|
|
except KeyError as exc:
|
|
raise ValueError(
|
|
"SIMULATIONAPP_PROFILE must be one of: off, standard, audit."
|
|
) from exc
|
|
|
|
|
|
PROFILE_MODE: ProfileMode = _read_startup_mode()
|
|
|
|
|
|
def _minimum_mode(value: str) -> ProfileMode:
|
|
normalized = value.strip().lower()
|
|
if normalized not in _MODE_RANK:
|
|
raise ValueError("minimum_mode must be one of: off, standard, audit.")
|
|
return cast(ProfileMode, normalized)
|
|
|
|
|
|
def _mode_enabled(minimum_mode: ProfileMode) -> bool:
|
|
# ``off`` is an unconditional zero-wrapper mode, even if a caller passes
|
|
# ``minimum_mode="off"`` by mistake.
|
|
return (
|
|
PROFILE_MODE != "off"
|
|
and _MODE_RANK[PROFILE_MODE] >= _MODE_RANK[minimum_mode]
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class _TimingStats:
|
|
calls: int = 0
|
|
inclusive_ns: int = 0
|
|
self_ns: int = 0
|
|
max_ns: int = 0
|
|
errors: int = 0
|
|
|
|
def record(self, inclusive_ns: int, self_ns: int, error: bool) -> None:
|
|
self.calls += 1
|
|
self.inclusive_ns += inclusive_ns
|
|
self.self_ns += self_ns
|
|
self.max_ns = max(self.max_ns, inclusive_ns)
|
|
if error:
|
|
self.errors += 1
|
|
|
|
def snapshot(self) -> dict[str, int]:
|
|
return {
|
|
"calls": self.calls,
|
|
"inclusiveNs": self.inclusive_ns,
|
|
"selfNs": self.self_ns,
|
|
"maxNs": self.max_ns,
|
|
"errors": self.errors,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class _PropertyStats(_TimingStats):
|
|
operation: str = ""
|
|
layer: str = "semantic"
|
|
medium: str = "unknown"
|
|
exact_input_unique: int = 0
|
|
exact_input_repeats: int = 0
|
|
iteration_calls: int = 0
|
|
iteration_total: int = 0
|
|
iteration_max: int = 0
|
|
iteration_converged: int = 0
|
|
iteration_nonconverged: int = 0
|
|
cache_lookups: int = 0
|
|
cache_hits: int = 0
|
|
cache_misses: int = 0
|
|
|
|
def snapshot(self, *, audit: bool) -> dict[str, object]:
|
|
result: dict[str, object] = super().snapshot()
|
|
result.update(
|
|
{
|
|
"operation": self.operation,
|
|
"layer": self.layer,
|
|
"medium": self.medium,
|
|
"cacheLookups": self.cache_lookups,
|
|
"cacheHits": self.cache_hits,
|
|
"cacheMisses": self.cache_misses,
|
|
}
|
|
)
|
|
if audit:
|
|
result.update(
|
|
{
|
|
"exactInputUnique": self.exact_input_unique,
|
|
"exactInputRepeats": self.exact_input_repeats,
|
|
"iterationCalls": self.iteration_calls,
|
|
"iterationTotal": self.iteration_total,
|
|
"iterationMax": self.iteration_max,
|
|
"iterationConverged": self.iteration_converged,
|
|
"iterationNonconverged": self.iteration_nonconverged,
|
|
}
|
|
)
|
|
return result
|
|
|
|
|
|
@dataclass
|
|
class _ActiveSpan:
|
|
trace: PerformanceTrace
|
|
name: str
|
|
started_ns: int
|
|
property_key: str | None = None
|
|
property_operation: str | None = None
|
|
property_outermost: bool = False
|
|
child_ns: int = 0
|
|
|
|
|
|
@dataclass
|
|
class PerformanceTrace:
|
|
"""Mutable counters owned by exactly one :func:`profile_run` context."""
|
|
|
|
mode: ProfileMode
|
|
_phases: dict[str, _TimingStats] = field(default_factory=dict, repr=False)
|
|
_properties: dict[str, _PropertyStats] = field(default_factory=dict, repr=False)
|
|
_property_outermost_ns: int = field(default=0, repr=False)
|
|
|
|
@property
|
|
def enabled(self) -> bool:
|
|
return self.mode != "off"
|
|
|
|
def _property_stats(
|
|
self,
|
|
key: str,
|
|
*,
|
|
operation: str,
|
|
layer: str,
|
|
medium: str,
|
|
) -> _PropertyStats:
|
|
stats = self._properties.get(key)
|
|
if stats is None:
|
|
stats = _PropertyStats(
|
|
operation=operation,
|
|
layer=layer,
|
|
medium=medium,
|
|
)
|
|
self._properties[key] = stats
|
|
return stats
|
|
|
|
def _record_span(self, frame: _ActiveSpan, elapsed_ns: int, error: bool) -> None:
|
|
self_ns = max(0, elapsed_ns - frame.child_ns)
|
|
if frame.property_key is None:
|
|
stats = self._phases.setdefault(frame.name, _TimingStats())
|
|
else:
|
|
layer, medium, operation = frame.property_key.split("|", 2)
|
|
stats = self._property_stats(
|
|
frame.property_key,
|
|
operation=operation,
|
|
layer=layer,
|
|
medium=medium,
|
|
)
|
|
if frame.property_outermost:
|
|
self._property_outermost_ns += elapsed_ns
|
|
stats.record(elapsed_ns, self_ns, error)
|
|
|
|
def _record_exact_input(
|
|
self,
|
|
key: str,
|
|
*,
|
|
operation: str,
|
|
layer: str,
|
|
medium: str,
|
|
fingerprint: object,
|
|
) -> None:
|
|
stats = self._property_stats(
|
|
key,
|
|
operation=operation,
|
|
layer=layer,
|
|
medium=medium,
|
|
)
|
|
shadow_key = (key, fingerprint)
|
|
shadow = _PROPERTY_SHADOW.get()
|
|
if shadow is None:
|
|
# A trace normally installs its own set in ``profile_run``. Keep
|
|
# the ContextVar default immutable so no task can accidentally
|
|
# share a process-global shadow set.
|
|
shadow = set()
|
|
_PROPERTY_SHADOW.set(shadow)
|
|
if shadow_key in shadow:
|
|
stats.exact_input_repeats += 1
|
|
else:
|
|
shadow.add(shadow_key)
|
|
stats.exact_input_unique += 1
|
|
|
|
def _record_iterations(
|
|
self,
|
|
key: str,
|
|
*,
|
|
operation: str,
|
|
layer: str,
|
|
medium: str,
|
|
iterations: int,
|
|
converged: bool,
|
|
) -> None:
|
|
stats = self._property_stats(
|
|
key,
|
|
operation=operation,
|
|
layer=layer,
|
|
medium=medium,
|
|
)
|
|
iteration_count = max(0, int(iterations))
|
|
stats.iteration_calls += 1
|
|
stats.iteration_total += iteration_count
|
|
stats.iteration_max = max(stats.iteration_max, iteration_count)
|
|
if converged:
|
|
stats.iteration_converged += 1
|
|
else:
|
|
stats.iteration_nonconverged += 1
|
|
|
|
def _record_cache(
|
|
self,
|
|
key: str,
|
|
*,
|
|
operation: str,
|
|
layer: str,
|
|
medium: str,
|
|
hits: int,
|
|
misses: int,
|
|
) -> None:
|
|
stats = self._property_stats(
|
|
key,
|
|
operation=operation,
|
|
layer=layer,
|
|
medium=medium,
|
|
)
|
|
hit_delta = max(0, hits)
|
|
miss_delta = max(0, misses)
|
|
stats.cache_hits += hit_delta
|
|
stats.cache_misses += miss_delta
|
|
stats.cache_lookups += hit_delta + miss_delta
|
|
|
|
def snapshot(self) -> dict[str, object]:
|
|
"""Return a detached, JSON-serializable copy of all counters."""
|
|
|
|
audit = self.mode == "audit"
|
|
return {
|
|
"mode": self.mode,
|
|
"phases": {
|
|
name: self._phases[name].snapshot()
|
|
for name in sorted(self._phases)
|
|
},
|
|
"properties": {
|
|
_public_property_key(key): self._properties[key].snapshot(audit=audit)
|
|
for key in sorted(self._properties)
|
|
},
|
|
"propertyOutermostNs": self._property_outermost_ns,
|
|
}
|
|
|
|
|
|
_CURRENT_TRACE: ContextVar[PerformanceTrace | None] = ContextVar(
|
|
"simulation_performance_trace",
|
|
default=None,
|
|
)
|
|
_ACTIVE_SPANS: ContextVar[tuple[_ActiveSpan, ...]] = ContextVar(
|
|
"simulation_performance_spans",
|
|
default=(),
|
|
)
|
|
_PROPERTY_SHADOW: ContextVar[set[object] | None] = ContextVar(
|
|
"simulation_property_shadow",
|
|
default=None,
|
|
)
|
|
|
|
|
|
def _public_property_key(key: str) -> str:
|
|
layer, medium, operation = key.split("|", 2)
|
|
return f"{layer}.{medium}.{operation}"
|
|
|
|
|
|
@contextmanager
|
|
def _tracked_span(
|
|
trace: PerformanceTrace,
|
|
name: str,
|
|
*,
|
|
property_key: str | None = None,
|
|
property_operation: str | None = None,
|
|
reset_property_shadow: bool = False,
|
|
) -> Generator[None, None, None]:
|
|
stack = _ACTIVE_SPANS.get()
|
|
property_outermost = property_key is not None and not any(
|
|
item.property_key is not None for item in stack
|
|
)
|
|
frame = _ActiveSpan(
|
|
trace=trace,
|
|
name=name,
|
|
started_ns=perf_counter_ns(),
|
|
property_key=property_key,
|
|
property_operation=property_operation,
|
|
property_outermost=property_outermost,
|
|
)
|
|
stack_token = _ACTIVE_SPANS.set((*stack, frame))
|
|
shadow_token: Token[set[object] | None] | None = None
|
|
if reset_property_shadow and trace.mode == "audit":
|
|
shadow_token = _PROPERTY_SHADOW.set(set())
|
|
error = False
|
|
try:
|
|
yield
|
|
except BaseException:
|
|
error = True
|
|
raise
|
|
finally:
|
|
elapsed_ns = max(0, perf_counter_ns() - frame.started_ns)
|
|
_ACTIVE_SPANS.reset(stack_token)
|
|
if shadow_token is not None:
|
|
_PROPERTY_SHADOW.reset(shadow_token)
|
|
if stack:
|
|
stack[-1].child_ns += elapsed_ns
|
|
trace._record_span(frame, elapsed_ns, error)
|
|
|
|
|
|
@contextmanager
|
|
def profile_run() -> Generator[PerformanceTrace, None, None]:
|
|
"""Create and bind an isolated trace for one simulation run.
|
|
|
|
The yielded trace remains usable after the context exits, which lets the
|
|
caller attach ``trace.snapshot()`` to a result without exposing live state.
|
|
"""
|
|
|
|
trace = PerformanceTrace(mode=PROFILE_MODE)
|
|
trace_token = _CURRENT_TRACE.set(trace)
|
|
spans_token = _ACTIVE_SPANS.set(())
|
|
shadow_token = _PROPERTY_SHADOW.set(set())
|
|
try:
|
|
if trace.enabled:
|
|
with _tracked_span(trace, "simulation.total"):
|
|
yield trace
|
|
else:
|
|
yield trace
|
|
finally:
|
|
_PROPERTY_SHADOW.reset(shadow_token)
|
|
_ACTIVE_SPANS.reset(spans_token)
|
|
_CURRENT_TRACE.reset(trace_token)
|
|
|
|
|
|
@contextmanager
|
|
def performance_span(
|
|
name: str,
|
|
minimum_mode: str = "standard",
|
|
reset_property_shadow: bool = False,
|
|
) -> Generator[None, None, None]:
|
|
"""Time a block in the current run, or act as a no-op outside one."""
|
|
|
|
minimum = _minimum_mode(minimum_mode)
|
|
trace = _CURRENT_TRACE.get()
|
|
if trace is None or not _mode_enabled(minimum):
|
|
yield
|
|
return
|
|
with _tracked_span(
|
|
trace,
|
|
name,
|
|
reset_property_shadow=reset_property_shadow,
|
|
):
|
|
yield
|
|
|
|
|
|
def profile_phase(
|
|
name: str,
|
|
minimum_mode: str = "standard",
|
|
reset_property_shadow: bool = False,
|
|
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
|
|
"""Decorate a simulation phase while preserving the off-mode callable."""
|
|
|
|
minimum = _minimum_mode(minimum_mode)
|
|
|
|
def decorate(function: Callable[_P, _R]) -> Callable[_P, _R]:
|
|
if not _mode_enabled(minimum):
|
|
return function
|
|
|
|
if inspect.iscoroutinefunction(function):
|
|
|
|
@wraps(function)
|
|
async def async_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> Any:
|
|
trace = _CURRENT_TRACE.get()
|
|
if trace is None:
|
|
return await function(*args, **kwargs)
|
|
with _tracked_span(
|
|
trace,
|
|
name,
|
|
reset_property_shadow=reset_property_shadow,
|
|
):
|
|
return await function(*args, **kwargs)
|
|
|
|
return cast(Callable[_P, _R], async_wrapper)
|
|
|
|
@wraps(function)
|
|
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
|
trace = _CURRENT_TRACE.get()
|
|
if trace is None:
|
|
return function(*args, **kwargs)
|
|
with _tracked_span(
|
|
trace,
|
|
name,
|
|
reset_property_shadow=reset_property_shadow,
|
|
):
|
|
return function(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return decorate
|
|
|
|
|
|
def _medium_name(args: tuple[object, ...]) -> str:
|
|
if not args:
|
|
return "unknown"
|
|
owner = args[0]
|
|
configured_name = getattr(owner, "name", None)
|
|
if isinstance(configured_name, str) and configured_name:
|
|
return configured_name
|
|
return type(owner).__name__
|
|
|
|
|
|
def _fingerprint(value: object) -> object:
|
|
"""Build a hashable, bit-exact token without retaining arbitrary objects."""
|
|
|
|
if value is None or isinstance(value, (bool, int, str, bytes)):
|
|
return (type(value).__name__, value)
|
|
if isinstance(value, float):
|
|
return ("float64", struct.pack("!d", value))
|
|
if isinstance(value, tuple):
|
|
return ("tuple", tuple(_fingerprint(item) for item in value))
|
|
if isinstance(value, list):
|
|
return ("list", tuple(_fingerprint(item) for item in value))
|
|
if isinstance(value, Mapping):
|
|
items = [(_fingerprint(key), _fingerprint(item)) for key, item in value.items()]
|
|
items.sort(key=repr)
|
|
return ("mapping", tuple(items))
|
|
return (
|
|
"object",
|
|
type(value).__module__,
|
|
type(value).__qualname__,
|
|
id(value),
|
|
)
|
|
|
|
|
|
def _input_fingerprint(
|
|
signature: inspect.Signature | None,
|
|
args: tuple[object, ...],
|
|
kwargs: dict[str, object],
|
|
) -> object:
|
|
if signature is not None:
|
|
try:
|
|
bound = signature.bind(*args, **kwargs)
|
|
bound.apply_defaults()
|
|
return tuple(
|
|
(name, _fingerprint(value))
|
|
for name, value in bound.arguments.items()
|
|
)
|
|
except TypeError:
|
|
pass
|
|
return (
|
|
_fingerprint(args),
|
|
tuple(sorted((name, _fingerprint(value)) for name, value in kwargs.items())),
|
|
)
|
|
|
|
|
|
def _cache_counts(function: Callable[..., object]) -> tuple[int, int] | None:
|
|
cache_info = getattr(function, "cache_info", None)
|
|
if not callable(cache_info):
|
|
return None
|
|
try:
|
|
info = cache_info()
|
|
return int(info.hits), int(info.misses)
|
|
except (AttributeError, TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
def _copy_cache_api(source: Callable[..., object], target: Callable[..., object]) -> None:
|
|
for attribute in ("cache_clear", "cache_info", "cache_parameters"):
|
|
value = getattr(source, attribute, None)
|
|
if value is not None:
|
|
setattr(target, attribute, value)
|
|
|
|
|
|
def profile_property(
|
|
operation: str,
|
|
layer: str = "semantic",
|
|
minimum_mode: str = "audit",
|
|
capture_inputs: bool = True,
|
|
track_cache: bool = False,
|
|
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
|
|
"""Decorate one thermodynamic property operation."""
|
|
|
|
minimum = _minimum_mode(minimum_mode)
|
|
|
|
def decorate(function: Callable[_P, _R]) -> Callable[_P, _R]:
|
|
if not _mode_enabled(minimum):
|
|
return function
|
|
try:
|
|
signature: inspect.Signature | None = inspect.signature(function)
|
|
except (TypeError, ValueError):
|
|
signature = None
|
|
|
|
def prepare(
|
|
args: tuple[object, ...],
|
|
kwargs: dict[str, object],
|
|
) -> tuple[PerformanceTrace | None, str, str, tuple[int, int] | None]:
|
|
trace = _CURRENT_TRACE.get()
|
|
medium = _medium_name(args)
|
|
key = f"{layer}|{medium}|{operation}"
|
|
if trace is not None and trace.mode == "audit" and capture_inputs:
|
|
trace._record_exact_input(
|
|
key,
|
|
operation=operation,
|
|
layer=layer,
|
|
medium=medium,
|
|
fingerprint=_input_fingerprint(signature, args, kwargs),
|
|
)
|
|
before = _cache_counts(function) if trace is not None and track_cache else None
|
|
return trace, medium, key, before
|
|
|
|
def finish_cache(
|
|
trace: PerformanceTrace | None,
|
|
medium: str,
|
|
key: str,
|
|
before: tuple[int, int] | None,
|
|
) -> None:
|
|
if trace is None or before is None:
|
|
return
|
|
after = _cache_counts(function)
|
|
if after is None:
|
|
return
|
|
trace._record_cache(
|
|
key,
|
|
operation=operation,
|
|
layer=layer,
|
|
medium=medium,
|
|
hits=after[0] - before[0],
|
|
misses=after[1] - before[1],
|
|
)
|
|
|
|
if inspect.iscoroutinefunction(function):
|
|
|
|
@wraps(function)
|
|
async def async_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> Any:
|
|
object_args = cast(tuple[object, ...], args)
|
|
object_kwargs = cast(dict[str, object], kwargs)
|
|
trace, medium, key, before = prepare(object_args, object_kwargs)
|
|
try:
|
|
if trace is None:
|
|
return await function(*args, **kwargs)
|
|
with _tracked_span(
|
|
trace,
|
|
f"property.{_public_property_key(key)}",
|
|
property_key=key,
|
|
property_operation=operation,
|
|
):
|
|
return await function(*args, **kwargs)
|
|
finally:
|
|
finish_cache(trace, medium, key, before)
|
|
|
|
_copy_cache_api(function, async_wrapper)
|
|
return cast(Callable[_P, _R], async_wrapper)
|
|
|
|
@wraps(function)
|
|
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
|
object_args = cast(tuple[object, ...], args)
|
|
object_kwargs = cast(dict[str, object], kwargs)
|
|
trace, medium, key, before = prepare(object_args, object_kwargs)
|
|
try:
|
|
if trace is None:
|
|
return function(*args, **kwargs)
|
|
with _tracked_span(
|
|
trace,
|
|
f"property.{_public_property_key(key)}",
|
|
property_key=key,
|
|
property_operation=operation,
|
|
):
|
|
return function(*args, **kwargs)
|
|
finally:
|
|
finish_cache(trace, medium, key, before)
|
|
|
|
_copy_cache_api(function, wrapper)
|
|
return wrapper
|
|
|
|
return decorate
|
|
|
|
|
|
def record_property_iterations(
|
|
operation: str,
|
|
iterations: int,
|
|
converged: bool,
|
|
) -> None:
|
|
"""Record inverse-property solver iterations in audit mode."""
|
|
|
|
trace = _CURRENT_TRACE.get()
|
|
if trace is None or trace.mode != "audit":
|
|
return
|
|
for frame in reversed(_ACTIVE_SPANS.get()):
|
|
if (
|
|
frame.property_key is not None
|
|
and frame.property_operation == operation
|
|
):
|
|
layer, medium, _unused_operation = frame.property_key.split("|", 2)
|
|
trace._record_iterations(
|
|
frame.property_key,
|
|
operation=operation,
|
|
layer=layer,
|
|
medium=medium,
|
|
iterations=iterations,
|
|
converged=converged,
|
|
)
|
|
return
|
|
key = f"semantic|unknown|{operation}"
|
|
trace._record_iterations(
|
|
key,
|
|
operation=operation,
|
|
layer="semantic",
|
|
medium="unknown",
|
|
iterations=iterations,
|
|
converged=converged,
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"PROFILE_MODE",
|
|
"PerformanceTrace",
|
|
"performance_span",
|
|
"profile_phase",
|
|
"profile_property",
|
|
"profile_run",
|
|
"record_property_iterations",
|
|
]
|