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 diff --git a/tests/general/test_signal_refinements.py b/tests/general/test_signal_refinements.py new file mode 100644 index 00000000..681c9970 --- /dev/null +++ b/tests/general/test_signal_refinements.py @@ -0,0 +1,163 @@ +"""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]) + + +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_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 ("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) + + +if __name__ == "__main__": + unittest.main() diff --git a/weightslab/__init__.py b/weightslab/__init__.py index e41340b6..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, 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, 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 @@ -148,6 +148,9 @@ def _clean(v: str) -> str: "get_samples_by_tag", "get_discarded_samples", "SignalContext", + "BatchSignalContext", + "StaleSignalError", + "drain_signals", "clear_all", "seed_everything", "run_pending_evaluation", @@ -166,6 +169,12 @@ def _clean(v: str) -> str: "write_history", "write_dataframe", + "classify_loss_shape", + "write_loss_shapes", + "write_signal_shapes", + "trajectory_stats", + "enable_loss_shape_signal", + "LOSS_SHAPES", "pointcloud_thumbnail", "pointcloud_boxes", diff --git a/weightslab/backend/logger.py b/weightslab/backend/logger.py index 4e6f95b1..a29bdf07 100644 --- a/weightslab/backend/logger.py +++ b/weightslab/backend/logger.py @@ -29,9 +29,12 @@ staging appends and flushes take the same lock. """ +import functools import json +import os import threading import time +from collections import defaultdict import duckdb import pandas as pd @@ -81,6 +84,20 @@ 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 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). + # 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=_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() @@ -167,7 +184,10 @@ 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)->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) @@ -198,11 +218,22 @@ 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): + # 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(), )) + self._qps_version[graph_name] += 1 # invalidate this signal's cached reads self._maybe_autoflush() + def _invalidate_qps_cache(self) -> None: + """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() + 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 +275,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 +475,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 +745,75 @@ 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). + + 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): + """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] 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): + """``(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`` = 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 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: + 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, 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/examples/Usecases/wl-detection-signals-classification/utils/criterions.py b/weightslab/examples/Usecases/wl-detection-signals-classification/utils/criterions.py index 42ee6690..31c443f5 100644 --- a/weightslab/examples/Usecases/wl-detection-signals-classification/utils/criterions.py +++ b/weightslab/examples/Usecases/wl-detection-signals-classification/utils/criterions.py @@ -268,6 +268,7 @@ def decode_predictions(outputs, grid_size, conf_thresh=0.3, max_det=10): # U_Shape -> learned then forgotten (catastrophic interference) # Spiked -> sudden jump at some step (data/aug/version change) +from weightslab.backend import ledgers # Allowed values for the categorical tag, in display order. LOSS_SHAPE_LABELS = [ "monotonic", "plateaued", "Flat_high", 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/Usecases/ws-signals-mnist/main.py b/weightslab/examples/Usecases/ws-signals-mnist/main.py new file mode 100644 index 00000000..a2b32824 --- /dev/null +++ b/weightslab/examples/Usecases/ws-signals-mnist/main.py @@ -0,0 +1,180 @@ +"""Per-sample signals on MNIST — readable, zero save_signals, with train + eval. + +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 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. +""" +import os, time, gc, json, shutil + +import numpy as np +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import Dataset +from torchvision import datasets, transforms + +import weightslab as wl +from weightslab.components.global_monitoring import guard_training_context, guard_testing_context + +LOSS = "loss_sample" +OUT = os.environ.get("WL_STRESS_OUT", "/tmp/wl_stress") +EPOCHS = int(os.environ.get("WL_STRESS_EPOCHS", "10")) +BATCH = 64 + + +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): + 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), + 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) + + +class MNISTIdx(Dataset): + """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 len(self.m) + + def __getitem__(self, i): + img, lab = self.m[i] + return self.t(img), self.base + i, lab + + def fast_get_label(self, i): + return int(self.m.targets[i]) + + +def log(msg): + print(msg, flush=True) + + +def main(): + 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) + 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) + 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(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) + + @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() + + @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=0) + 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.""" + 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) + # 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() + 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 + # 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("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: {df[sc].value_counts(dropna=True).to_dict()}") + log(f"[run] report -> {path}") + + +if __name__ == "__main__": + main() diff --git a/weightslab/src.py b/weightslab/src.py index 9b8e7d11..fead8219 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 @@ -71,6 +72,231 @@ 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, 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 + # 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: + return None + cols[inp] = np.array([at[sid] 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.""" + _warn_on_signal_cycles() # lazy, once/process — fires even with no wl.serve + 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 + + # 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() + 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, fresh_cache=fired) + 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 + 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(): + """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', {}) + 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): + """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 + _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} + + +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 @@ -78,18 +304,32 @@ 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, logits=None, preds=None, targets=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) + 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]: @@ -163,6 +403,86 @@ 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"), 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. + + 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 + + +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, + 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 + 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 + + 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)] = (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) + # ##################################################################################################################### # WEIGHTSLAB INTERNAL FUNCTIONS FOR LOGGING, SIGNAL EXTRACTION, WRAPPING, ETC. (not typically called directly by users) @@ -273,9 +593,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): @@ -466,7 +798,9 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw): # Extract torch function parameters _ = wl_kw.get('flag') - # preds_raw = a[0] if len(a) > 0 else None + # Kept in-memory for ctx.logits (logit-derived signals); NOT persisted unless + # ledger_store_preds_raw is set (see the gated save_signals below). + preds_raw = a[0] if len(a) > 0 else None # User parameters batch_ids = wl_kw.get('batch_ids') @@ -632,31 +966,73 @@ 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, - logger=ledgers.get_logger(), + _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 + # 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=_lg, dataframe=df_proxy, - origin=kwargs.get('origin', 'train') + 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, ) - 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()) + # 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 + 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, + logger=_lg, + 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 + 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())}") - _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: logger.error(f"Dynamic signal {name} failed: {e}") pass # User function error, skip @@ -675,6 +1051,14 @@ 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. Default off (dev disabled it); set the hyperparam to re-enable. + try: + _store_preds = bool(get_hyperparams().get('ledger_store_preds_raw', False)) + except Exception: + _store_preds = False + # Enqueue signals and data save_signals( signals=signals, @@ -685,6 +1069,17 @@ 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 ) + # 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: + _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: + 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 @@ -1074,6 +1469,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) @@ -1240,6 +1636,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_meta['min_step'] = min_step func._wl_signal_name = reg_name @@ -1686,7 +2089,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. @@ -1750,8 +2154,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: @@ -1844,6 +2250,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: + _dispatch_or_enqueue(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, @@ -3590,6 +4008,107 @@ def write_history( return path +LOSS_SHAPES = ["monotonic", "plateaued", "Flat_high", "high_variance", "U_Shape", "Spiked"] + + +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 < 2: + return None + n = y.size + rng = max(float(y.max() - y.min()), 1e-8) + tail = y[int(0.6 * n):] + 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 s["drop"] > drop_learned: + return "monotonic" + if s["drop"] > drop_plateau and s["tail_cv"] < tail_flat_cv: + return "plateaued" + if s["max_up_jump"] > spike_jump: + return "Spiked" + if s["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 + ``'_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 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", @@ -3597,9 +4116,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 @@ -3667,6 +4191,18 @@ 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() + + # 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",