Signal refinements: vectorized + reactive @wl.signal, overhead tuning, example

.gitignore

@@ -11,11 +11,15 @@ venv
 runs
 data
 MagicMock
+drop
+htmlcov
+htmlcov-report
 
 # Ignore hidden directories
 .history
 .lh
 .claude
+.coverage
 .chat
 .ai
 .pytest_cache

tests/general/test_signal_refinements.py

@@ -0,0 +1,145 @@
+"""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_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()

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, 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
@@ -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,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",

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,