expr_tracker 0.2.2__py3-none-any.whl → 0.2.4__py3-none-any.whl

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,
@@ -19,6 +19,7 @@ from .alerts import (
 )
 from .artifacts import Artifact
 from .run import Run
+from .spans import Span, span, start_span
 
 try:
     __version__ = _version("expr_tracker")
@@ -45,6 +46,7 @@ __all__ = [
     "ChannelConfig",
     "M",
     "Run",
+    "Span",
     "WebhookPolicy",
     "__version__",
     "add_alert_rule",
@@ -59,8 +61,11 @@ __all__ = [
     "list_alert_rules",
     "log",
     "log_artifact",
+    "plugins",
     "register_backend",
     "remove_alert_rule",
+    "span",
+    "start_span",
     "summary",
     "use_artifact",
 ]

expr_tracker/cli.py

@@ -13,6 +13,7 @@ from .alerts.engine import AlertEngine
 from .alerts.expr import EvalContext, compile_condition, parse_rule, validate
 from .alerts.models import AlertConfig, ChannelConfig, WebhookPolicy
 from .history import MetricSeries, read_history
+from .trace import write_trace
 
 
 @click.group()
@@ -54,6 +55,29 @@ def history(run: str, n: int, metrics: str | None, step_range: str | None, fmt:
         click.echo(_to_table(rows))
 
 
+@main.command()
+@click.argument("run", type=click.Path(exists=True))
+@click.option("-o", "--output", default="trace.json", help="Where to write the trace")
+@click.option(
+    "--stream",
+    default=None,
+    help="Comma separated streams; every stream by default. Use 'default' for the "
+    "unnamed one.",
+)
+@click.option("--step-range", default=None, help="start:end (end exclusive)")
+def trace(run: str, output: str, stream: str | None, step_range: str | None):
+    """Export recorded spans as a Chrome Trace.
+
+    Open the result at https://ui.perfetto.dev or chrome://tracing. It can be
+    loaded beside a torch.profiler trace.
+    """
+    wanted: str | list[str] | None = "*"
+    if stream:
+        wanted = [None if s in ("default", "") else s for s in stream.split(",")]
+    count = write_trace(run, output, wanted, step_range=_parse_range(step_range))
+    click.echo(f"wrote {count} span(s) to {output}")
+
+
 @main.group()
 def rules():
     """Alert rule tooling."""

expr_tracker/history/__init__.py

@@ -1,8 +1,15 @@
 from .codec import RecordCodec
 from .frame import project, to_output
-from .reader import JsonlReader, merge_steps, read_history, resolve_run_path
+from .naming import current_rank, metrics_filename, parse_stream, validate_stream
+from .reader import (
+    JsonlReader,
+    list_streams,
+    merge_steps,
+    read_history,
+    resolve_run_path,
+)
 from .series import MetricSeries
-from .store import HistoryStore, current_rank, resolve_commit
+from .store import HistoryStore, resolve_commit
 from .writer import JsonlWriter
 
 __all__ = [
@@ -12,10 +19,14 @@ __all__ = [
     "MetricSeries",
     "RecordCodec",
     "current_rank",
+    "list_streams",
     "merge_steps",
+    "metrics_filename",
+    "parse_stream",
     "project",
     "read_history",
     "resolve_commit",
     "resolve_run_path",
     "to_output",
+    "validate_stream",
 ]

expr_tracker/history/naming.py

@@ -0,0 +1,74 @@
+"""Run file naming: distributed rank shards and independent producer streams.
+
+Kept separate from the store and the reader because both need it, and the store
+already depends on the reader.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+
+STREAM_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_-]*$")
+RANK_PATTERN = re.compile(r"rank\d+")
+
+
+def current_rank() -> int:
+    """Distributed rank from the usual environment variables, 0 when unset."""
+    rank = os.getenv("RANK") or os.getenv("LOCAL_RANK") or "0"
+    try:
+        return int(rank)
+    except ValueError:
+        return 0
+
+
+def validate_stream(stream: str) -> str:
+    """Check a stream name is safe as a filename component.
+
+    The name becomes part of ``metrics.<stream>[.rankN].jsonl``, so it must not
+    contain a separator, and must not look like a rank shard.
+    """
+    if not STREAM_PATTERN.match(stream):
+        raise ValueError(
+            f"Invalid stream name {stream!r}; use letters, digits, '_' or '-', "
+            "starting with a letter or digit."
+        )
+    if RANK_PATTERN.fullmatch(stream):
+        raise ValueError(
+            f"Stream name {stream!r} collides with the rank shard naming; "
+            "pick another name."
+        )
+    return stream
+
+
+def metrics_filename(stream: str | None, rank_aware: bool = True) -> str:
+    """``metrics[.stream][.rankN].jsonl``.
+
+    Streams separate independent producers; non-zero ranks get their own shard so
+    concurrent appends cannot interleave. The two compose.
+    """
+    rank = current_rank() if rank_aware else 0
+    parts = ["metrics"]
+    if stream:
+        parts.append(stream)
+    if rank > 0:
+        parts.append(f"rank{rank}")
+    return ".".join(parts) + ".jsonl"
+
+
+def spans_filename(stream: str | None, rank_aware: bool = True) -> str:
+    """``spans[.stream][.rankN].jsonl``, alongside the metrics file."""
+    return metrics_filename(stream, rank_aware).replace("metrics", "spans", 1)
+
+
+def sidecar_filename(base: str, stream: str | None, suffix: str) -> str:
+    """``<base>[.stream].<suffix>``, so producers do not overwrite each other."""
+    return f"{base}.{stream}.{suffix}" if stream else f"{base}.{suffix}"
+
+
+def parse_stream(filename: str) -> str | None:
+    """The stream a metrics filename belongs to; ``None`` is the default producer."""
+    parts = filename.split(".")[1:-1]  # drop "metrics" and "jsonl"
+    if parts and RANK_PATTERN.fullmatch(parts[-1]):
+        parts = parts[:-1]
+    return parts[0] if parts else None

expr_tracker/history/reader.py

@@ -14,6 +14,7 @@ from pathlib import Path
 from loguru import logger
 
 from .frame import project, to_output
+from .naming import RANK_PATTERN, metrics_filename
 
 CHUNK_SIZE = 256 * 1024
 
@@ -317,6 +318,7 @@ def read_history(
     run: str | Path,
     n: int | None = 50,
     *,
+    stream: str | None = None,
     output_type: str = "dict",
     metrics: Sequence[str] | None = None,
     step_range: tuple[int | None, int | None] | None = None,
@@ -325,7 +327,7 @@ def read_history(
     dropna: bool = False,
 ):
     """Read any run's history without ``init()``; ``run`` is a directory or file."""
-    reader = JsonlReader(resolve_run_path(run))
+    reader = JsonlReader(resolve_run_path(run, stream))
     limit = None if n is None or n < 0 else max(0, n)
     if limit == 0:
         records: list[dict] = []
@@ -347,13 +349,41 @@ def read_history(
     return to_output(rows, output_type)
 
 
-def resolve_run_path(run: str | Path) -> Path:
+def resolve_run_path(run: str | Path, stream: str | None = None) -> Path:
+    """The metrics file of a run directory, or the file itself.
+
+    Names are matched exactly rather than by sort order: ``metrics.data.jsonl``
+    sorts before ``metrics.jsonl``, so picking the first match would silently
+    return a stream instead of the default producer.
+    """
     path = Path(run)
-    if path.is_dir():
-        candidates = sorted(path.glob("metrics*.jsonl"))
-        if not candidates:
-            raise FileNotFoundError(f"No metrics*.jsonl found under {path}")
-        return candidates[0]
-    if not path.exists():
-        raise FileNotFoundError(f"Run path does not exist: {path}")
-    return path
+    if not path.is_dir():
+        if not path.exists():
+            raise FileNotFoundError(f"Run path does not exist: {path}")
+        return path
+    wanted = metrics_filename(stream, rank_aware=False)
+    exact = path / wanted
+    if exact.is_file():
+        return exact
+    # Fall back to this stream's rank shards, which a worker rank writes instead
+    prefix = wanted[: -len(".jsonl")]
+    shards = sorted(path.glob(f"{prefix}.rank*.jsonl"))
+    if shards:
+        return shards[0]
+    raise FileNotFoundError(f"No {wanted} found under {path}; {_available(path)}")
+
+
+def _available(path: Path) -> str:
+    found = sorted(p.name for p in path.glob("metrics*.jsonl"))
+    return f"available: {', '.join(found)}" if found else "the directory has none"
+
+
+def list_streams(run: str | Path) -> list[str | None]:
+    """Every stream in a run directory; ``None`` is the default producer."""
+    streams: set[str | None] = set()
+    for path in Path(run).glob("metrics*.jsonl"):
+        parts = path.name.split(".")[1:-1]  # drop "metrics" and "jsonl"
+        if parts and RANK_PATTERN.fullmatch(parts[-1]):
+            parts = parts[:-1]
+        streams.add(parts[0] if parts else None)
+    return sorted(streams, key=lambda s: (s is not None, s or ""))

expr_tracker/history/store.py

@@ -13,18 +13,24 @@ from __future__ import annotations
 import atexit
 import json
 import operator
-import os
 import threading
 import time
 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
 
 from .codec import RESERVED_KEYS, RecordCodec, encode_line
 from .frame import project, to_output
+from .naming import (
+    metrics_filename,
+    sidecar_filename,
+    spans_filename,
+    validate_stream,
+)
 from .reader import JsonlReader, merge_steps
 from .series import DEFAULT_WINDOW, MetricSeries
 from .writer import (
@@ -65,6 +71,13 @@ class HistoryOptions:
     max_open_seconds: float | None = DEFAULT_MAX_OPEN_SECONDS
     step_policy: str = "monotonic"
     rank_aware: bool = True
+    # an independent producer writing its own file with its own step cursor
+    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
@@ -76,6 +89,8 @@ class HistoryOptions:
                 f"Unknown step_policy {self.step_policy!r}; "
                 f"expected one of {', '.join(STEP_POLICIES)}."
             )
+        if self.stream is not None:
+            object.__setattr__(self, "stream", validate_stream(self.stream))
         clamped = {
             "cache_bytes": max(0, int(self.cache_bytes)),
             "cache_rows": max(1, int(self.cache_rows)),
@@ -128,6 +143,9 @@ class HistoryStore:
         self.log_dir: Path | None = None
         self.log_fp: Path | None = None
         self.config_fp: Path | None = None
+        self.spans_fp: Path | None = None
+        self.span_writer: JsonlWriter | None = None
+        self.stream: str | None = None
         self.project: str = ""
         self.name: str = ""
         self._reset(HistoryOptions(), on_commit=None)
@@ -177,12 +195,29 @@ class HistoryStore:
         self.project, self.name = project, name
         self.log_dir = Path(dir or "./tracker/jsonl") / project / name
         self.log_dir.mkdir(parents=True, exist_ok=True)
-        self.config_fp = self.log_dir / "config.json"
-        self.log_fp = self.log_dir / _metrics_filename(opts.rank_aware)
+        self.stream = opts.stream
+        self.config_fp = self.log_dir / sidecar_filename("config", opts.stream, "json")
+        self.spans_fp = (
+            self.log_dir / spans_filename(opts.stream, opts.rank_aware)
+            if opts.spans
+            else None
+        )
+        self.log_fp = self.log_dir / metrics_filename(opts.stream, opts.rank_aware)
 
         self._close_previous()
         self._reset(opts, on_commit)
         self.writer = self._open_writer(opts)
+        self.span_writer = (
+            JsonlWriter(
+                self.spans_fp,
+                buffer_size=opts.buffer_size,
+                buffer_interval=opts.buffer_interval,
+                max_buffer_seconds=opts.max_buffer_seconds,
+                max_pending_records=opts.max_pending_records,
+            )
+            if self.spans_fp is not None
+            else None
+        )
         self._resume(opts.alert_window)
         self._write_config(config)
         self._register_atexit()
@@ -191,8 +226,9 @@ class HistoryStore:
     def _close_previous(self):
         """Re-initialising must not strand the previous writer's buffer or timer."""
         self._cancel_open_timer()
-        if self.writer is not None:
-            self.writer.close()
+        for writer in (self.writer, self.span_writer):
+            if writer is not None:
+                writer.close()
 
     def _open_writer(self, options: HistoryOptions) -> JsonlWriter:
         assert self.log_fp is not None
@@ -276,7 +312,24 @@ class HistoryStore:
         ``is None``, since step 0 is falsy.
         """
         self._require_init()
-        encoded = self._codec.encode(metrics)
+        return self.ingest(self._codec.encode(metrics), step=step, commit=commit)
+
+    def ingest(
+        self,
+        encoded: dict,
+        *,
+        step: int | None = None,
+        commit: bool | None = None,
+        accumulate: bool = False,
+    ) -> int | None:
+        """The single path metrics take into the open row.
+
+        ``log()`` and span durations both come through here, so the step policy,
+        the merge and the commit rules only exist once. ``accumulate`` adds to any
+        existing value instead of replacing it, which is what repeated spans in
+        one step need.
+        """
+        self._require_init()
         now = time.time()
         pending: list[dict] = []
         with self._lock:
@@ -287,7 +340,7 @@ class HistoryStore:
                 return None
             pending.append(self._switch_open_row(step, now))
             resolved = self._open_step
-            self._update_open_row(encoded)
+            self._update_open_row(encoded, accumulate)
             if resolve_commit(step, commit):
                 pending.append(self._close_open_row())
         for record in filter(None, pending):
@@ -295,6 +348,28 @@ class HistoryStore:
         self._schedule_open_timer()
         return resolved
 
+    def record_span(self, metrics: dict, record: dict | None = None) -> int | None:
+        """A finished span: durations join the open row, the tree goes to disk.
+
+        Durations never commit a step of their own -- they ride along with
+        whatever ``log()`` commits, so a span costs no extra row. The metrics are
+        built from floats and ints here, so they skip the encoder: spans are
+        recorded per sub-step and cannot afford it.
+        """
+        step = self.ingest(metrics, commit=False, accumulate=True)
+        if record is None or self.span_writer is None or step is None:
+            return step
+        try:
+            line = encode_line({"_step": step, **record})
+        except Exception as e:
+            logger.warning(f"Dropping span that cannot be serialized: {e}")
+            return step
+        # enqueue, not append: the batch goes out when the step commits, so a
+        # span does not pay for a flush decision of its own
+        if self.span_writer.enqueue(step, self.span_writer.lines, line):
+            self.span_writer.flush()
+        return step
+
     def _switch_open_row(self, step: int | None, now: float) -> dict | None:
         """Point the open row at ``step``, returning the row it displaced, if any.
 
@@ -331,12 +406,18 @@ class HistoryStore:
         if timer is not None:
             timer.cancel()
 
-    def _update_open_row(self, encoded: dict):
+    def _update_open_row(self, encoded: dict, accumulate: bool = False):
         for key in RESERVED_KEYS & set(encoded):
             self._codec.warn_once(
                 key, f"Metric name {key!r} is reserved and will be ignored."
             )
             encoded.pop(key)
+        if accumulate:
+            # Repeated spans in one step sum rather than overwrite
+            for key, value in encoded.items():
+                current = self._open_row.get(key)
+                self._open_row[key] = value + current if _addable(current) else value
+            return
         for key in encoded:
             if key in self._open_row:
                 self._codec.warn_once(
@@ -480,8 +561,9 @@ class HistoryStore:
                 record = self._close_open_row()
         if record is not None:
             self._emit(record)
-        if self.writer is not None:
-            self.writer.flush()
+        for writer in (self.writer, self.span_writer):
+            if writer is not None:
+                writer.flush()
 
     def finish(self):
         """Close the store. Never raises: a bad disk must not mask the real error."""
@@ -492,11 +574,12 @@ class HistoryStore:
             self.flush(commit_open=True)
         except Exception as e:
             logger.warning(f"Failed to flush history on finish: {e}")
-        if self.writer is not None:
-            try:
-                self.writer.close()
-            except Exception as e:
-                logger.warning(f"Failed to close {self.log_fp}: {e}")
+        for writer in (self.writer, self.span_writer):
+            if writer is not None:
+                try:
+                    writer.close()
+                except Exception as e:
+                    logger.warning(f"Failed to close {writer.path}: {e}")
         if self._atexit_registered:
             atexit.unregister(self._atexit)
             self._atexit_registered = False
@@ -758,19 +841,8 @@ class HistoryStore:
 # ---------------------------------------------------------------------- helpers
 
 
-def current_rank() -> int:
-    """Distributed rank from the usual environment variables, 0 when unset."""
-    rank = os.getenv("RANK") or os.getenv("LOCAL_RANK") or "0"
-    try:
-        return int(rank)
-    except ValueError:
-        return 0
-
-
-def _metrics_filename(rank_aware: bool) -> str:
-    """Non-zero ranks get their own shard so appends cannot interleave."""
-    rank = current_rank() if rank_aware else 0
-    return "metrics.jsonl" if rank <= 0 else f"metrics.rank{rank}.jsonl"
+def _addable(value) -> bool:
+    return isinstance(value, (int, float)) and not isinstance(value, bool)
 
 
 def _parse_cached(entries: Sequence[tuple[int, int, bytes]]) -> list[dict]: