From 3a4092f5aea45bb0e195d7d6454896ea60405c99 Mon Sep 17 00:00:00 2001
From: Alexandru-Andrei Rotaru <rotarualexandruandrei94@gmail.com>
Date: Tue, 7 Jul 2026 17:55:32 +0200
Subject: [PATCH 01/19] examples: add idiomatic per-sample signals MNIST
 example

A single self-contained PyTorch example showing the three ways to declare
per-sample signals in WeightsLab, plus an honest overhead breakdown:

  - base signals (from logits) pushed with wl.save_signals
  - a live derived @wl.signal (loss normalized by a running mean)
  - end-of-run derived signals from each sample's full loss trajectory:
    a six-class shape as a text categorical tag (100% coverage) + loss_cv /
    loss_drop, written back by sample_id

The run reports a per-step cost breakdown (baseline / loss-logging / signal
compute / persist) vs plain PyTorch, CUDA-synced, showing the signal math is
cheap and persistence dominates.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../examples/PyTorch/ws-signals-mnist/main.py | 451 ++++++++++++++++++
 1 file changed, 451 insertions(+)
 create mode 100644 weightslab/examples/PyTorch/ws-signals-mnist/main.py

diff --git a/weightslab/examples/PyTorch/ws-signals-mnist/main.py b/weightslab/examples/PyTorch/ws-signals-mnist/main.py
new file mode 100644
index 00000000..48a981c9
--- /dev/null
+++ b/weightslab/examples/PyTorch/ws-signals-mnist/main.py
@@ -0,0 +1,451 @@
+"""
+WeightsLab per-sample signals — a minimal, idiomatic MNIST example
+==================================================================
+
+Instrument an ordinary PyTorch training loop with WeightsLab so that, for
+*every training sample*, you capture a small graph of signals — and measure
+what that costs.
+
+    10 base signals    loss, confidence, margin, entropy, correctness, ...
+       (from logits)   pushed once per step with ``wl.save_signals``
+
+     1 live derived    loss normalized by a running mean — a dynamic @wl.signal
+       (@wl.signal)    the framework computes for you on every loss update
+
+     3 end-of-run      from each sample's *full* loss trajectory, computed once
+       derived signals after training:
+                         - loss_shape : one of six curve shapes, stored as a
+                                        text categorical tag (100% coverage)
+                         - loss_cv    : coefficient of variation (instability)
+                         - loss_drop  : net improvement first->last
+
+The output is a per-sample table mapping every ``sample_id`` to the latest value
+of every signal — the raw material for curation: surface the hard samples, the
+likely-mislabeled ones (``Flat_high``), and the ones the model forgets
+(``U_Shape``).
+
+The idiom
+---------
+Three ways to declare a signal, and this example shows all three:
+
+  * **Push it yourself** — value comes from something WeightsLab can't see (raw
+    logits): compute it in the loop and call ``wl.save_signals(...)``.
+
+  * **Declare a live @wl.signal** — value is a function of a metric WeightsLab
+    already logs (the per-sample loss). Write a function of ``ctx``; the
+    framework fires it per sample on every loss update and persists the result.
+
+  * **Derive at the end** — value needs a sample's *whole* trajectory. After
+    training, read each sample's loss history from the ledger, reduce it, and
+    write the result back (as a signal or a categorical tag) for every sample.
+
+Overhead
+--------
+The run reports the per-step cost of the signal machinery (watched loss +
+``save_signals``) measured against a plain-PyTorch baseline on the same loader,
+plus the wall-time of the end-of-run derived sweep. All timings are CUDA-synced.
+
+Usage
+-----
+    python mnist_signals_example.py --outdir ./out
+"""
+import argparse
+import os
+import shutil
+import time
+from collections import defaultdict
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.utils.data import Dataset
+
+import weightslab as wl
+from weightslab.components.global_monitoring import guard_training_context
+
+
+# ----------------------------------------------------------------------------
+# A small MNIST CNN that returns raw logits (no softmax) — so CrossEntropyLoss
+# is correct and the logit-derived signals below are meaningful. Kept inline so
+# this example is a single self-contained file.
+# ----------------------------------------------------------------------------
+class SmallCNN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.input_shape = (1, 1, 28, 28)
+        self.net = nn.Sequential(
+            nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),   # 28 -> 14
+            nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),  # 14 -> 7
+            nn.Flatten(),
+            nn.Linear(16 * 7 * 7, 64), nn.ReLU(),
+            nn.Linear(64, 10),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# ----------------------------------------------------------------------------
+# Config — small on purpose. A subset + a handful of epochs trains a real model
+# and yields a nice spread of per-sample loss-trajectory shapes in ~a minute.
+# ----------------------------------------------------------------------------
+SEED = 0
+SUBSET = 2000      # samples of MNIST train to use (None -> full 60k)
+EPOCHS = 20
+BATCH = 64
+LR = 0.01
+MEASURE_STEPS = 50  # steps used for the per-step overhead breakdown
+
+# The one metric everything subscribes to. It is the ``signal_name`` of the
+# watched loss below, and the ``subscribe_to`` of every live @wl.signal.
+LOSS_SIGNAL = "train/loss_sample"
+
+
+# ----------------------------------------------------------------------------
+# 1. Dataset — every sample carries a stable id so its signals can be tracked
+#    across epochs. The id is just the dataset index here; in a real pipeline
+#    it would be your annotation / image id.
+# ----------------------------------------------------------------------------
+class MNISTIdx(Dataset):
+    """MNIST that returns ``(image, uid, label)`` instead of ``(image, label)``."""
+
+    def __init__(self, root, train, transform, max_samples=None):
+        self.mnist = datasets.MNIST(root=root, train=train, download=True, transform=None)
+        self.transform = transform
+        self.max_samples = max_samples
+
+    def __len__(self):
+        if self.max_samples is not None:
+            return min(len(self.mnist), self.max_samples)
+        return len(self.mnist)
+
+    def __getitem__(self, idx):
+        image, label = self.mnist[idx]
+        return self.transform(image), idx, label
+
+
+# ----------------------------------------------------------------------------
+# 2. Base signals — pure functions of the batch logits. Nothing
+#    WeightsLab-specific here; these just turn logits + labels into per-sample
+#    scalars. They need raw logits, so we push them ourselves (below) with
+#    ``wl.save_signals`` rather than declaring them as @wl.signal.
+# ----------------------------------------------------------------------------
+def base_signals(logits, labels, loss_per_sample):
+    """10 base per-sample signals derived from raw logits ``[B, C]``."""
+    probs = torch.softmax(logits, dim=1)
+    top2 = probs.topk(2, dim=1).values
+    confidence = top2[:, 0]
+    margin = top2[:, 0] - top2[:, 1]
+    entropy = -(probs * (probs + 1e-12).log()).sum(dim=1)
+    correct = (logits.argmax(dim=1) == labels).float()
+    true_prob = probs.gather(1, labels.view(-1, 1)).squeeze(1)
+    return {
+        "sig/loss": loss_per_sample,
+        "sig/confidence": confidence,
+        "sig/margin": margin,
+        "sig/entropy": entropy,
+        "sig/correct": correct,
+        "sig/true_class_prob": true_prob,
+        "sig/nll": -(true_prob + 1e-12).log(),
+        "sig/logit_true": logits.gather(1, labels.view(-1, 1)).squeeze(1),
+        "sig/logit_max": logits.max(dim=1).values,
+        "sig/logit_std": logits.std(dim=1),
+    }
+
+
+# ----------------------------------------------------------------------------
+# 3. Live derived signal — a dynamic @wl.signal. It subscribes to the per-sample
+#    loss and fires once per sample every time the loss is logged, receiving a
+#    SignalContext ``ctx`` whose ``subscribed_value`` is that sample's current
+#    loss. It returns one scalar, stored per sample_id. No save_signals, no
+#    plumbing — the framework runs it. Here: loss normalized by a running mean
+#    ("how hard is this sample right now, relative to the average?").
+# ----------------------------------------------------------------------------
+_running_mean_loss = {"v": 1.0}
+
+
+@wl.signal(name="sig/loss_norm", subscribe_to=LOSS_SIGNAL)
+def loss_norm(ctx):
+    loss = ctx.subscribed_value
+    if loss is None:
+        return 0.0
+    m = _running_mean_loss
+    m["v"] = 0.98 * m["v"] + 0.02 * float(loss)
+    return float(loss) / (m["v"] + 1e-9)
+
+
+# ----------------------------------------------------------------------------
+# 4. Trajectory shape — the six curve archetypes, and the classifier that
+#    reduces a per-sample loss curve to one of them.
+# ----------------------------------------------------------------------------
+LOSS_SHAPE_LABELS = [
+    "monotonic", "plateaued", "Flat_high",
+    "high_variance", "U_Shape", "Spiked",
+]
+
+
+def classify_loss_shape(values):
+    """Reduce a per-sample loss trajectory (ordered by step) to a shape label.
+
+    Thresholds are scale-invariant (fractions of the trajectory's own range)
+    and illustrative — tune them for your own task. Returns None if there is
+    not enough history yet (< 5 points).
+    """
+    y = np.asarray(values, dtype=float)
+    if y.size < 5:
+        return None
+
+    n = y.size
+    first, last = float(y[0]), float(y[-1])
+    ymin, ymax = float(y.min()), float(y.max())
+    rng = max(ymax - ymin, 1e-8)
+    mean = float(y.mean())
+
+    cv = float(y.std()) / (abs(mean) + 1e-8)         # noisiness
+    drop = (first - last) / (abs(first) + 1e-8)      # net improvement
+    argmin = int(np.argmin(y))
+    rebound = (last - ymin) / rng                     # climb-back from trough
+    max_up_jump = float(np.diff(y).max()) / rng       # largest single-step rise
+
+    tail = y[int(0.6 * n):]
+    tail_flat = (float(tail.std()) / (abs(float(tail.mean())) + 1e-8)) < 0.1
+
+    # Learning shapes first (a big net drop is "being learned", even though a
+    # 2.0 -> 0.0 curve has a high coefficient of variation); then the
+    # not-learning / noisy shapes.
+    if 0.2 * n < argmin < 0.8 * n and rebound > 0.3:
+        return "U_Shape"        # learned then forgotten (min in the middle)
+    if drop > 0.4:
+        return "monotonic"      # steadily learned (large net drop)
+    if drop > 0.15 and tail_flat:
+        return "plateaued"      # improved, then stuck still-high
+    if max_up_jump > 0.5:
+        return "Spiked"         # sudden jump, no net learning
+    if cv > 0.5:
+        return "high_variance"  # noisy oscillation, not converging
+    return "Flat_high"          # never moved, stayed high (likely mislabel)
+
+
+# ----------------------------------------------------------------------------
+# 5. End-of-run derived signals — computed once, from each sample's *full* loss
+#    trajectory in the ledger. This is the "derive at the end" idiom: one query,
+#    then per sample reduce the curve and write the results back by sample_id.
+#    Guarantees 100% coverage, unlike the throttled live signal.
+# ----------------------------------------------------------------------------
+def finalize_derived_signals():
+    """Read every sample's loss trajectory and write the definitive derived
+    signals: shape (text categorical tag) + loss_cv + loss_drop (scalars)."""
+    series = defaultdict(list)
+    for sid, step, val, _ in wl.query_signal_history(LOSS_SIGNAL):
+        series[sid].append((step, val))
+
+    ids, cv, drop = [], [], []
+    by_label = defaultdict(list)
+    dist = defaultdict(int)
+    for sid, pts in series.items():
+        y = np.asarray([v for _, v in sorted(pts)], dtype=float)
+        ids.append(sid)
+        mean = float(y.mean())
+        cv.append(float(y.std()) / (abs(mean) + 1e-8))
+        drop.append((float(y[0]) - float(y[-1])) / (abs(float(y[0])) + 1e-8))
+        label = classify_loss_shape(y.tolist())
+        if label is not None:
+            by_label[label].append(sid)
+            dist[label] += 1
+
+    # Scalar derived signals: written by stored sample_id (save_signals writes
+    # to rows by the literal id — the same id space the ledger already uses).
+    wl.save_signals(
+        signals={"sig/loss_cv": torch.tensor(cv), "sig/loss_drop": torch.tensor(drop)},
+        batch_ids=ids, log=False,
+    )
+    # Definitive shape as a text categorical tag, for every sample.
+    for label, sids in by_label.items():
+        wl.set_categorical_tag(sids, "loss_shape", label)
+    return dict(dist)
+
+
+# ----------------------------------------------------------------------------
+# Overhead timing helpers
+# ----------------------------------------------------------------------------
+def _sync(device):
+    if device.type == "cuda":
+        torch.cuda.synchronize()
+
+
+def measure_overhead(model, loader, optimizer, criterion, device, n_steps):
+    """Per-step cost breakdown (ms/step, CUDA-synced), on the SAME loader.
+
+    For each measured step we time, separately:
+      - baseline : a plain forward+backward+step with an unwatched loss (the
+                   pure-PyTorch cost — no signals at all)
+      - watched  : forward+backward+step with the *watched* loss (adds the
+                   per-sample loss logging + any live @wl.signal subscribers)
+      - compute  : the base_signals() math (softmax/entropy/... on logits)
+      - persist  : the wl.save_signals() call (WeightsLab persistence to the
+                   ledger / H5)
+
+    'watched - baseline' isolates the logging cost; 'compute' is the signal math
+    (cheap); 'persist' is the storage cost.
+    """
+    plain = nn.CrossEntropyLoss()
+    acc = defaultdict(float)
+    it = iter(loader)
+    done = 0
+    while done < n_steps:
+        batch = next(it, None)
+        if batch is None:
+            it = iter(loader)
+            continue
+        inputs, ids, labels = batch
+        inputs, labels = inputs.to(device), labels.to(device)
+
+        # baseline: plain step (no WeightsLab)
+        _sync(device); t0 = time.perf_counter()
+        optimizer.zero_grad()
+        plain(model(inputs), labels).backward()
+        optimizer.step()
+        _sync(device); t1 = time.perf_counter()
+
+        # instrumented step, timing each stage
+        with guard_training_context:
+            optimizer.zero_grad()
+            logits = model(inputs)
+            preds = logits.argmax(dim=1, keepdim=True)
+            loss = criterion(logits, labels, batch_ids=ids, preds=preds)
+            loss.mean().backward()
+            optimizer.step()
+            _sync(device); t2 = time.perf_counter()
+            sig = base_signals(logits.detach(), labels, loss.detach())
+            _sync(device); t3 = time.perf_counter()
+            wl.save_signals(signals=sig, batch_ids=ids,
+                            preds_raw=logits.detach(), preds=preds, targets=labels)
+            _sync(device); t4 = time.perf_counter()
+
+        acc["baseline"] += t1 - t0
+        acc["watched"] += t2 - t1
+        acc["compute"] += t3 - t2
+        acc["persist"] += t4 - t3
+        done += 1
+    return {k: 1000.0 * v / max(done, 1) for k, v in acc.items()}
+
+
+# ----------------------------------------------------------------------------
+# Training — wrap the four moving parts, call the watched loss (which logs
+# LOSS_SIGNAL and fires the live @wl.signal subscribers), and push the
+# logit-derived base signals. Then derive from the trajectories at the end.
+# ----------------------------------------------------------------------------
+def train(outdir, data_root, device):
+    torch.manual_seed(SEED)
+    log_dir = os.path.join(outdir, "wl_logs")
+    # Start from a clean ledger so the trajectories reflect only this run
+    # (a stale ledger would reload old per-sample losses and freeze the curves).
+    shutil.rmtree(log_dir, ignore_errors=True)
+    os.makedirs(log_dir, exist_ok=True)
+
+    hyperparams = {
+        "experiment_name": "mnist_signals_example",
+        "device": str(device),
+        "root_log_dir": log_dir,
+        "serving_grpc": False,   # headless: no gRPC / CLI control plane
+        "serving_cli": False,
+        # Declare the loader here too, otherwise WeightsLab's config sync
+        # resets batch_size to its default. Keys map to loader kwargs.
+        "data": {"train_loader": {"batch_size": BATCH, "shuffle": True}},
+    }
+    wl.watch_or_edit(hyperparams, flag="hyperparameters", defaults=hyperparams)
+
+    dataset = MNISTIdx(data_root, train=True, transform=transforms.ToTensor(),
+                       max_samples=SUBSET)
+
+    # Wrap the four moving parts. After this, model/opt/loader behave exactly
+    # like their plain PyTorch counterparts — WeightsLab just observes.
+    model = wl.watch_or_edit(SmallCNN().to(device), flag="model", device=device)
+    optimizer = wl.watch_or_edit(optim.Adam(model.parameters(), lr=LR), flag="optimizer")
+    loader = wl.watch_or_edit(
+        dataset, flag="data", loader_name="train_loader",
+        batch_size=BATCH, shuffle=True, is_training=True, preload_labels=True,
+    )
+    # The watched loss: reduction="none" -> one value per sample. Passing
+    # batch_ids= on each call logs LOSS_SIGNAL per sample AND drives the live
+    # @wl.signal subscribers above. This is the only "signal" we hand-wire.
+    criterion = wl.watch_or_edit(
+        nn.CrossEntropyLoss(reduction="none"),
+        flag="loss", signal_name=LOSS_SIGNAL, per_sample=True, log=True,
+    )
+
+    # Declare the categorical tag up front so the studio shows all six choices;
+    # the live signal + the end sweep populate it.
+    wl.register_categorical_tag("loss_shape", LOSS_SHAPE_LABELS)
+
+    wl.serve(serving_grpc=False, serving_cli=False)
+    wl.start_training(timeout=3)
+
+    # ---- train: watched loss (logs the per-sample loss trajectory + fires the
+    #      live @wl.signal) and push the logit-derived base signals -----------
+    for epoch in range(1, EPOCHS + 1):
+        for inputs, ids, labels in loader:
+            inputs, labels = inputs.to(device), labels.to(device)
+            with guard_training_context:
+                optimizer.zero_grad()
+                logits = model(inputs)
+                preds = logits.argmax(dim=1, keepdim=True)
+
+                # Watched loss: logs LOSS_SIGNAL + fires the live @wl.signal.
+                loss_per_sample = criterion(logits, labels, batch_ids=ids, preds=preds)
+                loss_per_sample.mean().backward()
+                optimizer.step()
+
+                # Logit-derived base signals: push them ourselves.
+                wl.save_signals(
+                    signals=base_signals(logits.detach(), labels, loss_per_sample.detach()),
+                    batch_ids=ids,
+                    preds_raw=logits.detach(), preds=preds, targets=labels,
+                )
+        print(f"epoch {epoch:2d}/{EPOCHS}  loss={loss_per_sample.mean().item():.4f}")
+
+    # ---- overhead: per-step cost breakdown vs plain PyTorch ------------------
+    ovh = measure_overhead(model, loader, optimizer, criterion, device, MEASURE_STEPS)
+
+    # ---- derive from full trajectories, once, at the end --------------------
+    _sync(device)
+    t_fin = time.perf_counter()
+    dist = finalize_derived_signals()
+    fin_ms = 1000.0 * (time.perf_counter() - t_fin)
+
+    # ---- export: base + derived signals AND the text shape tag --------------
+    report_path = os.path.join(outdir, "per_sample_signals.csv")
+    report_path = wl.write_dataframe(report_path, format="csv", columns=["signals", "tags"])
+
+    # ---- summary ------------------------------------------------------------
+    instrumented = ovh["watched"] + ovh["compute"] + ovh["persist"]
+    logging_ms = ovh["watched"] - ovh["baseline"]
+    total_over = instrumented - ovh["baseline"]
+    print(f"\nDone. Per-sample report: {report_path}")
+    print(f"\n── per-step overhead (ms/step, CUDA-synced, {MEASURE_STEPS} steps) ──")
+    print(f"  baseline  (plain PyTorch)     : {ovh['baseline']:7.3f}")
+    print(f"  + loss logging + subscribers  : {logging_ms:7.3f}")
+    print(f"  + signal compute (the math)   : {ovh['compute']:7.3f}   <- signal computation")
+    print(f"  + persist (save_signals/H5)   : {ovh['persist']:7.3f}")
+    print(f"  = instrumented total          : {instrumented:7.3f}   "
+          f"(+{total_over:.3f} ms, {100.0 * total_over / max(ovh['baseline'], 1e-9):+.0f}%)")
+    print(f"  end-of-run derived sweep      : {fin_ms:7.1f} ms once ({sum(dist.values())} samples)")
+    print(f"\ntrajectory-shape distribution : {dist}")
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__.splitlines()[1])
+    ap.add_argument("--outdir", default="./out", help="where to write logs + report")
+    ap.add_argument("--data-root", default="./data", help="MNIST download location")
+    args = ap.parse_args()
+    os.makedirs(args.outdir, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"device={device}  subset={SUBSET or 'FULL(60k)'}  epochs={EPOCHS}  batch={BATCH}")
+    train(args.outdir, args.data_root, device)
+
+
+if __name__ == "__main__":
+    main()

From e0421fb012a919fe7b5d9ff2375b29e9d8d4e98b Mon Sep 17 00:00:00 2001
From: Alexandru-Andrei Rotaru <rotarualexandruandrei94@gmail.com>
Date: Tue, 7 Jul 2026 18:15:32 +0200
Subject: [PATCH 02/19] examples: tune ledger_flush_max_rows to cut per-step
 signal overhead

The ledger default (flush_max_rows=100) flushes every ~1.5 steps at batch 64,
which dominates the per-step cost. A sweep (batch 64, MNIST, single GPU):

  flush=4    +2086%      flush=512   +440%
  flush=64   +1577%      flush=4096  +210%
  flush=100  +631%       flush=8192  +200%  (H5 on/off no longer matters)

Setting flush_max_rows well above the batch size (here 8192) flushes only a few
times per epoch and cuts the measured overhead ~6x (185 -> 32 ms/step) with the
report unchanged (100% coverage). A value <= batch size is a landmine (it
flushes mid-batch). Flush frequency governs both the persist and the
watched-loss logging cost; H5-persistence on/off is irrelevant once flush is high.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 weightslab/examples/PyTorch/ws-signals-mnist/main.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/weightslab/examples/PyTorch/ws-signals-mnist/main.py b/weightslab/examples/PyTorch/ws-signals-mnist/main.py
index 48a981c9..f5912d96 100644
--- a/weightslab/examples/PyTorch/ws-signals-mnist/main.py
+++ b/weightslab/examples/PyTorch/ws-signals-mnist/main.py
@@ -98,6 +98,14 @@ def forward(self, x):
 LR = 0.01
 MEASURE_STEPS = 50  # steps used for the per-step overhead breakdown
 
+# The ledger flushes to storage every this-many buffered rows. This is the
+# single biggest overhead knob: the default (100) flushes every ~1.5 steps at
+# batch 64, which dominates the per-step cost. Set it well above the batch size
+# so the ledger flushes only a few times per epoch. A value <= batch size is a
+# landmine — it flushes mid-batch, every step. (Here: ~batch*128, a handful of
+# flushes over the run.)
+LEDGER_FLUSH_MAX_ROWS = 8192
+
 # The one metric everything subscribes to. It is the ``signal_name`` of the
 # watched loss below, and the ``subscribe_to`` of every live @wl.signal.
 LOSS_SIGNAL = "train/loss_sample"
@@ -351,6 +359,7 @@ def train(outdir, data_root, device):
         "root_log_dir": log_dir,
         "serving_grpc": False,   # headless: no gRPC / CLI control plane
         "serving_cli": False,
+        "ledger_flush_max_rows": LEDGER_FLUSH_MAX_ROWS,  # see note above
         # Declare the loader here too, otherwise WeightsLab's config sync
         # resets batch_size to its default. Keys map to loader kwargs.
         "data": {"train_loader": {"batch_size": BATCH, "shuffle": True}},

From 9adfe3e1b6a69b21ce02506443a8818e809cb97e Mon Sep 17 00:00:00 2001
From: Alexandru-Andrei Rotaru <rotarualexandruandrei94@gmail.com>
Date: Tue, 7 Jul 2026 20:14:42 +0200
Subject: [PATCH 03/19] feat(signals): vectorized @wl.signal(batched=True) path

The dynamic signal executor calls func(ctx) once per sample in a Python loop,
allocating a SignalContext per sample. For signals that read the ledger this is
catastrophic: each of the B calls does its own query_sample_history (lock +
flush + scan).

Add an opt-in batched contract:
  - BatchSignalContext carries the whole batch as arrays (sample_ids,
    subscribed_values) and exposes .history(name) as ONE query_per_sample for
    the batch (not one per sample).
  - executor branches on meta['batched']: build one context, call func once,
    expect a length-B array back.
  - fully backward compatible; the per-sample path is unchanged.

Measured (batch 64, MNIST, single GPU, flush 8192):
  L1 compute-only (subscribed_value):  31 -> 18 ms/step  (-43%)
  L2 history read (per-sample query):  254 -> 63 ms/step (-75%, 4x)
  subscriber invocations: per-sample 1728 -> batched 27 (one per step)
  values identical to the per-sample path (max abs diff 0).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 weightslab/src.py | 82 ++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 64 insertions(+), 18 deletions(-)

diff --git a/weightslab/src.py b/weightslab/src.py
index aa5a39b0..0c1de1dd 100644
--- a/weightslab/src.py
+++ b/weightslab/src.py
@@ -164,6 +164,36 @@ def is_dynamic(self) -> bool:
         return self.subscribed_value is not None
 
 
+class BatchSignalContext:
+    """Batched context for a vectorized ``@wl.signal(batched=True)``.
+
+    Instead of one :class:`SignalContext` per sample, a batched signal receives
+    the whole batch at once: ``sample_ids`` and ``subscribed_values`` are arrays
+    of length B, so the signal computes over all samples with vector ops and
+    returns one array of length B. It also exposes **batched** ledger reads
+    (:meth:`history` runs a single query for the whole batch instead of one query
+    per sample), which is where the large speed-ups come from.
+    """
+    def __init__(self, sample_ids, subscribed_values, logger=None, dataframe=None, origin=None):
+        self.sample_ids = [int(s) for s in sample_ids]          # length B
+        self.subscribed_values = np.asarray(subscribed_values, dtype=float)  # (B,)
+        self.logger = logger
+        self.dataframe = dataframe
+        self.origin = origin
+
+    def history(self, signal_name):
+        """Per-sample history of *signal_name* for every sample in the batch, in
+        a SINGLE ledger query. Returns ``{sample_id: [values in step order]}``
+        (empty list for samples with no history yet)."""
+        out = {s: [] for s in self.sample_ids}
+        if self.logger is None:
+            return out
+        # query_per_sample accepts a list of ids -> one scan for the whole batch.
+        for sid, step, val, _ in self.logger.query_per_sample(signal_name, sample_ids=self.sample_ids):
+            out.setdefault(int(sid), []).append(val)   # rows already ordered by seq (= step order)
+        return out
+
+
 # #####################################################################################################################
 # WEIGHTSLAB INTERNAL FUNCTIONS FOR LOGGING, SIGNAL EXTRACTION, WRAPPING, ETC. (not typically called directly by users)
 # #####################################################################################################################
@@ -626,27 +656,43 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw):
                          continue
 
                      try:
-                         batch_res = {}
-                         for i, uid in enumerate(ids_np):
-                             # Generic 'value' argument
-                             val = float(val_vec[i])
-
-                             # Unified Context Pattern
-                             ctx = SignalContext(
-                                 sample_id=int(uid),
-                                 subscribed_value=val,
+                         if meta.get('batched'):
+                             # Vectorized path: build ONE context for the whole
+                             # batch and call the signal once. It returns a length-B
+                             # array. Avoids B Python calls + B SignalContext allocs,
+                             # and lets the signal do batched ledger reads.
+                             bctx = BatchSignalContext(
+                                 sample_ids=[int(u) for u in ids_np],
+                                 subscribed_values=[float(v) for v in val_vec],
                                  logger=ledgers.get_logger(),
                                  dataframe=df_proxy,
-                                 origin=kwargs.get('origin', 'train')
+                                 origin=kwargs.get('origin', 'train'),
                              )
-                             try:
-                                 res = func(ctx) # Compute per sample result with unified context
-                             except TypeError:
-                                 # Fallback for legacy subscriber functions
-                                 res = func(sample_id=int(uid), value=val, dataframe=df_proxy)
-
-                             batch_res[uid] = res
-                         signal_value = list(batch_res.values())
+                             res = func(bctx)
+                             res = res.tolist() if hasattr(res, 'tolist') else list(res)
+                             signal_value = [float(x) for x in res]
+                         else:
+                             batch_res = {}
+                             for i, uid in enumerate(ids_np):
+                                 # Generic 'value' argument
+                                 val = float(val_vec[i])
+
+                                 # Unified Context Pattern
+                                 ctx = SignalContext(
+                                     sample_id=int(uid),
+                                     subscribed_value=val,
+                                     logger=ledgers.get_logger(),
+                                     dataframe=df_proxy,
+                                     origin=kwargs.get('origin', 'train')
+                                 )
+                                 try:
+                                     res = func(ctx) # Compute per sample result with unified context
+                                 except TypeError:
+                                     # Fallback for legacy subscriber functions
+                                     res = func(sample_id=int(uid), value=val, dataframe=df_proxy)
+
+                                 batch_res[uid] = res
+                             signal_value = list(batch_res.values())
                          dynamic_updates[name] = signal_value
                          if dynamic_updates and meta.get('log', True):
                              logger.debug(f"Dynamic updates computed for signal '{reg_name}': {list(dynamic_updates.keys())}")

From e8b0384363dd72517f8aacf07caa1463df1fcd51 Mon Sep 17 00:00:00 2001
From: Alexandru-Andrei Rotaru <rotarualexandruandrei94@gmail.com>
Date: Tue, 7 Jul 2026 20:27:11 +0200
Subject: [PATCH 04/19] feat(signals): multi-input signals via ctx.latest() /
 bctx.latest()

A dynamic signal subscribes to ONE trigger metric, but can now ingest any number
of OTHER signals by reading their current value from the ledger:

  @wl.signal(name="sig/hardness", subscribe_to="train/loss_sample", batched=True)
  def hardness(bctx):
      return bctx.subscribed_values * bctx.latest("sig/entropy") \
             * (1.0 - bctx.latest("sig/confidence"))

- BatchSignalContext.latest(name): most recent value of another signal for the
  whole batch in ONE query -> (B,) array.
- SignalContext.latest(name): same for a single sample.

Requirements: the ingested signals must be logged (a watched metric, another
@wl.signal, or save_signals(..., log=True)) and written earlier in the step than
the trigger, so latest() sees this step's values.

Verified: hardness == loss*entropy*(1-confidence) to 2e-7 over 640 rows.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 weightslab/src.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/weightslab/src.py b/weightslab/src.py
index 0c1de1dd..06b01f80 100644
--- a/weightslab/src.py
+++ b/weightslab/src.py
@@ -163,6 +163,15 @@ def is_dynamic(self) -> bool:
         """True if running during training (triggered by a metric)."""
         return self.subscribed_value is not None
 
+    def latest(self, signal_name, default=float("nan")):
+        """Most recent value of ANOTHER signal for THIS sample. Lets a signal
+        ingest several other signals: read each with ``ctx.latest(name)`` and
+        combine. Returns *default* if the signal has no value yet."""
+        if self.logger is None:
+            return default
+        rows = self.logger.query_per_sample(signal_name, sample_ids=[self.sample_id])
+        return rows[-1][2] if rows else default   # rows ordered by seq -> last is newest
+
 
 class BatchSignalContext:
     """Batched context for a vectorized ``@wl.signal(batched=True)``.
@@ -193,6 +202,18 @@ def history(self, signal_name):
             out.setdefault(int(sid), []).append(val)   # rows already ordered by seq (= step order)
         return out
 
+    def latest(self, signal_name, default=float("nan")):
+        """Most recent value of ANOTHER signal for each sample in the batch, in a
+        single batched query. ``(B,)`` array aligned to ``self.sample_ids``. Use
+        this to build a signal that ingests several other signals — call it once
+        per input signal, then combine the arrays with vector ops.
+        """
+        last = {}
+        if self.logger is not None:
+            for sid, step, val, _ in self.logger.query_per_sample(signal_name, sample_ids=self.sample_ids):
+                last[int(sid)] = val   # rows ordered by seq -> last assignment wins
+        return np.array([last.get(s, default) for s in self.sample_ids], dtype=float)
+
 
 # #####################################################################################################################
 # WEIGHTSLAB INTERNAL FUNCTIONS FOR LOGGING, SIGNAL EXTRACTION, WRAPPING, ETC. (not typically called directly by users)

From d0605b87f8ebb885ca8fea52989534296df30456 Mon Sep 17 00:00:00 2001
From: Alexandru-Andrei Rotaru <rotarualexandruandrei94@gmail.com>
Date: Tue, 7 Jul 2026 20:40:47 +0200
Subject: [PATCH 05/19] feat(signals): enforce fresh ingests
 (StaleSignalError), stop swallowing errors
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A signal that ingests other signals (subscribes to one trigger, reads others via
ctx.latest) could silently read STALE or missing inputs if they weren't logged,
or were written after the trigger. Make that a loud error instead:

- ctx.step / bctx.step carry the trigger's step.
- ctx.latest(name, require_fresh=True) / bctx.latest(...) raise StaleSignalError
  unless the ingested signal has a value at the current step.
- @wl.signal(..., ingests=["sig/a","sig/b"]) declares inputs; the executor
  validates their freshness for the whole batch BEFORE calling the signal.
- StaleSignalError now propagates out of the subscriber dispatch (previously ALL
  subscriber exceptions were debug-logged and swallowed — silent failures).

Verified: correct order (inputs logged before trigger) runs; wrong order raises
StaleSignalError naming the signal, step, and affected samples.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 weightslab/__init__.py |  4 ++-
 weightslab/src.py      | 70 ++++++++++++++++++++++++++++++++++++------
 2 files changed, 63 insertions(+), 11 deletions(-)

diff --git a/weightslab/__init__.py b/weightslab/__init__.py
index e41340b6..1aa33577 100644
--- a/weightslab/__init__.py
+++ b/weightslab/__init__.py
@@ -11,7 +11,7 @@
 import logging
 import threading
 
-from .src import watch_or_edit, start_training, serve, keep_serving, save_signals, save_instance_signals, save_group_signals, tag_samples, register_categorical_tag, set_categorical_tag, discard_samples, get_samples_by_tag, get_discarded_samples, signal, eval_fn, compute_signals, SignalContext, clear_all, run_pending_evaluation, trigger_pending_evaluation_async, query_signal_history, query_sample_history, query_instance_history, write_history, write_dataframe, get_current_experiment_hash, pointcloud_thumbnail, pointcloud_boxes
+from .src import watch_or_edit, start_training, serve, keep_serving, save_signals, save_instance_signals, save_group_signals, tag_samples, register_categorical_tag, set_categorical_tag, discard_samples, get_samples_by_tag, get_discarded_samples, signal, eval_fn, compute_signals, SignalContext, BatchSignalContext, StaleSignalError, clear_all, run_pending_evaluation, trigger_pending_evaluation_async, query_signal_history, query_sample_history, query_instance_history, write_history, write_dataframe, get_current_experiment_hash, pointcloud_thumbnail, pointcloud_boxes
 from .backend.ledgers import GLOBAL_LEDGER as ledger
 from .art import _BANNER
 from .utils.logs import setup_logging, set_log_directory, is_main_process
@@ -148,6 +148,8 @@ def _clean(v: str) -> str:
   	"get_samples_by_tag",
     "get_discarded_samples",
     "SignalContext",
+    "BatchSignalContext",
+    "StaleSignalError",
     "clear_all",
   	"seed_everything",
     "run_pending_evaluation",
diff --git a/weightslab/src.py b/weightslab/src.py
index 06b01f80..57d5f822 100644
--- a/weightslab/src.py
+++ b/weightslab/src.py
@@ -78,18 +78,28 @@ def _rebind_caller_local(original_obj: Any, new_obj: Any) -> None:
 _EVAL_WORKER_THREAD: Optional[threading.Thread] = None
 
 
+class StaleSignalError(RuntimeError):
+    """Raised when a signal ingests another signal that is not *fresh* — i.e. the
+    ingested signal has no value at the current step (it was not logged, or was
+    written after the trigger this step). See ``SignalContext.latest`` /
+    ``BatchSignalContext.latest`` (``require_fresh=True``) and
+    ``@wl.signal(ingests=[...])``."""
+    pass
+
+
 class SignalContext:
     """
     Unified context object for WeightsLab signals.
     Carries all available metadata for a single sample during computation.
     """
-    def __init__(self, sample_id, dataframe, data=None, subscribed_value=None, logger=None, origin=None):
+    def __init__(self, sample_id, dataframe, data=None, subscribed_value=None, logger=None, origin=None, step=None):
         self.sample_id = sample_id
         self.dataframe = dataframe
         self.data = data
         self.subscribed_value = subscribed_value
         self.logger = logger
         self.origin = origin
+        self.step = step          # step the trigger fired at (for freshness checks)
 
     @property
     def image(self) -> Optional[np.ndarray]:
@@ -163,13 +173,23 @@ def is_dynamic(self) -> bool:
         """True if running during training (triggered by a metric)."""
         return self.subscribed_value is not None
 
-    def latest(self, signal_name, default=float("nan")):
+    def latest(self, signal_name, default=float("nan"), require_fresh=False):
         """Most recent value of ANOTHER signal for THIS sample. Lets a signal
         ingest several other signals: read each with ``ctx.latest(name)`` and
-        combine. Returns *default* if the signal has no value yet."""
+        combine. Returns *default* if the signal has no value yet.
+
+        With ``require_fresh=True`` this raises if the ingested signal has no
+        value at the current step (``self.step``) — i.e. it was not logged, or
+        was written after (not before) the trigger this step.
+        """
         if self.logger is None:
             return default
         rows = self.logger.query_per_sample(signal_name, sample_ids=[self.sample_id])
+        if require_fresh and (not rows or self.step is None or rows[-1][1] != self.step):
+            raise StaleSignalError(
+                f"ingested signal '{signal_name}' is not fresh at step {self.step} for "
+                f"sample {self.sample_id} (log it with log=True and write it BEFORE the "
+                f"subscribing signal fires).")
         return rows[-1][2] if rows else default   # rows ordered by seq -> last is newest
 
 
@@ -183,12 +203,13 @@ class BatchSignalContext:
     (:meth:`history` runs a single query for the whole batch instead of one query
     per sample), which is where the large speed-ups come from.
     """
-    def __init__(self, sample_ids, subscribed_values, logger=None, dataframe=None, origin=None):
+    def __init__(self, sample_ids, subscribed_values, logger=None, dataframe=None, origin=None, step=None):
         self.sample_ids = [int(s) for s in sample_ids]          # length B
         self.subscribed_values = np.asarray(subscribed_values, dtype=float)  # (B,)
         self.logger = logger
         self.dataframe = dataframe
         self.origin = origin
+        self.step = step          # step the trigger fired at (for freshness checks)
 
     def history(self, signal_name):
         """Per-sample history of *signal_name* for every sample in the batch, in
@@ -202,17 +223,29 @@ def history(self, signal_name):
             out.setdefault(int(sid), []).append(val)   # rows already ordered by seq (= step order)
         return out
 
-    def latest(self, signal_name, default=float("nan")):
+    def latest(self, signal_name, default=float("nan"), require_fresh=False):
         """Most recent value of ANOTHER signal for each sample in the batch, in a
         single batched query. ``(B,)`` array aligned to ``self.sample_ids``. Use
         this to build a signal that ingests several other signals — call it once
         per input signal, then combine the arrays with vector ops.
+
+        With ``require_fresh=True`` this raises unless EVERY sample has a value of
+        *signal_name* at the current step (``self.step``) — catching a stale
+        ingest (input not logged, or written after the trigger this step).
         """
         last = {}
         if self.logger is not None:
             for sid, step, val, _ in self.logger.query_per_sample(signal_name, sample_ids=self.sample_ids):
-                last[int(sid)] = val   # rows ordered by seq -> last assignment wins
-        return np.array([last.get(s, default) for s in self.sample_ids], dtype=float)
+                last[int(sid)] = (step, val)   # rows ordered by seq -> last assignment wins
+        if require_fresh:
+            stale = [s for s in self.sample_ids
+                     if s not in last or self.step is None or last[s][0] != self.step]
+            if stale:
+                raise StaleSignalError(
+                    f"ingested signal '{signal_name}' is not fresh at step {self.step} for "
+                    f"{len(stale)}/{len(self.sample_ids)} sample(s) (e.g. {stale[:3]}). Log it "
+                    f"with log=True and write it BEFORE the subscribing signal fires.")
+        return np.array([last.get(s, (None, default))[1] for s in self.sample_ids], dtype=float)
 
 
 # #####################################################################################################################
@@ -677,6 +710,8 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw):
                          continue
 
                      try:
+                         _lg = ledgers.get_logger()
+                         ingests = meta.get('ingests') or []
                          if meta.get('batched'):
                              # Vectorized path: build ONE context for the whole
                              # batch and call the signal once. It returns a length-B
@@ -685,14 +720,26 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw):
                              bctx = BatchSignalContext(
                                  sample_ids=[int(u) for u in ids_np],
                                  subscribed_values=[float(v) for v in val_vec],
-                                 logger=ledgers.get_logger(),
+                                 logger=_lg,
                                  dataframe=df_proxy,
                                  origin=kwargs.get('origin', 'train'),
+                                 step=step,
                              )
+                             # Enforce that declared ingested signals are fresh.
+                             for dep in ingests:
+                                 bctx.latest(dep, require_fresh=True)
                              res = func(bctx)
                              res = res.tolist() if hasattr(res, 'tolist') else list(res)
                              signal_value = [float(x) for x in res]
                          else:
+                             # Validate declared ingests once for the whole batch.
+                             if ingests:
+                                 _vctx = BatchSignalContext(
+                                     sample_ids=[int(u) for u in ids_np],
+                                     subscribed_values=[float(v) for v in val_vec],
+                                     logger=_lg, step=step)
+                                 for dep in ingests:
+                                     _vctx.latest(dep, require_fresh=True)
                              batch_res = {}
                              for i, uid in enumerate(ids_np):
                                  # Generic 'value' argument
@@ -702,9 +749,10 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw):
                                  ctx = SignalContext(
                                      sample_id=int(uid),
                                      subscribed_value=val,
-                                     logger=ledgers.get_logger(),
+                                     logger=_lg,
                                      dataframe=df_proxy,
-                                     origin=kwargs.get('origin', 'train')
+                                     origin=kwargs.get('origin', 'train'),
+                                     step=step,
                                  )
                                  try:
                                      res = func(ctx) # Compute per sample result with unified context
@@ -718,6 +766,8 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw):
                          if dynamic_updates and meta.get('log', True):
                              logger.debug(f"Dynamic updates computed for signal '{reg_name}': {list(dynamic_updates.keys())}")
                              _log_signal(sum(signal_value)/len(signal_value), signal_value, name, step=step, **kwargs) # Log custom subscribed signals
+                     except StaleSignalError:
+                         raise  # correctness enforcement: surface, don't swallow
                      except Exception as e:
                          logger.debug(f"Dynamic signal {name} failed: {e}")
                          pass # User function error, skip

From ea3e9ea20257dae65eebbd823fe0bbd190d55286 Mon Sep 17 00:00:00 2001
From: Alexandru-Andrei Rotaru <rotarualexandruandrei94@gmail.com>
Date: Tue, 7 Jul 2026 21:07:00 +0200
Subject: [PATCH 06/19] =?UTF-8?q?feat(signals):=20reactive=20inputs=3D[...?=
 =?UTF-8?q?]=20=E2=80=94=20order-independent=20multi-input=20+=20chaining?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Unify subscribe_to and ingests into one dependency declaration. A signal with
inputs=[a, b, c] fires when ALL its inputs are present at a step, regardless of
the order they were logged — eliminating the ordering footgun instead of just
detecting it. Fired signals are themselves inputs, so signals chain.

- _react_dependents(): iterative, dedup'd (fires each signal at most once/step),
  reads inputs fresh from the logger, persists results (log=True), and cascades
  via a work-queue (no re-entrant dispatch).
- wired into wrappered_fwd (watched metric) and save_signals(log=True) so ANY
  logged signal can satisfy a dependent — that's what makes it order-independent.
- subscribe_to keeps the legacy dispatch (handles non-per-sample metrics); the
  two paths coexist without double-firing (inputs= vs subscribe_to).

Verified (batch 64): order-independent (entropy logged AFTER loss still fires the
loss*entropy signal), exactly-once/step, chaining (B<-A<-entropy), values correct,
and the existing subscribe_to signal still fires.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 weightslab/src.py | 119 +++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 118 insertions(+), 1 deletion(-)

diff --git a/weightslab/src.py b/weightslab/src.py
index 57d5f822..2e2dff30 100644
--- a/weightslab/src.py
+++ b/weightslab/src.py
@@ -71,6 +71,93 @@ def _rebind_caller_local(original_obj: Any, new_obj: Any) -> None:
 DATAFRAME_M = None
 # Global registry for custom signals
 _REGISTERED_SIGNALS = {}
+_REACTIVE_FIRED = {}   # signal_name -> step it last fired at (per-step dedup)
+
+
+def _gather_inputs_fresh(lg, inputs, ids, step):
+    """Return ``{name: (B,) array}`` if EVERY input signal has a value at *step*
+    for EVERY sample id, else ``None`` (not all inputs present yet)."""
+    cols = {}
+    for inp in inputs:
+        last = {}
+        for sid, s, val, _ in lg.query_per_sample(inp, sample_ids=ids):
+            last[int(sid)] = (s, val)
+        for sid in ids:
+            if sid not in last or last[sid][0] != step:
+                return None
+        cols[inp] = np.array([last[sid][1] for sid in ids], dtype=float)
+    return cols
+
+
+def _react_dependents(seed_names, batch_ids, step, origin='train'):
+    """Reactive signal firing. When the signals in *seed_names* were just logged,
+    fire any ``@wl.signal(inputs=[...])`` whose inputs are now ALL present at
+    *step* — order-independently, regardless of which input arrived last. Each
+    fired signal is itself a seed (chaining) and fires at most once per step.
+    Inputs are read from the logger, so ingested signals must be logged."""
+    if batch_ids is None:
+        return
+    try:
+        ids = [int(u) for u in (batch_ids.detach().cpu().numpy()
+                                if hasattr(batch_ids, 'detach') else batch_ids)]
+    except Exception:
+        return
+    lg = get_logger()
+    if lg is None:
+        return
+    try:
+        df_proxy = get_dataframe()
+    except Exception:
+        df_proxy = None
+
+    def dependents_of(nm):
+        for name, func in list(_REGISTERED_SIGNALS.items()):
+            meta = getattr(func, '_wl_signal_meta', {})
+            inputs = meta.get('inputs')
+            if inputs and nm in inputs:
+                yield name, func, meta, inputs
+
+    queue = list(seed_names)
+    while queue:
+        logged = queue.pop()
+        for name, func, meta, inputs in dependents_of(logged):
+            if _REACTIVE_FIRED.get(name) == step:
+                continue
+            every = meta.get('compute_every_n_steps', 1) or 1
+            if step % every != 0:
+                continue
+            cols = _gather_inputs_fresh(lg, inputs, ids, step)
+            if cols is None:
+                continue  # inputs not all fresh yet — a later log will re-trigger
+            _REACTIVE_FIRED[name] = step
+            try:
+                if meta.get('batched'):
+                    bctx = BatchSignalContext(sample_ids=ids, subscribed_values=cols[inputs[0]],
+                                              logger=lg, dataframe=df_proxy, origin=origin, step=step)
+                    bctx.inputs = cols
+                    res = func(bctx)
+                    vals = res.tolist() if hasattr(res, 'tolist') else list(res)
+                else:
+                    vals = []
+                    for i, sid in enumerate(ids):
+                        ctx = SignalContext(sample_id=sid, subscribed_value=float(cols[inputs[0]][i]),
+                                            logger=lg, dataframe=df_proxy, origin=origin, step=step)
+                        ctx.inputs = {k: float(cols[k][i]) for k in cols}
+                        vals.append(func(ctx))
+                vals = [float(x) for x in vals]
+            except StaleSignalError:
+                raise
+            except Exception as e:
+                logger.debug(f"reactive signal '{name}' failed: {e}")
+                continue
+            # Persist (log so downstream reactive signals can read it). _react=False:
+            # the queue below drives chaining, avoiding re-entrant dispatch.
+            try:
+                save_signals(signals={name: vals}, batch_ids=ids,
+                             log=meta.get('log', True), _react=False)
+            except Exception as e:
+                logger.debug(f"reactive persist '{name}' failed: {e}")
+            queue.append(name)   # its value may satisfy further reactive signals
 
 # Evaluation function registered via the @wl.eval_fn decorator.
 _REGISTERED_EVAL_FN: Optional[Any] = None
@@ -796,6 +883,16 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw):
             log=False # Already logged above, no need to log again in save_signals; set to False to avoid duplicate logging if save_signals is called separately without logging
         )
 
+    # Reactive signals: the watched metric was just logged (above), so it may
+    # satisfy an inputs=[...] signal that depends on it.
+    if not per_instance and batch_ids is not None:
+        try:
+            _react_dependents([reg_name], batch_ids, step, origin or 'train')
+        except StaleSignalError:
+            raise
+        except Exception as e:
+            logger.debug(f"reactive dispatch after {reg_name} failed: {e}")
+
     # Return the original output (dict for per-instance losses so caller can
     # use out['batch'] for backward, tensor for standard per-sample losses).
     return out_original if isinstance(out_original, dict) else out
@@ -1274,6 +1371,13 @@ def decorator(func):
         func._wl_signal_meta = kwargs
         func._wl_signal_meta['subscribe_to'] = subscribe_to
         func._wl_signal_meta['compute_every_n_steps'] = compute_every_n_steps
+        # Reactive dependency set. ``inputs=[...]`` is the unified form: the signal
+        # fires when ALL its inputs are present at a step, order-independently
+        # (subsumes subscribe_to + ingests). A single-input signal is the alias
+        # ``inputs=["x"]`` == ``subscribe_to="x"``. Signals that use the legacy
+        # ``subscribe_to`` keyword instead stay on the legacy dispatch.
+        _inp = kwargs.get('inputs')
+        func._wl_signal_meta['inputs'] = list(_inp) if _inp else None
         func._wl_signal_name = reg_name
 
         return func
@@ -1719,7 +1823,8 @@ def save_signals(
     targets: th.Tensor | np.ndarray | dict = None,
     preds: th.Tensor | np.ndarray | dict = None,
     step: int | None = None,
-    log: bool = False
+    log: bool = False,
+    _react: bool = True,
 ):
     """Save **per-sample** statistics to the tracked dataset.
 
@@ -1867,6 +1972,18 @@ def expand_dim(x):
         step=step
     )
 
+    # Reactive signals: these just-logged signals may satisfy an inputs=[...]
+    # signal. Only logged (queryable) signals can be inputs. _react=False on the
+    # reactive persist path prevents re-entrant dispatch.
+    if _react and log and isinstance(signals, dict):
+        try:
+            _react_dependents(list(signals.keys()), batch_ids, step,
+                              get_active_origin() or 'train')
+        except StaleSignalError:
+            raise
+        except Exception as e:
+            logger.debug(f"reactive dispatch after save_signals failed: {e}")
+
 
 def save_instance_signals(
     signals: dict,

From 093e6c41cecb78501fc89c3c8df3bdda01e65c99 Mon Sep 17 00:00:00 2001
From: Alexandru-Andrei Rotaru <rotarualexandruandrei94@gmail.com>
Date: Tue, 7 Jul 2026 21:27:11 +0200
Subject: [PATCH 07/19] feat(signals): in-core signal-worker thread
 (ledger_signal_worker)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Move reactive signal computation off the training thread. When
hyperparam ledger_signal_worker=True, the reactive dispatch is handed to a
background worker (WL-Signal-Worker) instead of running inline; the train thread
just detaches + enqueues (seed_names, ids, step). The seed signals are logged
synchronously, so the worker only defers the derived compute + persistence.

- _dispatch_or_enqueue() at both injection sites; drain_signals() (exported)
  joins the queue; write_dataframe() auto-drains so the report is complete.
- _gather_inputs_fresh selects the value AT the job's step (not "latest"), so a
  lagging worker still gathers that step's inputs.
- save_signals now honors an explicit step= (as its docstring says) instead of
  letting the live model age override it — required so reactive backfill logs
  derived values at the job's step; otherwise chained signals can't find them.

Verified worker on vs off: exactly-once/step, order-independent, chaining
(B<-A<-entropy) all identical; drain makes the report complete.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 weightslab/__init__.py |  3 +-
 weightslab/src.py      | 99 +++++++++++++++++++++++++++++++++++++-----
 2 files changed, 89 insertions(+), 13 deletions(-)

diff --git a/weightslab/__init__.py b/weightslab/__init__.py
index 1aa33577..61250afe 100644
--- a/weightslab/__init__.py
+++ b/weightslab/__init__.py
@@ -11,7 +11,7 @@
 import logging
 import threading
 
-from .src import watch_or_edit, start_training, serve, keep_serving, save_signals, save_instance_signals, save_group_signals, tag_samples, register_categorical_tag, set_categorical_tag, discard_samples, get_samples_by_tag, get_discarded_samples, signal, eval_fn, compute_signals, SignalContext, BatchSignalContext, StaleSignalError, clear_all, run_pending_evaluation, trigger_pending_evaluation_async, query_signal_history, query_sample_history, query_instance_history, write_history, write_dataframe, get_current_experiment_hash, pointcloud_thumbnail, pointcloud_boxes
+from .src import watch_or_edit, start_training, serve, keep_serving, save_signals, save_instance_signals, save_group_signals, tag_samples, register_categorical_tag, set_categorical_tag, discard_samples, get_samples_by_tag, get_discarded_samples, signal, eval_fn, compute_signals, SignalContext, BatchSignalContext, StaleSignalError, drain_signals, clear_all, run_pending_evaluation, trigger_pending_evaluation_async, query_signal_history, query_sample_history, query_instance_history, write_history, write_dataframe, get_current_experiment_hash, pointcloud_thumbnail, pointcloud_boxes
 from .backend.ledgers import GLOBAL_LEDGER as ledger
 from .art import _BANNER
 from .utils.logs import setup_logging, set_log_directory, is_main_process
@@ -150,6 +150,7 @@ def _clean(v: str) -> str:
     "SignalContext",
     "BatchSignalContext",
     "StaleSignalError",
+    "drain_signals",
     "clear_all",
   	"seed_everything",
     "run_pending_evaluation",
diff --git a/weightslab/src.py b/weightslab/src.py
index 2e2dff30..deb62e5e 100644
--- a/weightslab/src.py
+++ b/weightslab/src.py
@@ -14,6 +14,7 @@
 import tempfile
 import functools
 import threading
+import queue as _queue
 import traceback
 import numpy as np
 import torch as th
@@ -75,17 +76,20 @@ def _rebind_caller_local(original_obj: Any, new_obj: Any) -> None:
 
 
 def _gather_inputs_fresh(lg, inputs, ids, step):
-    """Return ``{name: (B,) array}`` if EVERY input signal has a value at *step*
-    for EVERY sample id, else ``None`` (not all inputs present yet)."""
+    """Return ``{name: (B,) array}`` if EVERY input signal has a value AT *step*
+    for EVERY sample id, else ``None``. Selects the value logged *at* `step`
+    (not merely the latest) so a lagging signal-worker job for an older step
+    still gathers that step's inputs correctly."""
     cols = {}
     for inp in inputs:
-        last = {}
+        at = {}
         for sid, s, val, _ in lg.query_per_sample(inp, sample_ids=ids):
-            last[int(sid)] = (s, val)
+            if s == step:
+                at[int(sid)] = val
         for sid in ids:
-            if sid not in last or last[sid][0] != step:
+            if sid not in at:
                 return None
-        cols[inp] = np.array([last[sid][1] for sid in ids], dtype=float)
+        cols[inp] = np.array([at[sid] for sid in ids], dtype=float)
     return cols
 
 
@@ -153,12 +157,77 @@ def dependents_of(nm):
             # Persist (log so downstream reactive signals can read it). _react=False:
             # the queue below drives chaining, avoiding re-entrant dispatch.
             try:
-                save_signals(signals={name: vals}, batch_ids=ids,
+                # step=step (not the live model step): attribute the derived
+                # value to the step it was derived from. Critical when the signal
+                # worker lags behind training, else chained signals can't find it.
+                save_signals(signals={name: vals}, batch_ids=ids, step=step,
                              log=meta.get('log', True), _react=False)
             except Exception as e:
                 logger.debug(f"reactive persist '{name}' failed: {e}")
             queue.append(name)   # its value may satisfy further reactive signals
 
+
+# --- optional signal-worker thread: run reactive dispatch off the train thread -
+_SIGNAL_WORKER = {"q": None, "thread": None}
+
+
+def _signal_worker_enabled():
+    try:
+        return bool(get_hyperparams().get("ledger_signal_worker", False))
+    except Exception:
+        return False
+
+
+def _ensure_signal_worker():
+    if _SIGNAL_WORKER["thread"] is not None and _SIGNAL_WORKER["thread"].is_alive():
+        return _SIGNAL_WORKER["q"]
+    q = _queue.Queue()
+
+    def _worker():
+        while True:
+            job = q.get()
+            if job is None:
+                q.task_done()
+                break
+            seed_names, ids, step, origin = job
+            try:
+                _react_dependents(seed_names, ids, step, origin)
+            except Exception as e:
+                logger.debug(f"signal worker job failed: {e}")
+            finally:
+                q.task_done()
+
+    t = threading.Thread(target=_worker, name="WL-Signal-Worker", daemon=True)
+    t.start()
+    _SIGNAL_WORKER["q"] = q
+    _SIGNAL_WORKER["thread"] = t
+    return q
+
+
+def _dispatch_or_enqueue(seed_names, batch_ids, step, origin):
+    """Reactive dispatch: inline on the train thread, or handed to the signal
+    worker thread when ``ledger_signal_worker`` is set. The seed signals are
+    already logged synchronously, so the worker only defers the derived compute
+    + persistence — the inputs it reads are already in the ledger."""
+    if not _signal_worker_enabled():
+        _react_dependents(seed_names, batch_ids, step, origin)
+        return
+    try:
+        ids = [int(u) for u in (batch_ids.detach().cpu().numpy()
+                                if hasattr(batch_ids, 'detach') else batch_ids)]
+    except Exception:
+        return
+    _ensure_signal_worker().put((list(seed_names), ids, step, origin))
+
+
+def drain_signals():
+    """Block until the signal-worker thread has processed all queued reactive
+    jobs. Called automatically by :func:`write_dataframe`; call it manually
+    before reading signals mid-run when ``ledger_signal_worker`` is enabled."""
+    q = _SIGNAL_WORKER["q"]
+    if q is not None:
+        q.join()
+
 # Evaluation function registered via the @wl.eval_fn decorator.
 _REGISTERED_EVAL_FN: Optional[Any] = None
 _EVAL_WORKER_LOCK = threading.Lock()
@@ -887,7 +956,7 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw):
     # satisfy an inputs=[...] signal that depends on it.
     if not per_instance and batch_ids is not None:
         try:
-            _react_dependents([reg_name], batch_ids, step, origin or 'train')
+            _dispatch_or_enqueue([reg_name], batch_ids, step, origin or 'train')
         except StaleSignalError:
             raise
         except Exception as e:
@@ -1888,8 +1957,10 @@ def save_signals(
     if DATAFRAME_M is None:
         DATAFRAME_M = get_dataframe()
 
-    # Get current model step
-    step = _get_step(step=step)
+    # Honor an explicit step (as documented); only infer from the model when it
+    # is omitted. Reactive backfill on the signal-worker thread depends on this —
+    # it logs derived values at the job's step, not the live (ahead) model step.
+    step = step if step is not None else _get_step()
 
     # Log if requested for each signals
     if log:
@@ -1977,8 +2048,8 @@ def expand_dim(x):
     # reactive persist path prevents re-entrant dispatch.
     if _react and log and isinstance(signals, dict):
         try:
-            _react_dependents(list(signals.keys()), batch_ids, step,
-                              get_active_origin() or 'train')
+            _dispatch_or_enqueue(list(signals.keys()), batch_ids, step,
+                                 get_active_origin() or 'train')
         except StaleSignalError:
             raise
         except Exception as e:
@@ -3801,6 +3872,10 @@ def write_dataframe(
     import os as _os
     import hashlib as _hashlib
 
+    # If reactive signals run on the worker thread, process every queued job
+    # before reading, so the report includes the latest steps' derived signals.
+    drain_signals()
+
     logger.debug(
         "write_dataframe called: path=%r, format=%r, columns=%r, "
         "sample_id=%r, instance_id=%r",

From b8e349246ee83b6c1e14270293b16657d80f0662 Mon Sep 17 00:00:00 2001
From: Alexandru Rotaru <rotarualexandruandrei94@gmail.com>
Date: Wed, 8 Jul 2026 13:42:46 +0200
Subject: [PATCH 08/19] perf(signals): step-scoped query cache + value-at-step
 gather; cycle warning
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Reactive-signal overhead was dominated by redundant, growing ledger reads.

Logger query cache (backend/logger.py):
- Memoize query_per_sample via functools.lru_cache keyed
  (signal, ids, hash, version[signal]); _stage_sample_row bumps the
  per-signal version (per-signal, not global, so persisting a derived
  signal doesn't invalidate the loss its siblings are still reading).
- Step-scoped: clear both caches when the training step advances. The
  loader reshuffles ids and the version bumps every step, so a key never
  recurs across steps — cross-step entries are pure dead weight. Bounds
  the cache to one step (currsize ~6 vs ~1600 accumulating), killing the
  churn while keeping the intra-step reuse (10 signals sharing the loss
  -> 1 query).
- query_per_sample_at_step(name, ids, step): O(batch) value-at-step read
  (WHERE step=? in-engine) behind its own step-cache.

Core (src.py):
- _gather_inputs_fresh uses the value-at-step read -> flat as history
  grows (single query 4.7->13ms/2.8x before; 2.1->2.4ms/1.1x after).
- Circular-dependency detection: _detect_signal_cycles (grey/black DFS
  over the inputs graph) + _warn_on_signal_cycles, lazy+idempotent, fired
  from start_training AND first reactive dispatch so it works headless
  (no wl.serve). Warns, does not raise; cyclic signals just stay dark.

Measured (100-epoch MNIST headless, no wl.serve): per-step time flat
(~157ms) vs rising 170->248ms; total wall 430s->286s (-33%). Reactive
verify, multi-level/multi-parent chain, and cycle tests all pass.
---
 weightslab/backend/logger.py | 97 ++++++++++++++++++++++++++++++++++--
 weightslab/src.py            | 76 ++++++++++++++++++++++++++--
 2 files changed, 165 insertions(+), 8 deletions(-)

diff --git a/weightslab/backend/logger.py b/weightslab/backend/logger.py
index 4e6f95b1..9c14b01f 100644
--- a/weightslab/backend/logger.py
+++ b/weightslab/backend/logger.py
@@ -29,9 +29,11 @@
   staging appends and flushes take the same lock.
 """
 
+import functools
 import json
 import threading
 import time
+from collections import defaultdict
 
 import duckdb
 import pandas as pd
@@ -81,6 +83,31 @@ def __init__(self, register: bool = True, db_path: str = ":memory:") -> None:
         self._stage_sample: list = []
         self._stage_instance: list = []
         self._seq = 0
+
+        # --- per-sample query cache ---------------------------------------
+        # Many consumers in a single step read the SAME (signal, ids): e.g. 10
+        # reactive signals all reading the loss, or a batched context re-reading
+        # a history. Each read is a lock + _flush_stage + DuckDB scan, so N
+        # identical reads cost Nx for one answer. Two memoized readers share one
+        # per-signal version map:
+        #   _qps_cache      -> query_per_sample        (full history)
+        #   _qps_step_cache -> query_per_sample_at_step (values AT one step)
+        # both keyed by (signal, ids, [step,] hash, version[signal]). Staging a
+        # per-sample row bumps ITS version (see _stage_sample_row), so a cached
+        # read is served until that signal changes, then recomputed. Per-signal
+        # (not global) versioning is essential: persisting a derived signal must
+        # NOT invalidate the loss every dependent is still reading this step.
+        #
+        # The caches are STEP-SCOPED: because the loader reshuffles ids and the
+        # version bumps every step, a cache key never recurs across steps — so
+        # cross-step entries are pure dead weight. _stage_sample_row clears both
+        # caches when the training step advances, bounding them to one step's
+        # worth of entries (the intra-step reuse — the actual win — is kept).
+        self._qps_version: dict = defaultdict(int)
+        self._qps_cache_step: int = -1
+        self._qps_cache = functools.lru_cache(maxsize=2048)(self._query_per_sample_uncached)
+        self._qps_step_cache = functools.lru_cache(maxsize=2048)(self._query_per_sample_at_step_uncached)
+
         self._ensure_tables()
         self._restore_runtime_state_from_db()
 
@@ -198,11 +225,30 @@ def _stage_signal_row(self, graph_name, exp_hash, step, metric_value, timestamp,
         self._maybe_autoflush()
 
     def _stage_sample_row(self, graph_name, exp_hash, sample_id, step, value):
+        # Step advanced -> last step's cache entries can never be hit again
+        # (ids reshuffle + version bump make every key unique per step). Drop
+        # them so the cache stays bounded to the current step instead of
+        # accumulating single-use entries.
+        if int(step) > self._qps_cache_step:
+            self._invalidate_qps_cache()
+            self._qps_cache_step = int(step)
         self._stage_sample.append((
             graph_name, exp_hash, str(sample_id), int(step), float(value), self._next_seq(),
         ))
+        # A new per-sample row for this signal invalidates its cached reads
+        # WITHIN the step (an input written after a first read must be re-read).
+        self._qps_version[graph_name] += 1
         self._maybe_autoflush()
 
+    def _invalidate_qps_cache(self) -> None:
+        """Drop all memoized per-sample query results (both the full-history and
+        the at-step readers) and reset versions. Used on step advance and by the
+        bulk delete/clear paths, where rows for many signals change at once and
+        per-signal version bumps would be impractical."""
+        self._qps_cache.cache_clear()
+        self._qps_step_cache.cache_clear()
+        self._qps_version.clear()
+
     def _stage_instance_row(self, graph_name, exp_hash, sample_id, annotation_id, step, value):
         self._stage_instance.append((
             graph_name, exp_hash, str(sample_id), int(annotation_id), int(step),
@@ -244,6 +290,7 @@ def clear_signal_histories(self):
             self._conn.execute("DELETE FROM per_instance")
             self._current_step_buffer.clear()
             self._buffered_step = None
+            self._invalidate_qps_cache()
 
     def _to_float(self, value):
         if isinstance(value, th.Tensor):
@@ -443,6 +490,7 @@ def remove_evaluation_hash(self, eval_hash: str) -> None:
             self._flush_stage()
             self._conn.execute("DELETE FROM signals WHERE experiment_hash = ?", [eval_hash])
             self._conn.execute("DELETE FROM per_sample WHERE experiment_hash = ?", [eval_hash])
+            self._invalidate_qps_cache()
 
         # Drop queued points that reference this hash.
         self._pending_queue = [
@@ -712,19 +760,62 @@ def query_per_sample(self, graph_name: str, sample_ids=None, exp_hash=None):
 
         Returns a list of ``(sample_id, step, value, experiment_hash)`` tuples,
         filtered by *sample_ids* and optionally *exp_hash* (``None`` = all hashes).
+
+        Served from the per-signal query cache (see ``self._qps_cache``): the
+        result is memoized until a new per-sample row for *graph_name* is staged,
+        so repeated identical reads within a step cost one scan, not N. A fresh
+        ``list`` copy is returned each call so callers may mutate it freely.
         """
+        ids_key = tuple(str(s) for s in sample_ids) if sample_ids is not None else None
+        cached = self._qps_cache(graph_name, ids_key, exp_hash, self._qps_version[graph_name])
+        return list(cached)
+
+    def _query_per_sample_uncached(self, graph_name, ids_key, exp_hash, _version):
+        """Actual DuckDB read behind :meth:`query_per_sample`. ``_version`` is a
+        cache-key discriminant only (unused in the body) — it changes when the
+        signal is written, forcing a recompute. Returns a tuple so the cached
+        value stays immutable across callers."""
         with self._lock:
             self._flush_stage()
             params = [graph_name]
             sql = "SELECT sample_id, step, value, experiment_hash FROM per_sample WHERE metric_name = ?"
             sql += self._hash_filter(exp_hash, params)
-            if sample_ids is not None:
+            if ids_key is not None:
                 sql += " AND sample_id IN (SELECT UNNEST(?))"
-                params.append([str(s) for s in sample_ids])
+                params.append(list(ids_key))
             sql += " ORDER BY seq"
             rows = self._conn.execute(sql, params).fetchall()
 
-        return [(sid, int(step), float(val), h) for (sid, step, val, h) in rows]
+        return tuple((sid, int(step), float(val), h) for (sid, step, val, h) in rows)
+
+    def query_per_sample_at_step(self, graph_name: str, sample_ids, step, exp_hash=None):
+        """Per-sample values of *graph_name* at EXACTLY *step* — returns a list
+        of ``(sample_id, value)``. Unlike :meth:`query_per_sample`, the result is
+        O(batch), not O(history): the DuckDB ``WHERE step = ?`` filter is applied
+        in-engine so only the current step's rows are materialized into Python.
+        This keeps the reactive freshness gather flat as training accumulates
+        history (it only ever needs the value at the firing step). Served from
+        the step cache; a fresh ``list`` copy is returned each call.
+        """
+        ids_key = tuple(str(s) for s in sample_ids) if sample_ids is not None else None
+        cached = self._qps_step_cache(graph_name, ids_key, int(step), exp_hash,
+                                      self._qps_version[graph_name])
+        return list(cached)
+
+    def _query_per_sample_at_step_uncached(self, graph_name, ids_key, step, exp_hash, _version):
+        """DuckDB read behind :meth:`query_per_sample_at_step`. ``_version`` is a
+        cache-key discriminant only. Returns an immutable tuple."""
+        with self._lock:
+            self._flush_stage()
+            params = [graph_name, int(step)]
+            sql = "SELECT sample_id, value FROM per_sample WHERE metric_name = ? AND step = ?"
+            sql += self._hash_filter(exp_hash, params)
+            if ids_key is not None:
+                sql += " AND sample_id IN (SELECT UNNEST(?))"
+                params.append(list(ids_key))
+            rows = self._conn.execute(sql, params).fetchall()
+
+        return tuple((sid, float(val)) for (sid, val) in rows)
 
     def query_per_instance(
         self,
diff --git a/weightslab/src.py b/weightslab/src.py
index deb62e5e..b52b94c2 100644
--- a/weightslab/src.py
+++ b/weightslab/src.py
@@ -79,13 +79,16 @@ def _gather_inputs_fresh(lg, inputs, ids, step):
     """Return ``{name: (B,) array}`` if EVERY input signal has a value AT *step*
     for EVERY sample id, else ``None``. Selects the value logged *at* `step`
     (not merely the latest) so a lagging signal-worker job for an older step
-    still gathers that step's inputs correctly."""
+    still gathers that step's inputs correctly.
+
+    Uses the ledger's value-at-step read (``query_per_sample_at_step``), which
+    returns O(batch) rows for the firing step rather than the full per-sample
+    history — so the gather stays flat as training accumulates steps. Repeated
+    reads of the same input across dependents (10 signals all reading the loss)
+    are collapsed by the ledger's per-signal, step-scoped query cache."""
     cols = {}
     for inp in inputs:
-        at = {}
-        for sid, s, val, _ in lg.query_per_sample(inp, sample_ids=ids):
-            if s == step:
-                at[int(sid)] = val
+        at = {int(sid): val for sid, val in lg.query_per_sample_at_step(inp, ids, step)}
         for sid in ids:
             if sid not in at:
                 return None
@@ -99,6 +102,7 @@ def _react_dependents(seed_names, batch_ids, step, origin='train'):
     *step* — order-independently, regardless of which input arrived last. Each
     fired signal is itself a seed (chaining) and fires at most once per step.
     Inputs are read from the logger, so ingested signals must be logged."""
+    _warn_on_signal_cycles()  # lazy, once/process — fires even with no wl.serve
     if batch_ids is None:
         return
     try:
@@ -167,6 +171,67 @@ def dependents_of(nm):
             queue.append(name)   # its value may satisfy further reactive signals
 
 
+def _detect_signal_cycles():
+    """Find circular dependencies in the reactive ``inputs`` graph.
+
+    Returns a list of cycles, each a list of signal names ending where it began
+    (e.g. ``['sig/X', 'sig/Y', 'sig/X']``). Only edges to OTHER registered
+    ``@wl.signal`` names are followed — base metrics (the watched loss, a plain
+    logged signal) are leaves, never part of a cycle. A signal in a cycle can
+    never fire (its inputs are never all fresh), so surfacing these turns a
+    silent no-fire into a diagnostic."""
+    graph = {}
+    for name, func in _REGISTERED_SIGNALS.items():
+        meta = getattr(func, '_wl_signal_meta', {})
+        inputs = meta.get('inputs') or []
+        graph[name] = [dep for dep in inputs if dep in _REGISTERED_SIGNALS]  # depends-on edges
+    WHITE, GREY, BLACK = 0, 1, 2
+    color = {n: WHITE for n in graph}
+    stack, cycles, seen = [], [], set()
+
+    def dfs(n):
+        color[n] = GREY
+        stack.append(n)
+        for dep in graph.get(n, ()):
+            if color.get(dep) == GREY:                 # back-edge → cycle
+                cyc = stack[stack.index(dep):] + [dep]
+                key = frozenset(cyc)
+                if key not in seen:
+                    seen.add(key)
+                    cycles.append(cyc)
+            elif color.get(dep) == WHITE:
+                dfs(dep)
+        stack.pop()
+        color[n] = BLACK
+
+    for n in graph:
+        if color[n] == WHITE:
+            dfs(n)
+    return cycles
+
+
+_SIGNAL_CYCLE_CHECKED = False
+
+
+def _warn_on_signal_cycles(force=False):
+    """Log a warning for each circular dependency in the reactive signal graph.
+    Idempotent (runs its scan at most once per process unless *force*). Triggered
+    from BOTH :func:`start_training` (early, before the first step) and lazily on
+    the first reactive dispatch — so it fires even in a headless SDK run that
+    never calls ``wl.serve``/``wl.start_training``. Warns rather than raises — the
+    cyclic signals simply stay dark while everything else trains normally."""
+    global _SIGNAL_CYCLE_CHECKED
+    if _SIGNAL_CYCLE_CHECKED and not force:
+        return
+    _SIGNAL_CYCLE_CHECKED = True
+    for cyc in _detect_signal_cycles():
+        logger.warning(
+            "WeightsLab reactive signals: circular dependency %s — these signals "
+            "will NEVER fire (each waits on the other's value, which never becomes "
+            "fresh). Break the cycle in their @wl.signal(inputs=[...]).",
+            " -> ".join(cyc))
+
+
 # --- optional signal-worker thread: run reactive dispatch off the train thread -
 _SIGNAL_WORKER = {"q": None, "thread": None}
 
@@ -1351,6 +1416,7 @@ def start_training(timeout: int = None) -> None:
         timeout: Maximum number of seconds to keep running. If ``None``, runs
             until interrupted.
     """
+    _warn_on_signal_cycles()  # surface circular reactive-signal deps (warn, don't raise)
     if timeout is not None and isinstance(timeout, int) and timeout > 0:
         logger.info(f"Starting WeightsLab training mode with a timeout of {timeout} seconds.")
         time.sleep(timeout)

From 473159879d7443ab0945c70ffbb5abe7edcdf5fd Mon Sep 17 00:00:00 2001
From: Alexandru Rotaru <rotarualexandruandrei94@gmail.com>
Date: Wed, 8 Jul 2026 20:09:42 +0200
Subject: [PATCH 09/19] perf(signals): staging-buffer value-at-step read +
 vectorized id->scalar dict
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two profile-driven per-step wins (−30% wall at N=60k, full reactive config;
124.8 → 87.2 ms median), correctness unchanged (signal stress verify all pass).

#2 (logger.py): serve query_per_sample_at_step from the in-memory staging
buffer. The reactive gather reads the CURRENT step's value, which was just
staged and isn't in DuckDB yet — so a reverse scan of the staging list (early
break once all ids found) skips the flush -> register(pandas) -> INSERT ->
unregister -> SELECT round-trip that profiling showed was ~a third of per-step
cost. Falls through to DuckDB only when an id isn't in the buffer (mid-step
flush moved it / older step). Also flattens the epoch-over-epoch drift: current
-step reads no longer scan the growing per_sample table (measured flat ~75 ms
over 98 epochs at 70k, vs the old code rising 110→137 in 7).

#3 (src.py): vectorize the id->scalar dict build — one .tolist() (a single
device sync) instead of a per-element .item() comprehension (B device syncs).

Note: a row-wise executemany flush was tried and REVERTED — it was ~6x slower;
register(pandas)+INSERT SELECT is DuckDB's fast bulk path (comment left in
_flush_stage to prevent re-attempts).
---
 weightslab/backend/logger.py | 46 +++++++++++++++++++++++++++++++++---
 weightslab/src.py            | 16 +++++++++++--
 2 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/weightslab/backend/logger.py b/weightslab/backend/logger.py
index 9c14b01f..8a2301da 100644
--- a/weightslab/backend/logger.py
+++ b/weightslab/backend/logger.py
@@ -194,7 +194,13 @@ def _maybe_autoflush(self) -> None:
             self._flush_stage()
 
     def _flush_stage(self) -> None:
-        """Bulk-insert all staged rows into DuckDB and clear the buffers."""
+        """Bulk-insert all staged rows into DuckDB and clear the buffers.
+
+        Uses ``register(pandas view) -> INSERT SELECT -> unregister``: this is
+        DuckDB's fast vectorized bulk-insert path. (A row-wise ``executemany``
+        from the staging tuples was measured ~6x SLOWER — DuckDB binds each row
+        individually — so despite the register/unregister showing up in profiles,
+        this stays the right approach for bulk.)"""
         with self._lock:
             if self._stage_signals:
                 df = pd.DataFrame(self._stage_signals, columns=_SIGNAL_COLS)
@@ -804,10 +810,44 @@ def query_per_sample_at_step(self, graph_name: str, sample_ids, step, exp_hash=N
 
     def _query_per_sample_at_step_uncached(self, graph_name, ids_key, step, exp_hash, _version):
         """DuckDB read behind :meth:`query_per_sample_at_step`. ``_version`` is a
-        cache-key discriminant only. Returns an immutable tuple."""
+        cache-key discriminant only. Returns an immutable tuple.
+
+        Fast path: the value at *step* was almost always just staged this step
+        and is still in the in-memory staging buffer (not yet flushed to DuckDB).
+        The reactive gather reads exactly this — the current step's value — so
+        scanning the small staging list lets us skip the expensive
+        flush -> register(pandas) -> INSERT -> unregister -> SELECT round-trip
+        entirely (profiling showed that DuckDB dance is ~a third of per-step
+        cost). We only fall through to DuckDB when some requested id is NOT in
+        the buffer (a mid-step flush moved it, or it's an older step)."""
+        step = int(step)
         with self._lock:
+            if ids_key is not None:
+                ids_set = set(ids_key)
+                at = {}
+                # Scan the staging buffer from the end: it's append-ordered by
+                # seq (== non-decreasing step), so the target step's rows are
+                # near the tail. Grab the latest value per id (first seen going
+                # backwards wins) and stop as soon as all ids are found; break
+                # out entirely once we pass below `step` (no earlier row can
+                # match). Touches ~one batch of rows, not the whole 8k buffer.
+                for row in reversed(self._stage_sample):
+                    s = row[3]
+                    if s < step:
+                        break
+                    if s == step and row[0] == graph_name \
+                            and (exp_hash is None or row[1] == exp_hash):
+                        sid = row[2]
+                        if sid in ids_set and sid not in at:
+                            at[sid] = row[4]
+                            if len(at) == len(ids_set):
+                                break
+                if len(at) == len(ids_set):
+                    return tuple((sid, float(val)) for sid, val in at.items())
+
+            # Fallback: not fully in the staging buffer -> flush + query DuckDB.
             self._flush_stage()
-            params = [graph_name, int(step)]
+            params = [graph_name, step]
             sql = "SELECT sample_id, value FROM per_sample WHERE metric_name = ? AND step = ?"
             sql += self._hash_filter(exp_hash, params)
             if ids_key is not None:
diff --git a/weightslab/src.py b/weightslab/src.py
index b52b94c2..48c69fdc 100644
--- a/weightslab/src.py
+++ b/weightslab/src.py
@@ -578,9 +578,21 @@ def _extract_scalar_from_tensor(batch_scalar: th.Tensor | np.ndarray, out: th.Te
                     batch_scalar = _tmp
                 except Exception:
                     pass
-            # Merged batch scalar with ids
+            # Merged batch scalar with ids. Vectorize the tensor->python
+            # conversion: one .tolist() (a single device sync) instead of a
+            # per-element .item() in a comprehension (B device syncs + B python
+            # calls). This is on the per-step hot path — profiling showed the
+            # per-element version was a top self-time cost.
             if isinstance(batch_scalar, (th.Tensor, np.ndarray)) and ids is not None and len(batch_scalar) == len(ids):
-                batch_scalar = {ids[i].item() if isinstance(ids, th.Tensor) else ids[i]: batch_scalar[i].item() for i in range(len(batch_scalar))}
+                _vals = (batch_scalar.detach().cpu().tolist() if isinstance(batch_scalar, th.Tensor)
+                         else np.asarray(batch_scalar).tolist())
+                if isinstance(ids, th.Tensor):
+                    _keys = ids.detach().cpu().tolist()
+                elif isinstance(ids, np.ndarray):
+                    _keys = ids.tolist()
+                else:
+                    _keys = list(ids)
+                batch_scalar = dict(zip(_keys, _vals))
         # 2. Otherwise fall back to extracting from 'out'
         elif out is not None:
             if isinstance(out, th.Tensor):

From 2858b72e8019e9165dc2c656c87ce3e8173e4eb0 Mon Sep 17 00:00:00 2001
From: Alexandru Rotaru <rotarualexandruandrei94@gmail.com>
Date: Thu, 9 Jul 2026 14:00:38 +0200
Subject: [PATCH 10/19] feat(signals): logit-derived signals + batched reactive
 persist + preds_raw gate
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Makes @wl.signal the efficient path (no user save_signals needed), verified by
the signal stress test (all checks pass).

- ctx.logits/preds/targets on SignalContext + BatchSignalContext: a
  subscribe_to signal can compute from the model output (entropy, confidence,
  ...) — logit-derived signals become @wl.signal, no save_signals push.
- Subscriber results logged id-keyed (dict), so reactive dependents can read
  them; reactive dispatch seeded with subscriber results so chains like
  conf_r<-entropy<-logits fire.
- Batched reactive persist: _react_dependents accumulates all fired signals
  in-memory (chains read from there) and persists once, not per signal.
  DuckDB execute 1404->936, register 702->468 (-33%); the @wl.signal path is
  now faster than batched save_signals.
- ledger_store_preds_raw hyperparam (default on): gate off to skip persisting
  the (B,C) logits every step (~7% saving); ctx.logits still works.
- Trim verbose comments ~70% throughout.
---
 weightslab/backend/logger.py |  91 +++++++--------------------
 weightslab/src.py            | 119 +++++++++++++++++++++--------------
 2 files changed, 96 insertions(+), 114 deletions(-)

diff --git a/weightslab/backend/logger.py b/weightslab/backend/logger.py
index 8a2301da..900f3d67 100644
--- a/weightslab/backend/logger.py
+++ b/weightslab/backend/logger.py
@@ -84,25 +84,12 @@ def __init__(self, register: bool = True, db_path: str = ":memory:") -> None:
         self._stage_instance: list = []
         self._seq = 0
 
-        # --- per-sample query cache ---------------------------------------
-        # Many consumers in a single step read the SAME (signal, ids): e.g. 10
-        # reactive signals all reading the loss, or a batched context re-reading
-        # a history. Each read is a lock + _flush_stage + DuckDB scan, so N
-        # identical reads cost Nx for one answer. Two memoized readers share one
-        # per-signal version map:
-        #   _qps_cache      -> query_per_sample        (full history)
-        #   _qps_step_cache -> query_per_sample_at_step (values AT one step)
-        # both keyed by (signal, ids, [step,] hash, version[signal]). Staging a
-        # per-sample row bumps ITS version (see _stage_sample_row), so a cached
-        # read is served until that signal changes, then recomputed. Per-signal
-        # (not global) versioning is essential: persisting a derived signal must
-        # NOT invalidate the loss every dependent is still reading this step.
-        #
-        # The caches are STEP-SCOPED: because the loader reshuffles ids and the
-        # version bumps every step, a cache key never recurs across steps — so
-        # cross-step entries are pure dead weight. _stage_sample_row clears both
-        # caches when the training step advances, bounding them to one step's
-        # worth of entries (the intra-step reuse — the actual win — is kept).
+        # Per-sample query cache. Many consumers read the SAME (signal, ids) in
+        # one step (e.g. reactive signals all reading the loss); memoize so N
+        # identical reads cost 1 scan. Keyed by (signal, ids, [step,] hash,
+        # version[signal]); staging a row bumps that signal's version to
+        # invalidate. Step-scoped: _stage_sample_row clears both caches when the
+        # step advances (keys never recur across steps, so old entries are dead).
         self._qps_version: dict = defaultdict(int)
         self._qps_cache_step: int = -1
         self._qps_cache = functools.lru_cache(maxsize=2048)(self._query_per_sample_uncached)
@@ -196,11 +183,8 @@ def _maybe_autoflush(self) -> None:
     def _flush_stage(self) -> None:
         """Bulk-insert all staged rows into DuckDB and clear the buffers.
 
-        Uses ``register(pandas view) -> INSERT SELECT -> unregister``: this is
-        DuckDB's fast vectorized bulk-insert path. (A row-wise ``executemany``
-        from the staging tuples was measured ~6x SLOWER — DuckDB binds each row
-        individually — so despite the register/unregister showing up in profiles,
-        this stays the right approach for bulk.)"""
+        Uses register(pandas)->INSERT SELECT->unregister (DuckDB's fast bulk
+        path). A row-wise executemany was measured ~6x slower — don't switch."""
         with self._lock:
             if self._stage_signals:
                 df = pd.DataFrame(self._stage_signals, columns=_SIGNAL_COLS)
@@ -231,26 +215,18 @@ def _stage_signal_row(self, graph_name, exp_hash, step, metric_value, timestamp,
         self._maybe_autoflush()
 
     def _stage_sample_row(self, graph_name, exp_hash, sample_id, step, value):
-        # Step advanced -> last step's cache entries can never be hit again
-        # (ids reshuffle + version bump make every key unique per step). Drop
-        # them so the cache stays bounded to the current step instead of
-        # accumulating single-use entries.
+        # New step -> last step's cache entries can't recur; drop them.
         if int(step) > self._qps_cache_step:
             self._invalidate_qps_cache()
             self._qps_cache_step = int(step)
         self._stage_sample.append((
             graph_name, exp_hash, str(sample_id), int(step), float(value), self._next_seq(),
         ))
-        # A new per-sample row for this signal invalidates its cached reads
-        # WITHIN the step (an input written after a first read must be re-read).
-        self._qps_version[graph_name] += 1
+        self._qps_version[graph_name] += 1   # invalidate this signal's cached reads
         self._maybe_autoflush()
 
     def _invalidate_qps_cache(self) -> None:
-        """Drop all memoized per-sample query results (both the full-history and
-        the at-step readers) and reset versions. Used on step advance and by the
-        bulk delete/clear paths, where rows for many signals change at once and
-        per-signal version bumps would be impractical."""
+        """Drop both query caches + versions (step advance; bulk delete/clear)."""
         self._qps_cache.cache_clear()
         self._qps_step_cache.cache_clear()
         self._qps_version.clear()
@@ -767,20 +743,15 @@ def query_per_sample(self, graph_name: str, sample_ids=None, exp_hash=None):
         Returns a list of ``(sample_id, step, value, experiment_hash)`` tuples,
         filtered by *sample_ids* and optionally *exp_hash* (``None`` = all hashes).
 
-        Served from the per-signal query cache (see ``self._qps_cache``): the
-        result is memoized until a new per-sample row for *graph_name* is staged,
-        so repeated identical reads within a step cost one scan, not N. A fresh
-        ``list`` copy is returned each call so callers may mutate it freely.
+        Cached (memoized until *graph_name* is next staged). Returns a fresh list.
         """
         ids_key = tuple(str(s) for s in sample_ids) if sample_ids is not None else None
         cached = self._qps_cache(graph_name, ids_key, exp_hash, self._qps_version[graph_name])
         return list(cached)
 
     def _query_per_sample_uncached(self, graph_name, ids_key, exp_hash, _version):
-        """Actual DuckDB read behind :meth:`query_per_sample`. ``_version`` is a
-        cache-key discriminant only (unused in the body) — it changes when the
-        signal is written, forcing a recompute. Returns a tuple so the cached
-        value stays immutable across callers."""
+        """DuckDB read behind :meth:`query_per_sample`. ``_version`` is a cache
+        key only (bumped on write -> recompute). Returns an immutable tuple."""
         with self._lock:
             self._flush_stage()
             params = [graph_name]
@@ -795,42 +766,26 @@ def _query_per_sample_uncached(self, graph_name, ids_key, exp_hash, _version):
         return tuple((sid, int(step), float(val), h) for (sid, step, val, h) in rows)
 
     def query_per_sample_at_step(self, graph_name: str, sample_ids, step, exp_hash=None):
-        """Per-sample values of *graph_name* at EXACTLY *step* — returns a list
-        of ``(sample_id, value)``. Unlike :meth:`query_per_sample`, the result is
-        O(batch), not O(history): the DuckDB ``WHERE step = ?`` filter is applied
-        in-engine so only the current step's rows are materialized into Python.
-        This keeps the reactive freshness gather flat as training accumulates
-        history (it only ever needs the value at the firing step). Served from
-        the step cache; a fresh ``list`` copy is returned each call.
-        """
+        """``(sample_id, value)`` for *graph_name* at exactly *step* — O(batch),
+        not O(history). Keeps the reactive gather flat as history grows. Cached."""
         ids_key = tuple(str(s) for s in sample_ids) if sample_ids is not None else None
         cached = self._qps_step_cache(graph_name, ids_key, int(step), exp_hash,
                                       self._qps_version[graph_name])
         return list(cached)
 
     def _query_per_sample_at_step_uncached(self, graph_name, ids_key, step, exp_hash, _version):
-        """DuckDB read behind :meth:`query_per_sample_at_step`. ``_version`` is a
-        cache-key discriminant only. Returns an immutable tuple.
-
-        Fast path: the value at *step* was almost always just staged this step
-        and is still in the in-memory staging buffer (not yet flushed to DuckDB).
-        The reactive gather reads exactly this — the current step's value — so
-        scanning the small staging list lets us skip the expensive
-        flush -> register(pandas) -> INSERT -> unregister -> SELECT round-trip
-        entirely (profiling showed that DuckDB dance is ~a third of per-step
-        cost). We only fall through to DuckDB when some requested id is NOT in
-        the buffer (a mid-step flush moved it, or it's an older step)."""
+        """DuckDB read behind :meth:`query_per_sample_at_step` (``_version`` = cache key).
+
+        Fast path: the current step's value is usually still in the in-memory
+        staging buffer, so scan it and skip the flush->register->INSERT->SELECT
+        round-trip. Fall through to DuckDB only if an id isn't staged."""
         step = int(step)
         with self._lock:
             if ids_key is not None:
                 ids_set = set(ids_key)
                 at = {}
-                # Scan the staging buffer from the end: it's append-ordered by
-                # seq (== non-decreasing step), so the target step's rows are
-                # near the tail. Grab the latest value per id (first seen going
-                # backwards wins) and stop as soon as all ids are found; break
-                # out entirely once we pass below `step` (no earlier row can
-                # match). Touches ~one batch of rows, not the whole 8k buffer.
+                # Scan from the tail (append-ordered by step); first value per id
+                # wins, stop when all found or once we drop below `step`.
                 for row in reversed(self._stage_sample):
                     s = row[3]
                     if s < step:
diff --git a/weightslab/src.py b/weightslab/src.py
index 48c69fdc..41c6248a 100644
--- a/weightslab/src.py
+++ b/weightslab/src.py
@@ -75,19 +75,18 @@ def _rebind_caller_local(original_obj: Any, new_obj: Any) -> None:
 _REACTIVE_FIRED = {}   # signal_name -> step it last fired at (per-step dedup)
 
 
-def _gather_inputs_fresh(lg, inputs, ids, step):
-    """Return ``{name: (B,) array}`` if EVERY input signal has a value AT *step*
-    for EVERY sample id, else ``None``. Selects the value logged *at* `step`
-    (not merely the latest) so a lagging signal-worker job for an older step
-    still gathers that step's inputs correctly.
-
-    Uses the ledger's value-at-step read (``query_per_sample_at_step``), which
-    returns O(batch) rows for the firing step rather than the full per-sample
-    history — so the gather stays flat as training accumulates steps. Repeated
-    reads of the same input across dependents (10 signals all reading the loss)
-    are collapsed by the ledger's per-signal, step-scoped query cache."""
+def _gather_inputs_fresh(lg, inputs, ids, step, fresh_cache=None):
+    """``{name: (B,) array}`` if every input has a value at *step* for every id,
+    else ``None`` (value AT step, not latest — a lagging worker still gathers it).
+    Reads via ``query_per_sample_at_step`` (O(batch), cached).
+
+    *fresh_cache*: in-memory values fired this pass — read them so the pass
+    persists once at the end, not per signal."""
     cols = {}
     for inp in inputs:
+        if fresh_cache is not None and inp in fresh_cache:
+            cols[inp] = fresh_cache[inp]
+            continue
         at = {int(sid): val for sid, val in lg.query_per_sample_at_step(inp, ids, step)}
         for sid in ids:
             if sid not in at:
@@ -125,6 +124,10 @@ def dependents_of(nm):
             if inputs and nm in inputs:
                 yield name, func, meta, inputs
 
+    # Fired this pass; kept in-memory so chains read without a ledger round-trip
+    # and the whole pass persists in ONE save_signals (not one write per signal).
+    fired = {}          # {name: (B,) array}
+    fired_log = {}      # {name: log flag}
     queue = list(seed_names)
     while queue:
         logged = queue.pop()
@@ -134,7 +137,7 @@ def dependents_of(nm):
             every = meta.get('compute_every_n_steps', 1) or 1
             if step % every != 0:
                 continue
-            cols = _gather_inputs_fresh(lg, inputs, ids, step)
+            cols = _gather_inputs_fresh(lg, inputs, ids, step, fresh_cache=fired)
             if cols is None:
                 continue  # inputs not all fresh yet — a later log will re-trigger
             _REACTIVE_FIRED[name] = step
@@ -158,28 +161,28 @@ def dependents_of(nm):
             except Exception as e:
                 logger.debug(f"reactive signal '{name}' failed: {e}")
                 continue
-            # Persist (log so downstream reactive signals can read it). _react=False:
-            # the queue below drives chaining, avoiding re-entrant dispatch.
-            try:
-                # step=step (not the live model step): attribute the derived
-                # value to the step it was derived from. Critical when the signal
-                # worker lags behind training, else chained signals can't find it.
-                save_signals(signals={name: vals}, batch_ids=ids, step=step,
-                             log=meta.get('log', True), _react=False)
-            except Exception as e:
-                logger.debug(f"reactive persist '{name}' failed: {e}")
-            queue.append(name)   # its value may satisfy further reactive signals
+            fired[name] = np.asarray(vals, dtype=float)
+            fired_log[name] = bool(meta.get('log', True))
+            queue.append(name)
+
+    # One persist for all signals fired this pass. step=step attributes values to
+    # the firing step (matters when the worker lags). _react=False: no re-dispatch.
+    if fired:
+        try:
+            _logged = {n: v.tolist() for n, v in fired.items() if fired_log.get(n, True)}
+            _unlogged = {n: v.tolist() for n, v in fired.items() if not fired_log.get(n, True)}
+            if _logged:
+                save_signals(signals=_logged, batch_ids=ids, step=step, log=True, _react=False)
+            if _unlogged:
+                save_signals(signals=_unlogged, batch_ids=ids, step=step, log=False, _react=False)
+        except Exception as e:
+            logger.debug(f"reactive batched persist failed: {e}")
 
 
 def _detect_signal_cycles():
-    """Find circular dependencies in the reactive ``inputs`` graph.
-
-    Returns a list of cycles, each a list of signal names ending where it began
-    (e.g. ``['sig/X', 'sig/Y', 'sig/X']``). Only edges to OTHER registered
-    ``@wl.signal`` names are followed — base metrics (the watched loss, a plain
-    logged signal) are leaves, never part of a cycle. A signal in a cycle can
-    never fire (its inputs are never all fresh), so surfacing these turns a
-    silent no-fire into a diagnostic."""
+    """Cycles in the reactive ``inputs`` graph, each as a name list ending where
+    it began. Only edges to other registered signals count (base metrics are
+    leaves). A signal in a cycle never fires — surfacing it beats silent no-fire."""
     graph = {}
     for name, func in _REGISTERED_SIGNALS.items():
         meta = getattr(func, '_wl_signal_meta', {})
@@ -214,12 +217,8 @@ def dfs(n):
 
 
 def _warn_on_signal_cycles(force=False):
-    """Log a warning for each circular dependency in the reactive signal graph.
-    Idempotent (runs its scan at most once per process unless *force*). Triggered
-    from BOTH :func:`start_training` (early, before the first step) and lazily on
-    the first reactive dispatch — so it fires even in a headless SDK run that
-    never calls ``wl.serve``/``wl.start_training``. Warns rather than raises — the
-    cyclic signals simply stay dark while everything else trains normally."""
+    """Warn once per process for each reactive cycle. Fired from start_training
+    and the first dispatch (so it works headless). Warns, doesn't raise."""
     global _SIGNAL_CYCLE_CHECKED
     if _SIGNAL_CYCLE_CHECKED and not force:
         return
@@ -313,7 +312,8 @@ class SignalContext:
     Unified context object for WeightsLab signals.
     Carries all available metadata for a single sample during computation.
     """
-    def __init__(self, sample_id, dataframe, data=None, subscribed_value=None, logger=None, origin=None, step=None):
+    def __init__(self, sample_id, dataframe, data=None, subscribed_value=None, logger=None, origin=None,
+                 step=None, logits=None, preds=None, targets=None):
         self.sample_id = sample_id
         self.dataframe = dataframe
         self.data = data
@@ -321,6 +321,9 @@ def __init__(self, sample_id, dataframe, data=None, subscribed_value=None, logge
         self.logger = logger
         self.origin = origin
         self.step = step          # step the trigger fired at (for freshness checks)
+        self.logits = logits      # this sample's raw model output row (subscribe_to path)
+        self.preds = preds
+        self.targets = targets
 
     @property
     def image(self) -> Optional[np.ndarray]:
@@ -424,13 +427,19 @@ class BatchSignalContext:
     (:meth:`history` runs a single query for the whole batch instead of one query
     per sample), which is where the large speed-ups come from.
     """
-    def __init__(self, sample_ids, subscribed_values, logger=None, dataframe=None, origin=None, step=None):
+    def __init__(self, sample_ids, subscribed_values, logger=None, dataframe=None, origin=None,
+                 step=None, logits=None, preds=None, targets=None):
         self.sample_ids = [int(s) for s in sample_ids]          # length B
         self.subscribed_values = np.asarray(subscribed_values, dtype=float)  # (B,)
         self.logger = logger
         self.dataframe = dataframe
         self.origin = origin
         self.step = step          # step the trigger fired at (for freshness checks)
+        # Raw model output for this batch (subscribe_to path): logit-derived
+        # signals compute from these, no save_signals. None on the inputs=[...] path.
+        self.logits = logits      # (B, C)
+        self.preds = preds
+        self.targets = targets
 
     def history(self, signal_name):
         """Per-sample history of *signal_name* for every sample in the batch, in
@@ -957,6 +966,9 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw):
                                  dataframe=df_proxy,
                                  origin=kwargs.get('origin', 'train'),
                                  step=step,
+                                 logits=preds_raw.detach() if hasattr(preds_raw, 'detach') else preds_raw,
+                                 preds=preds.detach() if hasattr(preds, 'detach') else preds,
+                                 targets=targets.detach() if hasattr(targets, 'detach') else targets,
                              )
                              # Enforce that declared ingested signals are fresh.
                              for dep in ingests:
@@ -979,6 +991,7 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw):
                                  val = float(val_vec[i])
 
                                  # Unified Context Pattern
+                                 _lrow = preds_raw[i] if hasattr(preds_raw, '__getitem__') and preds_raw is not None else None
                                  ctx = SignalContext(
                                      sample_id=int(uid),
                                      subscribed_value=val,
@@ -986,6 +999,7 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw):
                                      dataframe=df_proxy,
                                      origin=kwargs.get('origin', 'train'),
                                      step=step,
+                                     logits=_lrow.detach() if hasattr(_lrow, 'detach') else _lrow,
                                  )
                                  try:
                                      res = func(ctx) # Compute per sample result with unified context
@@ -998,7 +1012,11 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw):
                          dynamic_updates[name] = signal_value
                          if dynamic_updates and meta.get('log', True):
                              logger.debug(f"Dynamic updates computed for signal '{reg_name}': {list(dynamic_updates.keys())}")
-                             _log_signal(sum(signal_value)/len(signal_value), signal_value, name, step=step, **kwargs) # Log custom subscribed signals
+                             # Log per-sample keyed by id (dict), not a bare list —
+                             # else it reaches the dataframe but not the logger's
+                             # per_sample table, and reactive dependents can't read it.
+                             _sv = {int(ids_np[i]): float(signal_value[i]) for i in range(len(ids_np))}
+                             _log_signal(sum(signal_value)/len(signal_value), _sv, name, step=step, **kwargs) # Log custom subscribed signals
                      except StaleSignalError:
                          raise  # correctness enforcement: surface, don't swallow
                      except Exception as e:
@@ -1019,21 +1037,30 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw):
         preds = detach_to_cpu(preds)
         preds_raw = detach_to_cpu(preds_raw)
 
+        # Storing preds_raw (the (B,C) logits) every step is costly and only
+        # needed for FiftyOne/report drill-down — signals already read them via
+        # ctx.logits. Gate off to skip the store. Default on = current behavior.
+        try:
+            _store_preds = bool(get_hyperparams().get('ledger_store_preds_raw', True))
+        except Exception:
+            _store_preds = True
+
         # Enqueue signals and data
         save_signals(
             signals=signals,
             batch_ids=batch_ids,
-            preds_raw=preds_raw,
-            preds=preds,
-            targets=targets,
+            preds_raw=preds_raw if _store_preds else None,
+            preds=preds if _store_preds else None,
+            targets=targets if _store_preds else None,
             log=False # Already logged above, no need to log again in save_signals; set to False to avoid duplicate logging if save_signals is called separately without logging
         )
 
-    # Reactive signals: the watched metric was just logged (above), so it may
-    # satisfy an inputs=[...] signal that depends on it.
+    # Seed reactive dispatch with the watched metric AND subscribe_to results, so
+    # a chain like conf_r<-entropy<-logits fires without any user save_signals.
     if not per_instance and batch_ids is not None:
         try:
-            _dispatch_or_enqueue([reg_name], batch_ids, step, origin or 'train')
+            _seed = [reg_name] + [n for n in dynamic_updates.keys() if n != reg_name]
+            _dispatch_or_enqueue(_seed, batch_ids, step, origin or 'train')
         except StaleSignalError:
             raise
         except Exception as e:

From 090ab945b2f802b7bedd5e5f1a84d8f91858920f Mon Sep 17 00:00:00 2001
From: Alexandru Rotaru <rotarualexandruandrei94@gmail.com>
Date: Thu, 9 Jul 2026 16:07:18 +0200
Subject: [PATCH 11/19] examples: rewrite ws-signals-mnist to zero save_signals
 (ctx.logits + reactive)

460->180 lines. Per-step user code is just the watched loss; entropy is a
@wl.signal from ctx.logits, loss_norm/hardness are reactive, shapes classified
at the end. No save_signals.
---
 .../examples/PyTorch/ws-signals-mnist/main.py | 491 ++++--------------
 1 file changed, 96 insertions(+), 395 deletions(-)

diff --git a/weightslab/examples/PyTorch/ws-signals-mnist/main.py b/weightslab/examples/PyTorch/ws-signals-mnist/main.py
index f5912d96..24414fc7 100644
--- a/weightslab/examples/PyTorch/ws-signals-mnist/main.py
+++ b/weightslab/examples/PyTorch/ws-signals-mnist/main.py
@@ -1,459 +1,160 @@
-"""
-WeightsLab per-sample signals — a minimal, idiomatic MNIST example
-==================================================================
-
-Instrument an ordinary PyTorch training loop with WeightsLab so that, for
-*every training sample*, you capture a small graph of signals — and measure
-what that costs.
-
-    10 base signals    loss, confidence, margin, entropy, correctness, ...
-       (from logits)   pushed once per step with ``wl.save_signals``
-
-     1 live derived    loss normalized by a running mean — a dynamic @wl.signal
-       (@wl.signal)    the framework computes for you on every loss update
-
-     3 end-of-run      from each sample's *full* loss trajectory, computed once
-       derived signals after training:
-                         - loss_shape : one of six curve shapes, stored as a
-                                        text categorical tag (100% coverage)
-                         - loss_cv    : coefficient of variation (instability)
-                         - loss_drop  : net improvement first->last
-
-The output is a per-sample table mapping every ``sample_id`` to the latest value
-of every signal — the raw material for curation: surface the hard samples, the
-likely-mislabeled ones (``Flat_high``), and the ones the model forgets
-(``U_Shape``).
-
-The idiom
----------
-Three ways to declare a signal, and this example shows all three:
+"""Per-sample signals on MNIST — the idiomatic WeightsLab way, zero save_signals.
 
-  * **Push it yourself** — value comes from something WeightsLab can't see (raw
-    logits): compute it in the loop and call ``wl.save_signals(...)``.
+The whole per-step user code is the watched loss. Everything else is a
+``@wl.signal``:
 
-  * **Declare a live @wl.signal** — value is a function of a metric WeightsLab
-    already logs (the per-sample loss). Write a function of ``ctx``; the
-    framework fires it per sample on every loss update and persists the result.
+* the watched loss (``crit``) logs a per-sample ``loss_sample`` trajectory;
+* ``sig/entropy`` is computed from ``ctx.logits`` when the loss fires (a
+  logit-derived signal — no manual push);
+* ``sig/loss_norm`` / ``sig/hardness`` are reactive, derived from logged signals;
+* trajectory shapes are classified once at the end from the loss history.
 
-  * **Derive at the end** — value needs a sample's *whole* trajectory. After
-    training, read each sample's loss history from the ledger, reduce it, and
-    write the result back (as a signal or a categorical tag) for every sample.
-
-Overhead
---------
-The run reports the per-step cost of the signal machinery (watched loss +
-``save_signals``) measured against a plain-PyTorch baseline on the same loader,
-plus the wall-time of the end-of-run derived sweep. All timings are CUDA-synced.
-
-Usage
------
-    python mnist_signals_example.py --outdir ./out
+Run:  python main.py --outdir ./out
 """
 import argparse
-import os
-import shutil
-import time
 from collections import defaultdict
 
 import numpy as np
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torchvision import datasets, transforms
 from torch.utils.data import Dataset
+from torchvision import datasets, transforms
 
 import weightslab as wl
 from weightslab.components.global_monitoring import guard_training_context
 
+LOSS = "loss_sample"
+
 
-# ----------------------------------------------------------------------------
-# A small MNIST CNN that returns raw logits (no softmax) — so CrossEntropyLoss
-# is correct and the logit-derived signals below are meaningful. Kept inline so
-# this example is a single self-contained file.
-# ----------------------------------------------------------------------------
 class SmallCNN(nn.Module):
     def __init__(self):
         super().__init__()
         self.input_shape = (1, 1, 28, 28)
         self.net = nn.Sequential(
-            nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),   # 28 -> 14
-            nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),  # 14 -> 7
-            nn.Flatten(),
-            nn.Linear(16 * 7 * 7, 64), nn.ReLU(),
-            nn.Linear(64, 10),
-        )
+            nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Flatten(), nn.Linear(16 * 7 * 7, 64), nn.ReLU(), nn.Linear(64, 10))
 
     def forward(self, x):
         return self.net(x)
 
 
-# ----------------------------------------------------------------------------
-# Config — small on purpose. A subset + a handful of epochs trains a real model
-# and yields a nice spread of per-sample loss-trajectory shapes in ~a minute.
-# ----------------------------------------------------------------------------
-SEED = 0
-SUBSET = 2000      # samples of MNIST train to use (None -> full 60k)
-EPOCHS = 20
-BATCH = 64
-LR = 0.01
-MEASURE_STEPS = 50  # steps used for the per-step overhead breakdown
-
-# The ledger flushes to storage every this-many buffered rows. This is the
-# single biggest overhead knob: the default (100) flushes every ~1.5 steps at
-# batch 64, which dominates the per-step cost. Set it well above the batch size
-# so the ledger flushes only a few times per epoch. A value <= batch size is a
-# landmine — it flushes mid-batch, every step. (Here: ~batch*128, a handful of
-# flushes over the run.)
-LEDGER_FLUSH_MAX_ROWS = 8192
-
-# The one metric everything subscribes to. It is the ``signal_name`` of the
-# watched loss below, and the ``subscribe_to`` of every live @wl.signal.
-LOSS_SIGNAL = "train/loss_sample"
-
-
-# ----------------------------------------------------------------------------
-# 1. Dataset — every sample carries a stable id so its signals can be tracked
-#    across epochs. The id is just the dataset index here; in a real pipeline
-#    it would be your annotation / image id.
-# ----------------------------------------------------------------------------
 class MNISTIdx(Dataset):
-    """MNIST that returns ``(image, uid, label)`` instead of ``(image, label)``."""
-
-    def __init__(self, root, train, transform, max_samples=None):
-        self.mnist = datasets.MNIST(root=root, train=train, download=True, transform=None)
-        self.transform = transform
-        self.max_samples = max_samples
+    """Yields (image, uid, label). fast_get_label skips image decode at ledger init."""
+    def __init__(self, root, n):
+        self.m = datasets.MNIST(root, train=True, download=True, transform=None)
+        self.t = transforms.ToTensor(); self.n = n
 
     def __len__(self):
-        if self.max_samples is not None:
-            return min(len(self.mnist), self.max_samples)
-        return len(self.mnist)
-
-    def __getitem__(self, idx):
-        image, label = self.mnist[idx]
-        return self.transform(image), idx, label
-
-
-# ----------------------------------------------------------------------------
-# 2. Base signals — pure functions of the batch logits. Nothing
-#    WeightsLab-specific here; these just turn logits + labels into per-sample
-#    scalars. They need raw logits, so we push them ourselves (below) with
-#    ``wl.save_signals`` rather than declaring them as @wl.signal.
-# ----------------------------------------------------------------------------
-def base_signals(logits, labels, loss_per_sample):
-    """10 base per-sample signals derived from raw logits ``[B, C]``."""
-    probs = torch.softmax(logits, dim=1)
-    top2 = probs.topk(2, dim=1).values
-    confidence = top2[:, 0]
-    margin = top2[:, 0] - top2[:, 1]
-    entropy = -(probs * (probs + 1e-12).log()).sum(dim=1)
-    correct = (logits.argmax(dim=1) == labels).float()
-    true_prob = probs.gather(1, labels.view(-1, 1)).squeeze(1)
-    return {
-        "sig/loss": loss_per_sample,
-        "sig/confidence": confidence,
-        "sig/margin": margin,
-        "sig/entropy": entropy,
-        "sig/correct": correct,
-        "sig/true_class_prob": true_prob,
-        "sig/nll": -(true_prob + 1e-12).log(),
-        "sig/logit_true": logits.gather(1, labels.view(-1, 1)).squeeze(1),
-        "sig/logit_max": logits.max(dim=1).values,
-        "sig/logit_std": logits.std(dim=1),
-    }
-
+        return min(len(self.m), self.n)
 
-# ----------------------------------------------------------------------------
-# 3. Live derived signal — a dynamic @wl.signal. It subscribes to the per-sample
-#    loss and fires once per sample every time the loss is logged, receiving a
-#    SignalContext ``ctx`` whose ``subscribed_value`` is that sample's current
-#    loss. It returns one scalar, stored per sample_id. No save_signals, no
-#    plumbing — the framework runs it. Here: loss normalized by a running mean
-#    ("how hard is this sample right now, relative to the average?").
-# ----------------------------------------------------------------------------
-_running_mean_loss = {"v": 1.0}
+    def __getitem__(self, i):
+        img, lab = self.m[i]
+        return self.t(img), i, lab
 
+    def fast_get_label(self, i):
+        return int(self.m.targets[i])
 
-@wl.signal(name="sig/loss_norm", subscribe_to=LOSS_SIGNAL)
-def loss_norm(ctx):
-    loss = ctx.subscribed_value
-    if loss is None:
-        return 0.0
-    m = _running_mean_loss
-    m["v"] = 0.98 * m["v"] + 0.02 * float(loss)
-    return float(loss) / (m["v"] + 1e-9)
 
-
-# ----------------------------------------------------------------------------
-# 4. Trajectory shape — the six curve archetypes, and the classifier that
-#    reduces a per-sample loss curve to one of them.
-# ----------------------------------------------------------------------------
-LOSS_SHAPE_LABELS = [
-    "monotonic", "plateaued", "Flat_high",
-    "high_variance", "U_Shape", "Spiked",
-]
-
-
-def classify_loss_shape(values):
-    """Reduce a per-sample loss trajectory (ordered by step) to a shape label.
-
-    Thresholds are scale-invariant (fractions of the trajectory's own range)
-    and illustrative — tune them for your own task. Returns None if there is
-    not enough history yet (< 5 points).
-    """
+def classify_shape(values):
+    """Per-sample loss trajectory -> shape label (None if < 5 points)."""
     y = np.asarray(values, dtype=float)
     if y.size < 5:
         return None
-
     n = y.size
-    first, last = float(y[0]), float(y[-1])
-    ymin, ymax = float(y.min()), float(y.max())
-    rng = max(ymax - ymin, 1e-8)
-    mean = float(y.mean())
-
-    cv = float(y.std()) / (abs(mean) + 1e-8)         # noisiness
-    drop = (first - last) / (abs(first) + 1e-8)      # net improvement
+    rng = max(float(y.max() - y.min()), 1e-8)
+    drop = (float(y[0]) - float(y[-1])) / (abs(float(y[0])) + 1e-8)
+    cv = float(y.std()) / (abs(float(y.mean())) + 1e-8)
     argmin = int(np.argmin(y))
-    rebound = (last - ymin) / rng                     # climb-back from trough
-    max_up_jump = float(np.diff(y).max()) / rng       # largest single-step rise
-
+    rebound = (float(y[-1]) - float(y.min())) / rng
     tail = y[int(0.6 * n):]
-    tail_flat = (float(tail.std()) / (abs(float(tail.mean())) + 1e-8)) < 0.1
-
-    # Learning shapes first (a big net drop is "being learned", even though a
-    # 2.0 -> 0.0 curve has a high coefficient of variation); then the
-    # not-learning / noisy shapes.
+    tail_flat = float(tail.std()) / (abs(float(tail.mean())) + 1e-8) < 0.1
     if 0.2 * n < argmin < 0.8 * n and rebound > 0.3:
-        return "U_Shape"        # learned then forgotten (min in the middle)
+        return "U_Shape"
     if drop > 0.4:
-        return "monotonic"      # steadily learned (large net drop)
+        return "monotonic"
     if drop > 0.15 and tail_flat:
-        return "plateaued"      # improved, then stuck still-high
-    if max_up_jump > 0.5:
-        return "Spiked"         # sudden jump, no net learning
+        return "plateaued"
+    if float(np.diff(y).max()) / rng > 0.5:
+        return "Spiked"
     if cv > 0.5:
-        return "high_variance"  # noisy oscillation, not converging
-    return "Flat_high"          # never moved, stayed high (likely mislabel)
+        return "high_variance"
+    return "Flat_high"
 
 
-# ----------------------------------------------------------------------------
-# 5. End-of-run derived signals — computed once, from each sample's *full* loss
-#    trajectory in the ledger. This is the "derive at the end" idiom: one query,
-#    then per sample reduce the curve and write the results back by sample_id.
-#    Guarantees 100% coverage, unlike the throttled live signal.
-# ----------------------------------------------------------------------------
-def finalize_derived_signals():
-    """Read every sample's loss trajectory and write the definitive derived
-    signals: shape (text categorical tag) + loss_cv + loss_drop (scalars)."""
+def finalize_shapes():
+    """End sweep: classify each sample's loss trajectory into a categorical tag."""
     series = defaultdict(list)
-    for sid, step, val, _ in wl.query_signal_history(LOSS_SIGNAL):
+    for sid, step, val, _ in wl.query_signal_history(LOSS):
         series[sid].append((step, val))
-
-    ids, cv, drop = [], [], []
     by_label = defaultdict(list)
-    dist = defaultdict(int)
     for sid, pts in series.items():
-        y = np.asarray([v for _, v in sorted(pts)], dtype=float)
-        ids.append(sid)
-        mean = float(y.mean())
-        cv.append(float(y.std()) / (abs(mean) + 1e-8))
-        drop.append((float(y[0]) - float(y[-1])) / (abs(float(y[0])) + 1e-8))
-        label = classify_loss_shape(y.tolist())
+        label = classify_shape([v for _, v in sorted(pts)])
         if label is not None:
             by_label[label].append(sid)
-            dist[label] += 1
-
-    # Scalar derived signals: written by stored sample_id (save_signals writes
-    # to rows by the literal id — the same id space the ledger already uses).
-    wl.save_signals(
-        signals={"sig/loss_cv": torch.tensor(cv), "sig/loss_drop": torch.tensor(drop)},
-        batch_ids=ids, log=False,
-    )
-    # Definitive shape as a text categorical tag, for every sample.
     for label, sids in by_label.items():
         wl.set_categorical_tag(sids, "loss_shape", label)
-    return dict(dist)
-
-
-# ----------------------------------------------------------------------------
-# Overhead timing helpers
-# ----------------------------------------------------------------------------
-def _sync(device):
-    if device.type == "cuda":
-        torch.cuda.synchronize()
-
-
-def measure_overhead(model, loader, optimizer, criterion, device, n_steps):
-    """Per-step cost breakdown (ms/step, CUDA-synced), on the SAME loader.
-
-    For each measured step we time, separately:
-      - baseline : a plain forward+backward+step with an unwatched loss (the
-                   pure-PyTorch cost — no signals at all)
-      - watched  : forward+backward+step with the *watched* loss (adds the
-                   per-sample loss logging + any live @wl.signal subscribers)
-      - compute  : the base_signals() math (softmax/entropy/... on logits)
-      - persist  : the wl.save_signals() call (WeightsLab persistence to the
-                   ledger / H5)
-
-    'watched - baseline' isolates the logging cost; 'compute' is the signal math
-    (cheap); 'persist' is the storage cost.
-    """
-    plain = nn.CrossEntropyLoss()
-    acc = defaultdict(float)
-    it = iter(loader)
-    done = 0
-    while done < n_steps:
-        batch = next(it, None)
-        if batch is None:
-            it = iter(loader)
-            continue
-        inputs, ids, labels = batch
-        inputs, labels = inputs.to(device), labels.to(device)
-
-        # baseline: plain step (no WeightsLab)
-        _sync(device); t0 = time.perf_counter()
-        optimizer.zero_grad()
-        plain(model(inputs), labels).backward()
-        optimizer.step()
-        _sync(device); t1 = time.perf_counter()
-
-        # instrumented step, timing each stage
-        with guard_training_context:
-            optimizer.zero_grad()
-            logits = model(inputs)
-            preds = logits.argmax(dim=1, keepdim=True)
-            loss = criterion(logits, labels, batch_ids=ids, preds=preds)
-            loss.mean().backward()
-            optimizer.step()
-            _sync(device); t2 = time.perf_counter()
-            sig = base_signals(logits.detach(), labels, loss.detach())
-            _sync(device); t3 = time.perf_counter()
-            wl.save_signals(signals=sig, batch_ids=ids,
-                            preds_raw=logits.detach(), preds=preds, targets=labels)
-            _sync(device); t4 = time.perf_counter()
-
-        acc["baseline"] += t1 - t0
-        acc["watched"] += t2 - t1
-        acc["compute"] += t3 - t2
-        acc["persist"] += t4 - t3
-        done += 1
-    return {k: 1000.0 * v / max(done, 1) for k, v in acc.items()}
-
-
-# ----------------------------------------------------------------------------
-# Training — wrap the four moving parts, call the watched loss (which logs
-# LOSS_SIGNAL and fires the live @wl.signal subscribers), and push the
-# logit-derived base signals. Then derive from the trajectories at the end.
-# ----------------------------------------------------------------------------
-def train(outdir, data_root, device):
-    torch.manual_seed(SEED)
-    log_dir = os.path.join(outdir, "wl_logs")
-    # Start from a clean ledger so the trajectories reflect only this run
-    # (a stale ledger would reload old per-sample losses and freeze the curves).
-    shutil.rmtree(log_dir, ignore_errors=True)
-    os.makedirs(log_dir, exist_ok=True)
+    return {k: len(v) for k, v in by_label.items()}
 
-    hyperparams = {
-        "experiment_name": "mnist_signals_example",
-        "device": str(device),
-        "root_log_dir": log_dir,
-        "serving_grpc": False,   # headless: no gRPC / CLI control plane
-        "serving_cli": False,
-        "ledger_flush_max_rows": LEDGER_FLUSH_MAX_ROWS,  # see note above
-        # Declare the loader here too, otherwise WeightsLab's config sync
-        # resets batch_size to its default. Keys map to loader kwargs.
-        "data": {"train_loader": {"batch_size": BATCH, "shuffle": True}},
-    }
-    wl.watch_or_edit(hyperparams, flag="hyperparameters", defaults=hyperparams)
 
-    dataset = MNISTIdx(data_root, train=True, transform=transforms.ToTensor(),
-                       max_samples=SUBSET)
-
-    # Wrap the four moving parts. After this, model/opt/loader behave exactly
-    # like their plain PyTorch counterparts — WeightsLab just observes.
-    model = wl.watch_or_edit(SmallCNN().to(device), flag="model", device=device)
-    optimizer = wl.watch_or_edit(optim.Adam(model.parameters(), lr=LR), flag="optimizer")
-    loader = wl.watch_or_edit(
-        dataset, flag="data", loader_name="train_loader",
-        batch_size=BATCH, shuffle=True, is_training=True, preload_labels=True,
-    )
-    # The watched loss: reduction="none" -> one value per sample. Passing
-    # batch_ids= on each call logs LOSS_SIGNAL per sample AND drives the live
-    # @wl.signal subscribers above. This is the only "signal" we hand-wire.
-    criterion = wl.watch_or_edit(
-        nn.CrossEntropyLoss(reduction="none"),
-        flag="loss", signal_name=LOSS_SIGNAL, per_sample=True, log=True,
-    )
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--outdir", default="./out")
+    ap.add_argument("--subset", type=int, default=2000)
+    ap.add_argument("--epochs", type=int, default=12)
+    args = ap.parse_args()
 
-    # Declare the categorical tag up front so the studio shows all six choices;
-    # the live signal + the end sweep populate it.
-    wl.register_categorical_tag("loss_shape", LOSS_SHAPE_LABELS)
+    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    torch.manual_seed(0)
+    hp = {"experiment_name": "ws-signals-mnist", "device": str(dev),
+          "root_log_dir": args.outdir + "/wl_logs", "serving_grpc": False, "serving_cli": False,
+          "ledger_flush_max_rows": 8192,
+          "data": {"train_loader": {"batch_size": 64, "shuffle": True}}}
+    wl.watch_or_edit(hp, flag="hyperparameters", defaults=hp)
+
+    ds = MNISTIdx(args.outdir + "/data", args.subset)
+    model = wl.watch_or_edit(SmallCNN().to(dev), flag="model", device=dev)
+    opt = wl.watch_or_edit(optim.Adam(model.parameters(), lr=0.01), flag="optimizer")
+    loader = wl.watch_or_edit(ds, flag="data", loader_name="train_loader",
+                              batch_size=64, shuffle=True, is_training=True, preload_labels=True)
+    crit = wl.watch_or_edit(nn.CrossEntropyLoss(reduction="none"),
+                            flag="loss", signal_name=LOSS, per_sample=True, log=True)
+
+    # entropy from the model output when the loss fires — a @wl.signal, no push.
+    @wl.signal(name="sig/entropy", subscribe_to=LOSS, batched=True)
+    def entropy(b):
+        p = torch.softmax(b.logits, 1)
+        return (-(p * (p + 1e-12).log()).sum(1)).detach().cpu().numpy()
+
+    # reactive derived signals (off logged signals; no model output needed)
+    @wl.signal(name="sig/loss_norm", inputs=[LOSS], batched=True)
+    def loss_norm(b):
+        return b.inputs[LOSS] / (float(np.mean(b.inputs[LOSS])) + 1e-8)
+
+    @wl.signal(name="sig/hardness", inputs=[LOSS, "sig/entropy"], batched=True)
+    def hardness(b):
+        return b.inputs[LOSS] * b.inputs["sig/entropy"]
 
     wl.serve(serving_grpc=False, serving_cli=False)
-    wl.start_training(timeout=3)
+    wl.start_training(timeout=0)
 
-    # ---- train: watched loss (logs the per-sample loss trajectory + fires the
-    #      live @wl.signal) and push the logit-derived base signals -----------
-    for epoch in range(1, EPOCHS + 1):
-        for inputs, ids, labels in loader:
-            inputs, labels = inputs.to(device), labels.to(device)
+    for epoch in range(1, args.epochs + 1):
+        for img, ids, lab in loader:
+            img, lab = img.to(dev), lab.to(dev)
             with guard_training_context:
-                optimizer.zero_grad()
-                logits = model(inputs)
-                preds = logits.argmax(dim=1, keepdim=True)
-
-                # Watched loss: logs LOSS_SIGNAL + fires the live @wl.signal.
-                loss_per_sample = criterion(logits, labels, batch_ids=ids, preds=preds)
-                loss_per_sample.mean().backward()
-                optimizer.step()
-
-                # Logit-derived base signals: push them ourselves.
-                wl.save_signals(
-                    signals=base_signals(logits.detach(), labels, loss_per_sample.detach()),
-                    batch_ids=ids,
-                    preds_raw=logits.detach(), preds=preds, targets=labels,
-                )
-        print(f"epoch {epoch:2d}/{EPOCHS}  loss={loss_per_sample.mean().item():.4f}")
-
-    # ---- overhead: per-step cost breakdown vs plain PyTorch ------------------
-    ovh = measure_overhead(model, loader, optimizer, criterion, device, MEASURE_STEPS)
-
-    # ---- derive from full trajectories, once, at the end --------------------
-    _sync(device)
-    t_fin = time.perf_counter()
-    dist = finalize_derived_signals()
-    fin_ms = 1000.0 * (time.perf_counter() - t_fin)
-
-    # ---- export: base + derived signals AND the text shape tag --------------
-    report_path = os.path.join(outdir, "per_sample_signals.csv")
-    report_path = wl.write_dataframe(report_path, format="csv", columns=["signals", "tags"])
-
-    # ---- summary ------------------------------------------------------------
-    instrumented = ovh["watched"] + ovh["compute"] + ovh["persist"]
-    logging_ms = ovh["watched"] - ovh["baseline"]
-    total_over = instrumented - ovh["baseline"]
-    print(f"\nDone. Per-sample report: {report_path}")
-    print(f"\n── per-step overhead (ms/step, CUDA-synced, {MEASURE_STEPS} steps) ──")
-    print(f"  baseline  (plain PyTorch)     : {ovh['baseline']:7.3f}")
-    print(f"  + loss logging + subscribers  : {logging_ms:7.3f}")
-    print(f"  + signal compute (the math)   : {ovh['compute']:7.3f}   <- signal computation")
-    print(f"  + persist (save_signals/H5)   : {ovh['persist']:7.3f}")
-    print(f"  = instrumented total          : {instrumented:7.3f}   "
-          f"(+{total_over:.3f} ms, {100.0 * total_over / max(ovh['baseline'], 1e-9):+.0f}%)")
-    print(f"  end-of-run derived sweep      : {fin_ms:7.1f} ms once ({sum(dist.values())} samples)")
-    print(f"\ntrajectory-shape distribution : {dist}")
-
-
-def main():
-    ap = argparse.ArgumentParser(description=__doc__.splitlines()[1])
-    ap.add_argument("--outdir", default="./out", help="where to write logs + report")
-    ap.add_argument("--data-root", default="./data", help="MNIST download location")
-    args = ap.parse_args()
-    os.makedirs(args.outdir, exist_ok=True)
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print(f"device={device}  subset={SUBSET or 'FULL(60k)'}  epochs={EPOCHS}  batch={BATCH}")
-    train(args.outdir, args.data_root, device)
+                opt.zero_grad()
+                logits = model(img)
+                # the only per-step call: the watched loss logs loss_sample and
+                # fires the @wl.signal chain. No save_signals.
+                crit(logits, lab, batch_ids=ids, preds=logits.argmax(1, keepdim=True)).mean().backward()
+                opt.step()
+        print(f"epoch {epoch}/{args.epochs}")
+
+    print("trajectory-shape distribution:", finalize_shapes())
+    path = wl.write_dataframe(args.outdir + "/per_sample_signals.csv",
+                              format="csv", columns=["signals", "tags"])
+    print("per-sample report:", path)
 
 
 if __name__ == "__main__":

From 4a6871dbb9b6f1d48ac6d457831e353857814c34 Mon Sep 17 00:00:00 2001
From: Alexandru Rotaru <rotarualexandruandrei94@gmail.com>
Date: Thu, 9 Jul 2026 16:43:17 +0200
Subject: [PATCH 12/19] perf(signals): don't query the ledger for
 reactive-derived inputs mid-pass

A reactive-derived input only lives in the in-memory fresh_cache during the
pass; querying the ledger for it (when a dependent is attempted before it fires)
was a guaranteed-miss flush every step. Skip cheaply instead. Per-step 107->55ms
(-48%), 1 flush/step -> 0; verify all pass.
---
 weightslab/src.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/weightslab/src.py b/weightslab/src.py
index 41c6248a..c06408eb 100644
--- a/weightslab/src.py
+++ b/weightslab/src.py
@@ -87,6 +87,12 @@ def _gather_inputs_fresh(lg, inputs, ids, step, fresh_cache=None):
         if fresh_cache is not None and inp in fresh_cache:
             cols[inp] = fresh_cache[inp]
             continue
+        # A reactive-derived input lives only in fresh_cache this pass (it's not
+        # in the ledger at this step until the end-of-pass persist). If it's not
+        # there yet, skip — don't query the ledger (a guaranteed-miss flush).
+        _m = getattr(_REGISTERED_SIGNALS.get(inp), '_wl_signal_meta', {})
+        if _m.get('inputs'):
+            return None
         at = {int(sid): val for sid, val in lg.query_per_sample_at_step(inp, ids, step)}
         for sid in ids:
             if sid not in at:

From bb113180b31e033eb569299b09c4ba4385aecf87 Mon Sep 17 00:00:00 2001
From: Alexandru Rotaru <rotarualexandruandrei94@gmail.com>
Date: Thu, 9 Jul 2026 18:33:16 +0200
Subject: [PATCH 13/19] examples: unify signals-mnist with the stress harness
 (train+eval, live loss_shape signal, overhead readout)

---
 .../examples/PyTorch/ws-signals-mnist/main.py | 218 +++++++++++-------
 1 file changed, 135 insertions(+), 83 deletions(-)

diff --git a/weightslab/examples/PyTorch/ws-signals-mnist/main.py b/weightslab/examples/PyTorch/ws-signals-mnist/main.py
index 24414fc7..7c712d3c 100644
--- a/weightslab/examples/PyTorch/ws-signals-mnist/main.py
+++ b/weightslab/examples/PyTorch/ws-signals-mnist/main.py
@@ -1,18 +1,20 @@
-"""Per-sample signals on MNIST — the idiomatic WeightsLab way, zero save_signals.
+"""Per-sample signals on MNIST — readable, zero save_signals, with train + eval.
 
-The whole per-step user code is the watched loss. Everything else is a
-``@wl.signal``:
+Per-step user code is just the watched loss. Everything else is a @wl.signal:
+  entropy    from ctx.logits when the loss fires
+  loss_norm  reactive, from the logged loss
+  hardness   reactive, from loss + entropy
+  loss_shape reactive, classifies each sample's loss trajectory (live signal,
+             not an end-of-run function). Reads history, so it's throttled by
+             WL_SHAPE_EVERY (1 = every step, full coverage; higher = cheaper).
 
-* the watched loss (``crit``) logs a per-sample ``loss_sample`` trajectory;
-* ``sig/entropy`` is computed from ``ctx.logits`` when the loss fires (a
-  logit-derived signal — no manual push);
-* ``sig/loss_norm`` / ``sig/hardness`` are reactive, derived from logged signals;
-* trajectory shapes are classified once at the end from the loss history.
+Universal loss: the watched crit runs on the test split each epoch too, so test
+samples get a loss trajectory and a shape as well.
 
-Run:  python main.py --outdir ./out
+Env: WL_STRESS_EPOCHS, WL_STRESS_OUT, WL_SHAPE_EVERY.
 """
-import argparse
-from collections import defaultdict
+import os, time, gc, json, shutil
+from collections import defaultdict, Counter
 
 import numpy as np
 import torch
@@ -22,9 +24,49 @@
 from torchvision import datasets, transforms
 
 import weightslab as wl
-from weightslab.components.global_monitoring import guard_training_context
+from weightslab.components.global_monitoring import guard_training_context, guard_testing_context
 
 LOSS = "loss_sample"
+SHAPES = ["monotonic", "plateaued", "Flat_high", "high_variance", "U_Shape", "Spiked"]
+OUT = os.environ.get("WL_STRESS_OUT", "/tmp/wl_stress")
+EPOCHS = int(os.environ.get("WL_STRESS_EPOCHS", "10"))
+SHAPE_EVERY = int(os.environ.get("WL_SHAPE_EVERY", "1"))
+BATCH = 64
+
+
+def classify_shape(values):
+    """Loss trajectory -> shape index (or -1 with < 5 points)."""
+    y = np.asarray(values, dtype=float)
+    if y.size < 5:
+        return -1
+    n = y.size
+    rng = max(float(y.max() - y.min()), 1e-8)
+    drop = (float(y[0]) - float(y[-1])) / (abs(float(y[0])) + 1e-8)
+    cv = float(y.std()) / (abs(float(y.mean())) + 1e-8)
+    argmin = int(np.argmin(y))
+    rebound = (float(y[-1]) - float(y.min())) / rng
+    tail = y[int(0.6 * n):]
+    tail_flat = float(tail.std()) / (abs(float(tail.mean())) + 1e-8) < 0.1
+    if 0.2 * n < argmin < 0.8 * n and rebound > 0.3:
+        return SHAPES.index("U_Shape")
+    if drop > 0.4:
+        return SHAPES.index("monotonic")
+    if drop > 0.15 and tail_flat:
+        return SHAPES.index("plateaued")
+    if float(np.diff(y).max()) / rng > 0.5:
+        return SHAPES.index("Spiked")
+    if cv > 0.5:
+        return SHAPES.index("high_variance")
+    return SHAPES.index("Flat_high")
+
+
+def rss_gb():
+    try:
+        for ln in open("/proc/self/status"):
+            if ln.startswith("VmRSS:"):
+                return int(ln.split()[1]) / (1024 * 1024)
+    except Exception:
+        return -1.0
 
 
 class SmallCNN(nn.Module):
@@ -41,93 +83,69 @@ def forward(self, x):
 
 
 class MNISTIdx(Dataset):
-    """Yields (image, uid, label). fast_get_label skips image decode at ledger init."""
-    def __init__(self, root, n):
-        self.m = datasets.MNIST(root, train=True, download=True, transform=None)
-        self.t = transforms.ToTensor(); self.n = n
+    """Yields (image, uid, label). uid namespaced by split so train/test don't
+    collide in the shared ledger; fast_get_label skips decode at ledger init."""
+    def __init__(self, root, train, base):
+        self.m = datasets.MNIST(root, train=train, download=True, transform=None)
+        self.t = transforms.ToTensor(); self.base = base
 
     def __len__(self):
-        return min(len(self.m), self.n)
+        return len(self.m)
 
     def __getitem__(self, i):
         img, lab = self.m[i]
-        return self.t(img), i, lab
+        return self.t(img), self.base + i, lab
 
     def fast_get_label(self, i):
         return int(self.m.targets[i])
 
 
-def classify_shape(values):
-    """Per-sample loss trajectory -> shape label (None if < 5 points)."""
-    y = np.asarray(values, dtype=float)
-    if y.size < 5:
-        return None
-    n = y.size
-    rng = max(float(y.max() - y.min()), 1e-8)
-    drop = (float(y[0]) - float(y[-1])) / (abs(float(y[0])) + 1e-8)
-    cv = float(y.std()) / (abs(float(y.mean())) + 1e-8)
-    argmin = int(np.argmin(y))
-    rebound = (float(y[-1]) - float(y.min())) / rng
-    tail = y[int(0.6 * n):]
-    tail_flat = float(tail.std()) / (abs(float(tail.mean())) + 1e-8) < 0.1
-    if 0.2 * n < argmin < 0.8 * n and rebound > 0.3:
-        return "U_Shape"
-    if drop > 0.4:
-        return "monotonic"
-    if drop > 0.15 and tail_flat:
-        return "plateaued"
-    if float(np.diff(y).max()) / rng > 0.5:
-        return "Spiked"
-    if cv > 0.5:
-        return "high_variance"
-    return "Flat_high"
-
-
-def finalize_shapes():
-    """End sweep: classify each sample's loss trajectory into a categorical tag."""
-    series = defaultdict(list)
-    for sid, step, val, _ in wl.query_signal_history(LOSS):
-        series[sid].append((step, val))
-    by_label = defaultdict(list)
-    for sid, pts in series.items():
-        label = classify_shape([v for _, v in sorted(pts)])
-        if label is not None:
-            by_label[label].append(sid)
-    for label, sids in by_label.items():
-        wl.set_categorical_tag(sids, "loss_shape", label)
-    return {k: len(v) for k, v in by_label.items()}
+def log(msg):
+    print(msg, flush=True)
 
 
 def main():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--outdir", default="./out")
-    ap.add_argument("--subset", type=int, default=2000)
-    ap.add_argument("--epochs", type=int, default=12)
-    args = ap.parse_args()
-
+    shutil.rmtree(OUT, ignore_errors=True)
+    os.makedirs(OUT + "/wl_logs", exist_ok=True)
     dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     torch.manual_seed(0)
-    hp = {"experiment_name": "ws-signals-mnist", "device": str(dev),
-          "root_log_dir": args.outdir + "/wl_logs", "serving_grpc": False, "serving_cli": False,
-          "ledger_flush_max_rows": 8192,
-          "data": {"train_loader": {"batch_size": 64, "shuffle": True}}}
+    metrics = open(OUT + "/metrics.jsonl", "w")
+
+    # vanilla baseline (plain PyTorch) for the overhead comparison
+    vb = SmallCNN().to(dev); vo = optim.Adam(vb.parameters(), lr=0.01); vc = nn.CrossEntropyLoss()
+    xb = torch.randn(BATCH, 1, 28, 28, device=dev); yb = torch.randint(0, 10, (BATCH,), device=dev)
+    sync = (lambda: torch.cuda.synchronize()) if dev.type == "cuda" else (lambda: None)
+    for _ in range(10):
+        vo.zero_grad(); vc(vb(xb), yb).backward(); vo.step()
+    sync(); t0 = time.perf_counter()
+    for _ in range(50):
+        vo.zero_grad(); vc(vb(xb), yb).backward(); vo.step()
+    sync(); vanilla_ms = 1000 * (time.perf_counter() - t0) / 50
+    del vb, vo; gc.collect()
+    log(f"[run] vanilla baseline = {vanilla_ms:.2f} ms/step (dev {dev})")
+
+    # WeightsLab setup: both splits tracked, watched loss, signals
+    hp = {"experiment_name": "signals-mnist", "device": str(dev), "root_log_dir": OUT + "/wl_logs",
+          "serving_grpc": False, "serving_cli": False, "ledger_flush_max_rows": 8192,
+          "ledger_enable_h5_persistence": False, "experiment_dump_to_train_steps_ratio": 10_000_000,
+          "data": {"train_loader": {"batch_size": BATCH, "shuffle": True}}}
     wl.watch_or_edit(hp, flag="hyperparameters", defaults=hp)
-
-    ds = MNISTIdx(args.outdir + "/data", args.subset)
+    train_ds = MNISTIdx(OUT + "/data", train=True, base=0)
+    test_ds = MNISTIdx(OUT + "/data", train=False, base=1_000_000)
     model = wl.watch_or_edit(SmallCNN().to(dev), flag="model", device=dev)
     opt = wl.watch_or_edit(optim.Adam(model.parameters(), lr=0.01), flag="optimizer")
-    loader = wl.watch_or_edit(ds, flag="data", loader_name="train_loader",
-                              batch_size=64, shuffle=True, is_training=True, preload_labels=True)
+    loader = wl.watch_or_edit(train_ds, flag="data", loader_name="train_loader",
+                              batch_size=BATCH, shuffle=True, is_training=True, preload_labels=True)
+    test_loader = wl.watch_or_edit(test_ds, flag="data", loader_name="test_loader",
+                                   batch_size=256, shuffle=False, is_training=False, preload_labels=True)
     crit = wl.watch_or_edit(nn.CrossEntropyLoss(reduction="none"),
                             flag="loss", signal_name=LOSS, per_sample=True, log=True)
 
-    # entropy from the model output when the loss fires — a @wl.signal, no push.
     @wl.signal(name="sig/entropy", subscribe_to=LOSS, batched=True)
     def entropy(b):
         p = torch.softmax(b.logits, 1)
         return (-(p * (p + 1e-12).log()).sum(1)).detach().cpu().numpy()
 
-    # reactive derived signals (off logged signals; no model output needed)
     @wl.signal(name="sig/loss_norm", inputs=[LOSS], batched=True)
     def loss_norm(b):
         return b.inputs[LOSS] / (float(np.mean(b.inputs[LOSS])) + 1e-8)
@@ -136,25 +154,59 @@ def loss_norm(b):
     def hardness(b):
         return b.inputs[LOSS] * b.inputs["sig/entropy"]
 
+    @wl.signal(name="sig/loss_shape", inputs=[LOSS], batched=True, compute_every_n_steps=SHAPE_EVERY)
+    def loss_shape(b):
+        hist = b.history(LOSS)   # {uid: [loss values in step order]}
+        return np.array([classify_shape(hist[s]) for s in b.sample_ids], dtype=float)
+
     wl.serve(serving_grpc=False, serving_cli=False)
     wl.start_training(timeout=0)
-
-    for epoch in range(1, args.epochs + 1):
+    log(f"[run] tracked train={len(train_ds)} test={len(test_ds)} | shape_every={SHAPE_EVERY}")
+
+    def test_eval():
+        """Universal loss: watched crit over the test split each epoch."""
+        with torch.no_grad():
+            for tb in test_loader:
+                ti, tid, tl = tb[0].to(dev), tb[1], tb[2].to(dev)
+                with guard_testing_context:
+                    tlg = model(ti)
+                    crit(tlg, tl, batch_ids=tid, preds=tlg.argmax(1, keepdim=True))
+
+    step_times, gstep, t_run = [], 0, time.perf_counter()
+    for ep in range(1, EPOCHS + 1):
+        ep_ms = []
         for img, ids, lab in loader:
             img, lab = img.to(dev), lab.to(dev)
+            sync(); ts = time.perf_counter()
             with guard_training_context:
                 opt.zero_grad()
                 logits = model(img)
-                # the only per-step call: the watched loss logs loss_sample and
-                # fires the @wl.signal chain. No save_signals.
+                # only per-step call: the watched loss logs loss_sample and fires
+                # the @wl.signal chain. No save_signals.
                 crit(logits, lab, batch_ids=ids, preds=logits.argmax(1, keepdim=True)).mean().backward()
                 opt.step()
-        print(f"epoch {epoch}/{args.epochs}")
-
-    print("trajectory-shape distribution:", finalize_shapes())
-    path = wl.write_dataframe(args.outdir + "/per_sample_signals.csv",
-                              format="csv", columns=["signals", "tags"])
-    print("per-sample report:", path)
+            sync(); ep_ms.append(1000 * (time.perf_counter() - ts))
+            gstep += 1
+        test_eval()
+        wl_ms = float(np.mean(ep_ms)); step_times += ep_ms
+        rec = {"epoch": ep, "gstep": gstep, "wl_ms": round(wl_ms, 2), "vanilla_ms": round(vanilla_ms, 2),
+               "rss_gb": round(rss_gb(), 2), "elapsed_s": round(time.perf_counter() - t_run, 1)}
+        metrics.write(json.dumps(rec) + "\n"); metrics.flush()
+        log(f"[run] ep {ep:3d}/{EPOCHS} | WL {wl_ms:6.2f} ms vs vanilla {vanilla_ms:.2f} "
+            f"= +{100*(wl_ms/vanilla_ms-1):.0f}% | RSS {rec['rss_gb']:.2f} GB | {rec['elapsed_s']:.0f}s")
+
+    import pandas as pd
+    path = wl.write_dataframe(OUT + "/report.csv", format="csv", columns="signals")
+    df = pd.read_csv(path)
+    sc = [c for c in df.columns if c.endswith("sig/loss_shape")][0]
+    dist = Counter(SHAPES[int(v)] for v in df[sc].dropna() if v >= 0)
+    med = float(np.median(step_times))
+    metrics.close()
+    log("\n[run] ===== SUMMARY =====")
+    log(f"[run] vanilla {vanilla_ms:.2f} ms | WL median {med:.2f} ms/step (+{100*(med/vanilla_ms-1):.0f}%)")
+    log(f"[run] report {len(df)} rows | loss_shape covered {df[sc].notna().sum()}/{len(df)}")
+    log(f"[run] shape distribution: {dict(dist)}")
+    log(f"[run] report -> {path}")
 
 
 if __name__ == "__main__":

From 3ccc726fdff0712d13ab2e1e54ff7183233af438 Mon Sep 17 00:00:00 2001
From: Alexandru Rotaru <rotarualexandruandrei94@gmail.com>
Date: Thu, 9 Jul 2026 18:39:02 +0200
Subject: [PATCH 14/19] tests: unit tests for signal-refinements (query cache,
 value-at-step, cycle detection, ctx.logits, freshness, reactive-derived
 gather skip)

---
 tests/general/test_signal_refinements.py | 124 +++++++++++++++++++++++
 1 file changed, 124 insertions(+)
 create mode 100644 tests/general/test_signal_refinements.py

diff --git a/tests/general/test_signal_refinements.py b/tests/general/test_signal_refinements.py
new file mode 100644
index 00000000..5983606a
--- /dev/null
+++ b/tests/general/test_signal_refinements.py
@@ -0,0 +1,124 @@
+"""Unit tests for the signal-refinements changes: per-sample query cache,
+value-at-step read + staging fast path, cycle detection, ctx.logits, freshness,
+and the reactive-derived gather skip."""
+import unittest
+
+import numpy as np
+import torch
+
+import weightslab as wl
+from weightslab.backend.logger import LoggerQueue
+from weightslab.src import (
+    _REGISTERED_SIGNALS, _detect_signal_cycles, _gather_inputs_fresh,
+    BatchSignalContext, SignalContext, StaleSignalError,
+)
+
+
+def _lg():
+    return LoggerQueue(register=False)
+
+
+class TestQueryCache(unittest.TestCase):
+    def test_cached_and_fresh_copy(self):
+        lg = _lg()
+        lg._stage_sample_row("m", "h", "1", 0, 1.5)
+        lg._stage_sample_row("m", "h", "2", 0, 2.5)
+        r1 = lg.query_per_sample("m", sample_ids=["1", "2"])
+        r2 = lg.query_per_sample("m", sample_ids=["1", "2"])
+        self.assertEqual(len(r1), 2)
+        self.assertIsNot(r1, r2)                       # fresh list each call
+        self.assertGreaterEqual(lg._qps_cache.cache_info().hits, 1)  # 2nd read hit
+
+    def test_invalidated_on_write(self):
+        lg = _lg()
+        lg._stage_sample_row("m", "h", "1", 0, 1.0)
+        self.assertEqual(len(lg.query_per_sample("m", sample_ids=["1"])), 1)
+        lg._stage_sample_row("m", "h", "1", 1, 9.0)    # new row bumps version
+        self.assertEqual(len(lg.query_per_sample("m", sample_ids=["1"])), 2)
+
+    def test_step_scoped_clear(self):
+        lg = _lg()
+        lg._stage_sample_row("m", "h", "1", 0, 1.0)
+        lg.query_per_sample("m", sample_ids=["1"])
+        lg._stage_sample_row("m", "h", "2", 1, 2.0)    # step advance -> clear
+        self.assertEqual(lg._qps_cache_step, 1)
+
+
+class TestValueAtStep(unittest.TestCase):
+    def test_staging_fast_path(self):
+        lg = _lg()
+        lg._stage_sample_row("m", "h", "5", 3, 4.0)    # staged, not flushed
+        at = dict((int(s), v) for s, v in lg.query_per_sample_at_step("m", [5], 3))
+        self.assertEqual(at, {5: 4.0})
+
+    def test_absent_at_other_step(self):
+        lg = _lg()
+        lg._stage_sample_row("m", "h", "5", 3, 4.0)
+        self.assertEqual(lg.query_per_sample_at_step("m", [5], 99), [])
+
+
+class TestCycleDetection(unittest.TestCase):
+    def setUp(self): _REGISTERED_SIGNALS.clear()
+    def tearDown(self): _REGISTERED_SIGNALS.clear()
+
+    def test_self_loop_and_two_cycle(self):
+        @wl.signal(name="sig/self", inputs=["sig/self"], batched=True)
+        def _s(b): return b
+        @wl.signal(name="sig/X", inputs=["sig/Y"], batched=True)
+        def _x(b): return b
+        @wl.signal(name="sig/Y", inputs=["sig/X"], batched=True)
+        def _y(b): return b
+        keys = {frozenset(c) for c in _detect_signal_cycles()}
+        self.assertIn(frozenset(["sig/self"]), keys)
+        self.assertIn(frozenset(["sig/X", "sig/Y"]), keys)
+
+    def test_dag_has_no_cycle(self):
+        @wl.signal(name="sig/a", inputs=["train/loss"], batched=True)   # base leaf
+        def _a(b): return b
+        @wl.signal(name="sig/b", inputs=["sig/a"], batched=True)
+        def _b(b): return b
+        self.assertEqual(_detect_signal_cycles(), [])
+
+
+class TestContexts(unittest.TestCase):
+    def test_batch_context_carries_logits(self):
+        b = BatchSignalContext(sample_ids=[1, 2], subscribed_values=[0.1, 0.2],
+                               logits=torch.tensor([[1., 2., 3.], [4., 5., 6.]]))
+        self.assertEqual(tuple(b.logits.shape), (2, 3))
+
+    def test_sample_context_carries_logits(self):
+        c = SignalContext(sample_id=1, dataframe=None, logits=torch.tensor([1., 2., 3.]))
+        self.assertEqual(tuple(c.logits.shape), (3,))
+
+    def test_stale_signal_raises(self):
+        b = BatchSignalContext(sample_ids=[1], subscribed_values=[0.0], logger=_lg(), step=5)
+        with self.assertRaises(StaleSignalError):
+            b.latest("sig/never", require_fresh=True)
+
+
+class TestGatherFreshCache(unittest.TestCase):
+    def setUp(self): _REGISTERED_SIGNALS.clear()
+    def tearDown(self): _REGISTERED_SIGNALS.clear()
+
+    def test_reactive_derived_not_queried_when_absent(self):
+        @wl.signal(name="sig/derived", inputs=["x"], batched=True)   # reactive
+        def _d(b): return b
+        lg = _lg()
+        calls = {"n": 0}
+        orig = lg.query_per_sample_at_step
+        lg.query_per_sample_at_step = lambda *a, **k: (calls.__setitem__("n", calls["n"] + 1), orig(*a, **k))[1]
+        res = _gather_inputs_fresh(lg, ["sig/derived"], [1], 0, fresh_cache={})
+        self.assertIsNone(res)          # not fired yet -> skip
+        self.assertEqual(calls["n"], 0)  # and NOT a ledger query (no flush)
+
+    def test_reactive_derived_read_from_cache(self):
+        @wl.signal(name="sig/derived", inputs=["x"], batched=True)
+        def _d(b): return b
+        res = _gather_inputs_fresh(_lg(), ["sig/derived"], [1, 2], 0,
+                                   fresh_cache={"sig/derived": np.array([1.0, 2.0])})
+        self.assertIsNotNone(res)
+        np.testing.assert_array_equal(res["sig/derived"], [1.0, 2.0])
+
+
+if __name__ == "__main__":
+    unittest.main()

From 1414e0cb041a7f73aaf106a5c6678b6b4e4c5fca Mon Sep 17 00:00:00 2001
From: GuillaumePELLUET <guillaume@graybx.com>
Date: Fri, 10 Jul 2026 11:35:25 +0200
Subject: [PATCH 15/19] update gitignore - remove coverage files

---
 .gitignore | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/.gitignore b/.gitignore
index c7f25cb7..65a29764 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,11 +11,15 @@ venv
 runs
 data
 MagicMock
+drop
+htmlcov
+htmlcov-report
 
 # Ignore hidden directories
 .history
 .lh
 .claude
+.coverage
 .chat
 .ai
 .pytest_cache

From fea2758e3cf4199cdac6cff700c557155d1b6e9b Mon Sep 17 00:00:00 2001
From: GuillaumePELLUET <guillaume@graybx.com>
Date: Fri, 10 Jul 2026 11:38:16 +0200
Subject: [PATCH 16/19] remove last merged conflicts

---
 .../README.md                                 | 102 ------------------
 .../ws-signals-mnist/main.py                  |   0
 weightslab/trainer/services/data_service.py   |   6 --
 3 files changed, 108 deletions(-)
 delete mode 100644 weightslab/examples/Usecases/wl-loss_shapes_classification_per_sample/README.md
 rename weightslab/examples/{PyTorch => Usecases}/ws-signals-mnist/main.py (100%)

diff --git a/weightslab/examples/Usecases/wl-loss_shapes_classification_per_sample/README.md b/weightslab/examples/Usecases/wl-loss_shapes_classification_per_sample/README.md
deleted file mode 100644
index 3fee76a9..00000000
--- a/weightslab/examples/Usecases/wl-loss_shapes_classification_per_sample/README.md
+++ /dev/null
@@ -1,102 +0,0 @@
-# WeightsLab — Object Detection (pure PyTorch)
-
-A small, fully-runnable **object detection** example wired into WeightsLab. It
-trains a compact single-shot detector on the **Penn-Fudan Pedestrian** dataset
-(~170 real photos, one class: `person`) and streams per-sample / per-instance
-losses, IoU, and predicted bounding boxes to the WeightsLab UI.
-
-Everything here is plain PyTorch + torchvision — no detection framework
-(no Ultralytics/Detectron). The only pretrained piece is an ImageNet backbone.
-
-## Quick start
-
-From a WeightsLab install, the one-liner (installs this example's
-`requirements.txt`, then trains + serves until `Ctrl+C`):
-
-```bash
-weightslab start example --det
-```
-
-Or run it directly:
-
-```bash
-cd weightslab/examples/PyTorch/wl-detection
-pip install -r requirements.txt
-python main.py
-```
-
-The **first run downloads** the Penn-Fudan dataset (~50 MB, into `./data/`) and
-the MobileNetV3-Small ImageNet weights (~10 MB, cached by torch). Then open the
-UI (e.g. `http://localhost:5173`) to watch training.
-
-## What you'll see in the UI
-
-| Signal                  | Meaning                                              |
-| ----------------------- | ---------------------------------------------------- |
-| `train_loss/sample`     | Per-image training loss (the value being optimized)  |
-| `test_loss/sample`      | Per-image validation loss                            |
-| `train_iou/sample`      | Mean IoU per training image                          |
-| `test_iou/sample`       | Mean IoU per validation image                        |
-| `train_iou/instance`    | IoU per **ground-truth box** `(sample_id, annotation_id)` |
-| `test_iou/instance`     | Same, on validation                                  |
-
-Ground-truth and predicted **bounding boxes** are rendered as overlays on each
-sample (the dataset and model declare `task_type = "detection"`).
-
-## How it works
-
-```
-utils/data.py        PennFudanDetectionDataset — downloads Penn-Fudan, derives one
-                     bbox per pedestrian from the instance masks, returns the WL
-                     detection target [N, 6] = [x1, y1, x2, y2, class_id, conf]
-                     normalized to [0, 1]. ImageNet-normalized model inputs.
-                     `det_collate` keeps the variable box count as a per-sample list.
-
-utils/model.py       SmallDetector — ImageNet-pretrained MobileNetV3-Small backbone
-                     (frozen by default) + a small head that predicts ONE box per
-                     cell on an S x S grid: (objectness, tx, ty, tw, th, class...).
-                     `decode_grid` turns raw logits into xyxy boxes.
-
-utils/criterions.py  PerSampleDetectionLoss — YOLO-style objectness + coordinate +
-                     class loss, one differentiable scalar per sample (what WL
-                     backprops). PerSampleIoU / PerInstanceIoU — IoU metrics.
-                     decode_predictions — top-confidence boxes for the UI overlay.
-
-main.py              Wires it all to WeightsLab: watch_or_edit(...) for the logger,
-                     hyperparameters, data loaders, model, optimizer and the
-                     loss/metric signals; serve(); start_training(); train/test loop.
-```
-
-The detector is genuinely learnable: on a small subset, mean IoU rises from
-~0.39 to ~0.83 within ~60 steps.
-
-## Configuration (`config.yaml`)
-
-| Key                    | Default | Notes                                                       |
-| ---------------------- | ------- | ----------------------------------------------------------- |
-| `num_classes`          | `1`     | Penn-Fudan has one class (`person`).                        |
-| `image_size`           | `256`   | Square model input (UI shows the original image).           |
-| `grid_size`            | `8`     | Detector predicts on an `8 x 8` cell grid.                  |
-| `conf_thresh`          | `0.3`   | `objectness * class` threshold for displayed predictions.   |
-| `pretrained_backbone`  | `true`  | Load ImageNet weights for the MobileNetV3 backbone.         |
-| `freeze_backbone`      | `true`  | Train only the head (fast, less data-hungry). Set `false` to fine-tune the whole backbone once the head has warmed up. |
-| `data.*.batch_size`    | `8`     | Per-loader batch size.                                      |
-| `data.*.max_samples`   | `null`  | Cap a split for quick runs (`null` = full split).           |
-
-## Using your own dataset (e.g. traffic lights)
-
-The model, loss, metrics, `main.py`, and UI rendering are **dataset-agnostic** —
-only `utils/data.py` and a couple of config values change:
-
-1. Write a `Dataset` whose `get_items(idx, ...)` returns
-   `(image_tensor, uid, target, metadata)`, where `target` is an
-   `[N, 6]` float array `[x1, y1, x2, y2, class_id, confidence]` **normalized to
-   `[0, 1]`** (ground-truth confidence = `1.0`). Set `self.task_type = "detection"`,
-   `self.num_classes`, `self.class_names`, and expose `self.images` (a list of
-   image paths) so the UI can show the raw image.
-2. Reuse `det_collate` unchanged.
-3. In `config.yaml`, set `num_classes` to your class count (e.g. `3` for
-   `red / yellow / green`) and update `class_names` in the dataset / model.
-
-That's it — multi-class works out of the box (the classification head is already
-in the grid prediction; it's just trivial when `num_classes == 1`).
diff --git a/weightslab/examples/PyTorch/ws-signals-mnist/main.py b/weightslab/examples/Usecases/ws-signals-mnist/main.py
similarity index 100%
rename from weightslab/examples/PyTorch/ws-signals-mnist/main.py
rename to weightslab/examples/Usecases/ws-signals-mnist/main.py
diff --git a/weightslab/trainer/services/data_service.py b/weightslab/trainer/services/data_service.py
index aa033005..bd9071e4 100755
--- a/weightslab/trainer/services/data_service.py
+++ b/weightslab/trainer/services/data_service.py
@@ -2402,8 +2402,6 @@ def _agent_load_weights(self, step=None, exp_hash=None) -> str:
         except Exception as e:
             return f"Action: failed to load weights ({where}) from {exp_hash[:16]}: {e}"
 
-<<<<<<< HEAD
-=======
     # ------------------------------------------------------------------
     # Hyperparameter tuning (agent action)
     # ------------------------------------------------------------------
@@ -2655,7 +2653,6 @@ def _agent_show_config(self, param=None) -> str:
             text = text[:max_len] + "\n... (truncated; ask for a specific key, e.g. 'show the root log dir')"
         return f"Configuration ({hp_name}):\n{text}"
 
->>>>>>> 2db142c599a32f94669a3afad5f4098d5629ddde
     def _apply_agent_operation(self, df, func: str, params: dict) -> str:
         """
         Apply an agent-described operation to df in-place.
@@ -2712,8 +2709,6 @@ def _apply_agent_operation(self, df, func: str, params: dict) -> str:
                     exp_hash=params.get("hash") or params.get("exp_hash"),
                 )
 
-<<<<<<< HEAD
-=======
             # Set / scale a hyperparameter (batch size, learning rate, ratios, ...).
             elif action_name in ("set_hyperparam", "set_hyperparameter", "set_hp", "tune_hyperparam"):
                 return self._agent_set_hyperparam(
@@ -2729,7 +2724,6 @@ def _apply_agent_operation(self, df, func: str, params: dict) -> str:
                     param=params.get("param") or params.get("name") or params.get("key_path")
                 )
 
->>>>>>> 2db142c599a32f94669a3afad5f4098d5629ddde
             return f"Action triggered: {action_name} (Not implemented)"
 
         # --- 3. DATAFRAME MANIPULATION ---

From 5e0ed703eeefbea6c139cb42f37b752b02f1aaee Mon Sep 17 00:00:00 2001
From: Alexandru Rotaru <rotarualexandruandrei94@gmail.com>
Date: Fri, 10 Jul 2026 11:52:33 +0200
Subject: [PATCH 17/19] logger: make query-cache maxsize env-configurable
 (WL_QUERY_CACHE_MAXSIZE, default 2048)

Addresses PR #244 review (Guillaume): let users size the dedicated query cache.
---
 weightslab/backend/logger.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/weightslab/backend/logger.py b/weightslab/backend/logger.py
index 900f3d67..a29bdf07 100644
--- a/weightslab/backend/logger.py
+++ b/weightslab/backend/logger.py
@@ -31,6 +31,7 @@
 
 import functools
 import json
+import os
 import threading
 import time
 from collections import defaultdict
@@ -90,10 +91,12 @@ def __init__(self, register: bool = True, db_path: str = ":memory:") -> None:
         # version[signal]); staging a row bumps that signal's version to
         # invalidate. Step-scoped: _stage_sample_row clears both caches when the
         # step advances (keys never recur across steps, so old entries are dead).
+        # Cache size is env-configurable (WL_QUERY_CACHE_MAXSIZE, default 2048).
+        _qps_maxsize = int(os.environ.get("WL_QUERY_CACHE_MAXSIZE", "2048"))
         self._qps_version: dict = defaultdict(int)
         self._qps_cache_step: int = -1
-        self._qps_cache = functools.lru_cache(maxsize=2048)(self._query_per_sample_uncached)
-        self._qps_step_cache = functools.lru_cache(maxsize=2048)(self._query_per_sample_at_step_uncached)
+        self._qps_cache = functools.lru_cache(maxsize=_qps_maxsize)(self._query_per_sample_uncached)
+        self._qps_step_cache = functools.lru_cache(maxsize=_qps_maxsize)(self._query_per_sample_at_step_uncached)
 
         self._ensure_tables()
         self._restore_runtime_state_from_db()

From b3f09ebc995d8e652240ce9b1562b4c6a2975f8c Mon Sep 17 00:00:00 2001
From: Alexandru Rotaru <rotarualexandruandrei94@gmail.com>
Date: Fri, 10 Jul 2026 12:13:27 +0200
Subject: [PATCH 18/19] feat(signals): default loss-shape classifier + report
 summary (PR #244 review #3)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two-layer, reusable over any per-sample signal:
- classify_loss_shape(values, ...): default shape classifier, all thresholds
  named + configurable (natural defaults; tune without rewriting).
- write_signal_shapes(signal, tag, classifier): engine — classify every sample's
  trajectory of ANY signal (loss, accuracy, a second loss, any metric) into a
  categorical tag; returns the {label: count} distribution.
- write_loss_shapes(): convenience wrapper for the loss signal.
- write_dataframe(loss_shape_signal=...): runs the summary on report write (tags
  + logs the distribution), so a report written every N steps stays current.
Exported: classify_loss_shape, write_loss_shapes, write_signal_shapes, LOSS_SHAPES.
Tests added; existing signal/logger suites pass.
---
 tests/general/test_signal_refinements.py | 21 +++++++
 weightslab/__init__.py                   |  6 +-
 weightslab/src.py                        | 77 ++++++++++++++++++++++++
 3 files changed, 103 insertions(+), 1 deletion(-)

diff --git a/tests/general/test_signal_refinements.py b/tests/general/test_signal_refinements.py
index 5983606a..e5aa6ff2 100644
--- a/tests/general/test_signal_refinements.py
+++ b/tests/general/test_signal_refinements.py
@@ -120,5 +120,26 @@ def _d(b): return b
         np.testing.assert_array_equal(res["sig/derived"], [1.0, 2.0])
 
 
+class TestDefaultShapeClassifier(unittest.TestCase):
+    def test_labels(self):
+        self.assertEqual(wl.classify_loss_shape([2, 1.5, 1, 0.5, 0.05]), "monotonic")
+        self.assertEqual(wl.classify_loss_shape([2, 2, 2, 2, 2]), "Flat_high")
+        self.assertEqual(wl.classify_loss_shape([0.1, 0.1, 0.1, 3.0, 0.1]), "Spiked")
+
+    def test_too_short_is_none(self):
+        self.assertIsNone(wl.classify_loss_shape([1, 2, 3]))
+        self.assertIsNone(wl.classify_loss_shape([1, 0.5, 0.1], min_points=5))
+
+    def test_thresholds_configurable(self):
+        traj = [1.0, 0.9, 0.8, 0.75, 0.7]                       # net drop 0.3
+        self.assertNotEqual(wl.classify_loss_shape(traj), "monotonic")   # default 0.4 unmet
+        self.assertEqual(wl.classify_loss_shape(traj, drop_learned=0.25), "monotonic")
+
+    def test_exports(self):
+        for name in ("write_loss_shapes", "write_signal_shapes", "classify_loss_shape"):
+            self.assertTrue(hasattr(wl, name), name)
+        self.assertIn("monotonic", wl.LOSS_SHAPES)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/weightslab/__init__.py b/weightslab/__init__.py
index 61250afe..395e5776 100644
--- a/weightslab/__init__.py
+++ b/weightslab/__init__.py
@@ -11,7 +11,7 @@
 import logging
 import threading
 
-from .src import watch_or_edit, start_training, serve, keep_serving, save_signals, save_instance_signals, save_group_signals, tag_samples, register_categorical_tag, set_categorical_tag, discard_samples, get_samples_by_tag, get_discarded_samples, signal, eval_fn, compute_signals, SignalContext, BatchSignalContext, StaleSignalError, drain_signals, clear_all, run_pending_evaluation, trigger_pending_evaluation_async, query_signal_history, query_sample_history, query_instance_history, write_history, write_dataframe, get_current_experiment_hash, pointcloud_thumbnail, pointcloud_boxes
+from .src import watch_or_edit, start_training, serve, keep_serving, save_signals, save_instance_signals, save_group_signals, tag_samples, register_categorical_tag, set_categorical_tag, discard_samples, get_samples_by_tag, get_discarded_samples, signal, eval_fn, compute_signals, SignalContext, BatchSignalContext, StaleSignalError, drain_signals, clear_all, run_pending_evaluation, trigger_pending_evaluation_async, query_signal_history, query_sample_history, query_instance_history, write_history, write_dataframe, classify_loss_shape, write_loss_shapes, write_signal_shapes, LOSS_SHAPES, get_current_experiment_hash, pointcloud_thumbnail, pointcloud_boxes
 from .backend.ledgers import GLOBAL_LEDGER as ledger
 from .art import _BANNER
 from .utils.logs import setup_logging, set_log_directory, is_main_process
@@ -169,6 +169,10 @@ def _clean(v: str) -> str:
 
     "write_history",
     "write_dataframe",
+    "classify_loss_shape",
+    "write_loss_shapes",
+    "write_signal_shapes",
+    "LOSS_SHAPES",
 
     "pointcloud_thumbnail",
     "pointcloud_boxes",
diff --git a/weightslab/src.py b/weightslab/src.py
index 6d771f61..d3fe09ae 100644
--- a/weightslab/src.py
+++ b/weightslab/src.py
@@ -4006,6 +4006,70 @@ def write_history(
     return path
 
 
+LOSS_SHAPES = ["monotonic", "plateaued", "Flat_high", "high_variance", "U_Shape", "Spiked"]
+
+
+def classify_loss_shape(values, min_points=5, drop_learned=0.4, drop_plateau=0.15,
+                        tail_flat_cv=0.1, spike_jump=0.5, noisy_cv=0.5, u_rebound=0.3):
+    """Default per-sample loss-trajectory classifier -> one of ``LOSS_SHAPES``
+    (or ``None`` with < *min_points*). Every threshold is a named, scale-invariant
+    param with a natural default, so you can tune it without rewriting the
+    classifier, e.g. ``write_loss_shapes(classifier=lambda v:
+    wl.classify_loss_shape(v, drop_learned=0.5))``."""
+    y = np.asarray(values, dtype=float)
+    if y.size < min_points:
+        return None
+    n = y.size
+    rng = max(float(y.max() - y.min()), 1e-8)
+    drop = (float(y[0]) - float(y[-1])) / (abs(float(y[0])) + 1e-8)
+    cv = float(y.std()) / (abs(float(y.mean())) + 1e-8)
+    argmin = int(np.argmin(y))
+    rebound = (float(y[-1]) - float(y.min())) / rng
+    tail = y[int(0.6 * n):]
+    tail_flat = float(tail.std()) / (abs(float(tail.mean())) + 1e-8) < tail_flat_cv
+    if 0.2 * n < argmin < 0.8 * n and rebound > u_rebound:
+        return "U_Shape"
+    if drop > drop_learned:
+        return "monotonic"
+    if drop > drop_plateau and tail_flat:
+        return "plateaued"
+    if float(np.diff(y).max()) / rng > spike_jump:
+        return "Spiked"
+    if cv > noisy_cv:
+        return "high_variance"
+    return "Flat_high"
+
+
+def write_signal_shapes(signal_name, tag_name=None, classifier=None):
+    """Reusable engine: classify every sample's trajectory of *signal_name* into
+    a categorical tag and return the ``{label: count}`` distribution. Works for
+    ANY per-sample signal — loss, accuracy, a second loss, any metric. Reads the
+    full history once (cheap end-of-report reduction). *classifier* (trajectory
+    -> label|None) defaults to the loss-shaped one (for a decreasing metric);
+    pass your own for e.g. accuracy (increasing). *tag_name* defaults to
+    ``'<signal>_shape'``."""
+    clf = classifier or classify_loss_shape
+    if tag_name is None:
+        tag_name = signal_name.split("/")[-1] + "_shape"
+    series = {}
+    for sid, step, val, _ in query_signal_history(signal_name):
+        series.setdefault(sid, []).append((step, val))
+    by_label = {}
+    for sid, pts in series.items():
+        label = clf([v for _, v in sorted(pts)])
+        if label is not None:
+            by_label.setdefault(label, []).append(sid)
+    for label, sids in by_label.items():
+        set_categorical_tag(sids, tag_name, label)
+    return {k: len(v) for k, v in by_label.items()}
+
+
+def write_loss_shapes(loss_signal="loss_sample", classifier=None):
+    """Convenience wrapper over :func:`write_signal_shapes` for the loss signal
+    (tag ``loss_shape``, default loss classifier)."""
+    return write_signal_shapes(loss_signal, tag_name="loss_shape", classifier=classifier)
+
+
 def write_dataframe(
     path: str | None = None,
     format: str = "json",
@@ -4013,9 +4077,14 @@ def write_dataframe(
     sample_id=None,
     instance_id=None,
     verbose: bool = False,
+    loss_shape_signal: str | None = None,
 ) -> str:
     """Dump the WeightsLab sample dataframe to *path* as JSON or CSV.
 
+    When *loss_shape_signal* is set (e.g. ``"loss_sample"``), the default loss-
+    shape classifier runs first (tags each sample + logs a distribution summary),
+    so a report written every N steps carries an up-to-date shape summary.
+
     Parameters
     ----------
     path : str, optional
@@ -4087,6 +4156,14 @@ def write_dataframe(
     # before reading, so the report includes the latest steps' derived signals.
     drain_signals()
 
+    # Default report summary: tag each sample's loss shape + log the distribution.
+    if loss_shape_signal is not None:
+        try:
+            dist = write_loss_shapes(loss_shape_signal)
+            logger.info("[WeightsLab] loss-shape summary (%s): %s", loss_shape_signal, dist)
+        except Exception as e:
+            logger.debug("loss-shape summary skipped: %s", e)
+
     logger.debug(
         "write_dataframe called: path=%r, format=%r, columns=%r, "
         "sample_id=%r, instance_id=%r",

From c8610f71bda6c04cd1bd031718b1edace6fdaedf Mon Sep 17 00:00:00 2001
From: Alexandru Rotaru <rotarualexandruandrei94@gmail.com>
Date: Fri, 10 Jul 2026 14:03:52 +0200
Subject: [PATCH 19/19] refactor(signals): clean, layered, overridable shape
 API
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- trajectory_stats(values): generic reusable feature layer (renamed from
  loss_trajectory_stats — works on any trajectory).
- classify_loss_shape: built on trajectory_stats; every threshold a named param.
- write_signal_shapes: engine over any per-sample signal; write_loss_shapes: wrapper.
- enable_loss_shape_signal(loss_signal, every=N): live per-step @wl.signal toggle
  (the heavier live counterpart to the report-time write_loss_shapes).
- example ws-signals-mnist: drop the local classifier, use the core hook
  wl.write_dataframe(loss_shape_signal=...).
Exports + tests updated.
---
 tests/general/test_signal_refinements.py      | 20 ++++-
 weightslab/__init__.py                        |  4 +-
 .../Usecases/ws-signals-mnist/main.py         | 53 +++-----------
 weightslab/src.py                             | 73 ++++++++++++++-----
 4 files changed, 87 insertions(+), 63 deletions(-)

diff --git a/tests/general/test_signal_refinements.py b/tests/general/test_signal_refinements.py
index e5aa6ff2..681c9970 100644
--- a/tests/general/test_signal_refinements.py
+++ b/tests/general/test_signal_refinements.py
@@ -135,8 +135,26 @@ def test_thresholds_configurable(self):
         self.assertNotEqual(wl.classify_loss_shape(traj), "monotonic")   # default 0.4 unmet
         self.assertEqual(wl.classify_loss_shape(traj, drop_learned=0.25), "monotonic")
 
+    def test_trajectory_stats_reusable(self):
+        s = wl.trajectory_stats([2.0, 1.0, 0.5, 0.5, 0.4])
+        self.assertAlmostEqual(s["drop"], (2.0 - 0.4) / 2.0, places=5)
+        self.assertIn("cv", s); self.assertIn("tail_cv", s); self.assertIn("argmin_frac", s)
+        self.assertIsNone(wl.trajectory_stats([1.0]))   # < 2 points
+
+    def test_enable_live_signal_registers(self):
+        _REGISTERED_SIGNALS.clear()
+        try:
+            wl.enable_loss_shape_signal(loss_signal="loss_sample", name="sig/live_shape", every=5)
+            self.assertIn("sig/live_shape", _REGISTERED_SIGNALS)
+            meta = _REGISTERED_SIGNALS["sig/live_shape"]._wl_signal_meta
+            self.assertEqual(meta.get("inputs"), ["loss_sample"])
+            self.assertEqual(meta.get("compute_every_n_steps"), 5)
+        finally:
+            _REGISTERED_SIGNALS.clear()
+
     def test_exports(self):
-        for name in ("write_loss_shapes", "write_signal_shapes", "classify_loss_shape"):
+        for name in ("classify_loss_shape", "trajectory_stats", "write_loss_shapes",
+                     "write_signal_shapes", "enable_loss_shape_signal"):
             self.assertTrue(hasattr(wl, name), name)
         self.assertIn("monotonic", wl.LOSS_SHAPES)
 
diff --git a/weightslab/__init__.py b/weightslab/__init__.py
index 395e5776..6f8d222a 100644
--- a/weightslab/__init__.py
+++ b/weightslab/__init__.py
@@ -11,7 +11,7 @@
 import logging
 import threading
 
-from .src import watch_or_edit, start_training, serve, keep_serving, save_signals, save_instance_signals, save_group_signals, tag_samples, register_categorical_tag, set_categorical_tag, discard_samples, get_samples_by_tag, get_discarded_samples, signal, eval_fn, compute_signals, SignalContext, BatchSignalContext, StaleSignalError, drain_signals, clear_all, run_pending_evaluation, trigger_pending_evaluation_async, query_signal_history, query_sample_history, query_instance_history, write_history, write_dataframe, classify_loss_shape, write_loss_shapes, write_signal_shapes, LOSS_SHAPES, get_current_experiment_hash, pointcloud_thumbnail, pointcloud_boxes
+from .src import watch_or_edit, start_training, serve, keep_serving, save_signals, save_instance_signals, save_group_signals, tag_samples, register_categorical_tag, set_categorical_tag, discard_samples, get_samples_by_tag, get_discarded_samples, signal, eval_fn, compute_signals, SignalContext, BatchSignalContext, StaleSignalError, drain_signals, clear_all, run_pending_evaluation, trigger_pending_evaluation_async, query_signal_history, query_sample_history, query_instance_history, write_history, write_dataframe, classify_loss_shape, trajectory_stats, write_loss_shapes, write_signal_shapes, enable_loss_shape_signal, LOSS_SHAPES, get_current_experiment_hash, pointcloud_thumbnail, pointcloud_boxes
 from .backend.ledgers import GLOBAL_LEDGER as ledger
 from .art import _BANNER
 from .utils.logs import setup_logging, set_log_directory, is_main_process
@@ -172,6 +172,8 @@ def _clean(v: str) -> str:
     "classify_loss_shape",
     "write_loss_shapes",
     "write_signal_shapes",
+    "trajectory_stats",
+    "enable_loss_shape_signal",
     "LOSS_SHAPES",
 
     "pointcloud_thumbnail",
diff --git a/weightslab/examples/Usecases/ws-signals-mnist/main.py b/weightslab/examples/Usecases/ws-signals-mnist/main.py
index 7c712d3c..a2b32824 100644
--- a/weightslab/examples/Usecases/ws-signals-mnist/main.py
+++ b/weightslab/examples/Usecases/ws-signals-mnist/main.py
@@ -4,17 +4,16 @@
   entropy    from ctx.logits when the loss fires
   loss_norm  reactive, from the logged loss
   hardness   reactive, from loss + entropy
-  loss_shape reactive, classifies each sample's loss trajectory (live signal,
-             not an end-of-run function). Reads history, so it's throttled by
-             WL_SHAPE_EVERY (1 = every step, full coverage; higher = cheaper).
+
+loss_shape comes from WeightsLab's built-in classifier, applied on report write
+(wl.write_dataframe(loss_shape_signal=...)) — no classifier code lives here.
 
 Universal loss: the watched crit runs on the test split each epoch too, so test
 samples get a loss trajectory and a shape as well.
 
-Env: WL_STRESS_EPOCHS, WL_STRESS_OUT, WL_SHAPE_EVERY.
+Env: WL_STRESS_EPOCHS, WL_STRESS_OUT.
 """
 import os, time, gc, json, shutil
-from collections import defaultdict, Counter
 
 import numpy as np
 import torch
@@ -27,39 +26,11 @@
 from weightslab.components.global_monitoring import guard_training_context, guard_testing_context
 
 LOSS = "loss_sample"
-SHAPES = ["monotonic", "plateaued", "Flat_high", "high_variance", "U_Shape", "Spiked"]
 OUT = os.environ.get("WL_STRESS_OUT", "/tmp/wl_stress")
 EPOCHS = int(os.environ.get("WL_STRESS_EPOCHS", "10"))
-SHAPE_EVERY = int(os.environ.get("WL_SHAPE_EVERY", "1"))
 BATCH = 64
 
 
-def classify_shape(values):
-    """Loss trajectory -> shape index (or -1 with < 5 points)."""
-    y = np.asarray(values, dtype=float)
-    if y.size < 5:
-        return -1
-    n = y.size
-    rng = max(float(y.max() - y.min()), 1e-8)
-    drop = (float(y[0]) - float(y[-1])) / (abs(float(y[0])) + 1e-8)
-    cv = float(y.std()) / (abs(float(y.mean())) + 1e-8)
-    argmin = int(np.argmin(y))
-    rebound = (float(y[-1]) - float(y.min())) / rng
-    tail = y[int(0.6 * n):]
-    tail_flat = float(tail.std()) / (abs(float(tail.mean())) + 1e-8) < 0.1
-    if 0.2 * n < argmin < 0.8 * n and rebound > 0.3:
-        return SHAPES.index("U_Shape")
-    if drop > 0.4:
-        return SHAPES.index("monotonic")
-    if drop > 0.15 and tail_flat:
-        return SHAPES.index("plateaued")
-    if float(np.diff(y).max()) / rng > 0.5:
-        return SHAPES.index("Spiked")
-    if cv > 0.5:
-        return SHAPES.index("high_variance")
-    return SHAPES.index("Flat_high")
-
-
 def rss_gb():
     try:
         for ln in open("/proc/self/status"):
@@ -154,14 +125,9 @@ def loss_norm(b):
     def hardness(b):
         return b.inputs[LOSS] * b.inputs["sig/entropy"]
 
-    @wl.signal(name="sig/loss_shape", inputs=[LOSS], batched=True, compute_every_n_steps=SHAPE_EVERY)
-    def loss_shape(b):
-        hist = b.history(LOSS)   # {uid: [loss values in step order]}
-        return np.array([classify_shape(hist[s]) for s in b.sample_ids], dtype=float)
-
     wl.serve(serving_grpc=False, serving_cli=False)
     wl.start_training(timeout=0)
-    log(f"[run] tracked train={len(train_ds)} test={len(test_ds)} | shape_every={SHAPE_EVERY}")
+    log(f"[run] tracked train={len(train_ds)} test={len(test_ds)}")
 
     def test_eval():
         """Universal loss: watched crit over the test split each epoch."""
@@ -196,16 +162,17 @@ def test_eval():
             f"= +{100*(wl_ms/vanilla_ms-1):.0f}% | RSS {rec['rss_gb']:.2f} GB | {rec['elapsed_s']:.0f}s")
 
     import pandas as pd
-    path = wl.write_dataframe(OUT + "/report.csv", format="csv", columns="signals")
+    # loss_shape comes from WeightsLab's built-in classifier, run on report write.
+    path = wl.write_dataframe(OUT + "/report.csv", format="csv",
+                              columns=["signals", "tags"], loss_shape_signal=LOSS)
     df = pd.read_csv(path)
-    sc = [c for c in df.columns if c.endswith("sig/loss_shape")][0]
-    dist = Counter(SHAPES[int(v)] for v in df[sc].dropna() if v >= 0)
+    sc = [c for c in df.columns if c.endswith("loss_shape")][0]
     med = float(np.median(step_times))
     metrics.close()
     log("\n[run] ===== SUMMARY =====")
     log(f"[run] vanilla {vanilla_ms:.2f} ms | WL median {med:.2f} ms/step (+{100*(med/vanilla_ms-1):.0f}%)")
     log(f"[run] report {len(df)} rows | loss_shape covered {df[sc].notna().sum()}/{len(df)}")
-    log(f"[run] shape distribution: {dict(dist)}")
+    log(f"[run] shape distribution: {df[sc].value_counts(dropna=True).to_dict()}")
     log(f"[run] report -> {path}")
 
 
diff --git a/weightslab/src.py b/weightslab/src.py
index d3fe09ae..483f046d 100644
--- a/weightslab/src.py
+++ b/weightslab/src.py
@@ -4009,33 +4009,51 @@ def write_history(
 LOSS_SHAPES = ["monotonic", "plateaued", "Flat_high", "high_variance", "U_Shape", "Spiked"]
 
 
-def classify_loss_shape(values, min_points=5, drop_learned=0.4, drop_plateau=0.15,
-                        tail_flat_cv=0.1, spike_jump=0.5, noisy_cv=0.5, u_rebound=0.3):
-    """Default per-sample loss-trajectory classifier -> one of ``LOSS_SHAPES``
-    (or ``None`` with < *min_points*). Every threshold is a named, scale-invariant
-    param with a natural default, so you can tune it without rewriting the
-    classifier, e.g. ``write_loss_shapes(classifier=lambda v:
-    wl.classify_loss_shape(v, drop_learned=0.5))``."""
+def trajectory_stats(values):
+    """Scale-invariant summary stats of one sample's trajectory — the reusable
+    feature layer. Returns a dict (or ``None`` with < 2 points) so you can build
+    a custom classifier on top without re-deriving the features::
+
+        def my_clf(v):
+            s = wl.trajectory_stats(v)
+            return None if s is None else ("fast" if s["drop"] > 0.6 else "slow")
+    """
     y = np.asarray(values, dtype=float)
-    if y.size < min_points:
+    if y.size < 2:
         return None
     n = y.size
     rng = max(float(y.max() - y.min()), 1e-8)
-    drop = (float(y[0]) - float(y[-1])) / (abs(float(y[0])) + 1e-8)
-    cv = float(y.std()) / (abs(float(y.mean())) + 1e-8)
-    argmin = int(np.argmin(y))
-    rebound = (float(y[-1]) - float(y.min())) / rng
     tail = y[int(0.6 * n):]
-    tail_flat = float(tail.std()) / (abs(float(tail.mean())) + 1e-8) < tail_flat_cv
-    if 0.2 * n < argmin < 0.8 * n and rebound > u_rebound:
+    return {
+        "n": n,
+        "first": float(y[0]), "last": float(y[-1]),
+        "drop": (float(y[0]) - float(y[-1])) / (abs(float(y[0])) + 1e-8),
+        "cv": float(y.std()) / (abs(float(y.mean())) + 1e-8),
+        "argmin_frac": float(np.argmin(y)) / n,
+        "rebound": (float(y[-1]) - float(y.min())) / rng,
+        "max_up_jump": float(np.diff(y).max()) / rng,
+        "tail_cv": float(tail.std()) / (abs(float(tail.mean())) + 1e-8),
+    }
+
+
+def classify_loss_shape(values, min_points=5, drop_learned=0.4, drop_plateau=0.15,
+                        tail_flat_cv=0.1, spike_jump=0.5, noisy_cv=0.5, u_rebound=0.3):
+    """Default classifier -> one of ``LOSS_SHAPES`` (or ``None`` with <
+    *min_points*). Built on :func:`trajectory_stats` (reuse those features
+    for your own rules) with every threshold a named, tunable param. Override the
+    whole thing by passing your own ``classifier`` to :func:`write_signal_shapes`."""
+    s = trajectory_stats(values)
+    if s is None or s["n"] < min_points:
+        return None
+    if 0.2 < s["argmin_frac"] < 0.8 and s["rebound"] > u_rebound:
         return "U_Shape"
-    if drop > drop_learned:
+    if s["drop"] > drop_learned:
         return "monotonic"
-    if drop > drop_plateau and tail_flat:
+    if s["drop"] > drop_plateau and s["tail_cv"] < tail_flat_cv:
         return "plateaued"
-    if float(np.diff(y).max()) / rng > spike_jump:
+    if s["max_up_jump"] > spike_jump:
         return "Spiked"
-    if cv > noisy_cv:
+    if s["cv"] > noisy_cv:
         return "high_variance"
     return "Flat_high"
 
@@ -4070,6 +4088,25 @@ def write_loss_shapes(loss_signal="loss_sample", classifier=None):
     return write_signal_shapes(loss_signal, tag_name="loss_shape", classifier=classifier)
 
 
+def enable_loss_shape_signal(loss_signal="loss_sample", name="sig/loss_shape",
+                             every=1, classifier=None):
+    """Register a LIVE per-step ``@wl.signal`` that classifies each sample's
+    loss-trajectory-so-far into an int-coded shape (index into ``LOSS_SHAPES``,
+    ``-1`` before there's enough history), updated every *every* steps.
+
+    This is the live counterpart to :func:`write_loss_shapes` (report-time). It's
+    heavier — it reads history on every fire — so throttle with *every* or prefer
+    the report-time path for a definitive, full-coverage tag."""
+    clf = classifier or classify_loss_shape
+
+    @signal(name=name, inputs=[loss_signal], batched=True, compute_every_n_steps=every)
+    def _live_shape(b):
+        hist = b.history(loss_signal)
+        return np.array([(LOSS_SHAPES.index(lbl) if (lbl := clf(hist[s])) is not None else -1)
+                         for s in b.sample_ids], dtype=float)
+    return _live_shape
+
+
 def write_dataframe(
     path: str | None = None,
     format: str = "json",