diff --git a/docs/perf/o_change_register.md b/docs/perf/o_change_register.md new file mode 100644 index 00000000..270f3417 --- /dev/null +++ b/docs/perf/o_change_register.md @@ -0,0 +1,170 @@ +# Interactivity for 100GB+ datasets — register of O(data) operations + +Goal: every per-step and per-request operation should cost **O(change)** or +**O(page)**, never **O(dataset)**. Today several do, so cost grows with the +dataset while the actual work stays constant. + +> **Baseline caveat.** The measurements below were taken against a stale +> in-place copy of weightslab (`~/weightslab_src`, 1.3.3+multiview), which +> differs from `dev` in 51 files. Two serving costs listed there are **already +> fixed on dev** (see §B). Serving numbers must be re-measured on this branch +> before any serving change is attributed an improvement. The storage findings +> (§A) were re-verified against dev and still hold. + +Reference measurements (UltraEdit, 3,959,093 rows, ~19 cols, A10G box): + +| | measured | +|---|---| +| bare-torch step (no WL) | 1,162 ms → 20.66 samples/s | +| with WL, no UI client | ~5.5 samples/s (**3.8× slower**) | +| with WL + image requests | ~1.5 samples/s (**13.8× slower**) | +| per-step signal write itself | **4 ms (0.34%)** — already fine | +| grid page latency (64 imgs) | p50 5.2 s, p95 9.6 s | + +The signal path is not the problem. Storage write-amplification and the view +rebuild are. + +--- + +## A. STORAGE — `data/h5_dataframe_store.py` + +`upsert()` **receives** only dirty rows but **implements** a full table +replacement. + +| line | operation | cost | +|---|---|---| +| 693 | `_create_backup()` — full file copy **before every upsert** | O(file) | +| 708 | `existing = store.select(key)` — read entire table | O(N) | +| ~768 | `pd.concat([existing, delta])` | O(N) | +| ~772 | `existing[~existing.index.duplicated()]` — dedupe all rows | O(N) | +| ~785 | `_decategorize_for_storage(existing)` | O(N) | +| 801 | `store.remove(key)` — drop table | O(N) | +| 804 | `store.append(..., data_columns=True)` — rewrite + index **every** column | O(N·cols) | +| 846–883 | same read/remove/rewrite in the column-delete path | O(N) | + +**Amplification:** ~5 KB of changed signals per flush → ~200 MB written, +roughly **40,000×**. At `ledger_flush_interval=3.0s` vs ~1.5 s steps, that is a +full-table rewrite about every 2 steps. + +Fix direction: append new rows; modify existing rows in place +(`select_as_coordinates` + `table.modify_rows`). Backup incrementally, not per +upsert. No `data_columns=True` — no `store.select()` in this file uses `where=`, +so those per-column indexes are built and never read. + +*(A previous attempt to narrow `data_columns` broke the write path entirely — +678 upsert failures, zero persisted data. Any change here needs a +write→read→assert-contents check, not just a timing check.)* + +## B. SERVING — `trainer/services/data_service.py` + +`_pull_into_all_data_view_df()` (line 938) runs several full-frame passes. + +**Already fixed on dev — do not re-report as wins:** +- the collapse no longer re-enters `get_combined_df()`; the pulled frame is + passed in, so the frame is copied once, not twice +- `array_proxy` no longer does a per-cell `.apply(convert_to_proxy)` (was + 1,660 ms at 4M rows) + +Remaining, to be **re-measured on this branch**: + +| line | operation | cost (measured @4M) | +|---|---|---| +| 946 | `get_combined_df()` → `dataframe_manager:2140 self._df.copy()` | 101 ms–1.2 s | +| — | `get_collapse_annotations_to_samples_df(df)` — groupby collapse | 6,326 ms* | +| — | `safe_reset_index(df)` | 1,912 ms | +| — | `set_index([origin, sample_id])` | O(N) | +| 3636 | `updated_df.reindex(target_order)` | 290 ms | + +Callers — each one is a full O(N) rebuild: lines **440, 851, 3593, 4605, 4626**, +reached from `GetDataSamples`, `GetMetaData`, `EditDataSample`, `GetDataSplits`. + +Held under `_update_lock`, which the trainer also needs → measured lock holds of +39–126 s and the 3.7× training penalty while browsing. + +**The collapse is provably a no-op when `annotation_id.max() == 0`** (UltraEdit +is exactly 1:1) yet still costs 6.3 s per rebuild. + +Fix direction: serve a page from the source frame by index (O(page)); rebuild +the full view only for genuinely global operations (histogram, global sort); +apply deltas rather than rebuilding; never hold the writer lock across a +rebuild — build off-lock and swap the reference. + +## C. OTHER FULL SCANS + +| location | note | +|---|---| +| `dataframe_manager:1875` `data_snapshot.iterrows()` | input is O(change), but row-wise Python per flush | +| `dataframe_manager:2400` `.apply(lambda …)` | per cell | +| `data_service:1200` `_compute_natural_sort_stats` | builds a list of one Series per row (4M objects). Gated off (`compute_natural_sort=False`) — latent | +| `data_service:538` PreviewCache | bounded by `WL_MAX_PREVIEW_CACHE_SIZE` — OK | + +## D. ALREADY O(change) — keep + +- `self._pending` dirty-row set (`dataframe_manager:95, 751, 761`) +- flush work set: `work = list(self._pending)` (`:1827`) +- `_origin_revisions` per-origin version counters (`:94, 1235`) + +The bookkeeping needed for differential updates already exists; the storage and +view layers just don't use it. + +\* measured on the stale copy; re-measure on dev. + +## Measurement protocol + +Fixed workload: **1,000 train samples** = 41 steps at batch 24. Every change is +reported as: + +1. wall-clock for the 41 steps, vs the bare-torch floor +2. bytes written to H5 for those steps +3. grid-page latency (64 images) and training throughput **while** serving +4. **ledger contents verified** — signal columns present, measured-row count + +(4) is not optional: a previous "10× win" was writes silently failing. + +--- + +# E. Triage — which call sites need a full reconstruction + +`_slowUpdateInternals()` rebuilds the whole view: `copy → collapse → reset_index +→ set_index → reindex`. It has 18 call sites, and almost none of them need +that. Most just want **fresh values for rows the trainer touched**, which is +`O(change)`. + +`_fastUpdateInternals()` applies only dirty rows, via a maintained +`sample_id → position` map (`_rebuild_view_pos_map`), and returns `False` — +falling back to the full rebuild — whenever it cannot safely apply: + +- no view yet, or no position map (first build) +- a dirty `sample_id` absent from the map (new rows ⇒ structural change) +- backlog > `max_dirty` (a rebuild is genuinely cheaper) + +So the worst case is today's behaviour, never wrong data. + +| site | routing | why | +|---|---|---| +| `_bg_view_refresh` | **fast** | exists purely to refresh values after a stale read — the textbook differential case, and the one that holds `_lock` against the trainer | +| `_process_get_data_samples` | **fast** | grid fetch needs current values, not a new frame | +| `_compute_custom_signals` | **fast** | writes new signal *values*; schema unchanged | +| `GetDataSplits` | **fast** | read-only summary | +| `EditDataSample` ×5 | **fast** | per-sample value edits | +| `EditDataSample` ×3 (`df.modify`, `df.drop_column`) | **full** | changes the schema — differential cannot add/remove columns | +| `ApplyDataQuery` `@reset`/`@clear` | **full** | clears `_is_filtered` to restore the full universe; a differential updates values but cannot restore *dropped rows* | +| `ApplyDataQuery` filter + agent paths | **full** (deferred) | a forced rebuild preserves `_is_filtered` (`:3691`), so swapping in a differential changes which rows the user sees. Rare, user-initiated, low perf value, high blast radius — not worth the risk until the filter semantics are pinned down | +| `_compute_natural_sort_stats` | **full** | gated off (`compute_natural_sort=False`); latent | +| `_manual_save_data_state` | **full** | explicit user save; correctness over speed | + +**Kill-switch:** `WL_FAST_VIEW=0` disables the differential and the position-map +build, reproducing prior behaviour exactly. This is what makes a like-for-like +A/B possible from a single tree. + +## Why the no-client benchmark cannot show this + +A 41-step run with no UI client attached records **0 rebuild events** — nothing +calls `_slowUpdateInternals` at all, so the fast path has nothing to improve and +correctly measures as no change. The rebuild cost only materialises when a +client is attached, which is the case that measured **3.7× slower** with p50 +grid latency of 5.2 s. + +The A/B is therefore run under load: baseline → under-load → recovery phases +within one training process (`t_imgload.py`), so each arm is normalised against +its own idle throughput. diff --git a/weightslab/backend/dataloader_interface.py b/weightslab/backend/dataloader_interface.py index 9d796a97..19f7f6b4 100644 --- a/weightslab/backend/dataloader_interface.py +++ b/weightslab/backend/dataloader_interface.py @@ -43,6 +43,44 @@ _DENY_LIST_REFRESH_INTERVAL = 32 +def _close_inherited_h5_fds(worker_id: int = 0) -> None: + """DataLoader worker_init_fn: drop HDF5 handles inherited from the parent. + + torch forks workers, so every fd the parent had open at fork time is + duplicated into the child -- including the ledger store. The child never + uses them, but their mere existence makes HDF5 refuse the parent's + read-write open, which silently kills ledger persistence. + + Closes rather than just dropping the Python object: the fd is what holds + the file, and the child has no Python-level reference to it at all. + """ + import os + try: + fd_dir = "/proc/self/fd" + for entry in os.listdir(fd_dir): + try: + target = os.readlink(os.path.join(fd_dir, entry)) + except OSError: + continue + if target.endswith(".h5") or target.endswith(".h5.lock"): + try: + os.close(int(entry)) + except OSError: + pass + except Exception: + # Never let cleanup break a worker: a leaked handle degrades + # persistence, a raising worker_init_fn kills the run. + pass + + +def _with_worker_init(kwargs: dict, num_workers: int) -> dict: + """Attach the fd cleanup unless the caller supplied its own init.""" + if num_workers and not kwargs.get("worker_init_fn"): + kwargs = dict(kwargs) + kwargs["worker_init_fn"] = _close_inherited_h5_fds + return kwargs + + def _resolve_safe_num_workers(dataset: Any, num_workers: int, loader_name: Optional[str] = None) -> int: """Clamp worker count for datasets that cannot be pickled by Windows spawn.""" try: @@ -177,6 +215,12 @@ def _get_deny_list_revision(self) -> Optional[tuple[str, int]]: try: origin = self._get_current_origin() df_manager = get_dataframe() + if origin and df_manager is not None and hasattr(df_manager, "get_discard_revision"): + # Deliberately NOT get_origin_revision: that moves on every + # per-sample signal write, i.e. every training step, so the + # cache below could never hit and __len__ rescanned 3.96M rows + # per batch. Discard state is what the deny-list depends on. + return ("discard", int(df_manager.get_discard_revision(origin))) if origin and df_manager is not None and hasattr(df_manager, "get_origin_revision"): return ("origin", int(df_manager.get_origin_revision(origin))) except Exception: @@ -514,6 +558,7 @@ def __init__( self.tracked_dataset, batch_sampler=batch_sampler, num_workers=num_workers, + worker_init_fn=_close_inherited_h5_fds, pin_memory=pin_memory, collate_fn=collate_fn, persistent_workers=self._should_persist_workers(num_workers), @@ -1134,6 +1179,7 @@ def restore_iteration_state(self, state: dict) -> None: self.tracked_dataset, batch_sampler=sampler, num_workers=num_workers, + worker_init_fn=_close_inherited_h5_fds, pin_memory=pin_memory, collate_fn=collate_fn, # Ensure no conflicting args are passed alongside batch_sampler @@ -1220,6 +1266,7 @@ def set_batch_size(self, new_batch_size: int) -> None: self.tracked_dataset, batch_sampler=sampler, num_workers=num_workers, + worker_init_fn=_close_inherited_h5_fds, pin_memory=pin_memory, collate_fn=collate_fn, persistent_workers=self._should_persist_workers(num_workers), @@ -1232,6 +1279,7 @@ def set_batch_size(self, new_batch_size: int) -> None: batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, + worker_init_fn=_close_inherited_h5_fds, drop_last=drop_last, pin_memory=pin_memory, collate_fn=collate_fn, diff --git a/weightslab/backend/logger.py b/weightslab/backend/logger.py index 7ce7febe..0c2cbc3e 100644 --- a/weightslab/backend/logger.py +++ b/weightslab/backend/logger.py @@ -36,7 +36,7 @@ import os import threading import time -from collections import defaultdict +from collections import defaultdict, deque import duckdb import pandas as pd @@ -64,6 +64,18 @@ _STAGE_FLUSH_THRESHOLD = 50_000 # How often the background flush thread wakes up (see LoggerQueue._flush_loop). +def _default_history_tail() -> int: + """Recent points kept per sample for signal-DAG history reads. + + Bounded so history() costs O(batch) instead of scanning the whole + per_sample table (140ms at 20M rows, once per step, and growing). + """ + try: + return max(0, int(os.environ.get("WL_HISTORY_TAIL", "16"))) + except (TypeError, ValueError): + return 16 + + def _default_flush_interval_seconds() -> float: try: return float(os.environ.get("WL_LOGGER_FLUSH_INTERVAL_SECONDS", "2.0")) @@ -276,6 +288,10 @@ def __init__(self, register: bool = True, db_path: str = ":memory:") -> None: _qps_maxsize = int(os.environ.get("WL_QUERY_CACHE_MAXSIZE", "2048")) self._qps_version: dict = defaultdict(int) self._qps_cache_step: int = -1 + # {signal: {sample_id: deque}} -- the recent tail of each sample's + # per-sample values, maintained on write so history() never scans. + self._tail_len = _default_history_tail() + self._recent_tail: dict = defaultdict(dict) 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) @@ -757,6 +773,20 @@ def _next_seq(self) -> int: self._seq += 1 return s + def recent_per_sample(self, graph_name: str, sample_ids): + """Recent in-memory values per sample: ``{sample_id: [values]}``. + + O(batch). Samples not written in this process are absent rather than + empty-listed; callers treat both as "not enough history yet". + """ + tail = self._recent_tail.get(graph_name) or {} + out = {} + for s in sample_ids: + q = tail.get(str(s)) + if q: + out[s] = list(q) + return out + def _maybe_autoflush(self) -> None: if (len(self._stage_signals) + len(self._stage_sample) + len(self._stage_instance)) >= _STAGE_FLUSH_THRESHOLD: @@ -823,6 +853,13 @@ def _stage_sample_row(self, graph_name, exp_hash, sample_id, step, value): ) self._qps_version[graph_name] += 1 # invalidate this signal's cached reads self._loss_shape_dirty_samples[(graph_name, exp_hash)].add(str(sample_id)) + if self._tail_len: + _tail = self._recent_tail[graph_name] + _sid = str(sample_id) + _q = _tail.get(_sid) + if _q is None: + _q = _tail[_sid] = deque(maxlen=self._tail_len) + _q.append(float(value)) self._maybe_autoflush() def _invalidate_qps_cache(self) -> None: diff --git a/weightslab/data/dataframe_manager.py b/weightslab/data/dataframe_manager.py index 76b9a8f7..3bfcc31f 100644 --- a/weightslab/data/dataframe_manager.py +++ b/weightslab/data/dataframe_manager.py @@ -92,7 +92,15 @@ def __init__(self, flush_interval: float = 3.0, flush_max_rows: int = 100, enabl self._store: H5DataFrameStore | None = None self._array_store: H5ArrayStore | None = None self._origin_revisions: Dict[str, int] = {} + # Bumped ONLY when the `discarded` column changes. The deny-list cache + # keys on this instead of _origin_revisions, which training bumps every + # step (per-sample signal writes) and which therefore never lets a cache + # hit -- turning a per-batch len() into a 3.96M-row scan. + self._discard_revisions: Dict[str, int] = {} self._pending: set[int] = set() + # Parallel dirty set for the view. _pending is drained by the H5 flush, + # so the view cannot share it without one consumer starving the other. + self._view_pending: set = set() self._force_flush = False self._flush_interval = flush_interval self._flush_max_rows = flush_max_rows @@ -301,9 +309,66 @@ def _expand_dataframe_with_annotations(self, df: pd.DataFrame) -> pd.DataFrame: if not isinstance(work.index, pd.MultiIndex) and work.index.name != SID: work = work.copy() work.index.name = SID + fast = self._expand_fast_no_instances(work) + if fast is not None: + return fast + records = work.reset_index().to_dict("records") return self._expand_records_to_multi_index(records) + def _expand_fast_no_instances(self, work: pd.DataFrame): + """Vectorized expansion for frames with no per-instance targets. + + Returns the (sample_id, annotation_id=0) frame, or None to signal the + caller must use the record-by-record path. + """ + SID = SampleStats.Ex.SAMPLE_ID.value + ANNOT = SampleStats.Ex.INSTANCE_ID.value + TARGET = SampleStats.Ex.TARGET.value + try: + flat = work.reset_index() + if SID not in flat.columns: + return None + if TARGET in flat.columns: + tgt = flat[TARGET] + # Non-object dtype cannot hold a list => every target is scalar. + if tgt.dtype == object: + for v in tgt.to_numpy(): + if isinstance(v, (list, tuple, np.ndarray)) and len(v) > 0: + return None + # _normalize_sample_id always returns str(). astype(str) reproduces + # that for numeric dtypes only -- bytes would render as "b'x'". + sid_ser = flat[SID] + if pd.api.types.is_integer_dtype(sid_ser) or pd.api.types.is_float_dtype(sid_ser): + sids = sid_ser.astype(str).tolist() + else: + sids = [self._normalize_sample_id(v) for v in sid_ser.to_numpy()] + + out = flat.drop(columns=[c for c in (SID, ANNOT) if c in flat.columns]) + # The record path goes through python lists, so extension dtypes + # (string[pyarrow], categorical) come back as object. Match it. + for c in out.columns: + if isinstance(out[c].dtype, pd.api.extensions.ExtensionDtype): + out[c] = out[c].astype(object) + out.index = pd.MultiIndex.from_arrays( + [sids, np.zeros(len(out), dtype=np.int64)], names=[SID, ANNOT]) + return out + except Exception: + return None + + def _normalize_sample_id_index(self, values) -> "pd.Index": + """Vectorized _normalize_sample_id over an Index (~2x; 0.8s -> 0.5s at 2M). + + _normalize_sample_id is str() after unwrapping numpy scalars/bytes, so + astype(str) is exact for numeric dtypes; anything else keeps the loop. + """ + try: + if pd.api.types.is_integer_dtype(values) or pd.api.types.is_float_dtype(values): + return pd.Index(values.astype(str)) + except Exception: + pass + return pd.Index([self._normalize_sample_id(v) for v in values]) + def _normalize_sample_id(self, sample_id: Any) -> Any: """Normalize incoming sample IDs while preserving numeric IDs when possible.""" try: @@ -320,6 +385,22 @@ def _normalize_sample_id(self, sample_id: Any) -> Any: return str(sample_id) + def _level0_index(self): + """Level-0 (sample_id) values of the ledger index, cached. + + Keyed on the index object's identity: pandas Index is immutable, so a + reindex or rebuild yields a new object and invalidates this. Reusing the + object also reuses its hash engine, which is what makes a membership + probe O(1) instead of O(rows). + """ + idx = self._df.index + key = id(idx) + if getattr(self, "_lvl0_key", None) != key: + self._lvl0_key = key + self._lvl0 = (idx.get_level_values(0) + if isinstance(idx, pd.MultiIndex) else idx) + return self._lvl0 + def _coerce_sample_id_for_index(self, sample_id: Any) -> Any: """Coerce sample_id to match current dataframe index representation. @@ -332,8 +413,10 @@ def _coerce_sample_id_for_index(self, sample_id: Any) -> Any: # Check if multi-index if isinstance(self._df.index, pd.MultiIndex): - # Get level 0 (sample_id level) values - level_0_values = self._df.index.get_level_values(0) + # Cached: get_level_values(0) built a new Index over every row on each + # call, and a new object means a new hash engine, so this probe was + # O(rows) per sample. + level_0_values = self._level0_index() if sid in level_0_values: return sid sid_str = str(sid) @@ -355,6 +438,11 @@ def set_array_store(self, array_store: H5ArrayStore): if self._enable_h5_persistence: self._array_store = array_store + def _bump_discard_revisions(self, origins: Sequence[Any]) -> None: + for origin in origins or []: + key = str(origin) + self._discard_revisions[key] = self._discard_revisions.get(key, 0) + 1 + def _bump_origin_revisions(self, origins: Sequence[Any]) -> None: for origin in origins: if origin is None or pd.isna(origin): @@ -720,7 +808,7 @@ def upsert_df(self, df_local: List | pd.DataFrame, origin: str = None, force_flu # Normalize sample_id values in multi-index if isinstance(df_norm.index, pd.MultiIndex) and df_norm.index.nlevels >= 1: - level_0_normalized = pd.Index([self._normalize_sample_id(v) for v in df_norm.index.get_level_values(0)]) + level_0_normalized = self._normalize_sample_id_index(df_norm.index.get_level_values(0)) try: if df_norm.index.nlevels == 2: df_norm.index = pd.MultiIndex.from_arrays( @@ -784,7 +872,29 @@ def upsert_df(self, df_local: List | pd.DataFrame, origin: str = None, force_flu for col in all_cols: if col in self._df.columns and isinstance(self._df[col].dtype, pd.CategoricalDtype): self._df[col] = self._df[col].astype(object) - self._df.loc[existing_idx, all_cols] = df_norm.loc[existing_idx, all_cols] + # Label-aligned 2D .loc realigns the entire frame (22s at 4M rows + # when adding columns to every row). Resolve row positions once, + # then write each column positionally. Falls back if the index + # has duplicates/misses, where get_indexer returns -1. + _pos = self._df.index.get_indexer(existing_idx) + if len(_pos) and (_pos >= 0).all(): + for _c in all_cols: + _ci = self._df.columns.get_loc(_c) + _vals = df_norm.loc[existing_idx, _c].to_numpy() + # An object array (None for "no value yet") written into + # a float column upcasts the WHOLE column to object and + # it never returns -- an 8x penalty on every later sort. + # Coerce to the target dtype so None becomes NaN instead. + try: + _tgt = self._df[_c].dtype + if (_vals.dtype == object + and getattr(_tgt, "kind", "") in "fiu"): + _vals = pd.to_numeric(_vals, errors="coerce") + except Exception: + pass + self._df.iloc[_pos, _ci] = _vals + else: + self._df.loc[existing_idx, all_cols] = df_norm.loc[existing_idx, all_cols] # Append rows that do not exist yet. Use a boolean mask (not # .loc[difference]) so a duplicate key in df_norm can't be @@ -805,7 +915,9 @@ def upsert_df(self, df_local: List | pd.DataFrame, origin: str = None, force_flu self._auto_register_categorical_tags(df_norm) # Optimize memory by converting repetitive columns to categorical - self._df = self._optimize_dataframe_memory(self._df) + # Only columns just written can have changed dtype-wise; a full-frame + # nunique() over every object column was ~14s of a 670s startup. + self._df = self._optimize_dataframe_memory(self._df, columns=set(df_norm.columns)) # Mark dirty for flush (handle multi-index) if isinstance(df_norm.index, pd.MultiIndex): @@ -814,6 +926,8 @@ def upsert_df(self, df_local: List | pd.DataFrame, origin: str = None, force_flu sample_ids = df_norm.index.tolist() self.mark_dirty_batch(sample_ids, force_flush=force_flush) self._bump_origin_revisions(affected_origins) + if SampleStats.Ex.DISCARDED.value in set(df_norm.columns): + self._bump_discard_revisions(affected_origins) def mark_dirty(self, sample_id: int): """Mark sample as dirty for H5 flush. @@ -823,6 +937,7 @@ def mark_dirty(self, sample_id: int): with self._lock: normalized_id = self._coerce_sample_id_for_index(sample_id) self._pending.add(normalized_id) + self._view_pending.add(normalized_id) def drop_column(self, column: str): with self._lock: @@ -833,6 +948,7 @@ def drop_column(self, column: str): def mark_dirty_batch(self, sample_ids: List[int], force_flush: bool = False): with self._lock: self._pending.update(set(sample_ids)) + self._view_pending.update(set(sample_ids)) if force_flush: self._force_flush = True @@ -1337,9 +1453,55 @@ def update_values(self, origin: str, sample_id: int, updates: Dict[str, Any], an self._df = pd.concat([self._df, df_local]) self._bump_origin_revisions([origin]) - def get_origin_revision(self, origin: str) -> int: + def take_view_dirty(self, limit: int | None = None): + """Drain and return the sample_ids changed since the last view sync. + + Returns None when the backlog exceeds *limit*, meaning a differential + update would cost more than a rebuild — the caller should fall back. + """ + with self._lock: + if limit is not None and len(self._view_pending) > limit: + # Do NOT clear here. The caller is expected to rebuild, but that + # rebuild can bail (contended update lock) -- and these ids would + # then be lost with nothing to re-mark them. clear_view_dirty() + # is called once the rebuilt view is actually swapped in. + return None + out = list(self._view_pending) + self._view_pending.clear() + return out + + def clear_view_dirty(self): + """Drop the view-dirty backlog: a full rebuild has made the view current.""" + with self._lock: + self._view_pending.clear() + + def get_source_rows(self, sample_ids, columns=None): + """Rows for *sample_ids* straight from the source frame. O(len(ids)).""" with self._lock: - return int(self._origin_revisions.get(str(origin), 0)) + if self._df.empty or not len(sample_ids): + return None + idx = self._df.index + keys = idx.get_level_values(0) if isinstance(idx, pd.MultiIndex) else idx + want = set(str(s) for s in sample_ids) + mask = keys.astype(str).isin(want) + sub = self._df.loc[mask] + return sub[columns] if columns else sub + + def get_origin_revision(self, origin: str) -> int: + # No lock: a dict read is atomic under the GIL, and this is polled from + # __len__ on every batch. Taking self._lock here put the training thread + # in contention with the flush thread on the hot path. + return int(self._origin_revisions.get(str(origin), 0)) + + def get_discard_revision(self, origin: str) -> int: + """Revision of the `discarded` column for *origin*. + + Changes only on discard/restore, so a consumer that depends purely on + deny-list state can cache against it across training steps. Read without + the lock: called per batch, and a stale-by-one read is harmless (the next + batch picks the change up), whereas lock contention here is not. + """ + return int(self._discard_revisions.get(str(origin), 0)) def update_by_groups_bulk(self, origin: str, group_ids: List[Any], updates_list: List[Dict[str, Any]]): """Broadcast updates to multiple groups in one pass.""" @@ -1535,22 +1697,42 @@ def get_sample_column_values(self, sample_ids: List[Any], column: str) -> Dict[A coerced_ids = [self._coerce_sample_id_for_index(sid) for sid in sample_ids] + idx = self._df.index try: - if isinstance(self._df.index, pd.MultiIndex) and self._df.index.nlevels >= 2: - sample_level = self._df.index.get_level_values(0) - anno_level = self._df.index.get_level_values(1) - mask = sample_level.isin(coerced_ids) & (anno_level == 0) - slice_df = self._df[mask] - sids = slice_df.index.get_level_values(0) + # _positional NB_SEEN lookup: the wanted rows are exactly + # (sample_id, 0), so resolve their POSITIONS instead of scanning. + # The masked form below materialised both index levels and copied + # a boolean-masked frame over every row -- ~1.2s/step at 3.96M + # just to read a batch of integers. + if idx.has_duplicates: + raise ValueError("non-unique index; use the scan path") + if isinstance(idx, pd.MultiIndex) and idx.nlevels >= 2: + pos = idx.get_indexer([(cid, 0) for cid in coerced_ids]) else: - mask = self._df.index.isin(coerced_ids) - slice_df = self._df[mask] - sids = slice_df.index - - for sid, val in zip(sids, slice_df[column]): - values[self._normalize_sample_id(sid)] = val + pos = idx.get_indexer(list(coerced_ids)) + col = self._df[column].to_numpy() + for cid, p in zip(coerced_ids, pos): + if p >= 0: + values[self._normalize_sample_id(cid)] = col[p] except Exception: - pass + # Fallback: original scan, for a non-unique index or any dtype + # mismatch get_indexer will not tolerate. + try: + if isinstance(idx, pd.MultiIndex) and idx.nlevels >= 2: + sample_level = idx.get_level_values(0) + anno_level = idx.get_level_values(1) + mask = sample_level.isin(coerced_ids) & (anno_level == 0) + slice_df = self._df[mask] + sids = slice_df.index.get_level_values(0) + else: + mask = idx.isin(coerced_ids) + slice_df = self._df[mask] + sids = slice_df.index + + for sid, val in zip(sids, slice_df[column]): + values[self._normalize_sample_id(sid)] = val + except Exception: + pass return values @@ -1907,7 +2089,9 @@ def _apply_buffer_records(self, records: List[Dict[str, Any]]): self._apply_updates_frame_locked(instance_df, broadcast=False) self._apply_updates_frame_locked(sample_df, broadcast=True) # Keep newly-added signal columns float32 and empty object cells as None. - self._df = self._optimize_dataframe_memory(self._df) + self._df = self._optimize_dataframe_memory( + self._df, + columns=set(sample_df.columns) | set(instance_df.columns)) # Mark all as pending for h5 flush (outside lock) self.mark_dirty_batch(sample_ids) @@ -1945,13 +2129,22 @@ def _apply_buffer_records_nonblocking(self, records: List[Dict[str, Any]]): applied_index = written_s.append(written_i) if len(written_i) else written_s update_cols = sample_df.columns.union(instance_df.columns) # Keep newly-added signal columns float32 and empty object cells as None. - _df = self._optimize_dataframe_memory(self._df) + _df = self._optimize_dataframe_memory(self._df, columns=set(update_cols)) self._df = _df finally: self._lock.release() # Det→seg conversion / array normalization over all written rows. - if applied_index is not None and len(applied_index) > 0: + # This prepares cells for the H5 write, so it is only worth doing for + # columns that write will actually include. When predictions/targets are + # excluded from the save list (WEIGHTSLAB_SAVE_PREDICTIONS_IN_H5=0) the + # pass would otherwise call get_mask per row -- which reads the source + # image to size the mask -- for arrays that are never persisted. + _savable = set(_filter_columns_by_patterns( + list(update_cols), SAMPLES_STATS_TO_SAVE_TO_H5)) + _norm_cols = [c for c in update_cols + if c in self._array_columns and c in _savable] + if applied_index is not None and len(applied_index) > 0 and _norm_cols: if applied_index.has_duplicates: applied_index = applied_index[~applied_index.duplicated()] normalized_rows = self._df.loc[applied_index].apply( @@ -2008,6 +2201,33 @@ def _flush_to_h5_if_needed(self, force: bool = False, blocking: bool = False): # Everything below happens WITHOUT locks - fully async self._flush_snapshot_to_h5(data_snapshot, work) + def _rows_with_array_cells(self, data_snapshot: pd.DataFrame): + """Index labels of rows that may hold an array-valued cell. + + Column-wise scan replacing a full iterrows() pass: only object-dtype + columns can hold an ndarray/list/tuple/ArrayH5Proxy, and rows with none + of those are no-ops for the caller. + """ + cols = [c for c in self._array_columns if c in data_snapshot.columns] + if not cols: + return [] + hits = None + for col in cols: + ser = data_snapshot[col] + if ser.dtype != object: + continue # a numeric column cannot hold an array object + vals = ser.to_numpy() + mask = np.fromiter( + (isinstance(v, (np.ndarray, list, tuple, ArrayH5Proxy)) for v in vals), + dtype=bool, count=len(vals)) + if not mask.any(): + continue + found = ser.index[mask] + hits = found if hits is None else hits.union(found) + if hits is None: + return [] + return list(hits) + def _flush_snapshot_to_h5(self, data_snapshot: pd.DataFrame, work: List[int]): """Flush data snapshot to H5 - runs completely outside locks. @@ -2030,7 +2250,10 @@ def _flush_snapshot_to_h5(self, data_snapshot: pd.DataFrame, work: List[int]): arrays_to_store: Dict[str, Dict[str, np.ndarray]] = {} rowkey_to_index: Dict[str, Any] = {} - for idx, row in data_snapshot.iterrows(): + for idx in self._rows_with_array_cells(data_snapshot): + row = data_snapshot.loc[idx] + if isinstance(row, pd.DataFrame): # duplicate label guard + row = row.iloc[0] if is_multi: sample_id, annot = idx[0], int(idx[1]) else: @@ -2103,7 +2326,7 @@ def _flush_snapshot_to_h5(self, data_snapshot: pd.DataFrame, work: List[int]): except Exception as e: logger.error(f"[LedgeredDataFrameManager] Error flushing to H5: {e}") - def _optimize_dataframe_memory(self, df: pd.DataFrame, categorical_tags: Dict[str, List[str]] | None = None) -> pd.DataFrame: + def _optimize_dataframe_memory(self, df: pd.DataFrame, categorical_tags: Dict[str, List[str]] | None = None, columns=None) -> pd.DataFrame: """Optimize dataframe memory by converting repetitive string columns to categorical. Categorical dtype compresses repeated values: instead of storing each string, @@ -2130,12 +2353,39 @@ def _optimize_dataframe_memory(self, df: pd.DataFrame, categorical_tags: Dict[st if categorical_tags is None: categorical_tags = self._categorical_tags + # Honour `columns`: only the columns this flush actually wrote can have + # gained a float64 signal value or a fresh NaN, so scanning the rest is + # O(rows) of pure waste on every flush. + _scan_cols = (list(df.columns) if columns is None + else [c for c in df.columns if c in columns]) + + # === 0) Repair signal columns that were upcast to object === + # A single None written into a float column converts it permanently, and + # object columns sort ~8x slower. signals//* are numeric by definition, + # so any object one is damage rather than intent. Runs BEFORE the + # NaN->None pass below, which only touches object columns and would + # otherwise keep them that way. + for col in _scan_cols: + if not str(col).startswith("signals") or df[col].dtype != object: + continue + try: + coerced = pd.to_numeric(df[col], errors="coerce") + # Only if nothing was lost: a genuine non-numeric value means the + # column is not what we think it is, so leave it alone. + if coerced.notna().sum() == df[col].notna().sum(): + df[col] = coerced.astype(np.float32) + logger.debug( + "[LedgeredDataFrameManager] restored '%s' object -> float32", col) + except Exception as exc: + logger.debug("[LedgeredDataFrameManager] dtype repair skipped for '%s': %s", + col, exc) + # === 1) Downcast float64 signal columns to float32 === # Per-sample / per-instance signal scalars (loss & metric values) don't need # float64 precision, so this halves the cost of every ``signals//*`` column # with no practical loss for monitoring. Numeric dtype is preserved (NaNs # stay NaN). Done before categorical conversion below. - for col in df.columns: + for col in _scan_cols: if str(col).startswith("signals") and df[col].dtype == np.float64: try: df[col] = df[col].astype(np.float32) @@ -2159,7 +2409,12 @@ def _optimize_dataframe_memory(self, df: pd.DataFrame, categorical_tags: Dict[st # into a Categorical, its missing cells are stuck as NaN (categorical code # -1) and can no longer be set to a plain None — so we clean object cells to # None here first, while they are still plain object dtype. - for col in df.columns: + for col in _scan_cols: + # The docstring above says numeric/bool/categorical are skipped, but the + # loop had no dtype check: at registration every signals//* column is all + # NaN, so this did a full label-aligned write per float column. + if df[col].dtype != object: + continue na_mask = df[col].isna() if na_mask.any(): df.loc[na_mask, col] = None @@ -2171,6 +2426,17 @@ def _optimize_dataframe_memory(self, df: pd.DataFrame, categorical_tags: Dict[st SampleStats.Ex.TASK_TYPE.value, # Task type (e.g. 'classification', 'segmentation') ] + # nunique() over a 4M-row MultiIndex costs ~8s AND holds the GIL, stalling + # the training thread inside forward/backward. Only an object-dtype + # candidate reads it, so compute it on first use rather than every flush. + _n_rows_cache = [] + + def _n_rows(): + if not _n_rows_cache: + _n_rows_cache.append( + df.index.get_level_values(0).nunique() + if isinstance(df.index, pd.MultiIndex) else len(df)) + return _n_rows_cache[0] for col in categorical_candidates: if col not in df.columns: continue @@ -2182,15 +2448,14 @@ def _optimize_dataframe_memory(self, df: pd.DataFrame, categorical_tags: Dict[st # Only convert to categorical if: # 1. Column contains strings (object dtype) # 2. Number of unique values < 50% of total rows (good compression ratio) + if columns is not None and col not in columns: + continue if df[col].dtype == 'object': n_unique = df[col].nunique() # Use unique sample count, not row count — with MultiIndex each # sample has multiple annotation rows which would inflate n_rows # and make the ratio appear better than it really is. - if isinstance(df.index, pd.MultiIndex): - n_rows = df.index.get_level_values(0).nunique() - else: - n_rows = len(df) + n_rows = _n_rows() compression_ratio = n_unique / n_rows if n_rows > 0 else 1.0 if compression_ratio < 0.5 and n_unique > 1: # Worth compressing if < 50% unique diff --git a/weightslab/data/h5_array_store.py b/weightslab/data/h5_array_store.py index 02c3e7e3..405cf02b 100644 --- a/weightslab/data/h5_array_store.py +++ b/weightslab/data/h5_array_store.py @@ -514,6 +514,52 @@ def save_array( finally: self._rw_lock.release_write() + def _try_inplace_batch(self, prepared): + """Overwrite existing datasets in place; None means "cannot, fall back". + + Two passes under the write lock: check every destination exists with a + matching shape and dtype, and only then write. A partial in-place write + followed by a fallback would corrupt silently, so nothing is written + until the whole batch is known to fit. + """ + if not self._path.exists(): + return None + with self._local_lock: + self._rw_lock.acquire_write() + try: + with _InterProcessFileLock(self._lock_path, timeout=self._lock_timeout, + poll_interval=self._poll_interval): + with h5py.File(str(self._path), 'a') as f: + for group_name, key_data in prepared.items(): + grp = f.get(group_name) + if grp is None: + return None + for key_name, (array, _meta) in key_data.items(): + kg = grp.get(key_name) + if kg is None or 'data' not in kg: + return None + dset = kg['data'] + if dset.shape != array.shape or dset.dtype != array.dtype: + return None + for group_name, key_data in prepared.items(): + for key_name, (array, metadata) in key_data.items(): + kg = f[group_name][key_name] + kg['data'][...] = array + for mk, mv in metadata.items(): + kg.attrs[mk] = mv + return { + group_name: { + key_name: self._build_path_reference(group_name, key_name) + for key_name in key_data + } + for group_name, key_data in prepared.items() + } + except Exception as exc: + logger.debug(f"[H5ArrayStore] in-place batch fell back: {exc}") + return None + finally: + self._rw_lock.release_write() + def save_arrays_batch( self, arrays_dict: Dict[int, Dict[str, np.ndarray]], @@ -562,6 +608,15 @@ def save_arrays_batch( if not prepared: return {} + # O(change): if every array already exists with the same shape and dtype, + # overwrite the values in place. That is not a structural change, so it + # needs neither the temp file nor the full-file backup (9.5GB per flush + # at current ledger size). Returns None if anything would need creating + # or resizing, and the original two-phase path below runs unchanged. + inplace_refs = self._try_inplace_batch(prepared) + if inplace_refs is not None: + return inplace_refs + tmp_path = self._path.with_suffix(f".h5.writing_{uuid.uuid4().hex[:8]}") try: with h5py.File(str(tmp_path), 'w') as f_tmp: diff --git a/weightslab/data/h5_dataframe_store.py b/weightslab/data/h5_dataframe_store.py index 476b5655..431425d1 100644 --- a/weightslab/data/h5_dataframe_store.py +++ b/weightslab/data/h5_dataframe_store.py @@ -680,6 +680,120 @@ def load_all(self, origins: Iterable[str] = None, columns: Optional[Iterable[str return pd.DataFrame() raise + def ensure_index(self, origin: str, columns=("sample_id",)) -> bool: + """Build the on-disk column index deliberately (checkpoint / first query). + + Kept OFF the flush path: rebuilding it per upsert costs 92.7s at 4M rows + versus 6.9s without, for an index no hot-path read uses. + """ + key = self._key(origin) + try: + with self._local_lock: + with _InterProcessFileLock(self._lock_path, timeout=self._lock_timeout, + poll_interval=self._poll_interval): + with pd.HDFStore(str(self._path), mode="a") as store: + if key not in store: + return False + store.create_table_index(key, columns=list(columns), + optlevel=6, kind="medium") + return True + except Exception as exc: + logger.warning(f"[H5DataFrameStore] ensure_index({origin}) failed: {exc}") + return False + + # --- O(change) in-place update ------------------------------------------ + _POSMAP_CACHE: dict = {} + + def _posmap(self, store, key, force=False): + """(sample_id, annotation_id) -> row position. Rows are registered once + and never deleted, so positions are stable; built from ONE column (2.9s + at 4M rows) rather than reading the table.""" + ck = (str(self._path), key) + if not force and ck in self._POSMAP_CACHE: + return self._POSMAP_CACHE[ck] + try: + sids = store.select_column(key, "sample_id").values + try: + aids = store.select_column(key, "annotation_id").values + except Exception: + aids = np.zeros(len(sids), dtype="i8") + # Hoist the normalisation out of the insert loop: the per-row + # decode/str/int calls interleaved with dict inserts cost 3.6s at 4M + # rows, against ~1.7s for dict(zip(...)) over pre-normalised lists. + # One pass that is also the type check: str hits the identity branch, + # so a mixed-dtype column stays correct without a second full scan. + sid_list = [s if type(s) is str + else (s.decode() if isinstance(s, bytes) else str(s)) + for s in sids.tolist()] + m = dict(zip(zip(sid_list, aids.tolist()), range(len(sid_list)))) + self._POSMAP_CACHE[ck] = m + return m + except Exception as exc: + logger.debug(f"[H5DataFrameStore] posmap build failed: {exc}") + return None + + def _invalidate_posmap(self, key): + self._POSMAP_CACHE.pop((str(self._path), key), None) + + def _try_inplace(self, store, key, df_norm) -> bool: + """Overwrite existing rows' values in place. True if fully applied.""" + try: + import tables as _tables + except Exception: + return False + try: + node = store._handle.get_node(key) + tbl = getattr(node, "table", node) + if not isinstance(tbl, _tables.Table): + return False + # An indexed column cannot be modified in place (PyTables raises). + if any(tbl.cols._f_col(c).is_indexed for c in tbl.colnames): + return False + + cols = [c for c in df_norm.columns if c in tbl.colnames] + if len(cols) != len(df_norm.columns): + return False # new column => schema change + + pos = self._posmap(store, key) + if not pos: + return False + + idx = df_norm.index + if isinstance(idx, pd.MultiIndex): + pairs = [(str(a), int(b)) for a, b in zip(idx.get_level_values(0), + idx.get_level_values(1))] + else: + pairs = [(str(a), 0) for a in idx] + + coords = np.empty(len(pairs), dtype=np.int64) + for i, p in enumerate(pairs): + j = pos.get(p) + if j is None: + return False # unknown row => not an update + coords[i] = j + + order = np.argsort(coords) # PyTables wants ascending coords + coords_sorted = coords[order] + rec = tbl.read_coordinates(coords_sorted) + for c in cols: + vals = df_norm[c].to_numpy()[order] + tgt = rec[c].dtype + if tgt.kind == "S": + vals = np.array([("" if v is None else str(v)).encode()[:tgt.itemsize] + for v in vals], dtype=tgt) + else: + try: + vals = vals.astype(tgt, copy=False) + except Exception: + return False # dtype mismatch => fall back + rec[c] = vals + tbl.modify_coordinates(coords_sorted, rec) + tbl.flush() + return True + except Exception as exc: + logger.debug(f"[H5DataFrameStore] in-place update fell back: {exc}") + return False + def upsert(self, origin: str, df: pd.DataFrame) -> int: """Atomic upsert with corruption prevention via backup and checksum verification.""" df_norm = self._normalize_for_write(df) @@ -689,13 +803,26 @@ def upsert(self, origin: str, df: pd.DataFrame) -> int: key = self._key(origin) self._ensure_parent() - # Create backup BEFORE any writes - backup_path = self._create_backup() + # The backup is a full copy of the file (696MB at 4M rows). Only the + # read-merge-rewrite path below can destroy the table -- the in-place + # path just overwrites values in already-allocated rows -- so the copy + # is deferred until we know we are taking the destructive route. + backup_path = None with self._local_lock: with _InterProcessFileLock(self._lock_path, timeout=self._lock_timeout, poll_interval=self._poll_interval): try: with pd.HDFStore(str(self._path), mode="a") as store: + # O(change): if every row already exists and the schema is + # unchanged, overwrite values in place (0.1ms vs 42s). + if key in store and self._try_inplace(store, key, df_norm): + return len(df_norm) + + # Nothing has been written yet in this call; flush so the + # copy below captures a consistent on-disk file. + store.flush() + backup_path = self._create_backup() + existing = pd.DataFrame() # Try to load existing data. A ValueError can surface from a @@ -801,7 +928,8 @@ def upsert(self, origin: str, df: pd.DataFrame) -> int: store.remove(key) # Write new data - store.append(key, existing, format="table", data_columns=True) + store.append(key, existing, format="table", data_columns=True, index=False) + self._invalidate_posmap(key) # Force flush to disk store.flush() @@ -880,7 +1008,7 @@ def delete_column(self, column_name: str, origins: Optional[Iterable[str]] = Non # Remove old key and write updated dataframe store.remove(key) if not df.empty: - store.append(key, df, format="table", data_columns=True) + store.append(key, df, format="table", data_columns=True, index=False) modified_count += 1 logger.debug(f"[H5DataFrameStore] Deleted column {column_name} from {origin}") diff --git a/weightslab/examples/PyTorch/wl-video-generation/utils/data.py b/weightslab/examples/PyTorch/wl-video-generation/utils/data.py index ddb4d3ab..42112fba 100644 --- a/weightslab/examples/PyTorch/wl-video-generation/utils/data.py +++ b/weightslab/examples/PyTorch/wl-video-generation/utils/data.py @@ -31,7 +31,6 @@ """ import csv import logging -import os import subprocess import shutil import wave diff --git a/weightslab/src.py b/weightslab/src.py index 498b329c..c97c36d7 100644 --- a/weightslab/src.py +++ b/weightslab/src.py @@ -494,6 +494,16 @@ def history(self, signal_name): out = {s: [] for s in self.sample_ids} if self.logger is None: return out + # Read the bounded in-memory tail: O(batch). The full-history query this + # replaced scanned the entire per_sample table on every call (140ms at + # 20M rows, once per step, growing without bound). + if not os.environ.get("WL_HISTORY_FROM_DB"): + recent = self.logger.recent_per_sample(signal_name, self.sample_ids) + for s in self.sample_ids: + vals = recent.get(s) + if vals: + out[s] = list(vals) + 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) @@ -1023,9 +1033,23 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw): # 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. + # inputs= must be populated here too: on the + # subscribe_to path the context was built without it, + # so b.inputs was {} and any signal declaring + # inputs=[...] raised KeyError on every call. The + # subscribed signal IS the declared input here, and + # its per-sample values are already in val_vec. + _decl = meta.get('inputs') or [] + _sub = meta.get('subscribe_to') + _vals = [float(v) for v in val_vec] + _bin = {} + for _d in _decl: + if _sub is None or _d == _sub or _d == reg_name: + _bin[_d] = _vals bctx = BatchSignalContext( sample_ids=[int(u) for u in ids_np], - subscribed_values=[float(v) for v in val_vec], + subscribed_values=_vals, + inputs=_bin, logger=_lg, dataframe=df_proxy, origin=kwargs.get('origin', 'train'), @@ -4932,6 +4956,16 @@ def resolve_signal_classifier(signal_name): return _GLOBAL_CLASSIFIER or classify_loss_shape +_SHAPE_LABELS: dict = {} + + +def _label_counts(cache): + out = {} + for lab in cache.values(): + out[lab] = out.get(lab, 0) + 1 + return out + + def write_signal_shapes(signal_name, tag_name=None, classifier=None, exp_hash=None, sample_ids=None): """Reusable engine: classify each sample's own trajectory of *signal_name* into a categorical tag and return the ``{label: count}`` distribution. @@ -4950,17 +4984,29 @@ def write_signal_shapes(signal_name, tag_name=None, classifier=None, exp_hash=No clf = classifier or resolve_signal_classifier(signal_name) if tag_name is None: tag_name = signal_name + "_shape" if signal_name.endswith('_loss') else signal_name + "_loss_shape" + + # Labels for samples this pass does not reclassify are carried here, so an + # incremental call still returns a distribution over the WHOLE dataset, and + # a sample whose label is unchanged is not re-written to the ledger. + cache = _SHAPE_LABELS.setdefault(signal_name, {}) + if sample_ids is not None and not list(sample_ids): + return _label_counts(cache) + series = {} for sid, step, val, _ in query_signal_history(signal_name, exp_hash=exp_hash, sample_ids=sample_ids): 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) + if label is None: + continue + if cache.get(sid) == label: + continue # unchanged -> no ledger write needed + cache[sid] = label + 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()} + return _label_counts(cache) def write_loss_shapes(loss_signal="loss_sample", classifier=None): diff --git a/weightslab/trainer/services/data_service.py b/weightslab/trainer/services/data_service.py index dd065d33..131ae6d6 100755 --- a/weightslab/trainer/services/data_service.py +++ b/weightslab/trainer/services/data_service.py @@ -44,8 +44,8 @@ from weightslab.data import media_store from weightslab.trainer.trainer_tools import execute_df_operation, generate_overview, encode_image_to_raw_bytes from weightslab.data.data_utils import load_raw_image_array - # Image encoding / mask compression / proto helpers (extracted) + from weightslab.trainer.services.data_image_utils import ( rle_encode_mask, create_data_stat, @@ -494,6 +494,23 @@ def rewrite_boolean_keywords_to_bitwise(code: str) -> str: return code +def _histogram_category_cap() -> int: + """Max distinct bars a categorical histogram returns (WL_HIST_CATEGORY_CAP). + + Beyond this the response is neither renderable nor informative -- the + remainder is folded into a single "(other)" bar. + """ + try: + return max(1, int(os.environ.get("WL_HIST_CATEGORY_CAP", "200"))) + except Exception: + return 200 + + +def _fast_view_enabled() -> bool: + """Differential view refresh. On by default; set WL_FAST_VIEW=0 to opt out.""" + return os.environ.get("WL_FAST_VIEW", "1") not in ("0", "false", "False") + + class DataService: """ @@ -1109,6 +1126,23 @@ def _pull_into_all_data_view_df(self): # merge + proxy conversion) a second time over the whole dataset every refresh. df = self._df_manager.get_collapse_annotations_to_samples_df(df) + # The collapse yields object dtype for signals//* columns. They are + # numeric by definition, and an object column sorts ~8x slower + # (6.17s vs 0.76s at 3.96M) while also slowing histogram binning, + # groupby and every differential write. Coerce them back here, at the + # single point the view is materialised. + for _c in df.columns: + if not str(_c).startswith("signals") or df[_c].dtype != object: + continue + try: + _num = pd.to_numeric(df[_c], errors="coerce") + # Only when nothing is lost: a genuinely non-numeric value + # means the column is not what we assume, so leave it be. + if _num.notna().sum() == df[_c].notna().sum(): + df[_c] = _num.astype("float32") + except Exception as _exc: + logger.debug("[DataService] dtype restore skipped for %r: %s", _c, _exc) + # Ensure sample_id is a column if it was the index df = safe_reset_index(df) @@ -1131,7 +1165,9 @@ def _pull_into_all_data_view_df(self): return df except Exception as e: - logger.debug(f"[DataService] Error pulling data view: {e}") + # Was debug: a swallowed failure here silently freezes the view at + # the previous snapshot, which looks exactly like "no new data". + logger.error("[DataService] Error pulling data view: %s", e, exc_info=True) # Use getattr to safely check for attribute during __init__ current_df = getattr(self, "_all_datasets_df", None) return current_df if current_df is not None else pd.DataFrame() @@ -1448,8 +1484,9 @@ def _compute_custom_signals(self): except Exception as e: logger.error(f"[DataService] Failed to compute signals for loader '{loader_name}': {e}") - # Force view update - self._slowUpdateInternals(force=True) + # Refresh signal values; differential unless the schema actually changed. + if not self._fastUpdateInternals(): + self._slowUpdateInternals(force=True) def _process_sample_row(self, args): """Process a single dataframe row to create a DataRecord.""" @@ -1477,6 +1514,8 @@ def _process_sample_row(self, args): skip_prediction_for_request = metadata_only_request # ====== Step 2: Load dataset lazily (avoid unnecessary IO for metadata-only) ====== + # Views the client explicitly asked for; empty => send everything. + _wanted_stats = set(getattr(request, "stats_to_retrieve", None) or []) needs_dataset = bool(request.include_raw_data) or (not skip_label_for_request) dataset = self._get_dataset(origin) if needs_dataset else None @@ -2150,14 +2189,91 @@ def _json_default(o): target_height=target_height, ) - data_stats.append( - create_data_stat( - name='raw_data', - stat_type='bytes', - thumbnail=raw_data_bytes, - shape=raw_shape, + # An explicit stats_to_retrieve means the client knows which + # views it will draw; raw_data duplicates view rank 0, so + # only send it when actually asked for. + if not _wanted_stats or 'raw_data' in _wanted_stats: + data_stats.append( + create_data_stat( + name='raw_data', + stat_type='bytes', + thumbnail=raw_data_bytes, + shape=raw_shape, + ) ) - ) + + # Paired/multi-view datasets (e.g. a source+edited image + # pair) can optionally expose additional named views via + # extra_images() -- send each as its own 'image_' + # stat so the frontend renders it as its own grid column + # (isImageFieldName() already recognizes 'image_*'). This + # duck-typed hook is a no-op for datasets that don't + # define it. raw_data above already covers view rank 0 + # (e.g. 'source'); extra_images() may repeat that view + # under its own name too -- one small duplicated + # thumbnail, traded for not having to assume which named + # view is redundant across arbitrary datasets. + # Probe the UNWRAPPED dataset: `dataset` is WL's tracking + # wrapper and does not forward extra_images, so testing it + # silently disables every named view. + if hasattr(ds, "extra_images"): + try: + extra_views = ds.extra_images(ds_idx) or {} + except Exception as e: + extra_views = {} + logger.debug(f"extra_images failed for sample_id={sample_id}: {e}") + for view_name, view_pil in extra_views.items(): + if view_pil is None: + continue + # Filter BEFORE resize/encode -- that is the cost. + # Still ADVERTISE the view with an empty thumbnail: + # the panel builds its modality list from the stats + # present, so omitting it entirely would delete the + # toggle and make the view unrecoverable. + if (_wanted_stats + and ("image_%s" % view_name) not in _wanted_stats): + data_stats.append( + create_data_stat( + name="image_%s" % view_name, + stat_type='bytes', + thumbnail=b"", + shape=[], + ) + ) + continue + try: + resized_view = view_pil + if resized_view.size != (target_width, target_height): + _view_resample = ( + Image.Resampling.LANCZOS if is_full_resolution + else Image.Resampling.BILINEAR + ) + resized_view = resized_view.resize( + (target_width, target_height), _view_resample + ) + view_bytes, view_shape, _ = encode_image_to_raw_bytes( + np_img=None, + middle_pil=resized_view, + original_shape=[], + is_volumetric=False, + is_full_resolution=is_full_resolution, + target_width=target_width, + target_height=target_height, + ) + data_stats.append( + create_data_stat( + name=f"image_{view_name}", + stat_type='bytes', + thumbnail=view_bytes, + shape=view_shape, + ) + ) + del view_bytes, resized_view + except Exception as e: + logger.debug( + f"extra_images encode failed for sample_id={sample_id} " + f"view={view_name}: {e}" + ) # For video samples the bytes above are only the poster # frame, so advertise the clip's shape here. This lets the @@ -2470,16 +2586,49 @@ def _sample_id_sortable_series(self, values): return numeric return values.astype(str) + def _numeric_like_sort_cols(self, df: pd.DataFrame, by) -> set: + """Sort columns whose values are strings but mean numbers. + + '1916469' < '191647' lexicographically but not numerically, so sorting a + numeric-valued string column by raw string order is simply wrong. Only + object/string columns are candidates; genuinely numeric dtypes already + sort correctly and must not be touched. + """ + out = set() + by_list = [by] if isinstance(by, str) else list(by or []) + for col in by_list: + if col == SampleStatsEx.SAMPLE_ID.value: + out.add(col) + continue + try: + s = df[col] if col in df.columns else None + if s is None or pd.api.types.is_numeric_dtype(s) or hasattr(s, "cat"): + continue + probe = s.dropna() + if probe.empty: + continue + # Cheap decision on a sample: a full 4M-row coercion here would + # cost more than the sort it is meant to correct. + head = probe.head(2048) + coerced = pd.to_numeric(head, errors="coerce") + if coerced.notna().all(): + out.add(col) + except Exception: + continue + return out + def _sort_values_numeric_aware(self, df: pd.DataFrame, sort_params: dict) -> None: - """Sort dataframe while treating sample_id as numeric when possible.""" + """Sort dataframe, ordering numeric-valued string columns numerically.""" params = dict(sort_params) - if params.get("key") is None and self._sort_includes_sample_id(params.get("by")): - def _key(series: pd.Series): - if str(getattr(series, "name", "")) == SampleStatsEx.SAMPLE_ID.value: - return self._sample_id_sortable_series(series) - return series + if params.get("key") is None: + numeric_like = self._numeric_like_sort_cols(df, params.get("by")) + if numeric_like: + def _key(series: pd.Series): + if str(getattr(series, "name", "")) in numeric_like: + return self._sample_id_sortable_series(series) + return series - params["key"] = _key + params["key"] = _key df.sort_values(inplace=True, **params) @@ -3527,6 +3676,38 @@ def _apply_agent_operation(self, df, func: str, params: dict) -> str: # silently stops refreshing). orig_index_names = [n for n in df.index.names if n is not None] + # Fast path: nothing in `by` is an index level, so the frame can be + # sorted where it stands. Avoids reset_index + astype(int)/astype(str) + # over every sample_id + a set_index that re-factorizes 3.96M string + # keys -- measured as the bulk of a ~20s sort. + _fp_by = params.get("by") + _fp_list = [_fp_by] if isinstance(_fp_by, str) else list(_fp_by or []) + _fp_res = [ + (c if c in df.columns + else ("signals//" + c if ("signals//" + c) in df.columns else c)) + for c in _fp_list + ] + if (_fp_res + and all(c in df.columns for c in _fp_res) + and not any(c in orig_index_names for c in _fp_list) + and not any(c in orig_index_names for c in _fp_res)): + _fp_params = dict(params) + _fp_params["by"] = (_fp_res if isinstance(_fp_by, (list, tuple)) + else _fp_res[0]) + try: + # Through the helper, NOT df.sort_values directly: a + # numeric-valued string column (group_id, target, ...) + # otherwise sorts lexicographically -- '1916469' before + # '191647'. + self._sort_values_numeric_aware(df, _fp_params) + return "Applied operation: sort_values" + except (TypeError, ValueError, KeyError) as _fp_exc: + # Mixed dtypes or an unexpected key: fall through to the + # original reset/restore path rather than failing the query. + logger.debug( + "[sort] fast path declined (%s); using index round-trip", + type(_fp_exc).__name__) + def _restore_index(): cols = [n for n in orig_index_names if n in df.columns] if cols and not isinstance(df.index, pd.MultiIndex): @@ -3789,12 +3970,99 @@ def _bg_view_refresh(self) -> None: real rebuild+swap via force=True OFF the request path, then releases the guard so a later stale read can trigger another. Never raises into a request.""" try: - self._slowUpdateInternals(force=True) + if not self._fastUpdateInternals(): + self._slowUpdateInternals(force=True) except Exception: logger.exception("[ViewRefresh] background view refresh failed") finally: self._refresh_in_flight.release() + # Columns the trainer mutates. Structural columns (origin, edit_prompt, + # task_type, ...) never change after registration, so a differential sync + # only has to carry these. + _FAST_SYNC_PREFIXES = ("signals", "last_seen", "discarded", "prediction", "target") + + def _fast_sync_columns(self, view): + return [c for c in view.columns + if str(c).startswith(self._FAST_SYNC_PREFIXES)] + + + + def _fastUpdateInternals(self, max_dirty: int = 250_000) -> bool: + """O(change) view refresh. True if applied, False -> caller must rebuild. + + Falls back when there is no view yet, when a dirty row is absent from + the view (new rows => structural change), or when the backlog is large + enough that a rebuild is cheaper. + """ + if not _fast_view_enabled(): + return False # opt-out: behave exactly as before + view = self._all_datasets_df + dfm = self._df_manager + if view is None or getattr(view, "empty", True) or dfm is None: + return False + # A manager without dirty tracking cannot serve a delta -- rebuild. + if not (hasattr(dfm, "take_view_dirty") and hasattr(dfm, "get_source_rows")): + return False + # A column the ledger has but the view lacks can only arrive via a full + # rebuild -- the differential write below addresses existing columns + # only. Per-sample signal columns are created on their first write, so + # on a fresh ledger the view predates them and would never gain them. + try: + _src = getattr(dfm, "_df", None) + if _src is not None: + _have = set(view.columns) + _missing = [c for c in _src.columns + if str(c).startswith(self._FAST_SYNC_PREFIXES) + and c not in _have] + if _missing: + return False + except Exception: + pass + + dirty = dfm.take_view_dirty(limit=max_dirty) + if dirty is None: + return False # backlog too large; rebuild is cheaper + if not dirty: + return True # nothing changed since last sync + + sids = [str(s) for s in dirty] + # Address rows by LABEL. pandas keeps a hash engine on the index, built + # in C and cached, so this needs no precomputed position map -- and a + # label that is absent surfaces below as a no-match rather than as a + # silently wrong row. + keep = sids + + cols = self._fast_sync_columns(view) + if not cols: + return True + sub = dfm.get_source_rows(keep, columns=[c for c in cols if c in view.columns]) + if sub is None or sub.empty: + return True + if isinstance(sub.index, pd.MultiIndex): + sub = sub.droplevel(-1) + sub = sub[~sub.index.duplicated(keep="last")] + + # Only rows the view actually holds; a structural change (new sample) + # must still fall back to the full rebuild rather than be invented here. + SID = SampleStatsEx.SAMPLE_ID.value + _names = list(getattr(view.index, "names", []) or []) + view_keys = (view.index.get_level_values(SID) + if isinstance(view.index, pd.MultiIndex) and SID in _names + else view.index) + # Positions via the Index hash engine: vectorised and cached, so this + # costs nothing like the O(rows) dict the position map used to rebuild. + _pos = pd.Index(view_keys.astype(str)).get_indexer(sub.index.astype(str)) + _ok = _pos >= 0 + if not _ok.any(): + return True + if not _ok.all(): + return False + for c in sub.columns: + _ci = view.columns.get_loc(c) + view.iloc[_pos, _ci] = sub[c].to_numpy() + return True + def _slowUpdateInternals(self, force: bool = False, reset_view: bool = False) -> None: """Update the internal dataframe view with the latest data from the manager. @@ -3965,6 +4233,13 @@ def _slowUpdateInternals(self, force: bool = False, reset_view: bool = False) -> # Atomic swap to make the new view available to readers self._all_datasets_df = updated_df self._last_internals_update_time = current_time + # The rebuilt view reflects every row, so the differential backlog is + # satisfied. This is the only point at which that is true. + try: + if self._df_manager is not None: + self._df_manager.clear_view_dirty() + except Exception: + pass finally: held_ms = (time.time() - t_held_start) * 1000 @@ -4482,7 +4757,8 @@ def _process_get_data_samples(self, request, context): ) # Trigger update if needed (it has its own internal locking) - self._slowUpdateInternals() + if not self._fastUpdateInternals(): + self._slowUpdateInternals() # Atomic snapshot of the current authoritative dataframe current_df = self._all_datasets_df @@ -4793,37 +5069,39 @@ def ApplyDataQuery(self, request, context): operations = self._parse_direct_query(request.query) # Apply operations with lock - with self._watched_lock("_lock[ApplyDataQuery/ops]"): - # Skip the forced full-view rebuild for SORT-ONLY operations. Sorting just - # re-orders the existing snapshot, so a fresh collapse+combine (hundreds of - # ms on large views, and — being lock-held — contends with the training - # thread for multi-second stalls) is unnecessary. Filters/edits still refresh - # so they operate on the latest data. The view is frozen on direct queries - # anyway (_is_filtered=True), so it wasn't auto-refreshing mid-sort regardless. - _SORT_FUNCS = {"df.sort_values", "df.sort_index", "df.sort_view_slice"} - is_sort_only = bool(operations) and all( - op.get("function") in _SORT_FUNCS for op in operations) - if not is_sort_only: - self._slowUpdateInternals(force=True) # Refresh internals before applying non-sort operations - - # Work on a copy to allow concurrent readers to see a consistent state - df = self._all_datasets_df # Remove copy because memory waste and slowdown - messages = [] + _SORT_FUNCS = {"df.sort_values", "df.sort_index", "df.sort_view_slice"} + is_sort_only = bool(operations) and all( + op.get("function") in _SORT_FUNCS for op in operations) + def _run_ops(target): + out = [] for op in operations: - func = op.get("function") - params = op.get("params", {}) or {} - msg = self._apply_agent_operation(df, func, params) - messages.append(msg) - - final_message = " | ".join(messages) if messages else "No operation performed" - - # Atomic swap - self._all_datasets_df = df - - # Direct queries are manipulations -> Freeze the view - if operations: - self._is_filtered = True + out.append(self._apply_agent_operation( + target, op.get("function"), op.get("params", {}) or {})) + return " | ".join(out) if out else "No operation performed" + + if is_sort_only: + # Sorting 4M rows costs ~10s; under _lock that stalls the trainer + # for the whole duration (measured 11716ms). Sort a shallow copy + # off-lock, then hold the lock only for the reference swap. + base = self._all_datasets_df + df = base.copy(deep=False) if base is not None else base + final_message = _run_ops(df) + # No position map to rebuild: the differential refresh addresses + # rows by label through pandas' own index engine, so reordering + # the view invalidates nothing. + with self._watched_lock("_lock[ApplyDataQuery/swap]"): + self._all_datasets_df = df + if operations: + self._is_filtered = True + else: + with self._watched_lock("_lock[ApplyDataQuery/ops]"): + self._slowUpdateInternals(force=True) + df = self._all_datasets_df + final_message = _run_ops(df) + self._all_datasets_df = df + if operations: + self._is_filtered = True return self._build_success_response( df=df, @@ -5047,17 +5325,37 @@ def GetHistogram(self, request, context): if df is None or df.empty: return pb2.HistogramResponse( success=False, message="empty dataframe view", total_rows=0, bins=[]) - df = safe_reset_index(df) + # safe_reset_index copies AND consolidates the entire frame (~70% of + # this RPC at 3.96M x 19 by py-spy). It is only needed to reach fields + # that live in the index; when they are already columns, use the frame + # as it stands. + def _field(frame, name): + """Series for *name* whether it is a column or an index level.""" + if name in frame.columns: + return frame[name] + names = list(getattr(frame.index, "names", []) or []) + if name in names: + return pd.Series(frame.index.get_level_values(name), + index=frame.index) + if getattr(frame.index, "name", None) == name: + return pd.Series(frame.index, index=frame.index) + return None + + # Only reset when the histogrammed column itself cannot be reached. + # safe_reset_index copies AND block-consolidates the whole frame + # (~82% of this RPC by py-spy); get_level_values is ~0.02s. + if _field(df, column) is None: + df = safe_reset_index(df) n = len(df) if column not in df.columns: return pb2.HistogramResponse( success=False, message=f"column '{column}' not in view", total_rows=n, bins=[]) - origin = (df["origin"].astype(str).to_numpy() if "origin" in df.columns - else np.full(n, "")) - disc = (df["discarded"].astype(bool).to_numpy() if "discarded" in df.columns - else np.zeros(n, bool)) + _o = _field(df, "origin") + origin = _o.astype(str).to_numpy() if _o is not None else np.full(n, "") + _d = _field(df, "discarded") + disc = _d.astype(bool).to_numpy() if _d is not None else np.zeros(n, bool) # Detect whether column is categorical (string/object) or numeric. # A column is numeric if ANY value coerces to a finite number — even @@ -5068,7 +5366,9 @@ def GetHistogram(self, request, context): # as a spurious "unset" bar). We therefore treat as categorical only # a genuine pandas ``category`` dtype, or a column whose values do # not coerce to any numeric value at all (pure strings). - col_series = df[column] + col_series = _field(df, column) + if col_series is None: + col_series = df[column] numeric_vals = pd.to_numeric(col_series, errors="coerce") is_category_dtype = ( str(col_series.dtype) == "category" or hasattr(col_series, "cat") @@ -5087,20 +5387,44 @@ def GetHistogram(self, request, context): if is_categorical: # --- Categorical path --- labels = col_series.astype(str).where(col_series.notna(), "") - gf = pd.DataFrame({"l": labels, "o": origin, "d": disc}) - total_count = gf.groupby("l")["l"].count().rename("count") + # Count first with a single-key value_counts, then restrict the + # three-key breakdown to the rows that survive the cap. Grouping + # all 3.96M rows by (label, discarded, origin) when the column is + # free text builds 722,870 groups to then discard all but 200. + total_count = labels.value_counts().rename("count") + _cap_pre = _histogram_category_cap() + _keep = set(total_count.iloc[:_cap_pre].index) + _m = labels.isin(_keep).to_numpy() sub_map: dict = {} - for (lbl, d, o), c in gf.groupby(["l", "d", "o"]).size().items(): - sub_map.setdefault(str(lbl), []).append( - pb2.HistogramSubBar(origin=str(o), discarded=bool(d), count=int(c))) + if _m.any(): + gf = pd.DataFrame({"l": labels.to_numpy()[_m], + "o": origin[_m], "d": disc[_m]}) + for (lbl, d, o), c in gf.groupby(["l", "d", "o"]).size().items(): + sub_map.setdefault(str(lbl), []).append( + pb2.HistogramSubBar(origin=str(o), discarded=bool(d), count=int(c))) + # Cap the output: a free-text column (e.g. edit_prompt) has one + # category per sample -- 722,870 bars / 54 MB / 30.5s measured, + # which no viewer can draw. Keep the top-N by count and fold the + # remainder into one "(other)" bar so the response stays bounded. + _ordered = total_count.sort_values(ascending=False) + _cap = _histogram_category_cap() + _head, _tail = _ordered.iloc[:_cap], _ordered.iloc[_cap:] cat_bars = [ pb2.CategoricalHistogramBar( label=str(lbl), count=int(cnt), sub_bars=sub_map.get(str(lbl), []), ) - for lbl, cnt in total_count.sort_values(ascending=False).items() + for lbl, cnt in _head.items() ] + if len(_tail): + cat_bars.append(pb2.CategoricalHistogramBar( + label="(other: %d categories)" % len(_tail), + count=int(_tail.sum()), + sub_bars=[], + )) + logger.info("[HistCat] column=%s capped %d categories -> %d bars", + column, len(_ordered), len(cat_bars)) logger.info("[HistCat] column=%s rows=%d categories=%d", column, n, len(cat_bars)) return pb2.HistogramResponse( @@ -5113,6 +5437,11 @@ def GetHistogram(self, request, context): ) # --- Numeric path (unchanged) --- + # Each bar covers a fixed slice of the VIEW: total_rows / max_bins + # samples. That is what makes the chart show density -- a column only + # 0.2% populated shows a few filled bars and the rest empty. Binning + # over just the rows that carry a value makes the chart look equally + # full at any coverage, which reads as "everything has a value". bars = max(1, min(n, max_bins)) vals = numeric_vals.to_numpy() edges = (np.arange(bars + 1) * n) // bars @@ -5627,7 +5956,8 @@ def EditDataSample(self, request, context): with self._watched_lock("_lock[EditDataSample/__copy_metadata__]"): try: - self._slowUpdateInternals() + if not self._fastUpdateInternals(): + self._slowUpdateInternals() if self._all_datasets_df is None or self._all_datasets_df.empty: return pb2.DataEditsResponse( success=False, @@ -5795,7 +6125,8 @@ def EditDataSample(self, request, context): with self._watched_lock("_lock[EditDataSample/delete-col]"): try: - self._slowUpdateInternals() + if not self._fastUpdateInternals(): + self._slowUpdateInternals() if self._all_datasets_df is None or self._all_datasets_df.empty: return pb2.DataEditsResponse( success=False, @@ -5833,7 +6164,8 @@ def EditDataSample(self, request, context): # Kick a background view-refresh (non-blocking) — the in-memory view # is already consistent after the drop above, so blocking inline rebuild # is unnecessary and causes the gRPC response to stall for 5-10 s. - self._slowUpdateInternals() + if not self._fastUpdateInternals(): + self._slowUpdateInternals() return pb2.DataEditsResponse( success=True, @@ -5856,7 +6188,8 @@ def EditDataSample(self, request, context): with self._watched_lock("_lock[EditDataSample/__discard_by_tag__]"): try: - self._slowUpdateInternals() + if not self._fastUpdateInternals(): + self._slowUpdateInternals() if self._all_datasets_df is None or self._all_datasets_df.empty: return pb2.DataEditsResponse( success=False, @@ -5865,7 +6198,8 @@ def EditDataSample(self, request, context): tag_col = f"{SampleStatsEx.TAG.value}:{tag_name}" if tag_col not in self._all_datasets_df.columns: - self._slowUpdateInternals() + if not self._fastUpdateInternals(): + self._slowUpdateInternals() df = safe_reset_index(self._all_datasets_df) if tag_col not in df.columns: return pb2.DataEditsResponse( @@ -6079,7 +6413,8 @@ def GetDataSplits(self, request, context): # IMPORTANT: keep lock ordering consistent (_update_lock -> _lock). # Calling _slowUpdateInternals() while holding _lock can deadlock # with concurrent readers/writers under high UI refresh pressure. - self._slowUpdateInternals() + if not self._fastUpdateInternals(): + self._slowUpdateInternals() if context is not None and not context.is_active(): return pb2.DataSplitsResponse(success=False, split_names=[]) diff --git a/weightslab/trainer/trainer_tools.py b/weightslab/trainer/trainer_tools.py index 136dd34c..1b2d467c 100644 --- a/weightslab/trainer/trainer_tools.py +++ b/weightslab/trainer/trainer_tools.py @@ -434,10 +434,15 @@ def _get_input_tensor_for_sample(dataset, sample_id, device): def process_sample(sid, dataset, do_resize, resize_dims, experiment): try: - if hasattr(dataset, "_getitem_raw"): - tensor, idx, label = dataset._getitem_raw(id=sid) - else: - tensor, idx, label = dataset[sid] + # _getitem_raw returns (data, id, target, *metadata) by contract -- datasets + # implementing get_items() with metadata yield 4+ elements, so a fixed + # 3-way unpack raises "too many values to unpack" and kills every + # thumbnail. Unpack positionally instead. + _res = dataset._getitem_raw(id=sid) if hasattr(dataset, "_getitem_raw") else dataset[sid] + if not isinstance(_res, (tuple, list)): + _res = (_res, sid, None) + tensor = _res[0] + label = _res[2] if len(_res) > 2 else None if isinstance(tensor, torch.Tensor): img = tensor.detach().cpu() diff --git a/weightslab/utils/logs.py b/weightslab/utils/logs.py index dd75872f..717165b4 100644 --- a/weightslab/utils/logs.py +++ b/weightslab/utils/logs.py @@ -9,7 +9,7 @@ # Define the log format to include timestamp, level, module name, and function name -FORMAT = '%(asctime)s.%(msecs)03d %(levelname)s:%(name)s:%(funcName)s: %(message)s' +FORMAT = '%(asctime)s.%(msecs)03d %(levelname)s:%(name)s:%(filename)s:%(lineno)d:%(funcName)s: %(message)s' DATE_FORMAT = '%d/%m/%Y-%H:%M:%S' # Global variables to track the log file path and handler