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

expr_tracker/__init__.py

@@ -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",
@@ -61,6 +63,8 @@ __all__ = [
     "log_artifact",
     "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,7 +13,6 @@ from __future__ import annotations
 import atexit
 import json
 import operator
-import os
 import threading
 import time
 from collections import deque
@@ -25,6 +24,12 @@ 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 +70,10 @@ 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
     # alerts and output
     alert_window: int = DEFAULT_WINDOW
     print_to_screen: bool = False
@@ -76,6 +85,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 +139,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 +191,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 +222,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 +308,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 +336,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 +344,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 +402,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 +557,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 +570,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 +837,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]:

expr_tracker/run.py

@@ -17,6 +17,7 @@ from loguru import logger
 
 from .artifacts import Artifact, ArtifactStore, coerce_artifact
 from .history import HistoryStore, current_rank, resolve_commit
+from .history.naming import sidecar_filename
 from .summary import Summary
 
 _lock = threading.RLock()
@@ -72,6 +73,7 @@ class Run:
         alert: Any = None,
         alert_rules: Sequence[Any] = (),
         alert_on_rank: int | None = 0,
+        stream: str | None = None,
         **history_kwargs,
     ):
         self.project = project
@@ -96,11 +98,15 @@ class Run:
             dir=dir,
             print_to_screen=print_to_screen,
             on_commit=self._on_commit,
+            stream=stream,
             **{**backend_kwargs.get("jsonl", {}), **history_kwargs},
         )
         self._closers.append(self.history.finish)
         self.name = self.history.name
-        self.summary = Summary(self.history.log_dir / "summary.json")
+        self.stream = self.history.stream
+        self.summary = Summary(
+            self.history.log_dir / sidecar_filename("summary", self.stream, "json")
+        )
         # History survives a crash through its own atexit hook; the summary is
         # just as much a record of the run, so persist it the same way
         atexit.register(self._save_summary_at_exit)
@@ -150,9 +156,20 @@ class Run:
             self._rollback()
             raise
 
+    @property
+    def backend_run_name(self) -> str:
+        """A stream is an independent producer, so it gets its own backend run.
+
+        Grouping ties them back together: both wandb and trackio understand
+        ``group``, and neither can merge two step axes into one run.
+        """
+        return f"{self.name}-{self.stream}" if self.stream else self.name
+
     def _init_backend(
         self, name, instance, *, entity, dir, notes, tags, resume, kwargs
     ):
+        # A grouped run keeps the streams together in the backend UI
+        grouping = {"group": self.name} if self.stream else {}
         try:
             if name == "trackio":
                 # trackio.init has no entity/notes/tags/dir/id, so those ride along
@@ -165,23 +182,27 @@ class Run:
                 }
                 instance.init(
                     project=self.project,
-                    name=self.name,
+                    name=self.backend_run_name,
                     config=config,
                     resume=_trackio_resume(resume),
-                    **kwargs,
+                    **{**grouping, **kwargs},
                 )
             else:
                 instance.init(
                     project=self.project,
-                    name=self.name,
+                    name=self.backend_run_name,
                     entity=entity,
                     dir=dir,
                     notes=notes,
                     tags=tags,
                     resume=resume,
-                    id=self.name,
+                    id=self.backend_run_name,
                     config=self.config,
-                    **kwargs,
+                    **{
+                        **grouping,
+                        **({"job_type": self.stream} if self.stream else {}),
+                        **kwargs,
+                    },
                 )
         except Exception as e:
             logger.warning(f"Failed to initialize backend {name}: {e}")

expr_tracker/spans.py

@@ -0,0 +1,236 @@
+"""Spans: how long each part of a step took, and how the parts nest.
+
+A closed span adds its duration to the open row, so it becomes an ordinary metric
+that ``history()``, alert rules and plots already understand. The full tree, with
+timestamps and attributes, goes to ``spans.jsonl`` for trace viewers.
+
+Durations never commit a step of their own; they ride along with whatever
+``log()`` commits.
+"""
+
+from __future__ import annotations
+
+import functools
+import threading
+import time
+from contextvars import ContextVar
+from typing import Any
+
+from loguru import logger
+
+# Per-thread and per-task by construction: a new thread starts from the default,
+# and each asyncio task gets its own copy, which is exactly span nesting.
+_STACK: ContextVar[tuple[Span, ...]] = ContextVar("et_span_stack", default=())
+
+DURATION_SUFFIX = "duration_ms"
+COUNT_SUFFIX = "count"
+
+
+def current_span() -> Span | None:
+    stack = _STACK.get()
+    return stack[-1] if stack else None
+
+
+def active_path() -> str:
+    """The ``/`` joined names of the spans currently open on this thread."""
+    return "/".join(span.name for span in _STACK.get())
+
+
+class Span:
+    """One timed region. Prefer :func:`span`; construct it directly only when
+    you need to end it somewhere else."""
+
+    __slots__ = (
+        "_ended",
+        "_store",
+        "_token",
+        "attributes",
+        "duration_ms",
+        "error",
+        "name",
+        "path",
+        "start",
+        "started_at",
+        "track",
+    )
+
+    def __init__(self, name: str, store, **attributes: Any):
+        self.name = str(name).strip("/") or "span"
+        self.attributes = dict(attributes)
+        self._store = store
+        self.path = self.name
+        self.track = 0
+        self.start = 0.0
+        self.started_at = 0.0
+        self.duration_ms = 0.0
+        self.error: str | None = None
+        self._token = None
+        self._ended = False
+
+    # ------------------------------------------------------------------ lifecycle
+
+    def begin(self) -> Span:
+        parent = current_span()
+        self.path = f"{parent.path}/{self.name}" if parent else self.name
+        # Children share their parent's track, so a tree never straddles lanes
+        self.track = parent.track if parent else threading.get_ident()
+        self.started_at = time.time()
+        self.start = time.perf_counter()
+        self._token = _STACK.set((*_STACK.get(), self))
+        return self
+
+    def end(self, error: BaseException | None = None) -> float:
+        """Close the span and record it. Returns the duration in milliseconds."""
+        if self._ended:
+            return self.duration_ms
+        self._ended = True
+        self.duration_ms = (time.perf_counter() - self.start) * 1000.0
+        if error is not None:
+            self.error = type(error).__name__
+        if self._token is not None:
+            try:
+                _STACK.reset(self._token)
+            except ValueError:  # pragma: no cover - ended on another thread/task
+                _STACK.set(tuple(s for s in _STACK.get() if s is not self))
+        self._record()
+        return self.duration_ms
+
+    def set(self, **attributes: Any) -> Span:
+        """Attach attributes; they reach ``spans.jsonl``, not the metrics."""
+        self.attributes.update(attributes)
+        return self
+
+    # ------------------------------------------------------------------ context
+
+    def __enter__(self) -> Span:
+        return self
+
+    def __exit__(self, kind, value, traceback) -> bool:
+        self.end(value)
+        return False
+
+    async def __aenter__(self) -> Span:
+        return self
+
+    async def __aexit__(self, kind, value, traceback) -> bool:
+        self.end(value)
+        return False
+
+    def __repr__(self) -> str:  # pragma: no cover - debugging aid
+        state = f"{self.duration_ms:.3f}ms" if self._ended else "open"
+        return f"Span({self.path}, {state})"
+
+    # ------------------------------------------------------------------ output
+
+    def _record(self):
+        if self._store is None:
+            return
+        metrics = {
+            f"{self.path}/{DURATION_SUFFIX}": self.duration_ms,
+            f"{self.path}/{COUNT_SUFFIX}": 1,
+        }
+        record = {
+            "name": self.path,
+            "depth": self.path.count("/"),
+            "track": self.track,
+            "start": self.started_at,
+            "dur_ms": self.duration_ms,
+        }
+        if self.attributes:
+            record["args"] = self.attributes
+        if self.error:
+            record["error"] = self.error
+        try:
+            self._store.record_span(metrics, record)
+        except Exception as e:  # a measurement must never break the thing measured
+            logger.warning(f"Failed to record span {self.path!r}: {e}")
+
+
+class _NullSpan(Span):
+    """What you get without a run: still times, records nothing."""
+
+    def __init__(self, name: str, **attributes: Any):
+        super().__init__(name, None, **attributes)
+
+
+def start_span(name: str, **attributes: Any) -> Span:
+    """Begin a span that you will ``end()`` yourself, possibly in another scope."""
+    from .run import current_run
+
+    run = current_run()
+    store = run.history if run is not None else None
+    span = Span(name, store, **attributes) if store is not None else _NullSpan(name)
+    return span.begin()
+
+
+class span:
+    """Time a region, as a context manager, an async context manager or a decorator.
+
+    ```python
+    with et.span("forward"):
+        with et.span("attention"):
+            ...
+
+    @et.span("preprocess")
+    def preprocess(batch): ...
+
+    async with et.span("fetch"):
+        await loader.next()
+    ```
+
+    Nested names join with ``/``, so the duration lands on the metric
+    ``forward/attention/duration_ms`` and window functions work on it directly.
+    """
+
+    def __init__(self, name: str, **attributes: Any):
+        self._name = name
+        self._attributes = attributes
+        self._span: Span | None = None
+
+    def __enter__(self) -> Span:
+        self._span = start_span(self._name, **self._attributes)
+        return self._span
+
+    def __exit__(self, kind, value, traceback) -> bool:
+        if self._span is not None:
+            self._span.end(value)
+        return False
+
+    async def __aenter__(self) -> Span:
+        return self.__enter__()
+
+    async def __aexit__(self, kind, value, traceback) -> bool:
+        return self.__exit__(kind, value, traceback)
+
+    def __call__(self, fn):
+        import asyncio
+
+        if asyncio.iscoroutinefunction(fn):
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args, **kwargs):
+                current = start_span(self._name, **self._attributes)
+                error = None
+                try:
+                    return await fn(*args, **kwargs)
+                except BaseException as e:
+                    error = e
+                    raise
+                finally:
+                    current.end(error)
+
+            return async_wrapper
+
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            current = start_span(self._name, **self._attributes)
+            error = None
+            try:
+                return fn(*args, **kwargs)
+            except BaseException as e:
+                error = e
+                raise
+            finally:
+                current.end(error)
+
+        return wrapper

expr_tracker/trace.py

@@ -0,0 +1,189 @@
+"""Export recorded spans as a Chrome Trace, readable by Perfetto and chrome://tracing.
+
+Writing the format rather than a viewer means the timeline can sit next to a
+``torch.profiler`` trace in the same window, which is where the interesting
+question usually is: what were the GPUs doing while the data loader stalled.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from .history.naming import parse_stream, spans_filename
+
+MICROSECONDS = 1_000_000.0
+
+
+def span_files(run: str | Path, stream: str | list[str] | None = "*") -> list[Path]:
+    """Span files of a run: every stream by default, or the ones named."""
+    directory = Path(run)
+    if directory.is_file():
+        return [directory]
+    if stream == "*":
+        # Default producer first: "spans.data.jsonl" sorts before "spans.jsonl"
+        return sorted(
+            directory.glob("spans*.jsonl"),
+            key=lambda p: (parse_stream(p.name) is not None, p.name),
+        )
+    wanted = [stream] if isinstance(stream, (str, type(None))) else list(stream)
+    found = []
+    for name in wanted:
+        path = directory / spans_filename(name, rank_aware=False)
+        if not path.is_file():
+            raise FileNotFoundError(f"No {path.name} under {directory}")
+        found.append(path)
+    return found
+
+
+def read_spans(path: Path) -> list[dict]:
+    """Decode one span file, skipping anything malformed."""
+    if not path.is_file():
+        return []
+    spans = []
+    for line in path.read_text(encoding="utf-8").splitlines():
+        if not line.strip():
+            continue
+        try:
+            record = json.loads(line)
+        except ValueError:
+            continue
+        if isinstance(record, dict) and "start" in record and "dur_ms" in record:
+            spans.append(record)
+    return spans
+
+
+def _pack(spans: list[dict], indices: list[int], offset: int) -> dict[int, int]:
+    """Lay one track's spans onto lanes, each lane a valid stack."""
+    order = sorted(
+        indices, key=lambda i: (spans[i]["start"], -float(spans[i]["dur_ms"]))
+    )
+    lanes: list[list[float]] = []  # per lane, the end times currently open
+    placement: dict[int, int] = {}
+    for index in order:
+        start = float(spans[index]["start"])
+        end = start + float(spans[index]["dur_ms"]) / 1000.0
+        for lane, stack in enumerate(lanes):
+            while stack and stack[-1] <= start:
+                stack.pop()
+            if not stack or end <= stack[-1]:  # free, or nests inside the top
+                stack.append(end)
+                placement[index] = offset + lane
+                break
+        else:
+            lanes.append([end])
+            placement[index] = offset + len(lanes) - 1
+    return placement
+
+
+def assign_lanes(spans: list[dict]) -> dict[int, int]:
+    """Place spans on lanes so each lane is a valid stack, keyed by list index.
+
+    Threads come first: work on different threads is unrelated, and letting one
+    contain the other would draw a nesting that never happened. Within a thread,
+    anything still overlapping -- concurrent asyncio tasks -- gets its own lane,
+    because a Chrome Trace track must be properly nested to render at all.
+    """
+    by_track: dict[Any, list[int]] = {}
+    for index, record in enumerate(spans):
+        by_track.setdefault(record.get("track", 0), []).append(index)
+    placement: dict[int, int] = {}
+    offset = 0
+    for track in sorted(by_track, key=str):
+        lanes = _pack(spans, by_track[track], offset)
+        placement.update(lanes)
+        offset = max(lanes.values(), default=offset - 1) + 1
+    return placement
+
+
+def build_trace(
+    run: str | Path,
+    stream: str | list[str] | None = "*",
+    *,
+    step_range: tuple[int | None, int | None] | None = None,
+) -> dict[str, Any]:
+    """Chrome Trace Event Format for a run's spans, one process per stream."""
+    events: list[dict] = []
+    origin: float | None = None
+    for pid, path in enumerate(span_files(run, stream)):
+        spans = read_spans(path)
+        if step_range is not None:
+            low, high = step_range
+            spans = [
+                s
+                for s in spans
+                if (low is None or s.get("_step", 0) >= low)
+                and (high is None or s.get("_step", 0) < high)
+            ]
+        name = parse_stream(path.name) or "default"
+        events.append(
+            {
+                "name": "process_name",
+                "ph": "M",
+                "pid": pid,
+                "tid": 0,
+                "args": {"name": name},
+            }
+        )
+        if not spans:
+            continue
+        lanes = assign_lanes(spans)
+        tracks = sorted({s.get("track", 0) for s in spans}, key=str)
+        # Label each lane with the thread it came from
+        lane_track: dict[int, Any] = {}
+        for index, record in enumerate(spans):
+            lane_track.setdefault(lanes[index], record.get("track", 0))
+        for tid, track in sorted(lane_track.items()):
+            label = f"thread {track}" if len(tracks) > 1 else name
+            events.append(
+                {
+                    "name": "thread_name",
+                    "ph": "M",
+                    "pid": pid,
+                    "tid": tid,
+                    "args": {"name": label},
+                }
+            )
+        for index, record in enumerate(spans):
+            start = float(record["start"])
+            origin = start if origin is None else min(origin, start)
+            args = dict(record.get("args") or {})
+            args["step"] = record.get("_step")
+            if record.get("error"):
+                args["error"] = record["error"]
+            if len(tracks) > 1:
+                args["track"] = record.get("track", 0)
+            events.append(
+                {
+                    "name": record.get("name", "span"),
+                    "cat": name,
+                    "ph": "X",
+                    "pid": pid,
+                    "tid": lanes[index],
+                    "ts": start * MICROSECONDS,
+                    "dur": float(record["dur_ms"]) * 1000.0,
+                    "args": args,
+                }
+            )
+    # Relative timestamps keep the numbers readable and the streams aligned
+    if origin is not None:
+        for event in events:
+            if "ts" in event:
+                event["ts"] -= origin * MICROSECONDS
+    return {"traceEvents": events, "displayTimeUnit": "ms"}
+
+
+def write_trace(
+    run: str | Path,
+    output: str | Path,
+    stream: str | list[str] | None = "*",
+    *,
+    step_range: tuple[int | None, int | None] | None = None,
+) -> int:
+    """Write the trace and return how many spans it holds."""
+    trace = build_trace(run, stream, step_range=step_range)
+    path = Path(output)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(trace), encoding="utf-8")
+    return sum(1 for event in trace["traceEvents"] if event["ph"] == "X")

expr_tracker/tracker.py

@@ -10,6 +10,7 @@ from typing import Literal
 from .artifacts import Artifact
 from .history import read_history
 from .run import Run, current_run, require_run, set_run
+from .run import current_run as current
 
 Tracker = Run
 
@@ -70,9 +71,23 @@ def log(data: dict, step: int | None = None, commit: bool | None = None):
     require_run().log(data, step=step, commit=commit)
 
 
+class _Current:
+    """Sentinel: ``stream`` was not given, so use the running stream.
+
+    ``stream=None`` has to keep meaning the default, unnamed producer.
+    """
+
+    def __repr__(self) -> str:  # pragma: no cover - debugging aid
+        return "<current stream>"
+
+
+CURRENT_STREAM = _Current()
+
+
 def history(
     n: int | None = 50,
     *,
+    stream: str | _Current | None = CURRENT_STREAM,
     output_type: str = "dict",
     metrics: Sequence[str] | None = None,
     step_range: tuple[int | None, int | None] | None = None,
@@ -86,25 +101,32 @@ def history(
 
     ``output_type="pandas"`` returns a DataFrame (pandas is an optional extra).
     Passing ``run`` reads that run directory or file offline, without ``init()``.
+    ``stream`` selects an independent producer, ``None`` being the default one;
+    another process's stream is read from its file, so it reflects what that
+    process has flushed. Omit it to read whichever stream this process writes.
     """
-    if run is not None:
-        return read_history(
-            run,
+    active = getattr(current(), "stream", None)
+    if run is None and (stream is CURRENT_STREAM or stream == active):
+        return require_run().history.get(
             n,
             output_type=output_type,
             metrics=metrics,
             step_range=step_range,
             include_meta=include_meta,
+            include_open=include_open,
             fill_missing=fill_missing,
             dropna=dropna,
         )
-    return require_run().history.get(
+    if run is None:
+        run = require_run().history.log_dir
+    return read_history(
+        run,
         n,
+        stream=None if stream is CURRENT_STREAM else stream,
         output_type=output_type,
         metrics=metrics,
         step_range=step_range,
         include_meta=include_meta,
-        include_open=include_open,
         fill_missing=fill_missing,
         dropna=dropna,
     )

{expr_tracker-0.2.2.dist-info → expr_tracker-0.2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: expr_tracker
-Version: 0.2.2
+Version: 0.2.3
 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/
@@ -106,6 +106,8 @@ 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. |
+| [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`. |
 
@@ -148,6 +150,9 @@ uv run ruff format src tests
 | `test_hot_paths.py` | contracts and defaults of `et.log` / `et.history` / summary / alerts |
 | `test_value_encoding.py` | numpy, pydantic, datetime, Path, Enum round trips; output types; query bounds |
 | `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_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 |
 | `test_trackio.py` | trackio contract, resume mapping, real end-to-end |

{expr_tracker-0.2.2.dist-info → expr_tracker-0.2.3.dist-info}/RECORD

@@ -1,12 +1,14 @@
-expr_tracker/__init__.py,sha256=pegVYNdNFlGuiQmIrNjQvSn4xWEUf-esvK_PzwDvr24,1334
+expr_tracker/__init__.py,sha256=wWvM6nZSIlP_pWYcmQGEbMMD9GCBA0hpatui5TNwPHI,1418
 expr_tracker/_compat.py,sha256=qf8oYIK_lkoLzcPp6wPCYLJtVatCl7ROYRNPwc5F1Pg,981
 expr_tracker/artifacts.py,sha256=yRVXmuqdZUxiZa0c_CwZi4rBxaItwhu9bZ1x21yP1S0,12560
-expr_tracker/cli.py,sha256=YM1tuPX19Y_oNWtCi1DmAY3wyKd-uGILzM1twvv5il4,6305
+expr_tracker/cli.py,sha256=dtp0YsNtGg7T7M-elA511YX8NU_-OJRAwNbs8I1353o,7255
 expr_tracker/encoders.py,sha256=X5ySQasSrYPsdYCdk7nhpJFGFB9f_8HWQuGOQjTvh9E,11991
 expr_tracker/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-expr_tracker/run.py,sha256=7hHuLaa5Ox-cP891eMuxZ6_Uy9iTaaqMjjshl32kmVk,15744
+expr_tracker/run.py,sha256=nuXhNmepG688X86cI0FMM_hqjyygD1Oi8Ck8-lyHsPA,16664
+expr_tracker/spans.py,sha256=su1M0yxgi51911a6rBPaWhDoVBfIW1vla_okYCW9Z28,7365
 expr_tracker/summary.py,sha256=lpHSh3pci2Ym6wJhVZwNUJbUp2A64EN_rC7Ym2d7PPI,3150
-expr_tracker/tracker.py,sha256=NabstZ2Ysme64RkPooZj-r7AA8Wst49H9au5Q7yTn9o,5073
+expr_tracker/trace.py,sha256=lag-T-5F1LOz4MKsxeMfTW1FJN0KxxznJ48sYT7vX28,6966
+expr_tracker/tracker.py,sha256=hpZB8bFA2GqCbvHEDZkBS9cCzi2JgXeLUNTzJm-wr8c,5921
 expr_tracker/types.py,sha256=69gFtUaKGl4asLK7nJci9XyTyV0yTIPorB-bcfGrCBA,383
 expr_tracker/alerts/__init__.py,sha256=dl96OYZh75kV3NLBSA35zNvTFcL6XFMuPoYvRoz7hlQ,8192
 expr_tracker/alerts/dispatch.py,sha256=k2pDZVJy6ajlQ7kUdvaBBZ_AmdtZ8opaQWjADn_i8uI,11652
@@ -21,15 +23,16 @@ expr_tracker/alerts/expr/lexer.py,sha256=NrvDi-f_8F9V1_vQT-vAGENNA8nauD9IbJH4r21
 expr_tracker/alerts/expr/nodes.py,sha256=ZsXX4VJDolVQsVoEtIhkDj6QspMz7lTfRYaEZ8F6grg,9896
 expr_tracker/alerts/expr/parser.py,sha256=heGXAVjCZ1AM0u5-Usld4G_P38N0nUaXMlh8C7CHqjQ,6315
 expr_tracker/alerts/expr/rule.py,sha256=nO8vzhWyGDxSmSXqJrXOU_Yoy7Tzot3_c2W2bZ4s2Qw,3653
-expr_tracker/history/__init__.py,sha256=sppio0jJ2E8WBr3HcZ7UO5vo7GNqQA1rhEgN6wuhQ_o,522
+expr_tracker/history/__init__.py,sha256=oSAamhXglLh1oqIHI0-DqQHlR97yt1gLBWSTlyFbCo0,716
 expr_tracker/history/codec.py,sha256=iD0TmnmahpRmSRRXt1-lmTPPOzL4TdA4uRQlhMdg49o,3121
 expr_tracker/history/frame.py,sha256=PFDGqR0oHY94KosSXA7Fu-bphEhZ2yM8cXW9rlZ2IC0,2996
-expr_tracker/history/reader.py,sha256=YYrQluS5dopax84ltd5gr6xO4yrygRiOo3emIhjR6Kk,12822
+expr_tracker/history/naming.py,sha256=CT2xDpo0GHYnLyjgRlsPjIX3GAIDEdn4U9I1-cNahHw,2507
+expr_tracker/history/reader.py,sha256=5PY5nbNfX-Cl34jcyYTq2qm8nNgNDTqv7U5Xnb9Rgxw,14155
 expr_tracker/history/series.py,sha256=cFdUO1FRs5KW0jYsNPzda3tcu645NH1wwfXJitQfXFM,4042
-expr_tracker/history/store.py,sha256=Q7F0I_dF2zxHUTtSnzfVCviSeAvW49i3OqLMjoQh1sQ,33072
+expr_tracker/history/store.py,sha256=ZazPIcQ0Aj23DzIBN7B-xiIAk5ej6wfwsUdBYNywHNw,36047
 expr_tracker/history/writer.py,sha256=q5kNANcZCq5nykvOK5jZ2xc6Z1B4oZxEn3uSXzDCtlY,16096
-expr_tracker-0.2.2.dist-info/METADATA,sha256=525p8soaoyMXiBApvXQjcJ5JVTJOb1b0v6lakS6ZB5Y,8388
-expr_tracker-0.2.2.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
-expr_tracker-0.2.2.dist-info/entry_points.txt,sha256=GfZCAG2JFsd5CxTjRsGJAhYInpBCuLjrU3csmUVBDSs,45
-expr_tracker-0.2.2.dist-info/licenses/LICENSE,sha256=ORsYTnF-2SLUaM8FgXlJ-uXuN4ceGACAAU3JZiS8pvs,1069
-expr_tracker-0.2.2.dist-info/RECORD,,
+expr_tracker-0.2.3.dist-info/METADATA,sha256=HXNy4hcFe74JnO8dYmnPfoCqt2IJpjWJT34E6eK6yYE,9120
+expr_tracker-0.2.3.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+expr_tracker-0.2.3.dist-info/entry_points.txt,sha256=GfZCAG2JFsd5CxTjRsGJAhYInpBCuLjrU3csmUVBDSs,45
+expr_tracker-0.2.3.dist-info/licenses/LICENSE,sha256=ORsYTnF-2SLUaM8FgXlJ-uXuN4ceGACAAU3JZiS8pvs,1069
+expr_tracker-0.2.3.dist-info/RECORD,,