expr_tracker 0.2.3__py3-none-any.whl → 0.2.4__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- expr_tracker/__init__.py +2 -1
- expr_tracker/history/store.py +4 -0
- expr_tracker/plugins.py +269 -0
- expr_tracker/spans.py +142 -14
- expr_tracker/trace.py +1 -0
- {expr_tracker-0.2.3.dist-info → expr_tracker-0.2.4.dist-info}/METADATA +7 -3
- {expr_tracker-0.2.3.dist-info → expr_tracker-0.2.4.dist-info}/RECORD +10 -9
- {expr_tracker-0.2.3.dist-info → expr_tracker-0.2.4.dist-info}/WHEEL +1 -1
- {expr_tracker-0.2.3.dist-info → expr_tracker-0.2.4.dist-info}/entry_points.txt +0 -0
- {expr_tracker-0.2.3.dist-info → expr_tracker-0.2.4.dist-info}/licenses/LICENSE +0 -0
expr_tracker/__init__.py
CHANGED
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
from importlib.metadata import PackageNotFoundError
|
|
2
2
|
from importlib.metadata import version as _version
|
|
3
3
|
|
|
4
|
-
from . import tracker
|
|
4
|
+
from . import plugins, tracker
|
|
5
5
|
from .alerts import (
|
|
6
6
|
AlertConfig,
|
|
7
7
|
AlertLevel,
|
|
@@ -61,6 +61,7 @@ __all__ = [
|
|
|
61
61
|
"list_alert_rules",
|
|
62
62
|
"log",
|
|
63
63
|
"log_artifact",
|
|
64
|
+
"plugins",
|
|
64
65
|
"register_backend",
|
|
65
66
|
"remove_alert_rule",
|
|
66
67
|
"span",
|
expr_tracker/history/store.py
CHANGED
|
@@ -19,6 +19,7 @@ from collections import deque
|
|
|
19
19
|
from collections.abc import Callable, Sequence
|
|
20
20
|
from dataclasses import dataclass, fields
|
|
21
21
|
from pathlib import Path
|
|
22
|
+
from typing import Any
|
|
22
23
|
|
|
23
24
|
from loguru import logger
|
|
24
25
|
|
|
@@ -74,6 +75,9 @@ class HistoryOptions:
|
|
|
74
75
|
stream: str | None = None
|
|
75
76
|
# write the full span tree to spans.jsonl beside the metrics
|
|
76
77
|
spans: bool = True
|
|
78
|
+
# run-wide defaults for et.span(); a span's own arguments win
|
|
79
|
+
span_print_fn: Callable[[str], None] | None = None
|
|
80
|
+
span_plugins: Sequence[Any] = ()
|
|
77
81
|
# alerts and output
|
|
78
82
|
alert_window: int = DEFAULT_WINDOW
|
|
79
83
|
print_to_screen: bool = False
|
expr_tracker/plugins.py
ADDED
|
@@ -0,0 +1,269 @@
|
|
|
1
|
+
"""Span plugins: measure a resource across a span and attach it to the record.
|
|
2
|
+
|
|
3
|
+
A plugin is any object with ``start(span)`` and ``end(span) -> dict``; a plain
|
|
4
|
+
``fn(span) -> dict`` works as an end-only plugin. Whatever ``end`` returns is
|
|
5
|
+
merged into the span's metrics under the span's own path, so a plugin returning
|
|
6
|
+
``{"cpu_percent": 91.2}`` on a ``forward`` span writes ``forward/cpu_percent``.
|
|
7
|
+
|
|
8
|
+
Plugins are opt-in because they are not free. A bare span costs ~13us; see each
|
|
9
|
+
plugin's docstring for what it adds.
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
from __future__ import annotations
|
|
13
|
+
|
|
14
|
+
import atexit
|
|
15
|
+
import threading
|
|
16
|
+
import time
|
|
17
|
+
from collections import deque
|
|
18
|
+
from typing import Any
|
|
19
|
+
|
|
20
|
+
from loguru import logger
|
|
21
|
+
|
|
22
|
+
__all__ = ["CpuTime", "GpuStats", "TorchMemory", "stop_samplers"]
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class CpuTime:
|
|
26
|
+
"""CPU utilisation over the span, from :func:`time.process_time`.
|
|
27
|
+
|
|
28
|
+
Reports ``cpu_percent`` (process CPU time / wall time, so >100% means the
|
|
29
|
+
span used more than one core) and ``cpu_time_ms``. Exact, not sampled, and
|
|
30
|
+
the cheapest thing here at ~1.5us per span. Counts every thread in the
|
|
31
|
+
process, so concurrent work outside the span inflates it.
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
__slots__ = ("_stack",)
|
|
35
|
+
|
|
36
|
+
def __init__(self):
|
|
37
|
+
self._stack: dict[int, list[float]] = {}
|
|
38
|
+
|
|
39
|
+
def _frames(self) -> list[float]:
|
|
40
|
+
return self._stack.setdefault(threading.get_ident(), [])
|
|
41
|
+
|
|
42
|
+
def start(self, span) -> None:
|
|
43
|
+
self._frames().append(time.process_time())
|
|
44
|
+
|
|
45
|
+
def end(self, span) -> dict[str, float]:
|
|
46
|
+
frames = self._frames()
|
|
47
|
+
if not frames:
|
|
48
|
+
return {}
|
|
49
|
+
used_ms = (time.process_time() - frames.pop()) * 1000.0
|
|
50
|
+
percent = 100.0 * used_ms / span.duration_ms if span.duration_ms > 0 else 0.0
|
|
51
|
+
return {"cpu_time_ms": round(used_ms, 3), "cpu_percent": round(percent, 1)}
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
class TorchMemory:
|
|
55
|
+
"""Peak CUDA memory the allocator held during the span, via torch.
|
|
56
|
+
|
|
57
|
+
Reports ``gpu_mem_peak_mb`` and ``gpu_mem_delta_mb`` (end minus start, i.e.
|
|
58
|
+
what the span kept). Exact, ~10us per span.
|
|
59
|
+
|
|
60
|
+
``torch.cuda.max_memory_allocated`` is one global high-water mark, so a
|
|
61
|
+
naive reset would let a child span erase its parent's history. This keeps
|
|
62
|
+
its own stack and rolls each child's peak back into the enclosing span.
|
|
63
|
+
"""
|
|
64
|
+
|
|
65
|
+
__slots__ = ("_device", "_stack", "_torch")
|
|
66
|
+
|
|
67
|
+
def __init__(self, device: Any = None):
|
|
68
|
+
self._device = device
|
|
69
|
+
self._torch = None
|
|
70
|
+
self._stack: dict[int, list[tuple[float, float]]] = {}
|
|
71
|
+
|
|
72
|
+
def _cuda(self):
|
|
73
|
+
if self._torch is None:
|
|
74
|
+
import torch # raised to _safely() and logged once per span otherwise
|
|
75
|
+
|
|
76
|
+
self._torch = torch
|
|
77
|
+
if not self._torch.cuda.is_available():
|
|
78
|
+
raise RuntimeError("CUDA is not available")
|
|
79
|
+
return self._torch.cuda
|
|
80
|
+
|
|
81
|
+
def _frames(self) -> list[list[float]]:
|
|
82
|
+
return self._stack.setdefault(threading.get_ident(), [])
|
|
83
|
+
|
|
84
|
+
def _read(self, cuda) -> tuple[float, float]:
|
|
85
|
+
"""(allocated, peak) in bytes.
|
|
86
|
+
|
|
87
|
+
One nested-stats call costs ~17us; the two public helpers rebuild and
|
|
88
|
+
flatten the whole stats dict separately and cost ~90us each.
|
|
89
|
+
"""
|
|
90
|
+
nested = getattr(cuda, "memory_stats_as_nested_dict", None)
|
|
91
|
+
if nested is not None:
|
|
92
|
+
stats = nested(device=self._device).get("allocated_bytes")
|
|
93
|
+
if stats:
|
|
94
|
+
return float(stats["all"]["current"]), float(stats["all"]["peak"])
|
|
95
|
+
return (
|
|
96
|
+
float(cuda.memory_allocated(self._device)),
|
|
97
|
+
float(cuda.max_memory_allocated(self._device)),
|
|
98
|
+
)
|
|
99
|
+
|
|
100
|
+
def start(self, span) -> None:
|
|
101
|
+
cuda = self._cuda()
|
|
102
|
+
opening, outer_peak = self._read(cuda)
|
|
103
|
+
# [peak the enclosing span had reached, our opening bytes, carried peak]
|
|
104
|
+
self._frames().append([outer_peak, opening, 0.0])
|
|
105
|
+
cuda.reset_peak_memory_stats(self._device)
|
|
106
|
+
|
|
107
|
+
def end(self, span) -> dict[str, float]:
|
|
108
|
+
frames = self._frames()
|
|
109
|
+
if not frames:
|
|
110
|
+
return {}
|
|
111
|
+
outer_peak, opening, carried = frames.pop()
|
|
112
|
+
closing, raw_peak = self._read(self._cuda())
|
|
113
|
+
# The counter was reset for us, so it holds our peak alone; our own
|
|
114
|
+
# children left anything that predated them in `carried`
|
|
115
|
+
peak = max(raw_peak, carried)
|
|
116
|
+
if frames:
|
|
117
|
+
# Hand our parent back the peak it had before we reset the counter
|
|
118
|
+
frames[-1][2] = max(frames[-1][2], outer_peak)
|
|
119
|
+
return {
|
|
120
|
+
"gpu_mem_peak_mb": round(peak / 1048576, 3),
|
|
121
|
+
"gpu_mem_delta_mb": round((closing - opening) / 1048576, 3),
|
|
122
|
+
}
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
class _NvmlSampler:
|
|
126
|
+
"""One background thread polling NVML, shared by every :class:`GpuStats`.
|
|
127
|
+
|
|
128
|
+
Sampling in the background is the only honest way to report utilisation:
|
|
129
|
+
NVML's counter is a rolling average over its own window, so reading it once
|
|
130
|
+
at the end of a 3ms span says nothing about that span.
|
|
131
|
+
"""
|
|
132
|
+
|
|
133
|
+
def __init__(self, index: int, interval: float):
|
|
134
|
+
self.index = index
|
|
135
|
+
self.interval = interval
|
|
136
|
+
self.samples: deque[tuple[float, float, float]] = deque(
|
|
137
|
+
maxlen=max(64, int(60.0 / interval))
|
|
138
|
+
)
|
|
139
|
+
self._handle = None
|
|
140
|
+
self._thread: threading.Thread | None = None
|
|
141
|
+
self._stop = threading.Event()
|
|
142
|
+
self._lock = threading.Lock()
|
|
143
|
+
|
|
144
|
+
def handle(self):
|
|
145
|
+
if self._handle is None:
|
|
146
|
+
import pynvml
|
|
147
|
+
|
|
148
|
+
pynvml.nvmlInit()
|
|
149
|
+
self._handle = pynvml.nvmlDeviceGetHandleByIndex(self.index)
|
|
150
|
+
self._pynvml = pynvml
|
|
151
|
+
return self._handle
|
|
152
|
+
|
|
153
|
+
def read(self) -> tuple[float, float]:
|
|
154
|
+
handle = self.handle()
|
|
155
|
+
rates = self._pynvml.nvmlDeviceGetUtilizationRates(handle)
|
|
156
|
+
memory = self._pynvml.nvmlDeviceGetMemoryInfo(handle)
|
|
157
|
+
return float(rates.gpu), float(memory.used)
|
|
158
|
+
|
|
159
|
+
def ensure_running(self) -> None:
|
|
160
|
+
thread = self._thread
|
|
161
|
+
if thread is not None and thread.is_alive():
|
|
162
|
+
return # the common case, off the lock
|
|
163
|
+
with self._lock:
|
|
164
|
+
if self._thread is not None and self._thread.is_alive():
|
|
165
|
+
return
|
|
166
|
+
self.handle() # fail loudly here, not once per sample
|
|
167
|
+
self._stop.clear()
|
|
168
|
+
self._thread = threading.Thread(
|
|
169
|
+
target=self._run, name=f"et-nvml-{self.index}", daemon=True
|
|
170
|
+
)
|
|
171
|
+
self._thread.start()
|
|
172
|
+
|
|
173
|
+
def _run(self) -> None:
|
|
174
|
+
while not self._stop.wait(self.interval):
|
|
175
|
+
try:
|
|
176
|
+
util, used = self.read()
|
|
177
|
+
except Exception as e:
|
|
178
|
+
logger.warning(f"NVML sampling stopped: {e}")
|
|
179
|
+
return
|
|
180
|
+
with self._lock:
|
|
181
|
+
self.samples.append((time.time(), util, used))
|
|
182
|
+
|
|
183
|
+
def stop(self) -> None:
|
|
184
|
+
self._stop.set()
|
|
185
|
+
thread = self._thread
|
|
186
|
+
if thread is not None:
|
|
187
|
+
thread.join(timeout=self.interval * 2 + 1.0)
|
|
188
|
+
self._thread = None
|
|
189
|
+
|
|
190
|
+
def window(self, start: float, end: float) -> list[tuple[float, float, float]]:
|
|
191
|
+
"""Samples inside [start, end], newest first.
|
|
192
|
+
|
|
193
|
+
Scanned backwards so a span pays for its own window, not for the whole
|
|
194
|
+
minute of history the deque keeps.
|
|
195
|
+
"""
|
|
196
|
+
picked = []
|
|
197
|
+
with self._lock:
|
|
198
|
+
samples = self.samples
|
|
199
|
+
for index in range(len(samples) - 1, -1, -1):
|
|
200
|
+
sample = samples[index]
|
|
201
|
+
if sample[0] < start:
|
|
202
|
+
break
|
|
203
|
+
if sample[0] <= end:
|
|
204
|
+
picked.append(sample)
|
|
205
|
+
return picked
|
|
206
|
+
|
|
207
|
+
|
|
208
|
+
_SAMPLERS: dict[tuple[int, float], _NvmlSampler] = {}
|
|
209
|
+
_SAMPLERS_LOCK = threading.Lock()
|
|
210
|
+
|
|
211
|
+
|
|
212
|
+
def _sampler(index: int, interval: float) -> _NvmlSampler:
|
|
213
|
+
key = (index, interval)
|
|
214
|
+
with _SAMPLERS_LOCK:
|
|
215
|
+
sampler = _SAMPLERS.get(key)
|
|
216
|
+
if sampler is None:
|
|
217
|
+
sampler = _SAMPLERS[key] = _NvmlSampler(index, interval)
|
|
218
|
+
return sampler
|
|
219
|
+
|
|
220
|
+
|
|
221
|
+
def stop_samplers() -> None:
|
|
222
|
+
"""Stop every background NVML sampler. Called for you at run teardown."""
|
|
223
|
+
with _SAMPLERS_LOCK:
|
|
224
|
+
samplers = list(_SAMPLERS.values())
|
|
225
|
+
_SAMPLERS.clear()
|
|
226
|
+
for sampler in samplers:
|
|
227
|
+
sampler.stop()
|
|
228
|
+
|
|
229
|
+
|
|
230
|
+
class GpuStats:
|
|
231
|
+
"""GPU utilisation and device memory over the span, sampled via NVML.
|
|
232
|
+
|
|
233
|
+
Reports ``gpu_percent`` (mean utilisation across the span) and
|
|
234
|
+
``gpu_mem_used_mb`` (peak device memory). Needs ``nvidia-ml-py``:
|
|
235
|
+
``pip install expr-tracker[gpu]``.
|
|
236
|
+
|
|
237
|
+
A background thread does the sampling, so the span itself pays almost
|
|
238
|
+
nothing. Device memory covers the whole GPU, not just this process; use
|
|
239
|
+
:class:`TorchMemory` for what your own allocator holds.
|
|
240
|
+
|
|
241
|
+
Spans shorter than ``interval`` may contain no sample. They fall back to one
|
|
242
|
+
live reading (~14us), which for utilisation is NVML's own rolling average
|
|
243
|
+
rather than a measurement of the span.
|
|
244
|
+
"""
|
|
245
|
+
|
|
246
|
+
__slots__ = ("_sampler",)
|
|
247
|
+
|
|
248
|
+
def __init__(self, index: int = 0, interval: float = 0.1):
|
|
249
|
+
if interval <= 0:
|
|
250
|
+
raise ValueError("interval must be > 0")
|
|
251
|
+
self._sampler = _sampler(int(index), float(interval))
|
|
252
|
+
|
|
253
|
+
def start(self, span) -> None:
|
|
254
|
+
self._sampler.ensure_running()
|
|
255
|
+
|
|
256
|
+
def end(self, span) -> dict[str, float]:
|
|
257
|
+
samples = self._sampler.window(span.started_at, time.time())
|
|
258
|
+
if samples:
|
|
259
|
+
util = sum(s[1] for s in samples) / len(samples)
|
|
260
|
+
used = max(s[2] for s in samples)
|
|
261
|
+
else:
|
|
262
|
+
util, used = self._sampler.read()
|
|
263
|
+
return {
|
|
264
|
+
"gpu_percent": round(util, 1),
|
|
265
|
+
"gpu_mem_used_mb": round(used / 1048576, 3),
|
|
266
|
+
}
|
|
267
|
+
|
|
268
|
+
|
|
269
|
+
atexit.register(stop_samplers)
|
expr_tracker/spans.py
CHANGED
|
@@ -13,6 +13,7 @@ from __future__ import annotations
|
|
|
13
13
|
import functools
|
|
14
14
|
import threading
|
|
15
15
|
import time
|
|
16
|
+
from collections.abc import Callable, Sequence
|
|
16
17
|
from contextvars import ContextVar
|
|
17
18
|
from typing import Any
|
|
18
19
|
|
|
@@ -24,6 +25,41 @@ _STACK: ContextVar[tuple[Span, ...]] = ContextVar("et_span_stack", default=())
|
|
|
24
25
|
|
|
25
26
|
DURATION_SUFFIX = "duration_ms"
|
|
26
27
|
COUNT_SUFFIX = "count"
|
|
28
|
+
RESERVED_METRICS = frozenset({DURATION_SUFFIX, COUNT_SUFFIX})
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
_WARNED: set[tuple[str, str, str]] = set()
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _warn_once(*key: str) -> None:
|
|
35
|
+
"""A plugin that fails once fails every span; say so only the first time."""
|
|
36
|
+
if key in _WARNED:
|
|
37
|
+
return
|
|
38
|
+
if len(_WARNED) > 256: # a pathological plugin must not leak memory
|
|
39
|
+
_WARNED.clear()
|
|
40
|
+
_WARNED.add(key)
|
|
41
|
+
logger.warning(key[-1])
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def _safely(plugin: Any, hook: str, span: Span):
|
|
45
|
+
"""Call a plugin hook; a broken plugin must not break the measured code."""
|
|
46
|
+
method = getattr(plugin, hook, None)
|
|
47
|
+
if method is None and hook == "end" and callable(plugin):
|
|
48
|
+
method = plugin # a plain callable is an end-only plugin
|
|
49
|
+
if method is None:
|
|
50
|
+
return None
|
|
51
|
+
try:
|
|
52
|
+
return method(span)
|
|
53
|
+
except Exception as e:
|
|
54
|
+
name = type(plugin).__name__
|
|
55
|
+
_warn_once(name, hook, f"Span plugin {name} failed in {hook}(): {e}")
|
|
56
|
+
return None
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
def _pretty(value) -> str:
|
|
60
|
+
if isinstance(value, float):
|
|
61
|
+
return f"{value:.3f}".rstrip("0").rstrip(".")
|
|
62
|
+
return str(value)
|
|
27
63
|
|
|
28
64
|
|
|
29
65
|
def current_span() -> Span | None:
|
|
@@ -45,19 +81,37 @@ class Span:
|
|
|
45
81
|
"_store",
|
|
46
82
|
"_token",
|
|
47
83
|
"attributes",
|
|
84
|
+
"depth",
|
|
48
85
|
"duration_ms",
|
|
49
86
|
"error",
|
|
87
|
+
"metrics",
|
|
50
88
|
"name",
|
|
51
89
|
"path",
|
|
90
|
+
"plugins",
|
|
91
|
+
"print_depth",
|
|
92
|
+
"print_fn",
|
|
52
93
|
"start",
|
|
53
94
|
"started_at",
|
|
54
95
|
"track",
|
|
55
96
|
)
|
|
56
97
|
|
|
57
|
-
def __init__(
|
|
98
|
+
def __init__(
|
|
99
|
+
self,
|
|
100
|
+
name: str,
|
|
101
|
+
store,
|
|
102
|
+
*,
|
|
103
|
+
print_fn: Callable[[str], None] | None = None,
|
|
104
|
+
plugins: Sequence[Any] = (),
|
|
105
|
+
**attributes: Any,
|
|
106
|
+
):
|
|
58
107
|
self.name = str(name).strip("/") or "span"
|
|
59
108
|
self.attributes = dict(attributes)
|
|
60
109
|
self._store = store
|
|
110
|
+
self.print_fn = print_fn
|
|
111
|
+
self.plugins = tuple(plugins)
|
|
112
|
+
self.metrics: dict[str, Any] = {}
|
|
113
|
+
self.depth = 0
|
|
114
|
+
self.print_depth = 0
|
|
61
115
|
self.path = self.name
|
|
62
116
|
self.track = 0
|
|
63
117
|
self.start = 0.0
|
|
@@ -72,29 +126,78 @@ class Span:
|
|
|
72
126
|
def begin(self) -> Span:
|
|
73
127
|
parent = current_span()
|
|
74
128
|
self.path = f"{parent.path}/{self.name}" if parent else self.name
|
|
129
|
+
self.depth = parent.depth + 1 if parent else 0
|
|
75
130
|
# Children share their parent's track, so a tree never straddles lanes
|
|
76
131
|
self.track = parent.track if parent else threading.get_ident()
|
|
132
|
+
if self.print_fn is None and parent is not None:
|
|
133
|
+
self.print_fn = parent.print_fn # inherited, so a tree prints as one
|
|
134
|
+
if parent is not None and parent.print_fn is self.print_fn:
|
|
135
|
+
self.print_depth = parent.print_depth + 1
|
|
136
|
+
# else this span roots its own output: indent from 0, not from a parent
|
|
137
|
+
# that printed nothing
|
|
138
|
+
if not self.plugins and parent is not None:
|
|
139
|
+
self.plugins = parent.plugins
|
|
140
|
+
for plugin in self.plugins:
|
|
141
|
+
_safely(plugin, "start", self)
|
|
77
142
|
self.started_at = time.time()
|
|
78
143
|
self.start = time.perf_counter()
|
|
79
144
|
self._token = _STACK.set((*_STACK.get(), self))
|
|
145
|
+
self._announce("->", None)
|
|
80
146
|
return self
|
|
81
147
|
|
|
82
148
|
def end(self, error: BaseException | None = None) -> float:
|
|
83
149
|
"""Close the span and record it. Returns the duration in milliseconds."""
|
|
84
|
-
if self._ended:
|
|
150
|
+
if self._ended or self.started_at == 0.0:
|
|
151
|
+
# Never begun: there is no duration to report and nothing to record
|
|
85
152
|
return self.duration_ms
|
|
86
153
|
self._ended = True
|
|
87
154
|
self.duration_ms = (time.perf_counter() - self.start) * 1000.0
|
|
88
155
|
if error is not None:
|
|
89
156
|
self.error = type(error).__name__
|
|
157
|
+
for plugin in self.plugins:
|
|
158
|
+
measured = _safely(plugin, "end", self)
|
|
159
|
+
if isinstance(measured, dict):
|
|
160
|
+
self.metrics.update(self._admissible(plugin, measured))
|
|
90
161
|
if self._token is not None:
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
|
|
162
|
+
# Drop ourselves rather than resetting the token. A token restores
|
|
163
|
+
# the whole stack as it was, so a span ended out of order would
|
|
164
|
+
# resurrect spans that already closed underneath it.
|
|
165
|
+
self._token = None
|
|
166
|
+
_STACK.set(tuple(s for s in _STACK.get() if s is not self))
|
|
95
167
|
self._record()
|
|
168
|
+
self._announce("<-", self.duration_ms)
|
|
96
169
|
return self.duration_ms
|
|
97
170
|
|
|
171
|
+
def _admissible(self, plugin: Any, measured: dict) -> dict:
|
|
172
|
+
"""A plugin may add metrics to a span; it may not redefine its timing."""
|
|
173
|
+
if not RESERVED_METRICS.intersection(measured):
|
|
174
|
+
return measured
|
|
175
|
+
name = type(plugin).__name__
|
|
176
|
+
for key in RESERVED_METRICS.intersection(measured):
|
|
177
|
+
_warn_once(name, key, f"Span plugin {name} may not set {key!r}; ignored")
|
|
178
|
+
return {k: v for k, v in measured.items() if k not in RESERVED_METRICS}
|
|
179
|
+
|
|
180
|
+
def _announce(self, marker: str, duration_ms: float | None):
|
|
181
|
+
if self.print_fn is None:
|
|
182
|
+
return
|
|
183
|
+
indent = "\t" * self.print_depth
|
|
184
|
+
stamp = time.strftime("%H:%M:%S", time.localtime())
|
|
185
|
+
if duration_ms is None:
|
|
186
|
+
line = f"{indent}{marker} {self.name} {stamp}"
|
|
187
|
+
else:
|
|
188
|
+
extra = "".join(
|
|
189
|
+
f" {key}={_pretty(value)}" for key, value in self.metrics.items()
|
|
190
|
+
)
|
|
191
|
+
failed = f" !{self.error}" if self.error else ""
|
|
192
|
+
line = (
|
|
193
|
+
f"{indent}{marker} {self.name} {stamp} "
|
|
194
|
+
f"{duration_ms:.3f}ms{extra}{failed}"
|
|
195
|
+
)
|
|
196
|
+
try:
|
|
197
|
+
self.print_fn(line)
|
|
198
|
+
except Exception as e: # printing must never break the thing measured
|
|
199
|
+
_warn_once("print_fn", type(e).__name__, f"Span print handler failed: {e}")
|
|
200
|
+
|
|
98
201
|
def set(self, **attributes: Any) -> Span:
|
|
99
202
|
"""Attach attributes; they reach ``spans.jsonl``, not the metrics."""
|
|
100
203
|
self.attributes.update(attributes)
|
|
@@ -128,16 +231,19 @@ class Span:
|
|
|
128
231
|
metrics = {
|
|
129
232
|
f"{self.path}/{DURATION_SUFFIX}": self.duration_ms,
|
|
130
233
|
f"{self.path}/{COUNT_SUFFIX}": 1,
|
|
234
|
+
**{f"{self.path}/{key}": value for key, value in self.metrics.items()},
|
|
131
235
|
}
|
|
132
236
|
record = {
|
|
133
237
|
"name": self.path,
|
|
134
|
-
"depth": self.
|
|
238
|
+
"depth": self.depth,
|
|
135
239
|
"track": self.track,
|
|
136
240
|
"start": self.started_at,
|
|
137
241
|
"dur_ms": self.duration_ms,
|
|
138
242
|
}
|
|
139
243
|
if self.attributes:
|
|
140
244
|
record["args"] = self.attributes
|
|
245
|
+
if self.metrics:
|
|
246
|
+
record["metrics"] = self.metrics
|
|
141
247
|
if self.error:
|
|
142
248
|
record["error"] = self.error
|
|
143
249
|
try:
|
|
@@ -147,19 +253,34 @@ class Span:
|
|
|
147
253
|
|
|
148
254
|
|
|
149
255
|
class _NullSpan(Span):
|
|
150
|
-
"""What you get without a run: still times, records nothing."""
|
|
256
|
+
"""What you get without a run: still times and prints, records nothing."""
|
|
151
257
|
|
|
152
|
-
def __init__(self, name: str, **
|
|
153
|
-
super().__init__(name, None, **
|
|
258
|
+
def __init__(self, name: str, **kwargs: Any):
|
|
259
|
+
super().__init__(name, None, **kwargs)
|
|
154
260
|
|
|
155
261
|
|
|
156
|
-
def start_span(
|
|
262
|
+
def start_span(
|
|
263
|
+
name: str,
|
|
264
|
+
*,
|
|
265
|
+
print_fn: Callable[[str], None] | None = None,
|
|
266
|
+
plugins: Sequence[Any] = (),
|
|
267
|
+
**attributes: Any,
|
|
268
|
+
) -> Span:
|
|
157
269
|
"""Begin a span that you will ``end()`` yourself, possibly in another scope."""
|
|
158
270
|
from .run import current_run
|
|
159
271
|
|
|
160
272
|
run = current_run()
|
|
161
273
|
store = run.history if run is not None else None
|
|
162
|
-
|
|
274
|
+
if store is not None:
|
|
275
|
+
# The span's own arguments win over the run-wide defaults
|
|
276
|
+
print_fn = print_fn if print_fn is not None else store.options.span_print_fn
|
|
277
|
+
plugins = plugins or store.options.span_plugins
|
|
278
|
+
options = {"print_fn": print_fn, "plugins": plugins, **attributes}
|
|
279
|
+
span = (
|
|
280
|
+
Span(name, store, **options)
|
|
281
|
+
if store is not None
|
|
282
|
+
else _NullSpan(name, **options)
|
|
283
|
+
)
|
|
163
284
|
return span.begin()
|
|
164
285
|
|
|
165
286
|
|
|
@@ -182,9 +303,16 @@ class span:
|
|
|
182
303
|
``forward/attention/duration_ms`` and window functions work on it directly.
|
|
183
304
|
"""
|
|
184
305
|
|
|
185
|
-
def __init__(
|
|
306
|
+
def __init__(
|
|
307
|
+
self,
|
|
308
|
+
name: str,
|
|
309
|
+
*,
|
|
310
|
+
print_fn: Callable[[str], None] | None = None,
|
|
311
|
+
plugins: Sequence[Any] = (),
|
|
312
|
+
**attributes: Any,
|
|
313
|
+
):
|
|
186
314
|
self._name = name
|
|
187
|
-
self._attributes = attributes
|
|
315
|
+
self._attributes = {"print_fn": print_fn, "plugins": plugins, **attributes}
|
|
188
316
|
self._span: Span | None = None
|
|
189
317
|
|
|
190
318
|
def __enter__(self) -> Span:
|
expr_tracker/trace.py
CHANGED
|
@@ -149,6 +149,7 @@ def build_trace(
|
|
|
149
149
|
start = float(record["start"])
|
|
150
150
|
origin = start if origin is None else min(origin, start)
|
|
151
151
|
args = dict(record.get("args") or {})
|
|
152
|
+
args.update(record.get("metrics") or {})
|
|
152
153
|
args["step"] = record.get("_step")
|
|
153
154
|
if record.get("error"):
|
|
154
155
|
args["error"] = record["error"]
|
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
2
|
Name: expr_tracker
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.4
|
|
4
4
|
Summary: Local-first experiment tracking with queryable history and expression-based alerts on your training metrics
|
|
5
5
|
Project-URL: Homepage, https://hspk.github.io/expr_tracker/
|
|
6
6
|
Project-URL: Documentation, https://hspk.github.io/expr_tracker/
|
|
@@ -28,11 +28,14 @@ Requires-Dist: click>=8.1.0
|
|
|
28
28
|
Requires-Dist: loguru>=0.7.3
|
|
29
29
|
Requires-Dist: pydantic>=2.0
|
|
30
30
|
Provides-Extra: all
|
|
31
|
+
Requires-Dist: nvidia-ml-py>=12.0; extra == 'all'
|
|
31
32
|
Requires-Dist: pandas>=1.5; extra == 'all'
|
|
32
33
|
Requires-Dist: polars>=0.20; extra == 'all'
|
|
33
34
|
Requires-Dist: slark>=0.1.28; extra == 'all'
|
|
34
35
|
Requires-Dist: trackio>=0.4.0; extra == 'all'
|
|
35
36
|
Requires-Dist: wandb>=0.21.0; extra == 'all'
|
|
37
|
+
Provides-Extra: gpu
|
|
38
|
+
Requires-Dist: nvidia-ml-py>=12.0; extra == 'gpu'
|
|
36
39
|
Provides-Extra: lark
|
|
37
40
|
Requires-Dist: slark>=0.1.28; extra == 'lark'
|
|
38
41
|
Provides-Extra: pandas
|
|
@@ -106,7 +109,7 @@ install command; it never crashes a run.
|
|
|
106
109
|
| [Alerts](https://hspk.github.io/expr_tracker/guide/alerts/) | An expression DSL with rolling windows, three-valued logic (no false alarms during warm-up), a rule state machine, and a watchdog that catches a hung run. |
|
|
107
110
|
| [Channels](https://hspk.github.io/expr_tracker/guide/alerts/#channels) | Lark, Slack, DingTalk, WeCom, generic webhook, email — with rate limiting, dedup, retries and per-channel routing. |
|
|
108
111
|
| [Artifacts](https://hspk.github.io/expr_tracker/guide/artifacts/) | Versioned file sets, deduplicated by content, shared across a project's runs, with lineage. |
|
|
109
|
-
| [Spans](https://hspk.github.io/expr_tracker/guide/spans/) | Time the parts of a step, and their parts. Each duration becomes a metric, so alerts and queries work on it unchanged; `et trace` exports the timeline for Perfetto. |
|
|
112
|
+
| [Spans](https://hspk.github.io/expr_tracker/guide/spans/) | Time the parts of a step, and their parts. Each duration becomes a metric, so alerts and queries work on it unchanged; `et trace` exports the timeline for Perfetto. `print_fn` prints the tree live, and plugins attach CPU and GPU cost to each region. |
|
|
110
113
|
| [Streams](https://hspk.github.io/expr_tracker/guide/streams/) | Independent producers — a data worker and a training loop — each with their own step cursor and file inside one run. |
|
|
111
114
|
| [Distributed](https://hspk.github.io/expr_tracker/guide/distributed/) | Per-rank shards so concurrent appends cannot corrupt step order; only rank 0 alerts by default. |
|
|
112
115
|
| [CLI](https://hspk.github.io/expr_tracker/guide/cli/) | `et history`, `et rules explain`, `et rules test`, `et alert`. |
|
|
@@ -152,6 +155,7 @@ uv run ruff format src tests
|
|
|
152
155
|
| `test_expr_properties.py` | DSL properties: render round-trip stability, precedence, the whole `M` builder |
|
|
153
156
|
| `test_trace.py` | Chrome Trace export: lane layout, stream and step selection, the CLI |
|
|
154
157
|
| `test_spans.py` | nesting, aggregation, decorator and async forms, errors, thread and task isolation |
|
|
158
|
+
| `test_span_plugins.py` | `print_fn` output and indentation, the plugin protocol, failure isolation, CPU/GPU built-ins |
|
|
155
159
|
| `test_streams.py` | stream naming and validation, isolation, resolution order, backend grouping, two-process runs |
|
|
156
160
|
| `test_distributed.py` | rank shards, `alert_on_rank`, real multi-process runs |
|
|
157
161
|
| `test_wandb.py` | real wandb in offline mode: parameter mapping, step alignment, artifacts |
|
|
@@ -1,13 +1,14 @@
|
|
|
1
|
-
expr_tracker/__init__.py,sha256=
|
|
1
|
+
expr_tracker/__init__.py,sha256=AdY6R3bKLlZFhNtJPBAA3UHhO9nZ7YpojLrkW9OevlM,1442
|
|
2
2
|
expr_tracker/_compat.py,sha256=qf8oYIK_lkoLzcPp6wPCYLJtVatCl7ROYRNPwc5F1Pg,981
|
|
3
3
|
expr_tracker/artifacts.py,sha256=yRVXmuqdZUxiZa0c_CwZi4rBxaItwhu9bZ1x21yP1S0,12560
|
|
4
4
|
expr_tracker/cli.py,sha256=dtp0YsNtGg7T7M-elA511YX8NU_-OJRAwNbs8I1353o,7255
|
|
5
5
|
expr_tracker/encoders.py,sha256=X5ySQasSrYPsdYCdk7nhpJFGFB9f_8HWQuGOQjTvh9E,11991
|
|
6
|
+
expr_tracker/plugins.py,sha256=-ako71QBr5ARYS3aFUueHx3SMESk5ii2wj_yiz2_Mzc,9592
|
|
6
7
|
expr_tracker/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
7
8
|
expr_tracker/run.py,sha256=nuXhNmepG688X86cI0FMM_hqjyygD1Oi8Ck8-lyHsPA,16664
|
|
8
|
-
expr_tracker/spans.py,sha256=
|
|
9
|
+
expr_tracker/spans.py,sha256=sYj7BoFuNXEd4Ic11W_cqhLEjkh3ViBfTnW4vOvTTkg,12144
|
|
9
10
|
expr_tracker/summary.py,sha256=lpHSh3pci2Ym6wJhVZwNUJbUp2A64EN_rC7Ym2d7PPI,3150
|
|
10
|
-
expr_tracker/trace.py,sha256=
|
|
11
|
+
expr_tracker/trace.py,sha256=stOIuEzuFLHZrrqbaTWndOOnhSXEJzHF91zPGDpU5PA,7019
|
|
11
12
|
expr_tracker/tracker.py,sha256=hpZB8bFA2GqCbvHEDZkBS9cCzi2JgXeLUNTzJm-wr8c,5921
|
|
12
13
|
expr_tracker/types.py,sha256=69gFtUaKGl4asLK7nJci9XyTyV0yTIPorB-bcfGrCBA,383
|
|
13
14
|
expr_tracker/alerts/__init__.py,sha256=dl96OYZh75kV3NLBSA35zNvTFcL6XFMuPoYvRoz7hlQ,8192
|
|
@@ -29,10 +30,10 @@ expr_tracker/history/frame.py,sha256=PFDGqR0oHY94KosSXA7Fu-bphEhZ2yM8cXW9rlZ2IC0
|
|
|
29
30
|
expr_tracker/history/naming.py,sha256=CT2xDpo0GHYnLyjgRlsPjIX3GAIDEdn4U9I1-cNahHw,2507
|
|
30
31
|
expr_tracker/history/reader.py,sha256=5PY5nbNfX-Cl34jcyYTq2qm8nNgNDTqv7U5Xnb9Rgxw,14155
|
|
31
32
|
expr_tracker/history/series.py,sha256=cFdUO1FRs5KW0jYsNPzda3tcu645NH1wwfXJitQfXFM,4042
|
|
32
|
-
expr_tracker/history/store.py,sha256=
|
|
33
|
+
expr_tracker/history/store.py,sha256=uSbHcGdLXFVFbYQKm6IaQgpdCrNvvb9Axpxi9emvwBo,36228
|
|
33
34
|
expr_tracker/history/writer.py,sha256=q5kNANcZCq5nykvOK5jZ2xc6Z1B4oZxEn3uSXzDCtlY,16096
|
|
34
|
-
expr_tracker-0.2.
|
|
35
|
-
expr_tracker-0.2.
|
|
36
|
-
expr_tracker-0.2.
|
|
37
|
-
expr_tracker-0.2.
|
|
38
|
-
expr_tracker-0.2.
|
|
35
|
+
expr_tracker-0.2.4.dist-info/METADATA,sha256=aSjFfY1k1dJwzBRHWvm3aqisY6O0Gi6S7HvMAp37Fs4,9447
|
|
36
|
+
expr_tracker-0.2.4.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
|
|
37
|
+
expr_tracker-0.2.4.dist-info/entry_points.txt,sha256=GfZCAG2JFsd5CxTjRsGJAhYInpBCuLjrU3csmUVBDSs,45
|
|
38
|
+
expr_tracker-0.2.4.dist-info/licenses/LICENSE,sha256=ORsYTnF-2SLUaM8FgXlJ-uXuN4ceGACAAU3JZiS8pvs,1069
|
|
39
|
+
expr_tracker-0.2.4.dist-info/RECORD,,
|
|
File without changes
|
|
File without changes
|