expr-tracker 0.2.3__tar.gz → 0.2.4__tar.gz

{expr_tracker-0.2.3 → expr_tracker-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: expr_tracker
-Version: 0.2.3
+Version: 0.2.4
 Summary: Local-first experiment tracking with queryable history and expression-based alerts on your training metrics
 Project-URL: Homepage, https://hspk.github.io/expr_tracker/
 Project-URL: Documentation, https://hspk.github.io/expr_tracker/
@@ -28,11 +28,14 @@ Requires-Dist: click>=8.1.0
 Requires-Dist: loguru>=0.7.3
 Requires-Dist: pydantic>=2.0
 Provides-Extra: all
+Requires-Dist: nvidia-ml-py>=12.0; extra == 'all'
 Requires-Dist: pandas>=1.5; extra == 'all'
 Requires-Dist: polars>=0.20; extra == 'all'
 Requires-Dist: slark>=0.1.28; extra == 'all'
 Requires-Dist: trackio>=0.4.0; extra == 'all'
 Requires-Dist: wandb>=0.21.0; extra == 'all'
+Provides-Extra: gpu
+Requires-Dist: nvidia-ml-py>=12.0; extra == 'gpu'
 Provides-Extra: lark
 Requires-Dist: slark>=0.1.28; extra == 'lark'
 Provides-Extra: pandas
@@ -106,7 +109,7 @@ install command; it never crashes a run.
 | [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. |
 | [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. |
 | [Artifacts](https://hspk.github.io/expr_tracker/guide/artifacts/) | Versioned file sets, deduplicated by content, shared across a project's runs, with lineage. |
-| [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. |
+| [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. |
 | [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. |
 | [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. |
 | [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
 | `test_expr_properties.py` | DSL properties: render round-trip stability, precedence, the whole `M` builder |
 | `test_trace.py` | Chrome Trace export: lane layout, stream and step selection, the CLI |
 | `test_spans.py` | nesting, aggregation, decorator and async forms, errors, thread and task isolation |
+| `test_span_plugins.py` | `print_fn` output and indentation, the plugin protocol, failure isolation, CPU/GPU built-ins |
 | `test_streams.py` | stream naming and validation, isolation, resolution order, backend grouping, two-process runs |
 | `test_distributed.py` | rank shards, `alert_on_rank`, real multi-process runs |
 | `test_wandb.py` | real wandb in offline mode: parameter mapping, step alignment, artifacts |

{expr_tracker-0.2.3 → expr_tracker-0.2.4}/README.md

@@ -59,7 +59,7 @@ install command; it never crashes a run.
 | [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. |
 | [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. |
 | [Artifacts](https://hspk.github.io/expr_tracker/guide/artifacts/) | Versioned file sets, deduplicated by content, shared across a project's runs, with lineage. |
-| [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. |
+| [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. |
 | [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. |
 | [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. |
 | [CLI](https://hspk.github.io/expr_tracker/guide/cli/) | `et history`, `et rules explain`, `et rules test`, `et alert`. |
@@ -105,6 +105,7 @@ uv run ruff format src tests
 | `test_expr_properties.py` | DSL properties: render round-trip stability, precedence, the whole `M` builder |
 | `test_trace.py` | Chrome Trace export: lane layout, stream and step selection, the CLI |
 | `test_spans.py` | nesting, aggregation, decorator and async forms, errors, thread and task isolation |
+| `test_span_plugins.py` | `print_fn` output and indentation, the plugin protocol, failure isolation, CPU/GPU built-ins |
 | `test_streams.py` | stream naming and validation, isolation, resolution order, backend grouping, two-process runs |
 | `test_distributed.py` | rank shards, `alert_on_rank`, real multi-process runs |
 | `test_wandb.py` | real wandb in offline mode: parameter mapping, step alignment, artifacts |

{expr_tracker-0.2.3 → expr_tracker-0.2.4}/docs/design.md

@@ -111,11 +111,38 @@ so timing a region adds no row. The tree, with timestamps and attributes, goes t
 `spans[.stream][.rankN].jsonl` through a second `JsonlWriter`, enqueued rather
 than appended so a span does not pay for a flush decision of its own.
 
+Plugins hang off the same two hooks. `start` runs before the clock starts and
+`end` after it stops, so a plugin never inflates the duration it reports on; what
+`end` returns is merged into the span's metrics under the span's own path, which
+means a plugin needs no knowledge of the store, the codec or the accumulate rule.
+A plugin that raises is logged and dropped: a measurement must not break the
+thing it measures.
+
+Peak-memory plugins have to fight the fact that `torch.cuda.max_memory_allocated`
+is one global high-water mark. Resetting it per span is the only way to attribute
+a peak, but a child's reset destroys its parent's history. `TorchMemory` keeps a
+per-thread stack and, as each span closes, hands its parent back the peak the
+parent had reached before that child opened, so `max(counter, carried)` is
+correct at every level.
+
+GPU utilisation cannot be attributed by reading NVML at a span boundary: the
+counter is a rolling average over NVML's own window, so for a 3 ms span it
+describes the wrong interval entirely. `GpuStats` samples in a background thread
+and each span averages the samples that fall inside it, which also keeps the
+NVML call off the measured path.
+
 The nesting stack is a `ContextVar`, not a `threading.local`: a new thread starts
 from the default and each asyncio task gets its own copy, which is exactly the
 scoping spans need. (The run singleton went the other way, from `ContextVar` to a
 locked global, because there the goal was for worker threads to *share* it.)
 
+Closing a span filters itself out of the stack rather than resetting the token
+it got when it opened. A token restores the entire stack as it was, so with
+`start_span` — where spans need not close in the order they opened — an outer
+span closing first would be put back by its inner span's reset, and everything
+afterwards would nest under a span that had already ended. Filtering is exact,
+and at ten levels deep it still does not show up against the ~15 µs a span costs.
+
 ### A.6 Multiple processes
 
 The supported model is **rank 0 tracks**; there is no cross-rank merging.

expr_tracker-0.2.4/docs/guide/spans.md

@@ -0,0 +1,289 @@
+# Spans
+
+A step is rarely one thing. `et.span` times the parts, and their parts, and turns
+each duration into an ordinary metric — so `history()`, alert rules and plots all
+work on it with no extra machinery.
+
+```python
+for batch in loader:
+    with et.span("forward"):
+        with et.span("attention"):
+            attn = attention(batch)
+        with et.span("mlp"):
+            out = mlp(attn)
+    with et.span("backward"):
+        loss.backward()
+    et.log({"loss": loss.item()})
+```
+
+The step's row then carries:
+
+```jsonc
+{"_step": 42, "_time": ...,
+ "forward/duration_ms": 31.2,  "forward/count": 1,
+ "forward/attention/duration_ms": 18.4, "forward/attention/count": 1,
+ "forward/mlp/duration_ms": 9.1, "forward/mlp/count": 1,
+ "backward/duration_ms": 44.7, "backward/count": 1,
+ "loss": 0.31}
+```
+
+Nested names join with `/`, so `forward/norm` and `backward/norm` stay distinct,
+and the alert language reads them directly.
+
+## Forms
+
+```python
+with et.span("forward"): ...                # context manager
+
+async with et.span("fetch"): ...            # async
+
+@et.span("preprocess")                      # decorator, sync or async
+def preprocess(batch): ...
+
+span = et.start_span("epoch")               # manual, for crossing scopes
+...
+span.end()
+```
+
+## Repeated spans
+
+The same span several times in one step sums, and counts:
+
+```python
+for layer in layers:            # 32 layers
+    with et.span("layer"):
+        x = layer(x)
+```
+
+```jsonc
+{"layer/duration_ms": 412.8, "layer/count": 32}
+```
+
+The total is usually what you want; divide by the count for the mean.
+
+## A span never commits a step
+
+Durations join whatever `log()` commits, so a span costs no row of its own and
+you can time things before you know what to log:
+
+```python
+with et.span("forward"):
+    ...
+# nothing written yet
+et.log({"loss": loss})     # one row, with the metrics and the durations
+```
+
+## Attributes
+
+Attributes describe an individual span. They reach `spans.jsonl`, not the metrics,
+because they are usually not numbers:
+
+```python
+with et.span("load", batch_size=32) as span:
+    rows = read()
+    span.set(rows=len(rows))
+```
+
+## Alerting on durations
+
+Because a duration is a metric, the [expression language](../reference/expressions.md)
+already handles it:
+
+```python
+et.init(..., alert_rules=[
+    "mean(data/load_ms[50]) > 200 => warning: data loading is slowing down",
+    "forward/duration_ms > 3 * mean(forward/duration_ms[100]) => error: slow step",
+])
+```
+
+## The span file
+
+The full tree is written to `spans.jsonl` beside the metrics, one line per span:
+
+```jsonc
+{"_step": 42, "name": "forward/attention", "depth": 1,
+ "start": 1754323200.123456, "dur_ms": 18.4, "args": {"batch_size": 32}}
+```
+
+Children appear before their parents, because a parent closes last. It follows
+the [stream](streams.md): a data worker writes `spans.data.jsonl`.
+
+Turn it off if you only want the metrics:
+
+```python
+et.init(..., spans=False)
+```
+
+## Viewing the timeline
+
+`et trace` turns the span file into a Chrome Trace, which
+[Perfetto](https://ui.perfetto.dev) and `chrome://tracing` open directly:
+
+```bash
+et trace runs/llm/sft-1 -o trace.json
+et trace runs/llm/sft-1 --stream data --step-range 100:200
+```
+
+Each stream becomes a process and each thread a track, so a data worker and a
+training loop sit on one timeline and the gap where one waited for the other is
+visible. Spans keep their nesting, their step and their attributes.
+
+Exporting a standard format rather than drawing our own view means the result
+can be loaded beside a `torch.profiler` trace, which is usually where the real
+question is: what were the GPUs doing while the loader stalled.
+
+## Errors
+
+An exception is recorded and re-raised — the span never swallows it:
+
+```jsonc
+{"_step": 42, "name": "risky", "dur_ms": 3.1, "error": "ValueError"}
+```
+
+A failure inside the recording itself is logged and ignored: measuring something
+must not break it.
+
+## Threads and tasks
+
+The nesting stack is per thread and per asyncio task, so concurrent work does not
+nest inside unrelated spans:
+
+```python
+async def work(name):
+    async with et.span(name):        # "a" and "b", never "a/b"
+        await asyncio.sleep(1)
+
+await asyncio.gather(work("a"), work("b"))
+```
+
+## Printing a span as it runs
+
+Pass `print_fn` and the span announces itself, indented one tab per level:
+
+```python
+with et.span("step", print_fn=print):
+    with et.span("forward"):
+        with et.span("attention"):
+            ...
+```
+
+```
+-> step  16:41:38
+	-> forward  16:41:38
+		-> attention  16:41:38
+		<- attention  16:41:38  3.074ms
+	<- forward  16:41:38  5.481ms
+<- step  16:41:38  9.641ms
+```
+
+Children inherit the handler, so one argument on the outermost span prints the
+whole tree. A child that passes its own `print_fn` takes over its subtree; pass
+`print_fn=lambda line: None` to silence one. A failed span ends with `!ValueError`.
+
+Indentation is relative to the span that started printing, not to absolute
+nesting depth, so turning printing on deep inside a call stack still gives you a
+tree rooted at the left margin. A new thread begins its own tree, because the
+nesting stack does not cross threads. Concurrent asyncio tasks share one handler
+and interleave their lines; each is still indented by its own depth, and the
+`->` and `<-` markers pair them up.
+
+Any callable taking one string works — `print`, `logger.info`, a list's `append`
+in tests. It is called on the thread that ran the span, and an exception in it is
+logged and swallowed.
+
+Set it for the whole run instead of per span with `span_print_fn`:
+
+```python
+et.init(project="demo", span_print_fn=logger.info)
+```
+
+## Plugins
+
+A plugin measures a resource across a span. It is any object with
+`start(span)` and `end(span) -> dict`; a plain `fn(span) -> dict` is an
+end-only plugin. Whatever `end` returns is merged into the span's metrics under
+the span's own path:
+
+```python
+from expr_tracker.plugins import CpuTime, GpuStats, TorchMemory
+
+with et.span("forward", plugins=[CpuTime(), TorchMemory()]):
+    loss = model(batch)
+```
+
+```python
+et.history(1)[0]
+# {"forward/duration_ms": 41.2, "forward/count": 1,
+#  "forward/cpu_percent": 101.2, "forward/cpu_time_ms": 41.4,
+#  "forward/gpu_mem_peak_mb": 456.1, "forward/gpu_mem_delta_mb": 8.1, ...}
+```
+
+Like `print_fn`, plugins are inherited by children, can be overridden per span,
+and can be set run-wide with `span_plugins`. They are opt-in because they are not
+free. A plugin that raises is logged and skipped; it never breaks the span or the
+code being measured, and a plugin that fails on every span is logged only once.
+Plugin metrics also reach `spans.jsonl` and show up in the `et trace` viewer
+under a span's *Arguments*.
+
+`duration_ms` and `count` belong to the span. A plugin returning either has that
+key dropped with a warning, so nothing can quietly redefine what a span measured.
+
+### Built-ins
+
+| plugin | metrics | needs |
+| --- | --- | --- |
+| `CpuTime()` | `cpu_time_ms`, `cpu_percent` | nothing |
+| `TorchMemory(device=None)` | `gpu_mem_peak_mb`, `gpu_mem_delta_mb` | `torch` + CUDA |
+| `GpuStats(index=0, interval=0.1)` | `gpu_percent`, `gpu_mem_used_mb` | `pip install expr-tracker[gpu]` |
+
+`CpuTime` divides process CPU time by wall time, so `cpu_percent` above 100 means
+the span used more than one core. It counts every thread in the process, so
+unrelated concurrent work inflates it, and it is noise on spans under a
+millisecond.
+
+`TorchMemory` reports what your own allocator held, which is the number that
+predicts an OOM. `gpu_mem_peak_mb` is absolute, so memory a tensor was already
+holding when the span opened counts towards it; `gpu_mem_delta_mb` is what the
+span kept. `torch.cuda.max_memory_allocated` is a single global high-water mark,
+so a child span resetting it would erase its parent's history — the plugin keeps
+its own stack and hands each parent back the peak it had reached before the child
+opened.
+
+`GpuStats` reads the whole device, including other processes. A background thread
+samples NVML every `interval`, and a span reports the mean utilisation and the
+peak memory over its own window, so the span itself pays almost nothing. Spans
+shorter than `interval` contain no sample and fall back to one live reading,
+which for utilisation is NVML's own rolling average rather than a measurement of
+that span. Samplers are shared across plugins with the same device and interval.
+
+### CUDA is asynchronous
+
+A span measures the wall time of the Python block, and CUDA kernels are queued,
+not awaited. A span around a launch can close in microseconds while the GPU is
+still busy, and the wait then lands on whatever later line synchronises. If you
+want a span to mean GPU time, synchronise inside it:
+
+```python
+with et.span("forward", plugins=[TorchMemory()]):
+    out = model(batch)
+    torch.cuda.synchronize()
+```
+
+`gpu_mem_peak_mb` is unaffected: allocator bookkeeping is recorded at launch.
+
+## Cost
+
+| | per span |
+| --- | --- |
+| default | ~15 µs |
+| `spans=False` | ~6 µs |
+| no active run | ~3 µs |
+| `print_fn` | +6 µs |
+| `CpuTime()` | +8 µs |
+| `GpuStats()` | +4 µs, or +36 µs on a span shorter than `interval` |
+| `TorchMemory()` | +62 µs |
+
+For comparison, `et.log()` is ~24 µs. Twenty spans on a 100 ms step is 0.3% of
+the step. If your step is closer to a millisecond, set `spans=False` and keep the
+metrics, or time fewer regions. Put `TorchMemory` on the few spans whose memory
+you actually care about rather than run-wide.

{expr_tracker-0.2.3 → expr_tracker-0.2.4}/docs/reference/api.md

@@ -156,14 +156,35 @@ Registers a custom channel type.
 ## Spans
 
 ```python
-with et.span(name, **attributes) as span: ...   # also async, also a decorator
-span = et.start_span(name, **attributes)        # ends with span.end()
+with et.span(name, print_fn=None, plugins=(), **attributes) as span: ...
+span = et.start_span(name, print_fn=None, plugins=(), **attributes)
 span.set(**attributes)
 span.duration_ms
+span.metrics                                    # what the plugins measured
 ```
 
-A closed span adds `<path>/duration_ms` and `<path>/count` to the open row and
-appends the full record to `spans.jsonl`. See [Spans](../guide/spans.md).
+`et.span` is also an async context manager and a decorator. A closed span adds
+`<path>/duration_ms`, `<path>/count` and one key per plugin metric to the open
+row, and appends the full record to `spans.jsonl`.
+
+`print_fn(line)` announces the span's start and end, indented one tab per level.
+`plugins` measure a resource across the span; both are inherited by child spans
+and default to the run's `span_print_fn` and `span_plugins`.
+
+### Plugins
+
+```python
+from expr_tracker.plugins import CpuTime, GpuStats, TorchMemory
+```
+
+| plugin | metrics |
+| --- | --- |
+| `CpuTime()` | `cpu_time_ms`, `cpu_percent` |
+| `TorchMemory(device=None)` | `gpu_mem_peak_mb`, `gpu_mem_delta_mb` |
+| `GpuStats(index=0, interval=0.1)` | `gpu_percent`, `gpu_mem_used_mb` |
+
+Any object with `start(span)` and `end(span) -> dict` is a plugin, as is a plain
+`fn(span) -> dict`. See [Spans](../guide/spans.md).
 
 ### Trace export
 

{expr_tracker-0.2.3 → expr_tracker-0.2.4}/docs/reference/configuration.md

@@ -15,6 +15,8 @@ list of valid options rather than being silently ignored.
 | `rank_aware` | `True` | non-zero ranks write their own shard |
 | `stream` | `None` | an independent producer with its own file and step cursor |
 | `spans` | `True` | write the span tree to `spans.jsonl` beside the metrics |
+| `span_print_fn` | `None` | default `print_fn` for every span in the run |
+| `span_plugins` | `()` | default plugins for every span in the run |
 | `print_to_screen` | `False` | print every committed row |
 | `print_handle` | `print` | where those lines go |
 | `buffer_size` | 50 | flush after this many buffered rows |

{expr_tracker-0.2.3 → expr_tracker-0.2.4}/pyproject.toml

@@ -45,12 +45,14 @@ wandb = ["wandb>=0.21.0"]
 trackio = ["trackio>=0.4.0"]
 pandas = ["pandas>=1.5"]
 polars = ["polars>=0.20"]
+gpu = ["nvidia-ml-py>=12.0"]
 all = [
     "slark>=0.1.28",
     "wandb>=0.21.0",
     "trackio>=0.4.0",
     "pandas>=1.5",
     "polars>=0.20",
+    "nvidia-ml-py>=12.0",
 ]
 
 [project.urls]

{expr_tracker-0.2.3 → expr_tracker-0.2.4}/src/expr_tracker/__init__.py

@@ -1,7 +1,7 @@
 from importlib.metadata import PackageNotFoundError
 from importlib.metadata import version as _version
 
-from . import tracker
+from . import plugins, tracker
 from .alerts import (
     AlertConfig,
     AlertLevel,
@@ -61,6 +61,7 @@ __all__ = [
     "list_alert_rules",
     "log",
     "log_artifact",
+    "plugins",
     "register_backend",
     "remove_alert_rule",
     "span",

{expr_tracker-0.2.3 → expr_tracker-0.2.4}/src/expr_tracker/history/store.py

@@ -19,6 +19,7 @@ from collections import deque
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass, fields
 from pathlib import Path
+from typing import Any
 
 from loguru import logger
 
@@ -74,6 +75,9 @@ class HistoryOptions:
     stream: str | None = None
     # write the full span tree to spans.jsonl beside the metrics
     spans: bool = True
+    # run-wide defaults for et.span(); a span's own arguments win
+    span_print_fn: Callable[[str], None] | None = None
+    span_plugins: Sequence[Any] = ()
     # alerts and output
     alert_window: int = DEFAULT_WINDOW
     print_to_screen: bool = False