diff --git a/agent_config.yaml b/agent_config.yaml
index 3a3aa2f9..51b3c772 100644
--- a/agent_config.yaml
+++ b/agent_config.yaml
@@ -17,7 +17,16 @@ agent:
opencode_url: http://127.0.0.1:4096
# OpenCode model, "providerID/modelID" (can also be set as env variable
- # OPENCODE_MODEL). Empty string self-heals to whatever OpenCode's own
- # config was last set to, or a configured provider default, falling back
- # to the free-tier "opencode/deepseek-v4-flash-free" if neither resolves.
- opencode_model: "opencode/deepseek-v4-flash-free"
+ # OPENCODE_MODEL).
+ #
+ # Left EMPTY on purpose. Empty means "follow OpenCode's own config" -- which
+ # is what the studio's model picker (.wl-ag-model) writes on every pick, via
+ # PUT /config -- so choosing a model in the UI also changes the model
+ # weightslab's own queries use, and it is re-checked before every turn.
+ # When OpenCode names no model either, the fallback is "opencode/big-pickle".
+ #
+ # A value here SEEDS the choice: it is used when nothing has been chosen in
+ # OpenCode's config yet (and is published there, so the studio shows it), but
+ # a model picked in the UI afterwards wins. Set OPENCODE_MODEL instead to
+ # PIN a model that nothing can override.
+ opencode_model: ""
diff --git a/docs/agent.rst b/docs/agent.rst
index 596555f4..9519418d 100644
--- a/docs/agent.rst
+++ b/docs/agent.rst
@@ -211,16 +211,71 @@ If ``OPENCODE_URL`` is set and reachable, the UI server adopts it directly
instead of spawning a child; the backend SDK agent reads the same variable
(``agent.py``'s ``_load_config``) — set it once and both sides talk to the one
server, so a model you authenticate once is available everywhere.
-``OPENCODE_MODEL`` (or ``agent_config.yaml``'s ``agent.opencode_model``) picks
-the default model for the backend SDK agent, as an OpenCode
-``providerID/modelID`` string (e.g. ``openrouter/anthropic/claude-opus-4.6``).
-Leave it unset to fall back, in order, to: whatever model OpenCode's own
-``/config`` was last set to (the model picker's own pick, e.g. from the
-Weights Studio landing page), and otherwise the free-tier
-``opencode/deepseek-v4-flash-free`` automatically — a provider's own
-reported default used to be tried in between, but that could itself be an
-arbitrary, non-text-reasoning model whenever any provider had credentials
-configured, so it no longer overrides this.
+Which model gets used, and how the two sides agree
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+**OpenCode's own config is the shared source of truth.** ``GET /config``'s
+``model`` field is read by every client of that server -- the Weights Studio
+model picker, the OpenCode CLI, and the backend SDK agent -- and written by the
+picker (``PATCH /global/config``, ``global.config.update``) whenever you choose
+a model. Neither side has to know the other exists; they meet in that one
+field.
+
+The backend SDK agent resolves its model in this order:
+
+1. ``OPENCODE_MODEL`` -- a hard **pin**, for automation that must force a
+ model regardless of what anyone picked. The UI picker cannot move it; if the
+ two disagree, ``agent status`` logs which model this backend is actually
+ using and why.
+2. ``GET /config``'s ``model`` -- the shared choice above. Re-read **before
+ every turn**, so switching models in the UI mid-run takes effect
+ immediately; it is not latched at start-up.
+3. ``agent_config.yaml``'s ``agent.opencode_model`` -- a **seed**, not a pin:
+ which model to use when nothing has been chosen in OpenCode's config yet.
+ It is published there, so the studio shows it; once anything is chosen (in
+ the UI, by ``agent model``, or by the CLI), that choice wins and the seed
+ goes unused. Pinning the yaml value instead meant a run started *after*
+ picking a model in the studio quietly went back to the yaml one.
+4. ``opencode/big-pickle``, the built-in fallback, when nothing above resolves.
+ Published too, so a backend that started before the UI hands the picker the
+ model it is itself using.
+
+A provider's own reported default used to be tried just before the built-in
+fallback, but that could itself be an arbitrary, non-text-reasoning model
+whenever any provider had credentials configured, so it no longer overrides it.
+
+Whoever chooses last wins, and both surfaces follow: picking in the UI moves
+the backend's next query, and ``agent model `` from the
+CLI moves the UI's picker.
+
+Either start order therefore converges on one model:
+
+.. code-block:: text
+
+ Studio first: pick a model in the UI -> PATCH /global/config
+ -> weightslab start -> GET /config -> same model
+ (agent_config.yaml's seed is not used: something was chosen)
+
+ weightslab first: nothing chosen anywhere
+ -> agent_config.yaml's opencode_model, else
+ opencode/big-pickle -- and published
+ -> Studio starts -> GET /config -> same model
+
+The start-up banner states which of the three won:
+
+.. code-block:: text
+
+ Agent initialized from configuration C:\Users\you\wl_agent_config.yaml:
+ OpenCode URL=http://127.0.0.1:4096
+ Model=opencode/muse-spark-1.3-contributor-free (from agent_config.yaml's opencode_model (nothing chosen yet), published to OpenCode's config so the studio shows it)
+
+Other values in the parentheses are ``pinned by OPENCODE_MODEL``, ``from
+OpenCode's config, which the studio model picker writes``, ``chosen here and
+published to OpenCode's config`` (an ``agent model`` switch), ``built-in
+default, published to OpenCode's config for the studio``, and ``unresolved --
+OpenCode unreachable, retried on the first query``. It used to print ``Model=(server default)`` whenever nothing was
+pinned, which read as "my pick was ignored" even when the first query would
+have picked it up.
Credentials and provider setup live in OpenCode itself, never in WeightsLab:
@@ -335,7 +390,8 @@ configure ``agent_config.yaml`` and/or environment variables.
# 3. agent_config.yaml
agent:
opencode_url: http://127.0.0.1:4096
- opencode_model: "" # empty = use OpenCode's own configured default
+ opencode_model: "" # empty = follow OpenCode's config (the UI picker);
+ # a value here PINS the model instead
Then check it from the CLI:
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/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 00000000..32c86157
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,37 @@
+"""Shared pytest setup.
+
+Keeps the suite out of the developer's own WeightsLab state.
+
+``weightslab.utils.active_experiment`` records the active experiment directory
+in a real per-user file (``~/.weightslab/active_experiment.json``) so a
+training run started in another terminal lands in the experiment the UI
+established. That is deliberate at runtime -- and poison in tests: a run of
+this suite while a ``weightslab start`` was up resolved its root_log_dir into
+that live experiment, found the config and checkpoints of whatever was running
+there, and failed in setUp with an unrelated config (seen for real:
+tests/gRPC/test_grpc_tag_operations.py loading a segmentation example's
+hyperparameters).
+
+So every test session gets its own throwaway state directory. Tests that
+exercise the handoff itself set ``WEIGHTSLAB_STATE_DIR`` to their own temp
+directory anyway; this only changes the default.
+"""
+
+import os
+import tempfile
+
+import pytest
+
+
+@pytest.fixture(scope="session", autouse=True)
+def _isolate_weightslab_state():
+ previous = os.environ.get("WEIGHTSLAB_STATE_DIR")
+ with tempfile.TemporaryDirectory(prefix="wl-test-state-") as state_dir:
+ os.environ["WEIGHTSLAB_STATE_DIR"] = state_dir
+ try:
+ yield state_dir
+ finally:
+ if previous is None:
+ os.environ.pop("WEIGHTSLAB_STATE_DIR", None)
+ else:
+ os.environ["WEIGHTSLAB_STATE_DIR"] = previous
diff --git a/tests/test_src_functions.py b/tests/test_src_functions.py
index d99ed489..50c7a2b3 100644
--- a/tests/test_src_functions.py
+++ b/tests/test_src_functions.py
@@ -1,4 +1,6 @@
+import json
import os
+import shutil
import tempfile
import unittest
import numpy as np
@@ -6,6 +8,7 @@
import torch as th
import weightslab.src as src
+from weightslab.utils import active_experiment
from unittest.mock import MagicMock, patch
@@ -13,16 +16,27 @@
class TestResolveConfiguredRootLogDir(unittest.TestCase):
- """root_log_dir resolution: explicit config > WEIGHTSLAB_ROOT_LOG_DIR > temp dir."""
+ """root_log_dir resolution: explicit config > WEIGHTSLAB_ROOT_LOG_DIR >
+ the directory `weightslab start` recorded > temp dir."""
def setUp(self):
self._env_prev = os.environ.get("WEIGHTSLAB_ROOT_LOG_DIR")
+ # The marker is a real per-user file; point it at a scratch directory so
+ # these tests never read (or write) the developer's own active run.
+ self._state_prev = os.environ.get("WEIGHTSLAB_STATE_DIR")
+ self._state_dir = tempfile.mkdtemp()
+ os.environ["WEIGHTSLAB_STATE_DIR"] = self._state_dir
def tearDown(self):
if self._env_prev is None:
os.environ.pop("WEIGHTSLAB_ROOT_LOG_DIR", None)
else:
os.environ["WEIGHTSLAB_ROOT_LOG_DIR"] = self._env_prev
+ if self._state_prev is None:
+ os.environ.pop("WEIGHTSLAB_STATE_DIR", None)
+ else:
+ os.environ["WEIGHTSLAB_STATE_DIR"] = self._state_prev
+ shutil.rmtree(self._state_dir, ignore_errors=True)
def test_explicit_config_value_wins_over_env(self):
os.environ["WEIGHTSLAB_ROOT_LOG_DIR"] = "/env/dir"
@@ -49,6 +63,162 @@ def test_falls_back_to_tempdir_when_neither_set(self):
self.assertEqual(src._resolve_configured_root_log_dir(None), "/tmp/generated")
mk.assert_called_once()
+ def test_a_recorded_directory_whose_ui_has_exited_is_not_adopted(self):
+ # The handoff means "the UI is up over there, join its experiment". A
+ # record left by a `weightslab start` that has since exited must not
+ # redirect an unrelated run -- it did, and this repo's own gRPC tests
+ # resolved into a previous session's experiment and loaded its config.
+ os.environ.pop("WEIGHTSLAB_ROOT_LOG_DIR", None)
+ with tempfile.TemporaryDirectory() as ui_dir:
+ active_experiment.record_ui_experiment(ui_dir)
+ state = active_experiment.read_state()
+ state["ui"][-1]["pid"] = 2 ** 31 - 1 # cannot be running
+ active_experiment.state_path().write_text(json.dumps(state), encoding="utf-8")
+
+ with patch("weightslab.src.tempfile.mkdtemp", return_value="/tmp/generated") as mk:
+ self.assertEqual(src._resolve_configured_root_log_dir(None), "/tmp/generated")
+ mk.assert_called_once()
+
+ def test_adopts_the_directory_weightslab_start_recorded(self):
+ # `weightslab start` exports WEIGHTSLAB_ROOT_LOG_DIR into its OWN
+ # process only. A training run in another terminal never saw it and
+ # went to a temp dir, so the UI listed an empty reports/ while the run
+ # wrote elsewhere. The recorded directory closes that gap.
+ os.environ.pop("WEIGHTSLAB_ROOT_LOG_DIR", None)
+ with tempfile.TemporaryDirectory() as ui_dir:
+ active_experiment.record_ui_experiment(ui_dir)
+ with patch("weightslab.src.tempfile.mkdtemp", return_value="/tmp/generated") as mk:
+ resolved = src._resolve_configured_root_log_dir(None)
+ mk.assert_not_called()
+ self.assertEqual(os.path.realpath(resolved), os.path.realpath(ui_dir))
+
+ def test_explicit_config_still_wins_over_the_recorded_directory(self):
+ with tempfile.TemporaryDirectory() as ui_dir:
+ active_experiment.record_ui_experiment(ui_dir)
+ self.assertEqual(src._resolve_configured_root_log_dir("/explicit/dir"), "/explicit/dir")
+
+ def test_env_wins_over_the_recorded_directory(self):
+ with tempfile.TemporaryDirectory() as ui_dir, tempfile.TemporaryDirectory() as env_dir:
+ active_experiment.record_ui_experiment(ui_dir)
+ os.environ["WEIGHTSLAB_ROOT_LOG_DIR"] = env_dir
+ self.assertEqual(src._resolve_configured_root_log_dir(None), env_dir)
+
+ def test_a_recorded_directory_that_no_longer_exists_is_ignored(self):
+ os.environ.pop("WEIGHTSLAB_ROOT_LOG_DIR", None)
+ gone = tempfile.mkdtemp()
+ active_experiment.record_ui_experiment(gone)
+ shutil.rmtree(gone, ignore_errors=True)
+ with patch("weightslab.src.tempfile.mkdtemp", return_value="/tmp/generated") as mk:
+ self.assertEqual(src._resolve_configured_root_log_dir(None), "/tmp/generated")
+ mk.assert_called_once()
+
+
+class TestActiveExperimentMarker(unittest.TestCase):
+ """The cross-process handoff itself."""
+
+ def setUp(self):
+ self._state_prev = os.environ.get("WEIGHTSLAB_STATE_DIR")
+ self._state_dir = tempfile.mkdtemp()
+ os.environ["WEIGHTSLAB_STATE_DIR"] = self._state_dir
+
+ def tearDown(self):
+ if self._state_prev is None:
+ os.environ.pop("WEIGHTSLAB_STATE_DIR", None)
+ else:
+ os.environ["WEIGHTSLAB_STATE_DIR"] = self._state_prev
+ shutil.rmtree(self._state_dir, ignore_errors=True)
+
+ def test_two_uis_are_recorded_side_by_side(self):
+ # Two experiments at once (a cls UI and a seg UI) is supported; with one
+ # slot per side the second `weightslab start` erased the first.
+ with tempfile.TemporaryDirectory() as cls_dir, tempfile.TemporaryDirectory() as seg_dir:
+ active_experiment.record_ui_experiment(cls_dir, ui_port=8080, backend_port=50051)
+ # A second UI, standing in for another process.
+ state = active_experiment.read_state()
+ state["ui"].append({
+ "root_log_dir": seg_dir, "pid": os.getpid(),
+ "ui_port": 8081, "backend_port": 50052,
+ })
+ active_experiment.state_path().write_text(json.dumps(state), encoding="utf-8")
+
+ recorded = {e["root_log_dir"] for e in active_experiment.entries("ui")}
+ self.assertEqual(len(recorded), 2)
+
+ def test_two_live_uis_are_not_guessed_between(self):
+ with tempfile.TemporaryDirectory() as a_dir, tempfile.TemporaryDirectory() as b_dir:
+ active_experiment.record_ui_experiment(a_dir)
+ state = active_experiment.read_state()
+ state["ui"].append({"root_log_dir": b_dir, "pid": os.getpid()})
+ active_experiment.state_path().write_text(json.dumps(state), encoding="utf-8")
+
+ # Adopting either would put the run in the wrong experiment.
+ self.assertIsNone(active_experiment.live_ui_experiment_dir())
+
+ def test_a_backend_is_found_by_the_port_the_caller_talks_to(self):
+ with tempfile.TemporaryDirectory() as cls_dir, tempfile.TemporaryDirectory() as seg_dir:
+ state = {"backend": [
+ {"root_log_dir": cls_dir, "pid": os.getpid(), "grpc_port": 50051},
+ {"root_log_dir": seg_dir, "pid": os.getpid(), "grpc_port": 50052},
+ ]}
+ active_experiment.state_path().parent.mkdir(parents=True, exist_ok=True)
+ active_experiment.state_path().write_text(json.dumps(state), encoding="utf-8")
+
+ self.assertEqual(
+ os.path.realpath(active_experiment.live_backend_experiment_dir(50051)),
+ os.path.realpath(cls_dir))
+ self.assertEqual(
+ os.path.realpath(active_experiment.live_backend_experiment_dir(50052)),
+ os.path.realpath(seg_dir))
+ # No port, two candidates: no guess.
+ self.assertIsNone(active_experiment.live_backend_experiment_dir())
+ # A port nobody serves: no guess either.
+ self.assertIsNone(active_experiment.live_backend_experiment_dir(50099))
+
+ def test_the_older_single_object_marker_is_still_readable(self):
+ with tempfile.TemporaryDirectory() as ui_dir:
+ active_experiment.state_path().parent.mkdir(parents=True, exist_ok=True)
+ active_experiment.state_path().write_text(
+ json.dumps({"ui": {"root_log_dir": ui_dir, "pid": os.getpid()}}),
+ encoding="utf-8")
+ self.assertEqual(os.path.realpath(active_experiment.ui_experiment_dir()),
+ os.path.realpath(ui_dir))
+ self.assertEqual(os.path.realpath(active_experiment.live_ui_experiment_dir()),
+ os.path.realpath(ui_dir))
+
+ def test_ui_and_backend_entries_do_not_clobber_each_other(self):
+ with tempfile.TemporaryDirectory() as ui_dir, tempfile.TemporaryDirectory() as be_dir:
+ active_experiment.record_ui_experiment(ui_dir, ui_port=8080)
+ active_experiment.record_backend_experiment(be_dir)
+ self.assertEqual(os.path.realpath(active_experiment.ui_experiment_dir()),
+ os.path.realpath(ui_dir))
+ self.assertEqual(os.path.realpath(active_experiment.backend_experiment_dir()),
+ os.path.realpath(be_dir))
+ self.assertEqual(active_experiment.entries("ui")[-1]["ui_port"], 8080)
+
+ def test_missing_marker_reads_as_nothing_recorded(self):
+ self.assertEqual(active_experiment.read_state(), {})
+ self.assertIsNone(active_experiment.ui_experiment_dir())
+ self.assertIsNone(active_experiment.backend_experiment_dir())
+
+ def test_a_corrupt_marker_is_ignored_rather_than_raising(self):
+ path = active_experiment.state_path()
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text("{not json", encoding="utf-8")
+ self.assertEqual(active_experiment.read_state(), {})
+ self.assertIsNone(active_experiment.ui_experiment_dir())
+
+ def test_recording_nothing_is_a_no_op(self):
+ self.assertIsNone(active_experiment.record_ui_experiment(""))
+ self.assertEqual(active_experiment.read_state(), {})
+
+ def test_clear_removes_the_marker(self):
+ with tempfile.TemporaryDirectory() as ui_dir:
+ active_experiment.record_ui_experiment(ui_dir)
+ self.assertTrue(active_experiment.state_path().exists())
+ active_experiment.clear()
+ self.assertFalse(active_experiment.state_path().exists())
+ active_experiment.clear() # idempotent
+
class TestSrcTagAndDiscardFunctions(unittest.TestCase):
def setUp(self):
diff --git a/tests/trainer/services/test_agent_opencode_provider.py b/tests/trainer/services/test_agent_opencode_provider.py
index c2140950..868ba9c5 100644
--- a/tests/trainer/services/test_agent_opencode_provider.py
+++ b/tests/trainer/services/test_agent_opencode_provider.py
@@ -58,6 +58,118 @@ def _make_agent(df=None):
return agent_mod, agent
+def _agent_module():
+ with mock.patch.dict(sys.modules, _install_agent_dependency_stubs(), clear=False):
+ return importlib.import_module("weightslab.trainer.services.agent.agent")
+
+
+def _bare_agent(mod, model="opencode/ling-3.0-flash-fin-free"):
+ """An agent with just the OpenCode attributes _setup_providers touches.
+
+ Sidesteps the full DataManipulationAgent construction (schema build, ctx,
+ handler registry) -- these tests are about the model plumbing only.
+ """
+ agent = mod.DataManipulationAgent.__new__(mod.DataManipulationAgent)
+ agent.opencode_url = "http://127.0.0.1:4096"
+ agent.opencode_model = model
+ agent.opencode_workspace_dir = "."
+ agent._opencode_url_explicit = False
+ agent._opencode_model_explicit = False
+ agent._config_source_path = "(test)"
+ agent.preferred_provider = "opencode"
+ return agent
+
+
+def _fake_chat(publish_ok=True, configured="opencode/ling-3.0-flash-fin-free"):
+ chat = MagicMock()
+ chat.base_url = "http://127.0.0.1:4096"
+ chat.model_is_explicit = False
+ chat.publish_model.return_value = publish_ok
+ # What the shared config would hand back if it were consulted.
+ chat.resolve_model.return_value = (configured, "opencode-config")
+ return chat
+
+
+class TestModelSwitchAndReporting(unittest.TestCase):
+ """`agent model X` must actually switch to X, and `agent status` must name
+ the model the NEXT query will use.
+
+ Regression: the shared-config re-read that lets the studio's model picker
+ move the backend also ran on the explicit switch path, overwriting the model
+ the caller had just asked for -- `agent model opencode/big-pickle` replied
+ "Model switched to opencode/ling-3.0-flash-fin-free".
+ """
+
+ def test_change_model_switches_and_publishes_the_choice(self):
+ mod = _agent_module()
+ agent = _bare_agent(mod)
+ chat = _fake_chat(publish_ok=True)
+ with mock.patch.object(mod, "OpenCodeChat", return_value=chat):
+ ok, message = agent.change_model("opencode/big-pickle")
+
+ self.assertTrue(ok)
+ self.assertEqual(agent.opencode_model, "opencode/big-pickle")
+ self.assertIn("opencode/big-pickle", message)
+ self.assertNotIn("ling-3.0", message)
+ chat.publish_model.assert_called_once_with("opencode/big-pickle")
+ # The shared config must NOT be consulted on an explicit switch.
+ chat.resolve_model.assert_not_called()
+ # Published, so later turns keep following the config (the studio can
+ # still move it) rather than being pinned to this backend.
+ self.assertFalse(agent._opencode_model_explicit)
+ self.assertEqual(agent._opencode_model_source, "user-published")
+
+ def test_change_model_pins_when_the_config_refuses_the_write(self):
+ mod = _agent_module()
+ agent = _bare_agent(mod)
+ chat = _fake_chat(publish_ok=False)
+ with mock.patch.object(mod, "OpenCodeChat", return_value=chat):
+ ok, message = agent.change_model("opencode/big-pickle")
+
+ self.assertTrue(ok)
+ self.assertEqual(agent.opencode_model, "opencode/big-pickle")
+ self.assertTrue(agent._opencode_model_explicit)
+ self.assertTrue(chat.model_is_explicit)
+ self.assertEqual(agent._opencode_model_source, "user-pinned")
+ self.assertIn("pinned", message)
+
+ def test_change_model_rejects_an_empty_model(self):
+ mod = _agent_module()
+ agent = _bare_agent(mod)
+ ok, message = agent.change_model(" ")
+ self.assertFalse(ok)
+ self.assertIn("empty", message.lower())
+
+ def test_startup_without_a_request_follows_the_shared_config(self):
+ mod = _agent_module()
+ agent = _bare_agent(mod, model="")
+ chat = _fake_chat(configured="opencode/muse-spark-1.3-contributor-free")
+ with mock.patch.object(mod, "OpenCodeChat", return_value=chat):
+ agent._setup_providers()
+
+ chat.resolve_model.assert_called_once_with(publish_default=True)
+ self.assertEqual(agent.opencode_model, "opencode/muse-spark-1.3-contributor-free")
+ self.assertEqual(agent._opencode_model_source, "opencode-config")
+
+ def test_current_model_reports_a_studio_pick_made_after_start_up(self):
+ mod = _agent_module()
+ agent = _bare_agent(mod, model="opencode/big-pickle")
+ chat = _fake_chat(configured="openrouter/openai/gpt-5-mini")
+ agent._opencode_chat = chat
+
+ self.assertEqual(agent.current_model(), "openrouter/openai/gpt-5-mini")
+ self.assertEqual(agent.opencode_model, "openrouter/openai/gpt-5-mini")
+
+ def test_current_model_keeps_the_last_known_model_when_the_server_is_down(self):
+ mod = _agent_module()
+ agent = _bare_agent(mod, model="opencode/big-pickle")
+ chat = _fake_chat()
+ chat.resolve_model.side_effect = OSError("refused")
+ agent._opencode_chat = chat
+
+ self.assertEqual(agent.current_model(), "opencode/big-pickle")
+
+
@unittest.skip("Not ready yet -- OpenCodeChat is still a stub, and the test suite needs to be reworked to support it")
class TestOpenCodeConfigLoading(unittest.TestCase):
def test_opencode_url_and_model_default(self):
@@ -118,11 +230,16 @@ def test_opencode_chain_is_built(self):
mock_cls.return_value.as_runnable.return_value = fake_runnable
initialized = agent._setup_providers()
+ # seed_model carries agent_config.yaml's opencode_model, which SEEDS
+ # the shared choice instead of pinning it (a model picked in the studio
+ # wins over it); model/model_is_explicit carry OPENCODE_MODEL, which
+ # does pin.
mock_cls.assert_called_once_with(
agent.opencode_url, agent.opencode_model,
workspace_dir=agent.opencode_workspace_dir,
url_is_explicit=agent._opencode_url_explicit,
model_is_explicit=agent._opencode_model_explicit,
+ seed_model=agent._opencode_model_seed,
)
self.assertTrue(initialized)
self.assertIs(agent.chain_opencode, fake_runnable)
diff --git a/tests/trainer/services/test_data_service_discard_flag.py b/tests/trainer/services/test_data_service_discard_flag.py
new file mode 100644
index 00000000..97da1fa5
--- /dev/null
+++ b/tests/trainer/services/test_data_service_discard_flag.py
@@ -0,0 +1,308 @@
+"""The 'samples discard themselves in the studio while training' regression.
+
+Reported for detection/segmentation runs: as the model worked through the
+dataset, sample after sample greyed out in the UI, while the dataframe (checked
+from the notebook) said nothing was discarded.
+
+The chain, reproduced below:
+
+1. the trainer touches a sample -> its rows go dirty;
+2. `_fastUpdateInternals` syncs the columns the trainer mutates
+ (``signals*``, ``last_seen``, ``discarded``, ``prediction``, ``target``)
+ from the source into the view. It collapsed the source's per-annotation rows
+ with ``duplicated(keep="last")``, i.e. it kept the LAST annotation row --
+ whose *sample-level* columns are NaN, because the real values live on the
+ canonical row (annotation_id == 0);
+3. NaN therefore landed in the view's ``discarded``;
+4. GetDataSamples reported that flag as ``"1" if bool(value) else "0"`` -- and
+ ``bool(float("nan"))`` is True in Python.
+
+So every sample the model had seen was served as discarded. Only for
+annotation-expanded (detection/segmentation) ledgers, and only after training
+touched the sample: exactly the reported shape.
+"""
+
+import unittest
+
+import numpy as np
+import pandas as pd
+
+from weightslab.data.sample_stats import SampleStatsEx
+from weightslab.trainer.services.data_service import (
+ DataService,
+ is_set_flag,
+ set_flag_mask,
+)
+
+
+SID = SampleStatsEx.SAMPLE_ID.value
+ANNOT = SampleStatsEx.INSTANCE_ID.value
+DISCARDED = SampleStatsEx.DISCARDED.value
+
+
+def _source_rows():
+ """A segmentation ledger: sample-level values on annotation 0 only."""
+ return pd.DataFrame(
+ {
+ DISCARDED: [False, np.nan, np.nan, True, np.nan],
+ "last_seen": [11, 11, 11, 12, 12],
+ "prediction": ["a", None, None, "b", None],
+ },
+ index=pd.MultiIndex.from_tuples(
+ [("0", 0), ("0", 1), ("0", 2), ("1", 0), ("1", 1)],
+ names=[SID, ANNOT],
+ ),
+ )
+
+
+class _FakeManager:
+ """Minimal dataframe manager: dirty tracking + source row lookup."""
+
+ def __init__(self, source, dirty):
+ self._df = source
+ self._dirty = list(dirty)
+
+ def take_view_dirty(self, limit=None):
+ dirty, self._dirty = self._dirty, []
+ return dirty
+
+ def get_source_rows(self, sample_ids, columns=None):
+ wanted = [str(s) for s in sample_ids]
+ level = self._df.index.get_level_values(SID).astype(str)
+ rows = self._df[level.isin(wanted)]
+ return rows[columns] if columns else rows
+
+
+class TestFastViewSyncKeepsSampleLevelValues(unittest.TestCase):
+ def _service(self, source, dirty, view):
+ service = DataService.__new__(DataService)
+ service._all_datasets_df = view
+ service._df_manager = _FakeManager(source, dirty)
+ return service
+
+ def _collapsed_view(self, source):
+ """One row per sample, as the real view is built: annotation 0 wins."""
+ base = source[source.index.get_level_values(ANNOT) == 0].droplevel(ANNOT)
+ return base.copy()
+
+ def test_a_seen_sample_keeps_its_discarded_flag(self):
+ source = _source_rows()
+ view = self._collapsed_view(source)
+ service = self._service(source, ["0", "1"], view)
+
+ self.assertTrue(service._fastUpdateInternals())
+
+ # Sample 0 is NOT discarded and must stay that way; sample 1 is.
+ self.assertIs(bool(view.loc["0", DISCARDED]), False)
+ self.assertIs(bool(view.loc["1", DISCARDED]), True)
+ # The give-away of the old behaviour: a NaN in a column that had a value.
+ self.assertFalse(view[DISCARDED].isna().any(),
+ "sample-level flags were overwritten with an instance row's NaN")
+
+ def test_the_columns_the_trainer_owns_still_sync(self):
+ source = _source_rows()
+ view = self._collapsed_view(source)
+ source.loc[("0", 0), "last_seen"] = 99
+ service = self._service(source, ["0"], view)
+
+ self.assertTrue(service._fastUpdateInternals())
+ self.assertEqual(view.loc["0", "last_seen"], 99)
+
+ def test_falls_back_to_the_first_row_when_no_canonical_row_is_present(self):
+ # A slice of instance rows only (no annotation 0) must still sync
+ # something sane rather than raising or inventing a flag.
+ source = _source_rows().drop(index=("0", 0))
+ view = self._collapsed_view(_source_rows())
+ service = self._service(source, ["0"], view)
+
+ self.assertTrue(service._fastUpdateInternals())
+ self.assertTrue(pd.isna(view.loc["0", DISCARDED]) or view.loc["0", DISCARDED] is False)
+
+
+class TestDiscardedFlagIsNaNSafe(unittest.TestCase):
+ """The second half: how a nullable flag becomes what the studio renders.
+
+ is_set_flag / set_flag_mask are shared by the three places that read one:
+ GetDataSamples' `discarded` rendering flag, the boolean ``tag:*`` columns in
+ the metadata response, and the histogram's per-(origin, discarded) split.
+ """
+
+ @staticmethod
+ def _served_flag(value):
+ return "1" if is_set_flag(value) else "0"
+
+ def test_missing_is_not_discarded(self):
+ # bool(float("nan")) is True -- the whole bug in one line.
+ self.assertEqual(self._served_flag(np.nan), "0")
+ self.assertEqual(self._served_flag(None), "0")
+ self.assertEqual(self._served_flag(pd.NA), "0")
+
+ def test_real_values_still_come_through(self):
+ self.assertEqual(self._served_flag(True), "1")
+ self.assertEqual(self._served_flag(np.bool_(True)), "1")
+ self.assertEqual(self._served_flag(1), "1")
+ self.assertEqual(self._served_flag(False), "0")
+ self.assertEqual(self._served_flag(0), "0")
+
+
+ def test_a_string_flag_is_read_as_a_word_not_as_a_non_empty_string(self):
+ # bool("False") is True, and a boolean column that has been through the
+ # H5 store (categorical) can come back as these strings.
+ self.assertEqual(self._served_flag("False"), "0")
+ self.assertEqual(self._served_flag("false"), "0")
+ self.assertEqual(self._served_flag("0"), "0")
+ self.assertEqual(self._served_flag(""), "0")
+ self.assertEqual(self._served_flag("True"), "1")
+ self.assertEqual(self._served_flag("true"), "1")
+ self.assertEqual(self._served_flag("1"), "1")
+
+
+class TestSetFlagMask(unittest.TestCase):
+ """The column-wide form, used for tags and the histogram split."""
+
+ def test_a_sparse_tag_column_marks_only_the_tagged_samples(self):
+ # A tag is set on a few samples; every other row is NaN. astype(bool)
+ # turned those into True -- every sample wore every tag.
+ column = pd.Series([True, np.nan, False, None, True], dtype=object)
+ self.assertEqual(set_flag_mask(column).tolist(),
+ [True, False, False, False, True])
+
+ def test_it_handles_a_categorical_column(self):
+ # The H5 store optimises tag:* and discarded to categorical dtype.
+ column = pd.Series([True, None, False], dtype=object).astype("category")
+ self.assertEqual(set_flag_mask(column).tolist(), [True, False, False])
+
+ def test_it_handles_a_float_column_of_zeros_and_nans(self):
+ column = pd.Series([1.0, np.nan, 0.0])
+ self.assertEqual(set_flag_mask(column).tolist(), [True, False, False])
+
+ def test_an_absent_column_is_all_false(self):
+ self.assertEqual(set_flag_mask(None).tolist(), [])
+
+
+class TestFastViewSyncAddressing(unittest.TestCase):
+ """How dirty source rows are matched to view rows.
+
+ The view is indexed (origin, sample_id) precisely because one sample_id can
+ appear under two origins. Looking positions up in that non-unique level
+ raised InvalidIndexError, so the differential refresh failed on every call
+ and quietly fell back to the full rebuild.
+ """
+
+ def _service(self, source, dirty, view):
+ service = DataService.__new__(DataService)
+ service._all_datasets_df = view
+ service._df_manager = _FakeManager(source, dirty)
+ return service
+
+ def _source(self):
+ return pd.DataFrame(
+ {DISCARDED: [False, np.nan], "last_seen": [7, 7]},
+ index=pd.MultiIndex.from_tuples([("5", 0), ("5", 1)], names=[SID, ANNOT]),
+ )
+
+ def test_one_sample_id_under_two_origins_does_not_raise(self):
+ view = pd.DataFrame(
+ {DISCARDED: [False, False], "last_seen": [1, 2]},
+ index=pd.MultiIndex.from_tuples([("train_loader", "5"), ("test_loader", "5")],
+ names=["origin", SID]),
+ )
+ service = self._service(self._source(), ["5"], view)
+
+ self.assertTrue(service._fastUpdateInternals())
+ # The source can only speak per sample_id, so both rows take its values.
+ self.assertEqual(view["last_seen"].tolist(), [7, 7])
+ self.assertFalse(view[DISCARDED].isna().any())
+
+ def test_a_dirty_sample_the_view_does_not_hold_forces_a_rebuild(self):
+ view = pd.DataFrame(
+ {DISCARDED: [False], "last_seen": [1]},
+ index=pd.MultiIndex.from_tuples([("train_loader", "5")], names=["origin", SID]),
+ )
+ source = pd.concat([self._source(), pd.DataFrame(
+ {DISCARDED: [False], "last_seen": [3]},
+ index=pd.MultiIndex.from_tuples([("99", 0)], names=[SID, ANNOT]))])
+ service = self._service(source, ["5", "99"], view)
+
+ # 99 is new -> structural change -> only the full rebuild can add it.
+ self.assertFalse(service._fastUpdateInternals())
+
+ def test_nothing_to_do_when_no_dirty_sample_is_in_the_view(self):
+ view = pd.DataFrame(
+ {DISCARDED: [False], "last_seen": [1]},
+ index=pd.MultiIndex.from_tuples([("train_loader", "7")], names=["origin", SID]),
+ )
+ service = self._service(self._source(), ["5"], view)
+ self.assertTrue(service._fastUpdateInternals())
+ self.assertEqual(view["last_seen"].tolist(), [1])
+
+
+class TestDocumentedFlagDefaults(unittest.TestCase):
+ """`discarded` should never be NaN in the first place.
+
+ SampleStats.DEFAULTS documents it as False, right under the comment "None
+ are not accepted by PD H5 storage". It held NaN anyway: the existing
+ normalisation only visits columns an upsert ADDS, and only when the
+ incoming slice's dtype is already bool -- which it is not exactly when the
+ slice carries missing values. So a sample registered without the flag, and
+ every per-annotation row (sample-level values live on annotation 0), kept a
+ NaN, which `bool()` then read as True.
+
+ Belt and braces with the read-side fix above: the flag is defaulted at the
+ source, AND a NaN that reaches a reader anyway is read as not-set.
+ """
+
+ def _manager(self, frame):
+ from weightslab.data.dataframe_manager import LedgeredDataFrameManager
+ manager = LedgeredDataFrameManager.__new__(LedgeredDataFrameManager)
+ manager._df = frame
+ return manager
+
+ def test_missing_discarded_becomes_false(self):
+ frame = pd.DataFrame(
+ {DISCARDED: [True, np.nan, np.nan, False, np.nan]},
+ index=pd.MultiIndex.from_tuples(
+ [("0", 0), ("0", 1), ("0", 2), ("1", 0), ("1", 1)], names=[SID, ANNOT]),
+ )
+ manager = self._manager(frame)
+ manager._fill_documented_flag_defaults()
+
+ self.assertEqual(frame[DISCARDED].tolist(), [True, False, False, False, False])
+ self.assertFalse(frame[DISCARDED].isna().any())
+
+ def test_a_categorical_flag_column_is_widened_rather_than_raising(self):
+ # fillna on a Categorical raises unless the value is a known category,
+ # and the H5 store hands these columns back as categorical.
+ # NB: build the frame first, THEN cast -- handing a Series with its own
+ # RangeIndex to DataFrame(index=[...]) reindexes it to all-NaN.
+ frame = pd.DataFrame({DISCARDED: [True, None]}, index=pd.Index(["0", "1"], name=SID))
+ frame[DISCARDED] = frame[DISCARDED].astype("category")
+ self.assertIsInstance(frame[DISCARDED].dtype, pd.CategoricalDtype)
+ manager = self._manager(frame)
+ manager._fill_documented_flag_defaults()
+ self.assertEqual(frame[DISCARDED].tolist(), [True, False])
+
+ def test_sparse_tag_columns_are_left_alone(self):
+ # NaN and False mean the same thing for a boolean tag, and NaN is
+ # cheaper; for a CATEGORICAL tag, NaN means "unset", not a default.
+ frame = pd.DataFrame({"tag:hard": [True, np.nan], DISCARDED: [False, np.nan]},
+ index=pd.Index(["0", "1"], name=SID))
+ manager = self._manager(frame)
+ manager._fill_documented_flag_defaults()
+ self.assertTrue(pd.isna(frame["tag:hard"].iloc[1]))
+ self.assertIs(bool(frame[DISCARDED].iloc[1]), False)
+
+ def test_a_column_with_nothing_missing_is_not_rewritten(self):
+ frame = pd.DataFrame({DISCARDED: [True, False]}, index=pd.Index(["0", "1"], name=SID))
+ before = frame[DISCARDED].dtype
+ manager = self._manager(frame)
+ manager._fill_documented_flag_defaults()
+ self.assertEqual(frame[DISCARDED].dtype, before)
+
+ def test_an_empty_frame_is_a_no_op(self):
+ frame = pd.DataFrame()
+ self._manager(frame)._fill_documented_flag_defaults() # must not raise
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/trainer/services/test_opencode_chat.py b/tests/trainer/services/test_opencode_chat.py
index 81468e21..13d8caf5 100644
--- a/tests/trainer/services/test_opencode_chat.py
+++ b/tests/trainer/services/test_opencode_chat.py
@@ -398,15 +398,37 @@ def test_explicit_model_is_never_auto_replaced(self):
request_mock.assert_not_called()
self.assertEqual(chat.model, "openrouter/anthropic/claude-opus-4.6")
- def test_already_set_non_explicit_model_is_left_alone(self):
- # Already resolved once (e.g. a prior call) -- don't re-resolve or
- # re-request every single turn.
+ def test_already_set_non_explicit_model_follows_a_later_ui_pick(self):
+ # The studio's model picker writes the pick into OpenCode's own config
+ # (PUT /config). A backend that kept the model it resolved on its first
+ # turn went on answering with the old one for the rest of the run, so a
+ # NON-explicit model re-checks /config every turn and follows it.
chat = OpenCodeChat("http://127.0.0.1:4096", model="openrouter/openai/gpt-5", model_is_explicit=False)
+ with mock.patch.object(
+ chat, "_request",
+ return_value=self._fake_response({"model": "opencode/big-pickle"}),
+ ) as request_mock:
+ chat._ensure_model_resolved()
+ request_mock.assert_called_once_with("/config")
+ self.assertEqual(chat.model, "opencode/big-pickle")
+
+ def test_explicit_model_is_never_overridden_by_the_config(self):
+ # OPENCODE_MODEL / agent_config.yaml's opencode_model PIN the model:
+ # no request, and the UI picker cannot move it.
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="openrouter/openai/gpt-5", model_is_explicit=True)
with mock.patch.object(chat, "_request") as request_mock:
chat._ensure_model_resolved()
request_mock.assert_not_called()
self.assertEqual(chat.model, "openrouter/openai/gpt-5")
+ def test_a_failed_config_read_keeps_the_model_already_resolved(self):
+ # One local request failing must not drop a working model for the
+ # fallback mid-run.
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="anthropic/claude-haiku-4.5", model_is_explicit=False)
+ with mock.patch.object(chat, "_request", side_effect=OSError("refused")):
+ chat._ensure_model_resolved()
+ self.assertEqual(chat.model, "anthropic/claude-haiku-4.5")
+
def test_unset_model_resolves_from_config_own_model_field(self):
chat = OpenCodeChat("http://127.0.0.1:4096", model=None, model_is_explicit=False)
with mock.patch.object(chat, "_request", return_value=self._fake_response({"model": "anthropic/claude-haiku-4.5"})) as request_mock:
@@ -425,13 +447,13 @@ def test_falls_back_to_the_hardcoded_default_when_config_has_no_model(self):
with mock.patch.object(chat, "_request", return_value=self._fake_response({})) as request_mock:
chat._ensure_model_resolved()
request_mock.assert_called_once_with("/config")
- self.assertEqual(chat.model, "opencode/deepseek-v4-flash-free")
+ self.assertEqual(chat.model, "opencode/big-pickle")
def test_no_resolvable_model_falls_back_to_the_hardcoded_default(self):
chat = OpenCodeChat("http://127.0.0.1:4096", model=None, model_is_explicit=False)
with mock.patch.object(chat, "_request", side_effect=OSError("refused")):
chat._ensure_model_resolved() # must not raise
- self.assertEqual(chat.model, "opencode/deepseek-v4-flash-free")
+ self.assertEqual(chat.model, "opencode/big-pickle")
def test_config_field_that_is_not_provider_slash_model_falls_through(self):
# A malformed/unexpected `model` field (missing the "/") is treated
@@ -441,7 +463,155 @@ def test_config_field_that_is_not_provider_slash_model_falls_through(self):
with mock.patch.object(chat, "_request", return_value=self._fake_response({"model": "not-a-provider-model-pair"})) as request_mock:
chat._ensure_model_resolved()
request_mock.assert_called_once_with("/config")
- self.assertEqual(chat.model, "opencode/deepseek-v4-flash-free")
+ self.assertEqual(chat.model, "opencode/big-pickle")
+
+
+
+class PublishModelTests(unittest.TestCase):
+ """publish_model / resolve_model -- how the two sides converge on one model.
+
+ Verified against a live server before these were written: PATCH /config
+ (workspace scope) answers 200 and echoes the value back but does NOT change
+ what GET /config reports, while PATCH /global/config does. GET /config is
+ what both the studio picker and this backend read, so the write has to go
+ to the global scope, and be CONFIRMED rather than trusted.
+ """
+
+ @staticmethod
+ def _resp(payload):
+ return mock.MagicMock(
+ __enter__=mock.MagicMock(
+ return_value=mock.MagicMock(read=lambda: json.dumps(payload).encode())),
+ __exit__=mock.MagicMock(return_value=False),
+ )
+
+ def test_publish_writes_the_global_scope_first(self):
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="opencode/big-pickle",
+ model_is_explicit=False)
+ calls = []
+
+ def fake_request(path, method="GET", body=None, headers=None):
+ calls.append((path, method, body))
+ if method == "GET":
+ return self._resp({"model": "opencode/big-pickle"})
+ return self._resp({})
+
+ with mock.patch.object(chat, "_request", side_effect=fake_request):
+ self.assertTrue(chat.publish_model())
+ self.assertEqual(calls[0], ("/global/config", "PATCH", {"model": "opencode/big-pickle"}))
+ # Confirmed by reading the effective config back, not by the echo.
+ self.assertIn(("/config", "GET", None), calls)
+
+ def test_publish_falls_back_to_the_workspace_route(self):
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="opencode/big-pickle",
+ model_is_explicit=False)
+ seen = []
+
+ def fake_request(path, method="GET", body=None, headers=None):
+ seen.append((path, method))
+ if path == "/global/config":
+ raise OSError("no such route")
+ if method == "GET":
+ return self._resp({"model": "opencode/big-pickle"})
+ return self._resp({})
+
+ with mock.patch.object(chat, "_request", side_effect=fake_request):
+ self.assertTrue(chat.publish_model())
+ self.assertIn(("/config", "PATCH"), seen)
+
+ def test_publish_reports_failure_when_the_write_does_not_stick(self):
+ # Exactly the live failure mode: 200 back, effective config unchanged.
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="opencode/big-pickle",
+ model_is_explicit=False)
+
+ def fake_request(path, method="GET", body=None, headers=None):
+ if method == "GET":
+ return self._resp({}) # no model -- nothing was stored
+ return self._resp({"model": "opencode/big-pickle"}) # echo only
+
+ with mock.patch.object(chat, "_request", side_effect=fake_request):
+ self.assertFalse(chat.publish_model())
+
+ def test_resolve_publishes_the_fallback_so_the_studio_adopts_it(self):
+ # weightslab started FIRST: nothing chosen anywhere, so the built-in
+ # default is resolved AND published for the UI to read.
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="", model_is_explicit=False)
+ with mock.patch.object(chat, "_configured_model", side_effect=[None, "opencode/big-pickle"]), mock.patch.object(chat, "_request", return_value=self._resp({})), mock.patch.object(chat, "_ensure_reachable"):
+ model, source = chat.resolve_model(publish_default=True)
+ self.assertEqual(model, "opencode/big-pickle")
+ self.assertEqual(source, "default-published")
+
+ def test_resolve_adopts_the_studio_pick_without_publishing(self):
+ # studio started FIRST: its pick is already in OpenCode's config, so
+ # follow it and write nothing.
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="", model_is_explicit=False)
+ with mock.patch.object(chat, "_configured_model", return_value="openrouter/openai/gpt-5-mini"), mock.patch.object(chat, "publish_model") as publish, mock.patch.object(chat, "_ensure_reachable"):
+ model, source = chat.resolve_model(publish_default=True)
+ self.assertEqual((model, source), ("openrouter/openai/gpt-5-mini", "opencode-config"))
+ publish.assert_not_called()
+
+ def test_the_studio_pick_wins_over_a_yaml_seed(self):
+ # The reported bug: pick a model in the UI, then start a run --
+ # agent_config.yaml's opencode_model pinned the backend back to its own
+ # value, ignoring the pick. A yaml model is a SEED, not a pin.
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="", model_is_explicit=False,
+ seed_model="opencode/muse-spark-1.3-contributor-free")
+ with mock.patch.object(chat, "_configured_model", return_value="openrouter/openai/gpt-5-mini"), mock.patch.object(chat, "publish_model") as publish, mock.patch.object(chat, "_ensure_reachable"):
+ model, source = chat.resolve_model(publish_default=True)
+ self.assertEqual((model, source), ("openrouter/openai/gpt-5-mini", "opencode-config"))
+ publish.assert_not_called()
+
+ def test_the_yaml_seed_is_used_and_published_when_nothing_is_chosen(self):
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="", model_is_explicit=False,
+ seed_model="opencode/muse-spark-1.3-contributor-free")
+ with mock.patch.object(chat, "_configured_model", return_value=None), mock.patch.object(chat, "publish_model", return_value=True) as publish, mock.patch.object(chat, "_ensure_reachable"):
+ model, source = chat.resolve_model(publish_default=True)
+ self.assertEqual(
+ (model, source),
+ ("opencode/muse-spark-1.3-contributor-free", "config-seed-published"))
+ publish.assert_called_once()
+
+ def test_the_yaml_seed_beats_the_builtin_default(self):
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="", model_is_explicit=False,
+ seed_model="openrouter/openai/gpt-5-mini")
+ with mock.patch.object(chat, "_configured_model", return_value=None), mock.patch.object(chat, "publish_model", return_value=False), mock.patch.object(chat, "_ensure_reachable"):
+ model, _ = chat.resolve_model(publish_default=True)
+ self.assertEqual(model, "openrouter/openai/gpt-5-mini")
+
+ def test_the_env_pin_still_beats_the_studio_pick(self):
+ # OPENCODE_MODEL is per-process and deliberate: automation must be able
+ # to force a model regardless of what anyone picked in the UI.
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="openrouter/pinned/model",
+ model_is_explicit=True, seed_model="opencode/seed")
+ with mock.patch.object(chat, "_configured_model", return_value="openrouter/openai/gpt-5-mini"), mock.patch.object(chat, "publish_model", return_value=True), mock.patch.object(chat, "_ensure_reachable"):
+ model, source = chat.resolve_model(publish_default=True)
+ self.assertEqual((model, source), ("openrouter/pinned/model", "pinned-published"))
+
+ def test_resolve_publishes_a_pinned_model_so_the_studio_shows_it(self):
+ # A model pinned in agent_config.yaml / OPENCODE_MODEL is a deliberate
+ # choice too. It used to stay invisible to the studio, which then
+ # displayed an unrelated default while every backend query ran on the
+ # pinned model.
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="openrouter/pinned/model",
+ model_is_explicit=True)
+ with mock.patch.object(chat, "publish_model", return_value=True) as publish, mock.patch.object(chat, "_ensure_reachable"):
+ model, source = chat.resolve_model(publish_default=True)
+ self.assertEqual((model, source), ("openrouter/pinned/model", "pinned-published"))
+ publish.assert_called_once()
+
+ def test_a_pinned_model_stays_pinned_when_it_cannot_be_published(self):
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="openrouter/pinned/model",
+ model_is_explicit=True)
+ with mock.patch.object(chat, "publish_model", return_value=False), mock.patch.object(chat, "_ensure_reachable"):
+ model, source = chat.resolve_model(publish_default=True)
+ self.assertEqual((model, source), ("openrouter/pinned/model", "pinned"))
+
+ def test_a_model_read_from_the_config_is_never_republished(self):
+ chat = OpenCodeChat("http://127.0.0.1:4096", model="", model_is_explicit=False)
+ with mock.patch.object(chat, "_configured_model", return_value="opencode/big-pickle"), mock.patch.object(chat, "publish_model") as publish, mock.patch.object(chat, "_ensure_reachable"):
+ model, source = chat.resolve_model(publish_default=True)
+ self.assertEqual((model, source), ("opencode/big-pickle", "opencode-config"))
+ publish.assert_not_called()
if __name__ == "__main__":
diff --git a/tests/ui/test_server_experiment_reports.py b/tests/ui/test_server_experiment_reports.py
index bd0c0e7a..b1cf36ee 100644
--- a/tests/ui/test_server_experiment_reports.py
+++ b/tests/ui/test_server_experiment_reports.py
@@ -17,6 +17,7 @@
import json
import os
+import shutil
import tempfile
import threading
import time
@@ -25,6 +26,7 @@
import urllib.request
from weightslab.ui import server as ui_server
+from weightslab.utils import active_experiment
class _ServerTestCase(unittest.TestCase):
@@ -33,6 +35,14 @@ class _ServerTestCase(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.mkdtemp()
+ # The active-experiment marker is a real per-user file, and a LIVE
+ # backend recorded in it redirects these listings on purpose (see
+ # _experiment_dir_path). Point it at a scratch directory so the tests
+ # exercise the explicit experiment_dir below instead of whatever run
+ # the developer happens to have going.
+ self._state_prev = os.environ.get("WEIGHTSLAB_STATE_DIR")
+ self._state_dir = tempfile.mkdtemp()
+ os.environ["WEIGHTSLAB_STATE_DIR"] = self._state_dir
self.httpd = ui_server.serve_ui(
ui_host="127.0.0.1", ui_port=0,
backend_host="localhost", backend_port=50051,
@@ -47,6 +57,11 @@ def setUp(self):
def tearDown(self):
self.httpd.shutdown()
self.thread.join(timeout=5)
+ if self._state_prev is None:
+ os.environ.pop("WEIGHTSLAB_STATE_DIR", None)
+ else:
+ os.environ["WEIGHTSLAB_STATE_DIR"] = self._state_prev
+ shutil.rmtree(self._state_dir, ignore_errors=True)
def _get(self, path):
return urllib.request.urlopen(f"http://127.0.0.1:{self.port}{path}", timeout=5)
@@ -96,6 +111,49 @@ def test_entries_include_name_path_and_modified_at(self):
self.assertIsInstance(entry["modified_at"], (int, float))
+class TestReportsFollowTheRunningBackend(_ServerTestCase):
+ """The listing follows where the RUNNING backend actually writes.
+
+ Regression: `weightslab start` and a training run started in another
+ terminal resolved different directories (the run fell through to %TEMP%),
+ so right-clicking "Generate report" listed nothing even though reports had
+ just been generated.
+ """
+
+ def _write_report_in(self, directory, name):
+ reports_dir = os.path.join(directory, "reports")
+ os.makedirs(reports_dir, exist_ok=True)
+ with open(os.path.join(reports_dir, name), "w", encoding="utf-8") as f:
+ f.write("")
+
+ def test_a_live_backend_directory_is_listed_instead_of_the_uis_own(self):
+ backend_dir = tempfile.mkdtemp()
+ self.addCleanup(shutil.rmtree, backend_dir, ignore_errors=True)
+ self._write_report_in(backend_dir, "experiment_report_20260909_193124.html")
+ # This test process stands in for the running backend.
+ active_experiment.record_backend_experiment(backend_dir)
+
+ with self._get("/experiment-report/list") as r:
+ data = json.loads(r.read().decode())
+ self.assertEqual([e["name"] for e in data["reports"]],
+ ["experiment_report_20260909_193124.html"])
+
+ def test_a_finished_backend_does_not_hijack_the_listing(self):
+ backend_dir = tempfile.mkdtemp()
+ self.addCleanup(shutil.rmtree, backend_dir, ignore_errors=True)
+ self._write_report_in(backend_dir, "stale.html")
+ active_experiment.record_backend_experiment(backend_dir)
+ # Rewrite the record with a pid that cannot be running.
+ state = active_experiment.read_state()
+ state["backend"][-1]["pid"] = 2 ** 31 - 1
+ active_experiment.state_path().write_text(json.dumps(state), encoding="utf-8")
+
+ self._write_report("mine.html")
+ with self._get("/experiment-report/list") as r:
+ data = json.loads(r.read().decode())
+ self.assertEqual([e["name"] for e in data["reports"]], ["mine.html"])
+
+
class TestServeExperimentReport(_ServerTestCase):
def test_serves_html_content_with_correct_content_type(self):
diff --git a/tests/ui/test_server_shared_model.py b/tests/ui/test_server_shared_model.py
new file mode 100644
index 00000000..58ffa1f1
--- /dev/null
+++ b/tests/ui/test_server_shared_model.py
@@ -0,0 +1,184 @@
+"""Tests for the UI server's same-origin shared-model endpoints:
+
+- GET /agent-server/model -- the model OpenCode's config names, or null.
+- POST /agent-server/model -- set it (global scope, confirmed by reading back).
+
+Why they exist: the browser CAN call OpenCode directly, but only while
+OpenCode's ``--cors`` allowlist contains the page's exact origin. A LAN
+address, a tunnel hostname, or an ``opencode serve`` started by hand with no
+``--cors`` all make that cross-origin PATCH fail its preflight, so a model
+picked in the studio silently never reached OpenCode -- and therefore never
+reached weightslab's backend, which reads that same field to choose the model
+for its own queries. Proxying through this server is same-origin: no
+preflight, no allowlist.
+
+A fake OpenCode stands in for the real server, reproducing the two behaviours
+confirmed live: PATCH /global/config sticks, PATCH /config answers 200 and
+echoes the value back without changing what GET /config reports.
+"""
+
+import json
+import threading
+import time
+import unittest
+import unittest.mock
+import urllib.error
+import urllib.request
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+from weightslab.ui import server as ui_server
+
+
+class _FakeOpencode(BaseHTTPRequestHandler):
+ """model lives in the class so every request sees the same value."""
+
+ model = None
+ workspace_patch_sticks = False
+ reachable = True
+
+ def log_message(self, *args): # silence
+ pass
+
+ def _json(self, status, payload):
+ body = json.dumps(payload).encode()
+ self.send_response(status)
+ self.send_header("Content-Type", "application/json")
+ self.send_header("Content-Length", str(len(body)))
+ self.end_headers()
+ self.wfile.write(body)
+
+ def do_GET(self):
+ if not type(self).reachable:
+ self._json(500, {"error": "down"})
+ return
+ if self.path == "/config":
+ self._json(200, {"model": type(self).model} if type(self).model else {})
+ return
+ self._json(404, {})
+
+ def do_PATCH(self):
+ length = int(self.headers.get("Content-Length", "0") or 0)
+ body = json.loads(self.rfile.read(length).decode() or "{}") if length else {}
+ model = body.get("model")
+ if self.path == "/global/config":
+ type(self).model = model
+ self._json(200, {"model": model})
+ return
+ if self.path == "/config":
+ # Answers 200 and echoes the value, but only *stores* it when the
+ # server actually honours workspace scope -- the live one does not.
+ if type(self).workspace_patch_sticks:
+ type(self).model = model
+ self._json(200, {"model": model})
+ return
+ self._json(404, {})
+
+
+class TestSharedModelEndpoints(unittest.TestCase):
+ def setUp(self):
+ _FakeOpencode.model = None
+ _FakeOpencode.workspace_patch_sticks = False
+ _FakeOpencode.reachable = True
+
+ self.oc = ThreadingHTTPServer(("127.0.0.1", 0), _FakeOpencode)
+ self.oc_thread = threading.Thread(target=self.oc.serve_forever, daemon=True)
+ self.oc_thread.start()
+ oc_url = f"http://127.0.0.1:{self.oc.server_address[1]}"
+
+ # The UI server resolves OpenCode from the running session, then from
+ # OPENCODE_URL -- point it at the fake.
+ self._prev_url = ui_server.os.environ.get("OPENCODE_URL")
+ ui_server.os.environ["OPENCODE_URL"] = oc_url
+
+ self.httpd = ui_server.serve_ui(
+ ui_host="127.0.0.1", ui_port=0,
+ backend_host="localhost", backend_port=50051,
+ open_browser=False, block=False,
+ )
+ self.port = self.httpd.server_address[1]
+ self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True)
+ self.thread.start()
+ time.sleep(0.1)
+
+ def tearDown(self):
+ self.httpd.shutdown()
+ self.thread.join(timeout=5)
+ self.oc.shutdown()
+ self.oc_thread.join(timeout=5)
+ if self._prev_url is None:
+ ui_server.os.environ.pop("OPENCODE_URL", None)
+ else:
+ ui_server.os.environ["OPENCODE_URL"] = self._prev_url
+
+ # -- helpers ---------------------------------------------------------
+ def _get(self):
+ with urllib.request.urlopen(
+ f"http://127.0.0.1:{self.port}/agent-server/model", timeout=5) as r:
+ return json.loads(r.read().decode())
+
+ def _post(self, model):
+ req = urllib.request.Request(
+ f"http://127.0.0.1:{self.port}/agent-server/model", method="POST",
+ data=json.dumps({"model": model}).encode(),
+ headers={"Content-Type": "application/json"})
+ try:
+ with urllib.request.urlopen(req, timeout=5) as r:
+ return r.status, json.loads(r.read().decode())
+ except urllib.error.HTTPError as exc:
+ return exc.code, json.loads(exc.read().decode() or "{}")
+
+ # -- tests -----------------------------------------------------------
+ def test_get_reports_nothing_when_no_model_is_configured(self):
+ self.assertEqual(self._get(), {"ok": True, "model": None})
+
+ def test_get_reports_the_configured_model(self):
+ _FakeOpencode.model = "opencode/big-pickle"
+ self.assertEqual(self._get(), {"ok": True, "model": "opencode/big-pickle"})
+
+ def test_get_ignores_a_malformed_model_field(self):
+ _FakeOpencode.model = "not-a-provider-model-pair"
+ self.assertEqual(self._get(), {"ok": True, "model": None})
+
+ def test_post_writes_the_global_scope_and_confirms_it(self):
+ status, payload = self._post("openrouter/openai/gpt-5-mini")
+ self.assertEqual(status, 200)
+ self.assertTrue(payload["ok"])
+ self.assertEqual(payload["model"], "openrouter/openai/gpt-5-mini")
+ self.assertEqual(payload["via"], "/global/config")
+ # ...and it really is what OpenCode now reports.
+ self.assertEqual(_FakeOpencode.model, "openrouter/openai/gpt-5-mini")
+ self.assertEqual(self._get()["model"], "openrouter/openai/gpt-5-mini")
+
+ def test_post_rejects_a_model_without_a_provider(self):
+ status, payload = self._post("big-pickle")
+ self.assertEqual(status, 400)
+ self.assertFalse(payload["ok"])
+
+ def test_post_reports_failure_when_the_write_does_not_stick(self):
+ # Both routes answer 200 but nothing is stored -- exactly the live
+ # workspace-scope behaviour, generalised.
+ class _EchoOnly(_FakeOpencode):
+ pass
+
+ def do_PATCH(self): # noqa: N802 -- HTTP handler naming
+ length = int(self.headers.get("Content-Length", "0") or 0)
+ if length:
+ self.rfile.read(length)
+ self._json(200, {"model": "whatever"})
+
+ with unittest.mock.patch.object(_FakeOpencode, "do_PATCH", do_PATCH):
+ status, payload = self._post("opencode/big-pickle")
+ self.assertEqual(status, 200)
+ self.assertFalse(payload["ok"])
+ self.assertIn("did not accept", payload["error"])
+
+ def test_get_says_so_when_opencode_is_unreachable(self):
+ _FakeOpencode.reachable = False
+ payload = self._get()
+ self.assertFalse(payload["ok"])
+ self.assertIsNone(payload["model"])
+ self.assertIn("not reachable", payload["error"])
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/weightslab/backend/cli.py b/weightslab/backend/cli.py
index c3d8a853..11ccbeb7 100644
--- a/weightslab/backend/cli.py
+++ b/weightslab/backend/cli.py
@@ -801,12 +801,21 @@ def _handle_command(cmd: str) -> Any:
available = bool(agent.is_available())
except Exception:
available = False
+ # current_model() re-reads OpenCode's shared config, so a model
+ # picked in the studio after this backend started is reported
+ # here instead of the model resolved at start-up.
+ model = getattr(agent, 'opencode_model', None)
+ try:
+ if hasattr(agent, 'current_model'):
+ model = agent.current_model() or model
+ except Exception:
+ pass
return {
'ok': True,
'available': available,
'preferred_provider': getattr(agent, 'preferred_provider', None),
'opencode_url': getattr(agent, 'opencode_url', None),
- 'opencode_model': getattr(agent, 'opencode_model', None),
+ 'opencode_model': model,
'message': 'Agent available. Ready to help you.' if available else 'Agent not configured. Use agent init.',
}
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/cli.py b/weightslab/cli.py
index d4e5adb4..65a592e1 100644
--- a/weightslab/cli.py
+++ b/weightslab/cli.py
@@ -623,6 +623,31 @@ def example_start(args):
try:
env = os.environ.copy()
env['WEIGHTSLAB_SUPPRESS_BANNER'] = '1'
+ # `weightslab start` runs in its own terminal, so its
+ # WEIGHTSLAB_ROOT_LOG_DIR export never reaches this process -- read the
+ # directory it recorded and hand it to the example, so the run lands in
+ # the experiment the UI is showing instead of a throwaway temp dir.
+ # Anything already set in this shell wins: an explicit choice by the
+ # user must not be overridden by the last UI launch.
+ if not (env.get('WEIGHTSLAB_ROOT_LOG_DIR') or '').strip():
+ try:
+ from weightslab.utils.active_experiment import live_ui_experiment_dir
+ adopted = live_ui_experiment_dir()
+ except Exception as exc: # noqa: BLE001
+ logger.debug(f"Could not read the active experiment directory: {exc}")
+ adopted = None
+ if adopted:
+ env['WEIGHTSLAB_ROOT_LOG_DIR'] = adopted
+ logger.info(f" Using the experiment directory from `weightslab start`: {adopted}")
+ else:
+ logger.warning(
+ " No experiment directory found (no WEIGHTSLAB_ROOT_LOG_DIR here and no "
+ "`weightslab start` on record) — this example will write to a temporary "
+ "directory, and the UI will not find its reports or notebook. Start the UI "
+ "first with `weightslab start`, or set WEIGHTSLAB_ROOT_LOG_DIR."
+ )
+ else:
+ logger.info(f" Using WEIGHTSLAB_ROOT_LOG_DIR from this shell: {env['WEIGHTSLAB_ROOT_LOG_DIR']}")
result = subprocess.run([sys.executable, str(main_py)], cwd=str(example_dir), env=env)
except KeyboardInterrupt:
logger.info("Example stopped.")
@@ -892,6 +917,18 @@ def ui_start_native(args):
experiment_dir = _resolve_experiment_dir(getattr(args, "experiment_dir", None))
os.environ["WEIGHTSLAB_ROOT_LOG_DIR"] = str(experiment_dir)
os.environ["WL_LAST_EXPERIMENT_DIR"] = str(experiment_dir)
+ # An environment variable reaches only THIS process and its children. A
+ # training run started from another terminal is a different process tree,
+ # so also record the directory in the marker file every later weightslab
+ # process reads (weightslab.utils.active_experiment) -- without it, such a
+ # run fell through to a throwaway %TEMP% directory while this UI listed an
+ # empty reports/ from the directory established here.
+ try:
+ from weightslab.utils.active_experiment import record_ui_experiment
+ record_ui_experiment(experiment_dir)
+ except Exception as exc: # noqa: BLE001 -- advisory record, never fatal
+ logger.debug(f"Could not record the active experiment directory: {exc}")
+ # Re-recorded below with the resolved ports, once they are known.
_print_experiment_guidance(experiment_dir)
# If the agent has been initialized, provision OpenCode up front (in the
@@ -955,6 +992,15 @@ def ui_start_native(args):
ui_port = actual_port
logger.info(f"UI port source: {ui_port_source} (preferred {preferred_ui_port}, using {ui_port})")
os.environ["WL_LAST_UI_PORT"] = str(ui_port)
+ # Now that the ports are settled, stamp them on this UI's record: the
+ # backend port is what lets this UI ask for the experiment directory of
+ # ITS backend rather than of whichever backend started last -- the
+ # difference that matters when two experiments run side by side.
+ try:
+ from weightslab.utils.active_experiment import record_ui_experiment
+ record_ui_experiment(experiment_dir, ui_port=ui_port, backend_port=backend_port)
+ except Exception as exc: # noqa: BLE001
+ logger.debug(f"Could not record the UI ports: {exc}")
ui_server.serve_ui(
ui_host=ui_host,
diff --git a/weightslab/data/dataframe_manager.py b/weightslab/data/dataframe_manager.py
index 76b9a8f7..59afaf2c 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
@@ -800,12 +910,28 @@ def upsert_df(self, df_local: List | pd.DataFrame, origin: str = None, force_flu
if df_norm[col].dtype == bool:
self._df[col] = self._df[col].fillna(False).astype(bool)
+ # Columns SampleStatsEx documents a boolean default for must never
+ # hold NaN ("None are not accepted by PD H5 storage" -- see its
+ # DEFAULTS). Two ways they did anyway, and the loop above catches
+ # neither: it only visits columns being ADDED by this upsert, and
+ # only when the incoming slice's dtype is already bool -- which it
+ # is not precisely when the slice carries missing values. So a
+ # sample registered without the flag, and every per-annotation row
+ # (sample-level values live on annotation 0 only), kept a NaN.
+ #
+ # That is not cosmetic: bool(float("nan")) is True in Python, so a
+ # NaN flag read as a set one -- samples served to the studio as
+ # discarded while the dataframe said nothing was.
+ self._fill_documented_flag_defaults()
+
# Auto-register any string-valued tag: columns as categorical tags
# (e.g. dataset metadata declaring tag:weather = "rainy"/"sunny").
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 +940,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 +951,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 +962,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 +1467,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 +1711,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 +2103,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 +2143,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 +2215,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 +2264,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 +2340,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 +2367,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 +2423,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 +2440,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 +2462,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
@@ -2319,6 +2598,44 @@ def get_combined_df(
return df
+ def _fill_documented_flag_defaults(self) -> None:
+ """Give every boolean column SampleStatsEx documents a default its default.
+
+ Only the columns with a ``bool`` default in
+ ``SampleStatsEx.DEFAULTS`` (today: ``discarded``) -- a ``tag:*`` column
+ is deliberately left sparse, where NaN and False mean the same thing
+ and NaN costs nothing to store, and a *categorical* tag's NaN means
+ "unset", which is not a default at all.
+
+ Cheap by design: an ``isna().any()`` short-circuit per flag column, so
+ the common case (nothing missing) touches no rows.
+ """
+ if self._df is None or self._df.empty:
+ return
+ # DEFAULTS lives on SampleStats, the outer class -- SampleStatsEx is
+ # only its `Ex` enum, so reading it off that is a silent no-op.
+ from weightslab.data.sample_stats import SampleStats
+ defaults = getattr(SampleStats, "DEFAULTS", {}) or {}
+ for col, default in defaults.items():
+ if not isinstance(default, bool) or col not in self._df.columns:
+ continue
+ try:
+ series = self._df[col]
+ if not series.isna().any():
+ continue
+ if isinstance(series.dtype, pd.CategoricalDtype):
+ # fillna on a Categorical raises unless the value is a
+ # known category; widen first, then let the memory pass
+ # re-apply the dtype.
+ series = series.astype(object)
+ self._df[col] = series.fillna(default)
+ logger.debug(
+ "[LedgeredDataFrameManager] filled missing %r with its "
+ "documented default %r", col, default)
+ except Exception as exc: # noqa: BLE001 -- never fail an upsert on this
+ logger.debug(
+ "[LedgeredDataFrameManager] could not default %r: %s", col, exc)
+
def get_collapse_annotations_to_samples_df(self, df: pd.DataFrame | None = None) -> pd.DataFrame:
"""Collapse a (sample_id, annotation_id) multi-index df to one row per sample.
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 88067e30..bf3c3e3e 100644
--- a/weightslab/src.py
+++ b/weightslab/src.py
@@ -83,7 +83,16 @@ def _resolve_configured_root_log_dir(configured):
actually points at an existing directory; if it's set but stale/typo'd,
a warning is logged and resolution falls through to (3) instead of
silently training into a directory the UI never established.
- 3. A throwaway ``tempfile.mkdtemp()`` — last resort so serving never fails
+ 3. The directory a RUNNING ``weightslab start`` established, read from
+ the marker file (see weightslab.utils.active_experiment). Only while
+ that UI is alive: a record left by one that has since exited must not
+ redirect an unrelated run. The
+ environment variable only reaches processes started FROM that same
+ shell; a training run launched in another terminal (or by
+ ``weightslab start example``) is a different process tree and used to
+ fall straight through to (4), landing in %TEMP% while the UI listed an
+ empty reports/ from the directory it had established.
+ 4. A throwaway ``tempfile.mkdtemp()`` — last resort so serving never fails
for lack of a directory.
"""
if configured:
@@ -94,9 +103,28 @@ def _resolve_configured_root_log_dir(configured):
return env_dir
logger.warning(
f"WEIGHTSLAB_ROOT_LOG_DIR is set to '{env_dir}', but that directory "
- "does not exist. Falling back to a temporary directory instead."
+ "does not exist. Falling back to the experiment directory recorded "
+ "by `weightslab start`, then to a temporary directory."
)
- return tempfile.mkdtemp()
+ try:
+ from weightslab.utils.active_experiment import live_ui_experiment_dir
+ marker_dir = live_ui_experiment_dir()
+ except Exception: # noqa: BLE001 -- never block serving on the marker
+ marker_dir = None
+ if marker_dir:
+ logger.info(
+ "Using the experiment directory established by `weightslab start`: "
+ "%s (no root_log_dir configured and WEIGHTSLAB_ROOT_LOG_DIR is not "
+ "set in this process).", marker_dir)
+ return marker_dir
+ tmp_dir = tempfile.mkdtemp()
+ logger.warning(
+ "No root_log_dir configured, no WEIGHTSLAB_ROOT_LOG_DIR, and no "
+ "experiment directory recorded by `weightslab start` — this run will "
+ "write to the throwaway directory %s. Checkpoints, reports and the "
+ "notebook will NOT be where the UI looks for them. Start the UI first "
+ "(`weightslab start`), or set root_log_dir in your config.", tmp_dir)
+ return tmp_dir
# Get global dataframe proxy (auto-updated when ledger registers real manager)
@@ -494,6 +522,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 +1061,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'),
@@ -1446,6 +1498,14 @@ def new_forward(*a, **kw):
# _resolve_configured_root_log_dir for the resolution order).
_hp_cfg['root_log_dir'] = _resolve_configured_root_log_dir(
_hp_cfg.get('root_log_dir'))
+ # Publish where training ACTUALLY writes, so the UI lists this
+ # run's reports/notebooks even when the two resolved their
+ # directory by different routes.
+ try:
+ from weightslab.utils.active_experiment import record_backend_experiment
+ record_backend_experiment(_hp_cfg['root_log_dir'])
+ except Exception as _exc: # noqa: BLE001 -- advisory only
+ logger.debug("Could not record the backend experiment dir: %s", _exc)
try:
# Check if a checkpoint manager is already registered in ledger
try:
@@ -1738,6 +1798,25 @@ def serve(serving_cli: bool = True, serving_grpc: bool = True,
_notebook_service.configure_embedded_kernel(embed_kernel_decision)
if serving_grpc:
+ # Stamp the gRPC port on this backend's active-experiment record, so a
+ # UI proxying to that port finds THIS experiment's directory rather
+ # than whichever backend happened to start last (see
+ # active_experiment._sole_live_dir). Advisory: never fail serving on it.
+ try:
+ from weightslab.utils.active_experiment import record_backend_experiment
+ _grpc_port = (kwargs.get("grpc_port")
+ or int(os.getenv("GRPC_BACKEND_PORT", 50051)))
+ _root_log_dir = None
+ try:
+ _hp = ledgers.get_hyperparams()
+ _root_log_dir = _hp["root_log_dir"] if _hp is not None else None
+ except Exception: # noqa: BLE001 -- no hyperparameters registered
+ _root_log_dir = None
+ if _root_log_dir:
+ record_backend_experiment(_root_log_dir, grpc_port=int(_grpc_port))
+ except Exception as _exc: # noqa: BLE001
+ logger.debug("Could not record the backend gRPC port: %s", _exc)
+
grpc_serve(**kwargs)
if embed_kernel_decision:
@@ -4935,6 +5014,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.
@@ -4953,17 +5042,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/agent/agent.py b/weightslab/trainer/services/agent/agent.py
index 7158d484..b1701b94 100644
--- a/weightslab/trainer/services/agent/agent.py
+++ b/weightslab/trainer/services/agent/agent.py
@@ -735,6 +735,12 @@ def _load_config(self):
# actually chose is exempt from that self-healing.
self._opencode_model_explicit = "OPENCODE_MODEL" in os.environ
self.opencode_model = os.environ.get("OPENCODE_MODEL", "")
+ # agent_config.yaml's opencode_model lands here instead of pinning the
+ # model: it SEEDS the shared choice (used when nothing has been chosen
+ # yet, then published so the studio shows it), while a model picked in
+ # the UI afterwards wins. Pinning it meant a run started after picking
+ # a model in the studio silently went back to the yaml value.
+ self._opencode_model_seed = ""
# The same directory `weightslab start ` roots the browser
# landing-page agent at (WEIGHTSLAB_ROOT_LOG_DIR) -- the shared key
# opencode_process.py's lock file is discovered/published under, so
@@ -759,6 +765,9 @@ def _load_config(self):
inner_pkg / "agent_config.yaml",
Path.cwd() / "agent_config.yaml"
]
+ # Overwritten below only when a file is actually applied, so the banner
+ # can say "none found" instead of naming the last candidate it tried.
+ self._config_source_path = "(no agent config file found)"
for path in config_paths:
if not path.exists(): continue
try:
@@ -771,28 +780,36 @@ def _load_config(self):
if a_cfg.get("opencode_url"):
self._opencode_url_explicit = True
self.opencode_url = a_cfg.get("opencode_url", self.opencode_url)
- if a_cfg.get("opencode_model"):
- self._opencode_model_explicit = True
- self.opencode_model = a_cfg.get("opencode_model", self.opencode_model)
+ _cfg_model = str(a_cfg.get("opencode_model") or "").strip()
+ if _cfg_model:
+ # Seed, not pin -- see _opencode_model_seed above.
+ # OPENCODE_MODEL still wins: it is set per process, on purpose.
+ self._opencode_model_seed = _cfg_model
+ self._config_source_path = path
_LOGGER.info(f"Applied agent configuration from {path}")
_LOGGER.debug(f"Agent Config: {cfg}")
break
except Exception as e:
_LOGGER.warning(f"Error loading config from {path}: {e}")
- # Log the final configuration for transparency
- _LOGGER.info(
- "" + "\n" +
- "\n# #######################################" + "\n" +
- "# #######################################" + "\n" +
- f"Agent initialized from configuration {path}: " + "\n" +
- f"\tOpenCode URL={self.opencode_url}, Model={self.opencode_model or '(server default)'}" + "\n" +
- "# #######################################" + "\n" +
- "# #######################################" + "\n" + ""
- )
-
- def _setup_providers(self):
+ # The banner itself is emitted by _log_agent_configuration() AFTER
+ # _setup_providers has resolved the model, so it names the model
+ # actually in use instead of the pre-resolution blank -- which printed
+ # "(server default)" and read as "the studio's pick was ignored".
+
+ def _setup_providers(self, requested_model: Optional[str] = None):
+ """(Re)build the OpenCode provider.
+
+ `requested_model` is a model the USER just chose (CLI `agent model`,
+ `agent init --model`, the SetAgentModel RPC). It must survive this
+ call: the shared-config re-read below exists to follow the studio's
+ picker, and it used to overwrite the very model the caller had just
+ asked for -- `agent model X` answered "Model switched to ". A user choice is instead PUBLISHED to OpenCode's config, so it
+ becomes the shared choice the studio picker and every other client see
+ too, and is only pinned in-process when that write is refused.
+ """
self.chain_opencode = None
self._opencode_chat = None
initialized = False
@@ -806,18 +823,85 @@ def _setup_providers(self):
workspace_dir=self.opencode_workspace_dir,
url_is_explicit=self._opencode_url_explicit,
model_is_explicit=self._opencode_model_explicit,
+ seed_model=getattr(self, "_opencode_model_seed", ""),
)
self.chain_opencode = self._opencode_chat.as_runnable()
initialized = True
+ # Resolve up front rather than on the first query: the model is
+ # part of what the start-up banner reports, and a backend started
+ # before the studio publishes its fallback so the UI adopts the
+ # same model (see OpenCodeChat.resolve_model).
+ try:
+ if requested_model:
+ self._opencode_chat.model = requested_model
+ self.opencode_model = requested_model
+ if self._opencode_chat.publish_model(requested_model):
+ # Shared, not pinned: later turns keep re-reading the
+ # config, which now names this model -- so a studio
+ # pick after this still wins, as it should.
+ self._opencode_model_source = "user-published"
+ else:
+ # The config would not take it (read-only, older
+ # server). Pin it for this process so the switch the
+ # user asked for still takes effect.
+ self._opencode_chat.model_is_explicit = True
+ self._opencode_model_explicit = True
+ self._opencode_model_source = "user-pinned"
+ else:
+ resolved, source = self._opencode_chat.resolve_model(publish_default=True)
+ if resolved:
+ self.opencode_model = resolved
+ self._opencode_model_source = source
+ # _ensure_reachable may have moved us to a discovered server.
+ self.opencode_url = self._opencode_chat.base_url
+ except Exception as exc: # noqa: BLE001 -- server down; lazy path retries
+ _LOGGER.debug("[Agent] deferred OpenCode model resolution: %s", exc)
+ self._opencode_model_source = "unresolved"
_LOGGER.info(
f"[Agent] OpenCode enabled: {self.opencode_url} "
- f"(model={self.opencode_model or 'server default'})"
+ f"(model={self.opencode_model or 'unresolved'})"
)
except Exception as e:
_LOGGER.error(f"OpenCode error: {e}")
+ self._log_agent_configuration()
return initialized
+ # Human-readable provenance for the start-up banner.
+ _MODEL_SOURCE_LABELS = {
+ "pinned": "pinned by OPENCODE_MODEL",
+ "pinned-published": ("pinned by OPENCODE_MODEL, and published to "
+ "OpenCode's config so the studio shows it"),
+ "config-seed": ("from agent_config.yaml's opencode_model; nothing was "
+ "chosen in OpenCode's config yet"),
+ "config-seed-published": ("from agent_config.yaml's opencode_model "
+ "(nothing chosen yet), published to OpenCode's "
+ "config so the studio shows it"),
+ "opencode-config": "from OpenCode's config, which the studio model picker writes",
+ "default": "built-in default; OpenCode's config could not be updated",
+ "default-published": "built-in default, published to OpenCode's config for the studio",
+ "user-published": "chosen here and published to OpenCode's config",
+ "user-pinned": "chosen here; OpenCode's config refused the write, pinned to this backend",
+ "kept": "kept from this session",
+ "unresolved": "unresolved -- OpenCode unreachable, retried on the first query",
+ }
+
+ def _log_agent_configuration(self) -> None:
+ """Start-up banner: the model actually in use, and where it came from."""
+ source = getattr(self, "_opencode_model_source", "unresolved")
+ detail = self._MODEL_SOURCE_LABELS.get(source, source)
+ path = getattr(self, "_config_source_path", "(no agent config file found)")
+ _LOGGER.info(
+ "" + "\n" +
+ "\n# #######################################" + "\n" +
+ "# #######################################" + "\n" +
+ f"Agent initialized from configuration {path}: " + "\n" +
+ f"\tOpenCode URL={self.opencode_url}" + "\n" +
+ f"\tModel={self.opencode_model or '(unresolved)'} ({detail})" + "\n" +
+ "# #######################################" + "\n" +
+ "# #######################################" + "\n" + ""
+ )
+
def is_available(self) -> bool:
"""Return True if the OpenCode provider is ready to serve requests."""
return self.chain_opencode is not None
@@ -843,10 +927,12 @@ def initialize_with_cloud_key(self, api_key: str, provider: str, model: Optional
if model is not None and not model.strip():
return False, "Model cannot be empty."
- self.opencode_model = model.strip() if model and model.strip() else self.opencode_model
+ requested = model.strip() if model and model.strip() else None
+ if requested:
+ self.opencode_model = requested
self.preferred_provider = "opencode"
- success = self._setup_providers()
+ success = self._setup_providers(requested_model=requested)
if self.chain_opencode is None or not success:
return False, "Could not reach the OpenCode server. Please verify OPENCODE_URL and that it is running."
@@ -865,11 +951,56 @@ def change_model(self, model: str) -> "tuple[bool, str]":
if not model or not model.strip():
return False, "Model cannot be empty."
- self.opencode_model = model.strip()
- success = self._setup_providers()
+ requested = model.strip()
+ self.opencode_model = requested
+ success = self._setup_providers(requested_model=requested)
if self.chain_opencode is None or not success:
return False, "Could not reach the OpenCode server. Please verify OPENCODE_URL and that it is running."
- return True, f"Model switched to {self.opencode_model}. Ready to help you."
+ if self.opencode_model != requested:
+ # Never report a switch that did not happen.
+ return False, (f"Could not switch to {requested}: the model in use is "
+ f"{self.opencode_model}.")
+ shared = self._opencode_model_source == "user-published"
+ return True, (
+ f"Model switched to {self.opencode_model}. "
+ + ("Published to OpenCode's config, so the studio picker shows it too. "
+ if shared else
+ "OpenCode's config would not take it, so it is pinned to this backend only. ")
+ + "Ready to help you."
+ )
+
+ def current_model(self) -> Optional[str]:
+ """The model the NEXT query will actually use.
+
+ `self.opencode_model` is only what was resolved when the provider was
+ last (re)initialised. A model chosen in the studio afterwards lands in
+ OpenCode's config, which every turn re-reads -- so reporting the
+ snapshot made a UI pick look ignored (`agent status` kept naming the
+ old model while queries already used the new one). Re-resolves here:
+ one local GET, on a command the user typed.
+ """
+ chat = self._opencode_chat
+ if chat is None:
+ return self.opencode_model or None
+ try:
+ model, source = chat.resolve_model()
+ if model:
+ self.opencode_model = model
+ self._opencode_model_source = source
+ if source == "pinned":
+ # A pin wins for this backend, so say plainly when the studio
+ # is showing something else -- otherwise the two surfaces
+ # disagree with no explanation anywhere.
+ chosen = chat._configured_model()
+ if chosen and chosen != self.opencode_model:
+ _LOGGER.warning(
+ "[Agent] OpenCode's configured model is %s (the studio's "
+ "pick), but this backend is pinned to %s by "
+ "OPENCODE_MODEL. Unset that variable to follow the "
+ "picker.", chosen, self.opencode_model)
+ except Exception as exc: # noqa: BLE001 -- report the last known model
+ _LOGGER.debug("[Agent] current_model could not re-resolve: %s", exc)
+ return self.opencode_model or None
def _opencode_base_url(self) -> str:
"""Same self-heal `OpenCodeChat._ensure_reachable` gives every chat
diff --git a/weightslab/trainer/services/agent/opencode_chat.py b/weightslab/trainer/services/agent/opencode_chat.py
index e67581a4..139f0abf 100644
--- a/weightslab/trainer/services/agent/opencode_chat.py
+++ b/weightslab/trainer/services/agent/opencode_chat.py
@@ -64,7 +64,7 @@
# unset -- which is exactly the "OpenCode picks WHATEVER model happens to be
# configured, arbitrarily" failure this method exists to avoid in the first
# place.
-_DEFAULT_MODEL = "opencode/deepseek-v4-flash-free"
+_DEFAULT_MODEL = "opencode/big-pickle"
class OpenCodeError(RuntimeError):
@@ -79,7 +79,7 @@ class OpenCodeChat:
def __init__(self, base_url: str, model: Optional[str] = None, timeout: float = 60.0,
workspace_dir: Optional[str] = None, url_is_explicit: bool = True,
- model_is_explicit: bool = True):
+ model_is_explicit: bool = True, seed_model: Optional[str] = None):
self.base_url = (base_url or "http://127.0.0.1:4096").rstrip("/")
self.model = model
self.timeout = timeout
@@ -105,6 +105,13 @@ def __init__(self, base_url: str, model: Optional[str] = None, timeout: float =
# (confirmed live: an image-generation preview model, useless for
# this class's structured-JSON-reply use case).
self.model_is_explicit = model_is_explicit
+ # A model from agent_config.yaml SEEDS the shared choice rather than
+ # pinning it: it is what to use when nobody has chosen anything yet
+ # (and it is then published, so the studio shows it), but a model
+ # picked in the UI afterwards wins. It used to pin, so a run started
+ # after picking a model in the studio quietly went back to the yaml
+ # value. OPENCODE_MODEL stays a hard pin -- automation needs one.
+ self.seed_model = (seed_model or "").strip() or None
# -- wire helpers --------------------------------------------------- #
@@ -317,7 +324,11 @@ def _ensure_model_resolved(self) -> None:
this class's structured-JSON intent-parsing, since it isn't a
text-reasoning model at all).
+ Returns a short source label ("pinned" | "opencode-config" |
+ "config-seed" | "default" | "kept") for the banner/logs.
+
Resolution order:
+ 0. `OPENCODE_MODEL` (model_is_explicit) -- a hard pin, left alone.
1. `GET /config`'s own `model` field -- the one the model picker
writes back to opencode.json on every pick (opencodeClient.ts's
setDefaultModel), so it's "whatever the user last actually
@@ -337,22 +348,116 @@ def _ensure_model_resolved(self) -> None:
explicitly chosen, land on the known-good free model" now means
exactly that, with no provider-reported default able to override it.
- Resolved once and cached on self.model. An explicit model
- (model_is_explicit=True) is left alone -- deliberately chosen, not
- a placeholder to override.
+ Re-checked before every turn, NOT cached for the life of the
+ process: the studio's model picker (.wl-ag-model) writes the pick
+ into OpenCode's own config via `PUT /config`, and a backend that
+ latched onto the model it saw at startup went on answering with the
+ old one for the rest of the run. `GET /config` is a local request on
+ the same machine, so following it per turn costs nothing next to the
+ completion it precedes. A failed read keeps the current model instead
+ of falling back.
+
+ An explicit model (model_is_explicit=True -- OPENCODE_MODEL or
+ agent_config.yaml's `opencode_model`) is left alone: deliberately
+ chosen, not a placeholder to override, and NOT overridable from the
+ UI picker either.
"""
- if self.model_is_explicit or self.model:
- return
+ if self.model_is_explicit:
+ return "pinned"
+ model_id = self._configured_model()
+ if model_id:
+ if model_id != self.model:
+ _LOGGER.info("[OpenCodeChat] following OpenCode's configured "
+ "model: %s (was %s)", model_id, self.model or "unset")
+ self.model = model_id
+ return "opencode-config"
+ # Nothing chosen anywhere yet -- fall back to the configured seed
+ # before the built-in default, so a project's agent_config.yaml still
+ # decides which model a fresh setup starts on.
+ if self.seed_model:
+ self.model = self.seed_model
+ return "config-seed"
+ # /config could not be read (or names no model): keep whatever was
+ # resolved on an earlier turn rather than dropping a working model for
+ # the fallback because one local request happened to fail.
+ if not self.model:
+ self.model = _DEFAULT_MODEL
+ return "default"
+ return "kept"
+
+ def publish_model(self, model: Optional[str] = None) -> bool:
+ """Write `model` (default: the resolved one) into OpenCode's own config.
+
+ Same call the studio's model picker makes, so whichever side starts
+ first leaves ONE answer behind for the other to read, and both ends of
+ a session agree on the model without talking to each other.
+
+ GLOBAL scope, with the workspace route as fallback: verified live,
+ PATCH /config echoes the value back but does NOT change what GET
+ /config then reports, while PATCH /global/config
+ (`global.config.update`) does, immediately -- and GET /config is what
+ both sides read.
+
+ Best-effort: a read-only config or an older server must never stop the
+ agent from working with the model it already resolved in-process.
+ """
+ model = model or self.model
+ if not model:
+ return False
+ for path in ("/global/config", "/config"):
+ try:
+ with self._request(path, method="PATCH", body={"model": model}) as resp:
+ resp.read()
+ except Exception as exc: # noqa: BLE001 -- advisory write, never fatal
+ _LOGGER.debug("[OpenCodeChat] could not publish model %s via %s: %s",
+ model, path, exc)
+ continue
+ # Confirm against the effective config rather than trusting the
+ # echo: the workspace route answers 200 for a write it drops.
+ if self._configured_model() == model:
+ _LOGGER.info("[OpenCodeChat] published model %s to OpenCode (%s) at %s",
+ model, path, self.base_url)
+ return True
+ _LOGGER.debug("[OpenCodeChat] model %s could not be published to %s",
+ model, self.base_url)
+ return False
+
+ def _configured_model(self) -> Optional[str]:
+ """The model OpenCode itself reports (GET /config) -- the shared choice
+ the studio picker writes and this backend follows."""
try:
with self._request("/config") as resp:
config = json.loads(resp.read().decode("utf-8"))
- model_id = (config or {}).get("model")
- if isinstance(model_id, str) and "/" in model_id:
- self.model = model_id
- return
- except Exception: # noqa: BLE001 - fall through to the hardcoded default
- pass
- self.model = _DEFAULT_MODEL
+ except Exception: # noqa: BLE001
+ return None
+ model_id = (config or {}).get("model")
+ return model_id if isinstance(model_id, str) and "/" in model_id else None
+
+ def resolve_model(self, publish_default: bool = False):
+ """Resolve the model NOW instead of lazily on the first turn, and say
+ where it came from: ("pinned" | "pinned-published" | "opencode-config"
+ | "default" | "default-published" | "kept").
+
+ Called at agent start-up so the banner states the model actually in
+ use -- it used to print "(server default)" whenever nothing was pinned,
+ which read as "the studio's choice was ignored" even when the first
+ turn would have picked it up correctly.
+
+ With publish_default=True, the model is also written back to
+ OpenCode's config whenever this side is the one deciding it -- a pinned
+ model (OPENCODE_MODEL), a seed from agent_config.yaml, or the built-in
+ fallback. A backend started BEFORE the studio then hands the UI the
+ model it is itself using, instead of the studio showing an unrelated
+ default while every backend query ran on the pinned one.
+
+ A model that CAME from OpenCode's config is never re-published: there
+ is nothing to write, and doing so would fight the picker.
+ """
+ self._ensure_reachable()
+ source = self._ensure_model_resolved()
+ if publish_default and source in ("default", "config-seed", "pinned") and self.publish_model():
+ source = f"{source}-published"
+ return self.model, source
def _call(self, prompt_value):
from langchain_core.messages import AIMessage
diff --git a/weightslab/trainer/services/data_service.py b/weightslab/trainer/services/data_service.py
index dd065d33..a2b0b477 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,
@@ -146,6 +146,47 @@ def _media_chunk_bytes() -> int:
_NON_MASK_TASKS = ("classification", "tabular")
+def is_set_flag(value) -> bool:
+ """True only for a boolean column value that is really SET and true.
+
+ Two traps, both hit in production, both of them "a missing value read as
+ set":
+
+ * ``bool(float("nan")) is True`` in Python, and a nullable flag column is
+ full of NaN -- ``discarded`` for a sample nothing has written yet, or a
+ ``tag:*`` column for every sample that does not carry the tag. Reading it
+ with ``bool()`` / ``astype(bool)`` reported all of them as set: samples
+ greyed out as discarded while the dataframe said False, and every sample
+ wearing every tag.
+ * ``bool("False") is True`` as well, and a boolean column that has been
+ through the H5 store (where these columns become categorical) can come
+ back holding the STRINGS "True"/"False".
+
+ So: missing is false, a string is read as a word, everything else falls
+ back to ``bool()``.
+ """
+ try:
+ if value is None or pd.isna(value):
+ return False
+ except (TypeError, ValueError):
+ # pd.isna raises for some array-likes; those are not missing.
+ pass
+ if isinstance(value, str):
+ return value.strip().lower() in ("1", "true", "yes", "y", "t")
+ try:
+ return bool(value)
+ except Exception: # noqa: BLE001 -- an exotic value is not a set flag
+ return False
+
+
+def set_flag_mask(series) -> "np.ndarray":
+ """``is_set_flag`` over a whole column, as a numpy bool array."""
+ if series is None:
+ return np.zeros(0, dtype=bool)
+ return np.fromiter((is_set_flag(v) for v in series.tolist()),
+ dtype=bool, count=len(series))
+
+
def _is_non_mask_task(task_type) -> bool:
"""True when labels/predictions for this task must not be read as masks."""
return task_type in _NON_MASK_TASKS or is_generation_task(task_type)
@@ -494,6 +535,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 +1167,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 +1206,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 +1525,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 +1555,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
@@ -1569,10 +1649,10 @@ def _process_sample_row(self, args):
# 'discarded' drives the grayed-out cell rendering, so it rides with the
# image data as "1"/"0" (not treated as analytical metadata). This keeps
# the gray-out reliable on every grid (re)fetch / scroll.
- try:
- _discarded_str = "1" if bool(row.get(SampleStatsEx.DISCARDED.value)) else "0"
- except Exception:
- _discarded_str = "0"
+ # is_set_flag, not bool(): a nullable flag full of NaN read as
+ # set is what greyed out every sample the model had seen.
+ _discarded_str = "1" if is_set_flag(
+ row.get(SampleStatsEx.DISCARDED.value)) else "0"
data_stats.append(
create_data_stat(
SampleStatsEx.DISCARDED.value, 'string', shape=[1], value_string=_discarded_str, thumbnail=b""
@@ -2150,14 +2230,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 +2627,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 +3717,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 +4011,137 @@ 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
+ # Collapse the source's per-annotation rows to ONE row per sample the
+ # same way the view itself was built (see
+ # get_collapse_annotations_to_samples_df): the canonical row is
+ # annotation_id == 0, and sample-level columns live only there.
+ #
+ # This used to keep the LAST annotation row, which for a multi-instance
+ # sample carries NaN in every sample-level column -- so each sample the
+ # trainer touched had its view row's `discarded` / `prediction` /
+ # `target` overwritten with NaN. And `bool(float("nan"))` is True in
+ # Python, so GetDataSamples then reported discarded="1" and the studio
+ # greyed the sample out, progressively, exactly as the model worked
+ # through the dataset -- while the dataframe itself still said False.
+ if isinstance(sub.index, pd.MultiIndex):
+ ANNOT = SampleStatsEx.INSTANCE_ID.value
+ names = list(getattr(sub.index, "names", []) or [])
+ if ANNOT in names:
+ annot = sub.index.get_level_values(ANNOT)
+ try:
+ canonical = np.asarray(annot).astype(int) == 0
+ except (TypeError, ValueError):
+ canonical = np.array([str(a) in ("0", "0.0") for a in annot])
+ if canonical.any():
+ sub = sub[canonical]
+ sub = sub.droplevel(-1)
+ # Whatever is left, one row per sample: prefer the FIRST (the canonical
+ # row when the level was present, the first occurrence otherwise) --
+ # never the last, for the reason above.
+ sub = sub[~sub.index.duplicated(keep="first")]
+
+ # 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)
+ # Looked up the other way round -- `sub` (deduplicated just above, so
+ # unique) is the index being searched, and the VIEW's keys are the
+ # target. Searching the view's keys instead raised
+ # InvalidIndexError("Reindexing only valid with uniquely valued Index
+ # objects") whenever one sample_id appeared under two origins, which
+ # the view's own (origin, sample_id) index exists precisely to allow --
+ # and the differential refresh then failed every time, silently falling
+ # back to the full rebuild. This direction also updates BOTH rows of
+ # such a sample, which is the only thing the source (indexed by
+ # sample_id alone) can mean.
+ _view_keys = pd.Index(view_keys.astype(str))
+ _sub_keys = pd.Index(sub.index.astype(str))
+ _src = _sub_keys.get_indexer(_view_keys) # sub row per view row, -1 if none
+ _rows = np.flatnonzero(_src >= 0)
+ if _rows.size == 0:
+ return True
+ # A dirty sample the view does not hold is a structural change (a new
+ # sample): only the full rebuild can add it.
+ if len(set(_sub_keys)) != len(set(_view_keys[_rows])):
+ return False
+ _take = _src[_rows]
+ for c in sub.columns:
+ _ci = view.columns.get_loc(c)
+ view.iloc[_rows, _ci] = sub[c].to_numpy()[_take]
+ 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 +4312,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
@@ -4180,8 +4534,11 @@ def _build_metadata_only_response(self, df_slice: pd.DataFrame, requested_cols=N
_DataStat(name=col, type="string", shape=[1], value_string=v)
)
else:
- # Boolean tag: presence indicator "1" when True.
- bools = series.astype(bool).tolist()
+ # Boolean tag: presence indicator "1" when True. is_set_flag,
+ # not astype(bool): a tag column is NaN for every sample that
+ # does not carry the tag, and astype(bool) turns NaN into True
+ # -- which showed every tag on every sample.
+ bools = set_flag_mask(series)
for i, b in enumerate(bools):
if b:
row_stats[i].append(
@@ -4482,7 +4839,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 +5151,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 +5407,38 @@ 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")
+ # Same trap as above: NaN in a nullable flag is NOT "discarded".
+ disc = set_flag_mask(_d) 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 +5449,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 +5470,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 +5520,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 +6039,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 +6208,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 +6247,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 +6271,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 +6281,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 +6496,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/ui/server.py b/weightslab/ui/server.py
index 53c7eec7..0d7fed6b 100644
--- a/weightslab/ui/server.py
+++ b/weightslab/ui/server.py
@@ -1577,6 +1577,9 @@ class _UIRequestHandler(BaseHTTPRequestHandler):
grpc_auth_token: Optional[str] = None
rpc_timeout: float = 300.0
experiment_dir: Optional[str] = None
+ # The gRPC port this server proxies to; used to pick the right backend's
+ # experiment directory out of the active-experiment marker.
+ backend_port: Optional[int] = None
# -- logging: quiet by default, honour WEIGHTSLAB_UI_VERBOSE ------------- #
def log_message(self, fmt, *args): # noqa: D401
@@ -1612,6 +1615,9 @@ def do_GET(self): # noqa: N802
if path == "/agent-server/status":
self._send_json(HTTPStatus.OK, _opencode_session.status())
return
+ if path == "/agent-server/model":
+ self._get_shared_model()
+ return
if path == "/agent-server/loop/list":
self._send_json(HTTPStatus.OK, {"loops": _loop_registry.list()})
return
@@ -1645,6 +1651,8 @@ def do_POST(self): # noqa: N802
self._start_local_notebook()
elif path == "/agent-server/start":
self._start_agent_server()
+ elif path == "/agent-server/model":
+ self._set_shared_model()
elif path == "/agent-server/loop/start":
self._start_loop()
elif path.startswith("/agent-server/loop/") and path.endswith("/stop"):
@@ -1745,9 +1753,120 @@ def _collect_metadata(self):
# ------------------------------------------------------------------ #
# Local Jupyter Notebook launcher (landing-page button)
# ------------------------------------------------------------------ #
+ # ------------------------------------------------------------------ #
+ # Shared agent model, proxied SAME-ORIGIN
+ # ------------------------------------------------------------------ #
+ # OpenCode's own config holds the one model every client of that server
+ # agrees on: the studio's picker, the OpenCode CLI, `weightslab agent
+ # model`, and the backend SDK agent (which reads it to choose the model for
+ # its own queries). The browser could call OpenCode directly -- but only
+ # while OpenCode's --cors allowlist happens to contain the exact origin the
+ # page is served from, which quietly fails for a LAN address, a tunnel
+ # hostname, or an OpenCode somebody started by hand with no --cors at all.
+ # The page then showed a model nobody else was using, and picking one
+ # changed nothing outside the tab.
+ #
+ # Same-origin here means no preflight and no allowlist: this server talks
+ # to OpenCode over plain HTTP on the machine they share.
+ def _opencode_base_url(self) -> Optional[str]:
+ status = _opencode_session.status()
+ url = status.get("url") if isinstance(status, dict) else None
+ if isinstance(url, str) and url.strip():
+ return url.rstrip("/")
+ env_url = (os.environ.get("OPENCODE_URL") or "").strip()
+ return env_url.rstrip("/") if env_url else None
+
+ def _opencode_json(self, path: str, method: str = "GET", body: Optional[dict] = None,
+ timeout: float = 10.0):
+ """One request to the local OpenCode server; (status, parsed-json-or-None)."""
+ base = self._opencode_base_url()
+ if not base:
+ return None, None
+ import urllib.error
+ import urllib.request
+ data = json.dumps(body).encode("utf-8") if body is not None else None
+ headers = {"Content-Type": "application/json"} if data is not None else {}
+ req = urllib.request.Request(base + path, data=data, headers=headers, method=method)
+ try:
+ with urllib.request.urlopen(req, timeout=timeout) as resp:
+ raw = resp.read().decode("utf-8", "replace")
+ try:
+ return resp.status, json.loads(raw or "null")
+ except ValueError:
+ return resp.status, None
+ except urllib.error.HTTPError as exc:
+ return exc.code, None
+ except Exception as exc: # noqa: BLE001
+ logger.debug("[ui] OpenCode %s %s failed: %s", method, path, exc)
+ return None, None
+
+ def _get_shared_model(self):
+ status, payload = self._opencode_json("/config")
+ # A transport failure AND an error status both mean "we do not know
+ # which model is configured" -- reporting ok:True with model=None there
+ # would tell the page "nothing is configured", and it would then show
+ # its own default over whatever the backend is really using.
+ if status is None or not (200 <= int(status) < 300):
+ self._send_json(HTTPStatus.OK, {"ok": False, "model": None,
+ "error": "OpenCode server is not reachable"})
+ return
+ model = (payload or {}).get("model") if isinstance(payload, dict) else None
+ if not (isinstance(model, str) and "/" in model):
+ model = None
+ self._send_json(HTTPStatus.OK, {"ok": True, "model": model})
+
+ def _set_shared_model(self):
+ body = self._read_json_body()
+ model = str((body or {}).get("model") or "").strip()
+ if "/" not in model:
+ self._send_json(HTTPStatus.BAD_REQUEST,
+ {"ok": False, "error": "model must be \"providerID/modelID\""})
+ return
+ # Global scope first: PATCH /config (workspace scope) answers 200 and
+ # echoes the value back without changing what GET /config reports.
+ for target in ("/global/config", "/config"):
+ status, _ = self._opencode_json(target, method="PATCH", body={"model": model})
+ if status is None or not (200 <= int(status) < 300):
+ continue
+ _, payload = self._opencode_json("/config")
+ if isinstance(payload, dict) and payload.get("model") == model:
+ self._send_json(HTTPStatus.OK, {"ok": True, "model": model, "via": target})
+ return
+ self._send_json(HTTPStatus.OK, {
+ "ok": False, "model": None,
+ "error": "OpenCode did not accept the model (config may be read-only)"})
+
+ def _experiment_dir_path(self) -> str:
+ """The experiment directory to browse for this run's own files.
+
+ A RUNNING training backend's own resolved root_log_dir comes first: it
+ is where reports and notebooks are actually written, and it is not
+ always the directory this UI established -- training may have been
+ pointed elsewhere by a config file's `root_log_dir:`, or (before the
+ marker existed) have fallen through to a temp directory. Listing this
+ server's own directory then showed an empty reports/ right after a
+ report had been generated.
+
+ Only a LIVE backend counts: the marker outlives the process that wrote
+ it, and a finished run's directory must never hijack the listing of a
+ UI that was given an experiment directory of its own. Falls back to the
+ UI's own directory, the environment, then the working directory.
+ """
+ try:
+ from weightslab.utils.active_experiment import live_backend_experiment_dir
+ # By PORT: with two experiments up, "the live backend" is ambiguous
+ # and picking the wrong one shows the other experiment's reports.
+ backend_dir = live_backend_experiment_dir(
+ getattr(self, "backend_port", None))
+ except Exception: # noqa: BLE001 -- never break a listing on the marker
+ backend_dir = None
+ return (backend_dir
+ or self.experiment_dir
+ or os.environ.get("WEIGHTSLAB_ROOT_LOG_DIR")
+ or os.getcwd())
+
def _notebooks_dir_path(self) -> str:
- experiment_dir = self.experiment_dir or os.environ.get("WEIGHTSLAB_ROOT_LOG_DIR") or os.getcwd()
- return os.path.join(experiment_dir, "notebooks")
+ return os.path.join(self._experiment_dir_path(), "notebooks")
def _read_json_body(self) -> dict:
try:
@@ -2088,13 +2207,12 @@ def _start_local_notebook(self):
# (ApplyDataQuery -> the "generate_experiment_report" action, see
# data_service.py) -- these two endpoints only browse what's already on
# disk under /reports/, exactly like the local-notebook
- # endpoints above browse /notebooks/. Same assumption:
- # this UI server and the connected training backend share a filesystem
- # (the documented `weightslab start` usage), so root_log_dir resolved
- # here is the same directory the backend wrote reports into.
+ # endpoints above browse /notebooks/. The one assumption
+ # left is a shared filesystem (the documented `weightslab start` usage):
+ # WHICH directory is resolved by _experiment_dir_path(), which prefers the
+ # backend's own recorded root_log_dir over this server's.
def _agent_history_dir_path(self) -> str:
- experiment_dir = self.experiment_dir or os.environ.get("WEIGHTSLAB_ROOT_LOG_DIR") or os.getcwd()
- return os.path.join(experiment_dir, "agent")
+ return os.path.join(self._experiment_dir_path(), "agent")
def _dump_agent_history(self):
"""Write the agent conversation to the experiment directory.
@@ -2138,8 +2256,7 @@ def _dump_agent_history(self):
self._send_json(HTTPStatus.OK, {"ok": True, "path": path})
def _reports_dir_path(self) -> str:
- experiment_dir = self.experiment_dir or os.environ.get("WEIGHTSLAB_ROOT_LOG_DIR") or os.getcwd()
- return os.path.join(experiment_dir, "reports")
+ return os.path.join(self._experiment_dir_path(), "reports")
def _list_experiment_reports(self):
reports_dir = self._reports_dir_path()
@@ -2456,6 +2573,7 @@ def serve_ui(
{
"static_root": root,
"channel": channel,
+ "backend_port": backend_port,
"api_prefix": "/api",
"grpc_auth_token": grpc_auth_token,
"experiment_dir": experiment_dir,
diff --git a/weightslab/utils/active_experiment.py b/weightslab/utils/active_experiment.py
new file mode 100644
index 00000000..fb53f9b2
--- /dev/null
+++ b/weightslab/utils/active_experiment.py
@@ -0,0 +1,340 @@
+"""Cross-process handoff of the active experiment directory.
+
+``weightslab start`` establishes the experiment directory (checkpoints, logs,
+``notebooks/``, ``reports/``) and exports ``WEIGHTSLAB_ROOT_LOG_DIR`` -- but it
+can only export it into ITS OWN process. A training run launched from a second
+terminal, or by ``weightslab start example``, is a different process tree and
+never saw that variable: it fell through to a throwaway ``tempfile.mkdtemp()``,
+so the run wrote into ``%TEMP%\\tmpXXXXXXXX`` while the UI listed an empty
+``reports/`` (the "right-click Generate Report shows nothing, yet I generated
+reports" symptom) from the directory it had established itself.
+
+Hence this marker file: one small JSON document, per user, that both sides
+write to and read from. Each side is a LIST, because running two experiments
+side by side -- a classification UI on one port, a segmentation UI on another --
+is a supported thing to do. With a single slot per side, the second
+``weightslab start`` erased the first, and both UIs would then have listed the
+reports of whichever backend started last.
+
+ {
+ "ui": [{"root_log_dir": "...", "pid": 123, "ui_port": 8080,
+ "backend_port": 50051, "updated_at": "..."}, ...],
+ "backend": [{"root_log_dir": "...", "pid": 456, "grpc_port": 50051,
+ "updated_at": "..."}, ...]
+ }
+
+* ``ui`` is written by ``weightslab start`` -- the directory it established.
+ A later training process with nothing configured adopts it, which is what
+ makes the two halves land in the same experiment.
+* ``backend`` is written by ``wl.serve()`` -- the directory training ACTUALLY
+ resolved, whatever the route (an explicit ``root_log_dir:`` in a config file,
+ the environment, or the ``ui`` value above). The UI prefers it when listing
+ reports and notebooks, so those lists stay right even when training was
+ pointed somewhere the UI never chose.
+
+Entries are keyed by pid: a process replaces its own, dead ones are pruned on
+every write, and the list is bounded. Readers never guess -- a port pins the
+entry when the caller knows one (a UI asking for ITS backend), a single live
+entry is unambiguous, and two or more return nothing with a line in the log,
+leaving the caller its own directory. Showing one experiment's reports inside
+another experiment's UI, or writing a run into the wrong experiment, is worse
+than declining to answer.
+
+Neither side is required: every reader validates that the recorded directory
+still exists and falls back to its previous behaviour otherwise, and every
+write is best-effort -- a read-only home directory must never stop a run.
+
+Set ``WEIGHTSLAB_STATE_DIR`` to relocate the file (tests use it to stay out of
+the developer's real state).
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import tempfile
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+_FILE_NAME = "active_experiment.json"
+_SECTIONS = ("ui", "backend")
+# A long-lived machine should not accumulate entries nobody can attribute.
+_MAX_ENTRIES = 16
+
+
+def state_dir() -> Path:
+ """Directory holding the marker: ``$WEIGHTSLAB_STATE_DIR`` or ``~/.weightslab``."""
+ override = (os.environ.get("WEIGHTSLAB_STATE_DIR") or "").strip()
+ if override:
+ return Path(override).expanduser()
+ return Path.home() / ".weightslab"
+
+
+def state_path() -> Path:
+ """Absolute path of the marker file (may not exist yet)."""
+ return state_dir() / _FILE_NAME
+
+
+def read_state() -> dict:
+ """The whole marker, or ``{}`` when it is absent or unreadable."""
+ path = state_path()
+ try:
+ with open(path, "r", encoding="utf-8") as fh:
+ data = json.load(fh)
+ except FileNotFoundError:
+ return {}
+ except Exception as exc: # noqa: BLE001 -- a corrupt marker is not fatal
+ logger.debug("[active-experiment] could not read %s: %s", path, exc)
+ return {}
+ return data if isinstance(data, dict) else {}
+
+
+def entries(section: str, state: Optional[dict] = None) -> list:
+ """One section as a list, accepting the older single-object shape."""
+ value = (read_state() if state is None else state).get(section)
+ if isinstance(value, list):
+ return [e for e in value if isinstance(e, dict)]
+ if isinstance(value, dict):
+ return [value]
+ return []
+
+
+def _pid_is_running(pid) -> bool:
+ try:
+ pid = int(pid)
+ except (TypeError, ValueError):
+ return False
+ if pid <= 0:
+ return False
+ try:
+ import psutil
+ return psutil.pid_exists(pid)
+ except Exception: # noqa: BLE001 -- psutil missing/unusable: assume gone
+ return False
+
+
+class _MarkerLock:
+ """Brief exclusive hold on the marker, via an O_EXCL lock file.
+
+ Read-modify-write on a shared file loses updates when two processes do it
+ at once, and two processes doing it at once is precisely the case this
+ marker exists for: starting a classification UI and a segmentation UI
+ together dropped one of their port stamps, which is the field that tells
+ them apart afterwards.
+
+ Best-effort by design: if the lock cannot be taken (a stale file nobody
+ cleaned up, a read-only home), the write proceeds anyway -- an advisory
+ record must never block a run. A stale lock older than a few seconds is
+ broken on purpose, since nothing here holds it for more than a file write.
+ """
+
+ STALE_AFTER = 5.0
+
+ def __init__(self, path: Path, attempts: int = 60, delay: float = 0.02):
+ self._path = Path(str(path) + ".lock")
+ self._attempts = attempts
+ self._delay = delay
+ self._held = False
+
+ def __enter__(self):
+ import time
+ for _ in range(self._attempts):
+ try:
+ fd = os.open(str(self._path), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
+ os.write(fd, str(os.getpid()).encode())
+ os.close(fd)
+ self._held = True
+ return self
+ except FileExistsError:
+ try:
+ age = time.time() - os.path.getmtime(self._path)
+ if age > self.STALE_AFTER:
+ os.unlink(self._path)
+ continue
+ except OSError:
+ pass
+ time.sleep(self._delay)
+ except OSError:
+ break # cannot lock here at all; proceed unlocked
+ return self
+
+ def __exit__(self, *exc):
+ if self._held:
+ try:
+ os.unlink(self._path)
+ except OSError:
+ pass
+ return False
+
+
+def _write_section(section: str, root_log_dir, **meta) -> Optional[Path]:
+ """Record this process's entry in one section, keeping the others.
+
+ Replaces the entry for this pid (a process re-recording, e.g. once its port
+ is known), drops entries whose process is gone, and leaves every other live
+ entry in place -- that is what lets two experiments run side by side. Held
+ under _MarkerLock so two processes starting together cannot lose each
+ other's entry.
+ """
+ if section not in _SECTIONS:
+ raise ValueError(f"unknown section {section!r}")
+ if not root_log_dir:
+ return None
+
+ pid = os.getpid()
+ entry = {
+ "root_log_dir": str(Path(root_log_dir).expanduser().resolve()),
+ "pid": pid,
+ "updated_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+ }
+ entry.update({k: v for k, v in meta.items() if v is not None})
+
+ path = state_path()
+ try:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ with _MarkerLock(path):
+ # Re-read INSIDE the lock: another process may have added its own
+ # entry since this function started.
+ state = read_state()
+ kept = [e for e in entries(section, state)
+ if e.get("pid") != pid and _pid_is_running(e.get("pid"))]
+ state[section] = (kept + [entry])[-_MAX_ENTRIES:]
+ # Written via a temp file in the same directory, then replaced, so a
+ # concurrent reader never sees a half-written document.
+ fd, tmp_name = tempfile.mkstemp(dir=str(path.parent), prefix=".active-", suffix=".json")
+ try:
+ with os.fdopen(fd, "w", encoding="utf-8") as fh:
+ json.dump(state, fh, indent=2)
+ os.replace(tmp_name, path)
+ except Exception:
+ try:
+ os.unlink(tmp_name)
+ except OSError:
+ pass
+ raise
+ except Exception as exc: # noqa: BLE001 -- advisory record, never fatal
+ logger.debug("[active-experiment] could not record %s dir in %s: %s",
+ section, path, exc)
+ return None
+ logger.debug("[active-experiment] recorded %s root_log_dir=%s", section, entry["root_log_dir"])
+ return path
+
+
+def record_ui_experiment(root_log_dir, ui_port: Optional[int] = None,
+ backend_port: Optional[int] = None) -> Optional[Path]:
+ """Record the directory ``weightslab start`` just established.
+
+ ``backend_port`` is what makes two concurrent UIs distinguishable: it is
+ the gRPC port this UI proxies to, so it can later ask for the experiment
+ directory of ITS backend rather than of whichever one started last.
+ """
+ return _write_section("ui", root_log_dir, ui_port=ui_port, backend_port=backend_port)
+
+
+def record_backend_experiment(root_log_dir, grpc_port: Optional[int] = None) -> Optional[Path]:
+ """Record the directory training actually resolved (``wl.serve()``)."""
+ return _write_section("backend", root_log_dir, grpc_port=grpc_port)
+
+
+def _entry_dir(section: str, entry: dict) -> Optional[str]:
+ value = (entry or {}).get("root_log_dir")
+ if not isinstance(value, str) or not value:
+ return None
+ if not os.path.isdir(value):
+ # A run whose directory was deleted (or a marker copied between
+ # machines) must not redirect anything.
+ logger.debug("[active-experiment] %s dir %s no longer exists; ignoring", section, value)
+ return None
+ return value
+
+
+def _live_entries(section: str) -> list:
+ return [e for e in entries(section)
+ if _pid_is_running(e.get("pid")) and _entry_dir(section, e)]
+
+
+def _sole_live_dir(section: str, port_key: str = "", port: Optional[int] = None) -> Optional[str]:
+ """The directory of the ONE live entry that matches, or None.
+
+ With a port, an entry naming it wins outright -- that is how a UI finds ITS
+ backend rather than whichever backend started last. Otherwise a single live
+ entry is unambiguous and is used; two or more are not guessed between.
+ """
+ live = _live_entries(section)
+ if port is not None and port_key:
+ matching = [e for e in live if e.get(port_key) == port]
+ if matching:
+ # The same port twice can only be stale bookkeeping: newest wins.
+ return _entry_dir(section, matching[-1])
+ if len(live) == 1:
+ return _entry_dir(section, live[0])
+ if len(live) > 1:
+ logger.info(
+ "[active-experiment] %d live %s experiments recorded (%s); not "
+ "guessing between them -- name the directory explicitly "
+ "(WEIGHTSLAB_ROOT_LOG_DIR, or root_log_dir in the config).",
+ len(live), section,
+ ", ".join(str(e.get("root_log_dir")) for e in live))
+ return None
+
+
+def ui_experiment_dir() -> Optional[str]:
+ """Directory of the most recent recorded ``weightslab start``, live or not.
+
+ Raw record: it outlives the process that wrote it. Callers that REDIRECT a
+ run on this should use :func:`live_ui_experiment_dir` instead.
+ """
+ for entry in reversed(entries("ui")):
+ found = _entry_dir("ui", entry)
+ if found:
+ return found
+ return None
+
+
+def live_ui_experiment_dir() -> Optional[str]:
+ """Directory of a ``weightslab start`` that is *still running*.
+
+ The handoff exists for "the UI is up over there, put this run in its
+ experiment". A record left behind by a UI that has since exited must not
+ silently redirect an unrelated run months later -- which is exactly what
+ happened to this repo's own gRPC tests: they resolved into a previous
+ session's experiment directory and loaded ITS config. And with two UIs up
+ (two experiments side by side) there is no right answer to guess.
+ """
+ return _sole_live_dir("ui")
+
+
+def backend_experiment_dir() -> Optional[str]:
+ """Directory of the most recent recorded ``wl.serve()``, live or not."""
+ for entry in reversed(entries("backend")):
+ found = _entry_dir("backend", entry)
+ if found:
+ return found
+ return None
+
+
+def live_backend_experiment_dir(grpc_port: Optional[int] = None) -> Optional[str]:
+ """Directory of a backend that is *still running*.
+
+ Pass ``grpc_port`` -- the port the caller actually talks to -- and the
+ backend serving it is picked out by name. Without it, one live backend is
+ unambiguous and two are not guessed between: a UI showing the OTHER
+ experiment's reports is worse than a UI showing its own directory.
+
+ Dead entries are ignored: the record outlives the process that wrote it.
+ """
+ return _sole_live_dir("backend", "grpc_port", grpc_port)
+
+
+def clear() -> None:
+ """Remove the marker (best-effort). Used by tests and ``weightslab`` teardown."""
+ try:
+ state_path().unlink()
+ except FileNotFoundError:
+ pass
+ except Exception as exc: # noqa: BLE001
+ logger.debug("[active-experiment] could not clear marker: %s", exc)
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